openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-08 20:40:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Preserve MCP sharing discovery when deferring cached stdio startup (#60037)

Browse Source 
## 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:
Charlie Marsh
2026-11-07 02:26:13 +00:00
committed by copyberry
parent 0c58612c3c
commit 7ccf78dec4
5 changed files with 176 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager.rs
+3 -8
View File
@@ -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,
codex-rs/codex-mcp/src/connection_manager_tests.rs
+7
View File
@@ -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());
codex-rs/codex-mcp/src/rmcp_client.rs
+16 -9
View File
@@ -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))
codex-rs/codex-mcp/src/tool_catalog_cache.rs
+34 -22
View File
@@ -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> {