mirror of
https://github.com/openai/codex.git
synced 2026-11-08 20:40:33 +00:00
Preserve reasoning token accounting across resume and fork (#57839)
## 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:
@@ -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
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
|
||||
Reference in New Issue
Block a user