[codex] add configurable token budget compaction reminder (#29255)

## Why

The token-budget feature reports coarse remaining-context milestones,
but it does not give the model a configurable wrap-up prompt before
automatic compaction. A strict threshold-crossing check can also miss
resumed or reconfigured windows that are already inside the threshold.

## What changed

- Add structured `[features.token_budget]` configuration for an absolute
`reminder_threshold_tokens` and bounded `reminder_message_template`;
`{n_remaining}` is expanded when the reminder is delivered.
- Compute remaining tokens against the next effective auto-compaction
boundary, including scoped `body_after_prefix` accounting and the full
context-window limit.
- Make reminder delivery level-triggered before and after sampling, with
one-shot state owned by `AutoCompactWindow` and re-armed on compaction,
`new_context`, restore, or history replacement.
- Leave the existing initial full-window token-budget context, 25/50/75%
notices, and token-budget tools unchanged.
- Persist the resolved feature configuration in the session config lock
and regenerate the config schema.

## Validation

- `just test -p codex-core token_budget`
- `just test -p codex-core
token_budget_reminder_emits_after_crossing_compaction_threshold`
- `just test -p codex-core auto_compact_window`
- `just test -p codex-core
lock_contains_prompts_and_materializes_features`
- `just test -p codex-features`
- `just test -p codex-config`

codex-rs/core/src/config/config_tests.rs

@@ -460,6 +460,90 @@ direct_only_tool_namespaces = ["mcp__history", "mcp__notes"]
     Ok(())
 }
 
+#[tokio::test]
+async fn load_config_resolves_token_budget_config() -> std::io::Result<()> {
+    for (config_toml, expected) in [
+        (
+            "[features]\ntoken_budget = true\n",
+            TokenBudgetConfig::default(),
+        ),
+        (
+            r#"
+[features.token_budget]
+enabled = true
+reminder_threshold_tokens = 16000
+reminder_message_template = "Custom reminder: {n_remaining} tokens."
+"#,
+            TokenBudgetConfig {
+                reminder_threshold_tokens: Some(16_000),
+                reminder_message_template: "Custom reminder: {n_remaining} tokens.".to_string(),
+            },
+        ),
+    ] {
+        let codex_home = tempdir()?;
+        let config_toml = toml::from_str(config_toml).expect("TOML should deserialize");
+        let config = Config::load_from_base_config_with_overrides(
+            config_toml,
+            ConfigOverrides::default(),
+            codex_home.abs(),
+        )
+        .await?;
+
+        assert!(config.features.enabled(Feature::TokenBudget));
+        assert_eq!(config.token_budget, Some(expected));
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn load_config_rejects_invalid_token_budget_reminder_template() -> std::io::Result<()> {
+    for reminder_message_template in [
+        String::new(),
+        "x".repeat(TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES + 1),
+    ] {
+        let codex_home = tempdir()?;
+        let config_toml = toml::from_str(&format!(
+            "[features.token_budget]\nenabled = true\nreminder_message_template = {reminder_message_template:?}\n"
+        ))
+        .expect("TOML should deserialize");
+        let error = Config::load_from_base_config_with_overrides(
+            config_toml,
+            ConfigOverrides::default(),
+            codex_home.abs(),
+        )
+        .await
+        .expect_err("invalid reminder template should be rejected");
+
+        assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
+    }
+    Ok(())
+}
+
+#[tokio::test]
+async fn load_config_rejects_non_positive_token_budget_reminder_threshold() -> std::io::Result<()> {
+    for reminder_threshold_tokens in [-1, 0] {
+        let codex_home = tempdir()?;
+        let config_toml = toml::from_str(&format!(
+            "[features.token_budget]\nenabled = true\nreminder_threshold_tokens = {reminder_threshold_tokens}\n"
+        ))
+        .expect("TOML should deserialize");
+        let error = Config::load_from_base_config_with_overrides(
+            config_toml,
+            ConfigOverrides::default(),
+            codex_home.abs(),
+        )
+        .await
+        .expect_err("non-positive reminder threshold should be rejected");
+
+        assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
+        assert_eq!(
+            error.to_string(),
+            "features.token_budget.reminder_threshold_tokens must be positive"
+        );
+    }
+    Ok(())
+}
+
 #[tokio::test]
 async fn load_config_resolves_rollout_budget() -> std::io::Result<()> {
     let codex_home = tempdir()?;

codex-rs/core/src/config/mod.rs

@@ -68,6 +68,7 @@ use codex_features::Features;
 use codex_features::FeaturesToml;
 use codex_features::MultiAgentV2ConfigToml;
 use codex_features::NetworkProxyConfigToml;
+use codex_features::TokenBudgetConfigToml;
 use codex_git_utils::resolve_root_git_project_for_trust;
 use codex_install_context::InstallContext;
 use codex_login::AuthManagerConfig;
@@ -1026,6 +1027,8 @@ pub struct Config {
     /// Settings specific to the task-path-based multi-agent tool surface.
     pub multi_agent_v2: MultiAgentV2Config,
 
+    /// Context-window token budget configuration, when enabled.
+    pub token_budget: Option<TokenBudgetConfig>,
     /// Shared token budget for the root thread and its sub-agents.
     pub rollout_budget: Option<RolloutBudgetConfig>,
     /// Current-time reminder configuration, when enabled.
@@ -1075,6 +1078,27 @@ pub struct CodeModeConfig {
     pub direct_only_tool_namespaces: Vec<String>,
 }
 
+pub(crate) const DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE: &str = concat!(
+    "Your context window is nearly exhausted (only {n_remaining} tokens remaining) and will be automatically reset for you soon. ",
+    "Once reset, message items in current context window will be cleared in the new window, but notes and history items will be persistent across windows."
+);
+const TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES: usize = 1000;
+
+#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
+pub struct TokenBudgetConfig {
+    pub reminder_threshold_tokens: Option<i64>,
+    pub reminder_message_template: String,
+}
+
+impl Default for TokenBudgetConfig {
+    fn default() -> Self {
+        Self {
+            reminder_threshold_tokens: None,
+            reminder_message_template: DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string(),
+        }
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Serialize)]
 pub struct RolloutBudgetConfig {
     pub limit_tokens: i64,
@@ -2509,6 +2533,48 @@ fn resolve_multi_agent_v2_config(config_toml: &ConfigToml) -> MultiAgentV2Config
     }
 }
 
+fn resolve_token_budget_config(
+    config_toml: &ConfigToml,
+    features: &ManagedFeatures,
+) -> std::io::Result<Option<TokenBudgetConfig>> {
+    if !features.enabled(Feature::TokenBudget) {
+        return Ok(None);
+    }
+
+    let token_budget_config = token_budget_toml_config(config_toml.features.as_ref());
+    let reminder_threshold_tokens =
+        token_budget_config.and_then(|config| config.reminder_threshold_tokens);
+    if reminder_threshold_tokens.is_some_and(|tokens| tokens <= 0) {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "features.token_budget.reminder_threshold_tokens must be positive",
+        ));
+    }
+
+    let reminder_message_template = token_budget_config
+        .and_then(|config| config.reminder_message_template.clone())
+        .unwrap_or_else(|| DEFAULT_TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE.to_string());
+    if reminder_message_template.trim().is_empty() {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "features.token_budget.reminder_message_template must not be empty",
+        ));
+    }
+    if reminder_message_template.len() > TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            format!(
+                "features.token_budget.reminder_message_template must not exceed {TOKEN_BUDGET_REMINDER_MESSAGE_TEMPLATE_MAX_BYTES} bytes"
+            ),
+        ));
+    }
+
+    Ok(Some(TokenBudgetConfig {
+        reminder_threshold_tokens,
+        reminder_message_template,
+    }))
+}
+
 fn resolve_rollout_budget_config(
     config_toml: &ConfigToml,
     features: &ManagedFeatures,
@@ -2635,6 +2701,13 @@ fn multi_agent_v2_toml_config(features: Option<&FeaturesToml>) -> Option<&MultiA
     }
 }
 
+fn token_budget_toml_config(features: Option<&FeaturesToml>) -> Option<&TokenBudgetConfigToml> {
+    match features?.token_budget.as_ref()? {
+        FeatureToml::Enabled(_) => None,
+        FeatureToml::Config(config) => Some(config),
+    }
+}
+
 fn current_time_reminder_toml_config(
     features: Option<&FeaturesToml>,
 ) -> Option<&CurrentTimeReminderConfigToml> {
@@ -3287,6 +3360,7 @@ impl Config {
             resolve_experimental_request_user_input_enabled(&cfg);
         let code_mode = resolve_code_mode_config(&cfg);
         let multi_agent_v2 = resolve_multi_agent_v2_config(&cfg);
+        let token_budget = resolve_token_budget_config(&cfg, &features)?;
         let rollout_budget = resolve_rollout_budget_config(&cfg, &features)?;
         let current_time_reminder = resolve_current_time_reminder_config(&cfg, &features)?;
         let terminal_resize_reflow = resolve_terminal_resize_reflow_config(&cfg);
@@ -3830,6 +3904,7 @@ impl Config {
             background_terminal_max_timeout,
             ghost_snapshot,
             multi_agent_v2,
+            token_budget,
             rollout_budget,
             current_time_reminder,
             features,

codex-rs/core/src/context/mod.rs

@@ -73,6 +73,7 @@ pub(crate) use rollout_budget::RolloutBudgetContext;
 pub(crate) use subagent_notification::SubagentNotification;
 pub(crate) use token_budget_context::TokenBudgetContext;
 pub(crate) use token_budget_context::TokenBudgetRemainingContext;
+pub(crate) use token_budget_context::TokenBudgetReminder;
 pub(crate) use turn_aborted::TurnAborted;
 pub(crate) use user_instructions::UserInstructions;
 pub(crate) use user_shell_command::UserShellCommand;

codex-rs/core/src/context/token_budget_context.rs

@@ -99,3 +99,34 @@ impl ContextualUserFragment for TokenBudgetRemainingContext {
         }
     }
 }
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct TokenBudgetReminder {
+    message: String,
+}
+
+impl TokenBudgetReminder {
+    pub(crate) fn new(message_template: &str, n_remaining: i64) -> Self {
+        Self {
+            message: message_template.replace("{n_remaining}", &n_remaining.to_string()),
+        }
+    }
+}
+
+impl ContextualUserFragment for TokenBudgetReminder {
+    fn role(&self) -> &'static str {
+        "developer"
+    }
+
+    fn markers(&self) -> (&'static str, &'static str) {
+        Self::type_markers()
+    }
+
+    fn type_markers() -> (&'static str, &'static str) {
+        ("<token_budget>\n", "\n</token_budget>")
+    }
+
+    fn body(&self) -> String {
+        self.message.clone()
+    }
+}

codex-rs/core/src/session/config_lock.rs

@@ -10,6 +10,7 @@ use codex_features::FeatureToml;
 use codex_features::FeaturesToml;
 use codex_features::MultiAgentV2ConfigToml;
 use codex_features::RolloutBudgetConfigToml;
+use codex_features::TokenBudgetConfigToml;
 use codex_protocol::ThreadId;
 
 use crate::config::Config;
@@ -152,6 +153,12 @@ fn save_config_resolved_fields(
         resolved_config_to_toml(&config.multi_agent_v2, "features.multi_agent_v2")?;
     multi_agent_v2.enabled = Some(config.features.enabled(Feature::MultiAgentV2));
     features.multi_agent_v2 = Some(FeatureToml::Config(multi_agent_v2));
+    if let Some(token_budget) = config.token_budget.as_ref() {
+        let mut token_budget: TokenBudgetConfigToml =
+            resolved_config_to_toml(token_budget, "features.token_budget")?;
+        token_budget.enabled = Some(config.features.enabled(Feature::TokenBudget));
+        features.token_budget = Some(FeatureToml::Config(token_budget));
+    }
     if let Some(rollout_budget) = config.rollout_budget.as_ref() {
         let mut rollout_budget: RolloutBudgetConfigToml =
             resolved_config_to_toml(rollout_budget, "features.rollout_budget")?;
@@ -238,6 +245,14 @@ mod tests {
     async fn lock_contains_prompts_and_materializes_features() {
         let mut sc = crate::session::tests::make_session_configuration_for_tests().await;
         let mut config = (*sc.original_config_do_not_use).clone();
+        config.token_budget = Some(crate::config::TokenBudgetConfig {
+            reminder_threshold_tokens: Some(16_000),
+            reminder_message_template: "Locked reminder: {n_remaining} tokens.".to_string(),
+        });
+        config
+            .features
+            .enable(Feature::TokenBudget)
+            .expect("token_budget should be enableable in tests");
         config.rollout_budget = Some(crate::config::RolloutBudgetConfig {
             limit_tokens: 100_000,
             reminder_interval_tokens: 10_000,
@@ -318,6 +333,17 @@ mod tests {
             })
         ));
 
+        assert_eq!(
+            features.token_budget,
+            Some(FeatureToml::Config(TokenBudgetConfigToml {
+                enabled: Some(true),
+                reminder_threshold_tokens: Some(16_000),
+                reminder_message_template: Some(
+                    "Locked reminder: {n_remaining} tokens.".to_string()
+                ),
+            }))
+        );
+
         assert_eq!(
             features.rollout_budget,
             Some(FeatureToml::Config(RolloutBudgetConfigToml {

codex-rs/core/src/session/token_budget.rs

@@ -5,40 +5,61 @@ use codex_features::Feature;
 
 const TOKEN_BUDGET_USAGE_THRESHOLDS: [i64; 3] = [25, 50, 75];
 
-pub(super) async fn maybe_record_token_budget_remaining_context(
+pub(super) async fn maybe_record(
     sess: &Session,
     turn_context: &TurnContext,
     tokens_before_sampling: i64,
     tokens_after_sampling: i64,
+    tokens_until_compaction: i64,
 ) {
     if !turn_context.config.features.enabled(Feature::TokenBudget) {
         return;
     }
-    let Some(model_context_window) = turn_context.model_context_window() else {
-        return;
-    };
-    if model_context_window <= 0 || tokens_after_sampling <= tokens_before_sampling {
-        return;
+
+    let mut response_items = Vec::new();
+    if let Some(model_context_window) = turn_context.model_context_window()
+        && model_context_window > 0
+        && tokens_after_sampling > tokens_before_sampling
+    {
+        let tokens_before_sampling = tokens_before_sampling.max(0);
+        let tokens_after_sampling = tokens_after_sampling.max(0);
+        let crossed_threshold = TOKEN_BUDGET_USAGE_THRESHOLDS.iter().any(|threshold| {
+            tokens_before_sampling.saturating_mul(100)
+                < model_context_window.saturating_mul(*threshold)
+                && tokens_after_sampling.saturating_mul(100)
+                    >= model_context_window.saturating_mul(*threshold)
+        });
+        if crossed_threshold {
+            let tokens_left = model_context_window
+                .saturating_sub(tokens_after_sampling)
+                .max(0);
+            response_items.push(ContextualUserFragment::into(
+                crate::context::TokenBudgetRemainingContext::new(tokens_left),
+            ));
+        }
     }
 
-    let tokens_before_sampling = tokens_before_sampling.max(0);
-    let tokens_after_sampling = tokens_after_sampling.max(0);
-    let crossed_threshold = TOKEN_BUDGET_USAGE_THRESHOLDS.iter().any(|threshold| {
-        tokens_before_sampling.saturating_mul(100) < model_context_window.saturating_mul(*threshold)
-            && tokens_after_sampling.saturating_mul(100)
-                >= model_context_window.saturating_mul(*threshold)
-    });
-    if !crossed_threshold {
-        return;
+    if let Some(config) = turn_context.config.token_budget.as_ref().filter(|config| {
+        config
+            .reminder_threshold_tokens
+            .is_some_and(|threshold| tokens_until_compaction <= threshold)
+    }) {
+        let reminder_due = {
+            let mut state = sess.state.lock().await;
+            state.claim_token_budget_reminder()
+        };
+        if reminder_due {
+            response_items.push(ContextualUserFragment::into(
+                crate::context::TokenBudgetReminder::new(
+                    &config.reminder_message_template,
+                    tokens_until_compaction,
+                ),
+            ));
+        }
     }
 
-    let tokens_left = model_context_window
-        .saturating_sub(tokens_after_sampling)
-        .max(0);
-
-    let response_item = ContextualUserFragment::into(
-        crate::context::TokenBudgetRemainingContext::new(tokens_left),
-    );
-    sess.record_conversation_items(turn_context, std::slice::from_ref(&response_item))
-        .await;
+    if !response_items.is_empty() {
+        sess.record_conversation_items(turn_context, &response_items)
+            .await;
+    }
 }

codex-rs/core/src/session/turn.rs

@@ -313,11 +313,22 @@ pub(crate) async fn run_turn(
                 );
 
                 let tokens_after_sampling = token_status.active_context_tokens;
-                super::token_budget::maybe_record_token_budget_remaining_context(
+                let full_context_remaining = token_status
+                    .full_context_window_limit
+                    .map_or(i64::MAX, |limit| {
+                        limit.saturating_sub(tokens_after_sampling)
+                    });
+                let tokens_until_compaction = token_status
+                    .auto_compact_scope_limit
+                    .saturating_sub(token_status.auto_compact_scope_tokens)
+                    .min(full_context_remaining)
+                    .max(0);
+                super::token_budget::maybe_record(
                     sess.as_ref(),
                     turn_context.as_ref(),
                     tokens_before_sampling,
                     tokens_after_sampling,
+                    tokens_until_compaction,
                 )
                 .await;
 

codex-rs/core/src/state/auto_compact_window.rs

@@ -30,6 +30,7 @@ pub(super) struct AutoCompactWindow {
     /// not the growth itself; server-observed usage replaces estimated
     /// resume/recompute baselines when available.
     prefill_input_tokens: Option<AutoCompactWindowPrefill>,
+    token_budget_reminder_delivered: bool,
 }
 
 impl AutoCompactWindow {
@@ -44,6 +45,7 @@ impl AutoCompactWindow {
             },
             new_context_window_requested: false,
             prefill_input_tokens: None,
+            token_budget_reminder_delivered: false,
         }
     }
 
@@ -69,9 +71,14 @@ impl AutoCompactWindow {
         self.ids.previous_window_id = Some(self.ids.window_id);
         self.ids.window_id = Uuid::now_v7();
         self.new_context_window_requested = false;
+        self.token_budget_reminder_delivered = false;
         (self.window_number, self.ids)
     }
 
+    pub(super) fn claim_token_budget_reminder(&mut self) -> bool {
+        !std::mem::replace(&mut self.token_budget_reminder_delivered, true)
+    }
+
     pub(super) fn request_new_context_window(&mut self) {
         self.new_context_window_requested = true;
     }
@@ -154,6 +161,8 @@ mod tests {
         );
         assert_eq!(window.window_number(), 3);
         assert_eq!(window.ids().window_id, restored_window_id);
+        assert!(window.claim_token_budget_reminder());
+        assert!(!window.claim_token_budget_reminder());
         window.request_new_context_window();
         assert!(window.take_new_context_window_request());
         assert!(!window.take_new_context_window_request());
@@ -167,6 +176,7 @@ mod tests {
         assert_eq!(ids.window_id.get_version_num(), 7);
         assert_ne!(ids.window_id, restored_window_id);
         assert!(!window.take_new_context_window_request());
+        assert!(window.claim_token_budget_reminder());
 
         assert_eq!(
             window.snapshot(),

codex-rs/core/src/state/session.rs

@@ -148,6 +148,10 @@ impl SessionState {
         self.auto_compact_window.snapshot()
     }
 
+    pub(crate) fn claim_token_budget_reminder(&mut self) -> bool {
+        self.auto_compact_window.claim_token_budget_reminder()
+    }
+
     pub(crate) fn auto_compact_window_number(&self) -> u64 {
         self.auto_compact_window.window_number()
     }