From f5420174dafba153913a3e697f89002c338dfd7e Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 26 Aug 2026 10:32:33 +0000 Subject: [PATCH] Include originating item IDs in MCP request metadata (#40866) ## What changed - Add the originating Responses item ID to MCP tool request metadata as `_meta.itemId`. - Preserve the initiating `exec` item ID across code-mode cells so nested MCP calls carry the same origin. ## Testing - Cover metadata generation for direct and code-mode MCP tool calls. GitOrigin-RevId: 1161cac64f1cc3a6d4dac4368e63804bd9d2c417 --- codex-rs/core/src/mcp_tool_call.rs | 44 +++++++++------- codex-rs/core/src/mcp_tool_call_tests.rs | 21 ++++++-- codex-rs/core/src/tools/code_mode/delegate.rs | 51 ++++++++++++++++--- .../src/tools/code_mode/execute_handler.rs | 13 ++++- codex-rs/core/src/tools/code_mode/mod.rs | 16 +++++- codex-rs/core/src/tools/context.rs | 30 +++++++++++ codex-rs/core/src/tools/handlers/mcp.rs | 2 + codex-rs/core/tests/suite/code_mode.rs | 28 ++++++---- .../core/tests/suite/mcp_turn_metadata.rs | 19 ++++--- 9 files changed, 175 insertions(+), 49 deletions(-) diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index 75bf23c7b3..d364836c98 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -39,6 +39,7 @@ use codex_mcp::ToolInfo; use codex_mcp::auth_elicitation_completed_result; use codex_mcp::build_auth_elicitation_plan; use codex_mcp::mcp_permission_prompt_is_auto_approved; +use codex_protocol::ResponseItemId; use codex_protocol::approvals::ElicitationRequest; use codex_protocol::items::McpToolCallError; use codex_protocol::items::McpToolCallItem; @@ -116,6 +117,7 @@ pub(crate) async fn handle_mcp_tool_call( step_context: &Arc, cancellation_token: &CancellationToken, call_id: String, + originating_item_id: Option, tool_info: &ToolInfo, hook_tool_name: HookToolName, invocation_tool_name: ToolName, @@ -259,6 +261,7 @@ pub(crate) async fn handle_mcp_tool_call( &sess, step_context.as_ref(), &call_id, + originating_item_id.as_ref(), invocation, prepared_call, metadata, @@ -332,6 +335,7 @@ pub(crate) async fn handle_mcp_tool_call( &sess, step_context.as_ref(), &call_id, + originating_item_id.as_ref(), invocation, prepared_call, metadata, @@ -394,6 +398,7 @@ async fn handle_approved_mcp_tool_call( sess: &Arc, step_context: &StepContext, call_id: &str, + originating_item_id: Option<&ResponseItemId>, invocation: McpInvocation, prepared_call: PreparedMcpCall, metadata: McpToolApprovalMetadata, @@ -471,9 +476,10 @@ async fn handle_approved_mcp_tool_call( call_id, Some(&metadata), ); - let request_meta = with_mcp_tool_call_thread_id_meta( + let request_meta = with_mcp_tool_call_ids_meta( request_meta, &sess.thread_id.to_string(), + originating_item_id, ); let request_meta = augment_mcp_tool_request_meta_with_sandbox_state( step_context, @@ -1125,6 +1131,7 @@ const MCP_TOOL_OPENAI_OUTPUT_TEMPLATE_META_KEY: &str = "openai/outputTemplate"; const MCP_TOOL_UI_RESOURCE_URI_META_KEY: &str = "ui/resourceUri"; const MCP_TOOL_LINK_ID_META_KEY: &str = "link_id"; const MCP_TOOL_PLUGIN_ID_META_KEY: &str = "plugin_id"; +const MCP_TOOL_ITEM_ID_META_KEY: &str = "itemId"; const MCP_TOOL_THREAD_ID_META_KEY: &str = "threadId"; const MCP_TOOL_CONNECTED_ACCOUNT_EMAIL_META_KEY: &str = "connected_account_email"; const MCP_TOOL_RESOURCE_URI_META_KEY: &str = "resource_uri"; @@ -1229,28 +1236,27 @@ fn build_mcp_tool_call_request_meta( (!request_meta.is_empty()).then_some(serde_json::Value::Object(request_meta)) } -fn with_mcp_tool_call_thread_id_meta( +fn with_mcp_tool_call_ids_meta( meta: Option, thread_id: &str, + originating_item_id: Option<&ResponseItemId>, ) -> Option { - match meta { - Some(serde_json::Value::Object(mut map)) => { - map.insert( - MCP_TOOL_THREAD_ID_META_KEY.to_string(), - serde_json::Value::String(thread_id.to_string()), - ); - Some(serde_json::Value::Object(map)) - } - None => { - let mut map = serde_json::Map::new(); - map.insert( - MCP_TOOL_THREAD_ID_META_KEY.to_string(), - serde_json::Value::String(thread_id.to_string()), - ); - Some(serde_json::Value::Object(map)) - } - other => other, + let mut map = match meta { + Some(serde_json::Value::Object(map)) => map, + None => serde_json::Map::new(), + other => return other, + }; + map.insert( + MCP_TOOL_THREAD_ID_META_KEY.to_string(), + serde_json::Value::String(thread_id.to_string()), + ); + if let Some(item_id) = originating_item_id { + map.insert( + MCP_TOOL_ITEM_ID_META_KEY.to_string(), + serde_json::Value::String(item_id.to_string()), + ); } + Some(serde_json::Value::Object(map)) } #[derive(Clone, Copy)] diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index 51fbbb4bab..1e4cbb5066 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1606,30 +1606,43 @@ async fn codex_apps_auth_elicitation_enabled_by_default_requests_elicitation() { } #[test] -fn mcp_tool_call_thread_id_meta_is_added_to_request_meta() { +fn mcp_tool_call_ids_are_added_to_request_meta() { + let item_id = ResponseItemId::from_server("fc-live".to_string()); + assert_eq!( - with_mcp_tool_call_thread_id_meta( + with_mcp_tool_call_ids_meta( Some(serde_json::json!({ "source": "test-client", "threadId": "stale-thread", + "itemId": "stale-item", })), "thread-live", + Some(&item_id), ), Some(serde_json::json!({ "source": "test-client", "threadId": "thread-live", + "itemId": "fc-live", })) ); assert_eq!( - with_mcp_tool_call_thread_id_meta(/*meta*/ None, "thread-live"), + with_mcp_tool_call_ids_meta( + /*meta*/ None, + "thread-live", + /*originating_item_id*/ None, + ), Some(serde_json::json!({ "threadId": "thread-live", })) ); assert_eq!( - with_mcp_tool_call_thread_id_meta(Some(serde_json::json!("invalid-meta")), "thread-live"), + with_mcp_tool_call_ids_meta( + Some(serde_json::json!("invalid-meta")), + "thread-live", + /*originating_item_id*/ None, + ), Some(serde_json::json!("invalid-meta")) ); } diff --git a/codex-rs/core/src/tools/code_mode/delegate.rs b/codex-rs/core/src/tools/code_mode/delegate.rs index bb46a06e74..bfa6a1c0bb 100644 --- a/codex-rs/core/src/tools/code_mode/delegate.rs +++ b/codex-rs/core/src/tools/code_mode/delegate.rs @@ -7,6 +7,7 @@ use codex_code_mode::CodeModeNestedToolCall; use codex_code_mode::CodeModeSessionDelegate; use codex_code_mode::NotificationFuture; use codex_code_mode::ToolInvocationFuture; +use codex_protocol::ResponseItemId; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseItem; use serde_json::Value as JsonValue; @@ -24,7 +25,13 @@ use crate::tools::parallel::ToolCallRuntime; pub(super) struct CodeModeDispatchBroker { dispatch_tx: async_channel::Sender, dispatch_rx: async_channel::Receiver, - dispatch_gates: Arc>>>, + dispatch_gates: Arc>>, +} + +struct CellDispatchGate { + ready: watch::Sender, + // Keep the original exec item when later waits resume this cell. + originating_item_id: Option, } impl CodeModeDispatchBroker { @@ -37,8 +44,34 @@ impl CodeModeDispatchBroker { } } - pub(super) fn mark_cell_ready_for_dispatch(&self, cell_id: &CellId) { - dispatch_gate(&self.dispatch_gates, cell_id).send_replace(true); + pub(super) fn mark_cell_ready_for_dispatch( + &self, + cell_id: &CellId, + originating_item_id: Option, + ) { + let ready = { + let mut dispatch_gates = self + .dispatch_gates + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let gate = dispatch_gates + .entry(cell_id.clone()) + .or_insert_with(|| CellDispatchGate { + ready: watch::channel(false).0, + originating_item_id: None, + }); + gate.originating_item_id = originating_item_id; + gate.ready.clone() + }; + ready.send_replace(true); + } + + pub(super) fn cell_originating_item_id(&self, cell_id: &CellId) -> Option { + self.dispatch_gates + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(cell_id) + .and_then(|gate| gate.originating_item_id.clone()) } pub(super) fn close_cell(&self, cell_id: &CellId) { @@ -135,7 +168,7 @@ impl CodeModeDispatchBroker { } fn dispatch_gate( - dispatch_gates: &Mutex>>, + dispatch_gates: &Mutex>, cell_id: &CellId, ) -> watch::Sender { let mut dispatch_gates = match dispatch_gates.lock() { @@ -144,12 +177,16 @@ fn dispatch_gate( }; dispatch_gates .entry(cell_id.clone()) - .or_insert_with(|| watch::channel(false).0) + .or_insert_with(|| CellDispatchGate { + ready: watch::channel(false).0, + originating_item_id: None, + }) + .ready .clone() } fn remove_dispatch_gate( - dispatch_gates: &Mutex>>, + dispatch_gates: &Mutex>, cell_id: &CellId, ) { let mut dispatch_gates = match dispatch_gates.lock() { @@ -160,7 +197,7 @@ fn remove_dispatch_gate( } async fn wait_until_cell_ready_for_dispatch( - dispatch_gates: &Mutex>>, + dispatch_gates: &Mutex>, cell_id: &CellId, cancellation_token: &CancellationToken, ) -> bool { diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 648fff7ecb..5b90528e63 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/codex-rs/core/src/tools/code_mode/execute_handler.rs @@ -35,6 +35,7 @@ impl CodeModeExecuteHandler { session: std::sync::Arc, turn: std::sync::Arc, call_id: String, + originating_item_id: Option, code: String, telemetry: &mut CodeModeToolCallGuard, ) -> Result { @@ -103,7 +104,7 @@ impl CodeModeExecuteHandler { exec.session .services .code_mode_service - .mark_cell_ready_for_dispatch(&cell_id); + .mark_cell_ready_for_dispatch(&cell_id, originating_item_id); let response = started_cell .initial_response() .await @@ -155,6 +156,7 @@ impl CodeModeExecuteHandler { &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { + let originating_item_id = invocation.originating_item_id().await; let ToolInvocation { session, turn, @@ -174,7 +176,14 @@ impl CodeModeExecuteHandler { ); let result = match payload { ToolPayload::Custom { input } if is_exec_tool_name(&tool_name) => self - .execute(session, turn, call_id, input, &mut telemetry) + .execute( + session, + turn, + call_id, + originating_item_id, + input, + &mut telemetry, + ) .await .map(boxed_tool_output), _ => Err(FunctionCallError::RespondToModel(format!( diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index c6a183a31a..26559df62e 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/codex-rs/core/src/tools/code_mode/mod.rs @@ -176,8 +176,20 @@ impl CodeModeService { } } - pub(crate) fn mark_cell_ready_for_dispatch(&self, cell_id: &codex_code_mode::CellId) { - self.dispatch_broker.mark_cell_ready_for_dispatch(cell_id); + pub(crate) fn mark_cell_ready_for_dispatch( + &self, + cell_id: &codex_code_mode::CellId, + originating_item_id: Option, + ) { + self.dispatch_broker + .mark_cell_ready_for_dispatch(cell_id, originating_item_id); + } + + pub(crate) fn cell_originating_item_id( + &self, + cell_id: &codex_code_mode::CellId, + ) -> Option { + self.dispatch_broker.cell_originating_item_id(cell_id) } pub(crate) fn finish_cell_dispatch(&self, cell_id: &CellId) { diff --git a/codex-rs/core/src/tools/context.rs b/codex-rs/core/src/tools/context.rs index dafb4a3dcc..378f947058 100644 --- a/codex-rs/core/src/tools/context.rs +++ b/codex-rs/core/src/tools/context.rs @@ -6,11 +6,13 @@ use crate::session::turn_context::TurnContext; use crate::turn_diff_tracker::TurnDiffTracker; use crate::unified_exec::format_output_omission_marker; use crate::unified_exec::resolve_max_tokens; +use codex_protocol::ResponseItemId; use codex_protocol::mcp::CallToolResult; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; +use codex_protocol::models::ResponseItem; use codex_protocol::models::function_call_output_content_items_to_text; use codex_tools::LoadableToolSpec; use codex_tools::ToolName; @@ -66,6 +68,34 @@ pub struct ToolInvocation { pub payload: ToolPayload, } +impl ToolInvocation { + /// Returns the Responses item that requested this call or started its code-mode cell. + pub(crate) async fn originating_item_id(&self) -> Option { + if let ToolCallSource::CodeMode { cell_id, .. } = &self.source { + return self + .session + .services + .code_mode_service + .cell_originating_item_id(&codex_code_mode::CellId::new(cell_id.clone())); + } + + self.session + .clone_history() + .await + .raw_items() + .rev() + .find_map(|item| match item { + ResponseItem::FunctionCall { id, call_id, .. } + | ResponseItem::CustomToolCall { id, call_id, .. } + if call_id == &self.call_id => + { + id.clone() + } + _ => None, + }) + } +} + #[derive(Clone, Debug)] pub struct McpToolOutput { pub result: CallToolResult, diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index b62945b1f9..e1a96c6586 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -171,6 +171,7 @@ impl McpHandler { &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { + let originating_item_id = invocation.originating_item_id().await; let ToolInvocation { session, step_context, @@ -197,6 +198,7 @@ impl McpHandler { &step_context, &cancellation_token, call_id.clone(), + originating_item_id, &self.tool_info, self.hook_tool_name(), tool_name, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 29be9a8520..3f0b41fe81 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1209,14 +1209,11 @@ async fn code_mode_only_guides_all_tools_search_and_calls_deferred_app_tools() - let server = responses::start_mock_server().await; let apps_server = AppsTestServer::mount_searchable(&server).await?; - let resp_mock = responses::mount_sse_once( - &server, - sse(vec![ - ev_response_created("resp-1"), - ev_custom_tool_call( - "call-1", - "exec", - r#" + let originating_item_id = "ctc_code_mode_origin"; + let mut exec_call = ev_custom_tool_call( + "call-1", + "exec", + r#" const tool = ALL_TOOLS.find( ({ name }) => name === "mcp__codex_apps__calendar_timezone_option_99" ); @@ -1231,7 +1228,13 @@ if (!tool) { })); } "#, - ), + ); + exec_call["item"]["id"] = serde_json::json!(originating_item_id); + let resp_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + exec_call, ev_completed("resp-1"), ]), ) @@ -1325,6 +1328,13 @@ if (!tool) { "text": "called calendar_timezone_option_99 for at with ", }) ); + let apps_tool_calls = recorded_apps_tool_calls(&server).await; + assert!( + apps_tool_calls.iter().any(|call| { + call.pointer("/params/_meta/itemId") == Some(&serde_json::json!(originating_item_id)) + }), + "the nested MCP call should inherit its code cell's originating Responses item" + ); Ok(()) } diff --git a/codex-rs/core/tests/suite/mcp_turn_metadata.rs b/codex-rs/core/tests/suite/mcp_turn_metadata.rs index 57ef7cc3a5..d1dab757b5 100644 --- a/codex-rs/core/tests/suite/mcp_turn_metadata.rs +++ b/codex-rs/core/tests/suite/mcp_turn_metadata.rs @@ -170,6 +170,7 @@ async fn approved_mcp_tool_call_metadata_records_prior_user_input_request( let server = start_mock_server().await; let apps_server = AppsTestServer::mount(&server).await?; let call_id = "calendar-call-approval"; + let originating_item_id = "fc_calendar_approval_origin"; let calendar_args = serde_json::to_string(&json!({ "title": "Lunch", "starts_at": "2026-03-10T12:00:00Z" @@ -196,14 +197,16 @@ async fn approved_mcp_tool_call_metadata_records_prior_user_input_request( ev_completed("resp-permissions"), ])); } + let mut calendar_call = ev_function_call_with_namespace( + call_id, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + &calendar_args, + ); + calendar_call["item"]["id"] = json!(originating_item_id); response_sequence.push(sse(vec![ ev_response_created("resp-1"), - ev_function_call_with_namespace( - call_id, - SEARCH_CALENDAR_NAMESPACE, - SEARCH_CALENDAR_CREATE_TOOL, - &calendar_args, - ), + calendar_call, ev_completed("resp-1"), ])); if strict_auto_review { @@ -351,6 +354,10 @@ async fn approved_mcp_tool_call_metadata_records_prior_user_input_request( apps_tool_call.pointer("/params/_meta/callId"), Some(&json!(call_id)) ); + assert_eq!( + apps_tool_call.pointer("/params/_meta/itemId"), + Some(&json!(originating_item_id)) + ); assert_eq!( apps_tool_call .pointer("/params/_meta/x-codex-turn-metadata/user_input_requested_during_turn"),