mirror of
https://github.com/openai/codex.git
synced 2026-10-18 17:18:52 +00:00
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:
@@ -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(),
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user