diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index c678d708a6..c0acaf1b03 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -248,6 +248,15 @@ impl McpRuntime { } } + /// Waits for the selected server without capturing an execution binding. + pub async fn wait_for_server_startup(&self, server: &str) { + self.current + .load_full() + .connections + .wait_for_server_startup(server) + .await; + } + /// 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(); diff --git a/codex-rs/core/src/tools/handlers/mcp.rs b/codex-rs/core/src/tools/handlers/mcp.rs index 8515920290..1b1ecf361a 100644 --- a/codex-rs/core/src/tools/handlers/mcp.rs +++ b/codex-rs/core/src/tools/handlers/mcp.rs @@ -163,6 +163,10 @@ impl McpHandler { } impl CoreToolRuntime for McpHandler { + fn mcp_server_name(&self) -> Option<&str> { + Some(&self.tool_info.server_name) + } + fn telemetry_tags<'a>( &'a self, _invocation: &'a ToolInvocation, diff --git a/codex-rs/core/src/tools/parallel.rs b/codex-rs/core/src/tools/parallel.rs index 7bb590d334..deb0ab85ac 100644 --- a/codex-rs/core/src/tools/parallel.rs +++ b/codex-rs/core/src/tools/parallel.rs @@ -98,6 +98,7 @@ impl ToolCallRuntime { cancellation_token: CancellationToken, ) -> impl std::future::Future> { let supports_parallel = self.router.tool_supports_parallel(&call); + let mcp_server = self.router.mcp_server_name(&call).map(str::to_owned); let router = Arc::clone(&self.router); let session = Arc::clone(&self.session); let step_context = Arc::clone(&self.step_context); @@ -130,6 +131,15 @@ impl ToolCallRuntime { let mut dispatch_handle: AbortOnDropHandle> = AbortOnDropHandle::new(tokio::spawn(async move { + if let Some(server) = mcp_server { + session.refresh_mcp_if_dirty().await; + session + .services + .mcp_runtime + .wait_for_server_startup(&server) + .await; + } + let _guard = if supports_parallel { Either::Left(lock.read().await) } else { diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index fc6a3aaba3..3927941abb 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -51,6 +51,11 @@ pub use codex_tools::ToolExposure; /// Implementers provide the shared `ToolExecutor` behavior plus optional /// core-owned metadata for hooks, telemetry, tool search, and argument diffs. pub(crate) trait CoreToolRuntime: ToolExecutor { + /// Returns the MCP server that owns this runtime, when this is an MCP tool. + fn mcp_server_name(&self) -> Option<&str> { + None + } + fn matches_kind(&self, payload: &ToolPayload) -> bool { matches!( payload, @@ -283,6 +288,10 @@ impl ToolExecutor for ExposureOverride { } impl CoreToolRuntime for ExposureOverride { + fn mcp_server_name(&self) -> Option<&str> { + self.handler.mcp_server_name() + } + fn matches_kind(&self, payload: &ToolPayload) -> bool { self.handler.matches_kind(payload) } @@ -392,6 +401,10 @@ impl ToolRegistry { self.tools.get(name).map(Arc::clone) } + pub(crate) fn mcp_server_name(&self, name: &ToolName) -> Option<&str> { + self.tools.get(name)?.mcp_server_name() + } + #[cfg(test)] pub(crate) fn tool_names_for_test(&self) -> Vec { let mut names = self.tools.keys().cloned().collect::>(); diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index 5e48e65d23..73c98dec11 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/codex-rs/core/src/tools/router.rs @@ -149,6 +149,10 @@ impl ToolRouter { .unwrap_or(false) } + pub(crate) fn mcp_server_name(&self, call: &ToolCall) -> Option<&str> { + self.registry.mcp_server_name(&call.tool_name) + } + pub fn tool_waits_for_runtime_cancellation(&self, call: &ToolCall) -> bool { self.registry .waits_for_runtime_cancellation(&call.tool_name) diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index bd33221736..92625f0ea1 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -7,6 +7,8 @@ use crate::session::tests::make_session_and_context; use crate::tools::context::ToolPayload; use crate::tools::handlers::McpHandler; use crate::tools::registry::CoreToolRuntime; +use crate::tools::registry::ToolExposure; +use crate::tools::registry::override_tool_exposure; use crate::turn_diff_tracker::TurnDiffTracker; use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionRegistry; @@ -156,6 +158,17 @@ async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow }) .expect("test session should expose a parallel shell-like tool"); + assert_eq!( + router.mcp_server_name(&ToolCall { + tool_name: ToolName::plain(parallel_tool_name), + call_id: "call-local-tool".to_string(), + payload: ToolPayload::Function { + arguments: "{}".to_string(), + }, + }), + None + ); + assert!(!router.tool_supports_parallel(&ToolCall { tool_name: ToolName::namespaced("mcp__server__", parallel_tool_name), call_id: "call-namespaced-tool".to_string(), @@ -248,12 +261,15 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { "mcp__echo__", "query_with_delay", )), - mcp_runtime(mcp_tool_info( - "hello_echo", - /*supports_parallel_tool_calls*/ false, - "mcp__hello_echo__", - "query_with_delay", - )), + override_tool_exposure( + mcp_runtime(mcp_tool_info( + "hello_echo", + /*supports_parallel_tool_calls*/ false, + "mcp__hello_echo__", + "query_with_delay", + )), + ToolExposure::DirectModelOnly, + ), ], extension_tool_executors: Vec::new(), wait_for_environment_tool_config: None, @@ -271,6 +287,7 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { encrypted_function_args: None, }; assert!(router.tool_supports_parallel(&call)); + assert_eq!(router.mcp_server_name(&call), Some("echo")); let different_server_call = ToolCall { tool_name: ToolName::namespaced("mcp__hello_echo__", "query_with_delay"), @@ -281,6 +298,10 @@ async fn mcp_parallel_support_uses_handler_data() -> anyhow::Result<()> { encrypted_function_args: None, }; assert!(!router.tool_supports_parallel(&different_server_call)); + assert_eq!( + router.mcp_server_name(&different_server_call), + Some("hello_echo") + ); Ok(()) } diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index ecfd7fd772..1b3daf29be 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -196,6 +196,7 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: let second_process = process_label(&second_pid); let app_only_call_id = "cached-app-only-call"; + let unrelated_call_id = "cached-unrelated-plan"; let cached_response = mount_sse_once( &responses_server, responses::sse(vec![ @@ -207,6 +208,11 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: r#"{"message":"hello"}"#, ), responses::ev_function_call_with_namespace(app_only_call_id, NAMESPACE, "cwd", "{}"), + responses::ev_function_call( + unrelated_call_id, + "update_plan", + r#"{"plan":[{"step":"Continue while MCP starts","status":"in_progress"}]}"#, + ), responses::ev_completed("cached-call"), ]), ) @@ -220,12 +226,19 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: ]), ) .await; + let (unrelated_finished_tx, unrelated_finished_rx) = tokio::sync::oneshot::channel(); let second_for_turn = Arc::clone(&second_thread); let cached_turn = tokio::spawn(async move { second_for_turn .submit(user_turn("call the echo and cwd tools")) .await?; + let mut unrelated_finished_tx = Some(unrelated_finished_tx); let end = wait_for_event(&second_for_turn, |event| { + if matches!(event, EventMsg::PlanUpdate(_)) + && let Some(sender) = unrelated_finished_tx.take() + { + let _ = sender.send(()); + } matches!( event, EventMsg::McpToolCallEnd(end) if end.call_id == "cached-call" @@ -260,6 +273,10 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: &format!("Tools in the {NAMESPACE} namespace."), &format!("Echo from {first_process}."), ); + tokio::time::timeout(Duration::from_secs(2), unrelated_finished_rx) + .await + .context("an unrelated tool should complete while cached MCP startup is pending")? + .context("the unrelated tool should emit its plan update")?; fixture.codex.shutdown_and_wait().await?; fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) @@ -282,6 +299,13 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: output.contains(&second_process), "model-visible tool output should come from the live server: {output}" ); + assert_eq!( + cached_done_response + .single_request() + .function_call_output_text(unrelated_call_id) + .as_deref(), + Some("Plan updated") + ); second_thread.shutdown_and_wait().await?; responses_server.verify().await;