mirror of
https://github.com/openai/codex.git
synced 2026-10-27 18:48:37 +00:00
Add opt-in protocol discovery for local stdio MCP servers (#57210)
## Why Stdio servers without `CODEX_MCP_PROTOCOL_VERSION` use the legacy handshake, even when they support MCP `2026-07-28`. ## What changed - Add the disabled-by-default `mcp_connection_sharing` feature to enable discovery for eligible local stdio servers. - Recover from failed or late discovery with a fresh handshake when needed, waiting for the previous helper to release resources before restarting. - Retain a working legacy connection when an upgrade attempt fails for the same account, except on authentication errors. - Keep tool catalogs separate by protocol and invalidate legacy catalogs when discovery negotiates the modern protocol. ## Testing Add regression coverage for failed and late discovery, legacy fallback, protocol-specific catalog reuse, message limits, and Windows descendant cleanup. GitOrigin-RevId: 36c675ef3b81d32fccf06ccbddda2312cfe602c9
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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<McpServerConnectionIdentity>,
|
||||
client: AsyncManagedClient,
|
||||
legacy_fallback: Option<Arc<OnceLock<Arc<Self>>>>,
|
||||
// Startup-only budget; changing it must not replace a ready connection.
|
||||
startup_timeout: Duration,
|
||||
startup_trigger: Option<watch::Sender<bool>>,
|
||||
@@ -101,6 +105,54 @@ pub(crate) struct McpServerConnection {
|
||||
}
|
||||
|
||||
impl McpServerConnection {
|
||||
fn retained_legacy(&self) -> Option<&Arc<Self>> {
|
||||
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<Self>,
|
||||
replacement: &mut AsyncManagedClient,
|
||||
) -> Arc<OnceLock<Arc<Self>>> {
|
||||
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<McpServerConnection>,
|
||||
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<Arc<crate::McpConfig>>,
|
||||
@@ -290,11 +347,13 @@ impl McpConnectionSet {
|
||||
pub async fn new(
|
||||
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
|
||||
existing_ema: Option<&Self>,
|
||||
previous: Option<&Self>,
|
||||
previous: Option<PreviousMcpConnections<'_>>,
|
||||
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::<Vec<_>>();
|
||||
// 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()
|
||||
|
||||
@@ -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(),
|
||||
});
|
||||
|
||||
|
||||
@@ -190,6 +190,8 @@ pub struct McpConfig {
|
||||
pub non_prefixed_mcp_tool_servers: Vec<String>,
|
||||
/// 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.
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -166,6 +166,7 @@ impl ManagedClient {
|
||||
|
||||
pub(crate) type ManagedClientFuture =
|
||||
Shared<BoxFuture<'static, Result<ManagedClient, StartupOutcomeError>>>;
|
||||
type StartupShutdown = Shared<BoxFuture<'static, Result<(), String>>>;
|
||||
|
||||
#[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<StartupShutdown>,
|
||||
},
|
||||
}
|
||||
|
||||
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<anyhow::Error> for StartupOutcomeError {
|
||||
is_ema_authentication_required: is_authentication_required
|
||||
&& (error.downcast_ref::<EmaTokenSourceError>().is_some()
|
||||
|| error.downcast_ref::<EmaAuthFailure>().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",
|
||||
|
||||
@@ -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<CodexAuth>,
|
||||
auth_token: Option<String>,
|
||||
auth_manager: Option<Arc<AuthManager>>,
|
||||
auth_owner: Option<AuthOwner>,
|
||||
auth_generation: Arc<()>,
|
||||
plugins_available: bool,
|
||||
ready_selected_capability_roots: Vec<SelectedCapabilityRoot>,
|
||||
@@ -139,6 +141,23 @@ struct PublishedMcpRuntime {
|
||||
cached_binding: Mutex<Option<CachedMcpBinding>>,
|
||||
}
|
||||
|
||||
/// 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<Self> {
|
||||
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<Vec<ToolInfo>> {
|
||||
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(),
|
||||
|
||||
@@ -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<McpProtocolMode> {
|
||||
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
|
||||
}
|
||||
|
||||
@@ -184,6 +184,7 @@ impl McpToolCatalogCacheContext {
|
||||
cached_deadline.deadline
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
pub(crate) fn current_tools(&self) -> Option<Vec<ToolInfo>> {
|
||||
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<Arc<serde_json::Value>>,
|
||||
) {
|
||||
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<MutexGuard<'_, ToolCatalogCacheState>> {
|
||||
let mut state = lock_unpoisoned(&self.entry.state);
|
||||
if state.disabled_by_server || generation <= state.last_accepted_generation {
|
||||
return None;
|
||||
}
|
||||
state.last_accepted_generation = generation;
|
||||
state.optional_startup_deadline = None;
|
||||
Some(state)
|
||||
}
|
||||
}
|
||||
|
||||
fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
|
||||
@@ -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::<BTreeMap<_, _>>(),
|
||||
))
|
||||
|
||||
@@ -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"
|
||||
},
|
||||
|
||||
@@ -124,6 +124,7 @@ impl ManagedFeatures {
|
||||
let mut values = self.get().clone();
|
||||
for feature in [
|
||||
Feature::EnableMcpApps,
|
||||
Feature::McpConnectionSharing,
|
||||
Feature::SecretAuthStorage,
|
||||
Feature::McpOAuthRefreshCoordination,
|
||||
] {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
] {
|
||||
|
||||
@@ -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")??;
|
||||
|
||||
|
||||
@@ -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?;
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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<dyn AsyncRead + Send + Unpin>, 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<RxJsonRpcMessage<RoleClient>> {
|
||||
loop {
|
||||
let bytes = match self.stdout.fill_buf().await {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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>,
|
||||
child: Option<Box<Child>>,
|
||||
transport: StdioTransport,
|
||||
discovery: Option<Discovery>,
|
||||
}
|
||||
|
||||
pub(super) enum DiscoveryOutcome {
|
||||
Reply,
|
||||
RetryLegacy,
|
||||
Late {
|
||||
reply: Box<RxJsonRpcMessage<RoleClient>>,
|
||||
after_reoffer: bool,
|
||||
},
|
||||
}
|
||||
|
||||
struct Discovery {
|
||||
finished: Option<oneshot::Sender<DiscoveryOutcome>>,
|
||||
outcome: Option<DiscoveryOutcome>,
|
||||
retry_at: Option<tokio::time::Instant>,
|
||||
probe_id: Option<RequestId>,
|
||||
initialize_id: Option<RequestId>,
|
||||
}
|
||||
|
||||
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<RxJsonRpcMessage<RoleClient>>) -> 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::<i64>().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<RoleClient, ChildStdout, ChildStdin>),
|
||||
Legacy(AsyncRwTransport<RoleClient, Box<dyn AsyncRead + Unpin + Send>, 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<ChildStderr>)> {
|
||||
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<u32> {
|
||||
self.child.id()
|
||||
pub(super) fn watch_discovery(
|
||||
&mut self,
|
||||
timeout: Option<Duration>,
|
||||
) -> oneshot::Receiver<DiscoveryOutcome> {
|
||||
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<Box<Child>> {
|
||||
self.child.take()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -76,27 +210,87 @@ impl Transport<RoleClient> for LocalStdioTransport {
|
||||
&mut self,
|
||||
item: TxJsonRpcMessage<RoleClient>,
|
||||
) -> impl Future<Output = io::Result<()>> + 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<Output = Option<RxJsonRpcMessage<RoleClient>>> + Send {
|
||||
match &mut self.transport {
|
||||
StdioTransport::Legacy(transport) => transport.receive().boxed(),
|
||||
StdioTransport::V20260728(transport) => transport.receive().boxed(),
|
||||
async fn receive(&mut self) -> Option<RxJsonRpcMessage<RoleClient>> {
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<Self> {
|
||||
/// 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<Self> {
|
||||
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!(
|
||||
|
||||
@@ -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<StdioProcessSlot>,
|
||||
state: Mutex<ClientState>,
|
||||
stdio_process: Option<StdioServerProcessHandle>,
|
||||
transport_recipe: TransportRecipe,
|
||||
protocol_mode: McpProtocolMode,
|
||||
requires_read_only_tools: bool,
|
||||
@@ -516,7 +520,7 @@ impl RmcpClient {
|
||||
) -> io::Result<Self> {
|
||||
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<StdioServerTransport> {
|
||||
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<T, F, Fut>(
|
||||
&self,
|
||||
label: &str,
|
||||
|
||||
@@ -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<DiscoveryOutcome>,
|
||||
) -> anyhow::Result<Option<McpProtocolMode>> {
|
||||
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<McpProtocolMode> {
|
||||
if error.code != ErrorCode::UNSUPPORTED_PROTOCOL_VERSION {
|
||||
return None;
|
||||
}
|
||||
let supported = error.data.as_ref()?.get("supported")?;
|
||||
let versions = serde_json::from_value::<Vec<ProtocolVersion>>(supported.clone()).ok()?;
|
||||
restart_for_versions(&versions)
|
||||
}
|
||||
|
||||
fn restart_for_versions(supported: &[ProtocolVersion]) -> Option<McpProtocolMode> {
|
||||
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<DiscoveryOutcome>,
|
||||
) -> 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::<io::Error>().is_some_and(|e| {
|
||||
matches!(
|
||||
e.kind(),
|
||||
io::ErrorKind::BrokenPipe | io::ErrorKind::ConnectionReset
|
||||
)
|
||||
})
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
_ => matches!(discovery, Some(DiscoveryOutcome::RetryLegacy)),
|
||||
}
|
||||
}
|
||||
@@ -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<McpServerEnvVar>,
|
||||
cwd: Option<String>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
replacement_slot: Option<StdioProcessSlot>,
|
||||
}
|
||||
|
||||
/// A replacement holds the slot from OS spawn through registration. Closing it seals that handoff.
|
||||
pub(super) type StdioProcessSlot = Arc<AsyncMutex<StdioProcessRegistration>>;
|
||||
|
||||
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<BoxFuture<'static, Result<(), Arc<io::Error>>>>;
|
||||
|
||||
struct DiscoveryCleanup {
|
||||
process: Weak<StdioServerProcessHandle>,
|
||||
result: OnceLock<DiscoveryCleanupResult>,
|
||||
}
|
||||
|
||||
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<StdioServerProcessHandle>,
|
||||
discovery_cleanup: Option<Arc<DiscoveryCleanup>>,
|
||||
}
|
||||
|
||||
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<StdioServerProcessHandle>,
|
||||
discovery_cleanup: Option<Arc<DiscoveryCleanup>>,
|
||||
protocol_mode: McpProtocolMode,
|
||||
}
|
||||
|
||||
enum StdioServerTransportInner {
|
||||
@@ -137,7 +217,9 @@ impl Transport<RoleClient> 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<RoleClient> 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<Duration>,
|
||||
) -> Option<oneshot::Receiver<DiscoveryOutcome>> {
|
||||
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<StdioServerTransport>> {
|
||||
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<Option<Box<Child>>>,
|
||||
runtime: Handle,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[derive(Clone)]
|
||||
struct KillDiscoveryGroupOnDrop(Arc<Mutex<Option<Box<Child>>>>);
|
||||
|
||||
#[cfg(unix)]
|
||||
impl KillDiscoveryGroupOnDrop {
|
||||
fn new(child: Box<Child>) -> 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<OwnedHandle>,
|
||||
},
|
||||
Process(OwnedHandle),
|
||||
}
|
||||
|
||||
#[cfg(not(any(unix, windows)))]
|
||||
struct LocalProcessTerminator;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct StdioServerProcessHandle {
|
||||
inner: Arc<StdioServerProcessHandleInner>,
|
||||
}
|
||||
|
||||
struct StdioServerProcessHandleInner {
|
||||
program_name: String,
|
||||
kind: StdioServerProcessKind,
|
||||
terminated: AtomicBool,
|
||||
@@ -253,6 +435,12 @@ enum StdioServerProcessKind {
|
||||
Executor(Arc<dyn ExecProcess>),
|
||||
}
|
||||
|
||||
#[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<tokio::process::ChildStderr>,
|
||||
Option<u32>,
|
||||
)> {
|
||||
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<LocalProcessTerminator>,
|
||||
stderr_reader: Option<watch::Sender<()>>,
|
||||
) -> Self {
|
||||
Self {
|
||||
inner: Arc::new(StdioServerProcessHandleInner {
|
||||
program_name,
|
||||
kind: StdioServerProcessKind::Local(terminator),
|
||||
terminated: AtomicBool::new(false),
|
||||
stderr_reader,
|
||||
}),
|
||||
}
|
||||
) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
program_name,
|
||||
kind: StdioServerProcessKind::Local(terminator),
|
||||
terminated: AtomicBool::new(false),
|
||||
stderr_reader,
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) fn executor(program_name: String, process: Arc<dyn ExecProcess>) -> 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<dyn ExecProcess>) -> Arc<Self> {
|
||||
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,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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<()> {
|
||||
|
||||
@@ -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<bool>,
|
||||
}
|
||||
|
||||
// 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<Self> {
|
||||
@@ -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<std::os::windows::io::OwnedHandle> {
|
||||
pub fn open_process_handle(process_id: u32) -> io::Result<ProcessHandle> {
|
||||
Self::open_process(process_id, PROCESS_TERMINATE)
|
||||
}
|
||||
|
||||
fn open_process(process_id: u32, access: u32) -> io::Result<ProcessHandle> {
|
||||
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<bool> {
|
||||
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<Option<ProcessHandle>> {
|
||||
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<usize>) -> io::Result<&'a [usize]> {
|
||||
const HEADER_WORDS: usize =
|
||||
std::mem::offset_of!(JOBOBJECT_BASIC_PROCESS_ID_LIST, ProcessIdList)
|
||||
/ std::mem::size_of::<usize>();
|
||||
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::<JOBOBJECT_BASIC_PROCESS_ID_LIST>() };
|
||||
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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user