From cef2a64531430e96efd69ceed18b8c34dc2cb98f Mon Sep 17 00:00:00 2001 From: Owen Lin Date: Wed, 7 Oct 2026 20:35:07 +0000 Subject: [PATCH] Share MCP binding tool metadata with prepared calls (#51823) ## Why Prepared MCP calls cloned `ToolInfo` already stored in the binding's model-visible catalog, duplicating tool metadata. ## What changed Store the binding catalog in an `Arc<[ToolInfo]>` and let prepared calls reference its entries by index. Keep separately owned metadata for tools hidden from the model and calls prepared outside the binding catalog. ## Testing Extend the binding test to assert that advertised and prepared tool metadata share the same address, while preserving ready-client catalog selection and hidden-tool behavior when tool names collide. GitOrigin-RevId: 24b92f7e1dbb421c9ebce875b9b1024f7d697c5e --- codex-rs/codex-mcp/src/binding.rs | 52 ++++++++++--- .../src/connection_manager/tool_catalog.rs | 73 +++++++++++++------ .../codex-mcp/src/connection_manager_tests.rs | 34 ++++++++- 3 files changed, 123 insertions(+), 36 deletions(-) diff --git a/codex-rs/codex-mcp/src/binding.rs b/codex-rs/codex-mcp/src/binding.rs index f0a95b95c6..de9a72f49b 100644 --- a/codex-rs/codex-mcp/src/binding.rs +++ b/codex-rs/codex-mcp/src/binding.rs @@ -34,7 +34,7 @@ pub struct McpBinding { clients: Arc, config: Arc, plugins_available: bool, - tools: Vec, + tools: Arc<[ToolInfo]>, calls: HashMap<(String, String), PreparedMcpCall>, } @@ -56,7 +56,7 @@ impl McpBinding { clients: Arc, config: Arc, plugins_available: bool, - tools: Vec, + tools: impl Into>, calls: HashMap<(String, String), PreparedMcpCall>, ) -> Self { Self { @@ -64,7 +64,7 @@ impl McpBinding { clients, config, plugins_available, - tools, + tools: tools.into(), calls, } } @@ -173,13 +173,42 @@ pub struct PreparedMcpCall { client: Arc, config: Arc, catalog_snapshot: Arc, - tool_info: ToolInfo, + tool_info: PreparedToolInfo, server_name: String, server_metadata: McpServerMetadata, plugin_id: Option, selected_plugin_server: bool, } +#[derive(Clone)] +pub(crate) struct PreparedToolInfo { + catalog: Arc<[ToolInfo]>, + index: usize, +} + +impl PreparedToolInfo { + pub(crate) fn from_binding_catalog(catalog: Arc<[ToolInfo]>, index: usize) -> Self { + assert!( + index < catalog.len(), + "prepared tool index must refer to its binding catalog" + ); + Self { catalog, index } + } + + pub(crate) fn get(&self) -> &ToolInfo { + &self.catalog[self.index] + } +} + +impl From for PreparedToolInfo { + fn from(tool_info: ToolInfo) -> Self { + Self { + catalog: Arc::new([tool_info]), + index: 0, + } + } +} + impl PreparedMcpCall { #[expect( clippy::too_many_arguments, @@ -190,12 +219,13 @@ impl PreparedMcpCall { client: Arc, config: Arc, catalog_snapshot: Arc, - tool_info: ToolInfo, + tool_info: impl Into, server_metadata: McpServerMetadata, plugin_id: Option, selected_plugin_server: bool, ) -> Option { - let server_name = tool_info.server_name.clone(); + let tool_info = tool_info.into(); + let server_name = tool_info.get().server_name.clone(); config.permission_profile_for_server(&server_name)?; Some(Self { connections, @@ -211,7 +241,7 @@ impl PreparedMcpCall { } pub fn tool_info(&self) -> &ToolInfo { - &self.tool_info + self.tool_info.get() } /// Returns the configuration and approval authority captured with this client. @@ -261,7 +291,7 @@ impl PreparedMcpCall { pub fn tool_approval_mode(&self) -> AppToolApproval { self.server_metadata - .tool_approval_mode(&self.tool_info.tool.name) + .tool_approval_mode(&self.tool_info().tool.name) } /// Returns the explicit output budget captured with this call's effective server config. @@ -271,7 +301,7 @@ impl PreparedMcpCall { .server(&self.server_name)? .config() .tools - .get(self.tool_info.tool.name.as_ref())? + .get(self.tool_info().tool.name.as_ref())? .output_token_limit .map(std::num::NonZeroUsize::get) } @@ -316,7 +346,7 @@ impl PreparedMcpCall { } (server_timeout, requested_timeout) => server_timeout.or(requested_timeout), }; - let tool_name = self.tool_info.tool.name.to_string(); + let tool_name = self.tool_info().tool.name.to_string(); self.client .tool_catalog .run_with_snapshot(&self.catalog_snapshot, || async { @@ -324,7 +354,7 @@ impl PreparedMcpCall { let timeout_deadline = effective_timeout.map(|timeout| tokio::time::Instant::now() + timeout); let add_trusted_access_context = self.connections.add_trusted_access_context( - &self.tool_info, + self.tool_info(), &self.server_metadata, arguments.as_ref(), meta, diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index 54c6de2e8c..7ebf5af6f7 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -22,6 +22,7 @@ use super::catalog_telemetry::record_binding_catalog_size; use super::catalog_telemetry::tool_definition_json_bytes; use crate::binding::McpBinding; use crate::binding::PreparedMcpCall; +use crate::binding::PreparedToolInfo; use crate::binding_clients::McpBindingClients; use crate::client_tool_catalog::ClientToolCatalogRevision; use crate::client_tool_catalog::CodexAppsToolSnapshot; @@ -70,6 +71,11 @@ pub(crate) enum BindingCatalogRevision { Cached(u64), } +enum PendingCallToolInfo { + BindingCatalog(usize), + Owned(PreparedToolInfo), +} + impl McpConnectionSet { /// Observes opted-in servers at startup gates without waiting or cloning catalogs. pub(crate) fn record_startup_readiness( @@ -461,7 +467,7 @@ impl McpConnectionSet { &self.non_prefixed_mcp_tool_servers, ); let mut tools = Vec::with_capacity(listed_tools.len()); - let mut calls = std::collections::HashMap::with_capacity(listed_tools.len()); + let mut pending_calls = Vec::with_capacity(listed_tools.len()); for tool_info in listed_tools { let model_visible = crate::tool_is_model_visible(&tool_info); let Some((client, snapshot)) = clients.get(&tool_info.server_name) else { @@ -470,29 +476,46 @@ impl McpConnectionSet { } continue; }; - let Some(call) = self.prepare_call( - &tool_info, - Arc::clone(client), - Arc::clone(&config), - Arc::clone(snapshot), - ) else { + if config + .permission_profile_for_server(&tool_info.server_name) + .is_none() + { trace!( server_name = %tool_info.server_name, tool_name = %tool_info.tool.name, "omitting MCP tool without an exact ready client" ); continue; + } + let pending_tool_info = if model_visible { + let tool_index = tools.len(); + tools.push(tool_info); + PendingCallToolInfo::BindingCatalog(tool_index) + } else { + PendingCallToolInfo::Owned(tool_info.into()) + }; + pending_calls.push((pending_tool_info, Arc::clone(client), Arc::clone(snapshot))); + } + let tools: Arc<[ToolInfo]> = tools.into(); + let mut calls = std::collections::HashMap::with_capacity(pending_calls.len()); + for (pending_tool_info, client, snapshot) in pending_calls { + let tool_info = match pending_tool_info { + PendingCallToolInfo::BindingCatalog(tool_index) => { + PreparedToolInfo::from_binding_catalog(Arc::clone(&tools), tool_index) + } + PendingCallToolInfo::Owned(tool_info) => tool_info, + }; + let Some(call) = self.prepare_call(tool_info, client, Arc::clone(&config), snapshot) + else { + continue; }; calls.insert( ( - tool_info.server_name.clone(), - tool_info.tool.name.to_string(), + call.server_name().to_string(), + call.tool_info().tool.name.to_string(), ), call, ); - if model_visible { - tools.push(tool_info); - } } let clients = Arc::new(McpBindingClients::new( clients @@ -553,28 +576,36 @@ impl McpConnectionSet { tool_info .callable_name .clone_from(&advertised_tool.callable_name); - self.prepare_call(&tool_info, Arc::new(client), config, snapshot) + self.prepare_call(tool_info, Arc::new(client), config, snapshot) } fn prepare_call( self: &Arc, - tool_info: &ToolInfo, + tool_info: impl Into, client: Arc, config: Arc, tool_catalog_snapshot: Arc, ) -> Option { - let server_name = &tool_info.server_name; - let view = self.servers.get(server_name)?; + let tool_info = tool_info.into(); + let (server_metadata, plugin_id, selected_plugin_server) = { + let server_name = &tool_info.get().server_name; + let view = self.servers.get(server_name)?; + ( + view.metadata.clone(), + self.plugin_id_for_mcp_server_name(server_name) + .map(str::to_string), + self.is_selected_plugin_mcp_server(server_name), + ) + }; PreparedMcpCall::new( Arc::clone(self), client, config, tool_catalog_snapshot, - tool_info.clone(), - view.metadata.clone(), - self.plugin_id_for_mcp_server_name(server_name) - .map(str::to_string), - self.is_selected_plugin_mcp_server(server_name), + tool_info, + server_metadata, + plugin_id, + selected_plugin_server, ) } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 52be4fe101..f40a689beb 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2309,7 +2309,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { } #[tokio::test] -async fn capture_binding_uses_the_ready_clients_own_tools() { +async fn capture_binding_preserves_ready_client_catalog_and_shares_prepared_metadata() { let codex_home = tempdir().expect("tempdir"); let cache_context = create_codex_apps_tools_cache_context( codex_home.path().to_path_buf(), @@ -2323,9 +2323,21 @@ async fn capture_binding_uses_the_ready_clients_own_tools() { "shared_cached_tool", )], ); + let mut visible_collision = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "collision"); + visible_collision.connector_id = Some("a-visible".to_string()); + let mut hidden_collision = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "collision"); + hidden_collision.connector_id = Some("z-hidden".to_string()); + hidden_collision.tool.meta = Some(rmcp::model::MetaObject( + serde_json::json!({ "ui": { "visibility": [] } }) + .as_object() + .expect("metadata object") + .clone(), + )); let mut ready_client = create_test_managed_client(vec![ create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "client_local_tool"), create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "client_local_blocked"), + visible_collision, + hidden_collision, ]) .await; let tool_filter = ToolFilter { @@ -2387,11 +2399,25 @@ async fn capture_binding_uses_the_ready_clients_own_tools() { .iter() .map(|tool| tool.callable_name.as_str()) .collect::>(), - vec!["client_local_tool"] + vec!["client_local_tool", "collision"] + ); + let prepared = step + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "client_local_tool") + .expect("client-local tool should be callable"); + assert!(std::ptr::eq( + step.tools() + .first() + .expect("binding should advertise the client-local tool first"), + prepared.tool_info() + )); + assert_eq!( + step.tool_info(CODEX_APPS_MCP_SERVER_NAME, "collision") + .and_then(|tool| tool.connector_id.as_deref()), + Some("z-hidden") ); assert!( - step.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "client_local_tool") - .is_some() + step.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "collision") + .is_none() ); assert!( step.prepare_call(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool")