From 3bbf1fe75701c97fb190e0867002ba2d9dbda5db Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 27 Jul 2026 10:16:56 +0000 Subject: [PATCH] Expose cached MCP tools before server startup (#35590) ## Why Cached MCP definitions can be supplied to inference without waiting for the server to finish initializing. ## What changed - Publish cached tools while startup is still in progress, clearing their potentially stale read-only hint. - Wait for the selected server to start before executing a tool call, then prepare the call against the refreshed live binding. - Keep cached tools visible in a binding even when no live client is available, while rejecting attempts to prepare those calls. ## Testing - Cover cached-tool visibility before startup and replacement with live tool metadata afterward. - Verify cached definitions reach inference before MCP initialization and that calls unavailable in the live catalog return the expected model-visible error. GitOrigin-RevId: 3aae8f474c344ccdc5e08fe321bbad21d85bffd1 --- codex-rs/codex-mcp/src/connection_manager.rs | 7 +++ .../src/connection_manager/tool_catalog.rs | 14 ++++++ .../codex-mcp/src/connection_manager_tests.rs | 50 +++++++++++++++---- codex-rs/codex-mcp/src/runtime.rs | 15 ++++++ codex-rs/core/src/mcp_tool_call.rs | 6 ++- codex-rs/core/tests/suite/mcp_tool_cache.rs | 17 ++++--- 6 files changed, 89 insertions(+), 20 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 2e6be9864f84..15e79550ee54 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -539,6 +539,13 @@ impl McpConnectionSet { self.servers.contains_key(server_name) } + pub(crate) async fn wait_for_server_startup(&self, server_name: &str) -> bool { + let Some(view) = self.servers.get(server_name) else { + return false; + }; + view.connection.client().await.is_ok() + } + /// Stop all MCP clients owned by this manager and terminate stdio server processes. pub async fn shutdown(&self) { let connections = self 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 14b41bb2c255..1da31346481c 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -143,6 +143,19 @@ impl McpConnectionSet { .startup_complete .load(Ordering::Acquire) { + if view.connection.client.has_cached_tools() { + if let Some(server_tools) = + view.listed_tools(&self.tool_plugin_provenance).await + { + listed_tools.extend(server_tools.into_iter().map(|mut tool| { + if let Some(annotations) = tool.tool.annotations.as_mut() { + annotations.read_only_hint = None; + } + Self::with_server_metadata(tool, &view.metadata) + })); + } + continue; + } let _ = view.connection.client.client().await; } view.connection.client.reconnect_failed_startup().await; @@ -186,6 +199,7 @@ impl McpConnectionSet { continue; } let Some(client) = clients.client(&tool_info.server_name) else { + tools.push(tool_info); continue; }; let Some(call) = self.prepare_call(&tool_info, client, Arc::clone(&config), *revision) diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 25b3dc870d1d..82e97b1da206 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1513,20 +1513,21 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { } #[tokio::test] -async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() { +async fn capture_binding_exposes_cached_tools_before_startup() { let codex_home = tempdir().expect("tempdir"); let cache_context = create_codex_apps_tools_cache_context( codex_home.path().to_path_buf(), Some("account-one"), Some("user-one"), ); - store_current_tools( - &cache_context, - vec![create_test_tool( - CODEX_APPS_MCP_SERVER_NAME, - "shared_cached_tool", - )], + let mut cached_tool = create_test_tool(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool"); + cached_tool.tool.annotations = Some( + rmcp::model::ToolAnnotations::new() + .read_only(true) + .destructive(false) + .open_world(false), ); + store_current_tools(&cache_context, vec![cached_tool]); let startup_complete = Arc::new(std::sync::atomic::AtomicBool::new(false)); let startup_complete_for_client = Arc::clone(&startup_complete); let (startup_started, wait_for_startup) = tokio::sync::oneshot::channel(); @@ -1575,14 +1576,41 @@ async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() { }, ); let manager = Arc::new(manager); - let manager_for_capture = Arc::clone(&manager); - let capture = tokio::spawn(async move { capture_binding(&manager_for_capture).await }); + let cached_binding = capture_binding(&manager).await; + assert_eq!( + cached_binding + .tools() + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["shared_cached_tool"] + ); + assert_eq!( + cached_binding.tools()[0].tool.annotations, + Some( + rmcp::model::ToolAnnotations::new() + .destructive(false) + .open_world(false) + ) + ); + assert!( + cached_binding + .prepare_call(CODEX_APPS_MCP_SERVER_NAME, "shared_cached_tool") + .is_none() + ); + + let manager_for_startup = Arc::clone(&manager); + let startup = tokio::spawn(async move { + manager_for_startup + .wait_for_server_startup(CODEX_APPS_MCP_SERVER_NAME) + .await + }); wait_for_startup.await.expect("client startup should begin"); - assert!(!capture.is_finished()); release_startup.send(()).expect("release client startup"); + assert!(startup.await.expect("startup task")); - let step = capture.await.expect("capture task"); + let step = capture_binding(&manager).await; assert_eq!( step.tools() .iter() diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 9cad1c0251a5..27af83ef3f7b 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -240,6 +240,21 @@ impl McpRuntime { } } + /// Captures the current runtime after its selected server has finished startup. + pub async fn current_binding_for_call(&self, server: &str) -> Option> { + let current = self.current.load_full(); + let config = Arc::clone(current.config.as_ref()?); + if !current.connections.wait_for_server_startup(server).await { + return None; + } + Some(Arc::new( + current + .connections + .capture_binding_with_metadata(config, current.plugins_available) + .await, + )) + } + /// Returns the latest published configuration without waiting for clients. pub fn current_config(&self) -> Option> { self.current.load().config.clone() diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index e26192ff3f23..3ae964292345 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -141,7 +141,11 @@ pub(crate) async fn handle_mcp_tool_call( }; sess.refresh_mcp_if_dirty().await; - let current_binding = sess.services.mcp_runtime.current_binding().await; + let current_binding = sess + .services + .mcp_runtime + .current_binding_for_call(&server) + .await; let Some(prepared_call) = current_binding .as_ref() .and_then(|binding| binding.prepare_call(&server, &tool_name)) diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index 508b371b5341..ecfd7fd77203 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -248,30 +248,31 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: .await; anyhow::Ok(called_process) }); - fixture.codex.shutdown_and_wait().await?; - fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) - .await?; tokio::time::timeout(Duration::from_secs(2), async { while cached_response.requests().is_empty() { tokio::time::sleep(Duration::from_millis(10)).await; } }) .await - .context("live MCP definitions should reach inference after initialization")?; + .context("cached MCP definitions should reach inference before initialization")?; assert_definition( &cached_response, - &format!("Use the tools from {second_process}."), - &format!("Echo from {second_process}."), + &format!("Tools in the {NAMESPACE} namespace."), + &format!("Echo from {first_process}."), ); + fixture.codex.shutdown_and_wait().await?; + fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) + .await?; + let expected_error = format!("MCP tool `{SERVER_NAME}/cwd` is not available to the model"); assert_eq!(cached_turn.await??, second_process); let output = cached_done_response .single_request() .function_call_output_text(app_only_call_id) .expect("app-only tool error should be returned to the model"); assert!( - output.contains("is not available to the model") || output.contains("unsupported call"), - "app-only tools must be rejected before reaching the MCP server: {output}" + output.contains(&expected_error), + "model-visible tool output should contain the live visibility error: {output}" ); let output = cached_done_response .single_request()