openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-29 19:08:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve reasoning token accounting across resume and fork (#57839)

Browse Source 
## Why

Resuming or forking a thread loses the session flag indicating that reported token usage already includes retained reasoning. Counting that reasoning again can trigger automatic compaction too early.

## What changed

Persist `last_usage_includes_reasoning` with `TokenUsageInfo` and use it when estimating context usage. Record the flag with completed sampling and compaction usage, and mark local estimates as including reasoning. Older records without the field retain the existing accounting behavior.

## Testing

Add coverage for saved inclusive, exclusive, and legacy token counts. Extend the encrypted-reasoning compaction test to verify the compaction threshold for live, forked, and cold-resumed threads.

GitOrigin-RevId: 9f22262a693111ec0406dce11804eb4844ad8983
This commit is contained in:
andrewgu-oai
2026-10-28 21:22:06 +00:00
committed by copyberry
parent b91f2515f5
commit c62e9df39f
20 changed files with 191 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/compact.rs
+8 -5
View File
@@ -830,8 +830,7 @@ async fn drain_to_completed(
prompt: &Prompt,
phase: CompactionPhase,
) -> CodexResult<CompactionResponse> {
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,
codex-rs/core/src/context_manager/history.rs
+12 -10
View File
@@ -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<i64>,
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
codex-rs/core/src/context_manager/history_tests.rs
+29 -1
View File
@@ -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<bool>,
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)
);
codex-rs/core/src/session/async_compaction_tests.rs
+1
View File
@@ -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");
codex-rs/core/src/session/guardian_checkpoint_tests.rs
+1
View File
@@ -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;
codex-rs/core/src/session/mod.rs
+11 -11
View File
@@ -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;
codex-rs/core/src/session/tests.rs
+20 -3
View File
@@ -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,
}));
codex-rs/core/src/session/turn.rs
+3 -4
View File
@@ -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;
codex-rs/core/src/state/session.rs
+5 -14
View File
@@ -81,7 +81,6 @@ pub(crate) struct SessionState {
pub(crate) history_reset: CancellationToken,
pub(crate) latest_rate_limits: Option<RateLimitSnapshot>,
pub(crate) latest_token_usage_record: Option<TokenUsageRecord>,
pub(crate) server_reasoning_included: bool,
pub(crate) mcp_dependency_prompted: HashSet<String>,
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<i64>,
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<I>(&mut self, names: I)
codex-rs/core/src/tasks/mod_tests.rs
+9 -2
View File
@@ -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");
codex-rs/core/tests/suite/client.rs
+1
View File
@@ -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
},
codex-rs/core/tests/suite/compact.rs
+54 -17
View File
@@ -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");
}