From 20c3f9733f7d7ee51e0eb26d82265db78c648335 Mon Sep 17 00:00:00 2001 From: Tamir Duberstein Date: Tue, 25 Aug 2026 16:09:34 +0000 Subject: [PATCH] Restart pending MCP connections when startup timeout changes (#40636) ## Why Reusing an MCP connection that is still starting after its `startup_timeout_sec` changes leaves the old startup budget in effect. ## What changed - Track the effective startup timeout on each MCP connection. - Reuse a pending startup only when its timeout still matches the current configuration. - Continue reusing ready connections when only the startup timeout changes. ## Testing Added integration coverage that refreshes the timeout while one server is still initializing and another is ready, verifying that only the pending startup is replaced. GitOrigin-RevId: 5ed71a57d3bf0dd6e227aacd16acd3a35bc37ecf --- codex-rs/codex-mcp/src/connection_manager.rs | 8 ++ .../codex-mcp/src/connection_manager_tests.rs | 18 +++- .../core/tests/suite/mcp_tool_exposure.rs | 92 +++++++++++++++++++ 3 files changed, 116 insertions(+), 2 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 2b5d71d0d0..0128c3dac4 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -37,6 +37,7 @@ use crate::mcp::ToolPluginProvenance; use crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS; use crate::pagination::MAX_MCP_CATALOG_ITEMS; use crate::rmcp_client::AsyncManagedClient; +use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT; use crate::rmcp_client::DEFAULT_TOOL_TIMEOUT; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::StartupOutcomeError; @@ -81,6 +82,8 @@ static LIVE_CONNECTIONS: Gauge = Gauge::new("mcp.connections.live"); pub(crate) struct McpServerConnection { identity: Option, client: AsyncManagedClient, + // Startup-only budget; changing it must not replace a ready connection. + startup_timeout: Duration, startup_trigger: Option>, _diagnostics_guard: GaugeGuard, } @@ -299,6 +302,9 @@ impl McpConnectionSet { let metadata = McpServerMetadata::from(&server); let configured_config = server.config().clone(); let configured_tool_filter = ToolFilter::from_config(&configured_config); + let startup_timeout = configured_config + .startup_timeout_sec + .unwrap_or(DEFAULT_STARTUP_TIMEOUT); let configured_tool_timeout = Some( configured_config .tool_timeout_sec @@ -385,6 +391,7 @@ impl McpConnectionSet { let reusable_pending_startup = connection.identity.as_ref() == Some(&connection_identity) && !connection.client.startup_complete.load(Ordering::Acquire) + && connection.startup_timeout == startup_timeout && !connection.startup_is_dormant() && !connection.client.cancel_token.is_cancelled() && previous_view.catalog_item_limit == catalog_item_limit @@ -544,6 +551,7 @@ impl McpConnectionSet { connection: Arc::new(McpServerConnection { identity: Some(connection_identity), client: async_managed_client.clone(), + startup_timeout, startup_trigger, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 297d46b6ae..2347c9120e 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -122,6 +122,7 @@ impl McpConnectionSet { connection: Arc::new(McpServerConnection { identity: None, client, + startup_timeout: DEFAULT_STARTUP_TIMEOUT, startup_trigger: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), @@ -4262,6 +4263,9 @@ async fn manager_with_reusable_ready_server( connection: Arc::new(McpServerConnection { identity: Some(reusable_server_identity(config, runtime_context)), client: create_ready_async_managed_client(tools).await, + startup_timeout: config + .startup_timeout_sec + .unwrap_or(DEFAULT_STARTUP_TIMEOUT), startup_trigger: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), @@ -4396,6 +4400,9 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a startup_reconnect: None, cancel_token: CancellationToken::new(), }, + startup_timeout: config + .startup_timeout_sec + .unwrap_or(DEFAULT_STARTUP_TIMEOUT), startup_trigger: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), @@ -4492,6 +4499,7 @@ async fn reconciliation_reuses_an_unchanged_pending_server_without_waiting() -> .expect("test server should have one connection owner"); connection.client = pending_client; config.enabled_tools = Some(vec!["search".to_string()]); + config.startup_timeout_sec = Some(DEFAULT_STARTUP_TIMEOUT); let reconciled = tokio::time::timeout( Duration::from_millis(100), @@ -4867,11 +4875,14 @@ async fn reconciliation_reuses_ready_server_when_startup_timeout_changes() { vec![create_test_tool("docs", "search")], ) .await; - config.startup_timeout_sec = Some(Duration::from_secs(30)); + config.startup_timeout_sec = Some(Duration::from_secs(60)); let reconciled = reconcile_reusable_server(&previous, config, runtime_context).await; - assert!(previous.shares_test_connection_with(&reconciled, "docs")); + assert_eq!( + model_tool_names(&reconciled.list_all_tools().await), + HashSet::from([ToolName::namespaced("mcp__docs", "search")]) + ); } #[tokio::test] @@ -4930,6 +4941,9 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> { startup_reconnect: None, cancel_token: CancellationToken::new(), }, + startup_timeout: config + .startup_timeout_sec + .unwrap_or(DEFAULT_STARTUP_TIMEOUT), startup_trigger: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }); diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 0b07704e3e..a6fbecdafd 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -566,6 +566,98 @@ async fn root_reconciliation_reuses_pending_apps_startup() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn timeout_refresh_replaces_pending_startup_and_reuses_ready_connection() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let pending_mock = responses::start_mock_server().await; + let (pending_server, pending_startup) = + AppsTestServer::mount_with_startup_control(&pending_mock).await?; + let release_startup = pending_startup.hold_next_successful_initialize(); + let ready_mock = responses::start_mock_server().await; + let (ready_server, ready_startup) = + AppsTestServer::mount_with_startup_control(&ready_mock).await?; + let test = core_test_support::test_codex::test_codex() + .with_config(move |config| { + config + .mcp_servers + .set( + [("pending", pending_server), ("ready", ready_server)] + .into_iter() + .map(|(name, server)| { + ( + name.to_string(), + serde_json::from_value(json!({ + "url": format!("{}/api/codex/ps/mcp", server.chatgpt_base_url), + "startup_timeout_sec": 60, + })) + .expect("valid MCP config"), + ) + }) + .collect(), + ) + .expect("test config should allow MCP servers"); + }) + .build_with_auto_env(&server) + .await?; + tokio::time::timeout(Duration::from_secs(5), async { + while pending_startup.initialize_attempts() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("pending startup should begin before refresh"); + let ready_result = test + .codex + .call_mcp_tool( + "ready", + "calendar_list_events", + /*arguments*/ None, + /*meta*/ None, + ) + .await?; + + let mut refresh_config = test.config.clone(); + let mut servers = refresh_config.mcp_servers.get().clone(); + for config in servers.values_mut() { + config.startup_timeout_sec = None; + } + refresh_config.mcp_servers.set(servers)?; + test.codex.refresh_mcp_config(refresh_config).await; + // Publish without waiting for the held initialize to finish. + let error = test + .codex + .read_mcp_resource("unknown", ReadResourceRequestParams::new("test://resource")) + .await + .expect_err("the unknown server should not exist"); + assert_eq!(error.to_string(), "unknown MCP server 'unknown'"); + release_startup + .send(()) + .expect("the mock initialize should remain held until publication"); + + for name in ["pending", "ready"] { + let result = test + .codex + .call_mcp_tool( + name, + "calendar_list_events", + /*arguments*/ None, + /*meta*/ None, + ) + .await?; + assert_eq!(result, ready_result); + } + assert_eq!( + ( + pending_startup.initialize_attempts(), + ready_startup.initialize_attempts(), + ), + (2, 1) + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn out_of_band_resource_read_reconciles_the_published_mcp_runtime() -> Result<()> { let server = responses::start_mock_server().await;