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

@@ -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<GoalAccountingInner>,
     progress_accounting_lock: Semaphore,
+    descendant_token_usage: AtomicI64,
 }
 
 #[derive(Debug)]
@@ -21,6 +24,7 @@ struct GoalAccountingInner {
     turns: HashMap<String, GoalTurnAccounting>,
     wall_clock: GoalWallClockAccounting,
     budget_limit_reported_goal_id: Option<String>,
+    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<String>) {
         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,
         }
     }
 }

codex-rs/ext/goal/src/api.rs

@@ -345,7 +345,7 @@ impl GoalService {
         }
     }
 
-    fn runtime_for_thread(&self, thread_id: ThreadId) -> Option<Arc<GoalRuntimeHandle>> {
+    pub(crate) fn runtime_for_thread(&self, thread_id: ThreadId) -> Option<Arc<GoalRuntimeHandle>> {
         let key = thread_id.to_string();
         let mut runtimes = self.runtimes();
         let runtime = runtimes.get(&key).and_then(Weak::upgrade);

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>(|| {
                 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);
         })
     }
 }

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<Arc<GoalAccountingState>>,
 }
 
 pub(crate) enum ActiveGoalStopReason {
@@ -48,6 +49,7 @@ struct GoalRuntimeInner {
     metrics: GoalMetrics,
     thread_manager: Weak<ThreadManager>,
     accounting_state: Arc<GoalAccountingState>,
+    root_accounting_state: Option<Arc<GoalAccountingState>>,
     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<Arc<GoalAccountingState>> {
+        self.inner.root_accounting_state.clone()
+    }
+
     pub(crate) async fn goal_state_permit(&self) -> Result<SemaphorePermit<'_>, 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()),
             )