mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
@@ -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();
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user