From dac5f07403e048544cad6e9f9fbd52b91cb47a15 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Wed, 10 Jun 2026 21:17:10 -0700 Subject: [PATCH] [codex] Add context remaining tool (#27518) ## Why The token budget feature can inject remaining-context notices into model-visible context, but the model does not have a direct way to ask for that same remaining-token fragment on demand. This PR adds a small model tool for the token budget feature so the model can request the current remaining context window message without duplicating the fragment format. ## What changed - Adds a `get_context_remaining` direct-model tool behind `Feature::TokenBudget`. - Renders the tool output through `TokenBudgetRemainingContext`, matching the existing budget message shape. - Registers the tool alongside `new_context` in the token budget tool set. - Adds integration coverage that verifies the tool is exposed and returns the same `` remaining fragment already present in context. ## Validation - `just test -p codex-core token_budget` --- .../core/src/context/token_budget_context.rs | 18 ++- .../tools/handlers/get_context_remaining.rs | 54 ++++++++ .../handlers/get_context_remaining_spec.rs | 17 +++ codex-rs/core/src/tools/handlers/mod.rs | 3 + codex-rs/core/src/tools/spec_plan.rs | 2 + codex-rs/core/tests/suite/token_budget.rs | 130 ++++++++++++++++++ 6 files changed, 220 insertions(+), 4 deletions(-) create mode 100644 codex-rs/core/src/tools/handlers/get_context_remaining.rs create mode 100644 codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index d13c40448..df32ae4d7 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -39,12 +39,18 @@ impl ContextualUserFragment for TokenBudgetContext { #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetRemainingContext { - tokens_left: i64, + tokens_left: Option, } impl TokenBudgetRemainingContext { pub(crate) fn new(tokens_left: i64) -> Self { - Self { tokens_left } + Self { + tokens_left: Some(tokens_left), + } + } + + pub(crate) fn unknown() -> Self { + Self { tokens_left: None } } } @@ -62,7 +68,11 @@ impl ContextualUserFragment for TokenBudgetRemainingContext { } fn body(&self) -> String { - let tokens_left = self.tokens_left; - format!("You have {tokens_left} tokens left in this context window.") + match self.tokens_left { + Some(tokens_left) => { + format!("You have {tokens_left} tokens left in this context window.") + } + None => "You have unknown tokens left in this context window.".to_string(), + } } } diff --git a/codex-rs/core/src/tools/handlers/get_context_remaining.rs b/codex-rs/core/src/tools/handlers/get_context_remaining.rs new file mode 100644 index 000000000..a49ab6539 --- /dev/null +++ b/codex-rs/core/src/tools/handlers/get_context_remaining.rs @@ -0,0 +1,54 @@ +use crate::context::ContextualUserFragment; +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::get_context_remaining_spec::GET_CONTEXT_REMAINING_TOOL_NAME; +use crate::tools::handlers::get_context_remaining_spec::create_get_context_remaining_tool; +use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::ToolExecutor; +use codex_tools::ToolName; +use codex_tools::ToolSpec; + +pub struct GetContextRemainingHandler; + +impl ToolExecutor for GetContextRemainingHandler { + fn tool_name(&self) -> ToolName { + ToolName::plain(GET_CONTEXT_REMAINING_TOOL_NAME) + } + + fn spec(&self) -> ToolSpec { + create_get_context_remaining_tool() + } + + fn handle(&self, invocation: ToolInvocation) -> codex_tools::ToolExecutorFuture<'_> { + Box::pin(async move { + if !matches!(invocation.payload, ToolPayload::Function { .. }) { + return Err(FunctionCallError::RespondToModel( + "get_context_remaining handler received unsupported payload".to_string(), + )); + } + + let Some(model_context_window) = invocation.turn.model_context_window() else { + let fragment = crate::context::TokenBudgetRemainingContext::unknown().render(); + return Ok(boxed_tool_output(FunctionToolOutput::from_text( + fragment, + Some(true), + ))); + }; + let active_context_tokens = invocation.session.get_total_token_usage().await.max(0); + let tokens_left = model_context_window + .saturating_sub(active_context_tokens) + .max(0); + let fragment = crate::context::TokenBudgetRemainingContext::new(tokens_left).render(); + + Ok(boxed_tool_output(FunctionToolOutput::from_text( + fragment, + Some(true), + ))) + }) + } +} + +impl CoreToolRuntime for GetContextRemainingHandler {} diff --git a/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs b/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs new file mode 100644 index 000000000..988f2492f --- /dev/null +++ b/codex-rs/core/src/tools/handlers/get_context_remaining_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 GET_CONTEXT_REMAINING_TOOL_NAME: &str = "get_context_remaining"; + +pub fn create_get_context_remaining_tool() -> ToolSpec { + ToolSpec::Function(ResponsesApiTool { + name: GET_CONTEXT_REMAINING_TOOL_NAME.to_string(), + description: "Get the remaining tokens in the current 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/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index d36020f79..258c9914d 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -4,6 +4,8 @@ pub(crate) mod apply_patch; pub(crate) mod apply_patch_spec; mod dynamic; pub(crate) mod extension_tools; +mod get_context_remaining; +pub(crate) mod get_context_remaining_spec; mod list_available_plugins_to_install; pub(crate) mod list_available_plugins_to_install_spec; mod mcp; @@ -53,6 +55,7 @@ pub use apply_patch::ApplyPatchHandler; use codex_protocol::models::AdditionalPermissionProfile; use codex_protocol::protocol::AskForApproval; pub use dynamic::DynamicToolHandler; +pub use get_context_remaining::GetContextRemainingHandler; pub use list_available_plugins_to_install::ListAvailablePluginsToInstallHandler; pub use mcp::McpHandler; pub use mcp_resource::ListMcpResourceTemplatesHandler; diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index e4aaa738f..48ddcf7f0 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -9,6 +9,7 @@ use crate::tools::handlers::CodeModeWaitHandler; use crate::tools::handlers::DynamicToolHandler; use crate::tools::handlers::ExecCommandHandler; use crate::tools::handlers::ExecCommandHandlerOptions; +use crate::tools::handlers::GetContextRemainingHandler; use crate::tools::handlers::ListAvailablePluginsToInstallHandler; use crate::tools::handlers::ListMcpResourceTemplatesHandler; use crate::tools::handlers::ListMcpResourcesHandler; @@ -662,6 +663,7 @@ fn add_core_utility_tools(context: &CoreToolPlanContext<'_>, planned_tools: &mut if features.enabled(Feature::TokenBudget) { planned_tools.add_with_exposure(NewContextWindowHandler, ToolExposure::DirectModelOnly); + planned_tools.add_with_exposure(GetContextRemainingHandler, ToolExposure::DirectModelOnly); } if tool_suggest_enabled(turn_context) diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 7287ce39e..f50ec0989 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -179,6 +179,136 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "remaining-call"; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "noted"), + ev_completed_with_tokens("resp-1", /*total_tokens*/ 2_500), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_function_call(call_id, "get_context_remaining", "{}"), + ev_completed_with_tokens("resp-2", /*total_tokens*/ 2_500), + ]), + 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(10_000); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("spend some tokens").await?; + test.submit_turn("check remaining context").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 3); + assert!( + tool_names(&requests[1]) + .iter() + .any(|name| name == "get_context_remaining"), + "get_context_remaining should be exposed when token budget is enabled" + ); + + let full_context = "\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" + .to_string(); + let remaining_context = + "\nYou have 7000 tokens left in this context window.\n" + .to_string(); + assert_eq!( + token_budget_texts(&requests[1]), + vec![full_context, remaining_context.clone()] + ); + assert_eq!( + requests[2].function_call_output_content_and_success(call_id), + Some((Some(remaining_context), None)) + ); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_context_remaining_returns_unknown_when_window_is_unavailable() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let call_id = "remaining-call"; + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_function_call(call_id, "get_context_remaining", "{}"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "done"), + ev_completed("resp-2"), + ]), + ], + ) + .await; + let test = test_codex() + .with_model_info_override("gpt-5.2", |model_info| { + model_info.context_window = None; + model_info.max_context_window = None; + }) + .with_config(|config| { + config.model_context_window = None; + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }) + .build(&server) + .await?; + + test.submit_turn("check remaining context").await?; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2); + assert!( + tool_names(&requests[0]) + .iter() + .any(|name| name == "get_context_remaining"), + "get_context_remaining should be exposed when token budget is enabled" + ); + + assert_eq!(token_budget_texts(&requests[0]), Vec::::new()); + assert_eq!( + requests[1].function_call_output_content_and_success(call_id), + Some(( + Some( + "\nYou have unknown tokens left in this context window.\n" + .to_string() + ), + None, + )) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { skip_if_no_network!(Ok(()));