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

@@ -1154,6 +1154,16 @@ pub struct TokenBudgetConfig {
     pub auto_compact_fallback_buffer_tokens: Option<i64>,
 }
 
+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 {

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;

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()
+    }
+}

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<i64>,
     pub(crate) full_context_window_limit: Option<i64>,
-    pub(crate) tokens_until_compaction: Option<i64>,
+    pub(crate) base_window_tokens_remaining: Option<i64>,
     pub(crate) auto_compact_window_prefill_tokens: Option<i64>,
     pub(crate) full_context_window_limit_reached: bool,
     pub(crate) token_limit_reached: bool,
 }
 
-struct BodyAfterPrefixWindowStatus {
-    full_context_window_limit: Option<i64>,
-    auto_compact_window_prefill_tokens: Option<i64>,
+fn tokens_remaining(limit: Option<i64>, used: i64) -> Option<i64> {
+    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,

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<i64>,
+    base_window_tokens_remaining: Option<i64>,
+    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;
 }

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

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

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()
     }

codex-rs/core/src/tools/handlers/get_context_remaining.rs

@@ -82,7 +82,7 @@ impl ToolExecutor<ToolInvocation> for GetContextRemainingHandler {
             .await;
 
             Ok(boxed_tool_output(GetContextRemainingOutput::new(
-                token_status.tokens_until_compaction,
+                token_status.base_window_tokens_remaining,
             )))
         })
     }

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<String> {
     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()