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