diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 4913a5dc76..3668daa519 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server/src/bespoke_event_handling.rs b/codex-rs/app-server/src/bespoke_event_handling.rs index 8e5c1b9ade..a4130d93fd 100644 --- a/codex-rs/app-server/src/bespoke_event_handling.rs +++ b/codex-rs/app-server/src/bespoke_event_handling.rs @@ -3721,6 +3721,7 @@ mod tests { total_tokens: 200, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 10, cached_input_tokens: 5, diff --git a/codex-rs/app-server/tests/common/rollout.rs b/codex-rs/app-server/tests/common/rollout.rs index a7c3222929..44ebcfa967 100644 --- a/codex-rs/app-server/tests/common/rollout.rs +++ b/codex-rs/app-server/tests/common/rollout.rs @@ -125,6 +125,7 @@ pub fn create_fake_rollout_with_token_usage( total_tokens: 150, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 70, cached_input_tokens: 10, diff --git a/codex-rs/app-server/tests/suite/v2/thread_resume.rs b/codex-rs/app-server/tests/suite/v2/thread_resume.rs index a1f2dd3779..07f05dd7e3 100644 --- a/codex-rs/app-server/tests/suite/v2/thread_resume.rs +++ b/codex-rs/app-server/tests/suite/v2/thread_resume.rs @@ -4250,6 +4250,7 @@ async fn cold_paginated_resume_restores_usage_without_loading_turns() -> Result< total_tokens: 150, ..Default::default() }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 70, output_tokens: 20, @@ -4361,6 +4362,7 @@ async fn cold_paginated_resume_omits_usage_when_its_turn_is_ambiguous() -> Resul total_tokens: 150, ..Default::default() }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { total_tokens: 90, ..Default::default() @@ -4645,6 +4647,7 @@ async fn thread_resume_token_usage_replay_can_belong_to_interrupted_turn() -> Re total_tokens: 230, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 90, cached_input_tokens: 30, diff --git a/codex-rs/core/src/compact.rs b/codex-rs/core/src/compact.rs index 3c6642301f..2f6a91de70 100644 --- a/codex-rs/core/src/compact.rs +++ b/codex-rs/core/src/compact.rs @@ -830,8 +830,7 @@ async fn drain_to_completed( prompt: &Prompt, phase: CompactionPhase, ) -> CodexResult { - sess.set_server_reasoning_included(turn_context.model_info().use_responses_lite) - .await; + let mut includes_reasoning = turn_context.model_info().use_responses_lite; let mut stream = client_session .stream( prompt, @@ -862,7 +861,7 @@ async fn drain_to_completed( // Keep each attempt isolated until the response completes successfully. Ok(ResponseEvent::OutputItemDone(item)) => output.push(item), Ok(ResponseEvent::ServerReasoningIncluded(included)) => { - sess.set_server_reasoning_included(included).await; + includes_reasoning = included; } Ok(ResponseEvent::RateLimits(snapshot)) => { sess.update_rate_limits(turn_context, snapshot).await; @@ -896,8 +895,12 @@ async fn drain_to_completed( usage_metadata.as_ref(), ) .await; - sess.update_token_usage_info(turn_context, token_usage.as_ref()) - .await?; + sess.update_token_usage_info( + turn_context, + token_usage.as_ref(), + includes_reasoning, + ) + .await?; return Ok(CompactionResponse { response_id, output, diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 175f26e6ea..12fff896a8 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -246,9 +246,7 @@ impl ContextManager { reset_version: 0, user_message_revision: 0, guardian_review_context_revision: next_guardian_review_context_revision(), - token_info: TokenUsageInfo::new_or_append( - &None, &None, /*model_context_window*/ None, - ), + token_info: None, reference_context_item: None, world_state_baseline: None, } @@ -903,12 +901,14 @@ impl ContextManager { &mut self, usage: &TokenUsage, model_context_window: Option, + includes_reasoning: bool, ) { - self.token_info = TokenUsageInfo::new_or_append( + self.token_info = Some(TokenUsageInfo::new_or_append( &self.token_info, - &Some(usage.clone()), + usage, + includes_reasoning, model_context_window, - ); + )); } fn get_non_last_reasoning_items_tokens(&self) -> i64 { @@ -950,9 +950,7 @@ impl ContextManager { self.items[start..].iter().map(|envelope| &envelope.item) } - /// When true, the server already accounted for past reasoning tokens and - /// the client should not re-estimate them. - pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { + pub(crate) fn get_total_token_usage(&self) -> i64 { let last_tokens = self .token_info .as_ref() @@ -962,7 +960,11 @@ impl ContextManager { .items_after_last_model_generated_item() .map(estimate_item_token_count) .fold(0i64, i64::saturating_add); - if server_reasoning_included { + if self + .token_info + .as_ref() + .is_some_and(|info| info.last_usage_includes_reasoning == Some(true)) + { last_tokens.saturating_add(items_after_last_model_generated_tokens) } else { last_tokens diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index b5eda994c8..7e80afe42b 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -937,6 +937,33 @@ fn non_last_reasoning_tokens_ignore_entries_after_last_user() { assert_eq!(history.get_non_last_reasoning_items_tokens(), 32); } +#[test_case(None, 125; "legacy count")] +#[test_case(Some(false), 125; "exclusive count")] +#[test_case(Some(true), 100; "inclusive count")] +fn total_token_usage_respects_saved_reasoning_accounting( + includes_reasoning: Option, + expected: i64, +) { + let mut history = create_history_with_items(vec![ + reasoning_with_encrypted_content(/*len*/ 1_000), + user_msg("second"), + assistant_msg("counted reply"), + ]); + let usage = TokenUsage { + total_tokens: 100, + ..Default::default() + }; + let mut saved = json!({ "last_token_usage": usage, "total_token_usage": usage }); + if let Some(included) = includes_reasoning { + saved["last_usage_includes_reasoning"] = json!(included); + } + history.set_token_info(Some( + serde_json::from_value(saved).expect("saved token count"), + )); + + assert_eq!(history.get_total_token_usage(), expected); +} + #[test] fn items_after_last_model_generated_tokens_include_user_and_tool_output() { let history = create_history_with_items(vec![ @@ -1160,6 +1187,7 @@ fn total_token_usage_includes_all_items_after_last_model_generated_item() { ..Default::default() }, /*model_context_window*/ None, + /*includes_reasoning*/ true, ); let added_user = user_msg("new user message"); let added_tool_output = custom_tool_call_output("tool-tail", "new tool output"); @@ -1169,7 +1197,7 @@ fn total_token_usage_includes_all_items_after_last_model_generated_item() { ); assert_eq!( - history.get_total_token_usage(/*server_reasoning_included*/ true), + history.get_total_token_usage(), 100 + estimate_item_token_count(&added_user) + estimate_item_token_count(&added_tool_output) ); diff --git a/codex-rs/core/src/session/async_compaction_tests.rs b/codex-rs/core/src/session/async_compaction_tests.rs index e4a048124b..d73b96c657 100644 --- a/codex-rs/core/src/session/async_compaction_tests.rs +++ b/codex-rs/core/src/session/async_compaction_tests.rs @@ -65,6 +65,7 @@ async fn cancelled_turn_skips_launch_but_preserves_existing_work(trigger: Compac total_tokens: 70_000, ..Default::default() }), + /*includes_reasoning*/ false, ) .await .expect("record usage"); diff --git a/codex-rs/core/src/session/guardian_checkpoint_tests.rs b/codex-rs/core/src/session/guardian_checkpoint_tests.rs index 6ff5ab3589..c64bf40429 100644 --- a/codex-rs/core/src/session/guardian_checkpoint_tests.rs +++ b/codex-rs/core/src/session/guardian_checkpoint_tests.rs @@ -78,6 +78,7 @@ async fn guardian_checkpoint_preserves_live_context_without_storage( ..Default::default() }, Some(32_000), + /*includes_reasoning*/ true, ); } let expected = session.clone_history().await; diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b8f3a0f078..cc418ca5ed 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -417,7 +417,7 @@ use codex_utils_stream_parser::ProposedPlanSegment; /// Whether reported usage updates the installed context baseline or only spending. #[derive(Clone, Copy)] pub(crate) enum TokenUsageContext { - Sampling, + Sampling { includes_reasoning: bool }, SpendingOnly, } @@ -1569,7 +1569,7 @@ impl Session { pub(crate) async fn get_total_token_usage(&self) -> i64 { let state = self.state.lock().await; - state.get_total_token_usage(state.server_reasoning_included()) + state.get_total_token_usage() } pub(crate) async fn auto_compact_window_snapshot(&self) -> AutoCompactWindowSnapshot { @@ -5244,13 +5244,14 @@ impl Session { &self, turn_context: &TurnContext, token_usage: Option<&TokenUsage>, + includes_reasoning: bool, ) -> CodexResult<()> { let result = self .record_token_usage_info( turn_context, &turn_context.initial_settings, token_usage, - TokenUsageContext::Sampling, + TokenUsageContext::Sampling { includes_reasoning }, ) .await; self.send_token_count_event(turn_context).await; @@ -5312,10 +5313,11 @@ impl Session { let token_info = { let mut state = self.state.lock().await; match context { - TokenUsageContext::Sampling => { + TokenUsageContext::Sampling { includes_reasoning } => { state.update_token_info_from_usage( token_usage, turn_context.model_context_window(), + includes_reasoning, ); if matches!( turn_context.config.model_auto_compact_token_limit_scope, @@ -5328,6 +5330,7 @@ impl Session { TokenUsageContext::SpendingOnly => { let mut info = state.token_info().unwrap_or(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: turn_context.model_context_window(), }); @@ -5336,6 +5339,7 @@ impl Session { // Contributors receive the reported usage; stored context stays unchanged. info.last_token_usage = token_usage.clone(); + info.last_usage_includes_reasoning = None; info.model_context_window = turn_context .model_context_window() .or(info.model_context_window); @@ -5384,6 +5388,7 @@ impl Session { let mut state = self.state.lock().await; let mut info = state.token_info().unwrap_or(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: None, }); @@ -5402,9 +5407,9 @@ impl Session { info.model_context_window = Some(model_context_window); } - state.set_token_info(Some(info)); // The local estimate already includes every retained reasoning item. - state.set_server_reasoning_included(/*included*/ true); + info.last_usage_includes_reasoning = Some(true); + state.set_token_info(Some(info)); } self.set_auto_compact_window_estimated_prefill_for_scope( turn_context, @@ -5443,11 +5448,6 @@ impl Session { state.record_mcp_dependency_prompted(names); } - pub(crate) async fn set_server_reasoning_included(&self, included: bool) { - let mut state = self.state.lock().await; - state.set_server_reasoning_included(included); - } - pub(crate) async fn send_token_count_event(&self, turn_context: &TurnContext) { let (info, rate_limits) = { let state = self.state.lock().await; diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 579d109872..68022ae9d2 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -3420,6 +3420,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { total_tokens: 30, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 3, cached_input_tokens: 0, @@ -3441,6 +3442,7 @@ async fn record_initial_history_seeds_token_info_from_rollout() { total_tokens: 375, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage { input_tokens: 10, cached_input_tokens: 0, @@ -3590,6 +3592,7 @@ async fn recompute_token_usage_updates_model_context_window() { let mut state = session.state.lock().await; state.set_token_info(Some(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(258_400), })); @@ -3606,7 +3609,7 @@ async fn recompute_token_usage_updates_model_context_window() { assert_eq!(actual.model_context_window, Some(128_000)); } -#[test_case(crate::session::TokenUsageContext::Sampling; "sampling")] +#[test_case(crate::session::TokenUsageContext::Sampling { includes_reasoning: false }; "sampling")] #[test_case(crate::session::TokenUsageContext::SpendingOnly; "spending_only")] #[tokio::test] async fn record_token_usage_info_notifies_extension_contributors( @@ -3693,7 +3696,9 @@ async fn record_token_usage_info_notifies_extension_contributors( &turn_context, &turn_context.initial_settings, Some(&first_usage), - crate::session::TokenUsageContext::Sampling, + crate::session::TokenUsageContext::Sampling { + includes_reasoning: true, + }, ) .await .expect("first usage should be recorded"); @@ -3711,8 +3716,12 @@ async fn record_token_usage_info_notifies_extension_contributors( expected_total_usage.add_assign(&second_usage); let expected_session_info = TokenUsageInfo { total_token_usage: expected_total_usage.clone(), + last_usage_includes_reasoning: Some(matches!( + context, + crate::session::TokenUsageContext::SpendingOnly + )), last_token_usage: match context { - crate::session::TokenUsageContext::Sampling => second_usage.clone(), + crate::session::TokenUsageContext::Sampling { .. } => second_usage.clone(), crate::session::TokenUsageContext::SpendingOnly => first_usage.clone(), }, model_context_window: turn_context.model_context_window(), @@ -3724,6 +3733,7 @@ async fn record_token_usage_info_notifies_extension_contributors( turn_level_id: turn_context.sub_id.clone(), token_usage: TokenUsageInfo { total_token_usage: first_usage.clone(), + last_usage_includes_reasoning: Some(true), last_token_usage: first_usage, model_context_window: turn_context.model_context_window(), }, @@ -3736,6 +3746,12 @@ async fn record_token_usage_info_notifies_extension_contributors( turn_level_id: turn_context.sub_id.clone(), token_usage: TokenUsageInfo { total_token_usage: expected_total_usage, + last_usage_includes_reasoning: match context { + crate::session::TokenUsageContext::Sampling { includes_reasoning } => { + Some(includes_reasoning) + } + crate::session::TokenUsageContext::SpendingOnly => None, + }, last_token_usage: second_usage, model_context_window: turn_context.model_context_window(), }, @@ -12972,6 +12988,7 @@ async fn set_total_token_usage(sess: &Session, total_token_usage: TokenUsage) { let mut state = sess.state.lock().await; state.set_token_info(Some(TokenUsageInfo { total_token_usage, + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: None, })); diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index b56e7ca2aa..e4a092ed65 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -2760,8 +2760,7 @@ async fn try_run_sampling_request( let sampling_timing_guard = turn_context.turn_timing_state.begin_sampling(); let mut response_completed = false; let mut request_token_usage = None; - sess.set_server_reasoning_included(step_context.settings.model_info.use_responses_lite) - .await; + let mut includes_reasoning = step_context.settings.model_info.use_responses_lite; let stream = client_session .stream( &prompt, @@ -3215,7 +3214,7 @@ async fn try_run_sampling_request( .await; } ResponseEvent::ServerReasoningIncluded(included) => { - sess.set_server_reasoning_included(included).await; + includes_reasoning = included; } ResponseEvent::RateLimits(snapshot) => { // Update internal state with latest rate limits, but defer sending until @@ -3267,7 +3266,7 @@ async fn try_run_sampling_request( &turn_context, &step_context.settings, token_usage.as_ref(), - super::TokenUsageContext::Sampling, + super::TokenUsageContext::Sampling { includes_reasoning }, ) .await; should_emit_token_count = true; diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index ebfc5817cb..968f3425fb 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -81,7 +81,6 @@ pub(crate) struct SessionState { pub(crate) history_reset: CancellationToken, pub(crate) latest_rate_limits: Option, pub(crate) latest_token_usage_record: Option, - pub(crate) server_reasoning_included: bool, pub(crate) mcp_dependency_prompted: HashSet, pub(crate) additional_context: AdditionalContextStore, /// Settings used by the latest regular user turn, used for turn-to-turn @@ -136,7 +135,6 @@ impl SessionState { history_reset: CancellationToken::new(), latest_rate_limits: None, latest_token_usage_record: None, - server_reasoning_included: false, mcp_dependency_prompted: HashSet::new(), additional_context: AdditionalContextStore::default(), previous_turn_settings: None, @@ -296,8 +294,10 @@ impl SessionState { &mut self, usage: &TokenUsage, model_context_window: Option, + includes_reasoning: bool, ) { - self.history.update_token_info(usage, model_context_window); + self.history + .update_token_info(usage, model_context_window, includes_reasoning); } pub(crate) fn ensure_auto_compact_window_server_prefill_from_usage( @@ -379,17 +379,8 @@ impl SessionState { self.history.set_token_usage_full(context_window); } - pub(crate) fn get_total_token_usage(&self, server_reasoning_included: bool) -> i64 { - self.history - .get_total_token_usage(server_reasoning_included) - } - - pub(crate) fn set_server_reasoning_included(&mut self, included: bool) { - self.server_reasoning_included = included; - } - - pub(crate) fn server_reasoning_included(&self) -> bool { - self.server_reasoning_included + pub(crate) fn get_total_token_usage(&self) -> i64 { + self.history.get_total_token_usage() } pub(crate) fn record_mcp_dependency_prompted(&mut self, names: I) diff --git a/codex-rs/core/src/tasks/mod_tests.rs b/codex-rs/core/src/tasks/mod_tests.rs index ba03d83a15..eb4bc4b5e2 100644 --- a/codex-rs/core/src/tasks/mod_tests.rs +++ b/codex-rs/core/src/tasks/mod_tests.rs @@ -171,6 +171,7 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) { total_tokens: 1_000, ..TokenUsage::default() }), + /*includes_reasoning*/ false, ) .await .expect("earlier session usage should be recorded"); @@ -193,7 +194,11 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) { }; // Local pre-turn compaction records usage with the previous model's context. session - .update_token_usage_info(&previous_context, Some(&compaction_usage)) + .update_token_usage_info( + &previous_context, + Some(&compaction_usage), + /*includes_reasoning*/ false, + ) .await .expect("compaction usage should be recorded"); expected_usage.push((previous_context.model_info().slug.clone(), compaction_usage)); @@ -225,7 +230,9 @@ async fn turn_completion_metrics_follow_model_switch(scenario: UsageScenario) { &turn_context, &response_context.initial_settings, Some(&response_usage), - crate::session::TokenUsageContext::Sampling, + crate::session::TokenUsageContext::Sampling { + includes_reasoning: false, + }, ) .await .expect("response usage should be recorded"); diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 2e192498cb..d20e1dfb26 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -3576,6 +3576,7 @@ async fn token_count_includes_rate_limits_snapshot() { "reasoning_output_tokens": 0, "total_tokens": 123 }, + "last_usage_includes_reasoning": false, // Default model is gpt-5.4 in tests → 95% usable context window "model_context_window": 258400 }, diff --git a/codex-rs/core/tests/suite/compact.rs b/codex-rs/core/tests/suite/compact.rs index b3522a09b2..c480acee0c 100644 --- a/codex-rs/core/tests/suite/compact.rs +++ b/codex-rs/core/tests/suite/compact.rs @@ -5480,8 +5480,18 @@ async fn auto_compact_body_after_prefix_still_caps_at_context_window() { ); } +#[derive(Clone, Copy)] +enum ContinueThread { + Live, + Fork, + Restart, +} + +#[test_case::test_case(ContinueThread::Live; "live")] +#[test_case::test_case(ContinueThread::Fork; "fork")] +#[test_case::test_case(ContinueThread::Restart; "cold resume")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn auto_compact_accounts_for_encrypted_reasoning() { +async fn auto_compact_accounts_for_encrypted_reasoning(continuation: ContinueThread) { skip_if_no_network!(); let server = start_mock_server().await; @@ -5519,26 +5529,53 @@ async fn auto_compact_accounts_for_encrypted_reasoning() { ) .await; - let test = test_codex() - .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) - .with_model_info_override("gpt-5.5", |model| { - model.use_responses_lite = true; - }) - .with_config(|config| { - config.model = Some("gpt-5.5".to_string()); - set_test_compact_prompt(config); - config.model_auto_compact_token_limit = Some(81); - }) + let builder = || { + test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_model_info_override("gpt-5.5", |model| { + model.use_responses_lite = true; + }) + .with_config(|config| { + config.model = Some("gpt-5.5".to_string()); + set_test_compact_prompt(config); + config.model_auto_compact_token_limit = Some(81); + }) + }; + let mut test = builder() .build_with_auto_env(&server) .await .expect("build codex"); - for user in [ - "FIRST_REASONING", - "SECOND_REASONING", - "BELOW_COMPACT_LIMIT", - "AT_COMPACT_LIMIT", - ] { + for user in ["FIRST_REASONING", "SECOND_REASONING"] { + test.submit_text_turn(user).await.expect("submit user turn"); + } + + match continuation { + ContinueThread::Live => {} + ContinueThread::Restart => { + test = builder() + .restart_with_auto_env(&server, &test) + .await + .expect("restart"); + } + ContinueThread::Fork => { + let path = test + .session_configured + .rollout_path + .clone() + .expect("rollout path"); + let options = test.start_thread_options().await; + test.codex.shutdown_and_wait().await.expect("flush rollout"); + test.codex = test + .thread_manager + .fork_legacy_thread(usize::MAX, options, path) + .await + .expect("fork") + .thread; + } + } + + for user in ["BELOW_COMPACT_LIMIT", "AT_COMPACT_LIMIT"] { test.submit_text_turn(user).await.expect("submit user turn"); } diff --git a/codex-rs/ext/goal/tests/goal_extension_backend.rs b/codex-rs/ext/goal/tests/goal_extension_backend.rs index cc88125491..8faa045d7e 100644 --- a/codex-rs/ext/goal/tests/goal_extension_backend.rs +++ b/codex-rs/ext/goal/tests/goal_extension_backend.rs @@ -1852,6 +1852,7 @@ impl GoalExtensionHarness { let turn_store = ExtensionData::new(turn_id); let token_usage = TokenUsageInfo { total_token_usage: usage.clone(), + last_usage_includes_reasoning: None, last_token_usage: last_usage.clone(), model_context_window: None, }; diff --git a/codex-rs/external-agent-migration/src/sessions/export.rs b/codex-rs/external-agent-migration/src/sessions/export.rs index d9067e2e4a..d5e7fc2c3a 100644 --- a/codex-rs/external-agent-migration/src/sessions/export.rs +++ b/codex-rs/external-agent-migration/src/sessions/export.rs @@ -200,6 +200,7 @@ fn token_count_item(last_model_visible_tokens: i64) -> RolloutItem { RolloutItem::EventMsg(EventMsg::TokenCount(TokenCountEvent { info: Some(TokenUsageInfo { total_token_usage: usage.clone(), + last_usage_includes_reasoning: None, last_token_usage: usage, model_context_window: None, }), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 21d45f6b78..8c419c536c 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -2506,6 +2506,10 @@ pub struct TokenUsageRecord { pub struct TokenUsageInfo { pub total_token_usage: TokenUsage, pub last_token_usage: TokenUsage, + /// Whether the last context count includes reasoning retained from earlier turns. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[ts(optional)] + pub last_usage_includes_reasoning: Option, // TODO(aibrahim): make this not optional #[ts(type = "number | null")] pub model_context_window: Option, @@ -2514,33 +2518,26 @@ pub struct TokenUsageInfo { impl TokenUsageInfo { pub fn new_or_append( info: &Option, - last: &Option, + last: &TokenUsage, + includes_reasoning: bool, model_context_window: Option, - ) -> Option { - if info.is_none() && last.is_none() { - return None; - } - + ) -> Self { let mut info = match info { Some(info) => info.clone(), None => Self { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window, }, }; - if let Some(last) = last { - info.append_last_usage(last); - } + info.total_token_usage.add_assign(last); + info.last_token_usage = last.clone(); + info.last_usage_includes_reasoning = Some(includes_reasoning); if let Some(model_context_window) = model_context_window { info.model_context_window = Some(model_context_window); } - Some(info) - } - - pub fn append_last_usage(&mut self, last: &TokenUsage) { - self.total_token_usage.add_assign(last); - self.last_token_usage = last.clone(); + info } pub fn fill_to_context_window(&mut self, context_window: i64) { @@ -2552,6 +2549,7 @@ impl TokenUsageInfo { total_tokens: context_window, ..TokenUsage::default() }; + self.last_usage_includes_reasoning = None; self.last_token_usage = TokenUsage { total_tokens: delta, ..TokenUsage::default() @@ -2561,6 +2559,7 @@ impl TokenUsageInfo { pub fn full_context_window(context_window: i64) -> Self { let mut info = Self { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(context_window), }; @@ -6627,10 +6626,11 @@ mod tests { fn token_usage_info_new_or_append_updates_context_window_when_provided() { let initial = Some(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(258_400), }); - let last = Some(TokenUsage { + let last = TokenUsage { input_tokens: 10, cached_input_tokens: 0, cache_write_input_tokens: 0, @@ -6638,10 +6638,14 @@ mod tests { reasoning_output_tokens: 0, total_tokens: 10, codex_rollout_budget_units: None, - }); + }; - let info = TokenUsageInfo::new_or_append(&initial, &last, Some(128_000)) - .expect("new_or_append should return info"); + let info = TokenUsageInfo::new_or_append( + &initial, + &last, + /*includes_reasoning*/ true, + Some(128_000), + ); assert_eq!(info.model_context_window, Some(128_000)); } @@ -6650,10 +6654,11 @@ mod tests { fn token_usage_info_new_or_append_preserves_context_window_when_not_provided() { let initial = Some(TokenUsageInfo { total_token_usage: TokenUsage::default(), + last_usage_includes_reasoning: None, last_token_usage: TokenUsage::default(), model_context_window: Some(258_400), }); - let last = Some(TokenUsage { + let last = TokenUsage { input_tokens: 10, cached_input_tokens: 0, cache_write_input_tokens: 0, @@ -6661,11 +6666,11 @@ mod tests { reasoning_output_tokens: 0, total_tokens: 10, codex_rollout_budget_units: None, - }); + }; - let info = - TokenUsageInfo::new_or_append(&initial, &last, /*model_context_window*/ None) - .expect("new_or_append should return info"); + let info = TokenUsageInfo::new_or_append( + &initial, &last, /*includes_reasoning*/ true, /*model_context_window*/ None, + ); assert_eq!(info.model_context_window, Some(258_400)); } diff --git a/codex-rs/state/src/runtime/threads.rs b/codex-rs/state/src/runtime/threads.rs index 1fc8b1e927..222fc83b3b 100644 --- a/codex-rs/state/src/runtime/threads.rs +++ b/codex-rs/state/src/runtime/threads.rs @@ -3756,6 +3756,7 @@ mod tests { total_tokens: 321, codex_rollout_budget_units: None, }, + last_usage_includes_reasoning: None, last_token_usage: codex_protocol::protocol::TokenUsage::default(), model_context_window: None, }),