[codex] Add new context window tool (#27488)

## Why

The token budget feature tells the model how much room remains in the
current context window. When the model decides the current window is no
longer useful, it needs a way to ask Codex to start over with a fresh
context window without spending tokens on a compaction summary.

This PR adds that model-requestable escape hatch on top of #27438.

## What changed

- Added a direct-model-only `new_context` tool behind
`Feature::TokenBudget`.
- Stores the tool request on `AutoCompactWindow` and consumes it after
sampling so the next follow-up request in the same turn starts in the
new window.
- Starts the new window as a no-summary compaction checkpoint that
contains only fresh initial context, not preserved conversation history.
- Keeps the new window aligned with token-budget startup context,
including the `Current context window Z` message.
- Added integration coverage and a snapshot showing the same-turn
`new_context` flow into a fresh full-context follow-up request.

## Validation

- `just test -p codex-core token_budget`

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

@@ -3058,6 +3058,44 @@ impl Session {
         state.advance_auto_compact_window_id()
     }
 
+    pub(crate) async fn request_new_context_window(&self) {
+        let mut state = self.state.lock().await;
+        state.request_new_context_window();
+    }
+
+    pub(crate) async fn maybe_start_new_context_window(
+        &self,
+        turn_context: &TurnContext,
+    ) -> Option<u64> {
+        let window_id = {
+            let mut state = self.state.lock().await;
+            state.start_new_context_window_if_requested()
+        };
+        let window_id = window_id?;
+        let context_items = self.build_initial_context(turn_context).await;
+        let turn_context_item = turn_context.to_turn_context_item();
+        let replacement_history = context_items;
+        {
+            let mut state = self.state.lock().await;
+            state.replace_history(replacement_history.clone(), Some(turn_context_item.clone()));
+        };
+        self.persist_rollout_items(&[
+            RolloutItem::Compacted(CompactedItem {
+                message: String::new(),
+                replacement_history: Some(replacement_history),
+                window_id: Some(window_id),
+            }),
+            RolloutItem::TurnContext(turn_context_item),
+        ])
+        .await;
+        {
+            let mut state = self.state.lock().await;
+            state.queue_pending_session_start_source(codex_hooks::SessionStartSource::Compact);
+        }
+        self.recompute_token_usage(turn_context).await;
+        Some(window_id)
+    }
+
     pub(crate) async fn reference_context_item(&self) -> Option<TurnContextItem> {
         let state = self.state.lock().await;
         state.reference_context_item()

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

@@ -285,6 +285,15 @@ pub(crate) async fn run_turn(
                 )
                 .await;
 
+                let started_new_context_window = sess
+                    .maybe_start_new_context_window(turn_context.as_ref())
+                    .await
+                    .is_some();
+                if started_new_context_window && needs_follow_up {
+                    can_drain_pending_input = !model_needs_follow_up;
+                    continue;
+                }
+
                 // 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(

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

@@ -14,6 +14,7 @@ enum AutoCompactWindowPrefill {
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub(super) struct AutoCompactWindow {
     window_id: u64,
+    new_context_window_requested: bool,
     /// Absolute input-token baseline for the current compaction window.
     ///
     /// `body_after_prefix` subtracts this from later active-context usage. It is
@@ -26,6 +27,7 @@ impl AutoCompactWindow {
     pub(super) fn new() -> Self {
         Self {
             window_id: 0,
+            new_context_window_requested: false,
             prefill_input_tokens: None,
         }
     }
@@ -44,9 +46,20 @@ impl AutoCompactWindow {
 
     pub(super) fn advance_window_id(&mut self) -> u64 {
         self.window_id = self.window_id.saturating_add(1);
+        self.new_context_window_requested = false;
         self.window_id
     }
 
+    pub(super) fn request_new_context_window(&mut self) {
+        self.new_context_window_requested = true;
+    }
+
+    pub(super) fn take_new_context_window_request(&mut self) -> bool {
+        let requested = self.new_context_window_requested;
+        self.new_context_window_requested = false;
+        requested
+    }
+
     /// Records the request-input side of the first server usage sample. The
     /// sampled output from that response is body growth and should remain
     /// counted against the scoped auto-compact budget.
@@ -98,8 +111,13 @@ mod tests {
         assert_eq!(window.window_id(), 0);
         window.set_window_id(/*window_id*/ 3);
         assert_eq!(window.window_id(), 3);
+        window.request_new_context_window();
+        assert!(window.take_new_context_window_request());
+        assert!(!window.take_new_context_window_request());
+        window.request_new_context_window();
         assert_eq!(window.advance_window_id(), 4);
         assert_eq!(window.window_id(), 4);
+        assert!(!window.take_new_context_window_request());
 
         assert_eq!(
             window.snapshot(),

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

@@ -156,6 +156,20 @@ impl SessionState {
         self.auto_compact_window.advance_window_id()
     }
 
+    pub(crate) fn request_new_context_window(&mut self) {
+        self.auto_compact_window.request_new_context_window();
+    }
+
+    pub(crate) fn start_new_context_window_if_requested(&mut self) -> Option<u64> {
+        if !self.auto_compact_window.take_new_context_window_request() {
+            return None;
+        }
+
+        let window_id = self.auto_compact_window.advance_window_id();
+        self.auto_compact_window.clear_prefill();
+        Some(window_id)
+    }
+
     pub(crate) fn token_info(&self) -> Option<TokenUsageInfo> {
         self.history.token_info()
     }

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

@@ -13,6 +13,8 @@ pub(crate) mod multi_agents;
 pub(crate) mod multi_agents_common;
 pub(crate) mod multi_agents_spec;
 pub(crate) mod multi_agents_v2;
+mod new_context_window;
+pub(crate) mod new_context_window_spec;
 mod plan;
 pub(crate) mod plan_spec;
 mod request_permissions;
@@ -56,6 +58,7 @@ pub use mcp::McpHandler;
 pub use mcp_resource::ListMcpResourceTemplatesHandler;
 pub use mcp_resource::ListMcpResourcesHandler;
 pub use mcp_resource::ReadMcpResourceHandler;
+pub use new_context_window::NewContextWindowHandler;
 pub use plan::PlanHandler;
 pub use request_permissions::RequestPermissionsHandler;
 pub use request_plugin_install::RequestPluginInstallHandler;

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

@@ -0,0 +1,45 @@
+use crate::function_tool::FunctionCallError;
+use crate::tools::context::FunctionToolOutput;
+use crate::tools::context::ToolInvocation;
+use crate::tools::context::ToolPayload;
+use crate::tools::context::boxed_tool_output;
+use crate::tools::handlers::new_context_window_spec::NEW_CONTEXT_WINDOW_TOOL_NAME;
+use crate::tools::handlers::new_context_window_spec::create_new_context_window_tool;
+use crate::tools::registry::CoreToolRuntime;
+use crate::tools::registry::ToolExecutor;
+use codex_tools::ToolName;
+use codex_tools::ToolSpec;
+
+pub(crate) const NEW_CONTEXT_WINDOW_MESSAGE: &str =
+    "A new context window will start without summarizing conversation history.";
+
+pub struct NewContextWindowHandler;
+
+impl ToolExecutor<ToolInvocation> for NewContextWindowHandler {
+    fn tool_name(&self) -> ToolName {
+        ToolName::plain(NEW_CONTEXT_WINDOW_TOOL_NAME)
+    }
+
+    fn spec(&self) -> ToolSpec {
+        create_new_context_window_tool()
+    }
+
+    fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> {
+        Box::pin(async move {
+            if !matches!(invocation.payload, ToolPayload::Function { .. }) {
+                return Err(FunctionCallError::RespondToModel(
+                    "new_context handler received unsupported payload".to_string(),
+                ));
+            }
+
+            invocation.session.request_new_context_window().await;
+
+            Ok(boxed_tool_output(FunctionToolOutput::from_text(
+                NEW_CONTEXT_WINDOW_MESSAGE.to_string(),
+                Some(true),
+            )))
+        })
+    }
+}
+
+impl CoreToolRuntime for NewContextWindowHandler {}

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

@@ -0,0 +1,17 @@
+use codex_tools::JsonSchema;
+use codex_tools::ResponsesApiTool;
+use codex_tools::ToolSpec;
+use std::collections::BTreeMap;
+
+pub(crate) const NEW_CONTEXT_WINDOW_TOOL_NAME: &str = "new_context";
+
+pub fn create_new_context_window_tool() -> ToolSpec {
+    ToolSpec::Function(ResponsesApiTool {
+        name: NEW_CONTEXT_WINDOW_TOOL_NAME.to_string(),
+        description: "Start a new context window.".to_string(),
+        strict: false,
+        defer_loading: None,
+        parameters: JsonSchema::object(BTreeMap::new(), /*required*/ None, Some(false.into())),
+        output_schema: None,
+    })
+}

codex-rs/core/src/tools/spec_plan.rs

@@ -13,6 +13,7 @@ use crate::tools::handlers::ListAvailablePluginsToInstallHandler;
 use crate::tools::handlers::ListMcpResourceTemplatesHandler;
 use crate::tools::handlers::ListMcpResourcesHandler;
 use crate::tools::handlers::McpHandler;
+use crate::tools::handlers::NewContextWindowHandler;
 use crate::tools::handlers::PlanHandler;
 use crate::tools::handlers::ReadMcpResourceHandler;
 use crate::tools::handlers::RequestPermissionsHandler;
@@ -659,6 +660,10 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut
         planned_tools.add(RequestPermissionsHandler);
     }
 
+    if features.enabled(Feature::TokenBudget) {
+        planned_tools.add_with_exposure(NewContextWindowHandler, ToolExposure::DirectModelOnly);
+    }
+
     if tool_suggest_enabled(turn_context)
         && let Some(discoverable_tools) =
             context.discoverable_tools.filter(|tools| !tools.is_empty())

codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap

@@ -0,0 +1,15 @@
+---
+source: core/tests/suite/token_budget.rs
+assertion_line: 300
+expression: "context_snapshot::format_labeled_requests_snapshot(\"New context window tool installs fresh full context before the next follow-up request.\",\n&[(\"Final Follow-Up Request\", &requests[2])],\n&ContextSnapshotOptions::default(),)"
+---
+Scenario: New context window tool installs fresh full context before the next follow-up request.
+
+## Final Follow-Up Request
+00:message/developer[3]:
+    [01] <PERMISSIONS_INSTRUCTIONS>
+    [02] <SKILLS_INSTRUCTIONS>
+    [03] <token_budget>\nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n</token_budget>
+01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
+02:function_call/update_plan
+03:function_call_output:Plan updated

codex-rs/core/tests/suite/token_budget.rs

@@ -4,10 +4,13 @@ use codex_model_provider_info::built_in_model_providers;
 use codex_protocol::protocol::EventMsg;
 use codex_protocol::protocol::Op;
 use core_test_support::PathBufExt;
+use core_test_support::context_snapshot;
+use core_test_support::context_snapshot::ContextSnapshotOptions;
 use core_test_support::responses::ResponsesRequest;
 use core_test_support::responses::ev_assistant_message;
 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::mount_sse_sequence;
 use core_test_support::responses::sse;
@@ -17,6 +20,8 @@ use core_test_support::test_codex::local;
 use core_test_support::test_codex::test_codex;
 use core_test_support::wait_for_event;
 use pretty_assertions::assert_eq;
+use serde_json::Value;
+use serde_json::json;
 
 const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000;
 const EFFECTIVE_CONTEXT_WINDOW: i64 = CONFIGURED_CONTEXT_WINDOW * 95 / 100;
@@ -29,6 +34,17 @@ fn token_budget_texts(request: &ResponsesRequest) -> Vec<String> {
         .collect()
 }
 
+fn tool_names(request: &ResponsesRequest) -> Vec<String> {
+    request
+        .body_json()
+        .get("tools")
+        .and_then(Value::as_array)
+        .into_iter()
+        .flatten()
+        .filter_map(|tool| tool.get("name").and_then(Value::as_str).map(str::to_string))
+        .collect()
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> {
     skip_if_no_network!(Ok(()));
@@ -220,3 +236,84 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = start_mock_server().await;
+    let call_id = "new-window-call";
+    let continue_call_id = "continue-call";
+    let continue_args = json!({
+        "plan": [
+            {"step": "Continue in the new context window", "status": "in_progress"}
+        ],
+    })
+    .to_string();
+    let responses = mount_sse_sequence(
+        &server,
+        vec![
+            sse(vec![
+                ev_response_created("resp-1"),
+                ev_function_call(call_id, "new_context", "{}"),
+                ev_completed("resp-1"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-2"),
+                ev_function_call(continue_call_id, "update_plan", &continue_args),
+                ev_completed("resp-2"),
+            ]),
+            sse(vec![
+                ev_response_created("resp-3"),
+                ev_assistant_message("msg-3", "done"),
+                ev_completed("resp-3"),
+            ]),
+        ],
+    )
+    .await;
+    let test = test_codex()
+        .with_config(|config| {
+            config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW);
+            config
+                .features
+                .enable(Feature::TokenBudget)
+                .expect("test config should allow token budget");
+        })
+        .build(&server)
+        .await?;
+
+    test.submit_turn("request new context window").await?;
+
+    let requests = responses.requests();
+    assert_eq!(requests.len(), 3);
+    assert!(
+        tool_names(&requests[0])
+            .iter()
+            .any(|name| name == "new_context"),
+        "new_context should be exposed when token budget is enabled"
+    );
+    assert_eq!(
+        token_budget_texts(&requests[2]),
+        vec![format!(
+            "<token_budget>\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n</token_budget>"
+        )]
+    );
+    assert!(
+        !requests[2].body_contains_text("request new context window"),
+        "new_context should drop the prior window history before continuing the turn"
+    );
+    assert_eq!(
+        requests[2].function_call_output_text(continue_call_id),
+        Some("Plan updated".to_string())
+    );
+    insta::assert_snapshot!(
+        "token_budget_new_context_window_tool_full_context",
+        context_snapshot::format_labeled_requests_snapshot(
+            "New context window tool installs fresh full context before the next follow-up request.",
+            &[("Final Follow-Up Request", &requests[2])],
+            &ContextSnapshotOptions::default(),
+        )
+    );
+
+    Ok(())
+}