mirror of
https://github.com/codewhale-hq/Codewhale.git
synced 2026-11-11 01:51:34 -10:00
fix(fleet): token usage is tracked, never enforced; the reserved report turn survives (#6189, #6277)
Slice 5 of the fleet rework (design authority #6298): the machinery measures and reports usage, and no token budget may kill a worker, shrink the report turn, or gate a spawn. Step-cap and wall-time stops remain, each with the reserved hand-back turn. Removed: token_budget from WorkerRuntimeProfile / ChildLaunchManifest / SubAgentTask / SpawnRequest / SubAgentSpawnOptions (+ workflow_budget_scope), EngineConfig + Config subagent_token_budget, MIN_SUBAGENT_SPAWN_TOKEN_RESERVE, AgentUsageBudgetScope, budget fields on AgentRunUsage / AgentCoordSummary, SubAgentManager budget-scope machinery, the token arm of budget_failure_reason, and the token reserve math in budget_handback.rs (fixed 8_192 allowance now). Fixed in the same slice, each found by lane verification: - request_report no longer refuses the reserved hand-back turn when the task step cap is fully spent (#6277): the stop message itself promises the remaining turn is reserved for hand-back. - The outer task backstop no longer races the inner deadline (it replaced a real receipt with a generic error); the authored task wall time is folded into worker_profile.wall_deadline_ms so inner stages, the model-visible budget block, and the receipt name one effective limit, and the wall-time stop causes now carry the "raise the inherited wall-time limit" guidance. - cleanup_inner regains the lineage retention the budget cleanup deleted: ancestors of live workers are pinned for continuation lookup and descendant control (lineage_ancestors / lineage_records_in_use). - Stale assertions that pinned token budgets updated (completion_usage, launch_receipt, worker_is_not_stranded, effective_limits key list). CI-lane evidence (whale-h100, x86_64, RUST_MIN_STACK=16MiB as scripts/dev-cargo.sh): - cargo check -p codewhale-tui --locked --all-targets => 0 errors (2 runs) - cargo test -p codewhale-tui --lib --locked subagent => "781 passed; 5 failed" default threads; "783 passed; 3 failed" at --test-threads=8 (the failing set shifts with parallelism => load flake in the real-time fixtures, not a regression). Locally the full budget_handback filter is green: "15 passed; 0 failed". - cargo clippy --workspace --all-targets --locked -- -D warnings -A clippy::uninlined_format_args -A clippy::too_many_arguments -A clippy::unnecessary_map_or => clean (this is CI's exact invocation). Pre-existing reds, unchanged by this slice and named for the record: delivery_tests::enforced_readonly_python_queries_sqlite_under_a_live_peer_write_claim (sandbox mounts --tmpfs /tmp, so a workspace under /tmp cannot be --chdir'd) and tests::launch_receipt::issue_5305_untethered_runtime_fails_closed_before_admission.
This commit is contained in:
1 parent
ab1420c51e
commit
a7a8bdb338
23 files changed
+189
-2096
No files matched your search
@@ -2497,11 +2497,6 @@ pub struct SubagentsConfig {
|
||||
/// execution bounded.
|
||||
#[serde(default, alias = "max_total", alias = "admission_limit")]
|
||||
pub max_admitted: Option<usize>,
|
||||
/// Optional aggregate token budget shared by a root `agent` run and its
|
||||
/// descendants. When unset or 0, sub-agents keep legacy unlimited spend
|
||||
/// behavior unless an individual `agent` call supplies a per-run override.
|
||||
#[serde(default)]
|
||||
pub token_budget: Option<u64>,
|
||||
/// Deprecated pre-v0.8.61 alias for `launch_concurrency`. Honored only
|
||||
/// when `launch_concurrency` is unset, so the new key always wins.
|
||||
#[serde(default, rename = "interactive_max_launch")]
|
||||
@@ -2587,8 +2582,6 @@ pub struct SubagentProviderConfig {
|
||||
#[serde(default, alias = "max_total", alias = "admission_limit")]
|
||||
pub max_admitted: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub token_budget: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub api_timeout_secs: Option<u64>,
|
||||
#[serde(default)]
|
||||
pub heartbeat_timeout_secs: Option<u64>,
|
||||
@@ -7876,28 +7869,6 @@ impl Config {
|
||||
.clamp(max_concurrent, MAX_SUBAGENT_ADMISSION)
|
||||
}
|
||||
|
||||
/// Optional aggregate token budget for each root `agent` run.
|
||||
///
|
||||
/// Reads `[subagents] token_budget`. `None` and `0` both mean unlimited,
|
||||
/// preserving legacy behavior until a budget is explicitly configured.
|
||||
#[must_use]
|
||||
pub fn subagent_token_budget(&self) -> Option<u64> {
|
||||
self.subagents
|
||||
.as_ref()
|
||||
.and_then(|cfg| cfg.token_budget)
|
||||
.filter(|budget| *budget > 0)
|
||||
}
|
||||
|
||||
/// Return the provider-specific aggregate token budget for each root
|
||||
/// `agent` run.
|
||||
#[must_use]
|
||||
pub fn subagent_token_budget_for_provider(&self, provider: ApiProvider) -> Option<u64> {
|
||||
self.subagent_provider_config(provider)
|
||||
.and_then(|cfg| cfg.token_budget)
|
||||
.or_else(|| self.subagents.as_ref().and_then(|cfg| cfg.token_budget))
|
||||
.filter(|budget| *budget > 0)
|
||||
}
|
||||
|
||||
/// Default per-child model-turn budget from `[subagents]
|
||||
/// default_max_steps`, applied when an `agent` start carries no explicit
|
||||
/// `max_steps` (#5324). `None` or `0` mean unbounded; a positive value is
|
||||
|
||||
@@ -3281,29 +3281,6 @@ fn structured_role_pins_require_a_typed_model_field() {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subagent_token_budget_is_optional_and_zero_disables() {
|
||||
assert_eq!(Config::default().subagent_token_budget(), None);
|
||||
|
||||
let disabled = Config {
|
||||
subagents: Some(SubagentsConfig {
|
||||
token_budget: Some(0),
|
||||
..SubagentsConfig::default()
|
||||
}),
|
||||
..Config::default()
|
||||
};
|
||||
assert_eq!(disabled.subagent_token_budget(), None);
|
||||
|
||||
let configured = Config {
|
||||
subagents: Some(SubagentsConfig {
|
||||
token_budget: Some(50_000),
|
||||
..SubagentsConfig::default()
|
||||
}),
|
||||
..Config::default()
|
||||
};
|
||||
assert_eq!(configured.subagent_token_budget(), Some(50_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn subagent_admission_limit_defaults_and_clamps() {
|
||||
assert_eq!(
|
||||
@@ -3357,7 +3334,6 @@ max_concurrent = 20
|
||||
launch_concurrency = 20
|
||||
max_admitted = 200
|
||||
max_depth = 6
|
||||
token_budget = 100000
|
||||
api_timeout_secs = 900
|
||||
heartbeat_timeout_secs = 1200
|
||||
|
||||
@@ -3366,7 +3342,6 @@ max_concurrent = 4
|
||||
launch_concurrency = 3
|
||||
max_admitted = 12
|
||||
max_depth = 2
|
||||
token_budget = 25000
|
||||
api_timeout_secs = 180
|
||||
heartbeat_timeout_secs = 240
|
||||
"#,
|
||||
@@ -3385,10 +3360,6 @@ heartbeat_timeout_secs = 240
|
||||
config.subagent_max_spawn_depth_for_provider(ApiProvider::Zai),
|
||||
2
|
||||
);
|
||||
assert_eq!(
|
||||
config.subagent_token_budget_for_provider(ApiProvider::Zai),
|
||||
Some(25_000)
|
||||
);
|
||||
assert_eq!(
|
||||
config.subagent_api_timeout_secs_for_provider(ApiProvider::Zai),
|
||||
180
|
||||
|
||||
@@ -360,10 +360,6 @@ pub struct EngineConfig {
|
||||
/// `SubAgentRuntime::max_spawn_depth`. Override via
|
||||
/// `[subagents] max_depth = N` in `~/.codewhale/config.toml`.
|
||||
pub max_spawn_depth: u32,
|
||||
/// Optional aggregate token budget for each root sub-agent run.
|
||||
/// Descendant agents inherit the root pool unless a child starts a new
|
||||
/// budget scope with an explicit per-call override.
|
||||
pub subagent_token_budget: Option<u64>,
|
||||
/// Per-domain network policy decider (#135). Shared across the session so
|
||||
/// session-scoped approvals (`/network allow <host>`) persist for the
|
||||
/// remainder of the run.
|
||||
@@ -572,7 +568,6 @@ impl Default for EngineConfig {
|
||||
plan_state: new_shared_plan_state(),
|
||||
goal_state: new_shared_goal_state(),
|
||||
max_spawn_depth: crate::tools::subagent::DEFAULT_MAX_SPAWN_DEPTH,
|
||||
subagent_token_budget: None,
|
||||
network_policy: None,
|
||||
snapshots_enabled: true,
|
||||
snapshots_max_workspace_bytes:
|
||||
@@ -1647,7 +1642,6 @@ impl Engine {
|
||||
config.max_admitted_subagents,
|
||||
config.subagent_heartbeat_timeout,
|
||||
config.launch_concurrency,
|
||||
config.subagent_token_budget,
|
||||
// #5324: per-child budget defaults are operator config, not
|
||||
// per-call schema fields.
|
||||
api_config.subagent_default_max_steps(),
|
||||
@@ -3043,7 +3037,6 @@ impl Engine {
|
||||
self.config.max_admitted_subagents,
|
||||
self.config.subagent_heartbeat_timeout,
|
||||
self.config.launch_concurrency,
|
||||
self.config.subagent_token_budget,
|
||||
)
|
||||
};
|
||||
let launch_note = if launch_gate_applied {
|
||||
|
||||
@@ -533,8 +533,6 @@ pub(crate) async fn run_exec_agent(
|
||||
} else {
|
||||
execution_config.subagent_max_spawn_depth_for_provider(effective_provider)
|
||||
},
|
||||
subagent_token_budget: execution_config
|
||||
.subagent_token_budget_for_provider(effective_provider),
|
||||
network_policy,
|
||||
snapshots_enabled: !fleet_authority_active && execution_config.snapshots_config().enabled,
|
||||
snapshots_max_workspace_bytes: execution_config
|
||||
|
||||
@@ -468,10 +468,6 @@ pub fn fleet_task_to_worker_spec_with_profiles(
|
||||
coordination_contracts,
|
||||
expected_artifact: None,
|
||||
deliverables: Vec::new(),
|
||||
token_budget: task_spec
|
||||
.budget
|
||||
.as_ref()
|
||||
.and_then(|budget| budget.max_tokens),
|
||||
resume_identity: Some(session_name.clone()),
|
||||
generation: 1,
|
||||
resume_from_agent_id: None,
|
||||
|
||||
@@ -3301,7 +3301,6 @@ async fn run_fleet_command(workspace: &Path, config: &Config, args: FleetArgs) -
|
||||
.max(max_subagents),
|
||||
Duration::from_secs(config.subagent_heartbeat_timeout_secs_for_provider(provider)),
|
||||
config.launch_concurrency_for_provider(provider),
|
||||
config.subagent_token_budget_for_provider(provider),
|
||||
);
|
||||
// Probe the durable ledger *before* opening the manager: FleetManager::open
|
||||
// creates `.codewhale/fleet.jsonl` as a side effect, so a later probe would
|
||||
@@ -12436,7 +12435,6 @@ async fn build_direct_workflow_tool(
|
||||
.max(max_subagents),
|
||||
Duration::from_secs(config.subagent_heartbeat_timeout_secs_for_provider(provider)),
|
||||
config.launch_concurrency_for_provider(provider),
|
||||
config.subagent_token_budget_for_provider(provider),
|
||||
);
|
||||
let roster = Arc::new(crate::fleet::identity::load_effective_roster(
|
||||
&config.fleet_config(),
|
||||
|
||||
@@ -595,7 +595,6 @@ fn runtime_api_sub_agent_manager(workspace: &FsPath, workers: usize) -> SharedSu
|
||||
max_agents,
|
||||
Duration::from_secs(crate::config::DEFAULT_SUBAGENT_HEARTBEAT_TIMEOUT_SECS),
|
||||
max_agents,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -2686,10 +2686,6 @@ async fn agent_runs_runtime_api_exposes_persisted_worker_receipts() -> Result<()
|
||||
output_tokens: None,
|
||||
total_tokens: None,
|
||||
cost_microusd: None,
|
||||
token_budget: None,
|
||||
budget_spent_tokens: None,
|
||||
budget_remaining_tokens: None,
|
||||
budget_scope: None,
|
||||
note: "not reported".to_string(),
|
||||
},
|
||||
usage_source_fingerprints: Default::default(),
|
||||
|
||||
@@ -10370,7 +10370,6 @@ impl RuntimeThreadManager {
|
||||
plan_state: new_shared_plan_state(),
|
||||
goal_state,
|
||||
max_spawn_depth: cfg.subagent_max_spawn_depth_for_provider(provider),
|
||||
subagent_token_budget: cfg.subagent_token_budget_for_provider(provider),
|
||||
network_policy,
|
||||
snapshots_enabled: !isolated_chat && cfg.snapshots_config().enabled,
|
||||
snapshots_max_workspace_bytes: cfg
|
||||
|
||||
@@ -1,21 +1,14 @@
|
||||
//! A final report inside the existing worker's budget and turn loop.
|
||||
//! A bounded final report inside the existing worker's turn loop. Runs stop
|
||||
//! on wall time, steps, cancellation, or completion — never on token
|
||||
//! accounting (#6189); the report turn is sized by a fixed allowance, not by
|
||||
//! what a budget has left.
|
||||
use super::*;
|
||||
|
||||
pub(super) const MAX_HAND_BACK_TOKENS: u64 = 8_192;
|
||||
const MIN_HAND_BACK_TOKENS: u64 = 1_024;
|
||||
const MAX_HAND_BACK_OUTPUT: u32 = 1_024;
|
||||
const MIN_HAND_BACK_OUTPUT: u64 = 128;
|
||||
const MAX_HAND_BACK_TIME: Duration = Duration::from_secs(10);
|
||||
|
||||
pub(super) fn token_reserve(limit: u64) -> u64 {
|
||||
let reserve = (limit / 10).min(MAX_HAND_BACK_TOKENS);
|
||||
if reserve >= MIN_HAND_BACK_TOKENS {
|
||||
reserve
|
||||
} else {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn wall_deadlines(runtime: &SubAgentRuntime) -> (Option<Instant>, Option<Instant>) {
|
||||
let hard = runtime.worker_profile.wall_deadline_ms.map(|deadline| {
|
||||
Instant::now() + Duration::from_millis(deadline.saturating_sub(epoch_millis_now()))
|
||||
@@ -29,105 +22,9 @@ pub(super) fn wall_deadlines(runtime: &SubAgentRuntime) -> (Option<Instant>, Opt
|
||||
}
|
||||
|
||||
impl SubAgentManager {
|
||||
/// The same capped families used by measured budget accounting must also
|
||||
/// retain missing-usage evidence across sibling work and continuations.
|
||||
/// An untouched/legacy record is not itself evidence of a missing bill.
|
||||
fn budget_has_unreported_usage(&self, worker: &str) -> bool {
|
||||
let mut capped_ancestors = BTreeSet::new();
|
||||
let mut capped_scopes = BTreeSet::new();
|
||||
for id in self.budget_ancestors(worker) {
|
||||
let Some(record) = self.worker_records.get(&id) else {
|
||||
continue;
|
||||
};
|
||||
if record.spec.runtime_profile.token_budget.is_some() {
|
||||
capped_ancestors.insert(id);
|
||||
}
|
||||
if record.usage.token_budget.is_some()
|
||||
&& let Some(scope) = record.usage.budget_scope.as_deref()
|
||||
{
|
||||
capped_scopes.insert(scope);
|
||||
}
|
||||
}
|
||||
self.worker_records.values().any(|record| {
|
||||
if !record.has_unreported_usage {
|
||||
return false;
|
||||
}
|
||||
self.budget_ancestors(&record.spec.worker_id)
|
||||
.iter()
|
||||
.any(|id| {
|
||||
capped_ancestors.contains(id)
|
||||
|| self.worker_records.get(id).is_some_and(|ancestor| {
|
||||
ancestor
|
||||
.usage
|
||||
.budget_scope
|
||||
.as_deref()
|
||||
.is_some_and(|scope| capped_scopes.contains(scope))
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn reserved_handback_tokens(&self, owner: &str, scope: Option<&str>) -> u64 {
|
||||
self.handback_reservations
|
||||
.iter()
|
||||
.filter_map(|(worker, reservation)| reservation.upgrade().map(|value| (worker, value)))
|
||||
.filter(|(worker, _)| {
|
||||
let ancestors = self.budget_ancestors(worker);
|
||||
scope.map_or_else(
|
||||
|| ancestors.contains(owner),
|
||||
|scope| {
|
||||
ancestors.iter().any(|id| {
|
||||
self.worker_records.get(id).is_some_and(|record| {
|
||||
record.usage.budget_scope.as_deref() == Some(scope)
|
||||
})
|
||||
})
|
||||
},
|
||||
)
|
||||
})
|
||||
.fold(0_u64, |total, (_, tokens)| total.saturating_add(*tokens))
|
||||
}
|
||||
|
||||
/// Dispatch headroom, separate from actual billed usage. A single reserve
|
||||
/// is held back per applicable scope, not once per sibling. Active report
|
||||
/// reservations are subtracted while their provider response is pending.
|
||||
pub(super) fn available_worker_tokens(
|
||||
&self,
|
||||
worker: &str,
|
||||
preserve_report: bool,
|
||||
) -> Option<u64> {
|
||||
self.budget_ancestors(worker)
|
||||
.iter()
|
||||
.fold(None, |remaining, id| {
|
||||
let Some(record) = self.worker_records.get(id) else {
|
||||
return remaining;
|
||||
};
|
||||
let available = |spent: u64, limit: u64, scope: Option<&str>| {
|
||||
limit
|
||||
.saturating_sub(spent)
|
||||
.saturating_sub(self.reserved_handback_tokens(id, scope))
|
||||
.saturating_sub(if preserve_report {
|
||||
token_reserve(limit)
|
||||
} else {
|
||||
0
|
||||
})
|
||||
};
|
||||
let shared = self.budget_scope_state(id).map(|(spent, limit)| {
|
||||
available(spent, limit, record.usage.budget_scope.as_deref())
|
||||
});
|
||||
let local = record
|
||||
.spec
|
||||
.runtime_profile
|
||||
.token_budget
|
||||
.map(|limit| available(self.subtree_budget_spent(id), limit, None));
|
||||
narrow_optional_limit(remaining, narrow_optional_limit(shared, local))
|
||||
})
|
||||
}
|
||||
|
||||
pub(super) fn reserve_handback(
|
||||
&mut self,
|
||||
worker: &str,
|
||||
local_remaining: Option<u64>,
|
||||
allowance: u64,
|
||||
input_tokens: u64,
|
||||
output_cap: u32,
|
||||
) -> std::result::Result<(u32, Arc<u64>), &'static str> {
|
||||
@@ -138,25 +35,16 @@ impl SubAgentManager {
|
||||
{
|
||||
return Err("worker is no longer active");
|
||||
}
|
||||
if self.budget_has_unreported_usage(worker) {
|
||||
return Err("earlier provider token usage is unknown in an applicable budget scope");
|
||||
}
|
||||
self.handback_reservations
|
||||
.retain(|_, value| value.strong_count() > 0);
|
||||
if self.handback_reservations.contains_key(worker) {
|
||||
return Err("a hand-back turn is already in flight");
|
||||
}
|
||||
let available =
|
||||
narrow_optional_limit(self.available_worker_tokens(worker, false), local_remaining)
|
||||
.unwrap_or(allowance)
|
||||
.min(allowance);
|
||||
let output = available
|
||||
let output = MAX_HAND_BACK_TOKENS
|
||||
.saturating_sub(input_tokens)
|
||||
.min(u64::from(output_cap));
|
||||
if output < MIN_HAND_BACK_OUTPUT {
|
||||
return Err(
|
||||
"remaining token allowance cannot cover the estimated report input and output",
|
||||
);
|
||||
return Err("the fixed hand-back allowance cannot cover the report input and output");
|
||||
}
|
||||
let reservation = Arc::new(input_tokens.saturating_add(output));
|
||||
self.handback_reservations
|
||||
@@ -263,10 +151,6 @@ pub(super) async fn request_report(
|
||||
messages: &mut Vec<Message>,
|
||||
steps: &mut u32,
|
||||
max_steps: u32,
|
||||
token_budget: Option<u64>,
|
||||
tokens_used: u64,
|
||||
allowance: u64,
|
||||
usage_complete: bool,
|
||||
hard_deadline: Option<Instant>,
|
||||
cause: &str,
|
||||
) -> Outcome {
|
||||
@@ -281,34 +165,11 @@ pub(super) async fn request_report(
|
||||
if *steps == 0 {
|
||||
return fallback("no completed model turn was recorded");
|
||||
}
|
||||
if max_steps > 0 && *steps >= max_steps {
|
||||
return fallback("no model turn remains within the step cap");
|
||||
}
|
||||
if allowance < MIN_HAND_BACK_TOKENS {
|
||||
return fallback("token allowance is too small to reserve a report");
|
||||
}
|
||||
if narrow_optional_limit(
|
||||
runtime
|
||||
.manager
|
||||
.read()
|
||||
.await
|
||||
.remaining_worker_tokens(agent_id),
|
||||
token_budget.map(|limit| limit.saturating_sub(tokens_used)),
|
||||
) == Some(0)
|
||||
{
|
||||
return fallback("remaining token allowance is exhausted");
|
||||
}
|
||||
if !usage_complete
|
||||
&& (token_budget.is_some()
|
||||
|| runtime
|
||||
.manager
|
||||
.read()
|
||||
.await
|
||||
.remaining_worker_tokens(agent_id)
|
||||
.is_some())
|
||||
{
|
||||
return fallback("earlier provider token usage is unknown");
|
||||
}
|
||||
// The hand-back turn is a reserved allowance OUTSIDE the task step
|
||||
// budget: a worker stopped by its own step cap has, by definition, no
|
||||
// task step left, and that same stop message promises the remaining
|
||||
// turn is reserved for hand-back (#6277). Only the wall-time deadline
|
||||
// below can refuse the turn.
|
||||
let deadline = hard_deadline
|
||||
.unwrap_or_else(|| Instant::now() + MAX_HAND_BACK_TIME)
|
||||
.min(Instant::now() + MAX_HAND_BACK_TIME)
|
||||
@@ -328,12 +189,12 @@ pub(super) async fn request_report(
|
||||
let candidate = report_messages(assignment, messages, cause, evidence_bytes);
|
||||
let estimate =
|
||||
crate::compaction::estimate_input_tokens_conservative(&candidate, Some(&system)) as u64;
|
||||
if estimate.saturating_add(MIN_HAND_BACK_OUTPUT) <= allowance {
|
||||
if estimate.saturating_add(MIN_HAND_BACK_OUTPUT) <= MAX_HAND_BACK_TOKENS {
|
||||
break candidate;
|
||||
}
|
||||
if evidence_bytes <= 256 {
|
||||
return fallback(
|
||||
"token allowance cannot fit the report instructions and grounded evidence",
|
||||
"the fixed hand-back allowance cannot fit the report instructions and grounded evidence",
|
||||
);
|
||||
}
|
||||
evidence_bytes /= 2;
|
||||
@@ -346,8 +207,6 @@ pub(super) async fn request_report(
|
||||
.effective_route_envelope(&runtime.model, chrono::Utc::now());
|
||||
let (output_tokens, _reservation) = match runtime.manager.write().await.reserve_handback(
|
||||
agent_id,
|
||||
token_budget.map(|limit| limit.saturating_sub(tokens_used)),
|
||||
allowance,
|
||||
input_tokens,
|
||||
runtime
|
||||
.client
|
||||
|
||||
@@ -13,7 +13,6 @@ struct Fixture {
|
||||
report_started: Arc<Notify>,
|
||||
release_report: Arc<Notify>,
|
||||
cancel: CancellationToken,
|
||||
resume_runtime: SubAgentRuntime,
|
||||
completions: mpsc::Receiver<SubAgentCompletion>,
|
||||
mailbox: MailboxReceiver,
|
||||
}
|
||||
@@ -41,12 +40,7 @@ impl Drop for Fixture {
|
||||
}
|
||||
}
|
||||
|
||||
async fn fixture(
|
||||
mode: &'static str,
|
||||
first_tokens: u64,
|
||||
cap: Option<u64>,
|
||||
max_steps: u32,
|
||||
) -> Fixture {
|
||||
async fn fixture(mode: &'static str, first_tokens: u64, max_steps: u32) -> Fixture {
|
||||
let workspace = tempdir().unwrap();
|
||||
fs::write(
|
||||
workspace.path().join("README.md"),
|
||||
@@ -70,11 +64,7 @@ async fn fixture(
|
||||
requests.push(body);
|
||||
requests.len()
|
||||
};
|
||||
if mode == "resume-unknown" && call == 2 {
|
||||
report_started.notify_one();
|
||||
release_report.notified().await;
|
||||
}
|
||||
let choice = if call == 1 || (mode == "resume-unknown" && call == 3) {
|
||||
let choice = if call == 1 {
|
||||
json!({"index": 0, "message": {"role": "assistant",
|
||||
"content": if mode == "tool-only" { Value::Null } else { json!("RECORDED_FINDING: checksum validation is missing.") },
|
||||
"tool_calls": [{"id": "read-one", "type": "function", "function": {
|
||||
@@ -127,7 +117,6 @@ async fn fixture(
|
||||
));
|
||||
let mut spec = make_worker_spec("report-worker", workspace.path().to_path_buf());
|
||||
spec.max_steps = max_steps;
|
||||
spec.runtime_profile.token_budget = cap;
|
||||
spec.runtime_profile.max_steps = max_steps;
|
||||
spec.runtime_profile.wall_time_secs = Some(5);
|
||||
spec.runtime_profile.wall_deadline_ms = Some(epoch_millis_now() + 5_000);
|
||||
@@ -186,39 +175,8 @@ async fn fixture(
|
||||
{
|
||||
let mut guard = manager.write().await;
|
||||
guard.register_worker_for_session(spec, &runtime.context.state_namespace);
|
||||
if let Some(cap) = cap {
|
||||
guard.attach_shared_budget_scope("report-worker", "report-pool", cap);
|
||||
}
|
||||
guard.agents.insert("report-worker".to_string(), agent);
|
||||
if mode == "shared-unknown" {
|
||||
let mut sibling = make_worker_spec("settled-sibling", workspace.path().to_path_buf());
|
||||
sibling.runtime_profile.token_budget = cap;
|
||||
guard.register_worker_for_session(sibling, &runtime.context.state_namespace);
|
||||
guard.attach_shared_budget_scope("settled-sibling", "report-pool", cap.unwrap());
|
||||
guard.record_worker_usage(
|
||||
"settled-sibling",
|
||||
"missing-sibling",
|
||||
&Usage::default(),
|
||||
None,
|
||||
);
|
||||
guard.record_worker_usage(
|
||||
"settled-sibling",
|
||||
"known-sibling",
|
||||
&Usage {
|
||||
input_tokens: 7,
|
||||
output_tokens: 4,
|
||||
..Usage::default()
|
||||
},
|
||||
None,
|
||||
);
|
||||
guard
|
||||
.worker_records
|
||||
.get_mut("settled-sibling")
|
||||
.unwrap()
|
||||
.status = AgentWorkerStatus::Completed;
|
||||
}
|
||||
}
|
||||
let resume_runtime = runtime.clone();
|
||||
let task = tokio::spawn(run_subagent_task(SubAgentTask {
|
||||
manager_handle: Arc::clone(&manager),
|
||||
runtime,
|
||||
@@ -230,7 +188,6 @@ async fn fixture(
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps,
|
||||
token_budget: cap,
|
||||
wall_time: Duration::from_secs(5),
|
||||
input_rx,
|
||||
launch_gate: None,
|
||||
@@ -245,7 +202,6 @@ async fn fixture(
|
||||
report_started,
|
||||
release_report,
|
||||
cancel,
|
||||
resume_runtime,
|
||||
completions,
|
||||
mailbox: mailbox_rx,
|
||||
}
|
||||
@@ -256,7 +212,7 @@ async fn fixture(
|
||||
async fn budget_handback_turn_consolidates_tool_only_work_and_checks_declared_deliverables() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("tool-only", 15, Some(100_000), 2).await;
|
||||
let mut fixture = fixture("tool-only", 15, 1).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert_eq!(result.steps_taken, 2);
|
||||
@@ -317,43 +273,12 @@ async fn budget_handback_turn_consolidates_tool_only_work_and_checks_declared_de
|
||||
assert!(fixture.completions.try_recv().is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_turn_uses_remaining_token_reserve_before_hard_exhaustion() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("success", 92_000, Some(100_000), 8).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert_eq!(result.usage.as_ref().unwrap().total_tokens, Some(92_030));
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 2);
|
||||
assert!(result.result.as_deref().unwrap().contains("PARTIAL_REPORT"));
|
||||
assert!(
|
||||
fixture.requests.lock().unwrap()[1]
|
||||
.to_string()
|
||||
.contains("RECORDED_FINDING")
|
||||
);
|
||||
assert!(result.checkpoint.as_ref().unwrap().messages.iter().flat_map(|message| &message.content)
|
||||
.any(|block| matches!(block, ContentBlock::ToolResult { tool_use_id, content, is_error, .. }
|
||||
if tool_use_id == "read-one" && *is_error == Some(true)
|
||||
&& content.contains("not executed") && content.contains("budget_exhausted")
|
||||
&& !content.contains("crashed_and_repaired"))));
|
||||
assert!(
|
||||
!fixture
|
||||
.mailbox
|
||||
.drain()
|
||||
.iter()
|
||||
.any(|entry| matches!(entry.message, MailboxMessage::ToolCallStarted { .. })),
|
||||
"normal work stops at the reporting reserve"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_turn_rejects_provider_tools_and_preserves_fallback_verdicts() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("tool", 15, Some(100_000), 2).await;
|
||||
let mut fixture = fixture("tool", 15, 1).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 2);
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
@@ -419,7 +344,7 @@ async fn budget_handback_turn_failure_and_timeout_do_not_add_worker_retries() {
|
||||
("failure", "provider call failed"),
|
||||
("timeout", "report deadline expired"),
|
||||
] {
|
||||
let mut fixture = fixture(mode, 15, Some(100_000), 2).await;
|
||||
let mut fixture = fixture(mode, 15, 1).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 2);
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
@@ -442,36 +367,12 @@ async fn budget_handback_turn_failure_and_timeout_do_not_add_worker_retries() {
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_turn_tiny_exhausted_and_unknown_allowances_refuse_paid_summary() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
for (mode, tokens, cap, reason) in [
|
||||
("success", 15, 15, "too small"),
|
||||
("success", 100_000, 100_000, "remaining token allowance"),
|
||||
("unknown", 0, 100_000, "usage is unknown"),
|
||||
] {
|
||||
let mut fixture = fixture(mode, tokens, Some(cap), 2).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 1);
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert!(
|
||||
result.result.as_deref().unwrap().contains(reason),
|
||||
"{result:?}"
|
||||
);
|
||||
if mode == "unknown" {
|
||||
assert_eq!(result.usage.as_ref().unwrap().total_tokens, None);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_turn_missing_report_usage_is_not_claimed_as_zero_cost() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("report-unknown", 15, Some(100_000), 2).await;
|
||||
let mut fixture = fixture("report-unknown", 15, 1).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 2);
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
@@ -483,120 +384,10 @@ async fn budget_handback_turn_missing_report_usage_is_not_claimed_as_zero_cost()
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_resumed_known_response_keeps_interrupted_source_usage_unknown() {
|
||||
async fn budget_handback_inflight_wall_timeout_persists_unreported_usage() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("resume-unknown", 0, Some(100_000), 4).await;
|
||||
// The first response omits usage; the second ordinary request is held so
|
||||
// the real interrupt/followup path resumes that checkpoint, not a forged
|
||||
// continuation edge. A known response in the successor cannot repair the
|
||||
// source's positively known missing coverage.
|
||||
tokio::time::timeout(Duration::from_secs(2), fixture.report_started.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
let successor = {
|
||||
let mut manager = fixture.manager.write().await;
|
||||
let (_, interrupted) = manager
|
||||
.interrupt_child("report-worker", None, "fixture interruption".to_string())
|
||||
.unwrap();
|
||||
assert!(matches!(interrupted.status, SubAgentStatus::Interrupted(_)));
|
||||
assert!(interrupted.checkpoint.as_ref().unwrap().continuable);
|
||||
assert!(manager.worker_records["report-worker"].has_unreported_usage);
|
||||
assert_eq!(
|
||||
manager.worker_records["report-worker"].usage.total_tokens,
|
||||
None
|
||||
);
|
||||
// This fixture owns the original task handle instead of the manager.
|
||||
// Abort it just as production interrupt does before launching followup.
|
||||
fixture.task.take().unwrap().abort();
|
||||
let saved = serde_json::to_vec(&manager.worker_records).unwrap();
|
||||
manager.worker_records = serde_json::from_slice(&saved).unwrap();
|
||||
assert!(manager.worker_records["report-worker"].has_unreported_usage);
|
||||
manager
|
||||
.resume_from_checkpoint(
|
||||
Arc::clone(&fixture.manager),
|
||||
fixture.resume_runtime.clone(),
|
||||
"report-worker",
|
||||
"Continue the inspection.",
|
||||
)
|
||||
.unwrap()
|
||||
.agent_id
|
||||
};
|
||||
let result = tokio::time::timeout(Duration::from_secs(5), async {
|
||||
loop {
|
||||
let result = fixture.manager.read().await.get_result(&successor).unwrap();
|
||||
if result.status != SubAgentStatus::Running {
|
||||
return result;
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("resumed worker must settle");
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert_eq!(result.usage.as_ref().unwrap().total_tokens, Some(15));
|
||||
assert!(
|
||||
result
|
||||
.result
|
||||
.as_deref()
|
||||
.unwrap()
|
||||
.contains("usage is unknown in an applicable budget scope")
|
||||
);
|
||||
assert_eq!(
|
||||
fixture.requests.lock().unwrap().len(),
|
||||
3,
|
||||
"no successor reporting request"
|
||||
);
|
||||
let manager = fixture.manager.read().await;
|
||||
assert!(manager.worker_records["report-worker"].has_unreported_usage);
|
||||
assert!(!manager.worker_records[&successor].has_unreported_usage);
|
||||
assert_eq!(
|
||||
manager.worker_records["report-worker"].usage.total_tokens,
|
||||
None
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_shared_sibling_unknown_usage_blocks_only_paid_reporting() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("shared-unknown", 15, Some(100_000), 2).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert_eq!(result.usage.as_ref().unwrap().total_tokens, Some(15));
|
||||
assert!(
|
||||
result
|
||||
.result
|
||||
.as_deref()
|
||||
.unwrap()
|
||||
.contains("usage is unknown in an applicable budget scope")
|
||||
);
|
||||
assert_eq!(fixture.requests.lock().unwrap().len(), 1);
|
||||
let manager = fixture.manager.read().await;
|
||||
assert!(manager.worker_records["settled-sibling"].has_unreported_usage);
|
||||
assert!(!manager.worker_records["report-worker"].has_unreported_usage);
|
||||
assert_eq!(
|
||||
manager.worker_records["settled-sibling"].usage.total_tokens,
|
||||
Some(11)
|
||||
);
|
||||
assert_eq!(
|
||||
manager.aggregate_budget_spent("report-pool"),
|
||||
26,
|
||||
"known subtotals are never erased"
|
||||
);
|
||||
assert_eq!(
|
||||
manager.worker_records["report-worker"].verification.status,
|
||||
"deliverable_missing"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn budget_handback_inflight_wall_timeout_persists_unknown_scope_coverage() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("work-timeout", 15, Some(100_000), 4).await;
|
||||
let mut fixture = fixture("work-timeout", 15, 4).await;
|
||||
let result = fixture.finish().await;
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert!(
|
||||
@@ -608,26 +399,15 @@ async fn budget_handback_inflight_wall_timeout_persists_unknown_scope_coverage()
|
||||
);
|
||||
assert_eq!(
|
||||
fixture.requests.lock().unwrap().len(),
|
||||
2,
|
||||
"no report after an unmeasured in-flight request"
|
||||
3,
|
||||
"the reserved hand-back turn still dispatches after an unmeasured in-flight request; its own deadline bounds it"
|
||||
);
|
||||
let mut manager = fixture.manager.write().await;
|
||||
let manager = fixture.manager.read().await;
|
||||
assert!(manager.worker_records["report-worker"].has_unreported_usage);
|
||||
assert_eq!(
|
||||
manager.worker_records["report-worker"].usage.total_tokens,
|
||||
Some(15)
|
||||
);
|
||||
let saved = serde_json::to_vec(&manager.worker_records).unwrap();
|
||||
manager.worker_records = serde_json::from_slice(&saved).unwrap();
|
||||
let mut sibling = make_worker_spec("later-sibling", fixture.workspace.path().to_path_buf());
|
||||
sibling.runtime_profile.token_budget = Some(100_000);
|
||||
manager.register_worker(sibling);
|
||||
manager.attach_shared_budget_scope("later-sibling", "report-pool", 100_000);
|
||||
assert!(matches!(
|
||||
manager.reserve_handback("later-sibling", Some(100_000), 8_192, 500, 1_024),
|
||||
Err(reason) if reason.contains("usage is unknown")
|
||||
));
|
||||
assert_eq!(manager.aggregate_budget_spent("report-pool"), 15);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -635,8 +415,7 @@ fn budget_handback_coverage_marker_is_sticky_without_reclassifying_legacy_or_mea
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 4);
|
||||
for id in ["worker", "unrelated"] {
|
||||
let mut spec = make_worker_spec(id, tmp.path().to_path_buf());
|
||||
spec.runtime_profile.token_budget = Some(20_000);
|
||||
let spec = make_worker_spec(id, tmp.path().to_path_buf());
|
||||
manager.register_worker(spec);
|
||||
}
|
||||
let measured_zero = Usage {
|
||||
@@ -655,9 +434,7 @@ fn budget_handback_coverage_marker_is_sticky_without_reclassifying_legacy_or_mea
|
||||
let legacy: AgentWorkerRecord = serde_json::from_value(legacy).unwrap();
|
||||
assert!(!legacy.has_unreported_usage);
|
||||
manager.worker_records.insert("worker".to_string(), legacy);
|
||||
let lease = manager
|
||||
.reserve_handback("worker", Some(20_000), 2_000, 500, 1_024)
|
||||
.unwrap();
|
||||
let (_output, lease) = manager.reserve_handback("worker", 500, 1_024).unwrap();
|
||||
drop(lease);
|
||||
manager.record_worker_usage("worker", "missing", &Usage::default(), None);
|
||||
manager.record_worker_usage(
|
||||
@@ -677,10 +454,6 @@ fn budget_handback_coverage_marker_is_sticky_without_reclassifying_legacy_or_mea
|
||||
Some(15)
|
||||
);
|
||||
assert!(manager.worker_records["worker"].has_unreported_usage);
|
||||
assert!(matches!(
|
||||
manager.reserve_handback("worker", Some(20_000), 2_000, 500, 1_024),
|
||||
Err(reason) if reason.contains("usage is unknown")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -688,7 +461,7 @@ fn budget_handback_coverage_marker_is_sticky_without_reclassifying_legacy_or_mea
|
||||
async fn budget_handback_turn_cancellation_wins_once_and_releases_shared_reservation() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("hold", 15, Some(100_000), 2).await;
|
||||
let mut fixture = fixture("hold", 15, 2).await;
|
||||
tokio::time::timeout(Duration::from_secs(2), fixture.report_started.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -722,7 +495,7 @@ async fn budget_handback_turn_cancellation_wins_once_and_releases_shared_reserva
|
||||
async fn budget_handback_turn_cancellation_after_response_preserves_actual_usage() {
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let mut fixture = fixture("hold", 15, Some(100_000), 2).await;
|
||||
let mut fixture = fixture("hold", 15, 2).await;
|
||||
tokio::time::timeout(Duration::from_secs(2), fixture.report_started.notified())
|
||||
.await
|
||||
.unwrap();
|
||||
@@ -755,76 +528,21 @@ async fn budget_handback_turn_cancellation_after_response_preserves_actual_usage
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn budget_handback_reservations_intersect_shared_ancestor_and_source_caps_without_refunds() {
|
||||
fn handback_reservation_uses_the_fixed_allowance_and_refuses_a_second_turn() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 8);
|
||||
for (id, parent, cap) in [
|
||||
("parent", None, 20_000),
|
||||
("source", None, 15_000),
|
||||
("one", Some("parent"), 15_000),
|
||||
("two", Some("parent"), 15_000),
|
||||
] {
|
||||
let mut spec = make_worker_spec(id, tmp.path().to_path_buf());
|
||||
spec.parent_run_id = parent.map(str::to_string);
|
||||
spec.runtime_profile.token_budget = Some(cap);
|
||||
manager.register_worker(spec);
|
||||
manager.attach_shared_budget_scope(id, "shared", 20_000);
|
||||
}
|
||||
manager
|
||||
.worker_records
|
||||
.get_mut("parent")
|
||||
.unwrap()
|
||||
.usage
|
||||
.total_tokens = Some(18_000);
|
||||
let one = manager
|
||||
.reserve_handback("one", Some(2_000), 2_000, 900, 1_024)
|
||||
.unwrap();
|
||||
assert_eq!(manager.available_worker_tokens("two", false), Some(76));
|
||||
assert!(
|
||||
manager
|
||||
.reserve_handback("two", Some(2_000), 2_000, 100, 1_024)
|
||||
.is_err()
|
||||
);
|
||||
assert!(
|
||||
manager
|
||||
.resolve_spawn_budget_scope("new", Some("parent"), None)
|
||||
.is_err()
|
||||
);
|
||||
assert_eq!(
|
||||
manager.aggregate_budget_spent("shared"),
|
||||
18_000,
|
||||
"a reservation is not actual usage"
|
||||
);
|
||||
drop(one);
|
||||
assert_eq!(manager.available_worker_tokens("two", false), Some(2_000));
|
||||
let source = manager.worker_records.get_mut("source").unwrap();
|
||||
source.usage.total_tokens = Some(1_000);
|
||||
source.spec.runtime_profile.token_budget = Some(1_050);
|
||||
let profile = manager.worker_records["two"].spec.runtime_profile.clone();
|
||||
manager
|
||||
.worker_records
|
||||
.get_mut("two")
|
||||
.unwrap()
|
||||
.spec
|
||||
.launch_manifest = Some(
|
||||
serde_json::from_value(json!({
|
||||
"owner_session": "parent", "child_id": "two", "profile": profile,
|
||||
"prompt": "continue", "cwd": null, "worktree": false, "writable_roots": [],
|
||||
"writable_files": [], "coordination_contracts": [], "generation": 1,
|
||||
"resume_identity": null, "resume_from_agent_id": "source"
|
||||
}))
|
||||
.unwrap(),
|
||||
);
|
||||
assert_eq!(manager.available_worker_tokens("two", false), Some(50));
|
||||
assert!(
|
||||
manager
|
||||
.reserve_handback("two", None, 2_000, 100, 1_024)
|
||||
.is_err()
|
||||
);
|
||||
let saved = serde_json::to_vec(&manager.worker_records).unwrap();
|
||||
manager.worker_records = serde_json::from_slice(&saved).unwrap();
|
||||
assert_eq!(manager.available_worker_tokens("two", false), Some(50));
|
||||
assert_eq!(manager.aggregate_budget_spent("shared"), 19_000);
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 4);
|
||||
manager.register_worker(make_worker_spec("w", tmp.path().to_path_buf()));
|
||||
let (output, lease) = manager.reserve_handback("w", 500, 1_024).unwrap();
|
||||
assert_eq!(output, 1_024);
|
||||
assert!(matches!(
|
||||
manager.reserve_handback("w", 500, 1_024),
|
||||
Err(reason) if reason.contains("already in flight")
|
||||
));
|
||||
drop(lease);
|
||||
assert!(matches!(
|
||||
manager.reserve_handback("w", 8_200, 1_024),
|
||||
Err(reason) if reason.contains("fixed hand-back allowance")
|
||||
));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -849,10 +567,6 @@ async fn budget_handback_expired_original_deadline_refuses_the_model_call() {
|
||||
&mut vec![],
|
||||
&mut steps,
|
||||
2,
|
||||
Some(100_000),
|
||||
15,
|
||||
8_192,
|
||||
true,
|
||||
Some(Instant::now() - Duration::from_millis(1)),
|
||||
"wall-time budget exhausted",
|
||||
)
|
||||
@@ -890,7 +604,7 @@ fn git(root: &Path, args: &[&str]) {
|
||||
/// spawn-time delivery baseline is what makes the inventory attributable to
|
||||
/// this worker rather than the parent's own dirty files.
|
||||
#[tokio::test]
|
||||
async fn budget_death_preservation_note_names_surviving_workspace_changes() {
|
||||
async fn run_death_preservation_note_names_surviving_workspace_changes() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let root = tmp.path();
|
||||
git(root, &["init", "--quiet"]);
|
||||
|
||||
@@ -15,8 +15,6 @@ fn child(
|
||||
record.usage.input_tokens = units.map(|units| units * 8);
|
||||
record.usage.output_tokens = units.map(|units| units * 2);
|
||||
record.usage.total_tokens = units.map(|units| units * 10);
|
||||
// This repeated pool subtotal is deliberately not this child's spend.
|
||||
record.usage.budget_spent_tokens = Some(99_999);
|
||||
id
|
||||
}
|
||||
|
||||
@@ -413,11 +411,8 @@ async fn completion_usage_dozen_child_status_measures_bytes_and_keeps_descendant
|
||||
runtime_build_sha: "a".repeat(40),
|
||||
});
|
||||
record.spec.runtime_profile.max_steps = 12;
|
||||
record.spec.runtime_profile.token_budget = Some(12_000);
|
||||
record.spec.runtime_profile.wall_time_secs = Some(600);
|
||||
record.spec.runtime_profile.wall_deadline_ms = Some(record.updated_at_ms + 600_000);
|
||||
record.usage.token_budget = Some(12_000);
|
||||
record.usage.budget_remaining_tokens = Some(11_220);
|
||||
if index == 11 {
|
||||
record.verification.status = "deliverable_missing".into();
|
||||
record.verification.summary = "The claimed report.md is missing.".into();
|
||||
@@ -507,7 +502,8 @@ async fn completion_usage_dozen_child_status_measures_bytes_and_keeps_descendant
|
||||
let addressed: Value = serde_json::from_str(&addressed.content).unwrap();
|
||||
assert_eq!(addressed["compact"], true);
|
||||
assert_eq!(addressed["child_route"]["model_id"], "deepseek-v4-flash");
|
||||
assert_eq!(addressed["effective_limits"]["token_budget"], 12_000);
|
||||
assert_eq!(addressed["effective_limits"]["max_steps"], 12);
|
||||
assert_eq!(addressed["effective_limits"]["wall_time_secs"], 600);
|
||||
assert_eq!(addressed["max_spawn_depth"], 4);
|
||||
assert_eq!(addressed["usage"]["input_tokens"], 8);
|
||||
assert_eq!(addressed["usage"]["output_tokens"], 2);
|
||||
|
||||
@@ -42,7 +42,6 @@ fn worker(root: &Path, write: bool, paths: &[&str], scope: &[&str]) -> (SubAgent
|
||||
coordination_contracts: Vec::new(),
|
||||
expected_artifact: None,
|
||||
deliverables: paths.iter().map(|path| (*path).into()).collect(),
|
||||
token_budget: None,
|
||||
resume_identity: None,
|
||||
generation: 1,
|
||||
resume_from_agent_id: None,
|
||||
|
||||
@@ -123,12 +123,7 @@ pub(super) fn compact_row(manager: &SubAgentManager, agent: &SubAgent) -> Value
|
||||
object.insert("max_spawn_depth".into(), json!(record.spec.max_spawn_depth));
|
||||
if let Ok(profile) = serde_json::to_value(&record.spec.runtime_profile) {
|
||||
let mut limits = serde_json::Map::new();
|
||||
for key in [
|
||||
"max_steps",
|
||||
"token_budget",
|
||||
"wall_time_secs",
|
||||
"wall_deadline_ms",
|
||||
] {
|
||||
for key in ["max_steps", "wall_time_secs", "wall_deadline_ms"] {
|
||||
if let Some(value) = profile.get(key) {
|
||||
limits.insert(key.to_string(), value.clone());
|
||||
}
|
||||
@@ -142,8 +137,7 @@ pub(super) fn compact_row(manager: &SubAgentManager, agent: &SubAgent) -> Value
|
||||
"usage".into(),
|
||||
json!({
|
||||
"input_tokens": usage.input_tokens, "output_tokens": usage.output_tokens,
|
||||
"total_tokens": usage.total_tokens, "token_budget": usage.token_budget,
|
||||
"budget_remaining_tokens": usage.budget_remaining_tokens,
|
||||
"total_tokens": usage.total_tokens,
|
||||
}),
|
||||
);
|
||||
object.insert("last_activity_ms".into(), json!(record.updated_at_ms));
|
||||
|
||||
@@ -219,7 +219,6 @@ async fn lifecycle_compact_roster_bounds_every_state_and_pages_multibyte_names()
|
||||
}
|
||||
let record = guard.worker_records.get_mut(&id).unwrap();
|
||||
record.usage.total_tokens = Some(100);
|
||||
record.usage.budget_spent_tokens = Some(3700);
|
||||
record.verification.summary = "\"🐋".repeat(10_000);
|
||||
if i > 0 {
|
||||
record.parent_run_id = Some("agent_bounded-0".into());
|
||||
|
||||
@@ -1,42 +1,7 @@
|
||||
use super::tests::{make_snapshot, make_worker_spec, run_incomplete_response_worker, stub_runtime};
|
||||
use super::tests::{make_snapshot, make_worker_spec, stub_runtime};
|
||||
use super::*;
|
||||
use std::sync::atomic::Ordering;
|
||||
use tempfile::tempdir;
|
||||
|
||||
fn record(
|
||||
manager: &mut SubAgentManager,
|
||||
id: &str,
|
||||
parent: Option<&str>,
|
||||
cap: Option<u64>,
|
||||
spent: u64,
|
||||
) {
|
||||
let mut spec = make_worker_spec(id, manager.workspace.clone());
|
||||
spec.parent_run_id = parent.map(str::to_string);
|
||||
spec.runtime_profile.token_budget = cap;
|
||||
manager.register_worker(spec);
|
||||
manager
|
||||
.worker_records
|
||||
.get_mut(id)
|
||||
.unwrap()
|
||||
.usage
|
||||
.total_tokens = Some(spent);
|
||||
}
|
||||
|
||||
fn continuation(manager: &mut SubAgentManager, id: &str, source: &str) {
|
||||
let record = manager.worker_records.get_mut(id).unwrap();
|
||||
record.spec.launch_manifest = Some(
|
||||
serde_json::from_value(json!({
|
||||
"owner_session": "root", "child_id": id,
|
||||
"profile": record.spec.runtime_profile,
|
||||
"prompt": "continue", "cwd": null, "worktree": false,
|
||||
"writable_roots": [], "writable_files": [], "coordination_contracts": [],
|
||||
"token_budget": record.spec.runtime_profile.token_budget,
|
||||
"resume_identity": id, "generation": 1, "resume_from_agent_id": source
|
||||
}))
|
||||
.expect("continuation manifest"),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn depth_one_child_cannot_spawn_a_grandchild_even_with_a_wider_profile() {
|
||||
let root = stub_runtime().with_max_spawn_depth(1);
|
||||
@@ -99,20 +64,17 @@ fn operator_and_inherited_budgets_only_narrow_including_zero_sentinels() {
|
||||
assert_eq!(resolve_max_steps(FleetRole::Worker, Some(0), Some(7)), 7);
|
||||
let parent = WorkerRuntimeProfile {
|
||||
max_steps: 8,
|
||||
token_budget: Some(100),
|
||||
wall_time_secs: Some(40),
|
||||
wall_deadline_ms: Some(123_000),
|
||||
..WorkerRuntimeProfile::default()
|
||||
};
|
||||
let requested = WorkerRuntimeProfile {
|
||||
token_budget: Some(1_000),
|
||||
wall_time_secs: Some(4_000),
|
||||
wall_deadline_ms: Some(456_000),
|
||||
..WorkerRuntimeProfile::default()
|
||||
};
|
||||
let child = parent.derive_child(&requested);
|
||||
assert_eq!(child.max_steps, 8);
|
||||
assert_eq!(child.token_budget, Some(100));
|
||||
assert_eq!(child.wall_time_secs, Some(40));
|
||||
assert_eq!(child.wall_deadline_ms, Some(123_000));
|
||||
}
|
||||
@@ -122,12 +84,11 @@ fn child_runtime_budget_context_reports_every_resolved_limit() {
|
||||
let mut runtime = stub_runtime();
|
||||
runtime.worker_profile.wall_time_secs = Some(1_800);
|
||||
runtime.worker_profile.wall_deadline_ms = Some(epoch_millis_now() + 1_700_000);
|
||||
let context = child_runtime_budget_context(&runtime, 50, 49, true, Some(120_000));
|
||||
let context = child_runtime_budget_context(&runtime, 50, 49);
|
||||
assert!(context.contains("Runtime budget (host-enforced"));
|
||||
assert!(context.contains("wall clock: task work stops about"));
|
||||
assert!(context.contains("total run budget 30m 00s"));
|
||||
assert!(context.contains("49 model turns of task work (limit 50"));
|
||||
assert!(context.contains("about 120000 input+output tokens"));
|
||||
assert!(context.contains("reserved hand-back turn"));
|
||||
assert!(context.contains("Commit or checkpoint work-in-progress early"));
|
||||
}
|
||||
@@ -135,41 +96,36 @@ fn child_runtime_budget_context_reports_every_resolved_limit() {
|
||||
#[test]
|
||||
fn child_runtime_budget_context_names_unbounded_limits_honestly() {
|
||||
let runtime = stub_runtime();
|
||||
let context = child_runtime_budget_context(&runtime, 0, 0, false, None);
|
||||
let context = child_runtime_budget_context(&runtime, 0, 0);
|
||||
assert!(context.contains("wall clock: no wall-clock limit."));
|
||||
assert!(context.contains("model steps: no per-run step cap."));
|
||||
assert!(context.contains("token allowance: no per-run token cap."));
|
||||
assert!(!context.contains("reserved hand-back turn"));
|
||||
assert!(context.contains("There is no token cap"));
|
||||
assert!(context.contains("reserved hand-back turn"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn child_budget_pacing_notice_fires_at_three_quarters_of_each_bound() {
|
||||
let started_at = Instant::now() - Duration::from_secs(80);
|
||||
let deadline = started_at + Duration::from_secs(100);
|
||||
let notice =
|
||||
child_budget_pacing_notice(started_at, Some(deadline), 38, 50, Some(20), Some(100))
|
||||
.expect("all three bounds past 75%");
|
||||
let notice = child_budget_pacing_notice(started_at, Some(deadline), 38, 50)
|
||||
.expect("all three bounds past 75%");
|
||||
assert!(notice.contains("kind=\"budget_pacing\""));
|
||||
assert!(notice.contains("wall clock:"));
|
||||
assert!(notice.contains("model steps: 38 of 50 used"));
|
||||
assert!(notice.contains("token allowance: ~20 of ~100 tokens remain"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn child_budget_pacing_notice_stays_silent_with_headroom() {
|
||||
let started_at = Instant::now();
|
||||
let deadline = started_at + Duration::from_secs(100);
|
||||
assert!(
|
||||
child_budget_pacing_notice(started_at, Some(deadline), 10, 50, Some(80), Some(100))
|
||||
.is_none()
|
||||
);
|
||||
assert!(child_budget_pacing_notice(started_at, Some(deadline), 10, 50).is_none());
|
||||
// No bound at all means there is nothing to pace against.
|
||||
assert!(child_budget_pacing_notice(started_at, None, 9_999, 0, None, None).is_none());
|
||||
assert!(child_budget_pacing_notice(started_at, None, 9_999, 0).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn per_call_budget_fields_reject_empty_zero_null_negative_and_oversized_values() {
|
||||
for field in ["token_budget", "max_steps", "wall_time_secs"] {
|
||||
for field in ["max_steps", "wall_time_secs"] {
|
||||
for invalid in [
|
||||
json!(0),
|
||||
json!(-1),
|
||||
@@ -192,113 +148,6 @@ fn per_call_budget_fields_reject_empty_zero_null_negative_and_oversized_values()
|
||||
input[field] = json!(value);
|
||||
assert!(parse_spawn_request(&input).is_err(), "{input}");
|
||||
}
|
||||
let request =
|
||||
parse_spawn_request(&json!({"prompt": "inspect", "token_budget": 100, "max_tokens": 9}))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
request.token_budget,
|
||||
Some(9),
|
||||
"aliases cannot erase a tighter supplied cap"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_child_budget_keeps_parent_pool_and_default_is_a_ceiling() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager =
|
||||
SubAgentManager::new(tmp.path().to_path_buf(), 4).with_default_token_budget(Some(100));
|
||||
record(&mut manager, "parent", None, Some(100), 25);
|
||||
manager.attach_shared_budget_scope("parent", "pool", 100);
|
||||
let scope = manager
|
||||
.resolve_spawn_budget_scope("child", Some("parent"), Some(999))
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(scope.scope_id, "pool");
|
||||
assert_eq!((scope.limit, scope.spent, scope.remaining), (100, 25, 75));
|
||||
let root = manager
|
||||
.resolve_spawn_budget_scope("new_root", None, Some(999))
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
assert_eq!(root.limit, 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn descendant_and_resume_usage_consume_each_ancestor_once_at_equality() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 8);
|
||||
record(&mut manager, "root", None, Some(100), 10);
|
||||
manager.attach_shared_budget_scope("root", "pool", 100);
|
||||
record(&mut manager, "child", Some("root"), Some(30), 10);
|
||||
manager.attach_shared_budget_scope("child", "pool", 100);
|
||||
record(&mut manager, "grandchild", Some("child"), None, 10);
|
||||
manager.attach_shared_budget_scope("grandchild", "pool", 100);
|
||||
record(&mut manager, "resume", Some("root"), Some(20), 10);
|
||||
continuation(&mut manager, "resume", "child");
|
||||
manager.attach_shared_budget_scope("resume", "pool", 100);
|
||||
assert_eq!(manager.subtree_budget_spent("child"), 30);
|
||||
assert_eq!(manager.aggregate_budget_spent("pool"), 40);
|
||||
assert_eq!(manager.remaining_worker_tokens("resume"), Some(0));
|
||||
assert!(
|
||||
manager
|
||||
.token_budget_exhausted_detail("resume")
|
||||
.unwrap()
|
||||
.contains("30/30")
|
||||
);
|
||||
let encoded = serde_json::to_vec(&manager.worker_records).unwrap();
|
||||
let restored: HashMap<String, AgentWorkerRecord> = serde_json::from_slice(&encoded).unwrap();
|
||||
manager.worker_records = restored
|
||||
.into_iter()
|
||||
.map(|(id, record)| (id, normalize_worker_record(record)))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
manager.remaining_worker_tokens("resume"),
|
||||
Some(0),
|
||||
"reload never refunds earlier usage"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_fork_cannot_move_usage_out_of_either_source_or_current_parent_pool() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 8);
|
||||
record(&mut manager, "source", None, None, 10);
|
||||
manager.attach_shared_budget_scope("source", "source_pool", 40);
|
||||
record(&mut manager, "parent", None, None, 5);
|
||||
manager.attach_shared_budget_scope("parent", "parent_pool", 100);
|
||||
record(&mut manager, "fork", Some("parent"), None, 30);
|
||||
continuation(&mut manager, "fork", "source");
|
||||
manager.attach_shared_budget_scope("fork", "source_pool", 40);
|
||||
assert_eq!(manager.aggregate_budget_spent("source_pool"), 40);
|
||||
assert_eq!(manager.aggregate_budget_spent("parent_pool"), 35);
|
||||
assert_eq!(manager.remaining_worker_tokens("fork"), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_budget_lineage_cycles_are_finite_and_do_not_double_count() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 4);
|
||||
record(&mut manager, "a", Some("b"), Some(3), 1);
|
||||
record(&mut manager, "b", Some("a"), Some(3), 2);
|
||||
assert_eq!(manager.subtree_budget_spent("a"), 3);
|
||||
assert_eq!(manager.remaining_worker_tokens("b"), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cleanup_retains_completed_budget_evidence_while_a_pool_member_runs() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let mut manager = SubAgentManager::new(tmp.path().to_path_buf(), 4);
|
||||
record(&mut manager, "old", None, Some(100), 90);
|
||||
manager.attach_shared_budget_scope("old", "pool", 100);
|
||||
record(&mut manager, "live", None, None, 10);
|
||||
manager.attach_shared_budget_scope("live", "pool", 100);
|
||||
let old = manager.worker_records.get_mut("old").unwrap();
|
||||
old.status = AgentWorkerStatus::Completed;
|
||||
old.completed_at_ms = Some(0);
|
||||
old.updated_at_ms = 0;
|
||||
manager.cleanup(Duration::ZERO);
|
||||
assert!(manager.worker_records.contains_key("old"));
|
||||
assert_eq!(manager.aggregate_budget_spent("pool"), 100);
|
||||
assert_eq!(manager.remaining_worker_tokens("live"), Some(0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -319,35 +168,11 @@ fn budget_partial_handback_is_bounded_and_keeps_unknown_usage_honest() {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn token_equality_stops_before_tools_and_returns_measured_partial_work() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let (result, calls, mailbox, total_tokens) =
|
||||
run_incomplete_response_worker(tmp.path(), "max_tokens", 8, Some(15)).await;
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert_eq!(
|
||||
calls.load(Ordering::SeqCst),
|
||||
1,
|
||||
"no post-budget summary request"
|
||||
);
|
||||
assert_eq!(total_tokens, Some(15));
|
||||
assert!(
|
||||
!mailbox
|
||||
.iter()
|
||||
.any(|message| matches!(message, MailboxMessage::ToolCallStarted { .. }))
|
||||
);
|
||||
let partial = result.result.unwrap();
|
||||
assert!(partial.contains("15/15"));
|
||||
assert!(partial.contains("partial response diagnostics"));
|
||||
assert!(partial.contains("output was truncated"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn launch_narrows_all_limits_and_continuation_cannot_restart_deadline() {
|
||||
let tmp = tempdir().unwrap();
|
||||
let manager = Arc::new(RwLock::new(
|
||||
SubAgentManager::new(tmp.path().to_path_buf(), 4)
|
||||
.with_default_token_budget(Some(100))
|
||||
.with_default_max_steps(Some(4))
|
||||
.with_default_wall_time(Some(Duration::from_secs(10))),
|
||||
));
|
||||
@@ -356,7 +181,6 @@ async fn launch_narrows_all_limits_and_continuation_cannot_restart_deadline() {
|
||||
runtime.manager = Arc::clone(&manager);
|
||||
runtime.cancel_token.cancel(); // inspect admission; no provider request may run
|
||||
let options = SubAgentSpawnOptions {
|
||||
token_budget: Some(999),
|
||||
max_steps: Some(999),
|
||||
wall_time: Some(Duration::from_secs(999)),
|
||||
..Default::default()
|
||||
@@ -374,7 +198,6 @@ async fn launch_narrows_all_limits_and_continuation_cannot_restart_deadline() {
|
||||
)
|
||||
.unwrap();
|
||||
let profile = &guard.worker_records[&child.agent_id].spec.runtime_profile;
|
||||
assert_eq!(profile.token_budget, Some(100));
|
||||
assert_eq!(profile.max_steps, 4);
|
||||
assert!(profile.wall_time_secs.unwrap() <= 10);
|
||||
guard
|
||||
@@ -447,262 +270,6 @@ async fn resume_intersects_saved_write_shell_and_tool_permissions_with_current_c
|
||||
assert!(profile.denied_tools.contains(&"exec_shell".to_string()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn measured_budget_caps_actual_wire_output_and_accounts_overshoot_without_retry() {
|
||||
use axum::{Json, Router, routing::post};
|
||||
let tmp = tempdir().unwrap();
|
||||
let requests = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
|
||||
let observed = Arc::clone(&requests);
|
||||
let app = Router::new().route("/v1/chat/completions", post(move |Json(request): Json<Value>| {
|
||||
let observed = Arc::clone(&observed);
|
||||
async move {
|
||||
observed.lock().unwrap().push(request);
|
||||
Json(json!({
|
||||
"id": "budget-probe", "model": "deepseek-v4-flash",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "partial evidence"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
||||
}))
|
||||
}
|
||||
}));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
let config = crate::config::Config {
|
||||
api_key: Some("fixture-key".to_string()),
|
||||
base_url: Some(format!("http://{address}/v1")),
|
||||
..Default::default()
|
||||
};
|
||||
let mut runtime = stub_runtime();
|
||||
runtime.client = DeepSeekClient::new(&config).unwrap();
|
||||
runtime.context = ToolContext::new(tmp.path().to_path_buf());
|
||||
runtime.manager = Arc::new(RwLock::new(SubAgentManager::new(
|
||||
tmp.path().to_path_buf(),
|
||||
4,
|
||||
)));
|
||||
record(
|
||||
&mut *runtime.manager.write().await,
|
||||
"probe",
|
||||
None,
|
||||
Some(7),
|
||||
0,
|
||||
);
|
||||
let (_tx, rx) = mpsc::unbounded_channel();
|
||||
let result = tokio::time::timeout(
|
||||
Duration::from_secs(5),
|
||||
run_subagent(
|
||||
&runtime,
|
||||
"probe".to_string(),
|
||||
FleetRole::Scout,
|
||||
"report".to_string(),
|
||||
SubAgentAssignment::new("report".to_string(), None),
|
||||
Some(vec![]),
|
||||
false,
|
||||
Instant::now(),
|
||||
8,
|
||||
Some(7),
|
||||
None,
|
||||
rx,
|
||||
),
|
||||
)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
server.abort();
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
assert_eq!(
|
||||
result.usage.as_ref().and_then(|usage| usage.total_tokens),
|
||||
Some(15)
|
||||
);
|
||||
let requests = requests.lock().unwrap();
|
||||
assert_eq!(
|
||||
requests.len(),
|
||||
1,
|
||||
"no summary retry after budget exhaustion"
|
||||
);
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.get("max_tokens")
|
||||
.or_else(|| requests[0].get("max_completion_tokens"))
|
||||
.and_then(Value::as_u64),
|
||||
Some(7)
|
||||
);
|
||||
assert!(result.result.as_deref().unwrap().contains("15/7"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[allow(clippy::await_holding_lock)]
|
||||
async fn workflow_scope_precedes_child_admission_and_caps_the_first_request_after_a_yield() {
|
||||
use axum::{Json, Router, routing::post};
|
||||
let _retry = crate::retry_status::test_guard();
|
||||
crate::retry_status::clear();
|
||||
crate::retry_status::clear_rate_limit();
|
||||
let _env = crate::test_support::lock_test_env();
|
||||
let tmp = tempdir().unwrap();
|
||||
let _home = crate::test_support::EnvVarGuard::set("CODEWHALE_HOME", tmp.path().join("state"));
|
||||
let requests = Arc::new(std::sync::Mutex::new(Vec::<Value>::new()));
|
||||
let observed = Arc::clone(&requests);
|
||||
let app = Router::new().route("/{*path}", post(move |Json(request): Json<Value>| {
|
||||
let observed = Arc::clone(&observed);
|
||||
async move {
|
||||
observed.lock().unwrap().push(request);
|
||||
Json(json!({
|
||||
"id": "workflow-budget-probe", "model": "deepseek-v4-flash",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "partial evidence"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
|
||||
}))
|
||||
}
|
||||
}));
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let server = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
|
||||
let config = crate::config::Config {
|
||||
api_key: Some("fixture-key".to_string()),
|
||||
base_url: Some(format!("http://{address}/v1")),
|
||||
..Default::default()
|
||||
};
|
||||
let manager = Arc::new(RwLock::new(
|
||||
SubAgentManager::new(tmp.path().to_path_buf(), 4).with_launch_concurrency(1),
|
||||
));
|
||||
let mut runtime = stub_runtime().with_api_config(config.clone());
|
||||
runtime.client = DeepSeekClient::new(&config).unwrap();
|
||||
runtime.context = ToolContext::new(tmp.path().to_path_buf());
|
||||
runtime.manager = Arc::clone(&manager);
|
||||
let gate = manager.read().await.launch_gate.clone();
|
||||
let permit = gate.acquire_owned().await.unwrap();
|
||||
let request: codewhale_workflow_js::TaskRequest = serde_json::from_value(json!({
|
||||
"description": "report", "subagent_type": "scout", "allowed_tools": [],
|
||||
"worktree": false, "token_budget": 100,
|
||||
}))
|
||||
.unwrap();
|
||||
let identity = WorkflowTaskSpawnIdentity {
|
||||
workflow_run_id: "run-first-request-budget".to_string(),
|
||||
shared_token_budget: Some(12),
|
||||
workflow_phase_id: None,
|
||||
workflow_task_label: None,
|
||||
workflow_child_index: 0,
|
||||
fleet_authority_fingerprint: None,
|
||||
exact_fleet_binding: None,
|
||||
};
|
||||
let spawned = spawn_workflow_task(
|
||||
request.clone(),
|
||||
Arc::clone(&manager),
|
||||
runtime.clone(),
|
||||
identity.clone(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
// Force the already-scheduled child to poll while admission is held. No
|
||||
// caller-side attachment is performed, even after this scheduling boundary.
|
||||
tokio::task::yield_now().await;
|
||||
let handle = {
|
||||
let mut guard = manager.write().await;
|
||||
assert_eq!(
|
||||
guard.budget_scope_state(&spawned.result.agent_id),
|
||||
Some((0, 12))
|
||||
);
|
||||
assert_eq!(
|
||||
guard.remaining_worker_tokens(&spawned.result.agent_id),
|
||||
Some(12)
|
||||
);
|
||||
assert!(requests.lock().unwrap().is_empty());
|
||||
guard
|
||||
.agents
|
||||
.get_mut(&spawned.result.agent_id)
|
||||
.unwrap()
|
||||
.task_handle
|
||||
.take()
|
||||
.unwrap()
|
||||
};
|
||||
drop(permit);
|
||||
tokio::task::yield_now().await;
|
||||
tokio::time::timeout(Duration::from_secs(5), handle)
|
||||
.await
|
||||
.expect("budgeted child settles")
|
||||
.expect("child task succeeds");
|
||||
let count = manager.read().await.worker_records.len();
|
||||
let error = spawn_workflow_task(request, Arc::clone(&manager), runtime, identity)
|
||||
.await
|
||||
.expect_err("an exhausted host pool cannot admit another child");
|
||||
assert!(error.to_string().contains("15/12 tokens spent"), "{error}");
|
||||
let guard = manager.read().await;
|
||||
assert_eq!(
|
||||
guard.worker_records.len(),
|
||||
count,
|
||||
"refusal must precede registration"
|
||||
);
|
||||
assert_eq!(guard.budget_spent_for_scope("run-first-request-budget"), 15);
|
||||
assert_eq!(
|
||||
guard.get_result(&spawned.result.agent_id).unwrap().status,
|
||||
SubAgentStatus::BudgetExhausted
|
||||
);
|
||||
server.abort();
|
||||
let requests = requests.lock().unwrap();
|
||||
assert_eq!(requests.len(), 1, "no request after shared exhaustion");
|
||||
assert_eq!(
|
||||
requests[0]
|
||||
.get("max_tokens")
|
||||
.or_else(|| requests[0].get("max_completion_tokens"))
|
||||
.and_then(Value::as_u64),
|
||||
Some(12)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn workflow_scope_preserves_parent_source_and_per_call_lower_ceilings() {
|
||||
for (parent_cap, source_cap, call_cap, expected) in
|
||||
[(17, 80, 90, 10), (80, 16, 90, 11), (80, 80, 9, 9)]
|
||||
{
|
||||
let tmp = tempdir().unwrap();
|
||||
let manager = Arc::new(RwLock::new(SubAgentManager::new(
|
||||
tmp.path().to_path_buf(),
|
||||
4,
|
||||
)));
|
||||
let mut runtime = stub_runtime().child_runtime();
|
||||
runtime.context = ToolContext::new(tmp.path().to_path_buf());
|
||||
runtime.manager = Arc::clone(&manager);
|
||||
runtime.parent_agent_id = Some("parent".to_string());
|
||||
runtime.cancel_token.cancel(); // inspect admission without a provider request
|
||||
let mut guard = manager.write().await;
|
||||
record(&mut guard, "parent", None, Some(parent_cap), 7);
|
||||
guard.attach_shared_budget_scope("parent", "parent-pool", parent_cap);
|
||||
record(&mut guard, "source", None, Some(source_cap), 5);
|
||||
guard.attach_shared_budget_scope("source", "source-pool", source_cap);
|
||||
let child = guard
|
||||
.spawn_background_with_assignment_options(
|
||||
Arc::clone(&manager),
|
||||
runtime,
|
||||
FleetRole::Scout,
|
||||
"continue".to_string(),
|
||||
SubAgentAssignment::new("continue".to_string(), None),
|
||||
Some(vec![]),
|
||||
SubAgentSpawnOptions {
|
||||
token_budget: Some(call_cap),
|
||||
workflow_budget_scope: Some(("workflow-pool".to_string(), 100)),
|
||||
resume_from_agent_id: Some("source".to_string()),
|
||||
..Default::default()
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(guard.budget_scope_state(&child.agent_id), Some((0, 100)));
|
||||
assert_eq!(
|
||||
guard.remaining_worker_tokens(&child.agent_id),
|
||||
Some(expected)
|
||||
);
|
||||
guard
|
||||
.worker_records
|
||||
.get_mut(&child.agent_id)
|
||||
.unwrap()
|
||||
.usage
|
||||
.total_tokens = Some(expected);
|
||||
assert_eq!(guard.remaining_worker_tokens(&child.agent_id), Some(0));
|
||||
assert_eq!(guard.budget_spent_for_scope("parent-pool"), 7 + expected);
|
||||
assert_eq!(guard.budget_spent_for_scope("source-pool"), 5 + expected);
|
||||
assert_eq!(guard.budget_spent_for_scope("workflow-pool"), expected);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn root_fork_of_depth_two_leaf_cannot_regain_a_generation() {
|
||||
let tmp = tempdir().unwrap();
|
||||
@@ -787,8 +354,7 @@ fn hard_cap_truncates_at_valid_utf8_boundary() {
|
||||
}
|
||||
assert!(content.len() > target);
|
||||
let capped = hard_cap_tool_result(content, cap_tokens(u32::MAX));
|
||||
// Must be valid UTF-8 (no panic during to_string/display).
|
||||
let _ = format!("{capped}");
|
||||
// Must be valid UTF-8 (no panic during slicing or display).
|
||||
assert!(capped.ends_with("\n[truncated: true]"));
|
||||
// The marker is ASCII; verify the rest is still valid UTF-8.
|
||||
let without_marker = &capped[..capped.len() - "\n[truncated: true]".len()];
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -237,7 +237,6 @@ fn make_write_worker_spec(worker_id: &str, workspace: PathBuf, root: &str) -> Ag
|
||||
coordination_contracts: Vec::new(),
|
||||
expected_artifact: Some("tested patch".to_string()),
|
||||
deliverables: Vec::new(),
|
||||
token_budget: None,
|
||||
resume_identity: Some(worker_id.to_string()),
|
||||
generation: 1,
|
||||
resume_from_agent_id: None,
|
||||
@@ -553,7 +552,6 @@ fn worker_record_usage_accumulates_provider_tokens() {
|
||||
reloaded.usage_source_fingerprints,
|
||||
record.usage_source_fingerprints
|
||||
);
|
||||
assert_eq!(record.usage.token_budget, None);
|
||||
assert!(
|
||||
record.usage.note.contains("175 tokens"),
|
||||
"usage note includes reported total: {}",
|
||||
@@ -561,95 +559,6 @@ fn worker_record_usage_accumulates_provider_tokens() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn token_budget_scope_is_shared_across_nested_workers_and_blocks_when_spent() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let workspace = tmp.path().to_path_buf();
|
||||
let mut manager =
|
||||
SubAgentManager::new(workspace.clone(), 4).with_default_token_budget(Some(100));
|
||||
|
||||
manager.register_worker(make_worker_spec("agent_root", workspace.clone()));
|
||||
let root_scope = manager
|
||||
.resolve_spawn_budget_scope("agent_root", None, None)
|
||||
.expect("root budget resolves")
|
||||
.expect("root budget present");
|
||||
manager.attach_budget_scope("agent_root", root_scope);
|
||||
manager.record_worker_usage(
|
||||
"agent_root",
|
||||
"agent_root:response:1",
|
||||
&Usage {
|
||||
input_tokens: 40,
|
||||
output_tokens: 10,
|
||||
..Usage::default()
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
let mut child_spec = make_worker_spec("agent_child", workspace);
|
||||
child_spec.parent_run_id = Some("agent_root".to_string());
|
||||
let child_scope = manager
|
||||
.resolve_spawn_budget_scope("agent_child", Some("agent_root"), None)
|
||||
.expect("child inherits budget")
|
||||
.expect("child budget present");
|
||||
assert_eq!(child_scope.scope_id, "agent_root");
|
||||
assert_eq!(child_scope.limit, 100);
|
||||
assert_eq!(child_scope.spent, 50);
|
||||
for request in [20, 1_000] {
|
||||
let scope = manager
|
||||
.resolve_spawn_budget_scope("agent_override", Some("agent_root"), Some(request))
|
||||
.expect("a remaining inherited pool still permits a narrowed child")
|
||||
.expect("inherited pool remains present");
|
||||
assert_eq!(
|
||||
scope.scope_id, "agent_root",
|
||||
"a per-call cap cannot reset accounting"
|
||||
);
|
||||
assert_eq!(
|
||||
scope.limit, 100,
|
||||
"the local cap does not replace the aggregate ceiling"
|
||||
);
|
||||
assert_eq!(scope.spent, 50);
|
||||
assert_eq!(scope.remaining, 50);
|
||||
}
|
||||
manager.register_worker(child_spec);
|
||||
manager.attach_budget_scope("agent_child", child_scope);
|
||||
manager.record_worker_usage(
|
||||
"agent_child",
|
||||
"agent_child:response:1",
|
||||
&Usage {
|
||||
input_tokens: 30,
|
||||
output_tokens: 20,
|
||||
..Usage::default()
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
||||
let root = manager.get_worker_record("agent_root").expect("root");
|
||||
let child = manager.get_worker_record("agent_child").expect("child");
|
||||
assert_eq!(root.usage.budget_spent_tokens, Some(100));
|
||||
assert_eq!(child.usage.budget_spent_tokens, Some(100));
|
||||
assert_eq!(root.usage.budget_remaining_tokens, Some(0));
|
||||
assert_eq!(child.usage.budget_remaining_tokens, Some(0));
|
||||
assert_eq!(root.usage.status, "budget_exhausted");
|
||||
|
||||
let err = manager
|
||||
.resolve_spawn_budget_scope("agent_grandchild", Some("agent_child"), None)
|
||||
.expect_err("spent shared budget blocks further child spawn");
|
||||
assert!(
|
||||
err.to_string().contains("token budget exhausted"),
|
||||
"actionable exhaustion error: {err}"
|
||||
);
|
||||
|
||||
for request in [20, 1_000] {
|
||||
let error = manager
|
||||
.resolve_spawn_budget_scope("agent_override", Some("agent_child"), Some(request))
|
||||
.expect_err("neither a lower nor a higher local cap can escape an exhausted pool");
|
||||
assert!(
|
||||
error.to_string().contains("scope agent_root: 100/100"),
|
||||
"{error}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn agent_worker_profile_derives_from_parent_without_escalation() {
|
||||
let mut runtime = stub_runtime();
|
||||
@@ -1164,7 +1073,6 @@ fn headless_worker_registration_enforces_live_claims_and_projects_context() {
|
||||
coordination_contracts: Vec::new(),
|
||||
expected_artifact: None,
|
||||
deliverables: Vec::new(),
|
||||
token_budget: None,
|
||||
resume_identity: Some(format!("fleet-{id}")),
|
||||
generation: 1,
|
||||
resume_from_agent_id: None,
|
||||
@@ -2805,7 +2713,6 @@ async fn tool_free_subagent_omits_chat_tools_and_tool_choice() {
|
||||
Instant::now(),
|
||||
1,
|
||||
None,
|
||||
None,
|
||||
input_rx,
|
||||
)
|
||||
.await
|
||||
@@ -3459,7 +3366,6 @@ fn deliberate_spawn_requires_delegation_fields() {
|
||||
}))
|
||||
.expect("deliberate spawn with all fields");
|
||||
assert_eq!(ok.agent_type, FleetRole::Reviewer);
|
||||
assert_eq!(ok.token_budget, None);
|
||||
assert_eq!(ok.write_authority, Some(SpawnWriteAuthority::ReadOnly));
|
||||
assert_eq!(ok.expected_artifact.as_deref(), Some("review findings"));
|
||||
assert!(
|
||||
@@ -5374,35 +5280,6 @@ fn test_delegate_defaults_to_fork_context() {
|
||||
assert_eq!(parsed.fork_context, Some(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn spawn_request_parses_token_budget_override() {
|
||||
let parsed = parse_spawn_request(&json!({
|
||||
"prompt": "fan out safely",
|
||||
"token_budget": 12_345
|
||||
}))
|
||||
.expect("token budget parses");
|
||||
assert_eq!(parsed.token_budget, Some(12_345));
|
||||
|
||||
let parsed = parse_spawn_request(&json!({
|
||||
"prompt": "fleet-shaped alias",
|
||||
"max_tokens": 4_000
|
||||
}))
|
||||
.expect("max_tokens alias parses");
|
||||
assert_eq!(parsed.token_budget, Some(4_000));
|
||||
|
||||
let err = parse_spawn_request(&json!({
|
||||
"prompt": "bad budget",
|
||||
"token_budget": 0
|
||||
}))
|
||||
.expect_err("zero budget is invalid in tool input");
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("token_budget must be an integer greater than zero")
|
||||
&& err.to_string().contains("omit it to inherit"),
|
||||
"clear token budget error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn forked_subagent_messages_preserve_parent_prefix_then_append_task() {
|
||||
let parent_message = Message {
|
||||
@@ -5748,7 +5625,6 @@ fn subagent_tool_schemas_advertise_real_type_and_role_vocabulary() {
|
||||
"prompt",
|
||||
"resume_from",
|
||||
"thinking",
|
||||
"token_budget",
|
||||
"type",
|
||||
"until",
|
||||
"wall_time_secs",
|
||||
@@ -5795,8 +5671,8 @@ fn subagent_tool_schemas_advertise_real_type_and_role_vocabulary() {
|
||||
fn agent_tool_unadvertised_fields_remain_parse_accepted() {
|
||||
// #5324 compat: the fields removed from the advertised schema must stay
|
||||
// parse-accepted and honored unchanged — saved transcripts, ACP/MCP
|
||||
// clients and Fleet configs still replay them. Same contract as the
|
||||
// `token_budget` precedent (docs/SUBAGENTS.md).
|
||||
// clients and Fleet configs still replay them, the same contract as the
|
||||
// retired `token_budget` key (docs/SUBAGENTS.md).
|
||||
let request = parse_spawn_request(&json!({
|
||||
"prompt": "summarize the diff",
|
||||
"model": "deepseek-v4-flash",
|
||||
@@ -5856,14 +5732,14 @@ fn agent_tool_unadvertised_fields_remain_parse_accepted() {
|
||||
.expect_err("read_only plus a declared write scope stays refused");
|
||||
assert!(err.to_string().contains("read_only"), "{err}");
|
||||
|
||||
// token_budget keeps its own long-standing unadvertised-but-accepted
|
||||
// contract.
|
||||
let request = parse_spawn_request(&json!({
|
||||
// token_budget is retired from the runtime (#6189): runs are never
|
||||
// stopped by token accounting. Legacy input carrying it still parses;
|
||||
// the value is ignored, never enforced.
|
||||
parse_spawn_request(&json!({
|
||||
"prompt": "p",
|
||||
"token_budget": 5000,
|
||||
}))
|
||||
.expect("token_budget must stay parse-accepted");
|
||||
assert_eq!(request.token_budget, Some(5000));
|
||||
.expect("retired token_budget key stays parse-tolerated");
|
||||
|
||||
// The removed lifecycle extras are still read on their actions:
|
||||
// action aliases keep parsing, wait still reads timeout_secs, status
|
||||
@@ -9342,7 +9218,6 @@ async fn api_timeout_preserves_checkpoint_and_returns_needs_input_without_parkin
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 3,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -9513,7 +9388,6 @@ async fn subagent_retries_api_timeout_before_succeeding() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 3,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -9661,7 +9535,6 @@ async fn subagent_retries_transient_provider_header_timeout_before_succeeding()
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 3,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -9735,7 +9608,6 @@ async fn subagent_rate_limit_exhaustion_interrupts_with_checkpoint() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 3,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -10880,7 +10752,6 @@ fn explicit_state_roots_isolate_managers_for_the_same_execution_workspace() {
|
||||
2,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
let manager_b = new_shared_subagent_manager_with_state_root_and_timeout(
|
||||
workspace.clone(),
|
||||
@@ -10891,7 +10762,6 @@ fn explicit_state_roots_isolate_managers_for_the_same_execution_workspace() {
|
||||
2,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
for (manager, expected_root, worker_id) in [
|
||||
@@ -11559,7 +11429,7 @@ fn budget_exhaustion_is_a_high_priority_failure_event() {
|
||||
assert!(
|
||||
completion
|
||||
.payload
|
||||
.contains(r#""failure_class":"token_budget""#)
|
||||
.contains(r#""failure_class":"budget_exhausted""#)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -13944,7 +13814,6 @@ async fn turn_end_parking_preserves_a_step_zero_resumable_checkpoint() {
|
||||
false,
|
||||
Instant::now(),
|
||||
1,
|
||||
None,
|
||||
Some(parking),
|
||||
input_rx,
|
||||
)
|
||||
@@ -15436,7 +15305,6 @@ async fn run_subagent_task_claims_before_delivery_and_then_finalizes() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 1,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -15523,7 +15391,6 @@ async fn cancellation_wins_task_race_but_still_fans_in_exactly_once() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 1,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -15727,7 +15594,6 @@ async fn fatal_provider_failure_mid_run_parks_a_continuable_checkpoint() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 4,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -15828,7 +15694,6 @@ async fn fatal_provider_failure_mid_run_parks_a_continuable_checkpoint() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 2,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: resume_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -15988,7 +15853,6 @@ async fn repeated_typed_denials_stop_worker_as_failed_not_budget() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 20,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -16075,7 +15939,6 @@ async fn non_retryable_provider_failure_fans_in_to_every_terminal_sink() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 1,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -16255,7 +16118,7 @@ fn subagent_budget_exhaustion_completion_carries_budget_exhausted_sentinel() {
|
||||
assert_eq!(parsed["event"], "subagent.failed");
|
||||
assert_eq!(parsed["priority"], "high");
|
||||
assert_eq!(parsed["status"], "budget_exhausted");
|
||||
assert_eq!(parsed["failure_class"], "token_budget");
|
||||
assert_eq!(parsed["failure_class"], "budget_exhausted");
|
||||
assert_eq!(parsed["error_location"], "previous_line");
|
||||
}
|
||||
|
||||
@@ -16887,7 +16750,6 @@ async fn launch_gate_queues_extra_direct_children() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 1,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx,
|
||||
launch_gate: gate,
|
||||
@@ -17044,7 +16906,6 @@ async fn queued_turn_owned_child_parks_without_a_false_start_transition() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 1,
|
||||
token_budget: None,
|
||||
wall_time: Duration::from_secs(5),
|
||||
input_rx,
|
||||
launch_gate: Some(Arc::clone(&gate)),
|
||||
@@ -17195,7 +17056,6 @@ async fn launch_gate_wait_counts_against_child_wall_timeout() {
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps: 1,
|
||||
token_budget: None,
|
||||
wall_time: WALL_TIME,
|
||||
input_rx,
|
||||
launch_gate: Some(Arc::clone(&gate)),
|
||||
@@ -17417,7 +17277,6 @@ pub(super) async fn run_incomplete_response_worker(
|
||||
workspace: &Path,
|
||||
stop_reason: &str,
|
||||
max_steps: u32,
|
||||
token_budget: Option<u64>,
|
||||
) -> (
|
||||
SubAgentResult,
|
||||
Arc<AtomicUsize>,
|
||||
@@ -17471,7 +17330,6 @@ pub(super) async fn run_incomplete_response_worker(
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps,
|
||||
token_budget,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -17510,7 +17368,7 @@ fn assert_partial_tool_was_not_executed(messages: &[MailboxMessage]) {
|
||||
async fn output_limit_on_last_step_preserves_partial_text_and_exact_cause() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let (result, calls, mailbox, total_tokens) =
|
||||
run_incomplete_response_worker(tmp.path(), "length", 1, None).await;
|
||||
run_incomplete_response_worker(tmp.path(), "length", 1).await;
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
let SubAgentStatus::Failed(reason) = &result.status else {
|
||||
@@ -17542,21 +17400,17 @@ async fn output_limit_on_last_step_preserves_partial_text_and_exact_cause() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn output_limit_cause_wins_over_generic_token_budget() {
|
||||
async fn output_limit_incomplete_response_is_typed_and_records_usage() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let (result, calls, mailbox, total_tokens) =
|
||||
run_incomplete_response_worker(tmp.path(), "max_tokens", 4, Some(10)).await;
|
||||
run_incomplete_response_worker(tmp.path(), "max_tokens", 4).await;
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
let reason = &result
|
||||
.checkpoint
|
||||
.as_ref()
|
||||
.expect("budget checkpoint")
|
||||
.reason;
|
||||
let SubAgentStatus::Failed(reason) = &result.status else {
|
||||
panic!("incomplete output must fail, got {:?}", result.status);
|
||||
};
|
||||
assert!(reason.contains("output was truncated"), "{reason}");
|
||||
assert!(reason.contains("`max_tokens`"), "{reason}");
|
||||
assert!(reason.contains("token budget exhausted"), "{reason}");
|
||||
assert!(
|
||||
result
|
||||
.result
|
||||
@@ -17569,21 +17423,17 @@ async fn output_limit_cause_wins_over_generic_token_budget() {
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn non_output_incomplete_cause_wins_over_generic_token_budget() {
|
||||
async fn non_output_incomplete_response_is_typed_and_records_usage() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let (result, calls, mailbox, total_tokens) =
|
||||
run_incomplete_response_worker(tmp.path(), "incomplete:content_filter", 4, Some(10)).await;
|
||||
run_incomplete_response_worker(tmp.path(), "incomplete:content_filter", 4).await;
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
assert_eq!(result.status, SubAgentStatus::BudgetExhausted);
|
||||
let reason = &result
|
||||
.checkpoint
|
||||
.as_ref()
|
||||
.expect("budget checkpoint")
|
||||
.reason;
|
||||
let SubAgentStatus::Failed(reason) = &result.status else {
|
||||
panic!("incomplete response must fail, got {:?}", result.status);
|
||||
};
|
||||
assert!(reason.contains("response was incomplete"), "{reason}");
|
||||
assert!(reason.contains("`content_filter`"), "{reason}");
|
||||
assert!(reason.contains("token budget exhausted"), "{reason}");
|
||||
assert!(
|
||||
result
|
||||
.result
|
||||
@@ -17599,7 +17449,7 @@ async fn non_output_incomplete_cause_wins_over_generic_token_budget() {
|
||||
async fn output_limit_never_retries_or_executes_partial_tools() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let (result, calls, mailbox, total_tokens) =
|
||||
run_incomplete_response_worker(tmp.path(), "length", 2, Some(100)).await;
|
||||
run_incomplete_response_worker(tmp.path(), "length", 2).await;
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
let SubAgentStatus::Failed(reason) = &result.status else {
|
||||
@@ -17622,7 +17472,6 @@ async fn spawn_budget_capped_worker(
|
||||
workspace: &Path,
|
||||
prompt_tokens: u64,
|
||||
completion_tokens: u64,
|
||||
token_budget: Option<u64>,
|
||||
max_steps: u32,
|
||||
wall_time: Duration,
|
||||
) -> (
|
||||
@@ -17673,7 +17522,6 @@ async fn spawn_budget_capped_worker(
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps,
|
||||
token_budget,
|
||||
wall_time,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
@@ -17687,7 +17535,7 @@ async fn spawn_budget_capped_worker(
|
||||
async fn worker_stops_with_typed_wall_time_reason() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let (manager, agent_id, _calls, task_handle) =
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, None, 120, Duration::from_millis(1)).await;
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, 120, Duration::from_millis(1)).await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(5), task_handle)
|
||||
.await
|
||||
@@ -17706,257 +17554,10 @@ async fn worker_stops_with_typed_wall_time_reason() {
|
||||
.expect("wall-budget checkpoint")
|
||||
.reason;
|
||||
assert!(reason.contains("wall-time budget exhausted"), "{reason}");
|
||||
assert!(reason.contains("limit:"), "{reason}");
|
||||
assert!(reason.contains("wall-time limit"), "{reason}");
|
||||
assert!(reason.contains("operator"), "{reason}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_stops_when_per_worker_token_budget_exceeded() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
// 100 tokens/turn (60 in + 40 out) vs a 50-token cap: the worker must
|
||||
// stop with `BudgetExhausted` after its very first model turn instead of
|
||||
// running on to `max_steps`.
|
||||
let (manager, agent_id, calls, task_handle) =
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, Some(50), 4, DEFAULT_CHILD_WALL_TIME).await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(5), task_handle)
|
||||
.await
|
||||
.expect("budget-capped worker must terminate")
|
||||
.expect("task should finish");
|
||||
|
||||
assert_eq!(
|
||||
calls.load(Ordering::SeqCst),
|
||||
1,
|
||||
"worker must stop after the first over-budget turn, not run to max_steps"
|
||||
);
|
||||
|
||||
let result = {
|
||||
let manager = manager.read().await;
|
||||
manager.get_result(&agent_id).expect("agent registered")
|
||||
};
|
||||
assert!(
|
||||
matches!(result.status, SubAgentStatus::BudgetExhausted),
|
||||
"expected BudgetExhausted, got {:?}",
|
||||
result.status
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_without_per_worker_token_budget_runs_to_completion() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
// No per-worker cap: a final-text response completes the worker normally
|
||||
// even though each turn reports 100 tokens.
|
||||
let (manager, agent_id, calls, task_handle) =
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, None, 4, DEFAULT_CHILD_WALL_TIME).await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(5), task_handle)
|
||||
.await
|
||||
.expect("uncapped worker must terminate")
|
||||
.expect("task should finish");
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
|
||||
let result = {
|
||||
let manager = manager.read().await;
|
||||
manager.get_result(&agent_id).expect("agent registered")
|
||||
};
|
||||
assert!(
|
||||
matches!(result.status, SubAgentStatus::Completed),
|
||||
"uncapped worker should complete normally, got {:?}",
|
||||
result.status
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn per_worker_token_budget_does_not_double_count_scope_accounting() {
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
// The per-worker runtime cap stops the worker, but the scope-level
|
||||
// accounting (#3319 `aggregate_budget_spent` sums worker_records'
|
||||
// `total_tokens`) must reflect the tokens actually consumed exactly once
|
||||
// — never inflated by the runtime accumulator that triggered the stop.
|
||||
let (manager, agent_id, calls, task_handle) =
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, Some(50), 4, DEFAULT_CHILD_WALL_TIME).await;
|
||||
|
||||
tokio::time::timeout(Duration::from_secs(5), task_handle)
|
||||
.await
|
||||
.expect("budget-capped worker must terminate")
|
||||
.expect("task should finish");
|
||||
|
||||
assert_eq!(calls.load(Ordering::SeqCst), 1);
|
||||
|
||||
let (result, worker_record) = {
|
||||
let manager = manager.read().await;
|
||||
(
|
||||
manager.get_result(&agent_id).expect("agent registered"),
|
||||
manager.get_worker_record(&agent_id).expect("worker record"),
|
||||
)
|
||||
};
|
||||
assert!(
|
||||
matches!(result.status, SubAgentStatus::BudgetExhausted),
|
||||
"expected BudgetExhausted, got {:?}",
|
||||
result.status
|
||||
);
|
||||
// One turn of 60 in + 40 out = 100 tokens, counted exactly once.
|
||||
assert_eq!(
|
||||
worker_record.usage.total_tokens,
|
||||
Some(100),
|
||||
"scope accounting must equal the single turn's tokens, not double-count: {:?}",
|
||||
worker_record.usage
|
||||
);
|
||||
}
|
||||
|
||||
/// Variant of [`spawn_budget_capped_worker`] that attaches the worker to a
|
||||
/// shared workflow budget scope before its first model turn (no per-worker
|
||||
/// cap), returning the manager, agent id, call counter, and task handle.
|
||||
// Test helper: the eight parameters mirror the distinct knobs each test case
|
||||
// tunes; grouping them into a struct would add boilerplate at every call site
|
||||
// without improving readability.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn spawn_scope_budgeted_worker(
|
||||
manager: &Arc<RwLock<SubAgentManager>>,
|
||||
workspace: &Path,
|
||||
agent_id: &str,
|
||||
scope_id: &str,
|
||||
scope_limit: u64,
|
||||
prompt_tokens: u64,
|
||||
completion_tokens: u64,
|
||||
max_steps: u32,
|
||||
) -> (Arc<AtomicUsize>, tokio::task::JoinHandle<()>) {
|
||||
let agent_id = agent_id.to_string();
|
||||
let (task_input_tx, task_input_rx) = mpsc::unbounded_channel();
|
||||
let agent = SubAgent::new(
|
||||
agent_id.clone(),
|
||||
FleetRole::Worker,
|
||||
"Work within shared budget".to_string(),
|
||||
make_assignment(),
|
||||
"deepseek-v4-flash".to_string(),
|
||||
Some("Budget".to_string()),
|
||||
Some(vec![]),
|
||||
task_input_tx,
|
||||
workspace.to_path_buf(),
|
||||
"boot_scope_budget".to_string(),
|
||||
);
|
||||
{
|
||||
let mut manager = manager.write().await;
|
||||
manager.agents.insert(agent_id.clone(), agent);
|
||||
manager.register_worker(make_worker_spec(&agent_id, workspace.to_path_buf()));
|
||||
manager.attach_shared_budget_scope(&agent_id, scope_id, scope_limit);
|
||||
}
|
||||
|
||||
let (client, calls) =
|
||||
token_heavy_chat_client(prompt_tokens, completion_tokens, "partial answer").await;
|
||||
let mut runtime = stub_runtime();
|
||||
runtime.client = client;
|
||||
runtime.manager = Arc::clone(manager);
|
||||
runtime.context = ToolContext::new(workspace.to_path_buf());
|
||||
|
||||
let task = SubAgentTask {
|
||||
manager_handle: Arc::clone(manager),
|
||||
runtime,
|
||||
agent_id: agent_id.clone(),
|
||||
agent_type: FleetRole::Worker,
|
||||
prompt: "Work within shared budget".to_string(),
|
||||
assignment: make_assignment(),
|
||||
allowed_tools: Some(vec![]),
|
||||
fork_context: false,
|
||||
started_at: Instant::now(),
|
||||
max_steps,
|
||||
token_budget: None,
|
||||
wall_time: DEFAULT_CHILD_WALL_TIME,
|
||||
input_rx: task_input_rx,
|
||||
launch_gate: None,
|
||||
_foreground_child_registration: None,
|
||||
};
|
||||
let task_handle = tokio::spawn(run_subagent_task(task));
|
||||
(calls, task_handle)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shared_scope_budget_stops_admitted_children_mid_run() {
|
||||
// A workflow run's token_budget must be a collective ceiling for the
|
||||
// children it admitted, not just an admission gate for future spawns:
|
||||
// children that attach while the scope has room used to run uncapped, so
|
||||
// a fan-out could burn many times the budget and still report Completed.
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let manager = Arc::new(RwLock::new(SubAgentManager::new(
|
||||
tmp.path().to_path_buf(),
|
||||
4,
|
||||
)));
|
||||
let scope_id = "run-budget-ceiling";
|
||||
// Each model turn burns 100 tokens (60 in + 40 out); the run-level budget
|
||||
// leaves room for exactly one full turn across ALL children.
|
||||
let scope_limit = 150;
|
||||
|
||||
// First child: admitted while the scope is empty, burns its 100 and
|
||||
// completes normally.
|
||||
let (calls_a, handle_a) = spawn_scope_budgeted_worker(
|
||||
&manager,
|
||||
tmp.path(),
|
||||
"agent_scope_a",
|
||||
scope_id,
|
||||
scope_limit,
|
||||
60,
|
||||
40,
|
||||
4,
|
||||
)
|
||||
.await;
|
||||
tokio::time::timeout(Duration::from_secs(5), handle_a)
|
||||
.await
|
||||
.expect("first child must terminate")
|
||||
.expect("task should finish");
|
||||
assert_eq!(calls_a.load(Ordering::SeqCst), 1);
|
||||
let status_a = manager
|
||||
.read()
|
||||
.await
|
||||
.get_result("agent_scope_a")
|
||||
.expect("first child registered")
|
||||
.status;
|
||||
assert!(
|
||||
matches!(status_a, SubAgentStatus::Completed),
|
||||
"first child completes inside the shared budget, got {status_a:?}"
|
||||
);
|
||||
|
||||
// Second child: also admitted without a per-worker cap (remaining 50 >=
|
||||
// the spawn reserve). Its first turn pushes the shared scope to 200/150,
|
||||
// so it must stop with BudgetExhausted right after that turn instead of
|
||||
// completing or running on to max_steps.
|
||||
let (calls_b, handle_b) = spawn_scope_budgeted_worker(
|
||||
&manager,
|
||||
tmp.path(),
|
||||
"agent_scope_b",
|
||||
scope_id,
|
||||
scope_limit,
|
||||
60,
|
||||
40,
|
||||
4,
|
||||
)
|
||||
.await;
|
||||
tokio::time::timeout(Duration::from_secs(5), handle_b)
|
||||
.await
|
||||
.expect("second child must terminate")
|
||||
.expect("task should finish");
|
||||
assert_eq!(
|
||||
calls_b.load(Ordering::SeqCst),
|
||||
1,
|
||||
"second child must stop after the turn that crossed the shared budget"
|
||||
);
|
||||
let status_b = manager
|
||||
.read()
|
||||
.await
|
||||
.get_result("agent_scope_b")
|
||||
.expect("second child registered")
|
||||
.status;
|
||||
assert!(
|
||||
matches!(status_b, SubAgentStatus::BudgetExhausted),
|
||||
"second child must hit the shared ceiling, got {status_b:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
manager.read().await.budget_spent_for_scope(scope_id),
|
||||
200,
|
||||
"collective spend is accounted once per child"
|
||||
);
|
||||
}
|
||||
|
||||
/// Clears the process-wide rate-limit window on drop so a panicking test
|
||||
/// body cannot leak a live pause into concurrently running tests.
|
||||
struct ClearRateLimitOnDrop;
|
||||
@@ -17983,7 +17584,7 @@ async fn worker_is_not_stranded_by_transient_global_rate_limit_window() {
|
||||
|
||||
let tmp = tempdir().expect("tempdir");
|
||||
let (manager, agent_id, _calls, task_handle) =
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, Some(50), 4, DEFAULT_CHILD_WALL_TIME).await;
|
||||
spawn_budget_capped_worker(tmp.path(), 60, 40, 4, DEFAULT_CHILD_WALL_TIME).await;
|
||||
|
||||
// Simulate the concurrent test finishing: the window closes shortly
|
||||
// after the worker's first request has already observed it.
|
||||
@@ -18001,11 +17602,11 @@ async fn worker_is_not_stranded_by_transient_global_rate_limit_window() {
|
||||
let manager = manager.read().await;
|
||||
manager.get_result(&agent_id).expect("agent registered")
|
||||
};
|
||||
assert!(
|
||||
matches!(result.status, SubAgentStatus::BudgetExhausted),
|
||||
"expected BudgetExhausted, got {:?}",
|
||||
result.status
|
||||
);
|
||||
// Token budgets are tracked, never enforced (#6298 slice 5), so a
|
||||
// worker that used to stop at its token ceiling now runs to
|
||||
// completion. The regression under test is only that an already-cleared
|
||||
// rate-limit window cannot strand the worker past the 5s timeout.
|
||||
assert_eq!(result.status, SubAgentStatus::Completed);
|
||||
}
|
||||
|
||||
/// #4217: terminal worker records must age out of the persisted ledger so
|
||||
@@ -18553,7 +18154,6 @@ fn coordination_process_lock_rejects_second_process() {
|
||||
4,
|
||||
Duration::from_secs(30),
|
||||
4,
|
||||
None,
|
||||
);
|
||||
if role == "holder" {
|
||||
manager
|
||||
|
||||
@@ -172,13 +172,13 @@ fn compact_receipt_bounds_long_labels_and_declared_outputs_without_hiding_limits
|
||||
"follow_up": {"tool": "agent", "agent_id": "agent_bounded_receipt", "session_name": "🐋".repeat(4000)},
|
||||
"usage": {"status": "unknown", "note": "No provider receipt yet"},
|
||||
"worker_record": {"spec": {
|
||||
"runtime_profile": {"spawn_depth": 1, "max_spawn_depth": 2, "max_steps": 8, "token_budget": 500, "wall_time_secs": 30, "wall_deadline_ms": 10000},
|
||||
"runtime_profile": {"spawn_depth": 1, "max_spawn_depth": 2, "max_steps": 8, "wall_time_secs": 30, "wall_deadline_ms": 10000},
|
||||
"launch_manifest": {"deliverables": paths}
|
||||
}}
|
||||
});
|
||||
compact_spawn_receipt(&mut receipt, false);
|
||||
assert!(serde_json::to_vec(&receipt).unwrap().len() <= lifecycle::COMPACT_SPAWN_BYTES);
|
||||
assert_eq!(receipt["effective_limits"]["token_budget"], 500);
|
||||
assert_eq!(receipt["effective_limits"]["max_steps"], 8);
|
||||
assert_eq!(receipt["effective_limits"]["wall_deadline_ms"], 10000);
|
||||
assert_eq!(receipt["child_route"]["truncated"], true);
|
||||
let shown = receipt["deliverables"].as_array().unwrap();
|
||||
|
||||
@@ -3912,15 +3912,12 @@ impl SubAgentWorkflowDriver {
|
||||
}
|
||||
|
||||
fn current_budget_snapshot(&self) -> BudgetSnapshot {
|
||||
let spent = self
|
||||
.manager
|
||||
.try_read()
|
||||
.ok()
|
||||
.map(|manager| manager.budget_spent_for_scope(&self.run_id))
|
||||
.unwrap_or(0);
|
||||
// Token-budget enforcement was removed (#6189): nothing stops a run.
|
||||
// The snapshot survives for the declared spec ceiling only; spent is
|
||||
// no longer tracked per scope.
|
||||
BudgetSnapshot {
|
||||
total: self.total_budget,
|
||||
spent,
|
||||
spent: 0,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4621,7 +4618,6 @@ impl SubAgentWorkflowDriver {
|
||||
.map(str::to_string);
|
||||
let identity = WorkflowTaskSpawnIdentity {
|
||||
workflow_run_id: self.run_id.clone(),
|
||||
shared_token_budget: self.total_budget,
|
||||
workflow_phase_id,
|
||||
workflow_task_label,
|
||||
workflow_child_index,
|
||||
|
||||
@@ -14,10 +14,6 @@ fn usage_none() -> AgentRunUsage {
|
||||
output_tokens: None,
|
||||
total_tokens: None,
|
||||
cost_microusd: None,
|
||||
token_budget: None,
|
||||
budget_spent_tokens: None,
|
||||
budget_remaining_tokens: None,
|
||||
budget_scope: None,
|
||||
note: "no route audit for this worker".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -935,7 +935,6 @@ pub(crate) fn build_engine_config(app: &App, config: &Config) -> EngineConfig {
|
||||
},
|
||||
),
|
||||
max_spawn_depth: config.subagent_max_spawn_depth_for_provider(provider),
|
||||
subagent_token_budget: config.subagent_token_budget_for_provider(provider),
|
||||
allowed_tools: app.active_allowed_tools.clone(),
|
||||
disallowed_tools: None,
|
||||
max_tool_calls: None,
|
||||
|
||||
@@ -376,9 +376,6 @@ pub struct WorkerRuntimeProfile {
|
||||
pub max_spawn_depth: u32,
|
||||
#[serde(default)]
|
||||
pub spawn_depth: u32,
|
||||
/// Measured input plus output tokens for this worker and its descendants.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub token_budget: Option<u64>,
|
||||
/// Whole-run wall time, including queued, model and tool work.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub wall_time_secs: Option<u64>,
|
||||
@@ -470,7 +467,6 @@ impl WorkerRuntimeProfile {
|
||||
denied_tools: Vec::new(),
|
||||
max_spawn_depth: codewhale_config::DEFAULT_SPAWN_DEPTH,
|
||||
spawn_depth: 0,
|
||||
token_budget: None,
|
||||
wall_time_secs: None,
|
||||
wall_deadline_ms: None,
|
||||
max_steps: Self::default_max_steps(role.clone()),
|
||||
@@ -539,7 +535,6 @@ impl WorkerRuntimeProfile {
|
||||
denied_tools,
|
||||
max_spawn_depth,
|
||||
spawn_depth: self.spawn_depth.saturating_add(1),
|
||||
token_budget: narrow_optional_limit(self.token_budget, requested.token_budget),
|
||||
wall_time_secs: narrow_optional_limit(self.wall_time_secs, requested.wall_time_secs),
|
||||
wall_deadline_ms: narrow_optional_limit(
|
||||
self.wall_deadline_ms,
|
||||
@@ -614,7 +609,6 @@ pub struct ChildLaunchManifest {
|
||||
pub expected_artifact: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub deliverables: Vec<String>,
|
||||
pub token_budget: Option<u64>,
|
||||
pub resume_identity: Option<String>,
|
||||
#[serde(default)]
|
||||
pub generation: u32,
|
||||
|
||||
Reference in new issue
Block a user