From 4761851ff35c4ebdd35eb8801e1180a0a50fef60 Mon Sep 17 00:00:00 2001 From: Eric Traut Date: Thu, 27 Aug 2026 15:39:26 +0000 Subject: [PATCH] Account subagent token usage toward root goals (#41183) ## What changed - Roll token usage from spawned descendants, including nested subagents, into the root goal's usage. - Apply descendant usage during active and idle progress accounting so it contributes to token budgets. - Reset descendant accounting baselines when the active goal changes and preserve usage recorded concurrently with a checkpoint. ## Testing - Cover child and grandchild usage, budget exhaustion, unloaded parent runtimes, goal replacement, idle accounting, and concurrent checkpoints. GitOrigin-RevId: 8f97ec6778c55b9adf94b887b5fd03999ed0eb94 --- codex-rs/ext/goal/src/accounting.rs | 49 ++- codex-rs/ext/goal/src/api.rs | 2 +- codex-rs/ext/goal/src/extension.rs | 35 ++- codex-rs/ext/goal/src/runtime.rs | 9 +- codex-rs/ext/goal/tests/accounting.rs | 42 +++ .../ext/goal/tests/goal_extension_backend.rs | 279 +++++++++++++++++- 6 files changed, 398 insertions(+), 18 deletions(-) diff --git a/codex-rs/ext/goal/src/accounting.rs b/codex-rs/ext/goal/src/accounting.rs index ea942ea158..4012a9587b 100644 --- a/codex-rs/ext/goal/src/accounting.rs +++ b/codex-rs/ext/goal/src/accounting.rs @@ -4,6 +4,8 @@ use codex_state::ThreadGoalStatus; use std::collections::HashMap; use std::sync::Mutex; use std::sync::PoisonError; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use tokio::sync::Semaphore; @@ -13,6 +15,7 @@ use tokio::sync::SemaphorePermit; pub(crate) struct GoalAccountingState { inner: Mutex, progress_accounting_lock: Semaphore, + descendant_token_usage: AtomicI64, } #[derive(Debug)] @@ -21,6 +24,7 @@ struct GoalAccountingInner { turns: HashMap, wall_clock: GoalWallClockAccounting, budget_limit_reported_goal_id: Option, + last_accounted_descendant_token_usage: i64, } #[derive(Debug)] @@ -40,6 +44,7 @@ struct GoalWallClockAccounting { #[derive(Debug, Clone)] pub(crate) struct GoalProgressSnapshot { pub(crate) current_token_usage: TokenUsage, + current_descendant_token_usage: i64, pub(crate) expected_goal_id: String, pub(crate) time_delta_seconds: i64, pub(crate) token_delta: i64, @@ -47,8 +52,10 @@ pub(crate) struct GoalProgressSnapshot { #[derive(Debug, Clone)] pub(crate) struct IdleGoalProgressSnapshot { + current_descendant_token_usage: i64, pub(crate) expected_goal_id: String, pub(crate) time_delta_seconds: i64, + pub(crate) token_delta: i64, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -131,6 +138,14 @@ impl GoalAccountingState { }) } + pub(crate) fn record_descendant_token_usage(&self, usage: &TokenUsage) { + let delta = goal_token_delta_for_usage(usage); + if delta > 0 { + self.descendant_token_usage + .fetch_add(delta, Ordering::Relaxed); + } + } + pub(crate) fn mark_turn_goal_active(&self, turn_id: &str, goal_id: impl Into) { let mut inner = self.inner(); let goal_id = goal_id.into(); @@ -140,6 +155,10 @@ impl GoalAccountingState { if let Some(turn) = inner.turns.get_mut(turn_id) { turn.active_goal_id = Some(goal_id.clone()); if inner.current_turn_id.as_deref() == Some(turn_id) { + if inner.wall_clock.active_goal_id.as_deref() != Some(goal_id.as_str()) { + inner.last_accounted_descendant_token_usage = + self.descendant_token_usage.load(Ordering::Relaxed); + } inner.wall_clock.mark_active_goal(goal_id); } } @@ -155,9 +174,14 @@ impl GoalAccountingState { if inner.budget_limit_reported_goal_id.as_deref() != Some(goal_id.as_str()) { inner.budget_limit_reported_goal_id = None; } + let goal_changed = inner.wall_clock.active_goal_id.as_deref() != Some(goal_id.as_str()); let turn = inner.turns.get_mut(turn_id.as_str())?; turn.active_goal_id = Some(goal_id.clone()); - turn.reset_baseline_to_current(); + if goal_changed { + turn.reset_baseline_to_current(); + inner.last_accounted_descendant_token_usage = + self.descendant_token_usage.load(Ordering::Relaxed); + } inner.wall_clock.mark_active_goal(goal_id); Some(turn_id) } @@ -168,6 +192,10 @@ impl GoalAccountingState { if inner.budget_limit_reported_goal_id.as_deref() != Some(goal_id.as_str()) { inner.budget_limit_reported_goal_id = None; } + if inner.wall_clock.active_goal_id.as_deref() != Some(goal_id.as_str()) { + inner.last_accounted_descendant_token_usage = + self.descendant_token_usage.load(Ordering::Relaxed); + } inner.wall_clock.mark_active_goal(goal_id); } @@ -200,7 +228,12 @@ impl GoalAccountingState { return None; } let expected_goal_id = turn.active_goal_id()?; - let token_delta = turn.token_delta_since_last_accounting(); + let current_descendant_token_usage = self.descendant_token_usage.load(Ordering::Relaxed); + let descendant_token_delta = current_descendant_token_usage + .saturating_sub(inner.last_accounted_descendant_token_usage); + let token_delta = turn + .token_delta_since_last_accounting() + .saturating_add(descendant_token_delta); let time_delta_seconds = if inner.wall_clock.active_goal_id.as_deref() == Some(expected_goal_id.as_str()) { inner.wall_clock.time_delta_since_last_accounting() @@ -212,6 +245,7 @@ impl GoalAccountingState { } Some(GoalProgressSnapshot { current_token_usage: turn.current_token_usage.clone(), + current_descendant_token_usage, expected_goal_id, time_delta_seconds, token_delta, @@ -222,12 +256,17 @@ impl GoalAccountingState { let inner = self.inner(); let expected_goal_id = inner.wall_clock.active_goal_id.clone()?; let time_delta_seconds = inner.wall_clock.time_delta_since_last_accounting(); - if time_delta_seconds == 0 { + let current_descendant_token_usage = self.descendant_token_usage.load(Ordering::Relaxed); + let token_delta = current_descendant_token_usage + .saturating_sub(inner.last_accounted_descendant_token_usage); + if time_delta_seconds == 0 && token_delta <= 0 { return None; } Some(IdleGoalProgressSnapshot { + current_descendant_token_usage, expected_goal_id, time_delta_seconds, + token_delta, }) } @@ -246,6 +285,7 @@ impl GoalAccountingState { turn.active_goal_id = None; } } + inner.last_accounted_descendant_token_usage = snapshot.current_descendant_token_usage; inner.wall_clock.mark_accounted(snapshot.time_delta_seconds); if clear_active_goal { inner.wall_clock.clear_active_goal(); @@ -271,6 +311,7 @@ impl GoalAccountingState { ) { let clear_active_goal = should_clear_active_goal(status, budget_limited_goal_disposition); let mut inner = self.inner(); + inner.last_accounted_descendant_token_usage = snapshot.current_descendant_token_usage; inner.wall_clock.mark_accounted(snapshot.time_delta_seconds); if clear_active_goal { inner.wall_clock.clear_active_goal(); @@ -306,6 +347,7 @@ impl Default for GoalAccountingState { Self { inner: Mutex::new(GoalAccountingInner::default()), progress_accounting_lock: Semaphore::new(/*permits*/ 1), + descendant_token_usage: AtomicI64::new(0), } } } @@ -343,6 +385,7 @@ impl Default for GoalAccountingInner { turns: HashMap::new(), wall_clock: GoalWallClockAccounting::new(), budget_limit_reported_goal_id: None, + last_accounted_descendant_token_usage: 0, } } } diff --git a/codex-rs/ext/goal/src/api.rs b/codex-rs/ext/goal/src/api.rs index 149c95d0a1..94397d55b4 100644 --- a/codex-rs/ext/goal/src/api.rs +++ b/codex-rs/ext/goal/src/api.rs @@ -345,7 +345,7 @@ impl GoalService { } } - fn runtime_for_thread(&self, thread_id: ThreadId) -> Option> { + pub(crate) fn runtime_for_thread(&self, thread_id: ThreadId) -> Option> { let key = thread_id.to_string(); let mut runtimes = self.runtimes(); let runtime = runtimes.get(&key).and_then(Weak::upgrade); diff --git a/codex-rs/ext/goal/src/extension.rs b/codex-rs/ext/goal/src/extension.rs index 24d515c613..349652da4f 100644 --- a/codex-rs/ext/goal/src/extension.rs +++ b/codex-rs/ext/goal/src/extension.rs @@ -110,6 +110,30 @@ where let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { return; }; + let root_accounting_state = input + .session_source + .parent_thread_id() + .or_else(|| { + ThreadId::from_string(input.session_store.level_id()) + .ok() + .filter(|_| input.session_source.is_non_root_agent()) + }) + .and_then(|parent_thread_id| { + self.goal_service + .runtime_for_thread(parent_thread_id) + .or_else(|| { + ThreadId::from_string(input.session_store.level_id()) + .ok() + .and_then(|root_thread_id| { + self.goal_service.runtime_for_thread(root_thread_id) + }) + }) + }) + .map(|parent| { + parent + .root_accounting_state() + .unwrap_or_else(|| parent.accounting_state()) + }); let runtime = input.thread_store.get_or_init::(|| { GoalRuntimeHandle::new( thread_id, @@ -122,6 +146,7 @@ where analytics: self.analytics.clone(), enabled, tools_available_for_thread, + root_accounting_state, }, ) }); @@ -346,12 +371,12 @@ where return; } - let Some(_recorded) = runtime + if let Some(root_accounting_state) = runtime.root_accounting_state() { + root_accounting_state.record_descendant_token_usage(&token_usage.last_token_usage); + } + let _ = runtime .accounting_state() - .record_token_usage(turn_store.level_id(), &token_usage.total_token_usage) - else { - return; - }; + .record_token_usage(turn_store.level_id(), &token_usage.total_token_usage); }) } } diff --git a/codex-rs/ext/goal/src/runtime.rs b/codex-rs/ext/goal/src/runtime.rs index d0579090fd..d9f0f0df24 100644 --- a/codex-rs/ext/goal/src/runtime.rs +++ b/codex-rs/ext/goal/src/runtime.rs @@ -33,6 +33,7 @@ pub(crate) struct GoalRuntimeConfig { pub(crate) analytics: GoalAnalytics, pub(crate) enabled: bool, pub(crate) tools_available_for_thread: bool, + pub(crate) root_accounting_state: Option>, } pub(crate) enum ActiveGoalStopReason { @@ -48,6 +49,7 @@ struct GoalRuntimeInner { metrics: GoalMetrics, thread_manager: Weak, accounting_state: Arc, + root_accounting_state: Option>, enabled: AtomicBool, tools_available_for_thread: bool, goal_state_lock: Semaphore, @@ -100,6 +102,7 @@ impl GoalRuntimeHandle { metrics, thread_manager, accounting_state, + root_accounting_state: config.root_accounting_state, enabled: AtomicBool::new(config.enabled), tools_available_for_thread: config.tools_available_for_thread, goal_state_lock: Semaphore::new(/*permits*/ 1), @@ -127,6 +130,10 @@ impl GoalRuntimeHandle { Arc::clone(&self.inner.accounting_state) } + pub(crate) fn root_accounting_state(&self) -> Option> { + self.inner.root_accounting_state.clone() + } + pub(crate) async fn goal_state_permit(&self) -> Result, String> { self.inner .goal_state_lock @@ -546,7 +553,7 @@ impl GoalRuntimeHandle { .account_thread_goal_usage( self.thread_id(), snapshot.time_delta_seconds, - /*token_delta*/ 0, + snapshot.token_delta, mode, Some(snapshot.expected_goal_id.as_str()), ) diff --git a/codex-rs/ext/goal/tests/accounting.rs b/codex-rs/ext/goal/tests/accounting.rs index d5b601c997..9bd1186f08 100644 --- a/codex-rs/ext/goal/tests/accounting.rs +++ b/codex-rs/ext/goal/tests/accounting.rs @@ -3,9 +3,11 @@ #[path = "../src/accounting.rs"] mod accounting; +use accounting::BudgetLimitedGoalDisposition; use accounting::GoalAccountingState; use codex_protocol::config_types::ModeKind; use codex_protocol::protocol::TokenUsage; +use codex_state::ThreadGoalStatus; use pretty_assertions::assert_eq; #[test] @@ -51,6 +53,46 @@ fn goal_accounting_ignores_plan_mode_turns() { assert_eq!(None, recorded); } +#[test] +fn goal_accounting_preserves_concurrent_descendant_usage_across_checkpoints() { + let state = GoalAccountingState::default(); + state.start_turn("turn-1", ModeKind::Default, &TokenUsage::default()); + state.mark_current_turn_goal_active("goal-1"); + let first_usage = token_usage( + /*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8, + /*reasoning_output_tokens*/ 0, /*total_tokens*/ 28, + ); + let second_usage = token_usage( + /*input_tokens*/ 8, /*cached_input_tokens*/ 2, /*output_tokens*/ 4, + /*reasoning_output_tokens*/ 0, /*total_tokens*/ 12, + ); + std::thread::scope(|scope| { + scope.spawn(|| state.record_descendant_token_usage(&first_usage)); + scope.spawn(|| state.record_descendant_token_usage(&second_usage)); + }); + + let first = state + .progress_snapshot("turn-1") + .expect("descendant usage should create a progress snapshot"); + assert_eq!(33, first.token_delta); + + state.record_descendant_token_usage(&token_usage( + /*input_tokens*/ 6, /*cached_input_tokens*/ 1, /*output_tokens*/ 3, + /*reasoning_output_tokens*/ 0, /*total_tokens*/ 9, + )); + state.mark_progress_accounted_for_status( + "turn-1", + &first, + ThreadGoalStatus::Active, + BudgetLimitedGoalDisposition::KeepActive, + ); + + let second = state + .progress_snapshot("turn-1") + .expect("usage received during accounting should remain pending"); + assert_eq!(8, second.token_delta); +} + fn token_usage( input_tokens: i64, cached_input_tokens: i64, diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index 72832764e8..e029f121fe 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -409,6 +409,186 @@ async fn parallel_tool_finish_accounts_active_goal_progress_once() -> anyhow::Re Ok(()) } +#[tokio::test] +async fn spawned_descendant_usage_exhausts_root_goal_budget_once() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let child = harness.spawn_child(ThreadId::new()).await?; + child.start_turn("child-turn", &TokenUsage::default()).await; + let grandchild = child + .spawn_child_with_source(ThreadId::new(), SubAgentSource::Review) + .await?; + grandchild + .start_turn("grandchild-turn", &TokenUsage::default()) + .await; + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ "objective": "account for the entire agent tree", "token_budget": 62 }), + )) + .await?; + + harness + .record_token_usage( + "turn-1", + &token_usage( + /*input_tokens*/ 12, /*cached_input_tokens*/ 2, /*output_tokens*/ 4, + /*reasoning_output_tokens*/ 0, /*total_tokens*/ 16, + ), + ) + .await; + let first_child_usage = input_token_usage(/*input_tokens*/ 23); + child + .record_token_usage_with_last("child-turn", &first_child_usage, &first_child_usage) + .await; + child + .record_token_usage_with_last( + "child-turn", + &input_token_usage(/*input_tokens*/ 36), + &input_token_usage(/*input_tokens*/ 13), + ) + .await; + let grandchild_usage = input_token_usage(/*input_tokens*/ 12); + grandchild + .record_token_usage_with_last("grandchild-turn", &grandchild_usage, &grandchild_usage) + .await; + harness + .notify_tool_finish("turn-1", "call-shell", "shell") + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(62, goal.tokens_used); + assert_eq!(codex_state::ThreadGoalStatus::BudgetLimited, goal.status); + Ok(()) +} + +#[tokio::test] +async fn grandchild_usage_rolls_up_after_parent_runtime_unloads() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + + let tools = harness.tools(); + tool_by_name(&tools, "create_goal") + .handle(tool_call( + "create_goal", + "call-create-goal", + json!({ "objective": "account for evicted agent trees" }), + )) + .await?; + + let child = harness.spawn_child(ThreadId::new()).await?; + child.stop_thread().await; + let grandchild = child.spawn_child(ThreadId::new()).await?; + grandchild + .start_turn("grandchild-turn", &TokenUsage::default()) + .await; + let usage = input_token_usage(/*input_tokens*/ 23); + grandchild + .record_token_usage_with_last("grandchild-turn", &usage, &usage) + .await; + harness + .notify_tool_finish("turn-1", "call-shell", "shell") + .await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; + assert_eq!(23, goal.tokens_used); + Ok(()) +} + +#[tokio::test] +async fn subagent_usage_resets_when_root_goal_is_replaced() -> anyhow::Result<()> { + let runtime = test_runtime().await?; + let thread_id = test_thread_id()?; + seed_thread_metadata(runtime.as_ref(), thread_id).await?; + let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; + harness.start_turn("turn-1", &TokenUsage::default()).await; + let child = harness.spawn_child(ThreadId::new()).await?; + child.start_turn("child-turn", &TokenUsage::default()).await; + + let previous_usage = input_token_usage(/*input_tokens*/ 10); + child + .record_token_usage_with_last("child-turn", &previous_usage, &previous_usage) + .await; + + let tools = harness.tools(); + let create_tool = tool_by_name(&tools, "create_goal"); + create_tool + .handle(tool_call( + "create_goal", + "call-create-first-goal", + json!({ "objective": "first goal" }), + )) + .await?; + child + .record_token_usage_with_last( + "child-turn", + &input_token_usage(/*input_tokens*/ 35), + &input_token_usage(/*input_tokens*/ 25), + ) + .await; + let completion = tool_call( + "update_goal", + "call-complete-first-goal", + json!({ "status": "complete" }), + ); + let completed = tool_by_name(&tools, "update_goal") + .handle(completion.clone()) + .await?; + assert_eq!( + json!(25), + completed.code_mode_result(&completion.payload)["goal"]["tokensUsed"] + ); + + child + .record_token_usage_with_last( + "child-turn", + &input_token_usage(/*input_tokens*/ 55), + &input_token_usage(/*input_tokens*/ 20), + ) + .await; + create_tool + .handle(tool_call( + "create_goal", + "call-create-second-goal", + json!({ "objective": "replacement goal" }), + )) + .await?; + child + .record_token_usage_with_last( + "child-turn", + &input_token_usage(/*input_tokens*/ 64), + &input_token_usage(/*input_tokens*/ 9), + ) + .await; + harness.stop_turn("turn-1").await; + + let goal = runtime + .thread_goals() + .get_thread_goal(thread_id) + .await? + .ok_or_else(|| anyhow::anyhow!("replacement goal should exist"))?; + assert_eq!("replacement goal", goal.objective); + assert_eq!(9, goal.tokens_used); + Ok(()) +} + #[tokio::test] async fn budget_limited_goal_keeps_accruing_until_turn_stop() -> anyhow::Result<()> { let runtime = test_runtime().await?; @@ -943,8 +1123,7 @@ async fn external_goal_mutation_start_accounts_active_goal_progress() -> anyhow: } #[tokio::test] -async fn goal_service_external_set_active_resets_baseline_without_live_thread() -> anyhow::Result<()> -{ +async fn goal_service_external_set_active_preserves_concurrent_usage() -> anyhow::Result<()> { let runtime = test_runtime().await?; let thread_id = test_thread_id()?; seed_thread_metadata(runtime.as_ref(), thread_id).await?; @@ -969,6 +1148,8 @@ async fn goal_service_external_set_active_resets_baseline_without_live_thread() json!({ "objective": "old objective" }), )) .await?; + let child = harness.spawn_child(ThreadId::new()).await?; + child.start_turn("child-turn", &TokenUsage::default()).await; harness.sink.clear(); harness @@ -994,6 +1175,20 @@ async fn goal_service_external_set_active_resets_baseline_without_live_thread() }, ) .await?; + harness + .record_token_usage( + "turn-1", + &token_usage( + /*input_tokens*/ 125, /*cached_input_tokens*/ 0, + /*output_tokens*/ 0, /*reasoning_output_tokens*/ 0, + /*total_tokens*/ 125, + ), + ) + .await; + let child_usage = input_token_usage(/*input_tokens*/ 23); + child + .record_token_usage_with_last("child-turn", &child_usage, &child_usage) + .await; outcome.apply_runtime_effects(&harness.goal_service).await; harness @@ -1015,7 +1210,7 @@ async fn goal_service_external_set_active_resets_baseline_without_live_thread() .get_thread_goal(thread_id) .await? .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; - assert_eq!(30, goal.tokens_used); + assert_eq!(53, goal.tokens_used); Ok(()) } @@ -1076,6 +1271,12 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul let harness = GoalExtensionHarness::new(runtime.clone(), thread_id).await?; harness.resume_thread().await; + let child = harness.spawn_child(ThreadId::new()).await?; + child.start_turn("child-turn", &TokenUsage::default()).await; + let usage = input_token_usage(/*input_tokens*/ 23); + child + .record_token_usage_with_last("child-turn", &usage, &usage) + .await; tokio::time::sleep(Duration::from_millis(1_100)).await; harness .runtime_handle() @@ -1089,6 +1290,7 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul .await? .ok_or_else(|| anyhow::anyhow!("goal should exist"))?; assert_eq!(ThreadGoalStatus::Active, protocol_status(goal.status)); + assert_eq!(23, goal.tokens_used); assert!( goal.time_used_seconds >= 1, "resumed idle accounting should add elapsed wall-clock time" @@ -1098,7 +1300,7 @@ async fn thread_resume_rehydrates_active_goal_idle_accounting() -> anyhow::Resul event_id: format!("{thread_id}:external-goal-mutation"), turn_id: None, status: ThreadGoalStatus::Active, - tokens_used: 0, + tokens_used: 23, }], harness.sink.goal_events() ); @@ -1289,7 +1491,7 @@ fn tool_names(tools: &[Arc ToolExecutor>>]) -> Ve } struct GoalExtensionHarness { - registry: codex_extension_api::ExtensionRegistry<()>, + registry: Arc>, session_store: ExtensionData, thread_store: ExtensionData, goal_service: Arc, @@ -1316,8 +1518,8 @@ impl GoalExtensionHarness { max_goal_token_budget: None, }, ); - let registry = builder.build(); - let session_store = ExtensionData::new("session-1"); + let registry = Arc::new(builder.build()); + let session_store = ExtensionData::new(thread_id.to_string()); let thread_store = ExtensionData::new(thread_id.to_string()); let session_source = SessionSource::Cli; for contributor in registry.thread_lifecycle_contributors() { @@ -1343,6 +1545,49 @@ impl GoalExtensionHarness { }) } + async fn spawn_child(&self, thread_id: ThreadId) -> anyhow::Result { + let session_source = SubAgentSource::ThreadSpawn { + parent_thread_id: ThreadId::from_string(self.thread_store.level_id())?, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + }; + self.spawn_child_with_source(thread_id, session_source) + .await + } + + async fn spawn_child_with_source( + &self, + thread_id: ThreadId, + session_source: SubAgentSource, + ) -> anyhow::Result { + let session_store = ExtensionData::new(self.session_store.level_id()); + let thread_store = ExtensionData::new(thread_id.to_string()); + let session_source = SessionSource::SubAgent(session_source); + for contributor in self.registry.thread_lifecycle_contributors() { + contributor + .on_thread_start(ThreadStartInput { + config: &(), + session_source: &session_source, + persistent_thread_state_available: true, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store: &thread_store, + }) + .await; + } + Ok(Self { + registry: Arc::clone(&self.registry), + session_store, + thread_store, + goal_service: Arc::clone(&self.goal_service), + sink: Arc::clone(&self.sink), + }) + } + fn tools(&self) -> Vec ToolExecutor>>> { self.registry .tool_contributors() @@ -1388,10 +1633,20 @@ impl GoalExtensionHarness { } async fn record_token_usage(&self, turn_id: &str, usage: &TokenUsage) { + self.record_token_usage_with_last(turn_id, usage, &TokenUsage::default()) + .await; + } + + async fn record_token_usage_with_last( + &self, + turn_id: &str, + usage: &TokenUsage, + last_usage: &TokenUsage, + ) { let turn_store = ExtensionData::new(turn_id); let token_usage = TokenUsageInfo { total_token_usage: usage.clone(), - last_token_usage: TokenUsage::default(), + last_token_usage: last_usage.clone(), model_context_window: None, }; for contributor in self.registry.token_usage_contributors() { @@ -1603,6 +1858,14 @@ fn token_usage( } } +fn input_token_usage(input_tokens: i64) -> TokenUsage { + TokenUsage { + input_tokens, + total_tokens: input_tokens, + ..TokenUsage::default() + } +} + fn protocol_status(status: codex_state::ThreadGoalStatus) -> ThreadGoalStatus { match status { codex_state::ThreadGoalStatus::Active => ThreadGoalStatus::Active,