diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 3107f678b..c76e3c3c0 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -208,6 +208,7 @@ impl McpConnectionManager { }; let async_managed_client = AsyncManagedClient::new( server_name.clone(), + startup_submit_id.clone(), server, store_mode, keyring_backend_kind, @@ -257,6 +258,10 @@ impl McpConnectionManager { ) .await; + if matches!(&outcome, Err(StartupOutcomeError::Failed { .. })) { + async_managed_client.reconnect_failed_startup().await; + } + (server_name, outcome) }); } @@ -479,6 +484,7 @@ impl McpConnectionManager { pub async fn list_all_tools(&self) -> Vec { let mut tools = Vec::new(); for (server_name, managed_client) in &self.clients { + managed_client.reconnect_failed_startup().await; let has_cached_tools = managed_client.has_cached_tools(); let startup_complete = managed_client .startup_complete diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index d2869fe5c..1f61ab700 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -6,7 +6,10 @@ use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::elicitation_is_rejected_by_policy; use crate::rmcp_client::AsyncManagedClient; +use crate::rmcp_client::CODEX_APPS_RECONNECT_INITIAL_BACKOFF; +use crate::rmcp_client::CodexAppsStartupReconnect; use crate::rmcp_client::ManagedClient; +use crate::rmcp_client::ManagedClientFuture; use crate::rmcp_client::StartupOutcomeError; use crate::server::EffectiveMcpServer; use crate::server::McpServerMetadata; @@ -21,6 +24,7 @@ use codex_config::Constrained; use codex_config::McpServerConfig; use codex_config::McpServerToolConfig; use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::EnvironmentManager; use codex_protocol::ToolName; use codex_protocol::mcp::McpServerInfo; @@ -41,6 +45,7 @@ use rmcp::model::Tool; use std::collections::HashSet; use std::io; use std::sync::Arc; +use std::sync::atomic::AtomicUsize; use tempfile::tempdir; use tokio::io::DuplexStream; @@ -101,9 +106,8 @@ impl InProcessTransportFactory for TestInProcessTransportFactory { } } -async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { - let tool_filter = ToolFilter::default(); - let managed_client = ManagedClient { +async fn create_test_managed_client(tools: Vec) -> ManagedClient { + ManagedClient { client: Arc::new( RmcpClient::new_in_process_client(Arc::new(TestInProcessTransportFactory)) .await @@ -111,29 +115,73 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged ), server_info: create_test_server_info("Ready"), tools, - tool_filter: tool_filter.clone(), + tool_filter: ToolFilter::default(), tool_timeout: None, server_instructions: None, server_supports_sandbox_state_meta_capability: false, codex_apps_tools_cache_context: None, - }; + } +} +async fn create_ready_async_managed_client(tools: Vec) -> AsyncManagedClient { AsyncManagedClient { client: futures::future::ready::>(Ok( - managed_client, + create_test_managed_client(tools).await, )) .boxed() .shared(), is_codex_apps_mcp_server: false, cached_server_info: None, codex_apps_tools_cache_context: None, - tool_filter, + tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), } } +fn create_test_manager_with_failed_apps_startup( + cached_tools: Vec, + reconnect_factory: Arc ManagedClientFuture + Send + Sync>, +) -> McpConnectionManager { + let client: ManagedClientFuture = futures::future::ready(Err(StartupOutcomeError::Failed { + error: "startup failed".to_string(), + is_authentication_required: false, + })) + .boxed() + .shared(); + let codex_home = tempdir().expect("tempdir"); + let cache_context = create_codex_apps_tools_cache_context( + codex_home.path().to_path_buf(), + Some("reconnect-test-account"), + Some("reconnect-test-user"), + ); + cache_context.store_current_tools_for_test(cached_tools); + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionManager::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + manager.clients.insert( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + AsyncManagedClient { + client, + is_codex_apps_mcp_server: true, + cached_server_info: None, + codex_apps_tools_cache_context: Some(cache_context), + tool_filter: ToolFilter::default(), + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + cancel_token: CancellationToken::new(), + }, + ); + manager +} + fn model_tool_names(tools: &[ToolInfo]) -> HashSet { tools .iter() @@ -721,6 +769,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -760,6 +809,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -860,6 +910,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -898,6 +949,7 @@ async fn shutdown_cancels_pending_tool_listing() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, }, @@ -944,6 +996,7 @@ async fn shutdown_continues_after_caller_is_aborted() { codex_apps_tools_cache_context: None, tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -996,6 +1049,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -1045,6 +1099,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() codex_apps_tools_cache_context: Some(cache_context), tool_filter: ToolFilter::default(), startup_complete, + startup_reconnect: None, tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, @@ -1069,6 +1124,235 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() ); } +#[tokio::test] +async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() { + let recovered_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "drive_search", + )]) + .await; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_reconnect = Arc::clone(&attempts); + let reconnect_finished = Arc::new(tokio::sync::Notify::new()); + let reconnect_finished_for_factory = Arc::clone(&reconnect_finished); + let reconnect_factory = Arc::new(move || { + attempts_for_reconnect.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let reconnect_finished = Arc::clone(&reconnect_finished_for_factory); + let recovered_client = recovered_client.clone(); + async move { + reconnect_finished.notify_one(); + Ok(recovered_client) + } + .boxed() + .shared() + }); + let manager = create_test_manager_with_failed_apps_startup(Vec::new(), reconnect_factory); + + let reconnect_finished_wait = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert!(tools.is_empty()); + reconnect_finished_wait.await; + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); +} + +#[tokio::test(start_paused = true)] +async fn later_tool_list_retries_after_failed_reconnect_and_keeps_cached_tools() { + let recovered_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "drive_search", + )]) + .await; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_reconnect = Arc::clone(&attempts); + let reconnect_finished = Arc::new(tokio::sync::Notify::new()); + let reconnect_finished_for_factory = Arc::clone(&reconnect_finished); + let reconnect_factory = Arc::new(move || { + let attempt = attempts_for_reconnect.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let reconnect_finished = Arc::clone(&reconnect_finished_for_factory); + let recovered_client = recovered_client.clone(); + async move { + let result = if attempt < 2 { + Err(StartupOutcomeError::Failed { + error: "recreated startup failed".to_string(), + is_authentication_required: false, + }) + } else { + Ok(recovered_client) + }; + reconnect_finished.notify_one(); + result + } + .boxed() + .shared() + }); + let manager = create_test_manager_with_failed_apps_startup( + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "cached_drive_search", + )], + reconnect_factory, + ); + + let first_reconnect_finished = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + first_reconnect_finished.await; + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + + tokio::time::advance(CODEX_APPS_RECONNECT_INITIAL_BACKOFF).await; + let second_reconnect_finished = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + second_reconnect_finished.await; + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(CODEX_APPS_RECONNECT_INITIAL_BACKOFF).await; + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 2); + + tokio::time::advance(CODEX_APPS_RECONNECT_INITIAL_BACKOFF).await; + let third_reconnect_finished = reconnect_finished.notified(); + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + third_reconnect_finished.await; + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 3); + + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); +} + +#[tokio::test] +async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() { + let recovered_client = create_test_managed_client(vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "drive_search", + )]) + .await; + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_for_reconnect = Arc::clone(&attempts); + let reconnect_started = Arc::new(tokio::sync::Notify::new()); + let reconnect_started_for_factory = Arc::clone(&reconnect_started); + let release_reconnect = Arc::new(tokio::sync::Notify::new()); + let release_reconnect_for_factory = Arc::clone(&release_reconnect); + let reconnect_factory = Arc::new(move || { + let recovered_client = recovered_client.clone(); + let attempts = Arc::clone(&attempts_for_reconnect); + let reconnect_started = Arc::clone(&reconnect_started_for_factory); + let release_reconnect = Arc::clone(&release_reconnect_for_factory); + async move { + attempts.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + reconnect_started.notify_one(); + release_reconnect.notified().await; + Ok(recovered_client) + } + .boxed() + .shared() + }); + let manager = Arc::new(create_test_manager_with_failed_apps_startup( + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "cached_drive_search", + )], + reconnect_factory, + )); + let reconnect_started_wait = reconnect_started.notified(); + let first_tools = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()) + .await + .expect("cached tools should not wait for reconnect"); + + reconnect_started_wait.await; + let second_tools = tokio::time::timeout(Duration::from_millis(10), manager.list_all_tools()) + .await + .expect("concurrent cached tools should not wait for reconnect"); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); + assert_eq!( + first_tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + assert_eq!( + second_tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["cached_drive_search"] + ); + + release_reconnect.notify_one(); + tokio::task::yield_now().await; + let tools = manager.list_all_tools().await; + assert_eq!( + tools + .iter() + .map(|tool| tool.callable_name.as_str()) + .collect::>(), + vec!["drive_search"] + ); + assert_eq!(attempts.load(std::sync::atomic::Ordering::SeqCst), 1); +} + #[tokio::test] async fn list_all_tools_adds_server_metadata_to_tools() { let server_name = "docs"; diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index ddfb3ffdb..ee7dde1fe 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -12,6 +12,7 @@ use std::collections::HashMap; use std::env; use std::ffi::OsString; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; @@ -48,6 +49,9 @@ use codex_exec_server::HttpClient; use codex_exec_server::ReqwestHttpClient; use codex_protocol::mcp::McpServerInfo; use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::McpStartupStatus; +use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::ExecutorStdioServerLauncher; use codex_rmcp_client::LocalStdioServerLauncher; use codex_rmcp_client::RmcpClient; @@ -64,6 +68,7 @@ use rmcp::model::JsonObject; use rmcp::model::ProtocolVersion; use rmcp::model::Tool as RmcpTool; use rmcp::transport::auth::AuthError; +use tokio::time::Instant as TokioInstant; use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::instrument; @@ -80,6 +85,9 @@ pub(crate) const MCP_TOOLS_FETCH_UNCACHED_DURATION_METRIC: &str = pub(crate) const DEFAULT_STARTUP_TIMEOUT: Duration = Duration::from_secs(30); pub(crate) const DEFAULT_TOOL_TIMEOUT: Duration = Duration::from_secs(300); +pub(crate) const CODEX_APPS_RECONNECT_INITIAL_BACKOFF: Duration = Duration::from_secs(1); +const CODEX_APPS_RECONNECT_MAX_BACKOFF: Duration = Duration::from_secs(30); + const UNTRUSTED_CONNECTOR_META_KEYS: &[&str] = &[ "connector_id", "connector_name", @@ -128,56 +136,174 @@ impl ManagedClient { } } -#[derive(Clone)] -pub(crate) struct AsyncManagedClient { - pub(crate) client: Shared>>, - pub(crate) is_codex_apps_mcp_server: bool, - pub(crate) cached_server_info: Option, - pub(crate) codex_apps_tools_cache_context: Option, - pub(crate) tool_filter: ToolFilter, - pub(crate) startup_complete: Arc, - pub(crate) tool_plugin_provenance: Arc, - pub(crate) cancel_token: CancellationToken, +pub(crate) type ManagedClientFuture = + Shared>>; + +#[derive(Default)] +struct CodexAppsStartupReconnectState { + current_client: Option, + reconnect_in_flight: bool, + consecutive_failures: u32, + retry_not_before: Option, } -impl AsyncManagedClient { - // Keep this constructor flat so the startup inputs remain readable at the - // single call site instead of introducing a one-off params wrapper. - #[instrument(level = "trace", skip_all, fields(server_name = %server_name))] - #[allow(clippy::too_many_arguments)] - pub(crate) fn new( +#[derive(Clone)] +struct CodexAppsStartupStatusContext { + submit_id: String, + server_name: String, + tx_event: Sender, +} + +pub(crate) struct CodexAppsStartupReconnect { + factory: Arc ManagedClientFuture + Send + Sync>, + state: StdMutex, + startup_status_context: Option, +} + +impl CodexAppsStartupReconnect { + pub(crate) fn new(factory: Arc ManagedClientFuture + Send + Sync>) -> Self { + Self { + factory, + state: StdMutex::new(CodexAppsStartupReconnectState::default()), + startup_status_context: None, + } + } + + fn with_startup_status_context( + mut self, + submit_id: String, server_name: String, - server: EffectiveMcpServer, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, - cancel_token: CancellationToken, tx_event: Sender, - elicitation_requests: ElicitationRequestManager, - codex_apps_tools_cache_context: Option, - tool_plugin_provenance: Arc, - runtime_context: McpRuntimeContext, - runtime_auth_provider: Option, - client_elicitation_capability: ElicitationCapability, - supports_openai_form_elicitation: bool, ) -> Self { + self.startup_status_context = Some(CodexAppsStartupStatusContext { + submit_id, + server_name, + tx_event, + }); + self + } + + fn current_client(&self) -> Option { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .current_client + .clone() + } + + fn reconnect_in_background(self: &Arc) { + { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if state.current_client.is_some() || state.reconnect_in_flight { + return; + } + if state + .retry_not_before + .is_some_and(|retry_not_before| TokioInstant::now() < retry_not_before) + { + return; + } + state.reconnect_in_flight = true; + } + + let reconnect = Arc::clone(self); + tokio::spawn(async move { + let result = (reconnect.factory)().await; + let startup_status_context = reconnect.startup_status_context.clone(); + let recovered = { + let mut state = reconnect + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.reconnect_in_flight = false; + match result { + Ok(client) => { + state.current_client = Some(client); + state.consecutive_failures = 0; + state.retry_not_before = None; + true + } + Err(error) => { + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + let retry_after = codex_apps_reconnect_backoff(state.consecutive_failures); + state.retry_not_before = Some(TokioInstant::now() + retry_after); + warn!( + error = %error, + retry_after_ms = retry_after.as_millis(), + "Apps MCP startup reconnect failed; continuing with cached tools" + ); + false + } + } + }; + + if recovered && let Some(context) = startup_status_context { + let _ = context + .tx_event + .send(Event { + id: context.submit_id, + msg: EventMsg::McpStartupUpdate(McpStartupUpdateEvent { + server: context.server_name, + status: McpStartupStatus::Ready, + }), + }) + .await; + } + }); + } +} + +fn codex_apps_reconnect_backoff(consecutive_failures: u32) -> Duration { + let exponent = consecutive_failures.saturating_sub(1).min(5); + CODEX_APPS_RECONNECT_INITIAL_BACKOFF + .saturating_mul(1 << exponent) + .min(CODEX_APPS_RECONNECT_MAX_BACKOFF) +} + +#[derive(Clone)] +struct ManagedClientStartup { + server_name: String, + server: EffectiveMcpServer, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + tx_event: Sender, + elicitation_requests: ElicitationRequestManager, + codex_apps_tools_cache_context: Option, + runtime_context: McpRuntimeContext, + runtime_auth_provider: Option, + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, + cancel_token: CancellationToken, + startup_complete: Arc, +} + +impl ManagedClientStartup { + fn start(&self) -> ManagedClientFuture { + let Self { + server_name, + server, + store_mode, + keyring_backend_kind, + tx_event, + elicitation_requests, + codex_apps_tools_cache_context, + runtime_context, + runtime_auth_provider, + client_elicitation_capability, + supports_openai_form_elicitation, + cancel_token, + startup_complete, + } = self.clone(); let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; let tool_filter = server .configured_config() .map(ToolFilter::from_config) .unwrap_or_default(); - let cached_server_info = if is_codex_apps_mcp_server { - codex_apps_tools_cache_context - .as_ref() - .and_then(load_startup_cached_codex_apps_server_info) - } else { - None - }; - let startup_tool_filter = tool_filter.clone(); - let codex_apps_tools_cache_context_for_fut = codex_apps_tools_cache_context.clone(); - let startup_complete = Arc::new(AtomicBool::new(false)); - let startup_complete_for_fut = Arc::clone(&startup_complete); - let cancel_token_for_fut = cancel_token.clone(); - let fut = async move { + let cancel_token_for_fut = cancel_token; + async move { let outcome = match async { if let Err(error) = validate_mcp_server_name(&server_name) { return Err(error.into()); @@ -207,10 +333,10 @@ impl AsyncManagedClient { .configured_config() .and_then(|config| config.tool_timeout_sec) .unwrap_or(DEFAULT_TOOL_TIMEOUT), - tool_filter: startup_tool_filter, + tool_filter, tx_event, elicitation_requests, - codex_apps_tools_cache_context: codex_apps_tools_cache_context_for_fut, + codex_apps_tools_cache_context, client_elicitation_capability, supports_openai_form_elicitation, }, @@ -224,10 +350,91 @@ impl AsyncManagedClient { Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled), }; - startup_complete_for_fut.store(true, Ordering::Release); + startup_complete.store(true, Ordering::Release); outcome + } + .in_current_span() + .boxed() + .shared() + } +} + +#[derive(Clone)] +pub(crate) struct AsyncManagedClient { + pub(crate) client: ManagedClientFuture, + pub(crate) is_codex_apps_mcp_server: bool, + pub(crate) cached_server_info: Option, + pub(crate) codex_apps_tools_cache_context: Option, + pub(crate) tool_filter: ToolFilter, + pub(crate) startup_complete: Arc, + pub(crate) startup_reconnect: Option>, + pub(crate) tool_plugin_provenance: Arc, + pub(crate) cancel_token: CancellationToken, +} + +impl AsyncManagedClient { + // Keep this constructor flat so the startup inputs remain readable at the + // single call site instead of introducing a one-off params wrapper. + #[instrument(level = "trace", skip_all, fields(server_name = %server_name))] + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + server_name: String, + startup_submit_id: String, + server: EffectiveMcpServer, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + cancel_token: CancellationToken, + tx_event: Sender, + elicitation_requests: ElicitationRequestManager, + codex_apps_tools_cache_context: Option, + tool_plugin_provenance: Arc, + runtime_context: McpRuntimeContext, + runtime_auth_provider: Option, + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, + ) -> Self { + let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; + let reconnect_server_name = server_name.clone(); + let reconnect_tx_event = tx_event.clone(); + let tool_filter = server + .configured_config() + .map(ToolFilter::from_config) + .unwrap_or_default(); + let cached_server_info = if is_codex_apps_mcp_server { + codex_apps_tools_cache_context + .as_ref() + .and_then(load_startup_cached_codex_apps_server_info) + } else { + None }; - let client = fut.in_current_span().boxed().shared(); + let startup_complete = Arc::new(AtomicBool::new(false)); + let startup = Arc::new(ManagedClientStartup { + server_name, + server, + store_mode, + keyring_backend_kind, + tx_event, + elicitation_requests, + codex_apps_tools_cache_context: codex_apps_tools_cache_context.clone(), + runtime_context, + runtime_auth_provider, + client_elicitation_capability, + supports_openai_form_elicitation, + cancel_token: cancel_token.clone(), + startup_complete: Arc::clone(&startup_complete), + }); + let client = startup.start(); + let startup_reconnect = is_codex_apps_mcp_server.then(|| { + let startup = Arc::clone(&startup); + Arc::new( + CodexAppsStartupReconnect::new(Arc::new(move || startup.start())) + .with_startup_status_context( + startup_submit_id, + reconnect_server_name, + reconnect_tx_event, + ), + ) + }); if codex_apps_tools_cache_context .as_ref() .is_some_and(CodexAppsToolsCacheContext::has_current_tools) @@ -245,15 +452,35 @@ impl AsyncManagedClient { codex_apps_tools_cache_context, tool_filter, startup_complete, + startup_reconnect, tool_plugin_provenance, cancel_token, } } pub(crate) async fn client(&self) -> Result { + if let Some(client) = self + .startup_reconnect + .as_ref() + .and_then(|reconnect| reconnect.current_client()) + { + return Ok(client); + } self.client.clone().await } + pub(crate) async fn reconnect_failed_startup(&self) { + let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { + return; + }; + if !self.startup_complete.load(Ordering::Acquire) { + return; + } + if matches!(self.client().await, Err(StartupOutcomeError::Failed { .. })) { + startup_reconnect.reconnect_in_background(); + } + } + pub(crate) async fn shutdown(&self) { self.cancel_token.cancel(); match self.client().await { diff --git a/codex-rs/core/tests/common/apps_test_server.rs b/codex-rs/core/tests/common/apps_test_server.rs index 0849087d0..3daaada6c 100644 --- a/codex-rs/core/tests/common/apps_test_server.rs +++ b/codex-rs/core/tests/common/apps_test_server.rs @@ -7,6 +7,7 @@ use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; use serde_json::Value; use serde_json::json; +use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use wiremock::Mock; @@ -55,6 +56,23 @@ pub struct AppsTestServer { pub chatgpt_base_url: String, } +#[derive(Clone)] +pub struct AppsTestServerStartupControl { + initialize_attempts: Arc, + remaining_initialize_failures: Arc, +} + +impl AppsTestServerStartupControl { + pub fn fail_next_initialize_attempts(&self, attempts: usize) { + self.remaining_initialize_failures + .store(attempts, Ordering::SeqCst); + } + + pub fn initialize_attempts(&self) -> usize { + self.initialize_attempts.load(Ordering::SeqCst) + } +} + #[derive(Clone, Copy)] pub enum AppsTestToolLoading { Direct, @@ -130,6 +148,34 @@ impl AppsTestServer { }) } + pub async fn mount_with_startup_control( + server: &MockServer, + ) -> Result<(Self, AppsTestServerStartupControl)> { + mount_oauth_metadata(server).await; + mount_connectors_directory(server).await; + let control = AppsTestServerStartupControl { + initialize_attempts: Arc::new(AtomicUsize::new(0)), + remaining_initialize_failures: Arc::new(AtomicUsize::new(0)), + }; + mount_streamable_http_json_rpc_with_startup_control( + server, + CONNECTOR_NAME.to_string(), + CONNECTOR_DESCRIPTION.to_string(), + /*searchable*/ true, + /*include_app_only_tool*/ false, + AppsTestToolsListBehavior::AlwaysAvailable, + Some(Arc::clone(&control.initialize_attempts)), + Some(Arc::clone(&control.remaining_initialize_failures)), + ) + .await; + Ok(( + Self { + chatgpt_base_url: server.uri(), + }, + control, + )) + } + pub async fn mount_with_tools_available_after_initial_list( server: &MockServer, ) -> Result { @@ -312,6 +358,30 @@ async fn mount_streamable_http_json_rpc( searchable: bool, include_app_only_tool: bool, tools_list_behavior: AppsTestToolsListBehavior, +) { + mount_streamable_http_json_rpc_with_startup_control( + server, + connector_name, + connector_description, + searchable, + include_app_only_tool, + tools_list_behavior, + /*initialize_attempts*/ None, + /*remaining_initialize_failures*/ None, + ) + .await; +} + +#[allow(clippy::too_many_arguments)] +async fn mount_streamable_http_json_rpc_with_startup_control( + server: &MockServer, + connector_name: String, + connector_description: String, + searchable: bool, + include_app_only_tool: bool, + tools_list_behavior: AppsTestToolsListBehavior, + initialize_attempts: Option>, + remaining_initialize_failures: Option>, ) { Mock::given(method("POST")) .and(path_regex("^/api/codex/apps/?$")) @@ -322,6 +392,8 @@ async fn mount_streamable_http_json_rpc( include_app_only_tool, tools_list_behavior, tools_list_calls: AtomicUsize::new(0), + initialize_attempts, + remaining_initialize_failures, }) .mount(server) .await; @@ -334,6 +406,8 @@ struct CodexAppsJsonRpcResponder { include_app_only_tool: bool, tools_list_behavior: AppsTestToolsListBehavior, tools_list_calls: AtomicUsize, + initialize_attempts: Option>, + remaining_initialize_failures: Option>, } impl Respond for CodexAppsJsonRpcResponder { @@ -355,6 +429,24 @@ impl Respond for CodexAppsJsonRpcResponder { match method { "initialize" => { + if let Some(initialize_attempts) = &self.initialize_attempts { + initialize_attempts.fetch_add(1, Ordering::SeqCst); + } + if self + .remaining_initialize_failures + .as_ref() + .is_some_and(|remaining| { + remaining + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |remaining| { + remaining.checked_sub(1) + }) + .is_ok() + }) + { + return ResponseTemplate::new(400).set_body_json(json!({ + "error": "simulated non-retryable Apps MCP startup failure", + })); + } let id = body.get("id").cloned().unwrap_or(Value::Null); let protocol_version = body .pointer("/params/protocolVersion") diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index b86862edc..32ecf7cf2 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -1,5 +1,8 @@ use anyhow::Result; use codex_features::Feature; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_protocol::protocol::McpServerRefreshConfig; +use codex_protocol::protocol::Op; use core_test_support::apps_test_server::AppsTestServer; use core_test_support::apps_test_server::SEARCH_CALENDAR_CREATE_TOOL; use core_test_support::apps_test_server::SEARCH_CALENDAR_NAMESPACE; @@ -8,10 +11,13 @@ use core_test_support::responses; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::namespace_child_tool; use core_test_support::responses::sse; use core_test_support::skip_if_no_network; +use core_test_support::wait_for_mcp_server; use serde_json::Value; +use std::time::Duration; #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> { @@ -79,3 +85,104 @@ async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn later_follow_up_uses_background_recovered_apps_after_mid_thread_startup_failures() +-> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (apps_server, startup_control) = + AppsTestServer::mount_with_startup_control(&server).await?; + let response = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "initial turn"), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + ev_assistant_message("msg-2", "recovery-trigger turn"), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-3", "recovered follow-up turn"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + + let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone()) + .with_config(move |config| { + config + .features + .enable(Feature::CodeModeOnly) + .expect("test config should allow feature update"); + config.code_mode.direct_only_tool_namespaces = + vec![SEARCH_CALENDAR_NAMESPACE.to_string()]; + }); + let test = builder.build(&server).await?; + wait_for_mcp_server(&test.codex, CODEX_APPS_MCP_SERVER_NAME).await?; + test.submit_turn("use Calendar before refreshing MCP") + .await?; + + let initial_request = response.requests()[0].body_json(); + assert!( + namespace_child_tool( + &initial_request, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some(), + "Calendar should be available before the MCP refresh: {initial_request}" + ); + + tokio::fs::remove_dir_all(test.codex_home_path().join("cache/codex_apps_tools")).await?; + startup_control.fail_next_initialize_attempts(/*attempts*/ 1); + let runtime_mcp_config = test.codex.runtime_mcp_config(&test.config).await; + let refresh_config = McpServerRefreshConfig { + mcp_servers: serde_json::to_value(codex_mcp::configured_mcp_servers(&runtime_mcp_config))?, + mcp_oauth_credentials_store_mode: serde_json::to_value( + runtime_mcp_config.mcp_oauth_credentials_store_mode, + )?, + auth_keyring_backend_kind: serde_json::to_value( + runtime_mcp_config.auth_keyring_backend_kind, + )?, + }; + test.codex + .submit(Op::RefreshMcpServers { + config: refresh_config, + }) + .await?; + test.submit_turn("use Calendar after transient Apps startup failures") + .await?; + tokio::time::timeout(Duration::from_secs(1), async { + while startup_control.initialize_attempts() < 3 { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .expect("background Apps reconnect should complete"); + test.submit_turn("use Calendar after background Apps recovery") + .await?; + + let requests = response.requests(); + assert_eq!(requests.len(), 3); + let recovered_request = requests[2].body_json(); + assert!( + namespace_child_tool( + &recovered_request, + SEARCH_CALENDAR_NAMESPACE, + SEARCH_CALENDAR_CREATE_TOOL, + ) + .is_some(), + "Calendar should recover on the follow-up turn: {recovered_request}", + ); + assert_eq!(startup_control.initialize_attempts(), 3); + + Ok(()) +}