From 2744e8686c3defaa96f7aff6c51a9df961c51224 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Wed, 21 Oct 2026 01:44:41 +0000 Subject: [PATCH] Share prepared MCP tool definitions across sessions (#55724) ## Why Sessions using the same MCP catalog can repeat schema normalization and retain duplicate tool specifications. ## What changed Split immutable MCP tool metadata and normalized specifications from per-handler caches, and reuse matching definitions within each `McpManager`. Require matching tool metadata, agent-plugin mode, and schema budget before reuse. Keep shared definitions behind weak references and prune expired entries on cache misses so completed sessions do not keep their schemas alive. Build definitions outside the shared index lock. GitOrigin-RevId: 02345359a1b0d297ca56f2f469de88aba89106f1 --- codex-rs/core/src/mcp.rs | 7 + codex-rs/core/src/mcp_tool_exposure.rs | 95 ++++++++++++- codex-rs/core/src/mcp_tool_exposure_test.rs | 6 + codex-rs/core/src/tools/handlers/mcp.rs | 126 ++++++++++++------ .../src/tools/handlers/mcp_search_tests.rs | 2 +- codex-rs/core/src/tools/handlers/mod.rs | 1 + codex-rs/core/src/tools/spec_plan.rs | 1 + 7 files changed, 186 insertions(+), 52 deletions(-) diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 2cdc8a2659..a3046d892b 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -3,6 +3,7 @@ use std::collections::HashSet; use std::sync::Arc; use crate::config::Config; +use crate::mcp_tool_exposure::SharedMcpDefinitions; use codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID; use codex_config::McpServerConfig; use codex_config::types::PluginsConfigToml; @@ -114,6 +115,7 @@ pub struct McpManager { extensions: Arc>, codex_apps_tools_cache: ConnectorRuntimeManager, tool_catalog_cache: McpToolCatalogCache, + shared_definitions: Arc, } impl McpManager { @@ -136,6 +138,7 @@ impl McpManager { extensions, codex_apps_tools_cache, tool_catalog_cache: McpToolCatalogCache::default(), + shared_definitions: Arc::new(SharedMcpDefinitions::default()), } } @@ -147,6 +150,10 @@ impl McpManager { self.tool_catalog_cache.clone() } + pub(crate) fn shared_definitions(&self) -> &SharedMcpDefinitions { + &self.shared_definitions + } + /// Returns the MCP config after applying compatibility built-ins and /// runtime-only extension overlays. pub async fn runtime_config(&self, config: &Config) -> McpConfig { diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index e7aa97993a..6ecb7ccff2 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -3,6 +3,7 @@ use std::collections::HashSet; use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::Mutex; +use std::sync::Weak; use codex_connectors::AppToolPolicyEvaluator; use codex_connectors::AppToolPolicyInput; @@ -20,6 +21,7 @@ const MAX_AGENT_PLUGIN_MCP_TOTAL_BYTES: usize = 64_000; use crate::config::Config; use crate::tools::handlers::McpHandler; +use crate::tools::handlers::McpHandlerDefinition; use crate::tools::registry::ToolRegistry; #[derive(Default)] @@ -33,7 +35,75 @@ struct CachedMcpHandler { handler: Arc, } +/// Shares definitions for already-shared catalogs without keeping completed sessions alive. +#[derive(Default)] +pub(crate) struct SharedMcpDefinitions { + definitions: Mutex>>, +} + +#[derive(Eq, Hash, PartialEq)] +struct McpDefinitionKey { + input_schema: usize, + tool_name: ToolName, + agent_plugin: bool, + schema_max_bytes: Option, +} + +impl SharedMcpDefinitions { + fn get_or_create( + &self, + tool: &McpToolInfo, + agent_plugin: bool, + schema_max_bytes: Option, + ) -> Result, serde_json::Error> { + let key = McpDefinitionKey { + // A live definition owns this Arc. The address only selects a candidate; + // complete prepared metadata equality is still required below. + input_schema: Arc::as_ptr(&tool.tool.input_schema) as usize, + tool_name: tool.canonical_tool_name(), + agent_plugin, + schema_max_bytes, + }; + let budget = schema_max_bytes.map(NonZeroUsize::get); + let existing = self + .definitions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&key) + .and_then(Weak::upgrade); + if let Some(existing) = + existing.filter(|definition| definition.matches(tool, agent_plugin, budget)) + { + return Ok(existing); + } + + // Normalization can be expensive and must not hold the manager's index lock. + let definition = Arc::new(McpHandlerDefinition::new( + tool.clone(), + agent_plugin, + budget, + )?); + let mut definitions = self + .definitions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let current = definitions.get(&key).and_then(Weak::upgrade); + let winner = current + .as_ref() + .filter(|current| current.matches(tool, agent_plugin, budget)) + .cloned(); + if winner.is_none() { + definitions.insert(key, Arc::downgrade(&definition)); + } + // An upgraded or losing definition may own a large schema graph. + drop(definitions); + drop(current); + Ok(winner.unwrap_or(definition)) + } +} + impl McpHandlerCache { + #[allow(clippy::too_many_arguments)] pub(crate) fn append_mcp_tools( &self, binding: &Arc, @@ -41,6 +111,7 @@ impl McpHandlerCache { apps_enabled: bool, mcp_server_catalog: &codex_mcp::ResolvedMcpCatalog, search_tool_enabled: bool, + shared_definitions: &SharedMcpDefinitions, registry: &mut ToolRegistry, ) -> HashSet { let mut handlers = self @@ -54,12 +125,14 @@ impl McpHandlerCache { mcp_server_catalog, search_tool_enabled, &mut handlers, + shared_definitions, registry, ) } } #[instrument(level = "trace", skip_all)] +#[allow(clippy::too_many_arguments)] fn append_mcp_tools( all_mcp_tools: &[McpToolInfo], config: &Config, @@ -67,6 +140,7 @@ fn append_mcp_tools( mcp_server_catalog: &codex_mcp::ResolvedMcpCatalog, search_tool_enabled: bool, handlers: &mut HashMap, + shared_definitions: &SharedMcpDefinitions, registry: &mut ToolRegistry, ) -> HashSet { let current_tools: HashSet<_> = all_mcp_tools @@ -88,6 +162,7 @@ fn append_mcp_tools( }; let mut registered_tools = HashSet::new(); let mut agent_plugin_bytes = 0usize; + let mut pruned_shared_definitions = false; for tool in non_app_tools.chain(app_tools) { let tool_name = tool.canonical_tool_name(); let server = mcp_server_catalog.server(&tool.server_name); @@ -104,13 +179,19 @@ fn append_mcp_tools( Arc::clone(&cached.handler) } else { handlers.remove(&tool_name); - let handler = if agent_plugin { - McpHandler::new_agent_plugin(tool.clone()) - } else if let Some(budget) = tool_input_schema_max_bytes { - McpHandler::new_with_schema_max_bytes(tool.clone(), budget.get()) - } else { - McpHandler::new(tool.clone()) - }; + if !pruned_shared_definitions { + // Sweep once per miss-containing pass, never on the unchanged hot path. + // Idle managers can retain dead keys/control blocks until their next miss. + shared_definitions + .definitions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .retain(|_, definition| definition.strong_count() > 0); + pruned_shared_definitions = true; + } + let handler = shared_definitions + .get_or_create(tool, agent_plugin, tool_input_schema_max_bytes) + .map(McpHandler::from_definition); let handler = match handler { Ok(handler) => Arc::new(handler), diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs index c94be8ce55..5a9967d30b 100644 --- a/codex-rs/core/src/mcp_tool_exposure_test.rs +++ b/codex-rs/core/src/mcp_tool_exposure_test.rs @@ -104,6 +104,7 @@ fn runtimes_by_name_with_catalog( mcp_server_catalog, search_tool_enabled, &mut handlers, + &SharedMcpDefinitions::default(), &mut registry, ); registry @@ -254,6 +255,7 @@ async fn cached_app_handlers_still_obey_current_apps_enablement_and_tool_policy( &catalog, /*search_tool_enabled*/ false, &mut handlers, + &SharedMcpDefinitions::default(), &mut allowed_registry, ); let cached_handler = &allowed_registry @@ -270,6 +272,7 @@ async fn cached_app_handlers_still_obey_current_apps_enablement_and_tool_policy( &catalog, /*search_tool_enabled*/ false, &mut handlers, + &SharedMcpDefinitions::default(), &mut disabled_registry, ); let mut restricted_registry = ToolRegistry::default(); @@ -280,6 +283,7 @@ async fn cached_app_handlers_still_obey_current_apps_enablement_and_tool_policy( &catalog, /*search_tool_enabled*/ false, &mut handlers, + &SharedMcpDefinitions::default(), &mut restricted_registry, ); let mut restored_registry = ToolRegistry::default(); @@ -290,6 +294,7 @@ async fn cached_app_handlers_still_obey_current_apps_enablement_and_tool_policy( &catalog, /*search_tool_enabled*/ true, &mut handlers, + &SharedMcpDefinitions::default(), &mut restored_registry, ); let restored_handler = restored_registry @@ -434,6 +439,7 @@ async fn app_tool_registration_uses_trusted_catalog_metadata_and_preserves_sourc &ResolvedMcpCatalog::default(), /*search_tool_enabled*/ false, &mut handlers, + &SharedMcpDefinitions::default(), &mut registry, ); diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 98d61aff08..cadcbb855a 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -48,14 +48,52 @@ const MCP_TOOL_NAME_DELIMITER: &str = "__"; const MAX_AGENT_PLUGIN_MCP_NAMESPACE_DESCRIPTION_BYTES: usize = 1_000; const MAX_MCP_NAMESPACE_DESCRIPTION_BYTES: usize = 512 * 1024; -pub struct McpHandler { +/// Immutable tool metadata and normalized schema, shared only within an MCP manager. +pub(crate) struct McpHandlerDefinition { tool_info: ToolInfo, agent_plugin: bool, + schema_max_bytes: Option, spec: Arc, +} + +impl McpHandlerDefinition { + pub(crate) fn new( + tool_info: ToolInfo, + agent_plugin: bool, + schema_max_bytes: Option, + ) -> Result { + let spec = Arc::new(create_tool_spec( + &tool_info, + agent_plugin, + schema_max_bytes, + )?); + Ok(Self { + tool_info, + agent_plugin, + schema_max_bytes, + spec, + }) + } + + pub(crate) fn matches( + &self, + tool_info: &ToolInfo, + agent_plugin: bool, + schema_max_bytes: Option, + ) -> bool { + self.tool_info == *tool_info + && self.agent_plugin == agent_plugin + && self.schema_max_bytes == schema_max_bytes + } +} + +pub struct McpHandler { + definition: Arc, code_mode_tool_definitions: OnceLock<(Option, bool, Vec)>, } +#[cfg(test)] impl McpHandler { pub fn new(tool_info: ToolInfo) -> Result { Self::with_agent_plugin( @@ -63,17 +101,6 @@ impl McpHandler { ) } - pub fn new_with_schema_max_bytes( - tool_info: ToolInfo, - schema_max_bytes: usize, - ) -> Result { - Self::with_agent_plugin( - tool_info, - /*agent_plugin*/ false, - Some(schema_max_bytes), - ) - } - pub fn new_agent_plugin(tool_info: ToolInfo) -> Result { Self::with_agent_plugin( tool_info, /*agent_plugin*/ true, /*schema_max_bytes*/ None, @@ -85,25 +112,28 @@ impl McpHandler { agent_plugin: bool, schema_max_bytes: Option, ) -> Result { - let spec = Arc::new(create_tool_spec( - &tool_info, - agent_plugin, - schema_max_bytes, - )?); - Ok(Self { + Ok(Self::from_definition(Arc::new(McpHandlerDefinition::new( tool_info, agent_plugin, - spec, + schema_max_bytes, + )?))) + } +} + +impl McpHandler { + pub(crate) fn from_definition(definition: Arc) -> Self { + Self { + definition, code_mode_tool_definitions: OnceLock::new(), - }) + } } pub(crate) fn tool_info(&self) -> &ToolInfo { - &self.tool_info + &self.definition.tool_info } pub(crate) fn model_spec_bytes(&self) -> Result { - serde_json::to_vec(&self.spec).map(|spec| spec.len()) + serde_json::to_vec(&self.definition.spec).map(|spec| spec.len()) } fn hook_tool_name(&self) -> HookToolName { @@ -132,18 +162,19 @@ fn ensure_mcp_prefix(name: &str) -> String { impl ToolExecutor for McpHandler { fn tool_name(&self) -> ToolName { - self.tool_info.canonical_tool_name() + self.definition.tool_info.canonical_tool_name() } fn spec(&self) -> ToolSpec { - self.spec.as_ref().clone() + self.definition.spec.as_ref().clone() } fn supports_parallel_tool_calls(&self) -> bool { // Correctly implemented MCP servers should tolerate parallel calls to // tools that advertise themselves as read-only. - self.tool_info.supports_parallel_tool_calls + self.definition.tool_info.supports_parallel_tool_calls || self + .definition .tool_info .tool .annotations @@ -154,23 +185,27 @@ impl ToolExecutor for McpHandler { fn search_info(&self) -> Option { let source_name = self + .definition .tool_info .connector_name .as_deref() .map(str::trim) .filter(|connector_name| !connector_name.is_empty()) - .unwrap_or_else(|| self.tool_info.server_name.trim()); + .unwrap_or_else(|| self.definition.tool_info.server_name.trim()); let source_info = (!source_name.is_empty()).then(|| ToolSearchSourceInfo { name: source_name.to_string(), - description: effective_namespace_description(&self.tool_info, self.agent_plugin) - .map(str::trim) - .filter(|description| !description.is_empty()) - .map(str::to_string), + description: effective_namespace_description( + &self.definition.tool_info, + self.definition.agent_plugin, + ) + .map(str::trim) + .filter(|description| !description.is_empty()) + .map(str::to_string), }); ToolSearchInfo::from_shared_spec( - build_mcp_search_text(&self.tool_info, self.agent_plugin), - Arc::clone(&self.spec), + build_mcp_search_text(&self.definition.tool_info, self.definition.agent_plugin), + Arc::clone(&self.definition.spec), source_info, ) } @@ -188,7 +223,10 @@ impl McpHandler { &self, invocation: ToolInvocation, ) -> Result, FunctionCallError> { - let prepared_mcp_call = invocation.session.prepare_mcp_call(&self.tool_info).await; + let prepared_mcp_call = invocation + .session + .prepare_mcp_call(&self.definition.tool_info) + .await; // Use the executed call's binding; a later catalog refresh must not change eligibility. let result_metadata_capture_allowed = invocation .session @@ -245,7 +283,7 @@ impl McpHandler { &cancellation_token, call_id.clone(), originating_call, - &self.tool_info, + &self.definition.tool_info, prepared_mcp_call, self.hook_tool_name(), tool_name, @@ -269,7 +307,7 @@ impl McpHandler { impl CoreToolRuntime for McpHandler { fn immutable_spec(&self) -> Option<&Arc> { - Some(&self.spec) + Some(&self.definition.spec) } fn cached_code_mode_definitions( @@ -280,7 +318,7 @@ impl CoreToolRuntime for McpHandler { let (cached_budget, cached_tool_description_first, definitions) = self.code_mode_tool_definitions.get_or_init(|| { let mut definitions = codex_tools::collect_code_mode_tool_definitions( - std::iter::once(self.spec.as_ref()), + std::iter::once(self.definition.spec.as_ref()), code_mode_input_schema_max_bytes, tool_description_first, ); @@ -303,7 +341,7 @@ impl CoreToolRuntime for McpHandler { fn wait_until_ready<'a>(&'a self, session: &'a Arc) -> Option> { Some(Box::pin(async move { session - .wait_for_mcp_server(&self.tool_info.server_name) + .wait_for_mcp_server(&self.definition.tool_info.server_name) .await; })) } @@ -313,7 +351,7 @@ impl CoreToolRuntime for McpHandler { } fn mcp_server_name(&self) -> Option<&str> { - Some(&self.tool_info.server_name) + Some(&self.definition.tool_info.server_name) } fn on_tool_result_accepted(&self, invocation: &ToolInvocation, result: &dyn ToolOutput) { @@ -333,8 +371,8 @@ impl CoreToolRuntime for McpHandler { .get::() .is_some_and(|evidence| evidence.image_capture_enabled()); if !is_node_repl_backed_connector( - &self.tool_info.server_name, - self.tool_info.connector_id.as_deref(), + &self.definition.tool_info.server_name, + self.definition.tool_info.connector_id.as_deref(), ) || !result.success_for_logging() || evidence_mode == NodeReplReviewEvidenceMode::Disabled && !image_capture_enabled { @@ -430,7 +468,7 @@ impl CoreToolRuntime for McpHandler { .record( &format!( "{}.{}", - self.tool_info.server_name, self.tool_info.tool.name + self.definition.tool_info.server_name, self.definition.tool_info.tool.name ), cell_id, &invocation.call_id, @@ -439,8 +477,8 @@ impl CoreToolRuntime for McpHandler { } fn telemetry_tags(&self, _invocation: &ToolInvocation) -> ToolTelemetryTags { - let mut tags = vec![("mcp_server", self.tool_info.server_name.clone())]; - if let Some(origin) = &self.tool_info.server_origin { + let mut tags = vec![("mcp_server", self.definition.tool_info.server_name.clone())]; + if let Some(origin) = &self.definition.tool_info.server_origin { tags.push(("mcp_server_origin", origin.clone())); } tags @@ -840,7 +878,7 @@ mod tests { handler .immutable_spec() .expect("MCP spec should be immutable"), - &handler.spec, + &handler.definition.spec, )); let first = handler diff --git a/codex-rs/core/src/tools/handlers/mcp_search_tests.rs b/codex-rs/core/src/tools/handlers/mcp_search_tests.rs index 50c11dfd5a..d1b2be67ad 100644 --- a/codex-rs/core/src/tools/handlers/mcp_search_tests.rs +++ b/codex-rs/core/src/tools/handlers/mcp_search_tests.rs @@ -62,7 +62,7 @@ fn mcp_namespace_descriptions_preserve_complete_metadata() { }; assert_eq!(namespace.description, full_description); assert_eq!( - handler.tool_info.namespace_description, + handler.tool_info().namespace_description, Some(full_description) ); } diff --git a/codex-rs/core/src/tools/handlers/mod.rs b/codex-rs/core/src/tools/handlers/mod.rs index a4d92223c3..56ae313eef 100644 --- a/codex-rs/core/src/tools/handlers/mod.rs +++ b/codex-rs/core/src/tools/handlers/mod.rs @@ -62,6 +62,7 @@ pub use dynamic::DynamicToolHandler; pub use get_context_remaining::GetContextRemainingHandler; pub use list_available_plugins_to_install::ListAvailablePluginsToInstallHandler; pub use mcp::McpHandler; +pub(crate) use mcp::McpHandlerDefinition; pub use mcp_resource::ListMcpResourceTemplatesHandler; pub use mcp_resource::ListMcpResourcesHandler; pub use mcp_resource::ReadMcpResourceHandler; diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index 1b86b579e3..24b60f3ff2 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/codex-rs/core/src/tools/spec_plan.rs @@ -154,6 +154,7 @@ pub(crate) fn build_tool_router( apps_enabled, &mcp.config().mcp_server_catalog, model_info.supports_search_tool, + session.services.mcp_manager.shared_definitions(), &mut registry, ); apply_mcp_tool_exposure_policy(