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:
CodeWhale Bot committed 2026-09-17 06:16:32 -07:00
1 parent ab1420c51e
commit a7a8bdb338
23 files changed
+189 -2096

No files matched your search

-29
View File
@@ -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
-29
View File
@@ -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
-7
View File
@@ -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 {
-2
View File
@@ -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
-4
View File
@@ -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,
-2
View File
@@ -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(),
-1
View File
@@ -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,
)
}
-4
View File
@@ -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(),
-1
View File
@@ -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
+13 -154
View File
@@ -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,
+2 -8
View File
@@ -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());
+11 -445
View File
@@ -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
+29 -429
View File
@@ -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();
+4 -8
View File
@@ -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,
-4
View File
@@ -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(),
}
}
-1
View File
@@ -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,
-6
View File
@@ -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,