diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 52f31876d0..359e780dc0 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -393,7 +393,7 @@ impl McpConnectionManager { pub fn wait_for_servers_ready( &self, server_names: &[String], - ) -> impl std::future::Future> + Send + 'static { + ) -> impl std::future::Future> + Send + 'static + use<> { let clients = server_names .iter() .map(|server_name| (server_name.clone(), self.clients.get(server_name).cloned())) @@ -419,6 +419,17 @@ impl McpConnectionManager { } } + /// Returns an owned future that waits for every configured MCP server. + pub fn wait_for_all_servers_ready( + &self, + ) -> impl std::future::Future, Vec)> + Send + 'static + { + let mut server_names = self.clients.keys().cloned().collect::>(); + server_names.sort(); + let readiness = self.wait_for_servers_ready(&server_names); + async move { (server_names, readiness.await) } + } + pub async fn required_startup_failures( &self, required_servers: &[String], diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 7396e0e4bf..31faa4cc5e 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -144,10 +144,7 @@ impl CodexThread { } #[doc(hidden)] - pub async fn wait_for_mcp_servers_ready( - &self, - server_names: &[String], - ) -> Vec { + pub async fn wait_for_mcp_startup(&self) -> (Vec, Vec) { let readiness = { let manager = self .codex @@ -156,7 +153,7 @@ impl CodexThread { .mcp_connection_manager .read() .await; - manager.wait_for_servers_ready(server_names) + manager.wait_for_all_servers_ready() }; readiness.await } diff --git a/codex-rs/core/tests/common/lib.rs b/codex-rs/core/tests/common/lib.rs index e468a06ae6..565fddc089 100644 --- a/codex-rs/core/tests/common/lib.rs +++ b/codex-rs/core/tests/common/lib.rs @@ -246,9 +246,14 @@ where /// Waits for a configured MCP server to finish startup and requires it to be ready. pub async fn wait_for_mcp_server(codex: &CodexThread, server_name: &str) -> anyhow::Result<()> { - let server_names = vec![server_name.to_string()]; - let failures = codex.wait_for_mcp_servers_ready(&server_names).await; - if let Some(failure) = failures.first() { + let (server_names, failures) = codex.wait_for_mcp_startup().await; + if !server_names.iter().any(|name| name == server_name) { + anyhow::bail!("MCP server {server_name} was not initialized"); + } + if let Some(failure) = failures + .iter() + .find(|failure| failure.server == server_name) + { let error = &failure.error; anyhow::bail!("MCP server {server_name} failed to start: {error}"); }