diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index bac1caafd08c..24e1869e8aec 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -8,7 +8,7 @@ pub(crate) struct TokenBudgetContext { first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, - tokens_left: i64, + mcp_result: Option, } impl TokenBudgetContext { @@ -17,14 +17,14 @@ impl TokenBudgetContext { first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, - tokens_left: i64, + mcp_result: Option, ) -> Self { Self { thread_id, first_window_id, previous_window_id, window_id, - tokens_left, + mcp_result, } } } @@ -49,9 +49,13 @@ impl ContextualUserFragment for TokenBudgetContext { .previous_window_id .map_or_else(|| "none".to_string(), |window_id| window_id.to_string()); let window_id = self.window_id; - let tokens_left = self.tokens_left; + let mcp_result = self + .mcp_result + .as_deref() + .map(|result| format!("\n{result}")) + .unwrap_or_default(); format!( - "Thread id {thread_id}.\nFirst context window id {first_window_id}.\nPrevious context window id {previous_window_id}.\nCurrent context window id {window_id}.\nYou have {tokens_left} tokens left in this context window." + "Thread id {thread_id}.\nFirst context window id {first_window_id}.\nPrevious context window id {previous_window_id}.\nCurrent context window id {window_id}.{mcp_result}" ) } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 3cc746912a73..c85898555b3b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3216,15 +3216,38 @@ impl Session { } // This is full-context metadata. Steady-state context diffs should not re-emit it. if turn_context.config.features.enabled(Feature::TokenBudget) - && let Some(model_context_window) = turn_context.model_context_window() + && turn_context.model_context_window().is_some() { + let mcp_result = self + .call_tool( + "notes", + "thread_hint", + /*arguments*/ None, + Some(serde_json::json!({ + "threadId": self.thread_id().to_string(), + })), + ) + .await + .ok() + .and_then(|result| { + let text = result + .content + .iter() + .filter_map(|content| { + content.get("text").and_then(serde_json::Value::as_str) + }) + .filter(|text| !text.is_empty()) + .collect::>() + .join("\n"); + (!text.is_empty()).then_some(text) + }); developer_sections.push( crate::context::TokenBudgetContext::new( self.thread_id(), auto_compact_window_ids.first_window_id, auto_compact_window_ids.previous_window_id, auto_compact_window_ids.window_id, - model_context_window, + mcp_result, ) .render(), ); 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 3b933b4ed738..e5dd607170bb 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,6 +1,6 @@ --- source: core/tests/suite/token_budget.rs -assertion_line: 539 +assertion_line: 588 expression: snapshot --- Scenario: New context window tool installs fresh full context before the next follow-up request. @@ -9,7 +9,7 @@ Scenario: New context window tool installs fresh full context before the next fo 00:message/developer[3]: [01] [02] - [03] \nThread id .\nFirst context window id .\nPrevious context window id .\nCurrent context window id .\nYou have 121600 tokens left in this context window.\n + [03] \nThread id .\nFirst context window id .\nPrevious context window id .\nCurrent context window id .\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 7c4dfa1bedc3..80f852a33ea0 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -1,4 +1,6 @@ use anyhow::Result; +use codex_config::types::McpServerConfig; +use codex_config::types::McpServerTransportConfig; use codex_features::Feature; use codex_model_provider_info::built_in_model_providers; use codex_protocol::protocol::EventMsg; @@ -17,15 +19,18 @@ use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; +use core_test_support::stdio_server_bin; use core_test_support::test_codex::local; use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; +use core_test_support::wait_for_mcp_server; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use std::collections::HashMap; +use std::time::Duration; const CONFIGURED_CONTEXT_WINDOW: i64 = 128_000; -const EFFECTIVE_CONTEXT_WINDOW: i64 = CONFIGURED_CONTEXT_WINDOW * 95 / 100; fn token_budget_texts(request: &ResponsesRequest) -> Vec { request @@ -38,11 +43,10 @@ fn token_budget_texts(request: &ResponsesRequest) -> Vec { fn token_budget_window_ids( text: &str, thread_id: codex_protocol::ThreadId, - tokens_left: i64, ) -> (String, Option, String) { let captures = assert_regex_match( &format!( - r"^\nThread id {thread_id}\.\nFirst context window id ([0-9a-f-]{{36}})\.\nPrevious context window id (none|[0-9a-f-]{{36}})\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nYou have {tokens_left} tokens left in this context window\.\n$" + r"^\nThread id {thread_id}\.\nFirst context window id ([0-9a-f-]{{36}})\.\nPrevious context window id (none|[0-9a-f-]{{36}})\.\nCurrent context window id ([0-9a-f-]{{36}})\.\n$" ), text, ); @@ -112,11 +116,8 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_texts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); - let (first_window_id, previous_window_id, window_id) = token_budget_window_ids( - &initial_token_budget[0], - thread_id, - EFFECTIVE_CONTEXT_WINDOW, - ); + let (first_window_id, previous_window_id, window_id) = + token_budget_window_ids(&initial_token_budget[0], thread_id); assert_eq!(previous_window_id, None); assert_eq!(first_window_id, window_id); assert_eq!( @@ -128,6 +129,89 @@ async fn token_budget_context_is_only_emitted_with_full_context() -> Result<()> Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let rmcp_test_server_bin = stdio_server_bin()?; + let test = test_codex() + .with_config(move |config| { + config.model_context_window = Some(CONFIGURED_CONTEXT_WINDOW); + config + .features + .enable(Feature::TokenBudget) + .expect("test config should allow token budget"); + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + "notes".to_string(), + McpServerConfig { + transport: McpServerTransportConfig::Stdio { + command: rmcp_test_server_bin, + args: Vec::new(), + env: None, + env_vars: Vec::new(), + cwd: None, + }, + environment_id: "local".to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: Some(Duration::from_secs(10)), + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + }, + ); + config + .mcp_servers + .set(servers) + .expect("test mcp servers should accept any configuration"); + }) + .build(&server) + .await?; + wait_for_mcp_server(&test.codex, "notes").await?; + let responses = mount_sse_sequence( + &server, + vec![sse(vec![ + ev_response_created("resp-1"), + ev_completed("resp-1"), + ])], + ) + .await; + + test.submit_turn("inject the history hint").await?; + + let request = responses.single_request(); + let thread_id = test.session_configured.thread_id; + let token_budgets = token_budget_texts(&request); + assert_eq!(token_budgets.len(), 1); + let captures = assert_regex_match( + &format!( + r"^\nThread id {thread_id}\.\nFirst context window id ([0-9a-f-]{{36}})\.\nPrevious context window id none\.\nCurrent context window id ([0-9a-f-]{{36}})\.\nmanual history hint for thread {thread_id}\nunstructured notes/thread_hint fixture result\n$" + ), + &token_budgets[0], + ); + assert_eq!( + captures.get(1).expect("first window id capture").as_str(), + captures.get(2).expect("current window id capture").as_str() + ); + assert!( + !tool_names(&request) + .iter() + .any(|name| name == "mcp__notes__thread_hint"), + "thread_hint should be hidden from model tool exposure" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> Result<()> { skip_if_no_network!(Ok(())); @@ -177,7 +261,7 @@ async fn token_budget_remaining_context_emits_on_first_threshold_crossing() -> R let thread_id = test.session_configured.thread_id; let full_context = token_budget_texts(&requests[0]); assert_eq!(full_context.len(), 1); - token_budget_window_ids(&full_context[0], thread_id, /*tokens_left*/ 9_500); + token_budget_window_ids(&full_context[0], thread_id); let full_context = full_context[0].clone(); let threshold_25 = "\nYou have 7000 tokens left in this context window.\n" @@ -270,7 +354,7 @@ async fn get_context_remaining_returns_token_budget_remaining_fragment() -> Resu .to_string(); let token_budgets = token_budget_texts(&requests[1]); assert_eq!(token_budgets.len(), 2); - token_budget_window_ids(&token_budgets[0], thread_id, /*tokens_left*/ 9_500); + token_budget_window_ids(&token_budgets[0], thread_id); assert_eq!(token_budgets[1], remaining_context); assert_eq!( requests[2].function_call_output_content_and_success(call_id), @@ -394,22 +478,14 @@ async fn token_budget_context_uses_new_window_after_compaction() -> Result<()> { let initial_token_budget = token_budget_texts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); let (initial_first_window_id, initial_previous_window_id, initial_window_id) = - token_budget_window_ids( - &initial_token_budget[0], - thread_id, - EFFECTIVE_CONTEXT_WINDOW, - ); + token_budget_window_ids(&initial_token_budget[0], thread_id); let post_compaction_token_budget = token_budget_texts(&requests[2]); assert_eq!(post_compaction_token_budget.len(), 1); let ( post_compaction_first_window_id, post_compaction_previous_window_id, post_compaction_window_id, - ) = token_budget_window_ids( - &post_compaction_token_budget[0], - thread_id, - EFFECTIVE_CONTEXT_WINDOW, - ); + ) = token_budget_window_ids(&post_compaction_token_budget[0], thread_id); assert_eq!(initial_previous_window_id, None); assert_eq!(initial_first_window_id, initial_window_id); assert_eq!(post_compaction_first_window_id, initial_first_window_id); @@ -480,18 +556,12 @@ async fn new_context_tool_starts_new_window_before_follow_up() -> Result<()> { let thread_id = test.session_configured.thread_id; let initial_token_budget = token_budget_texts(&requests[0]); assert_eq!(initial_token_budget.len(), 1); - let (initial_first_window_id, _, initial_window_id) = token_budget_window_ids( - &initial_token_budget[0], - thread_id, - EFFECTIVE_CONTEXT_WINDOW, - ); + let (initial_first_window_id, _, initial_window_id) = + token_budget_window_ids(&initial_token_budget[0], thread_id); let new_window_token_budget = token_budget_texts(&requests[2]); assert_eq!(new_window_token_budget.len(), 1); - let (new_first_window_id, new_previous_window_id, new_window_id) = token_budget_window_ids( - &new_window_token_budget[0], - thread_id, - EFFECTIVE_CONTEXT_WINDOW, - ); + let (new_first_window_id, new_previous_window_id, new_window_id) = + token_budget_window_ids(&new_window_token_budget[0], thread_id); assert_eq!(new_first_window_id, initial_first_window_id); assert_eq!( new_previous_window_id.as_deref(), diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index 9ff36e349dea..722f5f7d7875 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -19,6 +19,7 @@ use rmcp::model::JsonObject; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::ListToolsResult; +use rmcp::model::Meta; use rmcp::model::PaginatedRequestParams; use rmcp::model::RawResource; use rmcp::model::RawResourceTemplate; @@ -70,9 +71,27 @@ impl TestToolServer { ); sandbox_meta_tool.annotations = Some(ToolAnnotations::new().read_only(true)); + #[expect(clippy::expect_used)] + let thread_hint_schema: JsonObject = serde_json::from_value(json!({ + "type": "object", + "properties": {}, + "additionalProperties": false + })) + .expect("thread_hint tool schema should deserialize"); + let mut thread_hint_tool = Tool::new( + Cow::Borrowed("thread_hint"), + Cow::Borrowed("Return an unstructured history hint for a thread."), + Arc::new(thread_hint_schema), + ); + thread_hint_tool.annotations = Some(ToolAnnotations::new().read_only(true)); + let mut thread_hint_meta = Meta::new(); + thread_hint_meta.insert("ui".to_string(), json!({ "visibility": [] })); + thread_hint_tool.meta = Some(thread_hint_meta); + let tools = vec![ Self::echo_tool(), Self::echo_dash_tool(), + thread_hint_tool, Self::client_capabilities_tool(), Self::cwd_tool(), Self::sync_tool(), @@ -537,6 +556,22 @@ impl ServerHandler for TestToolServer { .map_err(|err| McpError::internal_error(err.to_string(), None))?; Ok(Self::structured_result(json!({ "cwd": cwd }))) } + "thread_hint" => { + let thread_id = context + .meta + .0 + .get("threadId") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| { + McpError::invalid_params("missing threadId metadata".to_string(), None) + })?; + Ok(CallToolResult::success(vec![ + rmcp::model::Content::text(format!( + "manual history hint for thread {thread_id}" + )), + rmcp::model::Content::text("unstructured notes/thread_hint fixture result"), + ])) + } "echo" | "echo-tool" => { let args: EchoArgs = match request.arguments { Some(arguments) => serde_json::from_value(serde_json::Value::Object(