diff --git a/codex-rs/app-server/src/request_processors/config_processor.rs b/codex-rs/app-server/src/request_processors/config_processor.rs index bbfc952151..3726468eeb 100644 --- a/codex-rs/app-server/src/request_processors/config_processor.rs +++ b/codex-rs/app-server/src/request_processors/config_processor.rs @@ -72,6 +72,7 @@ const SUPPORTED_EXPERIMENTAL_FEATURE_ENABLEMENT: &[&str] = &[ BACKGROUND_PAGINATED_ROLLOUT_MIGRATION_FEATURE, "codex_apps_mcp_2026_07_28", "mcp_2026_07_28", + "mcp_connection_sharing", "memories", "mentions_v2", "remote_control", diff --git a/codex-rs/cli/tests/daemon_startup.rs b/codex-rs/cli/tests/daemon_startup.rs index afe2058826..ae2fe3eefc 100644 --- a/codex-rs/cli/tests/daemon_startup.rs +++ b/codex-rs/cli/tests/daemon_startup.rs @@ -135,7 +135,8 @@ async fn restrictive_launcher_uses_embedded_if_daemon_cannot_start() -> Result<( .kill_on_drop(true); let mut child = command.spawn()?; ensure!( - job.assign_and_resume_process(child.id().context("missing CLI pid")?)?, + job.assign_and_resume_process(child.id().context("missing CLI pid")?)? + .is_some(), "CLI job assignment failed" ); let status = child.wait().await?; diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 125d934cb0..2ad43785d2 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -31,12 +31,14 @@ pub(crate) use tool_catalog::BindingCatalogRevision; pub use tool_catalog::tool_is_model_visible; use codex_config::McpServerAuth; +use std::cell::LazyCell; use std::collections::HashMap; use std::sync::Arc; use std::sync::OnceLock; use std::sync::atomic::Ordering; use std::time::Duration; +use crate::McpProtocolMode; use crate::binding::call_tool_result_from_rmcp; use crate::catalog::McpServerSource; use crate::elicitation::ElicitationRequestManager; @@ -58,7 +60,6 @@ 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; @@ -85,6 +86,8 @@ use codex_protocol::protocol::McpStartupFailureReason; use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::determine_streamable_http_auth_status_from_credentials; +use futures::FutureExt; +use futures::future::join_all; use tokio::sync::watch; use tokio::task::JoinSet; use tracing::warn; @@ -94,6 +97,7 @@ static LIVE_CONNECTIONS: Gauge = Gauge::new("mcp.connections.live"); pub(crate) struct McpServerConnection { identity: Option, client: AsyncManagedClient, + legacy_fallback: Option>>>, // Startup-only budget; changing it must not replace a ready connection. startup_timeout: Duration, startup_trigger: Option>, @@ -101,6 +105,54 @@ pub(crate) struct McpServerConnection { } impl McpServerConnection { + fn retained_legacy(&self) -> Option<&Arc> { + self.legacy_fallback.as_deref().and_then(OnceLock::get) + } + + fn may_use_legacy_fallback(&self) -> bool { + self.legacy_fallback + .as_ref() + .is_some_and(|fallback| fallback.get().is_some() || self.client.client.peek().is_none()) + } + + /// Keep the working private connection if a singleton cannot start a second process. + fn retain_on_failed_discovery( + self: &Arc, + replacement: &mut AsyncManagedClient, + ) -> Arc>> { + let previous = Arc::clone(self.retained_legacy().unwrap_or(self)); + let fallback = Arc::new(OnceLock::new()); + let selected = Arc::clone(&fallback); + let cancellation = replacement.cancel_token.clone(); + replacement.client = replacement + .client + .clone() + .then(move |outcome| async move { + let Err(error @ StartupOutcomeError::Failed { .. }) = &outcome else { + return outcome; + }; + if error.is_authentication_required() { + return outcome; + } + error.wait_for_shutdown().await; + let Some(client) = previous.client.ready_client() else { + return outcome; + }; + if client.client.is_closed().await { + return outcome; + } + if cancellation.is_cancelled() || previous.client.cancel_token.is_cancelled() { + return Err(StartupOutcomeError::Cancelled); + } + // The connection owns the startup token and the capabilities for this transport. + selected.get_or_init(|| previous); + Ok(client) + }) + .boxed() + .shared(); + fallback + } + async fn reusable_client( &self, desired: &McpServerConnectionIdentity, @@ -186,7 +238,7 @@ impl Drop for McpServerConnection { #[derive(Clone)] struct McpServerView { connection: Arc, - protocol_mode: crate::McpProtocolMode, + protocol_mode: McpProtocolMode, startup_readiness: McpStartupReadiness, metadata: McpServerMetadata, tool_filter: ToolFilter, @@ -268,6 +320,11 @@ impl McpServerView { } } +pub(crate) struct PreviousMcpConnections<'a> { + pub(crate) connections: &'a McpConnectionSet, + pub(crate) same_owner: bool, +} + /// A published view over a set of running MCP server connections. pub(crate) struct McpConnectionSet { config: Option>, @@ -290,11 +347,13 @@ impl McpConnectionSet { pub async fn new( ema_auth: Option, existing_ema: Option<&Self>, - previous: Option<&Self>, + previous: Option>, publication_gate: McpPublicationGate, input: McpRuntimeInput, elicitation_router: ElicitationRequestRouter, ) -> Self { + let same_owner = previous.as_ref().is_some_and(|state| state.same_owner); + let previous = previous.map(|state| state.connections); let trusted_access = TrustedAccessContext::from_runtime(&input); let McpRuntimeInput { startup_policy, @@ -437,7 +496,7 @@ impl McpConnectionSet { let uses_ema_auth = config.xaa_enabled && matches!(configured_config.auth, McpServerAuth::EmaAuth); let protocol_mode = if uses_ema_auth { - crate::McpProtocolMode::V20260728 + McpProtocolMode::V20260728 } else { protocol_mode }; @@ -525,6 +584,13 @@ impl McpConnectionSet { } else { chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider) }; + let discover_local_stdio = config.discover_local_stdio + && server.can_discover_local_stdio(&server_name, runtime_auth_provider.as_ref()); + let protocol_mode = if discover_local_stdio { + McpProtocolMode::DiscoverLocalStdio + } else { + protocol_mode + }; if is_host_owned_codex_apps { event_stream_connection = Some(Arc::new(EventStreamConnectionSettings { server: server.clone(), @@ -561,24 +627,54 @@ impl McpConnectionSet { ); let expected_protocol_mode = match &configured_config.transport { McpServerTransportConfig::StreamableHttp { .. } => Some(protocol_mode), - McpServerTransportConfig::Stdio { .. } - if protocol_mode == crate::McpProtocolMode::Legacy => - { - Some(crate::McpProtocolMode::Legacy) - } - McpServerTransportConfig::Stdio { env, .. } => match env - .as_ref() - .and_then(|variables| variables.get("CODEX_MCP_PROTOCOL_VERSION")) - { - None => Some(crate::McpProtocolMode::Legacy), - Some(version) - if version == rmcp::model::ProtocolVersion::V_2026_07_28.as_str() => - { - Some(protocol_mode) - } - Some(_) => None, - }, + McpServerTransportConfig::Stdio { .. } => server.stdio_protocol_mode(protocol_mode), }; + let mut tool_catalog_cache_context = LazyCell::new(|| { + if server_name == CODEX_APPS_MCP_SERVER_NAME + || server.requires_read_only_mcp_tools() + || matches!(configured_config.auth, McpServerAuth::EmaAuth) + { + return None; + } + let environment = resolved_environment.as_ref().ok()?; + let cache_mode = match expected_protocol_mode? { + McpProtocolMode::DiscoverLocalStdio | McpProtocolMode::Legacy => { + McpProtocolMode::Legacy + } + McpProtocolMode::V20260728 => McpProtocolMode::V20260728, + }; + let mut context = tool_catalog_cache.context( + &server_name, + &configured_config, + &runtime_context, + environment.as_ref(), + (&client_elicitation_capability, &client_mcp_extensions), + Some((&connection_identity, cache_mode, server.is_agent_plugin())), + )?; + context.track_tool_list_changes = config.mcp_tool_list_changed; + Some(context) + }); + let defer_startup = allow_deferred_startup + && !tool_plugin_context.is_selected_plugin_mcp_server(&server_name) + && tool_catalog_cache_context + .as_ref() + .and_then(|context| { + context.current_revision_if(|tools| { + tools + .iter() + .any(|tool| configured_tool_filter.allows(&tool.tool.name)) + }) + }) + .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); if matches!(configured_config.auth, McpServerAuth::EmaAuth) && existing_ema.is_some() { if let Some(previous_view) = existing_ema .filter(|_| elicitation_previous.is_some()) @@ -640,9 +736,11 @@ impl McpConnectionSet { // Post-start publications may only retain the original EMA connection. continue; } + let mut legacy_to_upgrade = None; if let Some(previous_view) = reusable_previous.and_then(|previous| previous.servers.get(&server_name)) && previous_view.refresh_tool_list == config.mcp_tool_list_changed + && (same_owner || !previous_view.connection.may_use_legacy_fallback()) { let connection = Arc::clone(&previous_view.connection); let reusable_pending_startup = connection.identity.as_ref() @@ -668,17 +766,19 @@ impl McpConnectionSet { } else { None }; + let ready = connection + .reusable_client(&connection_identity) + .await + .filter(|_| previous_view.catalog_item_limit == catalog_item_limit); if reusable_pending_startup || unchanged_auth_failure.is_some() - || connection - .reusable_client(&connection_identity) - .await - .is_some_and(|client| { - previous_view.catalog_item_limit == catalog_item_limit - && expected_protocol_mode.is_some_and(|expected| { - client.client.protocol_mode() == expected - }) + || ready.as_ref().is_some_and(|client| { + expected_protocol_mode.is_some_and(|expected| { + client.client.protocol_mode() == expected + || (previous_view.protocol_mode == expected + && connection.retained_legacy().is_some()) }) + }) { let pending_client = reusable_pending_startup.then(|| connection.client.clone()); @@ -746,34 +846,22 @@ impl McpConnectionSet { } continue; } + if same_owner + && expected_protocol_mode == Some(McpProtocolMode::DiscoverLocalStdio) + && previous_view.protocol_mode == McpProtocolMode::Legacy + && connection.identity.as_ref() == Some(&connection_identity) + && !connection.client.cancel_token.is_cancelled() + && ready.is_some_and(|client| { + client.client.protocol_mode() == McpProtocolMode::Legacy + }) + { + legacy_to_upgrade = Some(connection); + } } let cancel_token = startup_cancellation_token.child_token(); - let mut tool_catalog_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME - || server.requires_read_only_mcp_tools() - || matches!(configured_config.auth, McpServerAuth::EmaAuth) - { - None - } else if let Ok(environment) = resolved_environment.as_ref() { - tool_catalog_cache.context( - &server_name, - &configured_config, - &runtime_context, - environment.as_ref(), - (&client_elicitation_capability, &client_mcp_extensions), - Some(( - &connection_identity, - protocol_mode, - server.is_agent_plugin(), - )), - ) - } else { - None - }; - if let Some(context) = &mut tool_catalog_cache_context { - context.track_tool_list_changes = config.mcp_tool_list_changed; - } + let tool_catalog_cache_context = tool_catalog_cache_context.take(); let has_runtime_auth = runtime_auth_provider.is_some(); - let async_managed_client = AsyncManagedClient::new( + let mut async_managed_client = AsyncManagedClient::new( ema_auth.clone(), server_name.clone(), startup_submit_id.clone(), @@ -803,17 +891,12 @@ impl McpConnectionSet { protocol_mode, catalog_item_limit, ); - let defer_startup = allow_deferred_startup - && !tool_plugin_context.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)) - }); + // A forced reconnect must also wait for the previously published helper to exit. + if let Some(previous) = existing_ema.and_then(|set| set.servers.get(&server_name)) { + async_managed_client.wait_for_failed_predecessor(&previous.connection.client); + } + let legacy_fallback = legacy_to_upgrade + .map(|connection| connection.retain_on_failed_discovery(&mut async_managed_client)); let (startup_trigger, startup_receiver) = if defer_startup { let (trigger, receiver) = watch::channel(false); (Some(trigger), Some(receiver)) @@ -826,6 +909,7 @@ impl McpConnectionSet { connection: Arc::new(McpServerConnection { identity: Some(connection_identity), client: async_managed_client.clone(), + legacy_fallback, startup_timeout, startup_trigger, _diagnostics_guard: LIVE_CONNECTIONS.track(), @@ -1117,9 +1201,7 @@ impl McpConnectionSet { .collect::>(); // Keep cleanup alive if an interrupt cancels the refresh that requested it. let shutdown_task = tokio::spawn(async move { - for connection in connections { - connection.shutdown().await; - } + join_all(connections.iter().map(|connection| connection.shutdown())).await; }); if let Err(error) = shutdown_task.await { warn!("MCP client shutdown task failed: {error}"); @@ -1252,6 +1334,8 @@ impl McpConnectionSet { .filter(|(name, _)| include_server(name)) .filter_map(|(name, view)| { view.connection + .retained_legacy() + .unwrap_or(&view.connection) .client .server_capabilities .lock() diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 06e4c3b8d5..08323befbd 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -148,6 +148,7 @@ impl McpConnectionSet { client, startup_timeout: DEFAULT_STARTUP_TIMEOUT, startup_trigger: None, + legacy_fallback: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), metadata: McpServerMetadata { @@ -617,6 +618,7 @@ async fn connection_statuses_observe_clients_without_starting_them() { error: "broken".to_string(), is_authentication_required: false, is_ema_authentication_required: false, + shutdown: None, }, ), ( @@ -625,6 +627,7 @@ async fn connection_statuses_observe_clients_without_starting_them() { error: "login".to_string(), is_authentication_required: true, is_ema_authentication_required: false, + shutdown: None, }, ), ( @@ -700,6 +703,7 @@ async fn connection_statuses_follow_latest_reconnect_outcome() { error: "retry failed".to_string(), is_authentication_required: attempt == 0, is_ema_authentication_required: false, + shutdown: None, }), _ => Ok(recovered), } @@ -866,6 +870,7 @@ fn create_test_manager_with_failed_apps_startup( error: "startup failed".to_string(), is_authentication_required: false, is_ema_authentication_required: false, + shutdown: None, })) .boxed() .shared(); @@ -3194,6 +3199,7 @@ async fn stable_catalog_revisions_ignore_terminal_optional_server_failures() { error: "optional startup failed".to_string(), is_authentication_required: false, is_ema_authentication_required: false, + shutdown: None, }, )) .boxed() @@ -3397,8 +3403,10 @@ async fn capture_binding_uses_cache_published_during_optional_startup() { let server_config: McpServerConfig = serde_json::from_value(serde_json::json!({ "command": "pending-mcp" })) .expect("server configuration"); - let cache_context = McpToolCatalogCache::default() - .context( + let cache = McpToolCatalogCache::default(); + let identity = reusable_server_identity("pending", &server_config, &runtime_context); + let catalog = |mode| { + cache.context( "pending", &server_config, &runtime_context, @@ -3407,9 +3415,17 @@ async fn capture_binding_uses_cache_published_during_optional_startup() { &ElicitationCapability::default(), &ClientMcpExtensions::default(), ), - /*connection_identity*/ None, + Some((&identity, mode, /*agent_plugin*/ false)), ) - .expect("shared catalog"); + }; + let legacy = catalog(McpProtocolMode::Legacy).expect("legacy catalog"); + legacy.publish_if_newest( + legacy.begin_fetch(), + &[create_test_tool("pending", "legacy_tool")], + /*meta*/ None, + ); + assert!(catalog(McpProtocolMode::DiscoverLocalStdio).is_none()); + let cache_context = catalog(McpProtocolMode::V20260728).expect("modern catalog"); let (mut client, started, _release) = create_gated_async_managed_client( create_test_managed_client(vec![create_test_tool("pending", "live_tool")]).await, ); @@ -3427,7 +3443,7 @@ async fn capture_binding_uses_cache_published_during_optional_startup() { assert!(futures::poll!(&mut binding).is_pending()); started .await - .expect("optional startup began without a cache"); + .expect("a catalog from the legacy handshake cannot satisfy discovery"); cache_context.publish_if_newest( cache_context.begin_fetch(), &[create_test_tool("pending", "peer_tool")], @@ -4144,6 +4160,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() error: "startup failed".to_string(), is_authentication_required: false, is_ema_authentication_required: false, + shutdown: None, }, )) .boxed() @@ -4304,6 +4321,7 @@ async fn later_tool_list_retries_after_failed_reconnect_and_keeps_cached_tools() error: "recreated startup failed".to_string(), is_authentication_required: false, is_ema_authentication_required: false, + shutdown: None, }) } else { Ok(recovered_client) @@ -5110,6 +5128,7 @@ fn mcp_init_error_display_prompts_for_login_when_auth_required() { error: "OAuth refresh token was rejected: invalid_grant".to_string(), is_authentication_required: true, is_ema_authentication_required: false, + shutdown: None, }, ] { let display = mcp_init_error_display( @@ -5142,6 +5161,7 @@ fn mcp_init_error_display_identifies_oauth_reauthentication() { error: "authorization required: Bearer error=\"invalid_token\"".to_string(), is_authentication_required: true, is_ema_authentication_required: false, + shutdown: None, }; let executor_config: McpServerConfig = serde_json::from_value(serde_json::json!({ "url": "https://example.com/mcp", @@ -5206,6 +5226,7 @@ fn mcp_startup_failure_reason_requires_existing_oauth_and_auth_failure() { error: "startup failed".to_string(), is_authentication_required, is_ema_authentication_required: false, + shutdown: None, }; assert_eq!( @@ -5388,6 +5409,7 @@ async fn manager_with_reusable_ready_server( .startup_timeout_sec .unwrap_or(DEFAULT_STARTUP_TIMEOUT), startup_trigger: None, + legacy_fallback: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), metadata: McpServerMetadata::from(&server), @@ -5426,7 +5448,10 @@ async fn reconcile_reusable_server_with_mcp_config( McpConnectionSet::new( /*ema_auth*/ None, /*existing_ema*/ None, - Some(previous), + Some(PreviousMcpConnections { + connections: previous, + same_owner: true, + }), McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, @@ -5839,6 +5864,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a .startup_timeout_sec .unwrap_or(DEFAULT_STARTUP_TIMEOUT), startup_trigger: None, + legacy_fallback: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }), metadata: McpServerMetadata::from(&server), @@ -6129,6 +6155,7 @@ async fn reconciliation_retries_non_oauth_authentication_failures() { error: "bearer token rejected".to_string(), is_authentication_required: true, is_ema_authentication_required: false, + shutdown: None, })) .boxed() .shared(); @@ -6389,7 +6416,10 @@ async fn reconciliation_replaces_connection_when_protocol_mode_changes() { let reconciled = McpConnectionSet::new( /*ema_auth*/ None, /*existing_ema*/ None, - Some(&previous), + Some(PreviousMcpConnections { + connections: &previous, + same_owner: true, + }), McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, @@ -6450,7 +6480,10 @@ async fn reconciliation_reuses_legacy_stdio_server_when_modern_protocol_is_enabl let reconciled = McpConnectionSet::new( /*ema_auth*/ None, /*existing_ema*/ None, - Some(&previous), + Some(PreviousMcpConnections { + connections: &previous, + same_owner: true, + }), McpPublicationGate::already_published(), McpRuntimeInput { network_policy_metadata: None, @@ -6667,6 +6700,7 @@ async fn reconciliation_replaces_closed_connections() -> anyhow::Result<()> { .startup_timeout_sec .unwrap_or(DEFAULT_STARTUP_TIMEOUT), startup_trigger: None, + legacy_fallback: None, _diagnostics_guard: LIVE_CONNECTIONS.track(), }); diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index fa5242e3e7..b34a822183 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -190,6 +190,8 @@ pub struct McpConfig { pub non_prefixed_mcp_tool_servers: Vec, /// Protocol mode for servers other than the host-owned Codex Apps registration. pub protocol_mode: McpProtocolMode, + /// Try protocol discovery before initializing eligible local stdio servers. + pub discover_local_stdio: bool, /// Refresh ordinary server catalogs after tools/list_changed notifications. pub mcp_tool_list_changed: bool, /// Independent protocol mode for the trusted, HTTP Codex Apps registration. diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 68730ba038..c066d8c8b3 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -130,6 +130,7 @@ async fn status_snapshot_only_downgrades_oauth_and_ema_authentication_failures() error: "OAuth refresh token was rejected".to_string(), is_authentication_required: true, is_ema_authentication_required: false, + shutdown: None, }; let mut manager = McpConnectionSet::empty(/*prefix_mcp_tool_names*/ true); let mut auth_status_entries = HashMap::new(); @@ -147,6 +148,7 @@ async fn status_snapshot_only_downgrades_oauth_and_ema_authentication_failures() error: "provider temporarily unavailable".to_string(), is_authentication_required: false, is_ema_authentication_required: false, + shutdown: None, }), ), ( @@ -236,6 +238,7 @@ pub(crate) fn test_mcp_config(codex_home: PathBuf) -> McpConfig { prefix_mcp_tool_names: true, non_prefixed_mcp_tool_servers: Vec::new(), protocol_mode: McpProtocolMode::Legacy, + discover_local_stdio: false, mcp_tool_list_changed: true, host_owned_apps_protocol_mode: McpProtocolMode::Legacy, client_elicitation_capability: ElicitationCapability::default(), diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 54d5d4dfe7..8e66a9ecfc 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -166,6 +166,7 @@ impl ManagedClient { pub(crate) type ManagedClientFuture = Shared>>; +type StartupShutdown = Shared>>; #[derive(Default)] struct CodexAppsStartupReconnectState { @@ -368,6 +369,7 @@ impl ManagedClientStartup { .as_ref() .map(McpToolCatalogCacheContext::begin_fetch); let refresh_start = is_codex_apps_mcp_server.then(Instant::now); + let mut discovery_client = None; let outcome = match async { if let Err(error) = validate_mcp_server_name(&server_name) { return Err(error.into()); @@ -399,6 +401,9 @@ impl ManagedClientStartup { ))); } }; + if client.protocol_mode() == McpProtocolMode::DiscoverLocalStdio { + discovery_client = Some(Arc::clone(&client)); + } start_server_task( server_name.clone(), client, @@ -422,8 +427,32 @@ impl ManagedClientStartup { .or_cancel(&cancel_token_for_fut) .await { - Ok(result) => result, - Err(CancelErr::Cancelled) => Err(StartupOutcomeError::Cancelled), + Ok(mut result) => { + if let Err(StartupOutcomeError::Failed { shutdown, .. }) = &mut result + && let Some(client) = discovery_client + { + // Report the startup error now. A retry waits for the helper to exit. + *shutdown = Some( + tokio::spawn(async move { client.shutdown_before_retry().await }) + .map(|joined| match joined { + Ok(result) => result.map_err(|error| error.to_string()), + Err(error) => Err(error.to_string()), + }) + .boxed() + .shared(), + ); + } + result + } + Err(CancelErr::Cancelled) => { + // A cancelled legacy replacement may still be inside the native spawn. + if let Some(client) = discovery_client + && let Err(error) = client.shutdown_before_retry().await + { + warn!("failed to stop cancelled MCP discovery: {error}"); + } + Err(StartupOutcomeError::Cancelled) + } }; // Log once per startup attempt, including discovery without startup notifications. if let Err(StartupOutcomeError::Failed { error, .. }) = &outcome { @@ -540,6 +569,9 @@ impl AsyncManagedClient { startup_complete: Arc::clone(&startup_complete), server_capabilities: Arc::clone(&server_capabilities), }); + // Startup can publish a legacy fallback. Readers wait for the negotiated catalog. + let tool_catalog_cache_context = tool_catalog_cache_context + .filter(|_| protocol_mode != McpProtocolMode::DiscoverLocalStdio); let client = startup.start(); let startup_reconnect = is_codex_apps_mcp_server.then(|| { let startup = Arc::clone(&startup); @@ -595,6 +627,33 @@ impl AsyncManagedClient { self.ready_client().map(|client| client.client) } + /// Waits for this server's failed predecessor when the replacement actually starts. + pub(crate) fn wait_for_failed_predecessor(&mut self, previous: &Self) { + let Some(Err( + error @ StartupOutcomeError::Failed { + shutdown: Some(shutdown), + .. + }, + )) = previous.client.peek() + else { + return; + }; + let shutdown = shutdown.clone(); + let error = error.clone(); + let client = self.client.clone(); + let startup_complete = Arc::clone(&self.startup_complete); + self.client = async move { + if let Err(stop_error) = shutdown.await { + warn!("cannot restart MCP server after failed cleanup: {stop_error}"); + startup_complete.store(true, Ordering::Release); + return Err(error); + } + client.await + } + .boxed() + .shared(); + } + pub(crate) async fn reconnect_failed_startup(&self) { let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { return; @@ -614,6 +673,7 @@ impl AsyncManagedClient { Err(StartupOutcomeError::Cancelled) => {} Err(error) => { warn!("failed to initialize MCP client during shutdown: {error:#}"); + error.wait_for_shutdown().await; } } } @@ -684,10 +744,22 @@ pub(crate) enum StartupOutcomeError { error: String, is_authentication_required: bool, is_ema_authentication_required: bool, + shutdown: Option, }, } impl StartupOutcomeError { + pub(crate) async fn wait_for_shutdown(&self) { + if let Self::Failed { + shutdown: Some(shutdown), + .. + } = self + && let Err(error) = shutdown.clone().await + { + warn!("failed to stop MCP server after startup failed: {error}"); + } + } + pub(crate) fn is_authentication_required(&self) -> bool { match self { Self::Cancelled => false, @@ -709,6 +781,7 @@ impl From for StartupOutcomeError { is_ema_authentication_required: is_authentication_required && (error.downcast_ref::().is_some() || error.downcast_ref::().is_some()), + shutdown: None, } } } @@ -1013,7 +1086,7 @@ async fn start_server_task( tx_event, elicitation_requests, codex_apps_tools_cache_context, - tool_catalog_cache_context, + mut tool_catalog_cache_context, tool_catalog_fetch_ticket, client_elicitation_capability, client_mcp_extensions, @@ -1050,6 +1123,13 @@ async fn start_server_task( &initialize_result, ); let initialize_result = initialize_result.map_err(StartupOutcomeError::from)?; + if client.protocol_mode() == McpProtocolMode::DiscoverLocalStdio + && !initialize_result.protocol_version.has_initialize() + && let Some(context) = tool_catalog_cache_context.take() + && let Some(ticket) = &tool_catalog_fetch_ticket + { + context.invalidate_if_newest(ticket); + } *server_capabilities .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) = @@ -1186,7 +1266,7 @@ fn record_protocol_discovery_metrics( let mode = match mode { McpProtocolMode::Legacy => "legacy", - McpProtocolMode::V20260728 => "auto", + McpProtocolMode::V20260728 | McpProtocolMode::DiscoverLocalStdio => "auto", }; let outcome = match result { Ok(result) if result.protocol_version == ProtocolVersion::V_2026_07_28 => "modern", diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 52ab17007d..e994fb0f38 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -53,6 +53,7 @@ use crate::binding::PreparedMcpCall; use crate::client_tool_catalog::CodexAppsToolSnapshot; use crate::connection_manager::BindingCatalogRevision; use crate::connection_manager::McpConnectionSet; +use crate::connection_manager::PreviousMcpConnections; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationReviewerHandle; @@ -131,6 +132,7 @@ struct PublishedMcpRuntime { auth: Option, auth_token: Option, auth_manager: Option>, + auth_owner: Option, auth_generation: Arc<()>, plugins_available: bool, ready_selected_capability_roots: Vec, @@ -139,6 +141,23 @@ struct PublishedMcpRuntime { cached_binding: Mutex>, } +/// The account that owns a private helper across credential refreshes. +#[derive(Clone, Copy, PartialEq, Eq)] +struct AuthOwner { + generation: u64, + is_fedramp: bool, +} + +impl AuthOwner { + fn capture(auth: Option<&CodexAuth>, manager: Option<&AuthManager>) -> Option { + let auth = auth?; + Some(Self { + generation: manager?.auth_change_state_for(auth)?.owner_generation, + is_fedramp: auth.is_fedramp_account(), + }) + } +} + fn ensure_host_owned_apps_registration( current: &PublishedMcpRuntime, server: &str, @@ -242,6 +261,7 @@ impl McpRuntime { auth: None, auth_token: None, auth_manager: None, + auth_owner: None, auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -323,22 +343,18 @@ impl McpRuntime { /// Reconciles configured servers and publishes their immutable runtime snapshot. pub async fn replace(&self, input: McpRuntimeInput) { - let current = self.current.load_full(); let mut reconnect = McpReconnectGuard { pending: &self.reconnect_pending, claimed: self.reconnect_pending.swap(false, Ordering::AcqRel), }; - self.publish( - input, - (!reconnect.claimed).then_some(current.connections.as_ref()), - ) - .await; + self.publish(input, /*reuse_previous*/ !reconnect.claimed) + .await; reconnect.claimed = false; } /// Starts fresh connections and returns their complete, refreshed Apps catalog. pub async fn replace_fresh(&self, input: McpRuntimeInput) -> anyhow::Result> { - self.publish(input, /*previous*/ None).await; + self.publish(input, /*reuse_previous*/ false).await; self.latest_hard_refresh_codex_apps_tools_cache().await } @@ -357,7 +373,7 @@ impl McpRuntime { } } - async fn publish(&self, input: McpRuntimeInput, previous: Option<&McpConnectionSet>) { + async fn publish(&self, input: McpRuntimeInput, reuse_previous: bool) { let current = self.current.load_full(); let retain_ema = current.config.is_some(); let ema_auth = (input.config.xaa_enabled @@ -382,7 +398,21 @@ impl McpRuntime { let auth = input.auth.clone(); let auth_manager = input.auth_manager.clone(); let auth_token = auth.as_ref().and_then(|auth| auth.get_token().ok()); - let auth_generation = if current.auth_matches(auth.as_ref()) { + let auth_matches = current.auth_matches(auth.as_ref()); + let auth_owner = AuthOwner::capture(auth.as_ref(), auth_manager.as_deref()); + let same_owner = match (current.auth_owner, auth_owner) { + (Some(previous), Some(next)) => { + previous == next + && current + .auth_manager + .as_ref() + .zip(auth_manager.as_ref()) + .is_some_and(|(previous, next)| Arc::ptr_eq(previous, next)) + } + (None, None) => auth_matches, + _ => false, + }; + let auth_generation = if auth_matches { Arc::clone(¤t.auth_generation) } else { Arc::new(()) @@ -395,7 +425,10 @@ impl McpRuntime { McpConnectionSet::new( ema_auth, retain_ema.then_some(current.connections.as_ref()), - previous, + reuse_previous.then_some(PreviousMcpConnections { + connections: ¤t.connections, + same_owner, + }), publication_gate, input, self.elicitation_router.clone(), @@ -421,6 +454,7 @@ impl McpRuntime { auth, auth_token, auth_manager, + auth_owner, auth_generation, plugins_available, ready_selected_capability_roots, @@ -1129,6 +1163,7 @@ mod tests { auth: None, auth_token: None, auth_manager: None, + auth_owner: None, auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), @@ -1204,6 +1239,7 @@ mod tests { auth: None, auth_token: None, auth_manager: None, + auth_owner: None, auth_generation: Arc::new(()), plugins_available: false, ready_selected_capability_roots: Vec::new(), diff --git a/codex-rs/codex-mcp/src/server.rs b/codex-rs/codex-mcp/src/server.rs index df532ee1ea..3fa6d853b5 100644 --- a/codex-rs/codex-mcp/src/server.rs +++ b/codex-rs/codex-mcp/src/server.rs @@ -1,8 +1,10 @@ use std::collections::HashMap; +use std::ffi::OsStr; use std::ffi::OsString; use std::path::PathBuf; use std::sync::Arc; +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::runtime::McpRuntimeContext; use codex_api::SharedAuthProvider; use codex_config::AppToolApproval; @@ -17,6 +19,7 @@ use codex_exec_server::Environment; use codex_login::CodexAuth; use codex_protocol::mcp::ClientMcpExtensions; use codex_rmcp_client::McpOAuthRefreshMode; +use codex_rmcp_client::McpProtocolMode; use codex_rmcp_client::StoredOAuthCredentialSnapshot; use codex_rmcp_client::StoredOAuthTokens; use codex_utils_path_uri::PathUri; @@ -82,6 +85,35 @@ impl EffectiveMcpServer { &self.config } + /// Checks whether Core can discover the protocol of a local stdio server. + /// + /// The reserved `codex_apps` server has a per-chat catalog. Servers with runtime or + /// enterprise-managed credentials need that chat's authentication state. + pub(crate) fn can_discover_local_stdio( + &self, + server_name: &str, + runtime_auth_provider: Option<&SharedAuthProvider>, + ) -> bool { + server_name != CODEX_APPS_MCP_SERVER_NAME + && self.config.is_local_environment() + && runtime_auth_provider.is_none() + && !matches!(self.config.auth, McpServerAuth::EmaAuth) + && self + .stdio_protocol_mode(McpProtocolMode::DiscoverLocalStdio) + .is_some() + } + + /// Returns the handshake a stdio launch would use, or `None` for HTTP or an invalid version. + pub(crate) fn stdio_protocol_mode(&self, mode: McpProtocolMode) -> Option { + let McpServerTransportConfig::Stdio { env, .. } = &self.config.transport else { + return None; + }; + let version = env + .as_ref() + .and_then(|env| env.get("CODEX_MCP_PROTOCOL_VERSION")); + mode.stdio_mode(version.map(OsStr::new)).ok() + } + pub fn enabled(&self) -> bool { self.config.enabled } diff --git a/codex-rs/codex-mcp/src/tool_catalog_cache.rs b/codex-rs/codex-mcp/src/tool_catalog_cache.rs index 8385633bd6..a83f859144 100644 --- a/codex-rs/codex-mcp/src/tool_catalog_cache.rs +++ b/codex-rs/codex-mcp/src/tool_catalog_cache.rs @@ -184,6 +184,7 @@ impl McpToolCatalogCacheContext { cached_deadline.deadline } + #[cfg(test)] pub(crate) fn current_tools(&self) -> Option> { self.current_tools_or(/*fallback*/ None) } @@ -235,25 +236,28 @@ 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; + } + } + pub(crate) fn publish_if_newest( &self, ticket: McpToolCatalogFetchTicket, tools: &[ToolInfo], meta: Option>, ) { - let mut state = lock_unpoisoned(&self.entry.state); - if state.disabled_by_server || ticket.generation <= state.last_accepted_generation { + let Some(mut state) = self.accept_generation(ticket.generation) else { return; - } - + }; let mut tools = tools.to_vec(); for tool in &mut tools { // Tool annotations affect approval and parallelism decisions, so only the live // connection may supply them. tool.tool.annotations = None; } - state.last_accepted_generation = ticket.generation; - state.optional_startup_deadline = None; state.snapshot = Some(ToolCatalogSnapshot { tools, meta, @@ -261,6 +265,16 @@ 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> { @@ -417,6 +431,13 @@ impl ToolCatalogTransportIdentity { else { return None; }; + // An unmarked helper can expose different tools after a legacy fallback. + if matches!( + connection_identity, + Some((_, McpProtocolMode::DiscoverLocalStdio, _)) + ) { + return None; + } if env_vars .iter() .any(codex_config::McpServerEnvVar::is_remote_source) @@ -438,6 +459,7 @@ impl ToolCatalogTransportIdentity { env_vars, cwd, &config.environment_id, + connection_identity.map(|(_, mode, _)| mode.preferred_protocol_version()), client_elicitation_capability, client_mcp_extensions.iter().collect::>(), )) diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index a61f58b8ea..7b101bd261 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -1143,6 +1143,9 @@ "mcp_2026_07_28": { "type": "boolean" }, + "mcp_connection_sharing": { + "type": "boolean" + }, "mcp_oauth_refresh_coordination": { "type": "boolean" }, @@ -7751,6 +7754,9 @@ "mcp_2026_07_28": { "type": "boolean" }, + "mcp_connection_sharing": { + "type": "boolean" + }, "mcp_oauth_refresh_coordination": { "type": "boolean" }, diff --git a/codex-rs/core/src/config/managed_features.rs b/codex-rs/core/src/config/managed_features.rs index dc57ca9216..2fd34d2f0c 100644 --- a/codex-rs/core/src/config/managed_features.rs +++ b/codex-rs/core/src/config/managed_features.rs @@ -124,6 +124,7 @@ impl ManagedFeatures { let mut values = self.get().clone(); for feature in [ Feature::EnableMcpApps, + Feature::McpConnectionSharing, Feature::SecretAuthStorage, Feature::McpOAuthRefreshCoordination, ] { diff --git a/codex-rs/core/src/config/mod.rs b/codex-rs/core/src/config/mod.rs index 1f192747a4..e30c8db539 100644 --- a/codex-rs/core/src/config/mod.rs +++ b/codex-rs/core/src/config/mod.rs @@ -1957,6 +1957,7 @@ impl Config { Vec::new() }, protocol_mode: self.mcp_protocol_mode(), + discover_local_stdio: self.features.enabled(Feature::McpConnectionSharing), host_owned_apps_protocol_mode: if self.features.enabled(Feature::CodexAppsMcp20260728) { McpProtocolMode::V20260728 } else { diff --git a/codex-rs/core/src/config/runtime_refresh.rs b/codex-rs/core/src/config/runtime_refresh.rs index a57a5d5aa9..a8b61d3598 100644 --- a/codex-rs/core/src/config/runtime_refresh.rs +++ b/codex-rs/core/src/config/runtime_refresh.rs @@ -67,6 +67,7 @@ impl Config { // These operational controls accept host defaults beneath explicit layer values. // Enterprise opt-in always comes from trusted layers instead. for feature in [ + Feature::McpConnectionSharing, Feature::SecretAuthStorage, Feature::McpOAuthRefreshCoordination, ] { diff --git a/codex-rs/core/tests/common/process.rs b/codex-rs/core/tests/common/process.rs index 82e42aac95..710164d7c1 100644 --- a/codex-rs/core/tests/common/process.rs +++ b/codex-rs/core/tests/common/process.rs @@ -40,7 +40,8 @@ async fn wait_for_process_exit_inner(pid: String) -> anyhow::Result<()> { pub async fn wait_for_process_exit(pid: &str) -> anyhow::Result<()> { let pid = pid.to_string(); - tokio::time::timeout(Duration::from_secs(2), wait_for_process_exit_inner(pid)) + // A stdio process group can keep its leader's PID for the two-second shutdown grace. + tokio::time::timeout(Duration::from_secs(5), wait_for_process_exit_inner(pid)) .await .context("timed out waiting for process to exit")??; diff --git a/codex-rs/core/tests/suite/mcp_tool_cache.rs b/codex-rs/core/tests/suite/mcp_tool_cache.rs index 602252bb8c..0a0440bc54 100644 --- a/codex-rs/core/tests/suite/mcp_tool_cache.rs +++ b/codex-rs/core/tests/suite/mcp_tool_cache.rs @@ -19,6 +19,7 @@ use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::McpServerContribution; use codex_extension_api::McpServerContributionContext; use codex_extension_api::McpServerContributor; +use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_protocol::mcp::McpServerConnectionStatus; use codex_protocol::models::PermissionProfile; @@ -700,6 +701,11 @@ 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, .. @@ -713,7 +719,7 @@ async fn cached_mcp_startup_is_eager_for_root_and_lazy_for_subagents() -> anyhow agent_nickname: None, agent_role: None, })), - ..StartThreadOptions::new(fixture.config.clone()) + ..StartThreadOptions::new(subagent_config) }) .await?; second_thread.submit(Op::Interrupt).await?; diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 4a6595122d..5fd61d1445 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -269,6 +269,8 @@ pub enum Feature { Mcp20260728, /// Honor ordinary MCP server tool-list notifications within a turn. McpToolListChanged, + /// Discover the MCP protocol for eligible local stdio servers. + McpConnectionSharing, /// Enable MCP protocol version 2026-07-28 for the host-owned Codex Apps server. CodexAppsMcp20260728, /// Let RMCP coordinate OAuth refresh through the configured credential store. @@ -1634,6 +1636,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::McpConnectionSharing, + key: "mcp_connection_sharing", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CodexAppsMcp20260728, key: "codex_apps_mcp_2026_07_28", diff --git a/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs index 3f6afd99f3..a1293e55a7 100644 --- a/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_mcp_2026_stdio_server.rs @@ -2,6 +2,7 @@ use std::io; use std::io::BufRead; use std::io::Write; +use anyhow::Context; use anyhow::Result; use serde_json::Value; use serde_json::json; @@ -41,7 +42,18 @@ fn main() -> Result<()> { let mode = std::env::args().nth(1).unwrap_or_default(); assert!(matches!( mode.as_str(), - "modern" | "legacy" | "legacy-fallback" | "oversized-stdout" | "oversized-stdout-legacy" + "modern" + | "modern-late" + | "legacy" + | "legacy-fallback" + | "legacy-reject-and-exit" + | "legacy-init-first" + | "legacy-reject-late-initialize" + | "legacy-late" + | "legacy-silent" + | "legacy-versions" + | "oversized-stdout" + | "oversized-stdout-legacy" )); assert!(std::env::var_os("CODEX_MCP_PROTOCOL_VERSION").is_none()); let is_legacy = mode.starts_with("legacy") || mode.ends_with("-legacy"); @@ -49,8 +61,9 @@ fn main() -> Result<()> { let mut stdout = io::stdout().lock(); let mut initialized = false; let mut pending_legacy_call = None; + let mut late_discovery_id = None; - for line in io::stdin().lock().lines() { + for (index, line) in io::stdin().lock().lines().enumerate() { let message: Value = serde_json::from_str(&line?)?; let Some(method) = message.get("method").and_then(Value::as_str) else { let pending_call = pending_legacy_call @@ -70,9 +83,60 @@ fn main() -> Result<()> { continue; }; + // Some older helpers assume their first request is initialize without inspecting its method. + let method = if mode == "legacy-init-first" && !initialized { + "initialize" + } else { + method + }; match method { + // Some legacy stdio servers exit if their first message is not initialize. + "server/discover" if mode == "legacy" => { + // Some wrappers print a banner before starting the protocol. + writeln!(stdout, "starting server")?; + write_message( + &mut stdout, + json!({"jsonrpc": "2.0", "id": "startup", "method": "ping"}), + )?; + return Ok(()); + } + "server/discover" if mode == "legacy-late" => { + late_discovery_id = Some(message["id"].clone()); + } + // The first helper makes its version offer only after it receives initialize. + "server/discover" + if mode == "modern-late" + && !std::path::Path::new("retry-modern").try_exists()? => {} + "initialize" if mode == "modern-late" => { + std::fs::write("retry-modern", [])?; + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", "id": null, + "error": {"code": -32022, "message": "unsupported protocol", "data": {"supported": [MODERN_VERSION]}}, + }), + )?; + return Ok(()); + } + "server/discover" if mode == "legacy-silent" => { + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", "method": "notifications/message", + "params": {"level": "info", "data": "starting server"}, + }), + )?; + } + "server/discover" if mode == "legacy-versions" => { + write_message( + &mut stdout, + json!({ + "jsonrpc": "2.0", "id": message["id"], + "error": {"code": -32022, "message": "unsupported protocol", "data": {"supported": [LEGACY_VERSION]}}, + }), + )?; + } "server/discover" if is_legacy => { - assert_eq!(mode, "legacy-fallback"); write_message( &mut stdout, json!({ @@ -81,6 +145,9 @@ fn main() -> Result<()> { "error": {"code": -32601, "message": "method not found"}, }), )?; + if mode == "legacy-reject-and-exit" { + return Ok(()); + } } "server/discover" => { assert!(!initialized); @@ -107,9 +174,28 @@ fn main() -> Result<()> { }), )?; } + "initialize" if mode == "legacy-reject-late-initialize" && index != 0 => { + write_message( + &mut stdout, + json!({"jsonrpc": "2.0", "id": message["id"], "error": { + "code": -32600, "message": "initialize must be the first request" + }}), + )?; + } "initialize" => { assert!(is_legacy); - assert_eq!(message["params"]["protocolVersion"], LEGACY_VERSION); + if mode == "legacy-late" { + let id = late_discovery_id + .take() + .context("initialize must follow discover on the same pipe")?; + write_message( + &mut stdout, + json!({"jsonrpc":"2.0","id":id,"error":{"code":-32601,"message":"Method not found"}}), + )?; + } + if message["method"] == "initialize" { + assert_eq!(message["params"]["protocolVersion"], LEGACY_VERSION); + } initialized = true; write_message( &mut stdout, @@ -130,6 +216,9 @@ fn main() -> Result<()> { if !is_legacy { assert_modern_metadata(&message); } + if mode == "oversized-stdout-legacy" { + write!(stdout, "\u{feff}")?; + } let description = if oversized { "x".repeat(MAX_MESSAGE_BYTES + 1) } else { diff --git a/codex-rs/rmcp-client/src/bounded_stdio_transport.rs b/codex-rs/rmcp-client/src/bounded_stdio_transport.rs index fdde9f698e..dc065995f5 100644 --- a/codex-rs/rmcp-client/src/bounded_stdio_transport.rs +++ b/codex-rs/rmcp-client/src/bounded_stdio_transport.rs @@ -10,6 +10,8 @@ use rmcp::service::RxJsonRpcMessage; use rmcp::service::TxJsonRpcMessage; use rmcp::transport::Transport; use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::process::ChildStdin; @@ -38,6 +40,20 @@ impl BoundedStdioTransport { } } + /// Pass both a partial line and buffered read-ahead to the SDK's legacy codec. + pub(super) fn into_legacy_pipes( + self, + ) -> io::Result<(Box, ChildStdin)> { + let stdin = self + .stdin + .try_lock() + .map_err(io::Error::other)? + .take() + .ok_or_else(|| io::Error::new(io::ErrorKind::BrokenPipe, "MCP stdin closed"))?; + let stdout = io::Cursor::new(self.pending_line).chain(self.stdout); + Ok((Box::new(stdout), stdin)) + } + async fn receive_message(&mut self) -> Option> { loop { let bytes = match self.stdout.fill_buf().await { diff --git a/codex-rs/rmcp-client/src/http_client_adapter.rs b/codex-rs/rmcp-client/src/http_client_adapter.rs index d6c6b02c3c..f2e2cf55a3 100644 --- a/codex-rs/rmcp-client/src/http_client_adapter.rs +++ b/codex-rs/rmcp-client/src/http_client_adapter.rs @@ -898,11 +898,11 @@ pub(super) fn legacy_discovery_fallback_response( } fn only_known_legacy_protocol_versions(versions: &[ProtocolVersion]) -> bool { - !versions.is_empty() - && versions.iter().all(|version| { - ProtocolVersion::KNOWN_VERSIONS.contains(version) - && version < &ProtocolVersion::V_2026_07_28 - }) + !versions.is_empty() && versions.iter().all(is_known_legacy_protocol_version) +} + +pub(super) fn is_known_legacy_protocol_version(version: &ProtocolVersion) -> bool { + ProtocolVersion::KNOWN_VERSIONS.contains(version) && version < &ProtocolVersion::V_2026_07_28 } fn explicitly_rejects_modern_protocol_version(message: &str) -> bool { diff --git a/codex-rs/rmcp-client/src/local_stdio_transport.rs b/codex-rs/rmcp-client/src/local_stdio_transport.rs index 381aad1a99..2e908278dc 100644 --- a/codex-rs/rmcp-client/src/local_stdio_transport.rs +++ b/codex-rs/rmcp-client/src/local_stdio_transport.rs @@ -6,40 +6,150 @@ use std::future::Future; use std::io; use std::time::Duration; -use codex_utils_pty::Command; use futures::FutureExt; +use rmcp::model::ClientRequest; +use rmcp::model::ErrorCode; +use rmcp::model::JsonRpcMessage; +use rmcp::model::RequestId; +use rmcp::model::ServerResult; use rmcp::service::RoleClient; use rmcp::service::RxJsonRpcMessage; use rmcp::service::TxJsonRpcMessage; use rmcp::transport::Transport; use rmcp::transport::async_rw::AsyncRwTransport; +use tokio::io::AsyncRead; use tokio::process::ChildStderr; use tokio::process::ChildStdin; -use tokio::process::ChildStdout; +use tokio::sync::oneshot; use crate::bounded_stdio_transport::BoundedStdioTransport; use crate::protocol_mode::McpProtocolMode; use codex_utils_pty::Child; pub(super) struct LocalStdioTransport { - child: Box, + child: Option>, transport: StdioTransport, + discovery: Option, +} + +pub(super) enum DiscoveryOutcome { + Reply, + RetryLegacy, + Late { + reply: Box>, + after_reoffer: bool, + }, +} + +struct Discovery { + finished: Option>, + outcome: Option, + retry_at: Option, + probe_id: Option, + initialize_id: Option, +} + +impl Drop for Discovery { + fn drop(&mut self) { + // A version re-offer can supersede an earlier reply. Publish only when startup ends. + if let Some(tx) = self.finished.take() + && let Some(outcome) = self.outcome.take() + { + let _ = tx.send(outcome); + } + } +} + +impl Discovery { + /// Track discovery replies and skip messages that do not answer the active handshake. + fn handle_reply(&mut self, message: &mut Option>) -> bool { + let Some(reply) = message.as_ref() else { + self.outcome.get_or_insert(DiscoveryOutcome::RetryLegacy); + return false; + }; + let (id, modern) = match reply { + JsonRpcMessage::Response(response) => ( + Some(&response.id), + matches!(response.result, ServerResult::DiscoverResult(_)), + ), + JsonRpcMessage::Error(error) => ( + error.id.as_ref(), + matches!( + error.error.code, + ErrorCode::UNSUPPORTED_PROTOCOL_VERSION + | ErrorCode::HEADER_MISMATCH + | ErrorCode::MISSING_REQUIRED_CLIENT_CAPABILITY + ), + ), + _ => return false, + }; + let matches_id = |sent: &RequestId| match (sent, id) { + (RequestId::Number(sent), Some(RequestId::String(received))) => { + received.parse::().ok() == Some(*sent) + } + (_, Some(received)) => sent == received, + (_, None) => false, + }; + let legacy = self.initialize_id.is_some(); + let matches_probe = self.probe_id.as_ref().is_some_and(matches_id); + let idless = self.probe_id.is_some() && id.is_none(); + if legacy + && (matches_probe || idless) + && matches!(reply, JsonRpcMessage::Error(error) if error.error.code == ErrorCode::METHOD_NOT_FOUND) + { + // If initialize also fails, try it as the first request on a fresh helper. + self.outcome.get_or_insert(DiscoveryOutcome::RetryLegacy); + } + // With one request rmcp accepts an ID-less error. With two, only a modern + // protocol error can be classified as a discovery reply without an ID. + if matches_probe || (idless && (!legacy || modern)) { + self.probe_id = None; + if legacy { + // Auto sent initialize before receiving the reply to its earlier probe. + if !modern && matches!(reply, JsonRpcMessage::Error(_)) { + return true; + } + // Stop the obsolete initialize and let the caller classify the reply. + self.outcome = message.take().map(|reply| DiscoveryOutcome::Late { + reply: Box::new(reply), + after_reoffer: matches!(self.outcome, Some(DiscoveryOutcome::Reply)), + }); + return false; + } + if modern { + self.retry_at = None; + self.outcome = Some(DiscoveryOutcome::Reply); + } + } else if self.probe_id.is_some() + && id.is_some() + && !self.initialize_id.as_ref().is_some_and(matches_id) + { + // rmcp would abort the handshake before reading the reply to discover or initialize. + return true; + } else if legacy && matches!(reply, JsonRpcMessage::Response(_)) { + // Once initialize wins, rmcp owns the session and can ignore the obsolete probe. + self.probe_id = None; + self.retry_at = None; + self.outcome = Some(DiscoveryOutcome::Reply); + } + false + } } enum StdioTransport { /// Preserve rmcp's existing framing for servers using the initialize handshake. - Legacy(AsyncRwTransport), + Legacy(AsyncRwTransport, ChildStdin>), /// Bound frames and skip messages unknown to the client during 2026-07-28 discovery. - V20260728(BoundedStdioTransport), + Modern(BoundedStdioTransport), + Closed, } impl LocalStdioTransport { - pub(super) fn spawn( - command: Command, + pub(super) fn new( + mut child: Child, program_name: String, protocol_mode: McpProtocolMode, ) -> io::Result<(Self, Option)> { - let mut child = command.spawn()?; let stdin = child .stdin .take() @@ -50,22 +160,46 @@ impl LocalStdioTransport { .ok_or_else(|| io::Error::other("MCP server stdout was not piped"))?; let stderr = child.stderr.take(); let transport = match protocol_mode { - McpProtocolMode::Legacy => StdioTransport::Legacy(AsyncRwTransport::new(stdout, stdin)), - McpProtocolMode::V20260728 => { - StdioTransport::V20260728(BoundedStdioTransport::new(stdin, stdout, program_name)) + McpProtocolMode::Legacy => { + StdioTransport::Legacy(AsyncRwTransport::new(Box::new(stdout), stdin)) + } + McpProtocolMode::V20260728 | McpProtocolMode::DiscoverLocalStdio => { + StdioTransport::Modern(BoundedStdioTransport::new(stdin, stdout, program_name)) } }; Ok(( Self { - child: Box::new(child), + child: Some(Box::new(child)), transport, + discovery: None, }, stderr, )) } - pub(super) fn id(&self) -> Option { - self.child.id() + pub(super) fn watch_discovery( + &mut self, + timeout: Option, + ) -> oneshot::Receiver { + let (tx, rx) = oneshot::channel(); + self.discovery = Some(Discovery { + finished: Some(tx), + outcome: None, + // Aim for three seconds of discovery. Reserve at least a quarter of short + // timeouts and half of longer ones for a fresh initialize. + retry_at: timeout.and_then(|limit| { + let probe_timeout = Duration::from_secs(3).clamp(limit / 2, limit - limit / 4); + tokio::time::Instant::now().checked_add(probe_timeout) + }), + probe_id: None, + initialize_id: None, + }); + rx + } + + #[cfg(unix)] + pub(super) fn take_child(&mut self) -> Option> { + self.child.take() } } @@ -76,27 +210,87 @@ impl Transport for LocalStdioTransport { &mut self, item: TxJsonRpcMessage, ) -> impl Future> + Send + 'static { + if self.discovery.as_ref().is_some_and(|probe| { + probe + .finished + .as_ref() + .is_none_or(oneshot::Sender::is_closed) + }) { + self.discovery = None; + } + if let Some(probe) = &mut self.discovery + && let JsonRpcMessage::Request(request) = &item + && matches!(request.request, ClientRequest::DiscoverRequest(_)) + { + probe.probe_id = Some(request.id.clone()); + } + if let Some(probe) = &mut self.discovery + && let JsonRpcMessage::Request(request) = &item + && matches!(request.request, ClientRequest::InitializeRequest(_)) + && matches!(self.transport, StdioTransport::Modern(_)) + { + let StdioTransport::Modern(transport) = + std::mem::replace(&mut self.transport, StdioTransport::Closed) + else { + unreachable!(); + }; + match transport.into_legacy_pipes() { + Ok((stdout, stdin)) => { + probe.initialize_id = Some(request.id.clone()); + self.transport = StdioTransport::Legacy(AsyncRwTransport::new(stdout, stdin)); + } + Err(error) => return futures::future::err(error).boxed(), + } + } match &mut self.transport { StdioTransport::Legacy(transport) => transport.send(item).boxed(), - StdioTransport::V20260728(transport) => transport.send(item).boxed(), + StdioTransport::Modern(transport) => transport.send(item).boxed(), + StdioTransport::Closed => futures::future::err(io::Error::new( + io::ErrorKind::BrokenPipe, + "MCP stdin closed", + )) + .boxed(), } } - fn receive(&mut self) -> impl Future>> + Send { - match &mut self.transport { - StdioTransport::Legacy(transport) => transport.receive().boxed(), - StdioTransport::V20260728(transport) => transport.receive().boxed(), + async fn receive(&mut self) -> Option> { + loop { + let receive = async { + match &mut self.transport { + StdioTransport::Legacy(transport) => transport.receive().await, + StdioTransport::Modern(transport) => transport.receive().await, + StdioTransport::Closed => None, + } + }; + let mut message = match self.discovery.as_ref().and_then(|probe| probe.retry_at) { + Some(retry_at) => tokio::time::timeout_at(retry_at, receive) + .await + .unwrap_or_default(), + None => receive.await, + }; + if self + .discovery + .as_mut() + .is_some_and(|probe| probe.handle_reply(&mut message)) + { + continue; + } + return message; } } async fn close(&mut self) -> io::Result<()> { match &mut self.transport { StdioTransport::Legacy(transport) => transport.close().await?, - StdioTransport::V20260728(transport) => transport.close().await?, + StdioTransport::Modern(transport) => transport.close().await?, + StdioTransport::Closed => {} } - match tokio::time::timeout(Duration::from_secs(3), self.child.wait()).await { + let Some(child) = &mut self.child else { + return Ok(()); + }; + match tokio::time::timeout(Duration::from_secs(3), child.wait()).await { Ok(status) => status.map(|_| ()), - Err(_) => self.child.kill().await, + Err(_) => child.kill().await, } } } diff --git a/codex-rs/rmcp-client/src/protocol_mode.rs b/codex-rs/rmcp-client/src/protocol_mode.rs index 656ba8f6e9..35d673a86c 100644 --- a/codex-rs/rmcp-client/src/protocol_mode.rs +++ b/codex-rs/rmcp-client/src/protocol_mode.rs @@ -12,6 +12,8 @@ pub enum McpProtocolMode { Legacy, /// Allow the MCP 2026-07-28 discovery and request lifecycle. V20260728, + /// Probe an eligible local stdio server even when its launch config omits a version. + DiscoverLocalStdio, } impl McpProtocolMode { @@ -19,27 +21,27 @@ impl McpProtocolMode { pub fn preferred_protocol_version(self) -> ProtocolVersion { match self { Self::Legacy => ProtocolVersion::V_2025_06_18, - Self::V20260728 => ProtocolVersion::V_2026_07_28, + Self::V20260728 | Self::DiscoverLocalStdio => ProtocolVersion::V_2026_07_28, } } pub(crate) fn client_lifecycle(self) -> ClientLifecycleMode { match self { Self::Legacy => ClientLifecycleMode::Initialize, - Self::V20260728 => ClientLifecycleMode::Auto { + Self::V20260728 | Self::DiscoverLocalStdio => ClientLifecycleMode::Auto { preferred_versions: vec![ProtocolVersion::V_2026_07_28], legacy_version: Some(ProtocolVersion::V_2025_06_18), }, } } - pub(crate) fn stdio_mode(self, requested_version: Option<&OsStr>) -> io::Result { + /// Selects the stdio handshake from the client policy and the server's declared version. + pub fn stdio_mode(self, requested_version: Option<&OsStr>) -> io::Result { match (self, requested_version) { (Self::Legacy, _) => Ok(Self::Legacy), - (_, None) => Ok(Self::Legacy), - (Self::V20260728, Some(version)) if version == OsStr::new("2026-07-28") => { - Ok(Self::V20260728) - } + (Self::DiscoverLocalStdio, None) => Ok(Self::DiscoverLocalStdio), + (Self::V20260728, None) => Ok(Self::Legacy), + (_, Some(version)) if version == OsStr::new("2026-07-28") => Ok(Self::V20260728), (_, Some(version)) => Err(io::Error::new( io::ErrorKind::InvalidInput, format!( diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index ed3a144d5d..46b86cac6b 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -103,14 +103,17 @@ use crate::ordered_byte_transport::OrderedByteTransport; use crate::protocol_mode::McpProtocolMode; use crate::request_cancellation::RequestCancellation; use crate::startup_error::is_authentication_required_error; +use crate::stdio_server_launcher::StdioProcessSlot; use crate::stdio_server_launcher::StdioServerCommand; use crate::stdio_server_launcher::StdioServerLauncher; -use crate::stdio_server_launcher::StdioServerProcessHandle; use crate::stdio_server_launcher::StdioServerTransport; +use crate::stdio_server_launcher::StdioTermination; use crate::tool_list_changes::track_tool_list_changes; use crate::utils::build_default_headers; use codex_config::types::OAuthCredentialsStoreMode; +#[path = "stdio_handshake.rs"] +mod stdio_handshake; #[path = "streamable_http_retry.rs"] mod streamable_http_retry; @@ -410,8 +413,9 @@ pub struct CancellableEventStreamRequest { /// MCP client implemented on top of the official `rmcp` SDK. /// https://github.com/modelcontextprotocol/rust-sdk pub struct RmcpClient { + // Drop our interest in retrying before the SDK transport uses its ordinary shutdown grace. + stdio_process: Option, state: Mutex, - stdio_process: Option, transport_recipe: TransportRecipe, protocol_mode: McpProtocolMode, requires_read_only_tools: bool, @@ -516,7 +520,7 @@ impl RmcpClient { ) -> io::Result { let requested_stdio_version = match protocol_mode { McpProtocolMode::Legacy => None, - McpProtocolMode::V20260728 => env + McpProtocolMode::V20260728 | McpProtocolMode::DiscoverLocalStdio => env .as_mut() .and_then(|env| env.remove(OsStr::new("CODEX_MCP_PROTOCOL_VERSION"))), }; @@ -536,7 +540,7 @@ impl RmcpClient { .await .map_err(io::Error::other)?; let stdio_process = match &transport { - PendingTransport::Stdio { transport } => Some(transport.process_handle()), + PendingTransport::Stdio { transport } => Some(transport.process_slot()), PendingTransport::InProcess { .. } | PendingTransport::StreamableHttp { .. } | PendingTransport::StreamableHttpWithOAuth { .. } @@ -897,7 +901,7 @@ impl RmcpClient { params.meta = crate::trace_context::with_current_trace(params.meta); self.refresh_oauth_if_needed().await?; let cancellation = self.request_cancellation(); - let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; + let requested_modern = self.protocol_mode != McpProtocolMode::Legacy; let result = self .run_service_operation("resources/read", timeout, move |service| { let params = params.clone(); @@ -1062,7 +1066,7 @@ impl RmcpClient { }; let mut rmcp_params = CallToolRequestParams::new(name); rmcp_params.arguments = arguments; - let requested_modern = self.protocol_mode == McpProtocolMode::V20260728; + let requested_modern = self.protocol_mode != McpProtocolMode::Legacy; let meta_for_session = &meta_for_session; let pause_state = if send_elicitation.is_some() { ElicitationPauseState::new() @@ -1309,18 +1313,34 @@ impl RmcpClient { /// Stop the MCP transport and any stdio server process owned by this client. pub async fn shutdown(&self) { + if let Err(error) = self.shutdown_for(StdioTermination::Shutdown).await { + warn!("failed to terminate MCP stdio server process: {error}"); + } + } + + /// Wait for a failed discovery process to release resources needed by a subsequent launch. + pub async fn shutdown_before_retry(&self) -> io::Result<()> { + self.shutdown_for(StdioTermination::Restart).await + } + + async fn shutdown_for(&self, reason: StdioTermination) -> io::Result<()> { let previous_state = { let mut guard = self.state.lock().await; std::mem::replace(&mut *guard, ClientState::Closed) }; - if let Some(process) = &self.stdio_process - && let Err(error) = process.terminate().await - { - warn!("failed to terminate MCP stdio server process: {error}"); + let process = match &self.stdio_process { + Some(slot) => Some(slot.lock().await.close()), + None => None, + }; + if matches!(reason, StdioTermination::Restart) { + // Let the helper see EOF before escalating to process-group termination. + drop(previous_state); } - - drop(previous_state); + if let Some(process) = process { + process.terminate(reason).await?; + } + Ok(()) } /// This should be called after every tool call so that if a given tool call triggered @@ -1539,12 +1559,18 @@ impl RmcpClient { )> { // Request IDs and remembered cancellations belong to this connection, including // when a failed initialization or expired HTTP session creates a new transport. - let send_elicitation = Arc::clone(&initialize_context.send_elicitation); - let client_service = ElicitationClientService::new( - initialize_context.client_info.clone(), - Box::new(move |id, request| send_elicitation(id, request)), - self.elicitation_pause_state.clone(), - ); + let client_service = || { + let send_elicitation = Arc::clone(&initialize_context.send_elicitation); + let mut info = initialize_context.client_info.clone(); + if self.protocol_mode == McpProtocolMode::DiscoverLocalStdio { + info.protocol_version = McpProtocolMode::Legacy.preferred_protocol_version(); + } + ElicitationClientService::new( + info, + Box::new(move |id, request| send_elicitation(id, request)), + self.elicitation_pause_state.clone(), + ) + }; let _initialize_deadline = match &self.transport_recipe { TransportRecipe::StreamableHttp { initialize_deadline, @@ -1561,9 +1587,21 @@ impl RmcpClient { TransportRecipe::InProcess { .. } | TransportRecipe::Stdio { .. } => None, }; let lifecycle = self.protocol_mode.client_lifecycle(); + let connect_stdio = |transport: StdioServerTransport, lifecycle| { + client_service() + .serve_with_lifecycle( + track_tool_list_changes( + OrderedByteTransport::new(transport), + Arc::clone(&self.tool_list_revision), + ), + lifecycle, + ) + .boxed() + }; + let mut discovery = None; let (transport, oauth_runtime) = match pending_transport { PendingTransport::InProcess { transport } => ( - client_service + client_service() .serve_with_lifecycle( track_tool_list_changes( OrderedByteTransport::new( @@ -1576,20 +1614,17 @@ impl RmcpClient { .boxed(), None, ), - PendingTransport::Stdio { transport } => ( - client_service - .serve_with_lifecycle( - track_tool_list_changes( - OrderedByteTransport::new(*transport), - Arc::clone(&self.tool_list_revision), - ), - lifecycle, - ) - .boxed(), - None, - ), + PendingTransport::Stdio { mut transport } => { + if transport.protocol_mode() == McpProtocolMode::DiscoverLocalStdio { + discovery = transport + .watch_discovery(timeout) + .map(|reply| (reply, transport.process_ref())); + } + let lifecycle = transport.protocol_mode().client_lifecycle(); + (connect_stdio(*transport, lifecycle), None) + } PendingTransport::StreamableHttp { transport } => ( - client_service + client_service() .serve_with_lifecycle( track_tool_list_changes( capture_event_notifications(transport), @@ -1604,7 +1639,7 @@ impl RmcpClient { transport, oauth_runtime, } => ( - client_service + client_service() .serve_with_lifecycle( track_tool_list_changes(transport, Arc::clone(&self.tool_list_revision)), lifecycle, @@ -1613,7 +1648,7 @@ impl RmcpClient { Some(oauth_runtime), ), PendingTransport::StreamableHttpWithAccessTokenOnly { transport } => ( - client_service + client_service() .serve_with_lifecycle( track_tool_list_changes(transport, Arc::clone(&self.tool_list_revision)), lifecycle, @@ -1623,18 +1658,38 @@ impl RmcpClient { ), }; + let handshake = async { + let mut result = transport.await; + if let (Some((mut reply, process)), Err(error)) = (discovery, &result) + && let Some(mode) = stdio_handshake::restart_mode(error, reply.try_recv().ok())? + { + process + .terminate(StdioTermination::Restart) + .await + .context("stop stdio server before handshake restart")?; + let fresh = self.launch_stdio_replacement(mode).await?; + // A confirmed modern helper doesn't need another legacy probe timer. + // Windows can still select initialize if it could not contain the new process. + let lifecycle = match fresh.protocol_mode() { + McpProtocolMode::DiscoverLocalStdio => { + rmcp::service::ClientLifecycleMode::Discover { + preferred_versions: vec![ProtocolVersion::V_2026_07_28], + } + } + mode => mode.client_lifecycle(), + }; + result = connect_stdio(fresh, lifecycle).await; + } + result.map_err(|source| HandshakeError { source }.into()) + }; let service_result = match timeout { - Some(duration) => match time::timeout(duration, transport).await { - Ok(result) => { - result.map_err(|source| anyhow::Error::from(HandshakeError { source })) - } - Err(_elapsed) => Err(anyhow!( - "timed out handshaking with MCP server after {duration:?}" - )), - }, - None => transport + Some(duration) => time::timeout(duration, handshake) .await - .map_err(|source| anyhow::Error::from(HandshakeError { source })), + .with_context(|| { + format!("timed out handshaking with MCP server after {duration:?}") + }) + .flatten(), + None => handshake.await, }; let service = match service_result { Ok(service) => service, @@ -1660,6 +1715,23 @@ impl RmcpClient { Ok((Arc::new(service), oauth_runtime)) } + async fn launch_stdio_replacement( + &self, + mode: McpProtocolMode, + ) -> Result { + let TransportRecipe::Stdio { command, launcher } = &self.transport_recipe else { + unreachable!("stdio discovery uses a stdio launch recipe"); + }; + let slot = self + .stdio_process + .as_ref() + .context("missing stdio process")?; + launcher + .launch(command.for_restart(mode, Arc::clone(slot))) + .await + .context("restart stdio server") + } + async fn run_service_operation( &self, label: &str, diff --git a/codex-rs/rmcp-client/src/stdio_handshake.rs b/codex-rs/rmcp-client/src/stdio_handshake.rs new file mode 100644 index 0000000000..edb513c30a --- /dev/null +++ b/codex-rs/rmcp-client/src/stdio_handshake.rs @@ -0,0 +1,115 @@ +//! Selects the handshake for a fresh local helper when automatic discovery cannot finish in place. +//! A late discovery response can still negotiate the modern protocol after rmcp sent initialize. + +use std::io; + +use rmcp::model::ErrorCode; +use rmcp::model::ErrorData; +use rmcp::model::JsonRpcMessage; +use rmcp::model::ProtocolVersion; +use rmcp::model::ServerResult; +use rmcp::service::ClientInitializeError; + +use super::HandshakeError; +use crate::McpProtocolMode; +use crate::http_client_adapter::is_known_legacy_protocol_version; +use crate::local_stdio_transport::DiscoveryOutcome; + +pub(super) fn restart_mode( + error: &ClientInitializeError, + discovery: Option, +) -> anyhow::Result> { + let Some(DiscoveryOutcome::Late { + reply, + after_reoffer, + }) = discovery + else { + return Ok(should_restart_legacy(error, discovery).then_some(McpProtocolMode::Legacy)); + }; + match *reply { + JsonRpcMessage::Response(response) => match response.result { + ServerResult::DiscoverResult(result) => { + let supported = result.supported_versions; + restart_for_versions(&supported).map(Some).ok_or_else(|| { + ClientInitializeError::NoCompatibleProtocolVersion { + client_supported: vec![ProtocolVersion::V_2026_07_28], + server_supported: supported, + } + }) + } + ServerResult::InitializeResult(result) + if is_known_legacy_protocol_version(&result.protocol_version) => + { + Ok(Some(McpProtocolMode::Legacy)) + } + result => Err(ClientInitializeError::ExpectedInitResult(Some(result))), + }, + JsonRpcMessage::Error(response) => { + let mode = offered_restart_mode(&response.error); + match mode { + // rmcp permits one re-offer. A fresh discover verifies it even without an ID. + Some(McpProtocolMode::DiscoverLocalStdio) if !after_reoffer => Ok(mode), + // Without an ID the error may have rejected initialize itself. + Some(McpProtocolMode::Legacy) if response.id.is_some() => Ok(mode), + _ => Err(ClientInitializeError::JsonRpcError(response.error)), + } + } + _ => unreachable!("a late discovery reply is a response or error"), + } + .map_err(|source| HandshakeError { source }.into()) +} + +fn offered_restart_mode(error: &ErrorData) -> Option { + if error.code != ErrorCode::UNSUPPORTED_PROTOCOL_VERSION { + return None; + } + let supported = error.data.as_ref()?.get("supported")?; + let versions = serde_json::from_value::>(supported.clone()).ok()?; + restart_for_versions(&versions) +} + +fn restart_for_versions(supported: &[ProtocolVersion]) -> Option { + if supported.contains(&ProtocolVersion::V_2026_07_28) { + Some(McpProtocolMode::DiscoverLocalStdio) + } else if supported.iter().any(is_known_legacy_protocol_version) { + Some(McpProtocolMode::Legacy) + } else { + None + } +} + +fn should_restart_legacy( + error: &ClientInitializeError, + discovery: Option, +) -> bool { + match error { + // A helper that assumes its first request is initialize can answer discover with this. + ClientInitializeError::ExpectedInitResult(Some(ServerResult::InitializeResult(result))) => { + !matches!(discovery, Some(DiscoveryOutcome::Reply)) + && is_known_legacy_protocol_version(&result.protocol_version) + } + ClientInitializeError::NoCompatibleProtocolVersion { + server_supported, .. + } => restart_for_versions(server_supported) == Some(McpProtocolMode::Legacy), + ClientInitializeError::LegacyFallbackFailed { discover, fallback } => { + match fallback.as_ref() { + ClientInitializeError::ConnectionClosed(_) => true, + // A legacy helper can reject initialize if discover already took its first request. + ClientInitializeError::JsonRpcError(_) => matches!( + discover.as_ref(), + ClientInitializeError::JsonRpcError(error) if error.code == ErrorCode::METHOD_NOT_FOUND + ), + ClientInitializeError::TransportError { error, .. } => { + error.error.downcast_ref::().is_some_and(|e| { + matches!( + e.kind(), + io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset + ) + }) + } + _ => false, + } + } + _ => matches!(discovery, Some(DiscoveryOutcome::RetryLegacy)), + } +} diff --git a/codex-rs/rmcp-client/src/stdio_server_launcher.rs b/codex-rs/rmcp-client/src/stdio_server_launcher.rs index 92ccfa2107..4925929571 100644 --- a/codex-rs/rmcp-client/src/stdio_server_launcher.rs +++ b/codex-rs/rmcp-client/src/stdio_server_launcher.rs @@ -21,6 +21,10 @@ use std::os::windows::io::OwnedHandle; use std::path::Path; use std::path::PathBuf; use std::sync::Arc; +#[cfg(unix)] +use std::sync::Mutex; +use std::sync::OnceLock; +use std::sync::Weak; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; #[cfg(unix)] @@ -29,6 +33,7 @@ use std::thread::sleep; use std::thread::spawn; use std::time::Duration; +use crate::local_stdio_transport::DiscoveryOutcome; use anyhow::Result; use anyhow::anyhow; use codex_config::types::McpServerEnvVar; @@ -39,6 +44,8 @@ use codex_exec_server::ExecProcess; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_utils_path_uri::LegacyAppPathString; use codex_utils_path_uri::PathUri; +#[cfg(unix)] +use codex_utils_pty::Child; use codex_utils_pty::Command; #[cfg(unix)] use codex_utils_pty::DescriptorPolicy; @@ -47,16 +54,23 @@ use codex_utils_pty::ProcessMode; use codex_utils_pty::process_group::kill_process_group; #[cfg(unix)] use codex_utils_pty::process_group::terminate_process_group; +#[cfg(unix)] +use codex_utils_pty::process_group::wait_for_process_group_exit; use futures::FutureExt; use futures::future::BoxFuture; +use futures::future::Shared; use rmcp::service::RoleClient; use rmcp::service::RxJsonRpcMessage; use rmcp::service::TxJsonRpcMessage; use rmcp::transport::Transport; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; +use tokio::runtime::Handle; +use tokio::sync::Mutex as AsyncMutex; +use tokio::sync::oneshot; use tokio::sync::watch; use tokio::time::Instant; +use tokio_util::task::AbortOnDropHandle; use tracing::info; use tracing::warn; @@ -93,6 +107,70 @@ pub struct StdioServerCommand { env_vars: Vec, cwd: Option, protocol_mode: McpProtocolMode, + replacement_slot: Option, +} + +/// A replacement holds the slot from OS spawn through registration. Closing it seals that handoff. +pub(super) type StdioProcessSlot = Arc>; + +pub(super) struct StdioProcessRegistration { + process: StdioProcessRef, + closed: bool, +} + +impl StdioProcessRegistration { + pub(super) fn close(&mut self) -> StdioProcessRef { + self.closed = true; + self.process.clone() + } +} + +type DiscoveryCleanupResult = Shared>>>; + +struct DiscoveryCleanup { + process: Weak, + result: OnceLock, +} + +impl DiscoveryCleanup { + fn start(&self, runtime: &Handle, reason: StdioTermination) -> &DiscoveryCleanupResult { + self.result.get_or_init(|| { + let process = self.process.upgrade(); + // The owner waits for this worker on shutdown. Its grace periods must also advance + // when the caller leaves a current-thread runtime idle. + runtime + .spawn_blocking(move || { + process + .ok_or_else(|| io::Error::other("discovery process already released"))? + .terminate_discovery(reason) + }) + .map(|joined| match joined { + Ok(result) => result.map_err(Arc::new), + Err(error) => Err(Arc::new(io::Error::other(error))), + }) + .boxed() + .shared() + }) + } +} + +#[derive(Clone)] +pub(super) struct StdioProcessRef { + process: Weak, + discovery_cleanup: Option>, +} + +impl StdioProcessRef { + pub(super) async fn terminate(&self, reason: StdioTermination) -> io::Result<()> { + if let Some(cleanup) = &self.discovery_cleanup { + let result = cleanup.start(&Handle::current(), reason); + return result.clone().await.map_err(io::Error::other); + } + if let Some(process) = self.process.upgrade() { + process.terminate().await?; + } + Ok(()) + } } /// Client-side rmcp transport for a launched MCP stdio server. @@ -102,7 +180,9 @@ pub struct StdioServerCommand { /// directly to `rmcp::service::serve_client`. pub struct StdioServerTransport { inner: StdioServerTransportInner, - process: StdioServerProcessHandle, + process: Arc, + discovery_cleanup: Option>, + protocol_mode: McpProtocolMode, } enum StdioServerTransportInner { @@ -137,7 +217,9 @@ impl Transport for StdioServerTransport { } async fn close(&mut self) -> std::result::Result<(), Self::Error> { - self.process.terminate().await?; + if self.discovery_cleanup.is_none() { + self.process.terminate().await?; + } match &mut self.inner { StdioServerTransportInner::Local(transport) => transport.close().await, StdioServerTransportInner::Executor(transport) => transport.close().await, @@ -145,9 +227,48 @@ impl Transport for StdioServerTransport { } } +impl Drop for StdioServerTransport { + fn drop(&mut self) { + // A client or handshake still holding this result may retry after the SDK drops a + // failed transport. Otherwise the process handle runs its ordinary shutdown on drop. + if let Some(cleanup) = &self.discovery_cleanup + && Arc::strong_count(cleanup) > 1 + && !self.process.terminated.load(Ordering::Acquire) + && let Ok(runtime) = Handle::try_current() + { + cleanup.start(&runtime, StdioTermination::Restart); + } + } +} + impl StdioServerTransport { - pub(crate) fn process_handle(&self) -> StdioServerProcessHandle { - self.process.clone() + pub(super) fn protocol_mode(&self) -> McpProtocolMode { + self.protocol_mode + } + + /// Observes the first discovery reply or a clean EOF during local stdio startup. + pub(super) fn watch_discovery( + &mut self, + timeout: Option, + ) -> Option> { + match &mut self.inner { + StdioServerTransportInner::Local(local) => Some(local.watch_discovery(timeout)), + StdioServerTransportInner::Executor(_) => None, + } + } + + pub(super) fn process_ref(&self) -> StdioProcessRef { + StdioProcessRef { + process: Arc::downgrade(&self.process), + discovery_cleanup: self.discovery_cleanup.clone(), + } + } + + pub(super) fn process_slot(&self) -> StdioProcessSlot { + Arc::new(AsyncMutex::new(StdioProcessRegistration { + process: self.process_ref(), + closed: false, + })) } } @@ -169,6 +290,20 @@ impl StdioServerCommand { env_vars, cwd, protocol_mode, + replacement_slot: None, + } + } + + /// Registers the replacement with its original client before the native spawn can finish. + pub(super) fn for_restart( + &self, + protocol_mode: McpProtocolMode, + slot: StdioProcessSlot, + ) -> Self { + Self { + protocol_mode, + replacement_slot: Some(slot), + ..self.clone() } } } @@ -202,11 +337,13 @@ impl StdioServerLauncher for LocalStdioServerLauncher { ) -> BoxFuture<'static, io::Result> { let fallback_cwd = self.fallback_cwd.clone(); async move { - // Keep synchronous program resolution and process creation from blocking the - // caller's startup deadline. - tokio::task::spawn_blocking(move || Self::launch_server(command, fallback_cwd)) - .await - .map_err(io::Error::other)? + // Cancel queued launches when startup is dropped. A worker that already started + // drops its transport and cleans up the child when it finishes. + AbortOnDropHandle::new(tokio::task::spawn_blocking(move || { + Self::launch_server(command, fallback_cwd) + })) + .await + .map_err(io::Error::other)? } .boxed() } @@ -217,30 +354,75 @@ impl StdioServerLauncher for LocalStdioServerLauncher { #[cfg(unix)] const PROCESS_GROUP_TERM_GRACE_PERIOD: Duration = Duration::from_secs(2); +#[cfg(unix)] +const RESTART_DRAIN_GRACE_PERIOD: Duration = Duration::from_millis(250); + // Keep queued stderr diagnostics before closing the reader, even when an // escaped descendant prevents the pipe from reaching EOF. const STDERR_READER_DRAIN_GRACE_PERIOD: Duration = Duration::from_millis(250); #[cfg(unix)] -struct LocalProcessTerminator { - process_group_id: u32, +enum LocalProcessTerminator { + ProcessGroup(u32), + // The SDK may drop the failed transport before the caller can decide whether to restart it. + DiscoveryLeader { + leader: Mutex>>, + runtime: Handle, + }, +} + +#[cfg(unix)] +#[derive(Clone)] +struct KillDiscoveryGroupOnDrop(Arc>>>); + +#[cfg(unix)] +impl KillDiscoveryGroupOnDrop { + fn new(child: Box) -> Self { + Self(Arc::new(Mutex::new(Some(child)))) + } + + fn kill(&self) -> io::Result<()> { + // Runtime teardown must wait for a timer that has already started sending the signal. + let mut leader = self + .0 + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(child) = leader.as_ref() + && let Some(pid) = child.id() + { + kill_process_group_and_warn(pid)?; + } + leader.take(); + Ok(()) + } +} + +#[cfg(unix)] +impl Drop for KillDiscoveryGroupOnDrop { + fn drop(&mut self) { + let _ = self.kill(); + } +} + +#[cfg(unix)] +fn kill_process_group_and_warn(pid: u32) -> io::Result<()> { + kill_process_group(pid) + .inspect_err(|error| warn!("Failed to kill MCP process group {pid}: {error}")) } #[cfg(windows)] enum LocalProcessTerminator { - Job(codex_utils_pty::JobObject), + Job { + job: codex_utils_pty::JobObject, + discovery_root: Option, + }, Process(OwnedHandle), } #[cfg(not(any(unix, windows)))] struct LocalProcessTerminator; -#[derive(Clone)] pub(crate) struct StdioServerProcessHandle { - inner: Arc, -} - -struct StdioServerProcessHandleInner { program_name: String, kind: StdioServerProcessKind, terminated: AtomicBool, @@ -253,6 +435,12 @@ enum StdioServerProcessKind { Executor(Arc), } +#[derive(Clone, Copy)] +pub(super) enum StdioTermination { + Shutdown, + Restart, +} + mod private { pub trait Sealed {} } @@ -271,6 +459,7 @@ impl LocalStdioServerLauncher { env_vars, cwd, protocol_mode, + replacement_slot, } = command; let program_name = program.to_string_lossy().into_owned(); let envs = create_env_for_mcp_server(env, &env_vars).map_err(io::Error::other)?; @@ -306,45 +495,56 @@ impl LocalStdioServerLauncher { } }; - let spawn_transport = |command: Command| -> io::Result<( - StdioServerTransportInner, - Option, - Option, - )> { - let (transport, stderr) = - LocalStdioTransport::spawn(command, program_name.clone(), protocol_mode)?; - let process_id = transport.id(); - Ok(( - StdioServerTransportInner::Local(transport), - stderr, - process_id, - )) - }; - let (transport, stderr, process_id) = spawn_transport(command)?; + // The async caller can be cancelled while this blocking worker is inside spawn. + let slot = replacement_slot.as_ref().map(|slot| slot.blocking_lock()); + if slot.as_ref().is_some_and(|slot| slot.closed) { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "MCP client is shut down", + )); + } + let child = command.spawn()?; #[cfg(windows)] - let (transport, stderr, process_id, job) = match job { - Some(job) => match process_id + let (child, job) = match job { + Some(job) => match child + .id() .ok_or_else(|| io::Error::other("missing suspended MCP server process id")) .and_then(|process_id| job.assign_and_resume_process(process_id)) { - Ok(true) => (transport, stderr, process_id, Some(job)), - Ok(false) => (transport, stderr, process_id, None), + Ok(Some(root)) => (child, Some((job, root))), + Ok(None) => (child, None), Err(error) => { warn!( "Windows MCP process job containment failed; retrying without it: {error}" ); - drop(stderr); - drop(transport); + drop(child); drop(job); - let (transport, stderr, process_id) = spawn_transport(build_command())?; - (transport, stderr, process_id, None) + (build_command().spawn()?, None) } }, - None => (transport, stderr, process_id, None), + None => (child, None), }; #[cfg(windows)] + let protocol_mode = if job.is_none() && protocol_mode == McpProtocolMode::DiscoverLocalStdio + { + // A probe can require a restart. Without a job we cannot stop its descendants. + McpProtocolMode::Legacy + } else { + protocol_mode + }; + let process_id = child.id(); + // A legacy replacement still belongs to the discovery startup that launched it. + let discovery_owned = (protocol_mode == McpProtocolMode::DiscoverLocalStdio + || replacement_slot.is_some()) + && process_id.is_some(); + let (transport, stderr) = + LocalStdioTransport::new(child, program_name.clone(), protocol_mode)?; + #[cfg(windows)] let terminator = match job { - Some(job) => Some(LocalProcessTerminator::Job(job)), + Some((job, root)) => Some(LocalProcessTerminator::Job { + job, + discovery_root: discovery_owned.then_some(root), + }), None => process_id.and_then(|process_id| { match codex_utils_pty::JobObject::open_process_handle(process_id) { Ok(handle) => Some(LocalProcessTerminator::Process(handle)), @@ -355,8 +555,23 @@ impl LocalStdioServerLauncher { } }), }; - #[cfg(not(windows))] - let terminator = process_id.map(LocalProcessTerminator::new); + #[cfg(unix)] + let (transport, terminator) = { + let mut transport = transport; + let terminator = process_id.map(|pid| { + if discovery_owned { + LocalProcessTerminator::DiscoveryLeader { + leader: Mutex::new(transport.take_child()), + runtime: Handle::current(), + } + } else { + LocalProcessTerminator::ProcessGroup(pid) + } + }); + (transport, terminator) + }; + #[cfg(not(any(unix, windows)))] + let terminator = process_id.map(|_| LocalProcessTerminator); let stderr_reader = stderr.map(|stderr| { let program_name = program_name.clone(); let (stop_tx, mut stop_rx) = watch::channel(()); @@ -395,31 +610,77 @@ impl LocalStdioServerLauncher { stop_tx }); let process = StdioServerProcessHandle::local(program_name, terminator, stderr_reader); - - Ok(StdioServerTransport { - inner: transport, + let discovery_cleanup = discovery_owned.then(|| { + Arc::new(DiscoveryCleanup { + process: Arc::downgrade(&process), + result: OnceLock::new(), + }) + }); + let transport = StdioServerTransport { + inner: StdioServerTransportInner::Local(transport), process, - }) + discovery_cleanup, + protocol_mode, + }; + if let Some(mut slot) = slot { + slot.process = transport.process_ref(); + } + + Ok(transport) } } impl LocalProcessTerminator { - #[cfg(not(windows))] - fn new(process_group_id: u32) -> Self { - #[cfg(unix)] - { - Self { process_group_id } - } - #[cfg(not(any(unix, windows)))] - { - let _ = process_group_id; - Self + #[cfg(unix)] + fn terminate_discovery(&self, reason: StdioTermination) -> io::Result<()> { + let Self::DiscoveryLeader { leader, .. } = self else { + return Err(io::Error::other("missing Unix MCP discovery leader")); + }; + let child = leader + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + .ok_or_else(|| io::Error::other("Unix MCP discovery leader already released"))?; + let pid = child + .id() + .ok_or_else(|| io::Error::other("missing Unix MCP discovery leader PID"))?; + let cleanup = KillDiscoveryGroupOnDrop::new(child); + // A failed transport closed stdin. Let the helper release its endpoint before TERM, + // and reserve its PID through the last group signal. + let term_grace = match reason { + StdioTermination::Restart => { + sleep(RESTART_DRAIN_GRACE_PERIOD); + RESTART_DRAIN_GRACE_PERIOD + } + StdioTermination::Shutdown => PROCESS_GROUP_TERM_GRACE_PERIOD, + }; + if terminate_process_group(pid)? { + // A wrapper may exit before a child handling SIGTERM has released its own endpoint. + wait_for_process_group_exit(pid, term_grace); } + cleanup.kill()?; + wait_for_process_group_exit(pid, Duration::from_secs(1)); + Ok(()) } #[cfg(unix)] fn terminate(&self) { - let process_group_id = self.process_group_id; + let (process_group_id, discovery) = match self { + Self::ProcessGroup(pid) => (*pid, None), + Self::DiscoveryLeader { leader, runtime } => { + let Some(leader) = leader + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + else { + return; + }; + let Some(pid) = leader.id() else { + return; + }; + (pid, Some((runtime, leader))) + } + }; let should_escalate = match terminate_process_group(process_group_id) { Ok(exists) => exists, Err(error) => { @@ -428,19 +689,46 @@ impl LocalProcessTerminator { } }; if should_escalate { + let discovery = discovery.map(|(runtime, leader)| { + let on_timeout = KillDiscoveryGroupOnDrop::new(leader); + let on_runtime_exit = on_timeout.clone(); + let (finished, wait) = oneshot::channel::<()>(); + std::mem::drop(runtime.spawn(async move { + let _ = wait.await; + drop(on_runtime_exit); + })); + (on_timeout, finished) + }); + // A caller can leave a current-thread runtime idle for the entire grace period. spawn(move || { sleep(PROCESS_GROUP_TERM_GRACE_PERIOD); - if let Err(error) = kill_process_group(process_group_id) { - warn!("Failed to kill MCP process group {process_group_id}: {error}"); + if let Some((on_timeout, finished)) = discovery { + drop(on_timeout); + drop(finished); + } else { + let _ = kill_process_group_and_warn(process_group_id); } }); } } + #[cfg(windows)] + fn terminate_discovery(&self, _reason: StdioTermination) -> io::Result<()> { + match self { + Self::Job { + job, + discovery_root: Some(root), + } => job.terminate_and_wait(root, Duration::from_secs(1)), + _ => Err(io::Error::other( + "cannot restart an uncontained Windows MCP server", + )), + } + } + #[cfg(windows)] fn terminate(&self) { let result = match self { - Self::Job(job) => job.terminate(), + Self::Job { job, .. } => job.terminate(), Self::Process(process_handle) => { codex_utils_pty::JobObject::terminate_process_handle(process_handle) } @@ -450,6 +738,11 @@ impl LocalProcessTerminator { } } + #[cfg(not(any(unix, windows)))] + fn terminate_discovery(&self, _reason: StdioTermination) -> io::Result<()> { + Err(io::Error::other("MCP stdio restart is unsupported")) + } + #[cfg(not(any(unix, windows)))] fn terminate(&self) {} } @@ -459,55 +752,71 @@ impl StdioServerProcessHandle { program_name: String, terminator: Option, stderr_reader: Option>, - ) -> Self { - Self { - inner: Arc::new(StdioServerProcessHandleInner { - program_name, - kind: StdioServerProcessKind::Local(terminator), - terminated: AtomicBool::new(false), - stderr_reader, - }), - } + ) -> Arc { + Arc::new(Self { + program_name, + kind: StdioServerProcessKind::Local(terminator), + terminated: AtomicBool::new(false), + stderr_reader, + }) } - pub(crate) fn executor(program_name: String, process: Arc) -> Self { - Self { - inner: Arc::new(StdioServerProcessHandleInner { - program_name, - kind: StdioServerProcessKind::Executor(process), - terminated: AtomicBool::new(false), - stderr_reader: None, - }), + pub(crate) fn executor(program_name: String, process: Arc) -> Arc { + Arc::new(Self { + program_name, + kind: StdioServerProcessKind::Executor(process), + terminated: AtomicBool::new(false), + stderr_reader: None, + }) + } + + fn begin_termination(&self) -> bool { + if self.terminated.swap(true, Ordering::AcqRel) { + return false; } + if let Some(stderr_reader) = &self.stderr_reader { + stderr_reader.send_replace(()); + } + true + } + + fn terminate_discovery(&self, reason: StdioTermination) -> io::Result<()> { + if !self.begin_termination() { + return Err(io::Error::new( + io::ErrorKind::Interrupted, + "MCP stdio server is already shutting down", + )); + } + let StdioServerProcessKind::Local(Some(terminator)) = &self.kind else { + return Err(io::Error::other("missing local MCP discovery process")); + }; + // Stop an uncontained replacement, but don't permit a retry on it. + terminator + .terminate_discovery(reason) + .inspect_err(|_| terminator.terminate()) } pub(crate) async fn terminate(&self) -> io::Result<()> { - if self.inner.terminated.swap(true, Ordering::AcqRel) { + if !self.begin_termination() { return Ok(()); } - - let result = match &self.inner.kind { + match &self.kind { StdioServerProcessKind::Local(Some(terminator)) => { terminator.terminate(); Ok(()) } StdioServerProcessKind::Local(None) => Ok(()), - StdioServerProcessKind::Executor(process) => match process.terminate().await { - Ok(()) => Ok(()), - Err(error) => { - self.inner.terminated.store(false, Ordering::Release); - Err(io::Error::other(error)) - } - }, - }; - if let Some(stderr_reader) = &self.inner.stderr_reader { - stderr_reader.send_replace(()); + StdioServerProcessKind::Executor(process) => { + process.terminate().await.map_err(|error| { + self.terminated.store(false, Ordering::Release); + io::Error::other(error) + }) + } } - result } } -impl Drop for StdioServerProcessHandleInner { +impl Drop for StdioServerProcessHandle { fn drop(&mut self) { if self.terminated.swap(true, Ordering::AcqRel) { return; @@ -521,7 +830,7 @@ impl Drop for StdioServerProcessHandleInner { StdioServerProcessKind::Executor(process) => { let process = Arc::clone(process); let program_name = self.program_name.clone(); - let Ok(handle) = tokio::runtime::Handle::try_current() else { + let Ok(handle) = Handle::try_current() else { warn!( "Could not schedule remote MCP server process termination on drop ({}): no Tokio runtime is available", self.program_name @@ -593,7 +902,8 @@ impl ExecutorStdioServerLauncher { env, env_vars, cwd, - protocol_mode: _, + protocol_mode, + replacement_slot: _, } = command; let Some(cwd) = cwd else { return Err(io::Error::other( @@ -644,6 +954,8 @@ impl ExecutorStdioServerLauncher { program_name, )), process, + discovery_cleanup: None, + protocol_mode, }) } diff --git a/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs b/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs index 785aa4e195..ff28266210 100644 --- a/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs +++ b/codex-rs/rmcp-client/tests/mcp_2026_message_limits.rs @@ -234,13 +234,14 @@ async fn legacy_http_keeps_existing_large_json_response_behavior() -> anyhow::Re async fn modern_local_and_executor_stdio_reject_oversized_lines() -> anyhow::Result<()> { let server = codex_utils_cargo_bin::cargo_bin("test_mcp_2026_stdio_server")?; - for (modern, executor) in [(false, false), (false, true), (true, false), (true, true)] { - let mode = if modern { - McpProtocolMode::V20260728 - } else { - McpProtocolMode::Legacy - }; - let env = modern.then(|| { + for (mode, modern, executor) in [ + (McpProtocolMode::Legacy, false, false), + (McpProtocolMode::Legacy, false, true), + (McpProtocolMode::V20260728, true, false), + (McpProtocolMode::V20260728, true, true), + (McpProtocolMode::DiscoverLocalStdio, false, false), + ] { + let env = (mode == McpProtocolMode::V20260728).then(|| { HashMap::from([( OsString::from("CODEX_MCP_PROTOCOL_VERSION"), OsString::from("2026-07-28"), diff --git a/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs b/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs index 20f1c4eb45..42f16ba344 100644 --- a/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs +++ b/codex-rs/rmcp-client/tests/mcp_2026_stdio.rs @@ -46,21 +46,21 @@ async fn exercise_stdio_server( OsString::from("2026-07-28"), ); } - let cwd = std::env::current_dir()?; - #[cfg(unix)] let root = tempfile::tempdir()?; + let cwd = if use_executor { + std::env::current_dir()? + } else { + root.path().to_path_buf() + }; #[cfg(unix)] - let (server, cwd) = if use_executor { - (server, cwd) + let server = if use_executor { + server } else { let wrapper = root.path().join("server"); let script = "#!/bin/sh\nprintf '%s' \"$0\" > argv0\nexec \"$MCP_SERVER\" \"$@\"\n"; codex_utils_cargo_bin::write_executable(&wrapper, script)?; env.insert(OsString::from("MCP_SERVER"), server.into_os_string()); - ( - std::path::PathBuf::from("./server"), - root.path().to_path_buf(), - ) + std::path::PathBuf::from("./server") }; let client = RmcpClient::new_stdio_client_with_protocol_mode( server.into(), @@ -86,7 +86,11 @@ async fn exercise_stdio_server( Implementation::new("codex-mcp-client", "0.0.0-test"), ) .with_protocol_version(ProtocolVersion::V_2025_06_18), - Some(Duration::from_secs(5)), + Some(Duration::from_secs(if server_mode.ends_with("-late") { + 30 + } else { + 5 + })), Box::new(move |_request_id, request| { let observed_elicitations = Arc::clone(&observed_elicitations); async move { @@ -240,6 +244,29 @@ async fn rollout_flag_alone_preserves_legacy_stdio_and_sep1034_defaults() -> any .await } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unmarked_stdio_helpers_recover_from_failed_or_late_discovery() -> anyhow::Result<()> { + for mode in [ + "legacy", + "legacy-init-first", + "legacy-reject-late-initialize", + "legacy-reject-and-exit", + "legacy-late", + "legacy-silent", + "legacy-versions", + "modern-late", + ] { + exercise_stdio_server( + mode, + McpProtocolMode::DiscoverLocalStdio, + /*opt_in*/ false, + /*use_executor*/ false, + ) + .await?; + } + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn modern_stdio_safely_falls_back_to_legacy_elicitation() -> anyhow::Result<()> { exercise_stdio_server( diff --git a/codex-rs/utils/pty/src/process_group.rs b/codex-rs/utils/pty/src/process_group.rs index f403b7e647..f9e3f4381c 100644 --- a/codex-rs/utils/pty/src/process_group.rs +++ b/codex-rs/utils/pty/src/process_group.rs @@ -288,6 +288,33 @@ pub fn kill_process_group(_process_group_id: u32) -> io::Result<()> { Ok(()) } +/// Wait up to `grace` for a signaled group's members to release their resources. +/// After that, allow the caller to proceed. An unreaped zombie or a recycled group ID +/// can remain visible after the original processes have released their resources. +#[cfg(unix)] +pub fn wait_for_process_group_exit(process_group_id: u32, grace: std::time::Duration) { + let deadline = std::time::Instant::now() + grace; + loop { + #[cfg(target_os = "macos")] + let exists = signal_process_group_with_member_fallback( + process_group_id, + /*signal*/ 0, + signal_process_group_id, + signal_process_id, + ); + #[cfg(not(target_os = "macos"))] + let exists = signal_process_group_id(process_group_id as libc::pid_t, /*signal*/ 0); + if matches!(exists, Ok(false)) { + break; + } + let remaining = deadline.saturating_duration_since(std::time::Instant::now()); + if remaining.is_zero() { + break; + } + std::thread::sleep(remaining.min(std::time::Duration::from_millis(10))); + } +} + #[cfg(unix)] /// Kill the process group for a tokio child (best-effort). pub fn kill_child_process_group(child: &mut Child) -> io::Result<()> { diff --git a/codex-rs/utils/pty/src/win/job.rs b/codex-rs/utils/pty/src/win/job.rs index f963c6f55f..c6b803651d 100644 --- a/codex-rs/utils/pty/src/win/job.rs +++ b/codex-rs/utils/pty/src/win/job.rs @@ -1,29 +1,44 @@ use filedescriptor::OwnedHandle; +use std::collections::HashMap; use std::io; use std::os::windows::io::AsRawHandle; use std::os::windows::io::FromRawHandle; +use std::os::windows::io::OwnedHandle as ProcessHandle; use std::os::windows::io::RawHandle; use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; use tokio::process::Child; use tokio::process::Command; use winapi::shared::ntdef::NT_SUCCESS; use winapi::shared::ntdef::NTSTATUS; +use winapi::shared::winerror::ERROR_INVALID_PARAMETER; +use winapi::shared::winerror::ERROR_MORE_DATA; +use winapi::shared::winerror::WAIT_TIMEOUT; +use winapi::um::jobapi::IsProcessInJob; use winapi::um::jobapi2::AssignProcessToJobObject; use winapi::um::jobapi2::CreateJobObjectW; +use winapi::um::jobapi2::QueryInformationJobObject; use winapi::um::jobapi2::SetInformationJobObject; use winapi::um::jobapi2::TerminateJobObject; use winapi::um::processthreadsapi::OpenProcess; use winapi::um::processthreadsapi::TerminateProcess; +use winapi::um::synchapi::WaitForSingleObject; use winapi::um::winbase::CREATE_NO_WINDOW; use winapi::um::winbase::CREATE_SUSPENDED; +use winapi::um::winbase::WAIT_OBJECT_0; use winapi::um::winnt::HANDLE; use winapi::um::winnt::JOB_OBJECT_LIMIT_BREAKAWAY_OK; use winapi::um::winnt::JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; +use winapi::um::winnt::JOBOBJECT_BASIC_PROCESS_ID_LIST; use winapi::um::winnt::JOBOBJECT_EXTENDED_LIMIT_INFORMATION; +use winapi::um::winnt::JobObjectBasicProcessIdList; use winapi::um::winnt::JobObjectExtendedLimitInformation; +use winapi::um::winnt::PROCESS_QUERY_LIMITED_INFORMATION; use winapi::um::winnt::PROCESS_SET_QUOTA; use winapi::um::winnt::PROCESS_SUSPEND_RESUME; use winapi::um::winnt::PROCESS_TERMINATE; +use winapi::um::winnt::SYNCHRONIZE; #[link(name = "ntdll")] unsafe extern "system" { @@ -39,6 +54,19 @@ pub struct JobObject { preserve_descendants: Mutex, } +// A failed restart must still stop any remaining members. +struct TerminateOnDrop<'a>(Option<&'a JobObject>); + +impl Drop for TerminateOnDrop<'_> { + fn drop(&mut self) { + if let Some(job) = self.0 + && let Err(error) = job.terminate() + { + log::warn!("Failed to terminate Windows job after an interrupted restart: {error}"); + } + } +} + impl JobObject { /// Creates a Job Object configured to terminate all members when its last handle closes. pub fn create() -> io::Result { @@ -67,19 +95,23 @@ impl JobObject { } /// Captures an owned process handle before its numeric identifier can be reused. - pub fn open_process_handle(process_id: u32) -> io::Result { + pub fn open_process_handle(process_id: u32) -> io::Result { + Self::open_process(process_id, PROCESS_TERMINATE) + } + + fn open_process(process_id: u32, access: u32) -> io::Result { let handle = unsafe { - OpenProcess(PROCESS_TERMINATE, /*bInheritHandle*/ 0, process_id) + OpenProcess(access, /*bInheritHandle*/ 0, process_id) }; if handle.is_null() { return Err(io::Error::last_os_error()); } - Ok(unsafe { std::os::windows::io::OwnedHandle::from_raw_handle(handle.cast()) }) + Ok(unsafe { ProcessHandle::from_raw_handle(handle.cast()) }) } /// Terminates the exact process identified by a previously captured handle. - pub fn terminate_process_handle(handle: &std::os::windows::io::OwnedHandle) -> io::Result<()> { + pub fn terminate_process_handle(handle: &ProcessHandle) -> io::Result<()> { let terminated = unsafe { TerminateProcess(handle.as_raw_handle().cast(), /*uExitCode*/ 1) }; @@ -132,22 +164,15 @@ impl JobObject { .kill_on_drop(true); } - /// Assigns and resumes a suspended child, returning whether assignment succeeded. + /// Assigns and resumes a suspended child, returning its handle if containment succeeded. /// /// Nested jobs can reject assignment. Such a child is resumed without /// containment so callers can preserve their existing compatibility fallback. - pub fn assign_and_resume_process(&self, process_id: u32) -> io::Result { - let process = unsafe { - OpenProcess( - PROCESS_SET_QUOTA | PROCESS_TERMINATE | PROCESS_SUSPEND_RESUME, - /*bInheritHandle*/ 0, - process_id, - ) - }; - if process.is_null() { - return Err(io::Error::last_os_error()); - } - let process = unsafe { OwnedHandle::from_raw_handle(process.cast()) }; + pub fn assign_and_resume_process(&self, process_id: u32) -> io::Result> { + let process = Self::open_process( + process_id, + PROCESS_SET_QUOTA | PROCESS_TERMINATE | PROCESS_SUSPEND_RESUME | SYNCHRONIZE, + )?; let assignment = self.assign_process(process.as_raw_handle()); let status = unsafe { NtResumeProcess(process.as_raw_handle().cast()) }; @@ -161,12 +186,12 @@ impl JobObject { } match assignment { - Ok(()) => Ok(true), + Ok(()) => Ok(Some(process)), Err(error) => { log::warn!( "Windows process job assignment unavailable for pid {process_id}: {error}" ); - Ok(false) + Ok(None) } } } @@ -251,6 +276,113 @@ impl JobObject { Ok(()) } } + + /// Stops a root from [`Self::assign_and_resume_process`] and the remaining members of its job. + pub fn terminate_and_wait(&self, root: &ProcessHandle, grace: Duration) -> io::Result<()> { + let mut cleanup = TerminateOnDrop(Some(self)); + let deadline = Instant::now() + grace; + Self::request_stop(root)?; + Self::wait_for_exit(root, deadline)?; + // TerminateJobObject can empty the job list before members release their handles. + // Stop by process handle, then rescan for children born during the previous pass. + let mut list = vec![0usize; 16]; + let mut stopped = HashMap::new(); + loop { + if Instant::now() >= deadline { + return Err(io::ErrorKind::TimedOut.into()); + } + let pids = self.member_ids(&mut list)?; + if pids.iter().all(|pid| stopped.contains_key(pid)) { + self.terminate()?; + cleanup.0 = None; + return Ok(()); + } + for &pid in pids { + if stopped.contains_key(&pid) { + continue; + } + let member = match Self::open_process( + pid as u32, + PROCESS_TERMINATE | PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE, + ) { + Ok(member) => member, + Err(error) if error.raw_os_error() == Some(ERROR_INVALID_PARAMETER as i32) => { + continue; + } + Err(error) => return Err(error), + }; + // A listed PID can be reused before OpenProcess. Only signal confirmed job members. + let mut in_job = 0; + if unsafe { + IsProcessInJob( + member.as_raw_handle().cast(), + self.handle.as_raw_handle().cast(), + &mut in_job, + ) + } == 0 + { + return Err(io::Error::last_os_error()); + } + if in_job != 0 { + Self::request_stop(&member)?; + stopped.insert(pid, member); + } + } + for member in stopped.values() { + Self::wait_for_exit(member, deadline)?; + } + } + } + + fn member_ids<'a>(&self, list: &'a mut Vec) -> io::Result<&'a [usize]> { + const HEADER_WORDS: usize = + std::mem::offset_of!(JOBOBJECT_BASIC_PROCESS_ID_LIST, ProcessIdList) + / std::mem::size_of::(); + loop { + let listed = unsafe { + QueryInformationJobObject( + self.handle.as_raw_handle().cast(), + JobObjectBasicProcessIdList, + list.as_mut_ptr().cast(), + u32::try_from(std::mem::size_of_val(list.as_slice())) + .map_err(io::Error::other)?, + std::ptr::null_mut(), + ) + }; + if listed != 0 { + let header = unsafe { &*list.as_ptr().cast::() }; + if header.NumberOfProcessIdsInList == header.NumberOfAssignedProcesses { + return Ok(&list[HEADER_WORDS..][..header.NumberOfProcessIdsInList as usize]); + } + } else { + let error = io::Error::last_os_error(); + if error.raw_os_error() != Some(ERROR_MORE_DATA as i32) { + return Err(error); + } + } + list.resize(list.len() * 2, 0); + } + } + + fn request_stop(process: &ProcessHandle) -> io::Result<()> { + // Windows also returns access denied when the process has already exited. Its handle + // must still signal before a replacement can start. + match Self::terminate_process_handle(process) { + Err(error) if error.kind() == io::ErrorKind::PermissionDenied => Ok(()), + result => result, + } + } + + fn wait_for_exit(process: &ProcessHandle, deadline: Instant) -> io::Result<()> { + let remaining = deadline.saturating_duration_since(Instant::now()); + // The largest DWORD denotes INFINITE, so cap this bounded wait below it. + let millis = remaining.as_millis().min(u128::from(u32::MAX - 1)) as u32; + match unsafe { WaitForSingleObject(process.as_raw_handle().cast(), millis) } { + WAIT_OBJECT_0 => Ok(()), + WAIT_TIMEOUT => Err(io::ErrorKind::TimedOut.into()), + _ => Err(io::Error::last_os_error()), + } + } } impl AsRawHandle for JobObject { diff --git a/codex-rs/utils/pty/src/windows_tests.rs b/codex-rs/utils/pty/src/windows_tests.rs index 6d6c389b14..1a44d06f3f 100644 --- a/codex-rs/utils/pty/src/windows_tests.rs +++ b/codex-rs/utils/pty/src/windows_tests.rs @@ -8,6 +8,7 @@ use crate::spawn_pty_process; use pretty_assertions::assert_eq; use std::collections::HashMap; use std::os::windows::io::AsRawHandle; +use std::os::windows::io::BorrowedHandle; use std::os::windows::io::FromRawHandle; use std::os::windows::io::OwnedHandle; use std::os::windows::process::CommandExt; @@ -20,7 +21,10 @@ use tokio::io::BufReader; use tokio::process::Command; use winapi::um::jobapi::IsProcessInJob; use winapi::um::processthreadsapi::OpenProcess; +use winapi::um::synchapi::WaitForSingleObject; +use winapi::um::winbase::WAIT_OBJECT_0; use winapi::um::winnt::PROCESS_QUERY_LIMITED_INFORMATION; +use winapi::um::winnt::SYNCHRONIZE; const READY_MARKER: &str = "__CODEX_CHILD_READY__"; const VALUE_MARKER: &str = "__CODEX_CHILD_VALUE__"; @@ -67,7 +71,10 @@ async fn piped_child_has_no_console_with_or_without_job_containment() -> anyhow: let pid = child .id() .ok_or_else(|| anyhow::anyhow!("missing child pid"))?; - anyhow::ensure!(job.assign_and_resume_process(pid)?, "job assignment failed"); + anyhow::ensure!( + job.assign_and_resume_process(pid)?.is_some(), + "job assignment failed" + ); } child .stdin @@ -289,7 +296,13 @@ async fn contained_spawn_owns_immediate_descendant() -> anyhow::Result<()> { tokio::time::timeout(Duration::from_secs(10), stdout.read_line(&mut child_pid)).await??; let child_pid: u32 = child_pid.trim().parse()?; - let process = unsafe { OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, child_pid) }; + let process = unsafe { + OpenProcess( + PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE, + /*bInheritHandle*/ 0, + child_pid, + ) + }; anyhow::ensure!(!process.is_null(), "failed to open immediate child process"); let process = unsafe { OwnedHandle::from_raw_handle(process.cast()) }; let mut in_job = 0; @@ -303,7 +316,15 @@ async fn contained_spawn_owns_immediate_descendant() -> anyhow::Result<()> { anyhow::ensure!(checked != 0, "failed to inspect child Job Object"); anyhow::ensure!(in_job != 0, "immediate child escaped its Job Object"); - job.terminate()?; + let root_handle = root + .raw_handle() + .ok_or_else(|| anyhow::anyhow!("missing contained root handle"))?; + let root_handle = unsafe { BorrowedHandle::borrow_raw(root_handle) }.try_clone_to_owned()?; + job.terminate_and_wait(&root_handle, Duration::from_secs(10))?; + let child_status = unsafe { + WaitForSingleObject(process.as_raw_handle().cast(), /*dwMilliseconds*/ 0) + }; + assert_eq!(child_status, WAIT_OBJECT_0); tokio::time::timeout(Duration::from_secs(10), root.wait()).await??; Ok(()) } @@ -346,7 +367,9 @@ async fn rejected_job_assignment_resumes_existing_job_member() -> anyhow::Result .ok_or_else(|| anyhow::anyhow!("missing suspended process id"))?; assert!( - !rejected_job.assign_and_resume_process(process_id)?, + rejected_job + .assign_and_resume_process(process_id)? + .is_none(), "unrelated nested job unexpectedly accepted the process" ); let stdout = root