diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 2efa07917f..e85cb67c9a 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -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, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index acfd39a7ba..56cf3da8d6 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -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()); diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index bf3451561f..c6440742da 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -495,6 +495,7 @@ impl AsyncManagedClient { client_mcp_extensions: ClientMcpExtensions, auth_changes: Option>, protocol_mode: McpProtocolMode, + defer_startup: bool, catalog_item_limit: usize, connection_pool: Option<(crate::McpConnectionPool, ConnectionKey)>, publish_apps: Option>>>, @@ -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)) diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index a83f859144..fce1e8605f 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -59,7 +59,8 @@ struct ToolCatalogCacheEntry { struct ToolCatalogCacheState { snapshot: Option, optional_startup_deadline: Option, - 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, } @@ -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>, ) { - 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> { - 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(mutex: &Mutex) -> MutexGuard<'_, T> { diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index 109877ca4e..93c54ea3ac 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -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(()) +}