mirror of
https://github.com/openai/codex.git
synced 2026-11-03 19:53:25 +00:00
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
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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"
|
||||
);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -134,6 +134,11 @@ pub struct McpConfig {
|
||||
pub mcp_oauth_callback_port: Option<u16>,
|
||||
/// Optional OAuth redirect URI override for MCP login.
|
||||
pub mcp_oauth_callback_url: Option<String>,
|
||||
/// 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(
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -56,11 +56,16 @@ struct ToolCatalogCacheEntry {
|
||||
#[derive(Default)]
|
||||
struct ToolCatalogCacheState {
|
||||
snapshot: Option<ToolCatalogSnapshot>,
|
||||
optional_startup_deadline: Option<Instant>,
|
||||
optional_startup_deadline: Option<OptionalStartupDeadline>,
|
||||
last_accepted_generation: u64,
|
||||
disabled_by_server: bool,
|
||||
}
|
||||
|
||||
struct OptionalStartupDeadline {
|
||||
grace: Duration,
|
||||
deadline: Instant,
|
||||
}
|
||||
|
||||
struct ToolCatalogSnapshot {
|
||||
tools: Vec<ToolInfo>,
|
||||
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<Vec<ToolInfo>> {
|
||||
|
||||
@@ -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<String>,
|
||||
|
||||
/// 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<u64>,
|
||||
|
||||
/// User-defined provider entries that extend the built-in list. Built-in
|
||||
/// IDs cannot be overridden.
|
||||
#[serde(default, deserialize_with = "deserialize_model_providers")]
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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#"
|
||||
|
||||
@@ -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<String>,
|
||||
|
||||
/// 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<String, ModelProviderInfo>,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -247,6 +247,7 @@ fn new_config(model: Option<String>, 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(),
|
||||
|
||||
Reference in New Issue
Block a user