diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 5558916e5..5b7b6b3e8 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -68,6 +68,7 @@ use serde_json::Value as JsonValue; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::Instrument; +use tracing::info_span; use tracing::instrument; use tracing::trace; use tracing::trace_span; @@ -105,6 +106,7 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { pub struct McpConnectionManager { clients: HashMap, server_metadata: HashMap, + required_servers: Vec, tool_plugin_provenance: Arc, host_owned_codex_apps_enabled: bool, prefix_mcp_tool_names: bool, @@ -113,7 +115,7 @@ pub struct McpConnectionManager { } impl McpConnectionManager { - #[allow(clippy::new_ret_no_self, clippy::too_many_arguments)] + #[allow(clippy::too_many_arguments)] pub async fn new( mcp_servers: &HashMap, store_mode: OAuthCredentialsStoreMode, @@ -121,6 +123,7 @@ impl McpConnectionManager { approval_policy: &Constrained, submit_id: String, tx_event: Sender, + startup_cancellation_token: CancellationToken, initial_permission_profile: PermissionProfile, runtime_context: McpRuntimeContext, codex_home: PathBuf, @@ -131,8 +134,13 @@ impl McpConnectionManager { tool_plugin_provenance: ToolPluginProvenance, auth: Option<&CodexAuth>, elicitation_reviewer: Option, - ) -> (Self, CancellationToken) { - let cancel_token = CancellationToken::new(); + ) -> Self { + let mut required_servers = mcp_servers + .iter() + .filter(|(_, server)| server.enabled() && server.required()) + .map(|(name, _)| name.clone()) + .collect::>(); + required_servers.sort(); let mut clients = HashMap::new(); let mut server_metadata = HashMap::new(); let mut join_set = JoinSet::new(); @@ -152,7 +160,7 @@ impl McpConnectionManager { .filter(|(_, server)| server.enabled()) { server_metadata.insert(server_name.clone(), McpServerMetadata::from(&server)); - let cancel_token = cancel_token.child_token(); + let cancel_token = startup_cancellation_token.child_token(); let _ = emit_update( startup_submit_id.as_str(), &tx_event, @@ -237,11 +245,12 @@ impl McpConnectionManager { let manager = Self { clients, server_metadata, + required_servers, tool_plugin_provenance, host_owned_codex_apps_enabled, prefix_mcp_tool_names, elicitation_requests: elicitation_requests.clone(), - startup_cancellation_token: cancel_token.clone(), + startup_cancellation_token: startup_cancellation_token.clone(), }; tokio::spawn(async move { let outcomes = join_set.join_all().await; @@ -265,7 +274,53 @@ impl McpConnectionManager { }) .await; }); - (manager, cancel_token) + manager + } + + /// Waits for every required server and reports their startup failures together. + /// + /// Callers must make the manager reachable to request handlers before awaiting this method, + /// because server initialization may require client elicitation. + pub async fn validate_required_servers(&self) -> Result<()> { + let failures = async { + let mut failures = Vec::new(); + for server_name in &self.required_servers { + let Some(async_managed_client) = self.clients.get(server_name).cloned() else { + failures.push(McpStartupFailure { + server: server_name.clone(), + error: format!("required MCP server `{server_name}` was not initialized"), + }); + continue; + }; + + match async_managed_client.client().await { + Ok(_) => {} + Err(error) => failures.push(McpStartupFailure { + server: server_name.clone(), + error: startup_outcome_error_message(error), + }), + } + } + failures + } + .instrument(info_span!( + "session_init.required_mcp_wait", + otel.name = "session_init.required_mcp_wait", + session_init.required_mcp_server_count = self.required_servers.len(), + )) + .await; + if failures.is_empty() { + return Ok(()); + } + + let details = failures + .iter() + .map(|failure| format!("{}: {}", failure.server, failure.error)) + .collect::>() + .join("; "); + Err(anyhow!( + "required MCP servers failed to initialize: {details}" + )) } pub fn new_uninitialized_with_permission_profile( @@ -276,6 +331,7 @@ impl McpConnectionManager { Self { clients: HashMap::new(), server_metadata: HashMap::new(), + required_servers: Vec::new(), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), host_owned_codex_apps_enabled: false, prefix_mcp_tool_names, @@ -374,31 +430,6 @@ impl McpConnectionManager { } } - pub async fn required_startup_failures( - &self, - required_servers: &[String], - ) -> Vec { - let mut failures = Vec::new(); - for server_name in required_servers { - let Some(async_managed_client) = self.clients.get(server_name).cloned() else { - failures.push(McpStartupFailure { - server: server_name.clone(), - error: format!("required MCP server `{server_name}` was not initialized"), - }); - continue; - }; - - match async_managed_client.client().await { - Ok(_) => {} - Err(error) => failures.push(McpStartupFailure { - server: server_name.clone(), - error: startup_outcome_error_message(error), - }), - } - } - failures - } - /// 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 { diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 60a0026ee..4b90337ac 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1167,13 +1167,15 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ), ]); - let (manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let manager = McpConnectionManager::new( &mcp_servers, OAuthCredentialsStoreMode::default(), HashMap::new(), &approval_policy, String::new(), tx_event, + cancel_token.clone(), PermissionProfile::default(), McpRuntimeContext::new( Arc::new(EnvironmentManager::without_environments()), @@ -1201,13 +1203,18 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { .wait_for_server_ready("stdio", Duration::from_millis(10)) .await ); - let failures = manager - .required_startup_failures(&["stdio".to_string()]) - .await; - assert_eq!(failures.len(), 1); - assert_eq!(failures[0].server, "stdio"); + let error = match manager + .clients + .get("stdio") + .expect("stdio client") + .client() + .await + { + Ok(_) => panic!("local stdio MCP startup should fail"), + Err(error) => error, + }; assert_eq!( - failures[0].error, + startup_outcome_error_message(error), "local stdio MCP server `stdio` requires a local environment" ); cancel_token.cancel(); diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index a74f95bae..adc2cf615 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -35,6 +35,7 @@ use rmcp::model::ElicitationCapability; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use serde_json::Value; +use tokio_util::sync::CancellationToken; use crate::codex_apps::codex_apps_tools_cache_key; use crate::connection_manager::McpConnectionManager; @@ -293,13 +294,15 @@ pub async fn read_mcp_resource( .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_statuses, &config.approval_policy, String::new(), tx_event, + cancel_token.clone(), PermissionProfile::default(), runtime_context, config.codex_home.clone(), @@ -363,13 +366,15 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_status_entries.clone(), &config.approval_policy, submit_id, tx_event, + cancel_token.clone(), PermissionProfile::default(), runtime_context, config.codex_home.clone(), diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 623cb4421..6d6ca52d8 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -17,6 +17,7 @@ use codex_protocol::models::PermissionProfile; use codex_tools::DiscoverableTool; use rmcp::model::ToolAnnotations; use serde::Deserialize; +use tokio_util::sync::CancellationToken; use tracing::warn; use crate::config::Config; @@ -285,13 +286,15 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (mut mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let cancel_token = CancellationToken::new(); + let mut mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, auth_status_entries, &config.permissions.approval_policy, INITIAL_SUBMIT_ID.to_owned(), tx_event, + cancel_token.clone(), PermissionProfile::default(), // Connector discovery is threadless. Use an actually configured env if // one exists, but do not reintroduce the old hidden-local fallback. diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index bd71c3a05..da54fe187 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -46,6 +46,7 @@ use std::path::Path; use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; +use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::Level; use tracing_subscriber::fmt::format::FmtSpan; @@ -1256,13 +1257,14 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) { let auth = session.services.auth_manager.auth().await; - let (manager, _cancel_token) = codex_mcp::McpConnectionManager::new( + let manager = codex_mcp::McpConnectionManager::new( &HashMap::new(), turn_context.config.mcp_oauth_credentials_store_mode, HashMap::new(), &turn_context.approval_policy, turn_context.sub_id.clone(), session.get_tx_event(), + CancellationToken::new(), turn_context.permission_profile(), codex_mcp::McpRuntimeContext::new(Arc::clone(&session.services.environment_manager), { #[allow(deprecated)] diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index f2643d0e5..779813641 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -338,18 +338,21 @@ impl Session { turn_context.cwd.to_path_buf(), ), }; - { + let mcp_startup_cancellation_token = { let mut guard = self.services.mcp_startup_cancellation_token.lock().await; guard.cancel(); - *guard = CancellationToken::new(); - } - let (refreshed_manager, cancel_token) = McpConnectionManager::new( + let cancellation_token = CancellationToken::new(); + *guard = cancellation_token.clone(); + cancellation_token + }; + let refreshed_manager = McpConnectionManager::new( &mcp_servers, store_mode, auth_statuses, &turn_context.approval_policy, turn_context.sub_id.clone(), self.get_tx_event(), + mcp_startup_cancellation_token, turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), @@ -366,14 +369,6 @@ impl Session { let current_manager = self.services.mcp_connection_manager.read().await; refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); } - { - let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - if guard.is_cancelled() { - cancel_token.cancel(); - } - *guard = cancel_token; - } - let mut old_manager = { let mut manager = self.services.mcp_connection_manager.write().await; std::mem::replace(&mut *manager, refreshed_manager) diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 50ff98423..98d19957c 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -469,10 +469,6 @@ impl Session { #[instrument(name = "session_init", level = "info", skip_all)] #[allow(clippy::too_many_arguments)] - #[expect( - clippy::await_holding_invalid_type, - reason = "session initialization must serialize access through session-owned manager guards" - )] pub(crate) async fn new( mut session_configuration: SessionConfiguration, config: Arc, @@ -1116,15 +1112,6 @@ impl Session { sess.send_event_raw(event).await; } - let mut required_mcp_servers: Vec = mcp_servers - .iter() - .filter(|(_, server)| server.enabled() && server.required()) - .map(|(name, _)| name.clone()) - .collect(); - required_mcp_servers.sort(); - let enabled_mcp_server_count = - mcp_servers.values().filter(|server| server.enabled()).count(); - let required_mcp_server_count = required_mcp_servers.len(); let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config.as_ref()).await; let host_owned_codex_apps_enabled = config .features @@ -1137,11 +1124,13 @@ impl Session { } else { ElicitationCapability::default() }; - { + let mcp_startup_cancellation_token = { let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; cancel_guard.cancel(); - *cancel_guard = CancellationToken::new(); - } + let cancel_token = CancellationToken::new(); + *cancel_guard = cancel_token.clone(); + cancel_token + }; let turn_environment = crate::environment_selection::resolve_environment_selections( sess.services.environment_manager.as_ref(), session_configuration.environment_selections(), @@ -1164,13 +1153,14 @@ impl Session { session_configuration.cwd().to_path_buf(), ), }; - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( + let mcp_connection_manager = McpConnectionManager::new( &mcp_servers, config.mcp_oauth_credentials_store_mode, - auth_statuses.clone(), + auth_statuses, &session_configuration.approval_policy, INITIAL_SUBMIT_ID.to_owned(), tx_event.clone(), + mcp_startup_cancellation_token, session_configuration.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), @@ -1185,43 +1175,11 @@ impl Session { .instrument(info_span!( "session_init.mcp_manager_init", otel.name = "session_init.mcp_manager_init", - session_init.enabled_mcp_server_count = enabled_mcp_server_count, - session_init.required_mcp_server_count = required_mcp_server_count, )) .await; - { - let mut manager_guard = sess.services.mcp_connection_manager.write().await; - *manager_guard = mcp_connection_manager; - } - { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - if cancel_guard.is_cancelled() { - cancel_token.cancel(); - } - *cancel_guard = cancel_token; - } - if !required_mcp_servers.is_empty() { - let failures = sess - .services - .mcp_connection_manager - .read() - .await - .required_startup_failures(&required_mcp_servers) - .instrument(info_span!( - "session_init.required_mcp_wait", - otel.name = "session_init.required_mcp_wait", - session_init.required_mcp_server_count = required_mcp_server_count, - )) - .await; - if !failures.is_empty() { - let details = failures - .iter() - .map(|failure| format!("{}: {}", failure.server, failure.error)) - .collect::>() - .join("; "); - anyhow::bail!("required MCP servers failed to initialize: {details}"); - } - } + sess.services + .install_mcp_connection_manager(mcp_connection_manager) + .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; let session_start_source = match &initial_history { diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 59def9c7a..9bc4d7391 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -15,6 +15,7 @@ use crate::tools::code_mode::CodeModeService; use crate::tools::network_approval::NetworkApprovalService; use crate::tools::sandboxing::ApprovalStore; use crate::unified_exec::UnifiedExecProcessManager; +use anyhow::Result; use arc_swap::ArcSwap; use arc_swap::ArcSwapOption; use codex_analytics::AnalyticsEventsClient; @@ -82,3 +83,23 @@ pub(crate) struct SessionServices { /// the same manager through child-thread spawn paths without reconstructing it. pub(crate) environment_manager: Arc, } + +impl SessionServices { + /// Installs the manager before validating required servers so startup-time elicitation can + /// resolve through the session's manager while validation waits. + #[expect( + clippy::await_holding_invalid_type, + reason = "required MCP validation keeps the installed manager reachable for startup-time elicitation" + )] + pub(crate) async fn install_mcp_connection_manager( + &self, + manager: McpConnectionManager, + ) -> Result<()> { + *self.mcp_connection_manager.write().await = manager; + self.mcp_connection_manager + .read() + .await + .validate_required_servers() + .await + } +}