From 8fcabd75a033ef281561a80e8595b2fc24d3743a Mon Sep 17 00:00:00 2001 From: fahad-OAI Date: Fri, 6 Nov 2026 21:24:06 +0000 Subject: [PATCH] Preserve app invocation context across MCP tool refreshes (#59991) ## What changed Keep the advertised `_codex_apps.invocation_context` when an app's tool metadata refreshes on the same server connection, while using refreshed values for other metadata. If the invocation context changes on a different connection, skip the call as unavailable. ## Testing Extend the code-mode compaction test to refresh app tools before resuming a cell and verify that calls retain the original invocation context while picking up the refreshed `resource_uri`. GitOrigin-RevId: 1b6072faed74623994ea82035a6ed977451984a4 --- codex-rs/codex-mcp/src/binding.rs | 6 ++++ codex-rs/codex-mcp/src/connection_manager.rs | 7 ++++ .../codex-mcp/src/connection_manager_tests.rs | 8 +---- codex-rs/core/src/mcp_tool_call.rs | 35 ++++++++++++++++--- codex-rs/core/tests/suite/code_mode.rs | 25 ++++++++++--- 5 files changed, 65 insertions(+), 16 deletions(-) diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index 36809350f2..3135437b68 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -284,6 +284,12 @@ impl PreparedMcpCall { self.tool_info.get() } + /// Whether the server connection that issued the step's tools still owns this call. + pub fn uses_connection_from(&self, binding: &McpBinding) -> bool { + self.connections + .shares_server_connection(&binding.connections, &self.server_name) + } + /// Returns the configuration and approval authority captured with this client. pub fn config(&self) -> &McpConfig { &self.config diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 1bf5651986..2efa07917f 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -1195,6 +1195,13 @@ impl McpConnectionSet { self.servers.contains_key(server_name) } + pub(crate) fn shares_server_connection(&self, other: &Self, server: &str) -> bool { + self.servers + .get(server) + .zip(other.servers.get(server)) + .is_some_and(|(left, right)| Arc::ptr_eq(&left.connection, &right.connection)) + } + pub(crate) async fn authentication_failed_servers(&self) -> Vec { let mut failed_servers = Vec::new(); for (server_name, view) in &self.servers { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 5e4225a417..acfd39a7ba 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -181,13 +181,7 @@ impl McpConnectionSet { } fn shares_test_connection_with(&self, other: &Self, name: &str) -> bool { - let Some(left) = self.servers.get(name) else { - return false; - }; - let Some(right) = other.servers.get(name) else { - return false; - }; - Arc::ptr_eq(&left.connection, &right.connection) + self.shares_server_connection(other, name) } } diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index c7b0682fe4..4832f72068 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -133,6 +133,15 @@ const MCP_RESULT_TELEMETRY_SERVER_USER_FLOW_SPAN_ATTR: &str = "codex.mcp.server_user_flow.triggered"; const MCP_RESULT_TELEMETRY_TARGET_ID_MAX_CHARS: usize = 256; const MCP_TOOL_CALL_EVENT_RESULT_MAX_BYTES: usize = DEFAULT_OUTPUT_BYTES_CAP; +const APPS_INVOCATION_CONTEXT: &str = "invocation_context"; + +fn apps_invocation_context(tool: &ToolInfo) -> Option<&JsonValue> { + tool.tool + .meta + .as_ref()? + .get(MCP_TOOL_CODEX_APPS_META_KEY)? + .get(APPS_INVOCATION_CONTEXT) +} fn parse_mcp_tool_arguments(arguments: &str) -> serde_json::Result> { if arguments.trim().is_empty() { @@ -151,7 +160,7 @@ pub(crate) async fn handle_mcp_tool_call( cancellation_token: &CancellationToken, call_id: String, originating_call: Option, - tool_info: &ToolInfo, + advertised_tool: &ToolInfo, prepared_call: Option, hook_tool_name: HookToolName, invocation_tool_name: ToolName, @@ -159,8 +168,8 @@ pub(crate) async fn handle_mcp_tool_call( arguments: String, ) -> HandledMcpToolCall { let turn_context = &step_context.turn; - let server = tool_info.server_name.clone(); - let tool_name = tool_info.tool.name.to_string(); + let server = advertised_tool.server_name.clone(); + let tool_name = advertised_tool.tool.name.to_string(); let arguments_value = match parse_mcp_tool_arguments(&arguments) { Ok(value) => value, Err(e) => { @@ -178,7 +187,15 @@ pub(crate) async fn handle_mcp_tool_call( arguments: arguments_value.clone(), }; - let Some(prepared_call) = prepared_call else { + let issued_context = apps_invocation_context(advertised_tool); + let context_changed = server == CODEX_APPS_MCP_SERVER_NAME + && prepared_call + .as_ref() + .is_some_and(|call| apps_invocation_context(call.tool_info()) != issued_context); + // Only the connection that advertised the context may receive it after a catalog refresh. + let Some(prepared_call) = prepared_call + .filter(|call| !context_changed || call.uses_connection_from(&step_context.mcp)) + else { let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, /*metadata*/ None); let result = notify_mcp_tool_call_skip( @@ -197,7 +214,7 @@ pub(crate) async fn handle_mcp_tool_call( .unwrap_or_else(|| JsonValue::Object(serde_json::Map::new())), }; }; - let metadata = match mcp_tool_metadata( + let mut metadata = match mcp_tool_metadata( prepared_call.tool_info(), prepared_call.plugin_id(), invocation.arguments.as_ref(), @@ -223,6 +240,14 @@ pub(crate) async fn handle_mcp_tool_call( }; } }; + if context_changed { + let apps_meta = metadata.codex_apps_meta.get_or_insert_default(); + if let Some(value) = issued_context { + apps_meta.insert(APPS_INVOCATION_CONTEXT.to_string(), value.clone()); + } else { + apps_meta.remove(APPS_INVOCATION_CONTEXT); + } + } let item_metadata = McpToolCallItemMetadata::from_tool_metadata(&server, Some(&metadata)); let connector_id = metadata.connector_id.clone(); let connector_name = metadata.connector_name.clone(); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 4a3498d504..4a4de6307e 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -2031,10 +2031,17 @@ fn result_metadata_apps_builder(base_url: String, account_email: &str) -> TestCo async fn code_mode_mcp_metadata_keeps_originating_window_after_compaction() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let apps_server = mount_result_metadata_app( - &server, /*result_metadata*/ None, /*is_error*/ false, - ) - .await?; + let tools = Arc::new(Mutex::new(vec![serde_json::json!({ + "name": RESULT_METADATA_TOOL, + "inputSchema": { "type": "object" }, + "annotations": { "readOnlyHint": true }, + "_meta": { + "connector_id": "test_connector", + "connector_name": "MessageSearch", + "_codex_apps": { "invocation_context": "original" }, + }, + })])); + let apps_server = AppsTestServer::mount_with_tools(&server, Arc::clone(&tools)).await?; let (reached_tx, reached_rx) = oneshot::channel(); let (release_tx, release_rx) = oneshot::channel(); let control = Arc::new(ResultMetadataTestControl { @@ -2118,6 +2125,11 @@ text("done");"#, "wait", ) .await; + tools.lock().unwrap()[0]["_meta"]["_codex_apps"] = serde_json::json!({ + "invocation_context": "refreshed", + "resource_uri": "connector://messages/search", + }); + test.codex.refresh_codex_apps_tools().await?; release_tx.send(()).unwrap(); test.submit_turn("Finish the code-mode cell").await?; let resumed_request = resumed.function_call.single_request(); @@ -2138,8 +2150,13 @@ text("done");"#, ); let calls = recorded_apps_tool_calls(&server).await; assert_eq!(calls.len(), 2); + assert_eq!( + calls[1]["params"]["_meta"]["_codex_apps"]["resource_uri"], + "connector://messages/search" + ); for call in calls { let meta = &call["params"]["_meta"]; + assert_eq!(meta["_codex_apps"]["invocation_context"], "original"); assert_eq!( serde_json::json!({ "sessionId": meta["sessionId"],