From 768330dd6ca22ac7e0ab3c1efa41752959613d79 Mon Sep 17 00:00:00 2001 From: rka-oai Date: Wed, 15 Jul 2026 05:40:55 +0000 Subject: [PATCH] Add a fallback phase before automatic context rollover (#33255) ## Why Token-budget sessions need a final opportunity to preserve important state before the current context is compacted. ## What changed - When `auto_compact_fallback_prompt` is configured, reserve `auto_compact_fallback_buffer_tokens` beyond the base auto-compaction limit. - At the base limit, inject the fallback prompt once as a developer message while keeping the normal tool surface available. Continue reporting zero base-window tokens during this reserve. - Roll over when the reserve or full model context is exhausted, and skip the fallback when a new context was already requested. ## Testing Added coverage for using the fallback reserve, rolling over after it is exhausted, and bypassing the fallback for an explicit `new_context` request. GitOrigin-RevId: 32983ecadc5cb4927e15ecd7b1e896cb52a9ad3b --- codex-rs/core/src/config/mod.rs | 10 + codex-rs/core/src/context/mod.rs | 1 + .../core/src/context/token_budget_context.rs | 31 +++ codex-rs/core/src/session/context_window.rs | 61 +++--- codex-rs/core/src/session/token_budget.rs | 49 +++-- codex-rs/core/src/session/turn.rs | 10 +- .../core/src/state/auto_compact_window.rs | 10 + codex-rs/core/src/state/session.rs | 4 + .../tools/handlers/get_context_remaining.rs | 2 +- codex-rs/core/tests/suite/token_budget.rs | 178 +++++++++++++++++- 10 files changed, 303 insertions(+), 53 deletions(-) diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index a2b05dffe3..54bedaa02b 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1154,6 +1154,16 @@ pub struct TokenBudgetConfig { pub auto_compact_fallback_buffer_tokens: Option, } +impl TokenBudgetConfig { + pub(crate) fn fallback_buffer_tokens(&self) -> i64 { + if self.auto_compact_fallback_prompt.is_some() { + self.auto_compact_fallback_buffer_tokens.unwrap_or(0) + } else { + 0 + } + } +} + impl Default for TokenBudgetConfig { fn default() -> Self { Self { diff --git a/codex-rs/core/src/context/mod.rs b/codex-rs/core/src/context/mod.rs index 56c4637141..8987f90e3c 100644 --- a/codex-rs/core/src/context/mod.rs +++ b/codex-rs/core/src/context/mod.rs @@ -71,6 +71,7 @@ pub(crate) use realtime_start_with_instructions::RealtimeStartWithInstructions; pub(crate) use recommended_plugins_instructions::RecommendedPluginsInstructions; pub(crate) use rollout_budget::RolloutBudgetContext; pub(crate) use subagent_notification::SubagentNotification; +pub(crate) use token_budget_context::AutoCompactFallbackPrompt; pub(crate) use token_budget_context::ContextWindowGuidance; pub(crate) use token_budget_context::TokenBudgetContext; pub(crate) use token_budget_context::TokenBudgetRemainingContext; diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index f3eef0b734..f5a4409821 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -169,3 +169,34 @@ impl ContextualUserFragment for TokenBudgetReminder { self.message.clone() } } + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) struct AutoCompactFallbackPrompt { + message: String, +} + +impl AutoCompactFallbackPrompt { + pub(crate) fn new(message: &str) -> Self { + Self { + message: message.to_string(), + } + } +} + +impl ContextualUserFragment for AutoCompactFallbackPrompt { + fn role(&self) -> &'static str { + "developer" + } + + fn markers(&self) -> (&'static str, &'static str) { + Self::type_markers() + } + + fn type_markers() -> (&'static str, &'static str) { + ("", "") + } + + fn body(&self) -> String { + self.message.clone() + } +} diff --git a/codex-rs/core/src/session/context_window.rs b/codex-rs/core/src/session/context_window.rs index cd182919b4..466175c17d 100644 --- a/codex-rs/core/src/session/context_window.rs +++ b/codex-rs/core/src/session/context_window.rs @@ -10,15 +10,14 @@ pub(crate) struct ContextWindowTokenStatus { pub(crate) auto_compact_scope_tokens: i64, pub(crate) auto_compact_scope_limit: Option, pub(crate) full_context_window_limit: Option, - pub(crate) tokens_until_compaction: Option, + pub(crate) base_window_tokens_remaining: Option, pub(crate) auto_compact_window_prefill_tokens: Option, pub(crate) full_context_window_limit_reached: bool, pub(crate) token_limit_reached: bool, } -struct BodyAfterPrefixWindowStatus { - full_context_window_limit: Option, - auto_compact_window_prefill_tokens: Option, +fn tokens_remaining(limit: Option, used: i64) -> Option { + limit.map(|limit| limit.saturating_sub(used).max(0)) } pub(crate) async fn context_window_token_status( @@ -27,7 +26,8 @@ pub(crate) async fn context_window_token_status( ) -> ContextWindowTokenStatus { let active_context_tokens = sess.get_total_token_usage().await; - let (auto_compact_scope_tokens, auto_compact_scope_limit, body_window) = + // Count either the full active context or only the tokens added after the initial prefix. + let (auto_compact_scope_tokens, auto_compact_scope_limit, auto_compact_window_prefill_tokens) = match turn_context.config.model_auto_compact_token_limit_scope { AutoCompactTokenLimitScope::Total => ( active_context_tokens, @@ -42,49 +42,48 @@ pub(crate) async fn context_window_token_status( .config .model_auto_compact_token_limit .or_else(|| turn_context.model_info.auto_compact_token_limit()); - let full_context_window_limit = turn_context.model_context_window(); - ( active_context_tokens.saturating_sub(baseline), scope_limit, - Some(BodyAfterPrefixWindowStatus { - full_context_window_limit, - auto_compact_window_prefill_tokens: window.prefill_input_tokens, - }), + window.prefill_input_tokens, ) } }; - let full_context_window_limit = body_window - .as_ref() - .and_then(|window| window.full_context_window_limit); - let auto_compact_window_prefill_tokens = body_window - .as_ref() - .and_then(|window| window.auto_compact_window_prefill_tokens); + // The model's full context window is a hard cap, independent of the auto-compaction scope. + let full_context_window_limit = turn_context.model_context_window(); + // Report remaining tokens against the base (unbuffered) window, capped by the full context. + let base_window_tokens_remaining = [ + tokens_remaining(auto_compact_scope_limit, auto_compact_scope_tokens), + tokens_remaining(full_context_window_limit, active_context_tokens), + ] + .into_iter() + .flatten() + .min(); + + // Only reserve the fallback buffer when there is a fallback prompt to use it. + let auto_compact_fallback_buffer_tokens = turn_context + .config + .token_budget + .as_ref() + .map_or(0, crate::config::TokenBudgetConfig::fallback_buffer_tokens); + let buffered_auto_compact_limit = auto_compact_scope_limit + .map(|limit| limit.saturating_add(auto_compact_fallback_buffer_tokens)); + + // Force compaction once the buffered window or the model's full context window is reached. let full_context_window_limit_reached = - full_context_window_limit.is_some_and(|full_context_window_limit| { - active_context_tokens >= full_context_window_limit - }); - let token_limit_reached = auto_compact_scope_limit + full_context_window_limit.is_some_and(|limit| active_context_tokens >= limit); + let token_limit_reached = buffered_auto_compact_limit .is_some_and(|limit| auto_compact_scope_tokens >= limit) || full_context_window_limit_reached; - let auto_compact_scope_remaining = auto_compact_scope_limit - .map(|limit| limit.saturating_sub(auto_compact_scope_tokens).max(0)); - let full_context_remaining = - full_context_window_limit.map(|limit| limit.saturating_sub(active_context_tokens).max(0)); - let tokens_until_compaction = match (auto_compact_scope_remaining, full_context_remaining) { - (Some(scope_remaining), Some(full_remaining)) => Some(scope_remaining.min(full_remaining)), - (scope_remaining, full_remaining) => scope_remaining.or(full_remaining), - }; - ContextWindowTokenStatus { active_context_tokens, auto_compact_scope_tokens, auto_compact_scope_limit, full_context_window_limit, - tokens_until_compaction, + base_window_tokens_remaining, auto_compact_window_prefill_tokens, full_context_window_limit_reached, token_limit_reached, diff --git a/codex-rs/core/src/session/token_budget.rs b/codex-rs/core/src/session/token_budget.rs index a63f0750e1..35914119ac 100644 --- a/codex-rs/core/src/session/token_budget.rs +++ b/codex-rs/core/src/session/token_budget.rs @@ -6,35 +6,56 @@ use codex_features::Feature; pub(super) async fn maybe_record( sess: &Session, turn_context: &TurnContext, - tokens_until_compaction: Option, + base_window_tokens_remaining: Option, + allow_auto_compact_fallback: bool, ) { if !turn_context.config.features.enabled(Feature::TokenBudget) { return; } - let Some(tokens_until_compaction) = tokens_until_compaction else { + let Some(base_window_tokens_remaining) = base_window_tokens_remaining else { return; }; - let Some(config) = turn_context.config.token_budget.as_ref().filter(|config| { - config - .reminder_threshold_tokens - .is_some_and(|threshold| tokens_until_compaction <= threshold) - }) else { + let Some(config) = turn_context.config.token_budget.as_ref() else { return; }; - let reminder_due = { + if config + .reminder_threshold_tokens + .is_some_and(|threshold| base_window_tokens_remaining <= threshold) + { + let reminder_due = { + let mut state = sess.state.lock().await; + state.claim_token_budget_reminder() + }; + if reminder_due { + let response_item = + ContextualUserFragment::into(crate::context::TokenBudgetReminder::new( + &config.reminder_message_template, + base_window_tokens_remaining, + )); + sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) + .await; + } + } + + if !allow_auto_compact_fallback || base_window_tokens_remaining != 0 { + return; + } + let Some(prompt) = config.auto_compact_fallback_prompt.as_deref() else { + return; + }; + + let fallback_due = { let mut state = sess.state.lock().await; - state.claim_token_budget_reminder() + state.claim_auto_compact_fallback() }; - if !reminder_due { + if !fallback_due { return; } - let response_item = ContextualUserFragment::into(crate::context::TokenBudgetReminder::new( - &config.reminder_message_template, - tokens_until_compaction, - )); + let response_item = + ContextualUserFragment::into(crate::context::AutoCompactFallbackPrompt::new(prompt)); sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item)) .await; } diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index fd7a10d440..54e4cf070d 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -337,17 +337,19 @@ pub(crate) async fn run_turn( "post sampling token usage" ); + let should_roll_over = needs_follow_up + && (sess.take_new_context_window_request().await || token_limit_reached); + let allow_auto_compact_fallback = !should_roll_over && !token_limit_reached; super::token_budget::maybe_record( sess.as_ref(), turn_context.as_ref(), - token_status.tokens_until_compaction, + token_status.base_window_tokens_remaining, + allow_auto_compact_fallback, ) .await; // as long as compaction works well in getting us way below the token limit, we shouldn't worry about being in an infinite loop. - if needs_follow_up - && (sess.take_new_context_window_request().await || token_limit_reached) - { + if should_roll_over { if let Err(err) = run_auto_compact( &sess, Arc::clone(&step_context), diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index 6301231df6..2571749ed6 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/codex-rs/core/src/state/auto_compact_window.rs @@ -42,6 +42,7 @@ pub(super) struct AutoCompactWindow { /// resume/recompute baselines when available. prefill_input_tokens: Option, token_budget_reminder_delivered: bool, + auto_compact_fallback_delivered: bool, } impl AutoCompactWindow { @@ -52,6 +53,7 @@ impl AutoCompactWindow { new_context_window_requested: false, prefill_input_tokens: None, token_budget_reminder_delivered: false, + auto_compact_fallback_delivered: false, } } @@ -78,6 +80,7 @@ impl AutoCompactWindow { self.ids.window_id = Uuid::now_v7(); self.new_context_window_requested = false; self.token_budget_reminder_delivered = false; + self.auto_compact_fallback_delivered = false; (self.window_number, self.ids) } @@ -85,6 +88,10 @@ impl AutoCompactWindow { !std::mem::replace(&mut self.token_budget_reminder_delivered, true) } + pub(super) fn claim_auto_compact_fallback(&mut self) -> bool { + !std::mem::replace(&mut self.auto_compact_fallback_delivered, true) + } + pub(super) fn request_new_context_window(&mut self) { self.new_context_window_requested = true; } @@ -169,6 +176,8 @@ mod tests { assert_eq!(window.ids().window_id, restored_window_id); assert!(window.claim_token_budget_reminder()); assert!(!window.claim_token_budget_reminder()); + assert!(window.claim_auto_compact_fallback()); + assert!(!window.claim_auto_compact_fallback()); window.request_new_context_window(); assert!(window.take_new_context_window_request()); assert!(!window.take_new_context_window_request()); @@ -183,6 +192,7 @@ mod tests { assert_ne!(ids.window_id, restored_window_id); assert!(!window.take_new_context_window_request()); assert!(window.claim_token_budget_reminder()); + assert!(window.claim_auto_compact_fallback()); assert_eq!( window.snapshot(), diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index d769ec3872..28a08ef4d6 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -163,6 +163,10 @@ impl SessionState { self.auto_compact_window.claim_token_budget_reminder() } + pub(crate) fn claim_auto_compact_fallback(&mut self) -> bool { + self.auto_compact_window.claim_auto_compact_fallback() + } + pub(crate) fn auto_compact_window_number(&self) -> u64 { self.auto_compact_window.window_number() } diff --git a/codex-rs/core/src/tools/handlers/get_context_remaining.rs b/codex-rs/core/src/tools/handlers/get_context_remaining.rs index ffcdeb1896..77f0429722 100644 --- a/codex-rs/core/src/tools/handlers/get_context_remaining.rs +++ b/codex-rs/core/src/tools/handlers/get_context_remaining.rs @@ -82,7 +82,7 @@ impl ToolExecutor for GetContextRemainingHandler { .await; Ok(boxed_tool_output(GetContextRemainingOutput::new( - token_status.tokens_until_compaction, + token_status.base_window_tokens_remaining, ))) }) } diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index e536d6b888..38c7c47e53 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -27,6 +27,7 @@ use core_test_support::responses::ev_completed; use core_test_support::responses::ev_completed_with_tokens; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; +use core_test_support::responses::ev_shell_command_call; use core_test_support::responses::mount_compact_json_once; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; @@ -46,6 +47,7 @@ use std::path::Path; use std::time::Duration; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; +const AUTO_COMPACT_FALLBACK_PROMPT: &str = "Save the important state before rollover."; fn token_budget_contexts(request: &ResponsesRequest) -> Vec { let context_window_prefix = format!("{CONTEXT_WINDOW_OPEN_TAG}\nThread id: "); @@ -859,7 +861,171 @@ async fn token_budget_mid_turn_auto_compaction_resets_before_active_follow_up() } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { +async fn token_budget_auto_compact_fallback_uses_buffer_until_new_context() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let trigger_call_id = "trigger-call"; + let fallback_call_id = "fallback-call"; + let new_context_call_id = "new-context-call"; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(trigger_call_id, "get_context_remaining", "{}"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500), + ]), + sse(vec![ + ev_response_created("fallback-tool-resp"), + ev_shell_command_call(fallback_call_id, "echo fallback-note > fallback-note.txt"), + ev_completed_with_tokens("fallback-tool-resp", /*total_tokens*/ 10_000), + ]), + sse(vec![ + ev_response_created("new-context-resp"), + ev_function_call(new_context_call_id, "new_context", "{}"), + ev_completed_with_tokens("new-context-resp", /*total_tokens*/ 10_500), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_provider.name = "OpenAI (test)".into(); + config.model_context_window = Some(50_000); + config.model_auto_compact_token_limit = Some(9_000); + config.token_budget = Some(TokenBudgetConfig { + auto_compact_fallback_prompt: Some(AUTO_COMPACT_FALLBACK_PROMPT.to_string()), + auto_compact_fallback_buffer_tokens: Some(4_000), + ..TokenBudgetConfig::default() + }); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("trigger auto compact fallback").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 4); + let fallback_request = &requests[1]; + assert!( + fallback_request + .message_input_texts("developer") + .iter() + .any(|text| text == AUTO_COMPACT_FALLBACK_PROMPT), + "fallback prompt should be injected as a developer message" + ); + assert!( + tool_names(fallback_request) + .iter() + .any(|name| name == "get_context_remaining"), + "fallback should preserve the normal tool surface" + ); + assert_eq!( + fallback_request.function_call_output_text(trigger_call_id), + Some("You have 0 tokens left in this context window.".to_string()) + ); + let thread_id = test.session_configured.thread_id; + let initial_context = token_budget_contexts(&requests[0]); + assert_eq!(token_budget_contexts(fallback_request), initial_context); + assert_eq!(token_budget_contexts(&requests[2]), initial_context); + assert!(requests[2].body_contains_text(AUTO_COMPACT_FALLBACK_PROMPT)); + assert!(requests[2].body_contains_text(fallback_call_id)); + let (_, _, initial_window_id) = token_budget_window_ids(&initial_context[0], thread_id); + let follow_up_context = token_budget_contexts(&requests[3]); + let (_, previous_window_id, follow_up_window_id) = + token_budget_window_ids(&follow_up_context[0], thread_id); + assert_eq!( + previous_window_id.as_deref(), + Some(initial_window_id.as_str()) + ); + assert_ne!(follow_up_window_id, initial_window_id); + assert!(!requests[3].body_contains_text(AUTO_COMPACT_FALLBACK_PROMPT)); + assert_eq!( + requests[3].function_call_output_text(fallback_call_id), + None + ); + assert_eq!( + std::fs::read_to_string(test.workspace_path("fallback-note.txt"))?.trim(), + "fallback-note" + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_auto_compact_fallback_rolls_over_after_buffer() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call("trigger-call", "get_context_remaining", "{}"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500), + ]), + sse(vec![ + ev_response_created("fallback-resp"), + ev_function_call("buffer-call", "get_context_remaining", "{}"), + ev_completed_with_tokens("fallback-resp", /*total_tokens*/ 13_500), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + let test = test_codex() + .with_config(|config| { + config.model_provider.name = "OpenAI (test)".into(); + config.model_context_window = Some(50_000); + config.model_auto_compact_token_limit = Some(9_000); + config.token_budget = Some(TokenBudgetConfig { + auto_compact_fallback_prompt: Some(AUTO_COMPACT_FALLBACK_PROMPT.to_string()), + auto_compact_fallback_buffer_tokens: Some(4_000), + ..TokenBudgetConfig::default() + }); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("exhaust the fallback buffer").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + let fallback_request = &requests[1]; + assert!( + fallback_request + .message_input_texts("developer") + .iter() + .any(|text| text == AUTO_COMPACT_FALLBACK_PROMPT) + ); + assert!(!requests[2].body_contains_text(AUTO_COMPACT_FALLBACK_PROMPT)); + assert!(!requests[2].body_contains_text("exhaust the fallback buffer")); + assert_eq!(requests[2].function_call_output_text("buffer-call"), None); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn new_context_tool_skips_auto_compact_fallback() -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; @@ -877,7 +1043,7 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { sse(vec![ ev_response_created("resp-1"), ev_function_call(call_id, "new_context", "{}"), - ev_completed("resp-1"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 9_500), ]), sse(vec![ ev_response_created("resp-2"), @@ -894,7 +1060,12 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { .await; let test = test_codex() .with_config(|config| { - config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config.model_context_window = Some(10_000); + config.token_budget = Some(TokenBudgetConfig { + auto_compact_fallback_prompt: Some(AUTO_COMPACT_FALLBACK_PROMPT.to_string()), + auto_compact_fallback_buffer_tokens: Some(4_000), + ..TokenBudgetConfig::default() + }); config .features .enable(Feature::TokenBudget) @@ -907,6 +1078,7 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { let requests = responses.requests(); assert_eq!(requests.len(), 3); + assert!(!requests[1].body_contains_text(AUTO_COMPACT_FALLBACK_PROMPT)); assert!( tool_names(&requests[0]) .iter()