Start cached MCP servers lazily for subagents (#37261)

## Why

Subagents can reuse cached MCP tool definitions without starting every optional
server before any of its tools are used.

## What changed

- Defer startup for subagent-owned, optional MCP servers when their cache has at
  least one visible, allowed tool, and trigger startup on the first operation
  that needs the live client.
- Keep root sessions, required servers, selected plugin servers, and explicit
  reconnects eager.
- Exclude empty catalogs from cache eligibility and keep dormant servers out of
  the initial startup summary.

## Testing

- Extend MCP tool-cache coverage for eager root startup, unused and filtered
  subagent servers, deferred startup status events, and startup that survives an
  interrupted first tool call.

GitOrigin-RevId: e7528ff08312583a6747cbc6550ea5b6eabda9a6
This commit is contained in:
jif
2026-08-06 11:18:24 +00:00
committed by copyberry
parent a17da5e6e4
commit 1151b23f01
13 changed files with 325 additions and 31 deletions
@@ -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(),
+55 -4
View File
@@ -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<McpServerConnectionIdentity>,
client: AsyncManagedClient,
startup_trigger: Option<watch::Sender<bool>>,
}
impl McpServerConnection {
@@ -102,6 +106,9 @@ impl McpServerConnection {
}
pub(crate) async fn client(&self) -> Result<ManagedClient, StartupOutcomeError> {
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);
}
@@ -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;
};
@@ -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);
+1
View File
@@ -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;
+3
View File
@@ -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(),
+3 -3
View File
@@ -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) {
+10
View File
@@ -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<McpConfig>,
pub plugins_available: bool,
pub ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
+1 -1
View File
@@ -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 {
+2
View File
@@ -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(),
+6
View File
@@ -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(),
@@ -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(),
+210 -5
View File
@@ -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(())
}