agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

[codex] prototype mcp_history thread hint injection (#29259)

Browse Source 
## Why

Prototype whether the harness can invoke the `mcp_history` MCP while
constructing full initial context and expose its thread hint to the
model without requiring a model-issued tool call.

The prototype builds on the context-window lineage added by #29256 and
is now based directly on `main`.

## What changed

- Call `mcp_history/thread_hint` with no arguments while building the
full `<token_budget>` context.
- Pass the current `threadId` through MCP request metadata, matching the
normal MCP tool-call path.
- Serialize only the unstructured `content` result and append it inside
`<token_budget>` when the call succeeds.
- Omit the additional context when the MCP call or content serialization
fails.

## Prototype limitations

- The direct call bypasses the normal model-initiated MCP approval,
lifecycle-event, telemetry, and result-sanitization path.
- The call has no prototype-specific timeout, result-size cap, or
per-window cache.
- MCP latency is added to full-context construction, including
applicable compaction paths.

## Validation

- `just test -p codex-core token_budget`
This commit is contained in:
pakrym-oai
2026-06-20 18:02:02 -07:00
committed by GitHub
Unverified
parent d1209bddfc
commit b6d6be2a84
5 changed files with 171 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/context/token_budget_context.rs
+9 -5
View File
@@ -8,7 +8,7 @@ pub(crate) struct TokenBudgetContext {
first_window_id: Uuid,
previous_window_id: Option<Uuid>,
window_id: Uuid,
tokens_left: i64,
mcp_result: Option<String>,
}
impl TokenBudgetContext {
@@ -17,14 +17,14 @@ impl TokenBudgetContext {
first_window_id: Uuid,
previous_window_id: Option<Uuid>,
window_id: Uuid,
tokens_left: i64,
mcp_result: Option<String>,
) -> 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}"
)
}
}
codex-rs/core/src/session/mod.rs
+25 -2
View File
@@ -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::<Vec<_>>()
.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(),
);
codex-rs/core/tests/suite/snapshots/all__suite__token_budget__token_budget_new_context_window_tool_full_context.snap
+2 -2
View File
@@ -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] <PERMISSIONS_INSTRUCTIONS>
[02] <SKILLS_INSTRUCTIONS>
[03] <token_budget>\nThread id <THREAD_ID>.\nFirst context window id <FIRST_WINDOW_ID>.\nPrevious context window id <FIRST_WINDOW_ID>.\nCurrent context window id <WINDOW_ID>.\nYou have 121600 tokens left in this context window.\n</token_budget>
[03] <token_budget>\nThread id <THREAD_ID>.\nFirst context window id <FIRST_WINDOW_ID>.\nPrevious context window id <FIRST_WINDOW_ID>.\nCurrent context window id <WINDOW_ID>.\n</token_budget>
01:message/user:<ENVIRONMENT_CONTEXT:cwd=<CWD>>
02:function_call/update_plan
03:function_call_output:Plan updated
codex-rs/core/tests/suite/token_budget.rs
+100 -30
View File
@@ -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<String> {
request
@@ -38,11 +43,10 @@ fn token_budget_texts(request: &ResponsesRequest) -> Vec<String> {
fn token_budget_window_ids(
text: &str,
thread_id: codex_protocol::ThreadId,
tokens_left: i64,
) -> (String, Option<String>, String) {
let captures = assert_regex_match(
&format!(
r"^<token_budget>\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</token_budget>$"
r"^<token_budget>\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</token_budget>$"
),
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"^<token_budget>\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_budget>$"
),
&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 =
"<token_budget>\nYou have 7000 tokens left in this context window.\n</token_budget>"
@@ -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(),
codex-rs/rmcp-client/src/bin/test_stdio_server.rs
+35
View File
@@ -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(