mirror of
https://github.com/openai/codex.git
synced 2026-11-07 20:30:33 +00:00
Preserve MCP sharing discovery when deferring cached stdio startup (#60037)
## Why Subagents reusing cached MCP tool definitions were forced to use the legacy protocol, preventing connection sharing from being negotiated on first use. ## What changed Keep stdio protocol discovery enabled during deferred startup while allowing dormant subagents to advertise cached legacy tools. Track modern discovery separately from catalog updates so discovery clients reject older legacy catalogs while legacy clients can still reuse them. Keep modern tools out of the legacy catalog when joining pooled connections. ## Testing Add coverage for deferred subagents negotiating sharing on first use and reusing the same server process, while a legacy sibling retains cached definitions without starting its server. Extend cache coverage to ensure older legacy fetches cannot undo modern discovery. GitOrigin-RevId: 5f8f4deb22cfb81a08c05ff3fcd7ba5c9199551c
This commit is contained in:
@@ -681,6 +681,8 @@ impl McpConnectionSet {
|
||||
Some((&connection_identity, cache_mode, server.is_agent_plugin())),
|
||||
)?;
|
||||
context.track_tool_list_changes = config.mcp_tool_list_changed;
|
||||
context.before_stdio_discovery =
|
||||
expected_protocol_mode == Some(McpProtocolMode::DiscoverLocalStdio);
|
||||
Some(context)
|
||||
});
|
||||
let defer_startup = allow_deferred_startup
|
||||
@@ -695,14 +697,6 @@ impl McpConnectionSet {
|
||||
})
|
||||
})
|
||||
.is_some();
|
||||
// A subagent using legacy tool definitions must launch with the same handshake.
|
||||
let expected_protocol_mode = if defer_startup
|
||||
&& expected_protocol_mode == Some(McpProtocolMode::DiscoverLocalStdio)
|
||||
{
|
||||
Some(McpProtocolMode::Legacy)
|
||||
} else {
|
||||
expected_protocol_mode
|
||||
};
|
||||
let protocol_mode = expected_protocol_mode.unwrap_or(protocol_mode);
|
||||
let connection_pool = runtime_context.connection_pool.as_ref().filter(|_| {
|
||||
// Sharing a signed-in helper requires tracking account changes.
|
||||
@@ -954,6 +948,7 @@ impl McpConnectionSet {
|
||||
})
|
||||
.map(|manager| manager.auth_change_state_receiver()),
|
||||
protocol_mode,
|
||||
defer_startup,
|
||||
catalog_item_limit,
|
||||
connection_pool,
|
||||
publish_apps,
|
||||
|
||||
@@ -2843,6 +2843,9 @@ async fn tool_catalog_cache_sanitizes_tools_and_tracks_environment_generation()
|
||||
|
||||
let older = first_context.begin_fetch();
|
||||
let newer = first_context.begin_fetch();
|
||||
let mut discovery = first_context.clone();
|
||||
discovery.before_stdio_discovery = true;
|
||||
discovery.record_modern_stdio_discovery(&discovery.begin_fetch());
|
||||
first_context.publish_if_newest(
|
||||
newer,
|
||||
&[create_test_tool("docs", "new")],
|
||||
@@ -2857,6 +2860,10 @@ async fn tool_catalog_cache_sanitizes_tools_and_tracks_environment_generation()
|
||||
first_context.current_tools().expect("cached tools")[0].callable_name,
|
||||
"new"
|
||||
);
|
||||
assert!(
|
||||
!discovery.has_tools(),
|
||||
"older legacy fetches cannot undo modern discovery"
|
||||
);
|
||||
|
||||
tokio::time::advance(Duration::from_secs(30 * 60 + 1)).await;
|
||||
assert!(!first_context.has_tools());
|
||||
|
||||
@@ -495,6 +495,7 @@ impl AsyncManagedClient {
|
||||
client_mcp_extensions: ClientMcpExtensions,
|
||||
auth_changes: Option<watch::Receiver<AuthChangeState>>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
defer_startup: bool,
|
||||
catalog_item_limit: usize,
|
||||
connection_pool: Option<(crate::McpConnectionPool, ConnectionKey)>,
|
||||
publish_apps: Option<watch::Sender<Option<std::sync::Weak<RmcpClient>>>>,
|
||||
@@ -523,6 +524,17 @@ impl AsyncManagedClient {
|
||||
.as_ref()
|
||||
.and_then(|auth| auth.as_ref().ok())
|
||||
.map(|auth| auth.authority.clone());
|
||||
// Startup may publish a legacy catalog after negotiation. Only a dormant chat can
|
||||
// advertise it beforehand, and the pool must never put modern tools in that catalog.
|
||||
let pool_catalog = tool_catalog_cache_context
|
||||
.as_ref()
|
||||
.filter(|_| protocol_mode != McpProtocolMode::DiscoverLocalStdio)
|
||||
.cloned();
|
||||
let advertised_catalog = if defer_startup {
|
||||
tool_catalog_cache_context.clone()
|
||||
} else {
|
||||
pool_catalog.clone()
|
||||
};
|
||||
let startup = Arc::new(ManagedClientStartup {
|
||||
ema_auth,
|
||||
originator: AuthStorageOriginator::current(),
|
||||
@@ -534,7 +546,7 @@ impl AsyncManagedClient {
|
||||
tx_event,
|
||||
elicitation_requests,
|
||||
codex_apps_tools_cache_context: codex_apps_tools_cache_context.clone(),
|
||||
tool_catalog_cache_context: tool_catalog_cache_context.clone(),
|
||||
tool_catalog_cache_context,
|
||||
runtime_context,
|
||||
resolved_environment,
|
||||
runtime_auth_provider,
|
||||
@@ -547,10 +559,6 @@ impl AsyncManagedClient {
|
||||
startup_complete: Arc::clone(&startup_complete),
|
||||
server_capabilities: Arc::clone(&server_capabilities),
|
||||
});
|
||||
// A discovery attempt may publish a legacy catalog. It cannot read that catalog
|
||||
// before negotiation or publish a modern joiner's tool list into it.
|
||||
let tool_catalog_cache_context = tool_catalog_cache_context
|
||||
.filter(|_| protocol_mode != McpProtocolMode::DiscoverLocalStdio);
|
||||
let (client, connection_slot) = if let Some((pool, key)) = connection_pool {
|
||||
let slot = Arc::new(ConnectionSlot::default());
|
||||
let connection_slot = Arc::clone(&slot);
|
||||
@@ -558,14 +566,13 @@ impl AsyncManagedClient {
|
||||
let cancel_token = cancel_token.clone();
|
||||
let startup_complete = Arc::clone(&startup_complete);
|
||||
let server_capabilities = Arc::clone(&server_capabilities);
|
||||
let tool_catalog_cache_context = tool_catalog_cache_context.clone();
|
||||
let client = async move {
|
||||
let outcome = pool
|
||||
.connect(
|
||||
slot,
|
||||
key,
|
||||
server_capabilities,
|
||||
tool_catalog_cache_context,
|
||||
pool_catalog,
|
||||
cancel_token,
|
||||
|cancellation, sharing_declined| {
|
||||
let mut shared_startup = (*startup).clone();
|
||||
@@ -614,7 +621,7 @@ impl AsyncManagedClient {
|
||||
cached_server_info,
|
||||
server_capabilities,
|
||||
codex_apps_tools_cache_context,
|
||||
tool_catalog_cache_context,
|
||||
tool_catalog_cache_context: advertised_catalog,
|
||||
startup_complete,
|
||||
startup_reconnect,
|
||||
connection_slot,
|
||||
@@ -1183,7 +1190,7 @@ async fn start_server_task(
|
||||
&& let Some(context) = tool_catalog_cache_context.take()
|
||||
&& let Some(ticket) = &tool_catalog_fetch_ticket
|
||||
{
|
||||
context.invalidate_if_newest(ticket);
|
||||
context.record_modern_stdio_discovery(ticket);
|
||||
}
|
||||
let connection_sharing = sharing_negotiation
|
||||
.map(|negotiation| negotiation.finish(&initialize_result.protocol_version))
|
||||
|
||||
@@ -59,7 +59,8 @@ struct ToolCatalogCacheEntry {
|
||||
struct ToolCatalogCacheState {
|
||||
snapshot: Option<ToolCatalogSnapshot>,
|
||||
optional_startup_deadline: Option<OptionalStartupDeadline>,
|
||||
last_accepted_generation: u64,
|
||||
last_catalog_generation: u64,
|
||||
last_modern_discovery_generation: u64,
|
||||
disabled_by_server: bool,
|
||||
}
|
||||
|
||||
@@ -78,6 +79,7 @@ struct ToolCatalogSnapshot {
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct McpToolCatalogCacheContext {
|
||||
pub(crate) track_tool_list_changes: bool,
|
||||
pub(crate) before_stdio_discovery: bool,
|
||||
entry: Arc<ToolCatalogCacheEntry>,
|
||||
}
|
||||
|
||||
@@ -110,6 +112,7 @@ impl McpToolCatalogCache {
|
||||
Some(McpToolCatalogCacheContext {
|
||||
entry,
|
||||
track_tool_list_changes: false,
|
||||
before_stdio_discovery: false,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -124,7 +127,16 @@ impl Default for ToolCatalogCacheEntry {
|
||||
}
|
||||
|
||||
impl McpToolCatalogCacheContext {
|
||||
fn accepts_snapshot(&self, snapshot: &ToolCatalogSnapshot) -> bool {
|
||||
fn accepts_snapshot(
|
||||
&self,
|
||||
state: &ToolCatalogCacheState,
|
||||
snapshot: &ToolCatalogSnapshot,
|
||||
) -> bool {
|
||||
if self.before_stdio_discovery
|
||||
&& state.last_modern_discovery_generation > state.last_catalog_generation
|
||||
{
|
||||
return false;
|
||||
}
|
||||
// Enabling refresh must not reuse a catalog published without notification tracking.
|
||||
snapshot
|
||||
.tool_list_revision
|
||||
@@ -147,11 +159,11 @@ impl McpToolCatalogCacheContext {
|
||||
let state = lock_unpoisoned(&self.entry.state);
|
||||
let snapshot = state.snapshot.as_ref()?;
|
||||
(!state.disabled_by_server
|
||||
&& self.accepts_snapshot(snapshot)
|
||||
&& self.accepts_snapshot(&state, snapshot)
|
||||
&& !snapshot.tools.is_empty()
|
||||
&& snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL
|
||||
&& accepts_tools(&snapshot.tools))
|
||||
.then_some(state.last_accepted_generation)
|
||||
.then_some(state.last_catalog_generation)
|
||||
}
|
||||
|
||||
pub(crate) fn optional_startup_deadline(
|
||||
@@ -162,7 +174,7 @@ impl McpToolCatalogCacheContext {
|
||||
let mut state = lock_unpoisoned(&self.entry.state);
|
||||
if state.disabled_by_server
|
||||
|| state.snapshot.as_ref().is_some_and(|snapshot| {
|
||||
self.accepts_snapshot(snapshot)
|
||||
self.accepts_snapshot(&state, snapshot)
|
||||
&& snapshot.published_at.elapsed() <= TOOL_CATALOG_CACHE_TTL
|
||||
})
|
||||
{
|
||||
@@ -207,7 +219,7 @@ impl McpToolCatalogCacheContext {
|
||||
|| state
|
||||
.snapshot
|
||||
.as_ref()
|
||||
.is_some_and(|snapshot| !self.accepts_snapshot(snapshot))
|
||||
.is_some_and(|snapshot| !self.accepts_snapshot(&state, snapshot))
|
||||
{
|
||||
return None;
|
||||
}
|
||||
@@ -236,10 +248,17 @@ impl McpToolCatalogCacheContext {
|
||||
state.snapshot = None;
|
||||
}
|
||||
|
||||
/// An observed protocol change invalidates a legacy catalog, including any older fetch in flight.
|
||||
pub(crate) fn invalidate_if_newest(&self, ticket: &McpToolCatalogFetchTicket) {
|
||||
if let Some(mut state) = self.accept_generation(ticket.generation) {
|
||||
state.snapshot = None;
|
||||
/// A modern handshake rules out older legacy catalogs for discovery. Clients forced to use
|
||||
/// the legacy protocol can still use those catalogs, including legacy fetches in flight.
|
||||
pub(crate) fn record_modern_stdio_discovery(&self, ticket: &McpToolCatalogFetchTicket) {
|
||||
let mut state = lock_unpoisoned(&self.entry.state);
|
||||
if ticket.generation
|
||||
> state
|
||||
.last_catalog_generation
|
||||
.max(state.last_modern_discovery_generation)
|
||||
{
|
||||
state.last_modern_discovery_generation = ticket.generation;
|
||||
state.optional_startup_deadline = None;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -249,9 +268,12 @@ impl McpToolCatalogCacheContext {
|
||||
tools: &[ToolInfo],
|
||||
meta: Option<Arc<serde_json::Value>>,
|
||||
) {
|
||||
let Some(mut state) = self.accept_generation(ticket.generation) else {
|
||||
let mut state = lock_unpoisoned(&self.entry.state);
|
||||
if state.disabled_by_server || ticket.generation <= state.last_catalog_generation {
|
||||
return;
|
||||
};
|
||||
}
|
||||
state.last_catalog_generation = ticket.generation;
|
||||
state.optional_startup_deadline = None;
|
||||
let mut tools = tools.to_vec();
|
||||
for tool in &mut tools {
|
||||
// Tool annotations affect approval and parallelism decisions, so only the live
|
||||
@@ -265,16 +287,6 @@ impl McpToolCatalogCacheContext {
|
||||
tool_list_revision: ticket.tool_list_revision,
|
||||
});
|
||||
}
|
||||
|
||||
fn accept_generation(&self, generation: u64) -> Option<MutexGuard<'_, ToolCatalogCacheState>> {
|
||||
let mut state = lock_unpoisoned(&self.entry.state);
|
||||
if state.disabled_by_server || generation <= state.last_accepted_generation {
|
||||
return None;
|
||||
}
|
||||
state.last_accepted_generation = generation;
|
||||
state.optional_startup_deadline = None;
|
||||
Some(state)
|
||||
}
|
||||
}
|
||||
|
||||
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
|
||||
|
||||
@@ -747,11 +747,6 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
||||
/*sandbox*/ None,
|
||||
)
|
||||
.await?;
|
||||
// The subagent can reuse the parent's legacy catalog even when discovery is enabled.
|
||||
let mut subagent_config = fixture.config.clone();
|
||||
subagent_config
|
||||
.features
|
||||
.enable(Feature::McpConnectionSharing)?;
|
||||
let NewThread {
|
||||
thread: second_thread,
|
||||
..
|
||||
@@ -765,7 +760,7 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
||||
agent_nickname: None,
|
||||
agent_role: None,
|
||||
})),
|
||||
..StartThreadOptions::new(subagent_config)
|
||||
..StartThreadOptions::new(fixture.config.clone())
|
||||
})
|
||||
.await?;
|
||||
second_thread.submit(Op::Interrupt).await?;
|
||||
@@ -1158,3 +1153,118 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow
|
||||
responses_server.verify().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn cached_stdio_negotiates_sharing_on_first_use() -> anyhow::Result<()> {
|
||||
skip_if_no_network!(Ok(()));
|
||||
core_test_support::skip_if_remote!(Ok(()), "the shared stdio fixture runs on the host");
|
||||
skip_if_wine_exec!(Ok(()), "requires the native stdio server");
|
||||
|
||||
let responses_server = responses::start_mock_server().await;
|
||||
let command = core_test_support::stdio_server_bin()?;
|
||||
let fixture = test_codex()
|
||||
.with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false)
|
||||
.with_config(move |config| {
|
||||
config
|
||||
.features
|
||||
.disable(Feature::Mcp20260728)
|
||||
.expect("legacy protocol");
|
||||
config
|
||||
.features
|
||||
.disable(Feature::McpConnectionSharing)
|
||||
.expect("disable discovery");
|
||||
config.permissions.approval_policy = Constrained::allow_any(AskForApproval::Never);
|
||||
config
|
||||
.permissions
|
||||
.set_permission_profile(PermissionProfile::Disabled)
|
||||
.expect("test permissions");
|
||||
config
|
||||
.mcp_servers
|
||||
.set(
|
||||
serde_json::from_value(json!({
|
||||
"cached_rmcp": {
|
||||
"command": command,
|
||||
"cwd": config.cwd,
|
||||
"env": { "MCP_TEST_DYNAMIC_SERVER_METADATA": "1" },
|
||||
"enabled_tools": ["echo"]
|
||||
}
|
||||
}))
|
||||
.expect("stdio config"),
|
||||
)
|
||||
.expect("set stdio config");
|
||||
})
|
||||
.build_with_auto_env(&responses_server)
|
||||
.await?;
|
||||
wait_for_mcp_server(&fixture.codex, SERVER_NAME).await?;
|
||||
let mut config = fixture.config.clone();
|
||||
config.features.enable(Feature::McpConnectionSharing)?;
|
||||
let start_child = |config| {
|
||||
fixture.thread_manager.start_thread(StartThreadOptions {
|
||||
session_source: Some(SessionSource::SubAgent(SubAgentSource::Other(
|
||||
"cached-sharing".to_string(),
|
||||
))),
|
||||
..StartThreadOptions::new(config)
|
||||
})
|
||||
};
|
||||
let first = start_child(config.clone()).await?.thread;
|
||||
let second = start_child(config).await?.thread;
|
||||
let legacy_sibling = start_child(fixture.config.clone()).await?.thread;
|
||||
for thread in [&first, &second, &legacy_sibling] {
|
||||
let (config, _) = thread.current_mcp_config_and_runtime_context().await;
|
||||
assert_eq!(
|
||||
thread
|
||||
.mcp_connection_statuses(&config)
|
||||
.await
|
||||
.get(SERVER_NAME),
|
||||
Some(&McpServerConnectionStatus::NotStarted)
|
||||
);
|
||||
}
|
||||
|
||||
let mut processes = Vec::new();
|
||||
for thread in [&fixture.codex, &first, &second] {
|
||||
let result = thread
|
||||
.call_mcp_tool(
|
||||
SERVER_NAME,
|
||||
"echo",
|
||||
Some(json!({"message":"identify process"})),
|
||||
/*meta*/ None,
|
||||
)
|
||||
.await?;
|
||||
let content = result.structured_content.expect("echo result");
|
||||
processes.push(content["echo"].as_str().expect("process label").to_owned());
|
||||
}
|
||||
assert_ne!(processes[0], processes[1]);
|
||||
assert_eq!(processes[1], processes[2]);
|
||||
|
||||
let response =
|
||||
responses::mount_sse_once(&responses_server, responses::sse_completed("unrelated")).await;
|
||||
legacy_sibling
|
||||
.start_or_steer_turn(user_turn("Continue without calling any tools."))
|
||||
.await?;
|
||||
wait_for_event(&legacy_sibling, |event| {
|
||||
matches!(event, EventMsg::TurnComplete(_))
|
||||
})
|
||||
.await;
|
||||
let (config, _) = legacy_sibling
|
||||
.current_mcp_config_and_runtime_context()
|
||||
.await;
|
||||
assert_eq!(
|
||||
legacy_sibling
|
||||
.mcp_connection_statuses(&config)
|
||||
.await
|
||||
.get(SERVER_NAME),
|
||||
Some(&McpServerConnectionStatus::NotStarted)
|
||||
);
|
||||
let legacy = &processes[0];
|
||||
assert_definition(
|
||||
&response,
|
||||
&format!("Use the tools from {legacy}."),
|
||||
&format!("Echo from {legacy}."),
|
||||
);
|
||||
|
||||
legacy_sibling.shutdown_and_wait().await?;
|
||||
first.shutdown_and_wait().await?;
|
||||
second.shutdown_and_wait().await?;
|
||||
fixture.codex.shutdown_and_wait().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user