openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-12 21:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Restart pending MCP connections when startup timeout changes (#40636)

Browse Source 
## 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
This commit is contained in:
Tamir Duberstein
2026-08-25 16:20:38 +00:00
committed by copyberry
parent f5853cbe90
commit 20c3f9733f
3 changed files with 116 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager.rs
+8
View File
@@ -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<McpServerConnectionIdentity>,
client: AsyncManagedClient,
// Startup-only budget; changing it must not replace a ready connection.
startup_timeout: Duration,
startup_trigger: Option<watch::Sender<bool>>,
_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(),
}),
codex-rs/codex-mcp/src/connection_manager_tests.rs
+16 -2
View File
@@ -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(),
});
codex-rs/core/tests/suite/mcp_tool_exposure.rs
+92
View File
@@ -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;