[codex] Add token budget context feature (#27438)

## Why

The model should be able to see bounded context-window budget metadata
when the `token_budget` feature is enabled. The full-window message is
only injected with full context, while normal turns get a smaller
follow-up only when reported usage first crosses a budget threshold.

## What changed

- Added the `TokenBudget` feature flag.
- Added `<token_budget>` developer fragments for full context-window
metadata and current-window remaining tokens.
- Inserted the threshold message during normal turn handling by
comparing token usage before and after sampling, avoiding persistent
threshold bookkeeping.
- Added core integration coverage for full-context-only metadata and
25/50/75 percent threshold messages.

## Verification

- `just test -p codex-core token_budget`
- `git diff --check`

codex-rs/core/src/compact.rs

@@ -294,6 +294,7 @@ async fn run_compact_task_inner_impl(
     let user_messages = collect_user_messages(history_items);
 
     let mut new_history = build_compacted_history(Vec::new(), &user_messages, &summary_text);
+    let window_id = sess.advance_auto_compact_window_id().await;
 
     if matches!(
         initial_context_injection,
@@ -310,7 +311,7 @@ async fn run_compact_task_inner_impl(
     let compacted_item = CompactedItem {
         message: summary_text.clone(),
         replacement_history: Some(new_history.clone()),
-        window_id: None,
+        window_id: Some(window_id),
     };
     sess.replace_compacted_history(new_history, reference_context_item, compacted_item)
         .await;

codex-rs/core/src/compact_remote.rs

@@ -249,6 +249,7 @@ async fn run_remote_compact_task_inner_impl(
             turn_metadata_header.as_deref(),
         )
         .await?;
+    let new_window_id = sess.advance_auto_compact_window_id().await;
     new_history = process_compacted_history(
         sess.as_ref(),
         turn_context.as_ref(),
@@ -264,7 +265,7 @@ async fn run_remote_compact_task_inner_impl(
     let compacted_item = CompactedItem {
         message: String::new(),
         replacement_history: Some(new_history.clone()),
-        window_id: None,
+        window_id: Some(new_window_id),
     };
     // Install is the semantic boundary where the compact endpoint's output becomes live
     // thread history. Keep it distinct from the later inference request so the reducer can

codex-rs/core/src/compact_remote_v2.rs

@@ -285,6 +285,7 @@ async fn run_remote_compact_task_inner_impl(
     let (compacted_history, retained_images) =
         build_v2_compacted_history(&prompt_input, compaction_output);
     analytics_details.retained_image_count = Some(retained_images);
+    let new_window_id = sess.advance_auto_compact_window_id().await;
     let new_history = process_compacted_history(
         sess.as_ref(),
         turn_context.as_ref(),
@@ -300,7 +301,7 @@ async fn run_remote_compact_task_inner_impl(
     let compacted_item = CompactedItem {
         message: String::new(),
         replacement_history: Some(new_history.clone()),
-        window_id: None,
+        window_id: Some(new_window_id),
     };
     compaction_trace.record_installed(&CompactionCheckpointTracePayload {
         input_history: &trace_input_history,

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

@@ -23,6 +23,7 @@ mod realtime_end_instructions;
 mod realtime_start_instructions;
 mod realtime_start_with_instructions;
 mod subagent_notification;
+mod token_budget_context;
 mod turn_aborted;
 mod user_instructions;
 mod user_shell_command;
@@ -60,6 +61,8 @@ pub(crate) use realtime_end_instructions::RealtimeEndInstructions;
 pub(crate) use realtime_start_instructions::RealtimeStartInstructions;
 pub(crate) use realtime_start_with_instructions::RealtimeStartWithInstructions;
 pub(crate) use subagent_notification::SubagentNotification;
+pub(crate) use token_budget_context::TokenBudgetContext;
+pub(crate) use token_budget_context::TokenBudgetRemainingContext;
 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

@@ -0,0 +1,68 @@
+use super::ContextualUserFragment;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct TokenBudgetContext {
+    window_id: u64,
+    tokens_left: i64,
+}
+
+impl TokenBudgetContext {
+    pub(crate) fn new(window_id: u64, tokens_left: i64) -> Self {
+        Self {
+            window_id,
+            tokens_left,
+        }
+    }
+}
+
+impl ContextualUserFragment for TokenBudgetContext {
+    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 {
+        let window_id = self.window_id;
+        let tokens_left = self.tokens_left;
+        format!(
+            "Current context window {window_id}.\nYou have {tokens_left} tokens left in this context window."
+        )
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub(crate) struct TokenBudgetRemainingContext {
+    tokens_left: i64,
+}
+
+impl TokenBudgetRemainingContext {
+    pub(crate) fn new(tokens_left: i64) -> Self {
+        Self { tokens_left }
+    }
+}
+
+impl ContextualUserFragment for TokenBudgetRemainingContext {
+    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 {
+        let tokens_left = self.tokens_left;
+        format!("You have {tokens_left} tokens left in this context window.")
+    }
+}

codex-rs/core/src/event_mapping.rs

@@ -30,6 +30,7 @@ const CONTEXTUAL_DEVELOPER_PREFIXES: &[&str] = &[
     COLLABORATION_MODE_OPEN_TAG,
     REALTIME_CONVERSATION_OPEN_TAG,
     "<personality_spec>",
+    "<token_budget>",
 ];
 
 pub(crate) fn is_contextual_user_message_content(message: &[ContentItem]) -> bool {

codex-rs/core/src/event_mapping_tests.rs

@@ -1,3 +1,5 @@
+use super::has_non_contextual_dev_message_content;
+use super::is_contextual_dev_message_content;
 use super::parse_turn_item;
 use crate::context::ContextualUserFragment;
 use crate::context::InternalContextSource;
@@ -16,6 +18,17 @@ use codex_protocol::models::WebSearchAction;
 use codex_protocol::user_input::UserInput;
 use pretty_assertions::assert_eq;
 
+#[test]
+fn recognizes_token_budget_as_contextual_developer_content() {
+    let content = vec![ContentItem::InputText {
+        text: "<token_budget>\nYou have 710 tokens left in this context window.\n</token_budget>"
+            .to_string(),
+    }];
+
+    assert!(is_contextual_dev_message_content(&content));
+    assert!(!has_non_contextual_dev_message_content(&content));
+}
+
 #[test]
 fn parses_user_message_with_text_and_two_images() {
     let img1 = "https://example.com/one.png".to_string();

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

@@ -207,6 +207,7 @@ mod review;
 mod rollout_reconstruction;
 #[allow(clippy::module_inception)]
 pub(crate) mod session;
+mod token_budget;
 pub(crate) mod turn;
 pub(crate) mod turn_context;
 use self::config_lock::export_config_lock_if_configured;
@@ -2712,15 +2713,13 @@ impl Session {
         &self,
         items: Vec<ResponseItem>,
         reference_context_item: Option<TurnContextItem>,
-        mut compacted_item: CompactedItem,
+        compacted_item: CompactedItem,
     ) {
         {
             let mut state = self.state.lock().await;
             state.replace_history(items, reference_context_item.clone());
         }
 
-        compacted_item.window_id = Some(self.advance_auto_compact_window_id().await);
-
         self.persist_rollout_items(&[RolloutItem::Compacted(compacted_item)])
             .await;
         if let Some(turn_context_item) = reference_context_item {
@@ -2805,6 +2804,7 @@ impl Session {
             collaboration_mode,
             base_instructions,
             session_source,
+            auto_compact_window_id,
         ) = {
             let state = self.state.lock().await;
             (
@@ -2813,6 +2813,7 @@ impl Session {
                 state.session_configuration.collaboration_mode.clone(),
                 state.session_configuration.base_instructions.clone(),
                 state.session_configuration.session_source.clone(),
+                state.auto_compact_window_id(),
             )
         };
         if let Some(model_switch_message) =
@@ -2962,6 +2963,18 @@ impl Session {
                 .render(),
             );
         }
+        // This is full-context metadata. Steady-state context diffs should not re-emit it.
+        if turn_context.features.enabled(Feature::TokenBudget)
+            && let Some(model_context_window) = turn_context.model_context_window()
+        {
+            developer_sections.push(
+                crate::context::TokenBudgetContext::new(
+                    auto_compact_window_id,
+                    model_context_window,
+                )
+                .render(),
+            );
+        }
         if turn_context.config.include_environment_context {
             let shell = self.user_shell();
             let subagents = self
@@ -3040,7 +3053,7 @@ impl Session {
         format!("{thread_id}:{window_id}")
     }
 
-    async fn advance_auto_compact_window_id(&self) -> u64 {
+    pub(crate) async fn advance_auto_compact_window_id(&self) -> u64 {
         let mut state = self.state.lock().await;
         state.advance_auto_compact_window_id()
     }

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

@@ -0,0 +1,44 @@
+use super::session::Session;
+use super::turn_context::TurnContext;
+use crate::context::ContextualUserFragment;
+use codex_features::Feature;
+
+const TOKEN_BUDGET_USAGE_THRESHOLDS: [i64; 3] = [25, 50, 75];
+
+pub(super) async fn maybe_record_token_budget_remaining_context(
+    sess: &Session,
+    turn_context: &TurnContext,
+    tokens_before_sampling: i64,
+    tokens_after_sampling: i64,
+) {
+    if !turn_context.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 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;
+    }
+
+    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;
+}

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

@@ -225,6 +225,7 @@ pub(crate) async fn run_turn(
         let turn_metadata_header = turn_context
             .turn_metadata_state
             .current_header_value_for_model_request(&window_id);
+        let tokens_before_sampling = sess.get_total_token_usage().await;
         match run_sampling_request(
             Arc::clone(&sess),
             Arc::clone(&turn_context),
@@ -275,6 +276,15 @@ pub(crate) async fn run_turn(
                     "post sampling token usage"
                 );
 
+                let tokens_after_sampling = token_status.active_context_tokens;
+                super::token_budget::maybe_record_token_budget_remaining_context(
+                    sess.as_ref(),
+                    turn_context.as_ref(),
+                    tokens_before_sampling,
+                    tokens_after_sampling,
+                )
+                .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 token_limit_reached && needs_follow_up {
                     if let Err(err) = run_auto_compact(