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(