diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index a5e03e2ce5..fa9b6eb61b 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -1,9 +1,8 @@ use std::collections::HashMap; use std::collections::HashSet; -use std::collections::hash_map::Entry; +use std::num::NonZeroUsize; use std::sync::Arc; use std::sync::Mutex; -use std::sync::Weak; use codex_connectors::AppToolPolicyEvaluator; use codex_connectors::AppToolPolicyInput; @@ -25,12 +24,14 @@ use crate::tools::registry::ToolRegistry; #[derive(Default)] pub(crate) struct McpHandlerCache { - cached: Mutex>, + handlers: Mutex>, } -struct CachedMcpHandlers { - binding: Weak, - handlers: HashMap>, +struct CachedMcpHandler { + tool_info: McpToolInfo, + agent_plugin: bool, + schema_max_bytes: Option, + handler: Arc, } impl McpHandlerCache { @@ -43,29 +44,17 @@ impl McpHandlerCache { search_tool_enabled: bool, registry: &mut ToolRegistry, ) -> HashSet { - let mut cached = self - .cached + let mut handlers = self + .handlers .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - if !cached - .as_ref() - .and_then(|cached| cached.binding.upgrade()) - .is_some_and(|cached_binding| Arc::ptr_eq(&cached_binding, binding)) - { - *cached = None; - } - - let cached = cached.get_or_insert_with(|| CachedMcpHandlers { - binding: Arc::downgrade(binding), - handlers: HashMap::new(), - }); append_mcp_tools( binding.tools(), config, apps_enabled, mcp_server_catalog, search_tool_enabled, - &mut cached.handlers, + &mut handlers, registry, ) } @@ -78,9 +67,15 @@ fn append_mcp_tools( apps_enabled: bool, mcp_server_catalog: &codex_mcp::ResolvedMcpCatalog, search_tool_enabled: bool, - handlers: &mut HashMap>, + handlers: &mut HashMap, registry: &mut ToolRegistry, ) -> HashSet { + let current_tools: HashSet<_> = all_mcp_tools + .iter() + .map(McpToolInfo::canonical_tool_name) + .collect(); + handlers.retain(|name, _| current_tools.contains(name)); + // Keep regular MCP tools first; Apps tools also require connector and policy checks. let non_app_tools = filter_non_codex_apps_mcp_tools_only(all_mcp_tools); let app_tools = apps_enabled @@ -100,25 +95,41 @@ fn append_mcp_tools( let agent_plugin = server.is_some_and(|server| server.source().is_agent_plugin()); let tool_input_schema_max_bytes = server.and_then(|server| server.config().tool_input_schema_max_bytes); - let handler = match handlers.entry(tool_name.clone()) { - Entry::Occupied(entry) => Arc::clone(entry.get()), - Entry::Vacant(entry) => { - 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()) - }; + // Handlers contain immutable tool metadata, not a connection or authorization snapshot. + // Preserve their identity across equivalent bindings so the search index can also be reused. + let handler = if let Some(cached) = handlers.get(&tool_name).filter(|cached| { + cached.tool_info == *tool + && cached.agent_plugin == agent_plugin + && cached.schema_max_bytes == tool_input_schema_max_bytes + }) { + 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()) + }; - match handler { - Ok(handler) => Arc::clone(entry.insert(Arc::new(handler))), - Err(err) => { - warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}"); - continue; - } + let handler = match handler { + Ok(handler) => Arc::new(handler), + Err(err) => { + warn!("Skipping MCP tool `{tool_name}`: failed to build tool spec: {err}"); + continue; } - } + }; + handlers.insert( + tool_name.clone(), + CachedMcpHandler { + tool_info: tool.clone(), + agent_plugin, + schema_max_bytes: tool_input_schema_max_bytes, + handler: Arc::clone(&handler), + }, + ); + handler }; let fits_agent_budget = if agent_plugin { diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index e373cd86c5..b24a4d860a 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -59,10 +59,18 @@ use std::collections::HashMap; use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; use tempfile::TempDir; use tokio::sync::Semaphore; +use tracing::Subscriber; +use tracing::span::Attributes; +use tracing::span::Id; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::Context as LayerContext; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; use wiremock::Mock; use wiremock::Request; use wiremock::ResponseTemplate; @@ -70,6 +78,27 @@ use wiremock::matchers::body_partial_json; use wiremock::matchers::method; use wiremock::matchers::path_regex; +#[derive(Clone, Default)] +struct McpCacheCounters { + binding_captures: Arc, + search_index_builds: Arc, +} + +impl Layer for McpCacheCounters { + fn on_new_span(&self, attributes: &Attributes<'_>, _id: &Id, _context: LayerContext<'_, S>) { + let metadata = attributes.metadata(); + match (metadata.target(), metadata.name()) { + ("codex_mcp::connection_manager::tool_catalog", "capture_binding_with_metadata") => { + self.binding_captures.fetch_add(1, Ordering::SeqCst); + } + ("codex_core::tools::handlers::tool_search", "new") => { + self.search_index_builds.fetch_add(1, Ordering::SeqCst); + } + _ => {} + } + } +} + struct McpResourceClientCapture { client: Arc>>, } @@ -1080,10 +1109,18 @@ async fn deferred_tool_world_state_is_disabled_by_default() -> Result<()> { Ok(()) } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +// Keep spawned tasks on the thread with the scoped tracing subscriber. +#[tokio::test(flavor = "current_thread")] async fn deferred_tool_world_state_tracks_initial_unchanged_and_removed_namespaces() -> Result<()> { skip_if_no_network!(Ok(())); + let counters = McpCacheCounters::default(); + // Keep concurrent tests without a subscriber from caching these callsites as disabled. + let _interest_cache_guard = + tracing::Dispatch::new(tracing::subscriber::NoSubscriber::default()); + let _tracing = tracing_subscriber::registry() + .with(counters.clone()) + .set_default(); let server = responses::start_mock_server().await; let apps_server = AppsTestServer::mount_searchable(&server).await?; let response = mount_sse_sequence(&server, completed_response_sequence(/*count*/ 3)).await; @@ -1094,7 +1131,29 @@ async fn deferred_tool_world_state_tracks_initial_unchanged_and_removed_namespac test.submit_turn("inspect initially available deferred tools") .await?; + let initial_captures = counters.binding_captures.load(Ordering::SeqCst); + let initial_index_builds = counters.search_index_builds.load(Ordering::SeqCst); + assert!( + initial_captures > 0, + "the initial turn must capture an MCP binding" + ); + assert!( + initial_index_builds > 0, + "the initial turn must build the search index" + ); + + // Publish a new catalog revision with the same metadata from the ready client. + test.codex.refresh_codex_apps_tools().await?; test.submit_turn("inspect unchanged deferred tools").await?; + assert!( + counters.binding_captures.load(Ordering::SeqCst) > initial_captures, + "the follow-up must capture a new binding after the refresh" + ); + assert_eq!( + counters.search_index_builds.load(Ordering::SeqCst), + initial_index_builds, + "equivalent bindings must preserve MCP handlers and reuse the search index" + ); let mut refresh_config = test.config.clone(); let user_config_path = refresh_config.codex_home.join("config.toml"); @@ -1113,6 +1172,14 @@ enabled = false let requests = response.requests(); assert_eq!(requests.len(), 3); + assert!( + requests[2].body_json()["tools"] + .as_array() + .expect("model request tools") + .iter() + .all(|tool| tool["type"] != "tool_search"), + "removing all deferred tools must stop advertising the cached search tool" + ); let tools_states = requests .iter() .map(tools_state_sections) @@ -1133,6 +1200,7 @@ enabled = false ) ); + test.codex.shutdown_and_wait().await?; Ok(()) }