From 87ab01834af30cfae99014ca93035d3068716b3f Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Wed, 10 Jun 2026 20:39:07 -0700 Subject: [PATCH] [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 | 38 ++++++++ codex-rs/core/src/session/turn.rs | 9 ++ .../core/src/state/auto_compact_window.rs | 18 ++++ codex-rs/core/src/state/session.rs | 14 +++ codex-rs/core/src/tools/handlers/mod.rs | 3 + .../src/tools/handlers/new_context_window.rs | 45 +++++++++ .../tools/handlers/new_context_window_spec.rs | 17 ++++ codex-rs/core/src/tools/spec_plan.rs | 5 + ..._new_context_window_tool_full_context.snap | 15 +++ codex-rs/core/tests/suite/token_budget.rs | 97 +++++++++++++++++++ 10 files changed, 261 insertions(+) create mode 100644 codex-rs/core/src/tools/handlers/new_context_window.rs create mode 100644 codex-rs/core/src/tools/handlers/new_context_window_spec.rs create mode 100644 codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 0c9dfeccb..9f83ff5d0 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/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 { + 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 { let state = self.state.lock().await; state.reference_context_item() diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 85200d740..a694ed480 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/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( diff --git a/codex-rs/core/src/state/auto_compact_window.rs b/codex-rs/core/src/state/auto_compact_window.rs index f6ebd325f..1e3eceada 100644 --- a/codex-rs/core/src/state/auto_compact_window.rs +++ b/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(), diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index ced1428a3..269d3e0f6 100644 --- a/codex-rs/core/src/state/session.rs +++ b/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 { + 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 { self.history.token_info() } diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index fdfb80996..d36020f79 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/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; diff --git a/codex-rs/core/src/tools/handlers/new_context_window.rs b/codex-rs/core/src/tools/handlers/new_context_window.rs new file mode 100644 index 000000000..e1cd65639 --- /dev/null +++ b/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 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 {} diff --git a/codex-rs/core/src/tools/handlers/new_context_window_spec.rs b/codex-rs/core/src/tools/handlers/new_context_window_spec.rs new file mode 100644 index 000000000..f58b66e58 --- /dev/null +++ b/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, + }) +} diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 697eba1d4..e4aaa738f 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/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()) diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap b/codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap new file mode 100644 index 000000000..d042616af --- /dev/null +++ b/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] + [02] + [03] \nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n +01:message/user:> +02:function_call/update_plan +03:function_call_output:Plan updated diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 92c2fa929..7287ce39e 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/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 { .collect() } +fn tool_names(request: &ResponsesRequest) -> Vec { + 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!( + "\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" + )] + ); + 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(()) +}