From 7516eb5c70f8478e20955b4203c5c892ef641055 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Thu, 11 Jun 2026 15:10:29 -0700 Subject: [PATCH] Include thread id in token budget context (#27663) ## Why The token budget full-context fragment identifies the current context window, but not the thread that owns that window. Including the thread id makes the initial context-window metadata self-contained, and `get_context_remaining` also needs to be usable from Code Mode without forcing callers to parse the model-facing fragment string. ## What changed - Include the session thread id in the initial `` context fragment. - Expose `get_context_remaining` as a Code Mode nested tool while keeping `new_context` direct-model-only. - Keep direct model-facing `get_context_remaining` output as the existing `` text fragment. - Return only `tokens_left` from the Code Mode structured result for `get_context_remaining`. - Update token-budget integration tests and add Code Mode coverage for the structured result. ## Verification - `just test -p codex-core token_budget` - `just test -p codex-core code_mode_get_context_remaining_returns_structured_result` - `just test -p core_test_support redacted_text_mode_normalizes_uuids` --- .../core/src/context/token_budget_context.rs | 8 ++- codex-rs/core/src/session/mod.rs | 1 + .../tools/handlers/get_context_remaining.rs | 59 ++++++++++++++++--- .../handlers/get_context_remaining_spec.rs | 21 ++++++- codex-rs/core/src/tools/spec_plan.rs | 2 +- codex-rs/core/tests/suite/code_mode.rs | 41 +++++++++++++ ..._new_context_window_tool_full_context.snap | 5 +- codex-rs/core/tests/suite/token_budget.rs | 33 +++++++---- 8 files changed, 142 insertions(+), 28 deletions(-) diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index df32ae4d7..37b355d51 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -1,14 +1,17 @@ use super::ContextualUserFragment; +use codex_protocol::ThreadId; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct TokenBudgetContext { + thread_id: ThreadId, window_id: u64, tokens_left: i64, } impl TokenBudgetContext { - pub(crate) fn new(window_id: u64, tokens_left: i64) -> Self { + pub(crate) fn new(thread_id: ThreadId, window_id: u64, tokens_left: i64) -> Self { Self { + thread_id, window_id, tokens_left, } @@ -29,10 +32,11 @@ impl ContextualUserFragment for TokenBudgetContext { } fn body(&self) -> String { + let thread_id = self.thread_id; 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." + "Thread id {thread_id}.\nCurrent context window {window_id}.\nYou have {tokens_left} tokens left in this context window." ) } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index e6a4d03dd..68c05e2c4 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2991,6 +2991,7 @@ impl Session { { developer_sections.push( crate::context::TokenBudgetContext::new( + self.thread_id(), auto_compact_window_id, model_context_window, ) diff --git a/codex-rs/core/src/tools/handlers/get_context_remaining.rs b/codex-rs/core/src/tools/handlers/get_context_remaining.rs index a49ab6539..b9a3d73ca 100644 --- a/codex-rs/core/src/tools/handlers/get_context_remaining.rs +++ b/codex-rs/core/src/tools/handlers/get_context_remaining.rs @@ -2,14 +2,59 @@ use crate::context::ContextualUserFragment; use crate::function_tool::FunctionCallError; use crate::tools::context::FunctionToolOutput; use crate::tools::context::ToolInvocation; +use crate::tools::context::ToolOutput; 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_protocol::models::ResponseInputItem; use codex_tools::ToolName; use codex_tools::ToolSpec; +use serde_json::Value as JsonValue; +use serde_json::json; + +#[derive(Debug, Clone)] +struct GetContextRemainingOutput { + tokens_left: Option, +} + +impl GetContextRemainingOutput { + fn new(tokens_left: Option) -> Self { + Self { tokens_left } + } + + fn fragment(&self) -> String { + match self.tokens_left { + Some(tokens_left) => { + crate::context::TokenBudgetRemainingContext::new(tokens_left).render() + } + None => crate::context::TokenBudgetRemainingContext::unknown().render(), + } + } +} + +impl ToolOutput for GetContextRemainingOutput { + fn log_preview(&self) -> String { + self.fragment() + } + + fn success_for_logging(&self) -> bool { + true + } + + fn to_response_item(&self, call_id: &str, payload: &ToolPayload) -> ResponseInputItem { + FunctionToolOutput::from_text(self.fragment(), Some(true)) + .to_response_item(call_id, payload) + } + + fn code_mode_result(&self, _payload: &ToolPayload) -> JsonValue { + json!({ + "tokens_left": self.tokens_left, + }) + } +} pub struct GetContextRemainingHandler; @@ -31,22 +76,18 @@ impl ToolExecutor for GetContextRemainingHandler { } 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), + return Ok(boxed_tool_output(GetContextRemainingOutput::new( + /*tokens_left*/ None, ))); }; 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), - ))) + Ok(boxed_tool_output(GetContextRemainingOutput::new(Some( + tokens_left, + )))) }) } } 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 index 988f2492f..4ff54243f 100644 --- a/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs +++ b/codex-rs/core/src/tools/handlers/get_context_remaining_spec.rs @@ -1,6 +1,8 @@ use codex_tools::JsonSchema; use codex_tools::ResponsesApiTool; use codex_tools::ToolSpec; +use serde_json::Value; +use serde_json::json; use std::collections::BTreeMap; pub(crate) const GET_CONTEXT_REMAINING_TOOL_NAME: &str = "get_context_remaining"; @@ -12,6 +14,23 @@ pub fn create_get_context_remaining_tool() -> ToolSpec { strict: false, defer_loading: None, parameters: JsonSchema::object(BTreeMap::new(), /*required*/ None, Some(false.into())), - output_schema: None, + output_schema: Some(get_context_remaining_output_schema()), + }) +} + +fn get_context_remaining_output_schema() -> Value { + json!({ + "type": "object", + "properties": { + "tokens_left": { + "anyOf": [ + { "type": "integer" }, + { "type": "null" } + ], + "description": "Remaining tokens in the current context window, or null when unavailable." + } + }, + "required": ["tokens_left"], + "additionalProperties": false }) } diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 51a0da7ad..79c86b94d 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -654,7 +654,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); + planned_tools.add(GetContextRemainingHandler); } if tool_suggest_enabled(turn_context) diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index fdd87274a..0f954bac7 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -801,6 +801,47 @@ text(JSON.stringify(result)); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_get_context_remaining_returns_structured_result() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn_with_config( + &server, + "use exec to get remaining context", + r#" +const result = await tools.get_context_remaining({}); +text(JSON.stringify(result)); +"#, + |config| { + config.model_context_window = Some(10_000); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + }, + ) + .await?; + + let req = second_mock.single_request(); + let (output, success) = custom_tool_output_body_and_success(&req, "call-1"); + assert_ne!( + success, + Some(false), + "exec get_context_remaining call failed unexpectedly: {output}" + ); + + let parsed: Value = serde_json::from_str(&output)?; + assert_eq!( + parsed, + serde_json::json!({ + "tokens_left": 9500, + }) + ); + + Ok(()) +} + #[cfg_attr(windows, ignore = "flaky on windows")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_nested_tool_calls_can_run_in_parallel() -> Result<()> { 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 index d042616af..e3dddb45c 100644 --- 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 @@ -1,7 +1,6 @@ --- 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(),)" +expression: snapshot --- Scenario: New context window tool installs fresh full context before the next follow-up request. @@ -9,7 +8,7 @@ Scenario: New context window tool installs fresh full context before the next fo 00:message/developer[3]: [01] [02] - [03] \nCurrent context window 1.\nYou have 121600 tokens left in this context window.\n + [03] \nThread id .\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 f50ec0989..c64251e24 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -79,8 +79,9 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> let requests = responses.requests(); assert_eq!(requests.len(), 2); + let thread_id = test.session_configured.thread_id; let expected = vec![format!( - "\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" + "\nThread id {thread_id}.\nCurrent context window 0.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" )]; assert_eq!( token_budget_texts(&requests[0]), @@ -142,8 +143,10 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R let requests = responses.requests(); assert_eq!(requests.len(), 5); - let full_context = "\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" - .to_string(); + let thread_id = test.session_configured.thread_id; + let full_context = format!( + "\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" + ); let threshold_25 = "\nYou have 7000 tokens left in this context window.\n" .to_string(); @@ -229,8 +232,10 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu "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 thread_id = test.session_configured.thread_id; + let full_context = format!( + "\nThread id {thread_id}.\nCurrent context window 0.\nYou have 9500 tokens left in this context window.\n" + ); let remaining_context = "\nYou have 7000 tokens left in this context window.\n" .to_string(); @@ -356,10 +361,11 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { let requests = responses.requests(); assert_eq!(requests.len(), 3); + let thread_id = test.session_configured.thread_id; 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" + "\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" )], "post-compaction full context should report context window 1" ); @@ -422,10 +428,11 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { .any(|name| name == "new_context"), "new_context should be exposed when token budget is enabled" ); + let thread_id = test.session_configured.thread_id; 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" + "\nThread id {thread_id}.\nCurrent context window 1.\nYou have {EFFECTIVE_CONTEXT_WINDOW} tokens left in this context window.\n" )] ); assert!( @@ -436,13 +443,15 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { requests[2].function_call_output_text(continue_call_id), Some("Plan updated".to_string()) ); + let snapshot = 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(), + ); + let snapshot = snapshot.replace(&thread_id.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(), - ) + snapshot ); Ok(())