diff --git a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs index b5ca558b2b..56b9c941c0 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs @@ -8,6 +8,7 @@ use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeInput; +use codex_mcp::McpStartupPolicy; use codex_mcp::ToolInfo; use codex_mcp::effective_mcp_servers; use codex_mcp::host_owned_codex_apps_enabled; @@ -91,6 +92,7 @@ impl AppsRequestProcessor { host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) .then(|| Arc::clone(&self.auth_manager)); let runtime = McpRuntime::new(McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::clone(&mcp_config), plugins_available: false, ready_selected_capability_roots: Vec::new(), diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index b2c8a01745..3beb4d4895 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -43,8 +43,10 @@ use crate::rmcp_client::prepare_codex_apps_tools_for_model; use crate::rmcp_client::prepare_regular_mcp_tools_for_model; use crate::runtime::McpPublicationGate; use crate::runtime::McpRuntimeInput; +use crate::runtime::McpStartupPolicy; use crate::server::McpServerConnectionIdentity; use crate::server::McpServerMetadata; +use crate::tool_catalog_cache::McpToolCatalogCacheContext; use crate::tools::ToolFilter; use crate::tools::ToolInfo; use crate::tools::filter_tools; @@ -65,12 +67,14 @@ use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::determine_streamable_http_auth_status_from_credentials; use tokio::sync::Mutex; use tokio::sync::RwLock; +use tokio::sync::watch; use tokio::task::JoinSet; use tracing::warn; pub(crate) struct McpServerConnection { identity: Option, client: AsyncManagedClient, + startup_trigger: Option>, } impl McpServerConnection { @@ -102,6 +106,9 @@ impl McpServerConnection { } pub(crate) async fn client(&self) -> Result { + if let Some(startup_trigger) = &self.startup_trigger { + startup_trigger.send_replace(true); + } self.client.client().await } @@ -110,10 +117,16 @@ impl McpServerConnection { } fn cancel_startup(&self) { - if !self.client.startup_complete.load(Ordering::Acquire) { + if !self.startup_is_dormant() && !self.client.startup_complete.load(Ordering::Acquire) { self.client.cancel_token.cancel(); } } + + fn startup_is_dormant(&self) -> bool { + self.startup_trigger + .as_ref() + .is_some_and(|startup_trigger| !*startup_trigger.borrow()) + } } impl Drop for McpServerConnection { @@ -170,6 +183,7 @@ impl McpConnectionSet { elicitation_router: ElicitationRequestRouter, ) -> Self { let McpRuntimeInput { + startup_policy, config, plugins_available: _, ready_selected_capability_roots: _, @@ -207,6 +221,9 @@ impl McpConnectionSet { required_servers.sort(); let mut reused_ready = Vec::new(); let mut join_set = JoinSet::new(); + // Explicit reconnects have no previous set and must replace their clients eagerly. + let allow_deferred_startup = + startup_policy == McpStartupPolicy::LazyWhenCached && previous.is_some(); let reusable_previous = previous.filter(|previous| { !previous.servers.is_empty() && previous.elicitation_requests.update( @@ -397,12 +414,32 @@ impl McpConnectionSet { protocol_mode, catalog_item_limit, ); + let defer_startup = allow_deferred_startup + && !configured_config.required + && !tool_plugin_provenance.is_selected_plugin_mcp_server(&server_name) + && async_managed_client + .tool_catalog_cache_context + .as_ref() + .and_then(McpToolCatalogCacheContext::current_tools) + .is_some_and(|tools| { + tools.into_iter().any(|tool| { + configured_tool_filter.allows(&tool.tool.name) + && tool_is_model_visible(&tool) + }) + }); + let (startup_trigger, startup_receiver) = if defer_startup { + let (trigger, receiver) = watch::channel(false); + (Some(trigger), Some(receiver)) + } else { + (None, None) + }; servers.insert( server_name.clone(), McpServerView { connection: Arc::new(McpServerConnection { identity: Some(connection_identity), client: async_managed_client.clone(), + startup_trigger, }), metadata, tool_filter: configured_tool_filter, @@ -413,7 +450,15 @@ impl McpConnectionSet { let tx_event = tx_event.clone(); let submit_id = startup_submit_id.clone(); let publication_gate = publication_gate.clone(); - join_set.spawn(async move { + let startup = async move { + if let Some(mut startup_receiver) = startup_receiver + && tokio::select! { + started = startup_receiver.wait_for(|started| *started) => started.is_err(), + () = cancel_token.cancelled() => true, + } + { + return (server_name, Err(StartupOutcomeError::Cancelled)); + } if !publication_gate.wait().await { return (server_name, Err(StartupOutcomeError::Cancelled)); } @@ -506,7 +551,13 @@ impl McpConnectionSet { } (server_name, outcome) - }); + }; + if defer_startup { + // Dormant servers must not hold the initial startup summary open. + tokio::spawn(startup); + } else { + join_set.spawn(startup); + } } let manager = Self { servers, @@ -690,7 +741,7 @@ impl McpConnectionSet { server_infos.insert(server_name.clone(), server_info); continue; } - match client.client().await { + match view.connection.client().await { Ok(managed_client) => { server_infos.insert(server_name.clone(), managed_client.server_info); } 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 341b5d7100..d9ac4e46d2 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -177,9 +177,6 @@ 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_startup_deadline = *self - .optional_startup_deadline - .get_or_init(|| tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE); join_all(self.servers.iter().map(|(server_name, view)| async move { if !view .connection @@ -199,6 +196,13 @@ impl McpConnectionSet { return; } if !must_wait_for_startup { + let optional_startup_deadline = if view.connection.startup_is_dormant() { + tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE + } else { + *self.optional_startup_deadline.get_or_init(|| { + tokio::time::Instant::now() + OPTIONAL_MCP_STARTUP_GRACE + }) + }; let startup_deadline = view .connection .client @@ -206,7 +210,7 @@ impl McpConnectionSet { .as_ref() .map(|cache| cache.optional_startup_deadline(optional_startup_deadline)) .unwrap_or(optional_startup_deadline); - if tokio::time::timeout_at(startup_deadline, view.connection.client.client()) + if tokio::time::timeout_at(startup_deadline, view.connection.client()) .await .is_err() { @@ -214,7 +218,7 @@ impl McpConnectionSet { } return; } - let _ = view.connection.client.client().await; + let _ = view.connection.client().await; } })) .await; @@ -241,7 +245,7 @@ impl McpConnectionSet { return Some((server_name.clone(), None, server_tools)); } view.connection.client.reconnect_failed_startup().await; - let Ok(mut client) = view.connection.client.client().await else { + let Ok(mut client) = view.connection.client().await else { trace!(server_name = %server_name, "omitting MCP server without an exact ready client"); return None; }; diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 2197468b60..87b9d2e138 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -115,6 +115,7 @@ impl McpConnectionSet { connection: Arc::new(McpServerConnection { identity: None, client, + startup_trigger: None, }), metadata: McpServerMetadata { environment_id: String::new(), @@ -1325,6 +1326,7 @@ async fn codex_apps_extension_does_not_share_host_owned_tools_cache() -> anyhow: /*previous*/ None, McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(config), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -1635,6 +1637,9 @@ async fn tool_catalog_cache_sanitizes_tools_and_tracks_environment_generation() let first_environment = resolve_environment(); let first_environment_weak = Arc::downgrade(&first_environment); let first_context = cache_context(&first_environment); + first_context.publish_if_newest(first_context.begin_fetch(), &[]); + assert!(!first_context.has_tools()); + let mut tool = create_test_tool("docs", "search"); tool.tool.annotations = Some(rmcp::model::ToolAnnotations::new().read_only(true)); first_context.publish_if_newest(first_context.begin_fetch(), &[tool]); @@ -2062,12 +2067,10 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() deadline_after_publication, "publishing a catalog must not install a stale startup deadline" ); - let cached = tokio::time::timeout( - Duration::from_millis(1), - capture_binding(&create_connection_set()), - ) - .await - .expect("cached tools should be immediately available to later threads"); + let cached_manager = create_connection_set(); + let cached = tokio::time::timeout(Duration::from_millis(1), capture_binding(&cached_manager)) + .await + .expect("cached tools should be immediately available to later threads"); assert_eq!( cached .tools() @@ -2079,12 +2082,9 @@ async fn capture_binding_shares_optional_startup_grace_across_connection_sets() tokio::time::advance(Duration::from_secs(30 * 60 + 1)).await; assert!( - tokio::time::timeout( - Duration::from_millis(1), - capture_binding(&create_connection_set()), - ) - .await - .is_err(), + tokio::time::timeout(Duration::from_millis(1), capture_binding(&cached_manager)) + .await + .is_err(), "an expired catalog should receive a fresh startup grace" ); @@ -3070,6 +3070,7 @@ async fn executor_owned_chatgpt_mcp_accepts_only_safe_explicit_authorization() - /*previous*/ None, McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(runtime_config.clone()), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -3182,6 +3183,7 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { /*previous*/ None, McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(crate::mcp::tests::test_mcp_config( codex_home.path().to_path_buf(), )), @@ -3482,6 +3484,7 @@ 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_trigger: None, }), metadata: McpServerMetadata::from(&server), tool_filter: ToolFilter::from_config(config), @@ -3503,6 +3506,7 @@ async fn reconcile_reusable_server( Some(previous), McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(crate::mcp::tests::test_mcp_config( codex_home.path().to_path_buf(), )), @@ -3613,6 +3617,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a startup_reconnect: None, cancel_token: CancellationToken::new(), }, + startup_trigger: None, }), metadata: McpServerMetadata::from(&server), tool_filter: ToolFilter::from_config(&config), @@ -3726,6 +3731,7 @@ async fn reconciliation_replaces_connection_when_protocol_mode_changes() { Some(&previous), McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(mcp_config), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -3783,6 +3789,7 @@ async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabl Some(&previous), McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(mcp_config), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -3917,6 +3924,7 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> { startup_reconnect: None, cancel_token: CancellationToken::new(), }, + startup_trigger: None, }); assert!(!client.is_closed().await); diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index 43fb01387e..add33ffede 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -15,6 +15,7 @@ pub use rmcp_client::MCP_SANDBOX_STATE_META_CAPABILITY; pub use runtime::McpRuntime; pub use runtime::McpRuntimeContext; pub use runtime::McpRuntimeInput; +pub use runtime::McpStartupPolicy; pub use runtime::SandboxState; pub use tool_catalog_cache::McpToolCatalogCache; pub use tools::ToolInfo; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 6bddbf24c8..837c480524 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -53,6 +53,7 @@ use crate::connection_manager::McpConnectionSet; use crate::runtime::McpPublicationGate; use crate::runtime::McpRuntimeContext; use crate::runtime::McpRuntimeInput; +use crate::runtime::McpStartupPolicy; use crate::server::EffectiveMcpServer; use crate::tools::ToolInfo; @@ -359,6 +360,7 @@ pub async fn read_mcp_resource( /*previous*/ None, McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(runtime_config), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -436,6 +438,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( /*previous*/ None, McpPublicationGate::already_published(), McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::new(runtime_config), plugins_available: false, ready_selected_capability_roots: Vec::new(), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 1c326bffc7..5835e34157 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -316,10 +316,10 @@ impl ManagedClientStartup { .startup_timeout_sec .unwrap_or(DEFAULT_STARTUP_TIMEOUT); let cancel_token_for_fut = cancel_token; - let tool_catalog_fetch_ticket = tool_catalog_cache_context - .as_ref() - .map(McpToolCatalogCacheContext::begin_fetch); async move { + let tool_catalog_fetch_ticket = tool_catalog_cache_context + .as_ref() + .map(McpToolCatalogCacheContext::begin_fetch); let refresh_start = is_codex_apps_mcp_server.then(Instant::now); let outcome = match async { if let Err(error) = validate_mcp_server_name(&server_name) { diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index e9feec69ec..6d9ba4316f 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -48,8 +48,18 @@ use crate::server::EffectiveMcpServer; use crate::tool_catalog_cache::McpToolCatalogCache; use crate::tools::ToolInfo; +/// Controls when one task starts its optional MCP servers. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum McpStartupPolicy { + /// Start configured servers when their task's MCP runtime is published. + Eager, + /// Start optional servers with cached tool definitions on first use. + LazyWhenCached, +} + /// Everything needed to materialize one exact MCP configuration. pub struct McpRuntimeInput { + pub startup_policy: McpStartupPolicy, pub config: Arc, pub plugins_available: bool, pub ready_selected_capability_roots: Vec, diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index 3cd4f6004c..67fb07db45 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -107,7 +107,7 @@ impl Default for ToolCatalogCacheEntry { impl McpToolCatalogCacheContext { pub(crate) fn has_tools(&self) -> bool { - self.current_tools().is_some() + self.current_tools().is_some_and(|tools| !tools.is_empty()) } pub(crate) fn optional_startup_deadline(&self, default_deadline: Instant) -> Instant { diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 0fadc2c924..af236b144e 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -34,6 +34,7 @@ use codex_mcp::MCP_TOOL_CODEX_APPS_META_KEY; use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; use codex_mcp::McpRuntimeInput; +use codex_mcp::McpStartupPolicy; use codex_mcp::ToolInfo; use codex_mcp::ToolPluginProvenance; use codex_mcp::effective_mcp_servers; @@ -246,6 +247,7 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( codex_mcp::host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()) .then(|| Arc::clone(&auth_manager)); let mcp_runtime = McpRuntime::new(McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: Arc::clone(&mcp_config), plugins_available: false, ready_selected_capability_roots: Vec::new(), diff --git a/codex-rs/core/src/session/mcp_runtime.rs b/codex-rs/core/src/session/mcp_runtime.rs index 77efaae831..62d5d4638a 100644 --- a/codex-rs/core/src/session/mcp_runtime.rs +++ b/codex-rs/core/src/session/mcp_runtime.rs @@ -8,6 +8,7 @@ use super::session::SessionConfiguration; use super::*; use crate::mcp::McpRuntimeProjection; use codex_mcp::ElicitationReviewerHandle; +use codex_mcp::McpStartupPolicy; use codex_mcp::PreparedMcpCall; use codex_protocol::capabilities::SelectedCapabilityRoot; @@ -174,6 +175,11 @@ impl Session { .then(|| Arc::clone(&self.services.auth_manager)); McpRuntimeInput { + startup_policy: if matches!(desired.session_source, SessionSource::SubAgent(_)) { + McpStartupPolicy::LazyWhenCached + } else { + McpStartupPolicy::Eager + }, config: mcp_config, plugins_available, ready_selected_capability_roots: ready_selected_capability_roots.to_vec(), diff --git a/codex-rs/core/tests/suite/mcp_auth_refresh.rs b/codex-rs/core/tests/suite/mcp_auth_refresh.rs index 83386949f3..8929eca9f5 100644 --- a/codex-rs/core/tests/suite/mcp_auth_refresh.rs +++ b/codex-rs/core/tests/suite/mcp_auth_refresh.rs @@ -16,6 +16,7 @@ use codex_mcp::EffectiveMcpServer; use codex_mcp::McpRuntime; use codex_mcp::McpRuntimeContext; use codex_mcp::McpRuntimeInput; +use codex_mcp::McpStartupPolicy; use codex_mcp::McpToolCatalogCache; use codex_protocol::mcp::ClientMcpExtensions; use codex_protocol::protocol::AskForApproval; @@ -89,6 +90,7 @@ async fn hosted_plugin_runtime_ps_mcp_tool_calls_use_current_auth_manager_token( let plugins_manager = codex_core_plugins::PluginsManager::new(home.path().to_path_buf()); let mcp_config = Arc::new(config.to_mcp_config(&plugins_manager).await); let runtime = McpRuntime::new(McpRuntimeInput { + startup_policy: McpStartupPolicy::Eager, config: mcp_config, plugins_available: false, ready_selected_capability_roots: Vec::new(), diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index ab6cf0a51f..7d257d5714 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -12,7 +12,10 @@ use codex_protocol::models::PermissionProfile; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpInvocation; +use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::Op; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::SubAgentSource; use codex_protocol::user_input::UserInput; use codex_utils_path_uri::PathUri; use core_test_support::responses; @@ -238,7 +241,7 @@ async fn mcp_calls_stay_bound_to_each_thread() -> anyhow::Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow::Result<()> { +async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow::Result<()> { skip_if_wine_exec!( Ok(()), "requires a Windows test_stdio_server in the Wine-exec environment" @@ -313,6 +316,18 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: &format!("Echo from {first_process}."), ); + let NewThread { + thread: eager_thread, + .. + } = fixture + .thread_manager + .start_thread(StartThreadOptions::new(fixture.config.clone())) + .await?; + let eager_pid = wait_for_new_pid(fs.as_ref(), &pid_file, Some(&first_pid)).await?; + wait_for_mcp_server(&eager_thread, SERVER_NAME).await?; + eager_thread.shutdown_and_wait().await?; + let cached_process = process_label(&eager_pid); + fs.remove( &barrier_file, RemoveOptions { @@ -333,10 +348,56 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: .. } = fixture .thread_manager - .start_thread(StartThreadOptions::new(fixture.config.clone())) + .start_thread(StartThreadOptions { + session_source: Some(SessionSource::SubAgent(SubAgentSource::ThreadSpawn { + parent_thread_id: fixture.session_configured.thread_id, + depth: 1, + agent_path: None, + agent_nickname: None, + agent_role: None, + })), + ..StartThreadOptions::new(fixture.config.clone()) + }) .await?; - let second_pid = wait_for_new_pid(fs.as_ref(), &pid_file, Some(&first_pid)).await?; - let second_process = process_label(&second_pid); + second_thread.submit(Op::Interrupt).await?; + + let unused_response = mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("unused"), + responses::ev_assistant_message("unused-message", "done"), + responses::ev_completed("unused"), + ]), + ) + .await; + second_thread + .submit(user_turn("Do not call any MCP tools.")) + .await?; + let mut reported_ready_before_startup = false; + wait_for_event(&second_thread, |event| { + if let EventMsg::McpStartupUpdate(update) = event + && update.server == SERVER_NAME + && matches!(update.status, McpStartupStatus::Ready) + { + reported_ready_before_startup = true; + } + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + assert!( + !reported_ready_before_startup, + "a dormant MCP server must not be reported as ready" + ); + assert_definition( + &unused_response, + &format!("Tools in the {NAMESPACE} namespace."), + &format!("Echo from {cached_process}."), + ); + assert_eq!( + fs.read_file_text(&pid_file, /*sandbox*/ None).await?.trim(), + eager_pid, + "cached tool definitions should not start an unused subagent-owned MCP process" + ); let app_only_call_id = "cached-app-only-call"; let unrelated_call_id = "cached-unrelated-plan"; @@ -376,18 +437,31 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: .submit(user_turn("call the echo and cwd tools")) .await?; let mut unrelated_finished_tx = Some(unrelated_finished_tx); + let mut saw_starting = false; + let mut saw_ready = false; let end = wait_for_event(&second_for_turn, |event| { if matches!(event, EventMsg::PlanUpdate(_)) && let Some(sender) = unrelated_finished_tx.take() { let _ = sender.send(()); } + if let EventMsg::McpStartupUpdate(update) = event + && update.server == SERVER_NAME + { + saw_starting |= matches!(update.status, McpStartupStatus::Starting); + saw_ready |= matches!(update.status, McpStartupStatus::Ready); + } matches!( event, EventMsg::McpToolCallEnd(end) if end.call_id == "cached-call" ) }) .await; + assert!( + saw_starting, + "deferred startup should emit its starting status" + ); + assert!(saw_ready, "deferred startup should emit its ready status"); let EventMsg::McpToolCallEnd(end) = end else { unreachable!("event predicate guarantees an MCP tool result"); }; @@ -414,8 +488,10 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: assert_definition( &cached_response, &format!("Tools in the {NAMESPACE} namespace."), - &format!("Echo from {first_process}."), + &format!("Echo from {cached_process}."), ); + let second_pid = wait_for_new_pid(fs.as_ref(), &pid_file, Some(&eager_pid)).await?; + let second_process = process_label(&second_pid); tokio::time::timeout(Duration::from_secs(2), unrelated_finished_rx) .await .context("an unrelated tool should complete while cached MCP startup is pending")? @@ -451,6 +527,135 @@ async fn regular_mcp_definition_cache_preserves_live_session_state() -> anyhow:: ); second_thread.shutdown_and_wait().await?; + let mut filtered_config = fixture.config.clone(); + let mut filtered_servers = filtered_config.mcp_servers.get().clone(); + filtered_servers + .get_mut(SERVER_NAME) + .expect("cached MCP server should remain configured") + .enabled_tools = Some(vec!["cwd".to_string()]); + filtered_config.mcp_servers.set(filtered_servers)?; + let NewThread { + thread: filtered_thread, + .. + } = fixture + .thread_manager + .start_thread(StartThreadOptions { + session_source: Some(SessionSource::SubAgent(SubAgentSource::Other( + "filtered-cached-startup".to_string(), + ))), + ..StartThreadOptions::new(filtered_config) + }) + .await?; + let filtered_pid = wait_for_new_pid(fs.as_ref(), &pid_file, Some(&second_pid)).await?; + wait_for_mcp_server(&filtered_thread, SERVER_NAME).await?; + filtered_thread.shutdown_and_wait().await?; + fs.remove( + &barrier_file, + RemoveOptions { + recursive: false, + force: false, + }, + /*sandbox*/ None, + ) + .await?; + let NewThread { + thread: interrupted_thread, + .. + } = fixture + .thread_manager + .start_thread(StartThreadOptions { + session_source: Some(SessionSource::SubAgent(SubAgentSource::Other( + "interrupted-cached-startup".to_string(), + ))), + ..StartThreadOptions::new(fixture.config.clone()) + }) + .await?; + mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("interrupted-startup"), + responses::ev_function_call_with_namespace( + "interrupted-startup-call", + NAMESPACE, + "echo", + r#"{"message":"interrupted"}"#, + ), + responses::ev_completed("interrupted-startup"), + ]), + ) + .await; + interrupted_thread + .submit(user_turn("Start the cached MCP tool.")) + .await?; + let interrupted_pid = wait_for_new_pid(fs.as_ref(), &pid_file, Some(&filtered_pid)).await?; + wait_for_event(&interrupted_thread, |event| { + matches!( + event, + EventMsg::McpStartupUpdate(update) + if update.server == SERVER_NAME + && matches!(update.status, McpStartupStatus::Starting) + ) + }) + .await; + interrupted_thread.submit(Op::Interrupt).await?; + wait_for_event(&interrupted_thread, |event| { + matches!(event, EventMsg::TurnAborted(_)) + }) + .await; + fs.write_file(&barrier_file, b"ready".to_vec(), /*sandbox*/ None) + .await?; + tokio::time::timeout( + Duration::from_secs(2), + wait_for_event(&interrupted_thread, |event| { + matches!( + event, + EventMsg::McpStartupUpdate(update) + if update.server == SERVER_NAME + && matches!(update.status, McpStartupStatus::Ready) + ) + }), + ) + .await + .context("deferred MCP startup should survive an interrupted first tool call")?; + mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("retried-startup"), + responses::ev_function_call_with_namespace( + "retried-startup-call", + NAMESPACE, + "echo", + r#"{"message":"retried"}"#, + ), + responses::ev_completed("retried-startup"), + ]), + ) + .await; + let retry_done = mount_sse_once( + &responses_server, + responses::sse(vec![ + responses::ev_response_created("retried-done"), + responses::ev_assistant_message("retried-message", "done"), + responses::ev_completed("retried-done"), + ]), + ) + .await; + interrupted_thread + .submit(user_turn("Retry the cached MCP tool.")) + .await?; + wait_for_event(&interrupted_thread, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + let retry_output = retry_done + .single_request() + .function_call_output_text("retried-startup-call") + .expect("the retried MCP tool call should return its output"); + assert!( + retry_output.contains(&process_label(&interrupted_pid)), + "the retry should use the server started by the interrupted call" + ); + interrupted_thread.shutdown_and_wait().await?; responses_server.verify().await; Ok(()) }