From 124e560b9357746b6a7e2bba1fd244115fd9cc5a Mon Sep 17 00:00:00 2001 From: trolle-oai Date: Thu, 27 Aug 2026 17:55:37 +0000 Subject: [PATCH] Make the optional MCP startup grace configurable (#41199) ## What changed - Add `mcp_optional_startup_grace_ms` with a default of 1,000 ms to control how long tool catalog capture waits for optional MCP servers. - Treat a value of `0` as disabling the shared grace, so optional servers use their configured `startup_timeout_sec` instead. - Apply updated grace values during runtime and MCP configuration refreshes, and reset cached startup deadlines when the configured duration changes. ## Testing - Cover custom, disabled, refreshed, and shared startup grace behavior in MCP connection and core integration tests. GitOrigin-RevId: 936196215ed6e0a29c9ec72cfee17663978df8ef --- .../src/connection_manager/tool_catalog.rs | 26 +- .../codex-mcp/src/connection_manager_tests.rs | 105 ++++++- codex-rs/codex-mcp/src/lib.rs | 1 + codex-rs/codex-mcp/src/mcp/mod.rs | 8 + codex-rs/codex-mcp/src/mcp/mod_tests.rs | 1 + codex-rs/codex-mcp/src/tool_catalog_cache.rs | 30 +- codex-rs/config/src/config_toml.rs | 6 + codex-rs/core/config.schema.json | 6 + codex-rs/core/src/config/config_tests.rs | 19 ++ codex-rs/core/src/config/mod.rs | 10 + codex-rs/core/src/session/mod.rs | 2 + .../tests/suite/mcp_optional_startup_grace.rs | 296 ++++++++++++++++++ codex-rs/core/tests/suite/mod.rs | 1 + codex-rs/thread-manager-sample/src/main.rs | 1 + 14 files changed, 495 insertions(+), 17 deletions(-) create mode 100644 codex-rs/core/tests/suite/mcp_optional_startup_grace.rs diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index ca92e0eeda..c09c1b9090 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -1,6 +1,5 @@ use std::sync::Arc; use std::sync::atomic::Ordering; -use std::time::Duration; use std::time::Instant; use anyhow::Context; @@ -32,8 +31,6 @@ use crate::tools::normalize_tools_for_model_with_prefix; const MCP_UI_META_KEY: &str = "ui"; const MCP_UI_VISIBILITY_META_KEY: &str = "visibility"; const MCP_UI_MODEL_VISIBILITY: &str = "model"; -const OPTIONAL_MCP_STARTUP_GRACE: Duration = Duration::from_secs(1); - /// Returns whether a tool may be included in model-facing tool declarations. /// /// Tools without visibility metadata remain visible. Tools with visibility @@ -177,6 +174,7 @@ impl McpConnectionSet { let revision = self.tool_catalog_revision.read().await; let mut listed_tools = Vec::new(); let mut clients = std::collections::HashMap::new(); + let optional_mcp_startup_grace = config.optional_mcp_startup_grace; join_all(self.servers.iter().map(|(server_name, view)| async move { if !view .connection @@ -196,12 +194,21 @@ impl McpConnectionSet { if !must_wait_for_startup && has_cached_tools { return; } - if !must_wait_for_startup { + if !must_wait_for_startup && optional_mcp_startup_grace.is_zero() { + if let Some(cache) = view.connection.client.tool_catalog_cache_context.as_ref() + { + cache.optional_startup_deadline( + tokio::time::Instant::now(), + optional_mcp_startup_grace, + ); + } + let _ = view.connection.client().await; + } else if !must_wait_for_startup { let optional_startup_deadline = if view.connection.startup_is_dormant() { - tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE + tokio::time::Instant::now() + optional_mcp_startup_grace } else { *self.optional_startup_deadline.get_or_init(|| { - tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE + tokio::time::Instant::now() + optional_mcp_startup_grace }) }; let startup_deadline = view @@ -209,7 +216,12 @@ impl McpConnectionSet { .client .tool_catalog_cache_context .as_ref() - .map(|cache| cache.optional_startup_deadline(optional_startup_deadline)) + .map(|cache| { + cache.optional_startup_deadline( + optional_startup_deadline, + optional_mcp_startup_grace, + ) + }) .unwrap_or(optional_startup_deadline); if tokio::time::timeout_at(startup_deadline, view.connection.client()) .await diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index f5f7046b42..e49845b876 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -2563,7 +2563,7 @@ async fn capture_binding_exposes_cached_tools_before_startup() { } #[tokio::test(start_paused = true)] -async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup_grace() { +async fn capture_binding_skips_pending_optional_servers_after_configured_shared_startup_grace() { let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); let mut manager = McpConnectionSet::new_uninitialized( @@ -2581,6 +2581,7 @@ async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup .expect("optional plugin MCP config"), )); plugin_config.mcp_server_catalog = catalog.build(); + plugin_config.optional_mcp_startup_grace = Duration::from_millis(250); manager.tool_plugin_provenance = Arc::new(crate::tool_plugin_provenance(&plugin_config)); for server_name in ["pending-one", "pending-two"] { manager.insert_test_client( @@ -2602,9 +2603,16 @@ async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup let manager = Arc::new(manager); assert_eq!(manager.stable_catalog_revision().await, None); - let binding = tokio::time::timeout(Duration::from_millis(1500), capture_binding(&manager)) - .await - .expect("all optional servers should share a single startup grace"); + let binding = tokio::time::timeout( + Duration::from_millis(500), + manager.capture_binding_with_metadata( + Arc::new(plugin_config), + /*plugins_available*/ false, + /*required_servers*/ &[], + ), + ) + .await + .expect("all optional servers should share the configured startup grace"); assert!(binding.tools().is_empty()); let binding = tokio::time::timeout(Duration::from_millis(1), capture_binding(&manager)) @@ -2644,6 +2652,49 @@ async fn capture_binding_skips_pending_optional_servers_after_one_shared_startup assert!(binding.is_err(), "explicitly requested servers must wait"); } +#[tokio::test(start_paused = true)] +async fn capture_binding_waits_for_optional_startup_when_shared_grace_is_disabled() { + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + let (client, startup_started, release_startup) = create_gated_async_managed_client( + create_test_managed_client(vec![create_test_tool("optional", "echo")]).await, + ); + manager.insert_test_client("optional", client); + + let mut config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + config.optional_mcp_startup_grace = Duration::ZERO; + config + .server_permission_profiles + .insert("optional".to_string(), PermissionProfile::default()); + let manager = Arc::new(manager); + let mut capture = tokio::spawn(async move { + manager + .capture_binding_with_metadata( + Arc::new(config), + /*plugins_available*/ false, + /*required_servers*/ &[], + ) + .await + }); + + startup_started.await.expect("client startup should begin"); + assert!( + tokio::time::timeout(Duration::from_millis(1), &mut capture) + .await + .is_err(), + "disabled shared grace should keep waiting for optional startup" + ); + release_startup.send(()).expect("release client startup"); + + let binding = capture.await.expect("capture binding task"); + assert!(binding.prepare_call("optional", "echo").is_some()); +} + #[tokio::test] async fn stable_catalog_revision_ignores_terminal_optional_server_failures() { let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); @@ -2745,13 +2796,57 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() .expect("the next thread must not restart the same server's startup grace"); assert!(second.tools().is_empty()); + let mut disabled_config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + disabled_config.optional_mcp_startup_grace = Duration::ZERO; + let disabled_manager = create_connection_set(); + assert!( + tokio::time::timeout( + Duration::from_millis(1), + disabled_manager.capture_binding_with_metadata( + Arc::new(disabled_config), + /*plugins_available*/ false, + /*required_servers*/ &[], + ), + ) + .await + .is_err(), + "disabled grace should keep waiting for the pending optional server" + ); + + let restored_started = tokio::time::Instant::now(); + let restored = tokio::time::timeout( + Duration::from_millis(1500), + capture_binding(&create_connection_set()), + ) + .await + .expect("restoring the startup grace should create a fresh deadline"); + assert!(restored.tools().is_empty()); + assert_eq!(restored_started.elapsed(), Duration::from_secs(1)); + + let mut updated_config = crate::mcp::tests::test_mcp_config(std::env::temp_dir()); + updated_config.optional_mcp_startup_grace = Duration::from_millis(250); + let updated_manager = create_connection_set(); + let updated_started = tokio::time::Instant::now(); + let updated = tokio::time::timeout( + Duration::from_millis(500), + updated_manager.capture_binding_with_metadata( + Arc::new(updated_config), + /*plugins_available*/ false, + /*required_servers*/ &[], + ), + ) + .await + .expect("a changed startup grace should receive its newly configured deadline"); + assert!(updated.tools().is_empty()); + assert_eq!(updated_started.elapsed(), Duration::from_millis(250)); + cache_context.publish_if_newest( cache_context.begin_fetch(), &[create_test_tool("pending", "cached_tool")], ); let deadline_after_publication = tokio::time::Instant::now() + Duration::from_secs(1); assert_eq!( - cache_context.optional_startup_deadline(deadline_after_publication), + cache_context.optional_startup_deadline(deadline_after_publication, Duration::from_secs(1)), deadline_after_publication, "publishing a catalog must not install a stale startup deadline" ); diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 254f992899..f2bf80fbb4 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -43,6 +43,7 @@ pub use catalog::ResolvedMcpCatalog; pub use catalog::ResolvedMcpServer; pub use mcp::CODEX_APPS_MCP_SERVER_NAME; +pub use mcp::DEFAULT_OPTIONAL_MCP_STARTUP_GRACE; pub use mcp::McpConfig; pub use mcp::ToolPluginProvenance; pub use server::EffectiveMcpServer; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index b1261863b7..1042230797 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -134,6 +134,11 @@ pub struct McpConfig { pub mcp_oauth_callback_port: Option, /// Optional OAuth redirect URI override for MCP login. pub mcp_oauth_callback_url: Option, + /// How long a tool catalog capture waits for optional MCP servers to initialize. + /// + /// A zero duration disables the shared grace and waits for each server's + /// configured startup timeout instead. + pub optional_mcp_startup_grace: Duration, /// Whether skill MCP dependency installation prompts are enabled. pub skill_mcp_dependency_install_enabled: bool, /// Approval policy used for MCP tool calls and MCP elicitation requests. @@ -175,6 +180,9 @@ pub struct McpConfig { pub connector_snapshot: ConnectorSnapshot, } +/// Default amount of time a tool catalog capture waits for optional MCP servers. +pub const DEFAULT_OPTIONAL_MCP_STARTUP_GRACE: Duration = Duration::from_secs(1); + impl McpConfig { /// Resolves enabled runtime servers against the exact attachment permissions being published. pub fn set_server_permission_profiles( diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 350e8f8878..dd0856b664 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -27,6 +27,7 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig { auth_keyring_backend_kind: AuthKeyringBackendKind::default(), mcp_oauth_callback_port: None, mcp_oauth_callback_url: None, + optional_mcp_startup_grace: DEFAULT_OPTIONAL_MCP_STARTUP_GRACE, skill_mcp_dependency_install_enabled: true, approval_policy: Constrained::allow_any(AskForApproval::OnRequest), permission_profile: PermissionProfile::default(), diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index e6e99f4acb..77fb60e804 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -56,11 +56,16 @@ struct ToolCatalogCacheEntry { #[derive(Default)] struct ToolCatalogCacheState { snapshot: Option, - optional_startup_deadline: Option, + optional_startup_deadline: Option, last_accepted_generation: u64, disabled_by_server: bool, } +struct OptionalStartupDeadline { + grace: Duration, + deadline: Instant, +} + struct ToolCatalogSnapshot { tools: Vec, published_at: Instant, @@ -114,7 +119,11 @@ impl McpToolCatalogCacheContext { self.current_tools().is_some_and(|tools| !tools.is_empty()) } - pub(crate) fn optional_startup_deadline(&self, default_deadline: Instant) -> Instant { + pub(crate) fn optional_startup_deadline( + &self, + default_deadline: Instant, + startup_grace: Duration, + ) -> Instant { let mut state = lock_unpoisoned(&self.entry.state); if state.disabled_by_server || state @@ -124,9 +133,20 @@ impl McpToolCatalogCacheContext { { return default_deadline; } - *state - .optional_startup_deadline - .get_or_insert(default_deadline) + let cached_deadline = + state + .optional_startup_deadline + .get_or_insert(OptionalStartupDeadline { + grace: startup_grace, + deadline: default_deadline, + }); + if cached_deadline.grace != startup_grace { + *cached_deadline = OptionalStartupDeadline { + grace: startup_grace, + deadline: default_deadline, + }; + } + cached_deadline.deadline } pub(crate) fn current_tools(&self) -> Option> { diff --git a/codex-rs/config/src/config_toml.rs b/codex-rs/config/src/config_toml.rs index af5df4f1fc..ee0fd5a5d0 100644 --- a/codex-rs/config/src/config_toml.rs +++ b/codex-rs/config/src/config_toml.rs @@ -287,6 +287,12 @@ pub struct ConfigToml { /// to 127.0.0.1 (using `mcp_oauth_callback_port` when provided). pub mcp_oauth_callback_url: Option, + /// Milliseconds to wait for optional MCP servers while building the initial tool catalog. + /// + /// Defaults to 1000. Set to 0 to disable the shared grace and wait for each + /// server's configured `startup_timeout_sec` instead. + pub mcp_optional_startup_grace_ms: Option, + /// User-defined provider entries that extend the built-in list. Built-in /// IDs cannot be overridden. #[serde(default, deserialize_with = "deserialize_model_providers")] diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 0af689ee60..19484e26f1 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -6284,6 +6284,12 @@ "default": null, "description": "Preferred backend for storing MCP OAuth credentials. keyring: Use an OS-specific keyring service. https://github.com/openai/codex/blob/main/codex-rs/rmcp-client/src/oauth.rs#L2 file: Use a file in the Codex home directory. auto (default): Use the OS-specific keyring service if available, otherwise use a file." }, + "mcp_optional_startup_grace_ms": { + "description": "Milliseconds to wait for optional MCP servers while building the initial tool catalog.\n\nDefaults to 1000. Set to 0 to disable the shared grace and wait for each server's configured `startup_timeout_sec` instead.", + "format": "uint64", + "minimum": 0.0, + "type": "integer" + }, "mcp_servers": { "additionalProperties": { "$ref": "#/definitions/RawMcpServerConfig" diff --git a/codex-rs/core/src/config/config_tests.rs b/codex-rs/core/src/config/config_tests.rs index ce24e1b6b2..3a4554456f 100644 --- a/codex-rs/core/src/config/config_tests.rs +++ b/codex-rs/core/src/config/config_tests.rs @@ -228,6 +228,25 @@ async fn load_config_normalizes_relative_cwd_override() -> std::io::Result<()> { Ok(()) } +#[tokio::test] +async fn load_config_applies_optional_mcp_startup_grace() -> std::io::Result<()> { + let codex_home = tempdir()?; + let config_toml: ConfigToml = toml::from_str("mcp_optional_startup_grace_ms = 2500") + .expect("optional MCP startup grace should parse from config.toml"); + let config = Config::load_from_base_config_with_overrides( + config_toml, + ConfigOverrides::default(), + codex_home.abs(), + ) + .await?; + + assert_eq!( + config.mcp_optional_startup_grace, + Duration::from_millis(2500) + ); + Ok(()) +} + #[tokio::test] async fn test_toml_parsing() { let history_with_persistence = r#" diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index a97badf61c..96da2f76f8 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -80,6 +80,7 @@ use codex_http_client::OutboundProxyPolicy; use codex_install_context::InstallContext; use codex_login::AuthManagerConfig; use codex_login::AuthRouteConfig; +use codex_mcp::DEFAULT_OPTIONAL_MCP_STARTUP_GRACE; use codex_mcp::McpConfig; use codex_mcp::McpPluginAttribution; use codex_mcp::McpProtocolMode; @@ -140,6 +141,7 @@ use std::num::NonZeroUsize; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +use std::time::Duration; use crate::config::permissions::BUILT_IN_READ_ONLY_PROFILE; use crate::config::permissions::BUILT_IN_WORKSPACE_PROFILE; @@ -847,6 +849,9 @@ pub struct Config { /// to 127.0.0.1 (using `mcp_oauth_callback_port` when provided). pub mcp_oauth_callback_url: Option, + /// How long to wait for optional MCP servers while building the initial tool catalog. + pub mcp_optional_startup_grace: Duration, + /// Combined provider map (defaults plus user-defined providers). pub model_providers: HashMap, @@ -1731,6 +1736,7 @@ impl Config { auth_keyring_backend_kind: self.auth_keyring_backend_kind(), mcp_oauth_callback_port: self.mcp_oauth_callback_port, mcp_oauth_callback_url: self.mcp_oauth_callback_url.clone(), + optional_mcp_startup_grace: self.mcp_optional_startup_grace, skill_mcp_dependency_install_enabled: self .features .enabled(Feature::SkillMcpDependencyInstall), @@ -4165,6 +4171,10 @@ impl Config { ), mcp_oauth_callback_port: cfg.mcp_oauth_callback_port, mcp_oauth_callback_url: cfg.mcp_oauth_callback_url.clone(), + mcp_optional_startup_grace: cfg + .mcp_optional_startup_grace_ms + .map(Duration::from_millis) + .unwrap_or(DEFAULT_OPTIONAL_MCP_STARTUP_GRACE), model_providers, project_doc_max_bytes: cfg.project_doc_max_bytes.unwrap_or(AGENTS_MD_MAX_BYTES), project_doc_fallback_filenames: cfg diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 9867606a20..363eda5446 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -1800,6 +1800,7 @@ impl Session { config.tool_suggest = resolve_tool_suggest_config_from_layer_stack(&config.config_layer_stack); config.mcp_servers = next_config.mcp_servers.clone(); + config.mcp_optional_startup_grace = next_config.mcp_optional_startup_grace; config.mcp_oauth_credentials_store_mode = next_config.mcp_oauth_credentials_store_mode; if let Err(err) = config.features.set_enabled( Feature::SecretAuthStorage, @@ -1847,6 +1848,7 @@ impl Session { .config_layer_stack .with_user_layer_from(&config.config_layer_stack); config.mcp_servers = next_config.mcp_servers; + config.mcp_optional_startup_grace = next_config.mcp_optional_startup_grace; config.mcp_oauth_credentials_store_mode = next_config.mcp_oauth_credentials_store_mode; if let Err(err) = config.features.set_enabled( Feature::SecretAuthStorage, diff --git a/codex-rs/core/tests/suite/mcp_optional_startup_grace.rs b/codex-rs/core/tests/suite/mcp_optional_startup_grace.rs new file mode 100644 index 0000000000..f5ef343ae4 --- /dev/null +++ b/codex-rs/core/tests/suite/mcp_optional_startup_grace.rs @@ -0,0 +1,296 @@ +//! Covers optional MCP startup deadlines in the initial model-visible tool catalog. + +use std::time::Duration; + +use anyhow::Context; +use core_test_support::apps_test_server::AppsTestServer; +use core_test_support::responses; +use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::json; +use test_case::test_case; + +const SERVER_NAME: &str = "optional_startup"; +const TOOL_NAMESPACE: &str = "mcp__optional_startup"; +const TOOL_NAME: &str = "calendar_create_event"; + +#[derive(Clone, Copy)] +enum StartupGraceScenario { + ShortGraceOmitsPending, + CustomGraceIncludesReady, + DisabledGraceWaitsForStartup, + DisabledGraceRespectsStartupTimeout, +} + +#[derive(Clone, Copy)] +enum ConfigRefreshPath { + Runtime, + Mcp, +} + +#[test_case(StartupGraceScenario::ShortGraceOmitsPending; "custom grace omits a pending server")] +#[test_case(StartupGraceScenario::CustomGraceIncludesReady; "custom grace admits a ready server")] +#[test_case(StartupGraceScenario::DisabledGraceWaitsForStartup; "zero grace waits for server startup")] +#[test_case(StartupGraceScenario::DisabledGraceRespectsStartupTimeout; "zero grace respects server startup timeout")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn optional_mcp_startup_grace_controls_initial_turn_tool_catalog( + scenario: StartupGraceScenario, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let startup_grace = match scenario { + StartupGraceScenario::ShortGraceOmitsPending => Duration::from_millis(50), + StartupGraceScenario::CustomGraceIncludesReady => Duration::from_millis(250), + StartupGraceScenario::DisabledGraceWaitsForStartup + | StartupGraceScenario::DisabledGraceRespectsStartupTimeout => Duration::ZERO, + }; + let startup_timeout = match scenario { + StartupGraceScenario::DisabledGraceRespectsStartupTimeout => Duration::from_millis(250), + StartupGraceScenario::ShortGraceOmitsPending + | StartupGraceScenario::CustomGraceIncludesReady + | StartupGraceScenario::DisabledGraceWaitsForStartup => Duration::from_secs(1), + }; + let responses_server = responses::start_mock_server().await; + let mcp_server = responses::start_mock_server().await; + let (http_server, startup_control) = + AppsTestServer::mount_with_startup_control(&mcp_server).await?; + let release_startup = startup_control.hold_next_successful_initialize(); + let server_url = format!("{}/api/codex/ps/mcp", http_server.chatgpt_base_url); + let response = responses::mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("resp-optional-startup"), + responses::ev_assistant_message("msg-optional-startup", "done"), + responses::ev_completed("resp-optional-startup"), + ]), + ) + .await; + + let fixture = test_codex() + .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) + .with_config(move |config| { + config.mcp_optional_startup_grace = startup_grace; + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + SERVER_NAME.to_string(), + serde_json::from_value(json!({ + "url": server_url, + "http_headers": { "Authorization": "Bearer synthetic-test-token" }, + "enabled_tools": [TOOL_NAME], + "startup_timeout_sec": startup_timeout.as_secs_f64(), + })) + .expect("synthetic optional MCP server configuration"), + ); + config + .mcp_servers + .set(servers) + .expect("optional MCP server should satisfy test configuration"); + }) + .build_with_auto_env(&responses_server) + .await?; + + tokio::time::timeout(Duration::from_secs(5), async { + while startup_control.initialize_attempts() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .context("optional MCP initialization should begin before the first turn")?; + + let mut turn = Box::pin(fixture.submit_turn("show optional MCP tools")); + match scenario { + StartupGraceScenario::ShortGraceOmitsPending => { + tokio::time::timeout(Duration::from_millis(500), &mut turn) + .await + .context("the configured grace should omit the pending server")??; + release_startup + .send(()) + .expect("pending optional MCP startup should remain in flight"); + } + StartupGraceScenario::CustomGraceIncludesReady => { + assert!( + tokio::time::timeout(Duration::from_millis(25), &mut turn) + .await + .is_err(), + "the first turn should wait while the custom startup grace remains available" + ); + release_startup + .send(()) + .expect("optional MCP startup should remain in flight"); + tokio::time::timeout(Duration::from_secs(2), &mut turn) + .await + .context("a server ready within the configured grace should reach the model")??; + } + StartupGraceScenario::DisabledGraceWaitsForStartup => { + assert!( + tokio::time::timeout(Duration::from_millis(150), &mut turn) + .await + .is_err(), + "zero grace should wait for startup rather than omitting the optional server" + ); + release_startup + .send(()) + .expect("zero-grace optional MCP startup should remain in flight"); + tokio::time::timeout(Duration::from_secs(2), &mut turn) + .await + .context("zero grace should include a server ready before its startup timeout")??; + } + StartupGraceScenario::DisabledGraceRespectsStartupTimeout => { + assert!( + tokio::time::timeout(Duration::from_millis(100), &mut turn) + .await + .is_err(), + "zero grace should keep waiting until the server-specific startup timeout" + ); + tokio::time::timeout(Duration::from_secs(1), &mut turn) + .await + .context("zero grace should stop waiting once the server startup times out")??; + release_startup + .send(()) + .expect("the test MCP initialization request should remain gated"); + } + } + + let request = response.single_request().body_json(); + let server_tool = responses::namespace_child_tool(&request, TOOL_NAMESPACE, TOOL_NAME); + match scenario { + StartupGraceScenario::ShortGraceOmitsPending + | StartupGraceScenario::DisabledGraceRespectsStartupTimeout => assert!( + server_tool.is_none(), + "a pending optional MCP tool must be absent after its configured grace expires" + ), + StartupGraceScenario::CustomGraceIncludesReady + | StartupGraceScenario::DisabledGraceWaitsForStartup => assert!( + server_tool.is_some(), + "an optional MCP tool ready before its startup deadline must reach the model" + ), + } + + fixture.codex.shutdown_and_wait().await?; + Ok(()) +} + +#[test_case(ConfigRefreshPath::Runtime; "runtime configuration refresh")] +#[test_case(ConfigRefreshPath::Mcp; "MCP configuration refresh")] +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn running_thread_uses_refreshed_optional_mcp_startup_grace( + refresh_path: ConfigRefreshPath, +) -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let responses_server = responses::start_mock_server().await; + let mcp_server = responses::start_mock_server().await; + let (http_server, startup_control) = + AppsTestServer::mount_with_startup_control(&mcp_server).await?; + let release_startup = startup_control.hold_next_successful_initialize(); + let server_url = format!("{}/api/codex/ps/mcp", http_server.chatgpt_base_url); + let initial_response = responses::mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("resp-initial-grace"), + responses::ev_assistant_message("msg-initial-grace", "done"), + responses::ev_completed("resp-initial-grace"), + ]), + ) + .await; + + let fixture = test_codex() + .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) + .with_config(move |config| { + config.mcp_optional_startup_grace = Duration::from_millis(50); + let mut servers = config.mcp_servers.get().clone(); + servers.insert( + SERVER_NAME.to_string(), + serde_json::from_value(json!({ + "url": server_url, + "http_headers": { "Authorization": "Bearer synthetic-test-token" }, + "enabled_tools": [TOOL_NAME], + "startup_timeout_sec": 5, + })) + .expect("synthetic optional MCP server configuration"), + ); + config + .mcp_servers + .set(servers) + .expect("optional MCP server should satisfy test configuration"); + }) + .build_with_auto_env(&responses_server) + .await?; + + tokio::time::timeout(Duration::from_secs(5), async { + while startup_control.initialize_attempts() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .context("optional MCP initialization should begin before the initial turn")?; + + tokio::time::timeout( + Duration::from_millis(500), + fixture.submit_turn("show initial optional MCP tools"), + ) + .await + .context("the initial startup grace should omit the pending server")??; + assert!( + responses::namespace_child_tool( + &initial_response.single_request().body_json(), + TOOL_NAMESPACE, + TOOL_NAME, + ) + .is_none(), + "the pending optional MCP tool should be absent before configuration refresh" + ); + + let refreshed_grace = Duration::from_millis(250); + let mut refreshed_config = fixture.config.clone(); + refreshed_config.mcp_optional_startup_grace = refreshed_grace; + match refresh_path { + ConfigRefreshPath::Runtime => { + fixture.codex.refresh_runtime_config(refreshed_config).await; + } + ConfigRefreshPath::Mcp => { + fixture.codex.refresh_mcp_config(refreshed_config).await; + } + } + assert_eq!( + fixture.codex.config().await.mcp_optional_startup_grace, + refreshed_grace, + "the existing thread should retain the refreshed optional MCP startup grace" + ); + + let refreshed_response = responses::mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("resp-refreshed-grace"), + responses::ev_assistant_message("msg-refreshed-grace", "done"), + responses::ev_completed("resp-refreshed-grace"), + ]), + ) + .await; + let mut refreshed_turn = Box::pin(fixture.submit_turn("show refreshed optional MCP tools")); + assert!( + tokio::time::timeout(Duration::from_millis(25), &mut refreshed_turn) + .await + .is_err(), + "the refreshed startup grace should wait for the existing pending MCP server" + ); + release_startup + .send(()) + .expect("the existing optional MCP startup should remain in flight"); + tokio::time::timeout(Duration::from_secs(2), &mut refreshed_turn) + .await + .context("an optional MCP ready within the refreshed grace should reach the model")??; + assert!( + responses::namespace_child_tool( + &refreshed_response.single_request().body_json(), + TOOL_NAMESPACE, + TOOL_NAME, + ) + .is_some(), + "the refreshed optional MCP startup grace should expose the ready tool" + ); + + fixture.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 5aba823765..28cb2c6b46 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -93,6 +93,7 @@ mod json_result; mod live_cli; mod mcp_auth_elicitation; mod mcp_auth_refresh; +mod mcp_optional_startup_grace; #[cfg(unix)] mod mcp_refresh_cleanup; mod mcp_startup_refresh_http_proxy; diff --git a/codex-rs/thread-manager-sample/src/main.rs b/codex-rs/thread-manager-sample/src/main.rs index 32461981b9..e65038840a 100644 --- a/codex-rs/thread-manager-sample/src/main.rs +++ b/codex-rs/thread-manager-sample/src/main.rs @@ -247,6 +247,7 @@ fn new_config(model: Option, arg0_paths: Arg0DispatchPaths) -> anyhow::R mcp_oauth_credentials_store_mode: OAuthCredentialsStoreMode::File, mcp_oauth_callback_port: None, mcp_oauth_callback_url: None, + mcp_optional_startup_grace: std::time::Duration::from_secs(1), model_providers, project_doc_max_bytes: 32 * 1024, project_doc_fallback_filenames: Vec::new(),