diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index f824f3b7664d..1bc9bb8b0c7c 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -8,6 +8,8 @@ use std::collections::HashMap; use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::time::Duration; use arc_swap::ArcSwap; @@ -70,6 +72,7 @@ pub struct McpRuntimeInput { /// their exact connections and configuration for as long as they are needed. pub struct McpRuntime { current: ArcSwap, + reconnect_pending: AtomicBool, elicitation_router: ElicitationRequestRouter, } @@ -82,6 +85,19 @@ struct PublishedMcpRuntime { ready_selected_capability_roots: Vec, } +struct McpReconnectGuard<'a> { + pending: &'a AtomicBool, + claimed: bool, +} + +impl Drop for McpReconnectGuard<'_> { + fn drop(&mut self) { + if self.claimed { + self.pending.store(true, Ordering::Release); + } + } +} + #[derive(Clone)] pub(crate) struct McpPublicationGate { published: Option>, @@ -132,6 +148,7 @@ impl McpRuntime { plugins_available: false, ready_selected_capability_roots: Vec::new(), }), + reconnect_pending: AtomicBool::new(false), elicitation_router: ElicitationRequestRouter::default(), } } @@ -145,8 +162,16 @@ impl McpRuntime { /// Reconciles configured servers and publishes their immutable runtime snapshot. pub async fn replace(&self, input: McpRuntimeInput) { let current = self.current.load_full(); - self.publish(input, Some(current.connections.as_ref())) - .await; + let mut reconnect = McpReconnectGuard { + pending: &self.reconnect_pending, + claimed: self.reconnect_pending.swap(false, Ordering::AcqRel), + }; + self.publish( + input, + (!reconnect.claimed).then_some(current.connections.as_ref()), + ) + .await; + reconnect.claimed = false; } /// Starts fresh connections and returns their complete, refreshed Apps catalog. @@ -182,6 +207,11 @@ impl McpRuntime { let _ = publish.send(true); } + /// Ensures the next refresh creates fresh connections for every configured server. + pub fn reconnect_on_next_refresh(&self) { + self.reconnect_pending.store(true, Ordering::Release); + } + /// Captures the latest published configuration and live client handles. pub async fn current_binding(&self) -> Option> { let current = self.current.load_full(); diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 485769eba893..33f2acc6734c 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -422,6 +422,7 @@ pub async fn dynamic_tool_response(sess: &Arc, id: String, response: Dy } pub fn refresh_mcp_servers(sess: &Session) { + sess.services.mcp_runtime.reconnect_on_next_refresh(); sess.request_mcp_runtime_refresh(); } diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 92bc90577905..8a3c687ab43c 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -833,15 +833,41 @@ async fn later_follow_up_uses_background_recovered_apps_after_mid_thread_startup tokio::fs::remove_dir_all(test.codex_home_path().join("cache/codex_apps_tools")).await?; startup_control.fail_next_initialize_attempts(/*attempts*/ 1); - test.codex - .set_openai_form_elicitation_support(/*supported*/ true) - .await?; test.codex.submit(Op::RefreshMcpServers).await?; - test.submit_turn("use Calendar after transient Apps startup failures") + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "use Calendar after transient Apps startup failures".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: Default::default(), + }) .await?; tokio::time::timeout(Duration::from_secs(5), async { - while startup_control.initialize_attempts() < 3 { - tokio::time::sleep(Duration::from_millis(1)).await; + let mut turn_complete = false; + let mut apps_ready = false; + while !turn_complete || !apps_ready { + let event = test + .codex + .next_event() + .await + .expect("event stream should stay open"); + match event.msg { + EventMsg::TurnComplete(_) => turn_complete = true, + EventMsg::McpStartupUpdate(update) + if update.server == CODEX_APPS_MCP_SERVER_NAME + && matches!( + update.status, + codex_protocol::protocol::McpStartupStatus::Ready + ) => + { + apps_ready = true; + } + _ => {} + } } }) .await