From ff17e5d9b381b26f8bc05d845684c13e38d8a2cc Mon Sep 17 00:00:00 2001 From: bromano Date: Mon, 15 Jun 2026 23:24:50 +0000 Subject: [PATCH] [codex-mcp] Avoid blocking on optional MCP startup --- codex-rs/codex-mcp/src/connection_manager.rs | 68 +++++---- .../codex-mcp/src/connection_manager_tests.rs | 144 +++++++++++++++++- 2 files changed, 182 insertions(+), 30 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 359648656807..81edc4687c80 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -446,41 +446,51 @@ impl McpConnectionManager { /// Returns all tools with model-visible names normalized. #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] pub async fn list_all_tools(&self) -> Vec { - let mut tools = Vec::new(); - for (server_name, managed_client) in &self.clients { - let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some(); - let startup_complete = managed_client - .startup_complete - .load(std::sync::atomic::Ordering::Acquire); - trace!( - server_name = %server_name, - has_cached_tool_info_snapshot, - startup_complete, - "waiting for MCP server tools while building tool list" - ); - let Some(server_tools) = managed_client - .listed_tools() + let tools: Vec = + futures::future::join_all(self.clients.iter().map(|(server_name, managed_client)| { + let required = self.required_servers.binary_search(server_name).is_ok(); + let has_cached_tool_info_snapshot = + managed_client.cached_tool_info_snapshot.is_some(); + let startup_complete = managed_client.startup_complete.load(Ordering::Acquire); + async move { + trace!( + server_name = %server_name, + required, + has_cached_tool_info_snapshot, + startup_complete, + "waiting for MCP server tools while building tool list" + ); + if !required && !startup_complete && !has_cached_tool_info_snapshot { + return None; + } + if required && !startup_complete && has_cached_tool_info_snapshot { + let _ = managed_client.client().await; + } + let server_tools = managed_client.listed_tools().await; + if let Some(server_tools) = &server_tools { + trace!( + server_name = %server_name, + required, + tool_count = server_tools.len(), + "listed MCP server tools while building tool list" + ); + } + server_tools + } .instrument(trace_span!( "list_tools_for_server", server_name = %server_name, + required, has_cached_tool_info_snapshot, startup_complete )) - .await - else { - continue; - }; - trace!( - server_name = %server_name, - tool_count = server_tools.len(), - "listed MCP server tools while building tool list" - ); - tools.extend( - server_tools - .into_iter() - .map(|tool| self.with_server_metadata(tool)), - ); - } + })) + .await + .into_iter() + .flatten() + .flatten() + .map(|tool| self.with_server_metadata(tool)) + .collect(); normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names) } diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index bfbc09c473c0..27feb589fe53 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -833,6 +833,76 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() assert_eq!(tool.callable_name, "calendar_create_event"); } +#[tokio::test] +async fn list_all_tools_polls_clients_concurrently() { + let barrier = Arc::new(tokio::sync::Barrier::new(3)); + let (first_started_tx, first_started_rx) = tokio::sync::oneshot::channel(); + let first_barrier = Arc::clone(&barrier); + let first_pending_client = async move { + let _ = first_started_tx.send(()); + first_barrier.wait().await; + Err(StartupOutcomeError::Cancelled) + } + .boxed() + .shared(); + let (second_started_tx, second_started_rx) = tokio::sync::oneshot::channel(); + let second_barrier = Arc::clone(&barrier); + let second_pending_client = async move { + let _ = second_started_tx.send(()); + second_barrier.wait().await; + Err(StartupOutcomeError::Cancelled) + } + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.required_servers = vec!["first".to_string(), "second".to_string()]; + manager.clients.insert( + "first".to_string(), + AsyncManagedClient { + client: first_pending_client, + cached_tool_info_snapshot: None, + cached_server_info: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + manager.clients.insert( + "second".to_string(), + AsyncManagedClient { + client: second_pending_client, + cached_tool_info_snapshot: None, + cached_server_info: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + let list_task = tokio::spawn(async move { manager.list_all_tools().await }); + + let (first_started, second_started) = + tokio::time::timeout(Duration::from_secs(1), async move { + tokio::join!(first_started_rx, second_started_rx) + }) + .await + .expect("all client futures should be polled before either is released"); + first_started.expect("first client future should start"); + second_started.expect("second client future should start"); + barrier.wait().await; + + let tools = tokio::time::timeout(Duration::from_secs(1), list_task) + .await + .expect("tool listing should finish after clients are released") + .expect("tool listing task should not panic"); + assert!(tools.is_empty()); +} + #[tokio::test] async fn list_available_server_infos_uses_cache_while_client_is_pending() { let pending_client = futures::future::pending::>() @@ -957,7 +1027,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { } #[tokio::test] -async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_snapshot() { +async fn list_all_tools_skips_optional_client_pending_without_cached_tool_info_snapshot() { let pending_client = futures::future::pending::>() .boxed() .shared(); @@ -980,6 +1050,75 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_ }, ); + let timeout_result = + tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await; + let tools = timeout_result.expect("optional client tool listing should not block on startup"); + assert!(tools.is_empty()); +} + +#[tokio::test] +async fn list_all_tools_blocks_while_required_client_is_pending_without_cached_tool_info_snapshot() +{ + let pending_client = futures::future::pending::>() + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager + .required_servers + .push(CODEX_APPS_MCP_SERVER_NAME.to_string()); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: pending_client, + cached_tool_info_snapshot: None, + cached_server_info: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + + let timeout_result = + tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await; + assert!(timeout_result.is_err()); +} + +#[tokio::test] +async fn list_all_tools_blocks_while_required_client_is_pending_with_cached_tool_info_snapshot() { + let pending_client = futures::future::pending::>() + .boxed() + .shared(); + let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager + .required_servers + .push(CODEX_APPS_MCP_SERVER_NAME.to_string()); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client: pending_client, + cached_tool_info_snapshot: Some(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + )]), + cached_server_info: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + let timeout_result = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()).await; assert!(timeout_result.is_err()); @@ -1004,6 +1143,9 @@ async fn shutdown_cancels_pending_tool_listing() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); + manager + .required_servers + .push(CODEX_APPS_MCP_SERVER_NAME.to_string()); manager.clients.insert( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient {