mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +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:
@@ -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<_, _>>(),
|
||||
))
|
||||
|
||||
Reference in New Issue
Block a user