openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-08 15:50:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Account subagent token usage toward root goals (#41183)

Browse Source 
## 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
This commit is contained in:
Eric Traut
2026-08-27 15:39:26 +00:00
committed by copyberry
parent 5f49aba876
commit 4761851ff3
6 changed files with 398 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
codex-rs/ext/goal/src/accounting.rs
View File

@@ -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,
}
}
}

2
codex-rs/ext/goal/src/api.rs
View File

@@ -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);

35
codex-rs/ext/goal/src/extension.rs
View File

@@ -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);
})
}
}

9
codex-rs/ext/goal/src/runtime.rs
View File

@@ -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()),
)

42
codex-rs/ext/goal/tests/accounting.rs
View File

@@ -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,

279
codex-rs/ext/goal/tests/goal_extension_backend.rs
View File

@@ -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<dyn for<'call> ToolExecutor<ToolCall<'call>>>]) -> Ve
}
struct GoalExtensionHarness {
registry: codex_extension_api::ExtensionRegistry<()>,
registry: Arc<codex_extension_api::ExtensionRegistry<()>>,
session_store: ExtensionData,
thread_store: ExtensionData,
goal_service: Arc<GoalService>,
@@ -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<Self> {
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<Self> {
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<Arc<dyn for<'call> ToolExecutor<ToolCall<'call>>>> {
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,