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

@@ -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<ExtensionRegistry<Config>>,
     codex_apps_tools_cache: ConnectorRuntimeManager<ToolInfo>,
     tool_catalog_cache: McpToolCatalogCache,
+    shared_definitions: Arc<SharedMcpDefinitions>,
 }
 
 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 {

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<McpHandler>,
 }
 
+/// Shares definitions for already-shared catalogs without keeping completed sessions alive.
+#[derive(Default)]
+pub(crate) struct SharedMcpDefinitions {
+    definitions: Mutex<HashMap<McpDefinitionKey, Weak<McpHandlerDefinition>>>,
+}
+
+#[derive(Eq, Hash, PartialEq)]
+struct McpDefinitionKey {
+    input_schema: usize,
+    tool_name: ToolName,
+    agent_plugin: bool,
+    schema_max_bytes: Option<NonZeroUsize>,
+}
+
+impl SharedMcpDefinitions {
+    fn get_or_create(
+        &self,
+        tool: &McpToolInfo,
+        agent_plugin: bool,
+        schema_max_bytes: Option<NonZeroUsize>,
+    ) -> Result<Arc<McpHandlerDefinition>, 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<McpBinding>,
@@ -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<ToolName> {
         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<ToolName, CachedMcpHandler>,
+    shared_definitions: &SharedMcpDefinitions,
     registry: &mut ToolRegistry,
 ) -> HashSet<ToolName> {
     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),

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,
     );
 

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<usize>,
     spec: Arc<ToolSpec>,
+}
+
+impl McpHandlerDefinition {
+    pub(crate) fn new(
+        tool_info: ToolInfo,
+        agent_plugin: bool,
+        schema_max_bytes: Option<usize>,
+    ) -> Result<Self, serde_json::Error> {
+        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<usize>,
+    ) -> bool {
+        self.tool_info == *tool_info
+            && self.agent_plugin == agent_plugin
+            && self.schema_max_bytes == schema_max_bytes
+    }
+}
+
+pub struct McpHandler {
+    definition: Arc<McpHandlerDefinition>,
     code_mode_tool_definitions:
         OnceLock<(Option<usize>, bool, Vec<codex_code_mode::ToolDefinition>)>,
 }
 
+#[cfg(test)]
 impl McpHandler {
     pub fn new(tool_info: ToolInfo) -> Result<Self, serde_json::Error> {
         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, serde_json::Error> {
-        Self::with_agent_plugin(
-            tool_info,
-            /*agent_plugin*/ false,
-            Some(schema_max_bytes),
-        )
-    }
-
     pub fn new_agent_plugin(tool_info: ToolInfo) -> Result<Self, serde_json::Error> {
         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<usize>,
     ) -> Result<Self, serde_json::Error> {
-        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<McpHandlerDefinition>) -> 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<usize, serde_json::Error> {
-        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<ToolInvocation> 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<ToolInvocation> for McpHandler {
 
     fn search_info(&self) -> Option<ToolSearchInfo> {
         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<Box<dyn crate::tools::context::ToolOutput>, 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<ToolSpec>> {
-        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<Session>) -> Option<BoxFuture<'a, ()>> {
         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::<NodeReplReviewEvidence>()
             .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

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)
     );
 }

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;

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(