From ace5d9ce391f61c90ed478f37ab51338e8ccaa45 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Sat, 20 Jun 2026 13:13:46 -0700 Subject: [PATCH 1/3] prototype mcp_history thread hint injection --- codex-rs/core/src/context/token_budget_context.rs | 10 +++++++++- codex-rs/core/src/session/mod.rs | 13 +++++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index bac1caafd08c..4530db627b9f 100644 --- a/codex-rs/core/src/context/token_budget_context.rs +++ b/codex-rs/core/src/context/token_budget_context.rs @@ -9,6 +9,7 @@ pub(crate) struct TokenBudgetContext { previous_window_id: Option, window_id: Uuid, tokens_left: i64, + mcp_result: Option, } impl TokenBudgetContext { @@ -18,6 +19,7 @@ impl TokenBudgetContext { previous_window_id: Option, window_id: Uuid, tokens_left: i64, + mcp_result: Option, ) -> Self { Self { thread_id, @@ -25,6 +27,7 @@ impl TokenBudgetContext { previous_window_id, window_id, tokens_left, + mcp_result, } } } @@ -50,8 +53,13 @@ impl ContextualUserFragment for TokenBudgetContext { .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!("\nHistory MCP result: {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}.\nYou have {tokens_left} tokens left in this context window.{mcp_result}" ) } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 3cc746912a73..a0148d5a2279 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3218,6 +3218,18 @@ impl Session { if turn_context.config.features.enabled(Feature::TokenBudget) && let Some(model_context_window) = turn_context.model_context_window() { + let mcp_result = self + .call_tool( + "mcp_history", + "thread_hint", + /*arguments*/ None, + Some(serde_json::json!({ + "threadId": self.thread_id().to_string(), + })), + ) + .await + .ok() + .and_then(|result| serde_json::to_string(&result.content).ok()); developer_sections.push( crate::context::TokenBudgetContext::new( self.thread_id(), @@ -3225,6 +3237,7 @@ impl Session { auto_compact_window_ids.previous_window_id, auto_compact_window_ids.window_id, model_context_window, + mcp_result, ) .render(), ); From b496dfeca3bf07c90fbd1310a9159cbf45213315 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Sat, 20 Jun 2026 13:56:11 -0700 Subject: [PATCH 2/3] inject plain mcp history text --- .../core/src/context/token_budget_context.rs | 8 +- codex-rs/core/src/session/mod.rs | 16 ++- ..._new_context_window_tool_full_context.snap | 4 +- codex-rs/core/tests/suite/token_budget.rs | 130 ++++++++++++++---- .../rmcp-client/src/bin/test_stdio_server.rs | 37 +++++ 5 files changed, 154 insertions(+), 41 deletions(-) diff --git a/codex-rs/core/src/context/token_budget_context.rs b/codex-rs/core/src/context/token_budget_context.rs index 4530db627b9f..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,6 @@ pub(crate) struct TokenBudgetContext { first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, - tokens_left: i64, mcp_result: Option, } @@ -18,7 +17,6 @@ impl TokenBudgetContext { first_window_id: Uuid, previous_window_id: Option, window_id: Uuid, - tokens_left: i64, mcp_result: Option, ) -> Self { Self { @@ -26,7 +24,6 @@ impl TokenBudgetContext { first_window_id, previous_window_id, window_id, - tokens_left, mcp_result, } } @@ -52,14 +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!("\nHistory MCP result: {result}")) + .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.{mcp_result}" + "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 a0148d5a2279..b99ea80dfc2c 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3216,7 +3216,7 @@ 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( @@ -3229,14 +3229,24 @@ impl Session { ) .await .ok() - .and_then(|result| serde_json::to_string(&result.content).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..86d40abccfdb 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( + "mcp_history".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, "mcp_history").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 mcp_history/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__mcp_history__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..9ffaff4f159e 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,24 @@ 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 mcp_history/thread_hint fixture result", + ), + ])) + } "echo" | "echo-tool" => { let args: EchoArgs = match request.arguments { Some(arguments) => serde_json::from_value(serde_json::Value::Object( From 8dc8374bc4759036597c8ac302d1a4dbc710a411 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Sat, 20 Jun 2026 14:07:30 -0700 Subject: [PATCH 3/3] call notes history MCP server --- codex-rs/core/src/session/mod.rs | 2 +- codex-rs/core/tests/suite/token_budget.rs | 8 ++++---- codex-rs/rmcp-client/src/bin/test_stdio_server.rs | 4 +--- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index b99ea80dfc2c..c85898555b3b 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3220,7 +3220,7 @@ impl Session { { let mcp_result = self .call_tool( - "mcp_history", + "notes", "thread_hint", /*arguments*/ None, Some(serde_json::json!({ diff --git a/codex-rs/core/tests/suite/token_budget.rs b/codex-rs/core/tests/suite/token_budget.rs index 86d40abccfdb..80f852a33ea0 100644 --- a/codex-rs/core/tests/suite/token_budget.rs +++ b/codex-rs/core/tests/suite/token_budget.rs @@ -144,7 +144,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { .expect("test config should allow token budget"); let mut servers = config.mcp_servers.get().clone(); servers.insert( - "mcp_history".to_string(), + "notes".to_string(), McpServerConfig { transport: McpServerTransportConfig::Stdio { command: rmcp_test_server_bin, @@ -176,7 +176,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { }) .build(&server) .await?; - wait_for_mcp_server(&test.codex, "mcp_history").await?; + wait_for_mcp_server(&test.codex, "notes").await?; let responses = mount_sse_sequence( &server, vec![sse(vec![ @@ -194,7 +194,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { 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 mcp_history/thread_hint fixture result\n$" + 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], ); @@ -205,7 +205,7 @@ async fn token_budget_context_injects_plain_thread_hint_text() -> Result<()> { assert!( !tool_names(&request) .iter() - .any(|name| name == "mcp__mcp_history__thread_hint"), + .any(|name| name == "mcp__notes__thread_hint"), "thread_hint should be hidden from model tool exposure" ); 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 9ffaff4f159e..722f5f7d7875 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -569,9 +569,7 @@ impl ServerHandler for TestToolServer { rmcp::model::Content::text(format!( "manual history hint for thread {thread_id}" )), - rmcp::model::Content::text( - "unstructured mcp_history/thread_hint fixture result", - ), + rmcp::model::Content::text("unstructured notes/thread_hint fixture result"), ])) } "echo" | "echo-tool" => {