Reconnect MCP servers on explicit refresh (#35151)

## Why

Runtime updates reuse unchanged MCP connections, so an explicit
`RefreshMcpServers` operation could leave existing server connections intact.

## What changed

- Mark explicit MCP refreshes to rebuild connections for every configured
  server on the next runtime replacement.
- Preserve the reconnect request if that replacement is cancelled before it
  finishes.
- Continue reusing connections for ordinary runtime updates.

## Testing

Update the transient Apps startup recovery test to wait for both turn
completion and the MCP server's ready event before checking the recovered tool.

GitOrigin-RevId: cb9fd978b3c07543058a2393a154fb70911531de
This commit is contained in:
jif
2026-07-24 10:22:58 +00:00
committed by copyberry
parent 6c729ef1c1
commit f201c30c52
3 changed files with 65 additions and 8 deletions
+32 -2
View File
@@ -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<PublishedMcpRuntime>,
reconnect_pending: AtomicBool,
elicitation_router: ElicitationRequestRouter,
}
@@ -82,6 +85,19 @@ struct PublishedMcpRuntime {
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
}
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<watch::Receiver<bool>>,
@@ -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<Arc<McpBinding>> {
let current = self.current.load_full();
+1
View File
@@ -422,6 +422,7 @@ pub async fn dynamic_tool_response(sess: &Arc<Session>, id: String, response: Dy
}
pub fn refresh_mcp_servers(sess: &Session) {
sess.services.mcp_runtime.reconnect_on_next_refresh();
sess.request_mcp_runtime_refresh();
}
+32 -6
View File
@@ -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