openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-12 21:20:34 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
0702ebed4d6878c40bb90e01ef33039b21c6e1fc
codex/codex-rs/codex-mcp/src/connection_manager.rs
T
sayan-oai 561f0f5cf0 Keep same-named MCP servers separate across environments (#60905) 
## Why

MCP servers discovered on different executors can share a name. Resolving them by name alone exposes only one server and prevents tools and resources from being addressed independently.

## What changed

- Add the disabled-by-default `multi_environment_mcp` feature to expose environment-scoped servers, including multiple installations of the same plugin.
- Carry exact `McpServerKey` identities through catalogs, connections, tool calls, resources, and approval tracking. Give environment-scoped tools stable namespaces derived from their environment.
- Include `environment_id` in scoped resource listings and accept it in resource requests. Require it for environment-owned servers when multiple environments are selected.
- Preserve server ownership through hooks and use the owning environment's permissions for scoped elicitation review.

## Testing

Add coverage for same-named servers across environments, independent calls and approvals, resource selection, plugin discovery, and rejection of hook changes to the tool environment.

GitOrigin-RevId: 5593372b06a15c16aa3220faa3b19f5d3b20e67e
2026-11-11 23:36:19 +00:00

1602 lines
68 KiB
Rust
Raw Blame History

//! Aggregates MCP server connections for Codex.
//!
//! [`McpConnectionSet`] is the private connection set behind
//! [`crate::McpRuntime`] and [`crate::McpBinding`]. It coordinates startup status
//! events, keeps server metadata, and aggregates tools and resources across
//! running RMCP clients.
#[path = "connection_manager/catalog_telemetry.rs"]
mod catalog_telemetry;
#[path = "connection_manager/network_policy.rs"]
mod network_policy;
#[path = "connection_manager/required.rs"]
mod required;
pub use network_policy::NetworkPolicyMetadata;
pub use network_policy::NetworkPolicyMetadataProvider;
#[path = "connection_manager/resources.rs"]
mod resources;
#[path = "connection_manager/startup.rs"]
mod startup;
#[path = "connection_manager/status.rs"]
mod status;
#[path = "connection_manager/tool_catalog.rs"]
mod tool_catalog;
use startup::chatgpt_auth_provider_for_server;
use startup::emit_update;
use startup::mcp_init_error_display;
use startup::mcp_startup_failure_reason;
use startup::should_share_codex_apps_tools_cache;
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::McpServerKey;
use crate::binding::call_tool_result_from_rmcp;
use crate::catalog::McpServerSource;
use crate::connection_pool::ConnectionKey;
use crate::connection_pool::ConnectionSharing;
use crate::elicitation::ElicitationRequestManager;
use crate::elicitation::ElicitationRequestRouter;
use crate::event_stream::EventStreamConnectionSettings;
use crate::mcp::CODEX_APPS_MCP_SERVER_NAME;
use crate::mcp::ToolPluginContext;
use crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS;
use crate::pagination::MAX_MCP_CATALOG_ITEMS;
use crate::rmcp_client::AsyncManagedClient;
use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT;
use crate::rmcp_client::DEFAULT_TOOL_TIMEOUT;
use crate::rmcp_client::ManagedClient;
use crate::rmcp_client::StartupOutcomeError;
use crate::rmcp_client::prepare_codex_apps_tools_for_model;
use crate::rmcp_client::prepare_regular_mcp_tools_for_model;
use crate::runtime::McpPublicationGate;
use crate::runtime::McpRuntimeInput;
use crate::runtime::McpStartupPolicy;
use crate::server::McpServerConnectionIdentity;
use crate::server::McpServerMetadata;
use crate::tools::ToolFilter;
use crate::tools::ToolInfo;
use crate::tools::filter_tools;
use crate::trusted_access::ENTITLEMENT_CONTEXT_KEY;
use crate::trusted_access::TrustedAccessContext;
use anyhow::Context;
use anyhow::Result;
use anyhow::anyhow;
use anyhow::bail;
use codex_config::McpServerTransportConfig;
use codex_config::McpStartupReadiness;
use codex_diagnostics::Gauge;
use codex_diagnostics::GaugeGuard;
use codex_otel::auth_storage::AuthStorageOriginator;
use codex_protocol::mcp::CallToolResult;
use codex_protocol::mcp::ClientMcpExtensions;
use codex_protocol::mcp::ENTERPRISE_MANAGED_AUTHORIZATION_EXTENSION_ID;
use codex_protocol::mcp::McpServerInfo;
use codex_protocol::protocol::Event;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::McpStartupCompleteEvent;
use codex_protocol::protocol::McpStartupFailure;
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;
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>>,
// Only the initial host-owned Apps attempt publishes here. A retry uses normal readiness.
initialized_apps:
Option<watch::Receiver<Option<std::sync::Weak<codex_rmcp_client::RmcpClient>>>>,
_diagnostics_guard: GaugeGuard,
}
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,
) -> Option<ManagedClient> {
let current = self.identity.as_ref()?;
if !current.has_same_connection_config(desired) {
return None;
}
if !self.client.startup_complete.load(Ordering::Acquire) {
return None;
}
let client = self.client.client().await.ok()?;
if client.client.is_closed().await {
return None;
}
if current == desired {
if matches!(desired.oauth_credentials(), Ok(None))
&& tokio::time::timeout(Duration::ZERO, client.client.managed_oauth_credentials())
.await
.is_ok_and(|credentials| matches!(credentials, Some(Some(_))))
{
return None;
}
return Some(client);
}
let Ok(desired_credentials) = desired.oauth_credentials() else {
return Some(client);
};
let reusable = match client.client.managed_oauth_credentials().await {
Some(live_credentials) => live_credentials.as_ref() == desired_credentials,
None => current
.oauth_credentials()
.is_ok_and(|startup_credentials| startup_credentials == desired_credentials),
};
if reusable { Some(client) } else { None }
}
pub(crate) async fn client(&self) -> Result<ManagedClient, StartupOutcomeError> {
if let Some(startup_trigger) = &self.startup_trigger {
startup_trigger.send_replace(true);
}
self.client.client().await
}
async fn client_for_call(&self) -> Result<ManagedClient, StartupOutcomeError> {
let result = self.client().await;
if matches!(result, Err(StartupOutcomeError::Failed { .. }))
&& !self.client.cancel_token.is_cancelled()
&& let Some(reconnect) = &self.client.startup_reconnect
{
if let Some(attempt) = reconnect.reconnect_in_background() {
return attempt.await;
}
// Recovery may have finished between reading the client and joining it.
return self.client().await;
}
result
}
async fn shutdown(&self) {
self.client.shutdown().await;
}
fn cancel_startup(&self) {
if !self.startup_is_dormant() && !self.client.startup_complete.load(Ordering::Acquire) {
self.client.cancel_token.cancel();
}
}
fn startup_is_dormant(&self) -> bool {
self.startup_trigger
.as_ref()
.is_some_and(|startup_trigger| !*startup_trigger.borrow())
}
}
impl Drop for McpServerConnection {
fn drop(&mut self) {
self.client.cancel_token.cancel();
}
}
#[derive(Clone)]
struct McpServerView {
connection: Arc<McpServerConnection>,
protocol_mode: McpProtocolMode,
startup_readiness: McpStartupReadiness,
metadata: McpServerMetadata,
tool_filter: ToolFilter,
tool_timeout: Option<Duration>,
catalog_item_limit: usize,
refresh_tool_list: bool,
refresh_timeout: Duration,
}
impl McpServerView {
fn allows_cached_startup(&self) -> bool {
self.startup_readiness == McpStartupReadiness::Catalog
|| self.connection.startup_is_dormant()
}
fn cached_startup_tools(&self, fallback: Option<Vec<ToolInfo>>) -> Option<Vec<ToolInfo>> {
self.connection
.client
.cached_tools_or(fallback)
.filter(|tools| self.accepts_cached_tools(tools))
}
fn accepts_cached_tools(&self, tools: &[ToolInfo]) -> bool {
self.connection.client.is_codex_apps_mcp_server
|| match self.startup_readiness {
McpStartupReadiness::Connection => !tools.is_empty(),
McpStartupReadiness::Catalog => tools.iter().any(|tool| {
self.tool_filter.allows(&tool.tool.name)
&& (self.connection.startup_is_dormant() || tool_is_model_visible(tool))
}),
}
}
async fn refresh_notified_catalog(&self, server_name: &str, client: &ManagedClient) {
if !self.refresh_tool_list || self.connection.client.is_codex_apps_mcp_server {
return;
}
let timeout = self.refresh_timeout;
let result = client
.tool_catalog
.refresh_after_notifications(&client.client.tool_list_revision(), timeout, || async {
let relist = crate::rmcp_client::list_tools_for_client_uncached(
server_name,
/*is_codex_apps_mcp_server*/ false,
"list_changed",
&client.client,
Some(timeout),
self.catalog_item_limit,
client.server_instructions.as_deref(),
);
client.with_active_shared_lease(relist.boxed()).await?
})
.await;
if result.is_err() {
tracing::warn!(
server_name,
"MCP tool catalog refresh failed; retaining the previous catalog"
);
}
}
async fn listed_catalog(
&self,
server_name: &str,
tool_plugin_context: &ToolPluginContext,
exposure: crate::McpServerExposure,
) -> Result<(Vec<ToolInfo>, Option<Arc<serde_json::Value>>), StartupOutcomeError> {
let (tools, meta) = if let Some(client) = self.connection.client.ready_client() {
self.refresh_notified_catalog(server_name, &client).await;
client.listed_catalog().await
} else {
self.connection.client.listed_catalog().await?
};
Ok((
self.prepare_tools_for_model(tools, tool_plugin_context, exposure),
meta,
))
}
/// Publishes raw client tools with this view's policy, owner, and plugin attribution.
/// Shared client catalogs stay unannotated so later views can reuse them safely.
fn prepare_tools_for_model(
&self,
tools: Vec<ToolInfo>,
tool_plugin_context: &ToolPluginContext,
exposure: crate::McpServerExposure,
) -> Vec<ToolInfo> {
let tools = filter_tools(tools, &self.tool_filter)
.into_iter()
.map(|mut tool| {
tool.server_key = self.metadata.server_key.clone();
if exposure.environment_id(tool.server_key()).is_some() {
let environment = &self.metadata.environment_id;
let owner = format!("Runs in environment {environment}.");
tool.namespace_description = Some(match tool.namespace_description {
Some(description) => format!("{owner}\n{description}"),
None => owner,
});
}
tool.supports_parallel_tool_calls = self.metadata.supports_parallel_tool_calls;
tool.server_origin = self
.metadata
.origin
.as_ref()
.map(|origin| origin.as_str().to_string());
tool
})
.collect();
if self.connection.client.is_codex_apps_mcp_server {
prepare_codex_apps_tools_for_model(tools, tool_plugin_context)
} else {
prepare_regular_mcp_tools_for_model(tools, tool_plugin_context)
}
}
}
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>>,
network_policy_metadata: Option<Arc<dyn NetworkPolicyMetadataProvider>>,
servers: HashMap<McpServerKey, McpServerView>,
pub(crate) event_stream_connection: Option<Arc<EventStreamConnectionSettings>>,
disabled_servers: Vec<String>,
required_servers: Vec<McpServerKey>,
optional_startup_deadline: OnceLock<tokio::time::Instant>,
tool_plugin_context: Arc<ToolPluginContext>,
prefix_mcp_tool_names: bool,
non_prefixed_mcp_tool_servers: Vec<String>,
elicitation_requests: ElicitationRequestManager,
pub(crate) trusted_access: Option<TrustedAccessContext>,
}
impl McpConnectionSet {
fn server(&self, key: &McpServerKey) -> Option<&McpServerView> {
self.servers.get(key)
}
pub(crate) fn resolve_server_name(&self, name: &str) -> Option<&McpServerKey> {
self.config
.as_ref()
.and_then(|config| config.mcp_server_catalog.resolve_server_name(name))
.or_else(|| {
self.servers
.get_key_value(&McpServerKey::unscoped(name))
.map(|(key, _)| key)
})
}
fn server_by_name(&self, name: &str) -> Option<&McpServerView> {
self.server(self.resolve_server_name(name)?)
}
/// Name-only status and resource APIs still describe the original first server.
fn legacy_servers(&self) -> impl Iterator<Item = (&str, &McpServerView)> {
self.servers
.iter()
.filter_map(|(key, view)| self.legacy_server_name(key).map(|name| (name, view)))
}
pub(crate) fn legacy_server_name<'a>(&'a self, key: &'a McpServerKey) -> Option<&'a str> {
match &self.config {
Some(config) => config.mcp_server_catalog.legacy_server_name(key),
None => Some(key.name()),
}
}
/// Creates an MCP connection manager. Threadless callers can pass no `tx_event`; startup
/// notifications are then skipped and interactive elicitations are declined.
pub async fn new(
ema_auth: Option<crate::ema_auth::EmaAuthContext>,
existing_ema: 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,
config,
network_policy_metadata,
plugins_available: _,
ready_selected_capability_roots: _,
mcp_servers,
submit_id,
tx_event,
startup_cancellation_token,
runtime_context,
codex_apps_tools_cache,
tool_catalog_cache,
codex_apps_tools_cache_key,
client_mcp_extensions,
auth,
auth_manager,
runtime_auth,
elicitation_reviewer,
elicitation_lifecycle,
} = input;
let store_mode = config.mcp_oauth_credentials_store_mode;
let keyring_backend_kind = config.auth_keyring_backend_kind;
let oauth_refresh_mode = config.oauth_refresh_mode;
let codex_home = config.codex_home.clone();
let prefix_mcp_tool_names = config.prefix_mcp_tool_names;
let non_prefixed_mcp_tool_servers = config.non_prefixed_mcp_tool_servers.clone();
let default_protocol_mode = config.protocol_mode;
let host_owned_apps_protocol_mode = config.host_owned_apps_protocol_mode;
let client_elicitation_capability = config.client_elicitation_capability.clone();
let tool_plugin_context = crate::mcp::tool_plugin_context(&config);
let auth = auth.as_ref();
let mut servers = HashMap::new();
let mut event_stream_connection = None;
let disabled_servers = mcp_servers
.iter()
.filter(|(_, server)| !server.enabled())
.filter_map(|(key, _)| {
config
.mcp_server_catalog
.legacy_server_name(key)
.map(str::to_owned)
})
.collect();
let mut required_servers = mcp_servers
.iter()
.filter(|(_, server)| server.enabled() && server.required())
.map(|(server_name, _)| server_name.clone())
.collect::<Vec<_>>();
required_servers.sort();
let mut reused_ready = Vec::new();
let mut eager_starts = Vec::new();
let mut eager_workers = Vec::new();
let mut deferred_workers = Vec::new();
// Retained EMA clients keep callbacks attached to this manager even when
// an explicit reconnect replaces the ordinary clients.
let elicitation_previous = previous
.or_else(|| {
existing_ema.filter(|previous| {
previous
.servers
.values()
.any(|view| view.connection.client.ema_authority.is_some())
})
})
.filter(|previous| {
!previous.servers.is_empty()
&& previous.elicitation_requests.update(
Arc::clone(&config),
elicitation_reviewer.clone(),
elicitation_lifecycle.clone(),
)
});
let reusable_previous = previous.filter(|_| elicitation_previous.is_some());
let elicitation_requests = if let Some(previous) = elicitation_previous {
previous.elicitation_requests.clone()
} else {
ElicitationRequestManager::new(
Arc::clone(&config),
elicitation_reviewer,
elicitation_lifecycle,
elicitation_router,
)
};
let tool_plugin_context = Arc::new(tool_plugin_context);
let startup_submit_id = submit_id;
let static_chatgpt_auth_provider = auth
.filter(|auth| auth.uses_codex_backend())
.map(codex_model_provider::auth_provider_from_auth);
let codex_apps_auth_provider = auth_manager.as_ref().and_then(|auth_manager| {
auth.filter(|auth| auth.uses_codex_backend()).map(|auth| {
codex_model_provider::auth_provider_from_auth_manager(
Arc::clone(auth_manager),
auth,
)
})
});
for (server_key, server) in mcp_servers
.into_iter()
.filter(|(_, server)| server.enabled())
{
let server = server.with_read_only_mcp_tools(config.requires_read_only_mcp_tools);
let registration = config.mcp_server_catalog.server(&server_key);
let server_name = registration
.map_or(server_key.name(), |server| server.name())
.to_owned();
let report_startup = config
.mcp_server_catalog
.resolve_server_name(&server_name)
.is_none_or(|key| key == &server_key);
let client_mcp_extensions = crate::client_capabilities::server_mcp_extensions(
&client_mcp_extensions,
&server_name,
registration,
network_policy_metadata.is_some(),
&config.chatgpt_base_url,
);
let is_host_owned_codex_apps = registration.is_some_and(|server| {
server
.source()
.is_host_owned_apps(&server_name, server.config())
});
let host_plugin_root = registration.and_then(|server| match server.source() {
McpServerSource::Plugin(plugin) => plugin.host_root(),
McpServerSource::SelectedPlugin(_)
| McpServerSource::Config
| McpServerSource::Compatibility { .. }
| McpServerSource::Extension { .. } => None,
});
let catalog_item_limit = if is_host_owned_codex_apps {
MAX_CODEX_APPS_TOOL_CATALOG_ITEMS
} else {
MAX_MCP_CATALOG_ITEMS
};
let metadata = McpServerMetadata::new(server_key.clone(), &server);
let configured_config = server.config().clone();
let protocol_mode = if matches!(
&configured_config.transport,
McpServerTransportConfig::StreamableHttp { .. }
) {
registration
.and_then(crate::ResolvedMcpServer::protocol_mode)
.unwrap_or(if is_host_owned_codex_apps {
host_owned_apps_protocol_mode
} else {
default_protocol_mode
})
} else {
default_protocol_mode
};
let uses_ema_auth =
config.xaa_enabled && matches!(configured_config.auth, McpServerAuth::EmaAuth);
let protocol_mode = if uses_ema_auth {
McpProtocolMode::V20260728
} else {
protocol_mode
};
let client_mcp_extensions = if uses_ema_auth {
ClientMcpExtensions::new(
client_mcp_extensions
.iter()
.map(|(id, settings)| (id.to_string(), settings.clone()))
.chain(std::iter::once((
ENTERPRISE_MANAGED_AUTHORIZATION_EXTENSION_ID.to_string(),
serde_json::json!({}),
))),
)
} else {
client_mcp_extensions
};
let configured_tool_filter = ToolFilter::from_config(&configured_config);
let startup_timeout = configured_config
.startup_timeout_sec
.unwrap_or(DEFAULT_STARTUP_TIMEOUT);
let configured_tool_timeout = Some(
configured_config
.tool_timeout_sec
.unwrap_or(DEFAULT_TOOL_TIMEOUT),
);
let resolved_environment =
runtime_context.resolve_server_environment(&server_name, &configured_config);
// For built-in Codex Apps, `CODEX_CONNECTORS_TOKEN` is a debug
// override: it supplies runtime auth but bypasses the shared tools
// cache.
let uses_env_bearer_token = match &configured_config.transport {
McpServerTransportConfig::StreamableHttp {
bearer_token_env_var,
..
} => bearer_token_env_var.is_some(),
McpServerTransportConfig::Stdio { .. } => false,
};
// Filtered catalogs must not read or populate an unrestricted shared cache.
let eligible_for_apps_cache = is_host_owned_codex_apps
&& !server.requires_read_only_mcp_tools()
&& should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token);
let shares_codex_apps_tools_cache =
eligible_for_apps_cache && codex_apps_tools_cache.shares_snapshot();
let openai_runtime = runtime_auth.as_ref().filter(|_| {
!uses_env_bearer_token
&& configured_config.is_local_environment()
&& matches!(configured_config.auth, codex_config::McpServerAuth::ChatGpt)
&& auth.is_some_and(codex_login::CodexAuth::uses_codex_backend)
&& matches!(&configured_config.transport,
McpServerTransportConfig::StreamableHttp { url, .. }
if codex_http_client::classify_openai_api_url(url)
!= codex_http_client::OpenAiDestination::Other)
});
let runtime_auth_mode = openai_runtime.map(|(_, mode)| *mode);
let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| {
// Only equivalent discovery inputs may share executable Apps tools.
let mut transport = configured_config.transport.clone();
if let McpServerTransportConfig::StreamableHttp {
http_headers: Some(headers),
..
} = &mut transport
{
// mcp_server_config_for_url in codex-rs/codex-mcp/src/mcp/mod.rs
// adds thread attribution that threadless discovery does not carry.
headers.retain(|name, _| !name.eq_ignore_ascii_case("originator"));
}
let mut scope = serde_json::json!([
transport,
&configured_config.auth,
runtime_auth_mode,
protocol_mode.preferred_protocol_version().as_str(),
catalog_item_limit,
&client_elicitation_capability,
client_mcp_extensions
.iter()
.collect::<std::collections::BTreeMap<_, _>>(),
]);
scope.sort_all_objects();
codex_apps_tools_cache
.context(codex_home.clone(), codex_apps_tools_cache_key.clone())
.with_live_scope(scope.to_string())
});
// The reserved Codex Apps registration follows the shared
// AuthManager across refreshes. In the hosted-plugin path, this
// is the ChatGPT /ps/mcp connection. User-configured MCP
// registrations keep their existing configured auth path.
let chatgpt_auth_provider = if server_name == CODEX_APPS_MCP_SERVER_NAME {
codex_apps_auth_provider
.clone()
.or_else(|| static_chatgpt_auth_provider.clone())
} else {
static_chatgpt_auth_provider.clone()
};
// If Codex Apps has an env bearer token, that is its auth path. Do
// not also attach the ambient CodexAuth provider.
let runtime_auth_provider =
if server_name == CODEX_APPS_MCP_SERVER_NAME && uses_env_bearer_token {
None
} else {
chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider).map(
|provider| match (openai_runtime, auth) {
(Some(runtime), Some(auth)) => {
crate::auth::McpAuthProvider::runtime(runtime.clone(), auth.clone())
}
_ => crate::auth::McpAuthProvider::Credentials(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(),
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
runtime_context: runtime_context.clone(),
resolved_environment: resolved_environment.clone(),
auth_provider: runtime_auth_provider.clone(),
auth_manager: auth_manager.clone(),
auth: auth.cloned(),
protocol_mode,
client_mcp_extensions: client_mcp_extensions.clone(),
}));
}
let connection_identity = McpServerConnectionIdentity::new(
&server_name,
&server,
host_plugin_root,
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
&resolved_environment,
&runtime_context,
runtime_auth_provider.as_ref(),
runtime_auth_mode,
auth,
eligible_for_apps_cache
.then(|| (codex_home.clone(), codex_apps_tools_cache_key.clone())),
client_elicitation_capability.clone(),
client_mcp_extensions.clone(),
previous
.and_then(|previous| previous.servers.get(&server_key))
.and_then(|view| view.connection.identity.as_ref()),
);
let expected_protocol_mode = match &configured_config.transport {
McpServerTransportConfig::StreamableHttp { .. } => Some(protocol_mode),
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;
context.before_stdio_discovery =
expected_protocol_mode == Some(McpProtocolMode::DiscoverLocalStdio);
Some(context)
});
let allow_deferred_startup = match startup_policy {
McpStartupPolicy::Eager => false,
McpStartupPolicy::DeferCachedStdio => configured_config.can_defer_stdio_startup(),
McpStartupPolicy::DeferCachedServers => true,
};
// Explicit reconnects have no previous set and must replace their clients eagerly.
let defer_startup = previous.is_some()
&& allow_deferred_startup
&& !tool_plugin_context.is_selected_plugin_server(&server_key)
&& 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();
let protocol_mode = expected_protocol_mode.unwrap_or(protocol_mode);
let connection_pool = runtime_context.connection_pool.as_ref().filter(|_| {
// Sharing a signed-in helper requires tracking account changes.
discover_local_stdio
&& protocol_mode != McpProtocolMode::Legacy
&& (auth.is_none() || auth_manager.is_some())
});
if matches!(configured_config.auth, McpServerAuth::EmaAuth) && existing_ema.is_some() {
if let Some(previous_view) = existing_ema
.filter(|_| elicitation_previous.is_some())
.and_then(|previous| previous.servers.get(&server_key))
{
let connection = &previous_view.connection;
let authority_is_current = || {
connection
.client
.ema_authority
.as_ref()
.is_some_and(|authority| authority.check_new_request().is_ok())
};
let same_identity = connection.identity.as_ref() == Some(&connection_identity)
&& previous_view.catalog_item_limit == catalog_item_limit
&& previous_view.protocol_mode == protocol_mode;
let pending = same_identity
&& !connection.client.startup_complete.load(Ordering::Acquire)
&& connection.startup_timeout == startup_timeout
&& !connection.startup_is_dormant()
&& !connection.client.cancel_token.is_cancelled();
let ready = same_identity
&& !pending
&& connection
.reusable_client(&connection_identity)
.await
.is_some_and(|client| {
expected_protocol_mode == Some(client.client.protocol_mode())
});
if authority_is_current() && (pending || ready) {
servers.insert(
server_key.clone(),
McpServerView {
connection: Arc::clone(connection),
protocol_mode,
startup_readiness: configured_config.startup_readiness,
refresh_tool_list: config.mcp_tool_list_changed,
refresh_timeout: startup_timeout,
metadata,
tool_filter: configured_tool_filter,
tool_timeout: configured_tool_timeout,
catalog_item_limit,
},
);
if pending {
let client = connection.client.clone();
eager_workers
.push(async move { (server_key, client.client().await) }.boxed());
} else {
reused_ready.push(server_key);
}
}
}
// 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_key))
&& 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);
// A failed pool acquisition can leave this chat on its private helper.
let retained_legacy = connection.retained_legacy().is_some();
let can_reuse = retained_legacy
|| connection.client.connection_slot.is_some() == connection_pool.is_some();
let reusable_pending_startup = connection.identity.as_ref()
== Some(&connection_identity)
&& !connection.client.startup_complete.load(Ordering::Acquire)
&& connection.startup_timeout == startup_timeout
&& !connection.startup_is_dormant()
&& !connection.client.cancel_token.is_cancelled()
&& previous_view.catalog_item_limit == catalog_item_limit
&& expected_protocol_mode.is_some()
&& previous_view.protocol_mode == protocol_mode;
let unchanged_auth_failure = if connection.identity.as_ref()
== Some(&connection_identity)
&& connection_identity.oauth_store_was_contended
&& previous_view.protocol_mode == protocol_mode
&& connection.client.startup_complete.load(Ordering::Acquire)
{
connection
.client()
.await
.err()
.filter(StartupOutcomeError::is_authentication_required)
} else {
None
};
let ready = connection
.reusable_client(&connection_identity)
.await
.filter(|_| previous_view.catalog_item_limit == catalog_item_limit);
if can_reuse
&& (reusable_pending_startup
|| unchanged_auth_failure.is_some()
|| ready.as_ref().is_some_and(|client| {
expected_protocol_mode.is_some_and(|expected| {
client.client.protocol_mode() == expected
|| (previous_view.protocol_mode == expected && retained_legacy)
})
}))
{
let pending_client =
reusable_pending_startup.then(|| connection.client.clone());
servers.insert(
server_key.clone(),
McpServerView {
connection,
protocol_mode,
startup_readiness: configured_config.startup_readiness,
refresh_tool_list: config.mcp_tool_list_changed,
refresh_timeout: startup_timeout,
metadata,
tool_filter: configured_tool_filter,
tool_timeout: configured_tool_timeout,
catalog_item_limit,
},
);
if let Some(error) = unchanged_auth_failure {
if report_startup {
eager_starts.push(server_name.clone());
}
let reason = connection_identity
.oauth_credentials()
.ok()
.flatten()
.map(|_| McpStartupFailureReason::ReauthenticationRequired);
let status = McpStartupStatus::Failed {
error: mcp_init_error_display(
&server_name,
Some(&configured_config),
&error,
reason,
),
reason,
};
let tx_event = tx_event.clone().filter(|_| report_startup);
let submit_id = startup_submit_id.clone();
eager_workers.push(
async move {
if let Some(tx_event) = tx_event.as_ref() {
let _ = emit_update(
submit_id.as_str(),
tx_event,
McpStartupUpdateEvent {
server: server_name.clone(),
status,
},
)
.await;
}
(server_key, Err(error))
}
.boxed(),
);
} else if let Some(client) = pending_client {
eager_workers
.push(async move { (server_key, client.client().await) }.boxed());
} else {
reused_ready.push(server_key);
}
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
&& matches!(client.connection_sharing, ConnectionSharing::Private)
})
{
legacy_to_upgrade = Some(connection);
}
}
let cancel_token = startup_cancellation_token.child_token();
let tool_catalog_cache_context = tool_catalog_cache_context.take();
let has_runtime_auth = runtime_auth_provider.is_some();
let connection_pool = connection_pool.map(|pool| {
(
pool.clone(),
ConnectionKey {
server_name: server_name.clone(),
identity: connection_identity.clone(),
account: codex_apps_tools_cache_key.clone(),
auth_state: auth_manager.as_ref().map(|manager| {
(
Arc::clone(manager),
*manager.auth_change_state_receiver().borrow(),
)
}),
startup_timeout,
catalog_item_limit,
},
)
});
let (publish_apps, initialized_apps) =
if is_host_owned_codex_apps && config.apps_early_resource_discovery {
let (tx, rx) = watch::channel(None);
(Some(tx), Some(rx))
} else {
(None, None)
};
let mut async_managed_client = AsyncManagedClient::new(
ema_auth.clone(),
server_key.clone(),
startup_submit_id.clone(),
server,
store_mode,
keyring_backend_kind,
oauth_refresh_mode,
cancel_token.clone(),
tx_event.clone(),
elicitation_requests.clone(),
codex_apps_tools_cache_context,
tool_catalog_cache_context,
runtime_context.clone(),
resolved_environment,
runtime_auth_provider,
client_elicitation_capability.clone(),
client_mcp_extensions.clone(),
auth_manager
.as_ref()
.filter(|_| {
matches!(
&configured_config.transport,
McpServerTransportConfig::Stdio { .. }
)
})
.map(|manager| manager.auth_change_state_receiver()),
protocol_mode,
defer_startup,
catalog_item_limit,
connection_pool,
publish_apps,
);
// 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_key)) {
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))
} else {
(None, None)
};
servers.insert(
server_key.clone(),
McpServerView {
connection: Arc::new(McpServerConnection {
identity: Some(connection_identity),
client: async_managed_client.clone(),
legacy_fallback,
startup_timeout,
startup_trigger,
initialized_apps,
_diagnostics_guard: LIVE_CONNECTIONS.track(),
}),
protocol_mode,
startup_readiness: configured_config.startup_readiness,
refresh_tool_list: config.mcp_tool_list_changed,
refresh_timeout: startup_timeout,
metadata,
tool_filter: configured_tool_filter,
tool_timeout: configured_tool_timeout,
catalog_item_limit,
},
);
let tx_event = tx_event.clone().filter(|_| report_startup);
let submit_id = startup_submit_id.clone();
if !defer_startup && report_startup {
eager_starts.push(server_name.clone());
}
let startup = async move {
if let Some(mut startup_receiver) = startup_receiver
&& tokio::select! {
started = startup_receiver.wait_for(|started| *started) => started.is_err(),
() = cancel_token.cancelled() => true,
}
{
return (server_key, Err(StartupOutcomeError::Cancelled));
}
if defer_startup && let Some(tx_event) = tx_event.as_ref() {
let _ = emit_update(
submit_id.as_str(),
tx_event,
McpStartupUpdateEvent {
server: server_name.clone(),
status: McpStartupStatus::Starting,
},
)
.await;
}
let mut outcome = async_managed_client.client().await;
if cancel_token.is_cancelled() {
outcome = Err(StartupOutcomeError::Cancelled);
}
if let Some(tx_event) = tx_event.as_ref() {
let auth_state = match &outcome {
Err(error)
if error.is_authentication_required()
&& !has_runtime_auth
&& !matches!(configured_config.auth, McpServerAuth::EmaAuth) =>
{
match &configured_config.transport {
McpServerTransportConfig::StreamableHttp {
url,
bearer_token_env_var,
http_headers,
env_http_headers,
..
} => {
match determine_streamable_http_auth_status_from_credentials(
configured_config
.oauth_credential_name(&server_name)
.as_ref(),
url,
bearer_token_env_var.as_deref(),
http_headers.clone(),
env_http_headers.clone(),
store_mode,
keyring_backend_kind,
) {
Ok(auth_state) => auth_state,
Err(error) => {
warn!(
"failed to read stored auth status for MCP server `{server_name}`: {error:?}"
);
None
}
}
}
McpServerTransportConfig::Stdio { .. } => None,
}
}
Ok(_) | Err(_) => None,
};
if cancel_token.is_cancelled() {
outcome = Err(StartupOutcomeError::Cancelled);
}
let status = match &outcome {
Ok(_) => McpStartupStatus::Ready,
Err(StartupOutcomeError::Cancelled) => McpStartupStatus::Cancelled,
Err(error) => {
let reason = mcp_startup_failure_reason(auth_state, error);
let error_str = mcp_init_error_display(
server_name.as_str(),
Some(&configured_config),
error,
reason,
);
McpStartupStatus::Failed {
error: error_str,
reason,
}
}
};
let _ = emit_update(
submit_id.as_str(),
tx_event,
McpStartupUpdateEvent {
server: server_name.clone(),
status,
},
)
.await;
}
if cancel_token.is_cancelled() {
outcome = Err(StartupOutcomeError::Cancelled);
}
if matches!(&outcome, Err(StartupOutcomeError::Failed { .. })) {
async_managed_client.reconnect_failed_startup().await;
}
(server_key, outcome)
};
let startup = AuthStorageOriginator::current().scope(startup);
if defer_startup {
// Dormant servers must not hold the initial startup summary open.
deferred_workers.push(startup);
} else {
let trace = codex_rmcp_client::McpRequestTrace::current();
eager_workers.push(trace.scope(startup).boxed());
}
}
let summary_config = Arc::clone(&config);
let manager = Self {
config: Some(config),
network_policy_metadata,
servers,
event_stream_connection,
disabled_servers,
required_servers,
optional_startup_deadline: OnceLock::new(),
tool_plugin_context,
prefix_mcp_tool_names,
non_prefixed_mcp_tool_servers,
elicitation_requests: elicitation_requests.clone(),
trusted_access,
};
tokio::spawn(async move {
if !publication_gate.wait().await {
return;
}
// Announce the complete eager batch before a fast helper can report Ready.
// Dormant helpers announce themselves separately when their first caller arrives.
if let Some(tx_event) = &tx_event {
for server in eager_starts {
let _ = emit_update(
startup_submit_id.as_str(),
tx_event,
McpStartupUpdateEvent {
server,
status: McpStartupStatus::Starting,
},
)
.await;
}
}
for worker in deferred_workers {
tokio::spawn(worker);
}
let outcomes = JoinSet::from_iter(eager_workers).join_all().await;
let legacy_name = |key: McpServerKey| {
summary_config
.mcp_server_catalog
.legacy_server_name(&key)
.map(str::to_owned)
};
if let Some(tx_event) = tx_event {
let mut summary = McpStartupCompleteEvent {
ready: reused_ready.into_iter().filter_map(&legacy_name).collect(),
..Default::default()
};
for server_name in &summary.ready {
let _ = emit_update(
startup_submit_id.as_str(),
&tx_event,
McpStartupUpdateEvent {
server: server_name.clone(),
status: McpStartupStatus::Ready,
},
)
.await;
}
for (key, outcome) in outcomes {
let Some(server_name) = legacy_name(key) else {
continue;
};
match outcome {
Ok(_) => summary.ready.push(server_name),
Err(StartupOutcomeError::Cancelled) => summary.cancelled.push(server_name),
Err(StartupOutcomeError::Failed { error, .. }) => {
summary.failed.push(McpStartupFailure {
server: server_name,
error,
})
}
}
}
let _ = tx_event
.send(Event {
id: startup_submit_id,
msg: EventMsg::McpStartupComplete(summary),
})
.await;
}
});
manager
}
pub fn empty(prefix_mcp_tool_names: bool) -> Self {
Self {
config: None,
network_policy_metadata: None,
servers: HashMap::new(),
event_stream_connection: None,
disabled_servers: Vec::new(),
required_servers: Vec::new(),
optional_startup_deadline: OnceLock::new(),
tool_plugin_context: Arc::new(ToolPluginContext::default()),
prefix_mcp_tool_names,
non_prefixed_mcp_tool_servers: Vec::new(),
elicitation_requests: ElicitationRequestManager::default(),
trusted_access: None,
}
}
pub fn has_servers(&self) -> bool {
!self.servers.is_empty()
}
pub(crate) fn contains_server(&self, server_name: &str) -> bool {
self.server_by_name(server_name).is_some()
}
pub(crate) fn shares_server_connection(&self, other: &Self, server: &McpServerKey) -> bool {
self.servers
.get(server)
.zip(other.servers.get(server))
.is_some_and(|(left, right)| Arc::ptr_eq(&left.connection, &right.connection))
}
pub(crate) async fn authentication_failed_servers(&self) -> Vec<McpServerKey> {
let mut failed_servers = Vec::new();
for (server_name, view) in &self.servers {
if view
.connection
.client
.startup_complete
.load(Ordering::Acquire)
&& let Err(error) = view.connection.client().await
&& error.is_authentication_required()
{
failed_servers.push(server_name.clone());
}
}
failed_servers
}
pub(crate) async fn updated_oauth_credentials_after_auth_failure(
&self,
config: &crate::McpConfig,
) -> Vec<McpServerKey> {
let mut candidates = Vec::new();
for server_name in self.authentication_failed_servers().await {
if let Some(view) = self.servers.get(&server_name)
&& let Some(identity) = view.connection.identity.as_ref()
&& let Some(server) = config.mcp_server_catalog.server(&server_name)
{
candidates.push((
server_name,
server.name().to_owned(),
identity.clone(),
server.config().clone(),
));
}
}
if candidates.is_empty() {
return Vec::new();
}
let originator = AuthStorageOriginator::current();
match tokio::task::spawn_blocking(move || {
originator.sync_scope(|| {
candidates
.into_iter()
.filter_map(|(server_key, server_name, identity, config)| {
identity
.oauth_credentials_changed(&server_name, &config)
.then_some(server_key)
})
.collect()
})
})
.await
{
Ok(recovered_servers) => recovered_servers,
Err(error) => {
warn!(%error, "failed to inspect stored MCP OAuth credentials");
Vec::new()
}
}
}
pub(crate) async fn wait_for_server_startup(&self, server_name: &McpServerKey) -> bool {
let Some(view) = self.server(server_name) else {
return false;
};
view.connection.client.ready_transport().is_some()
|| view.connection.client_for_call().await.is_ok()
}
/// Release this manager's clients, terminating shared stdio servers after their final release.
pub async fn shutdown(&self) {
let connections = self
.servers
.values()
.map(|view| Arc::clone(&view.connection))
.collect::<Vec<_>>();
Self::shutdown_connections(connections).await;
}
/// Releases pool leases retained by older bindings.
pub(crate) async fn shutdown_shared_connections(&self) {
let connections = self
.servers
.values()
.filter(|view| view.connection.client.connection_slot.is_some())
.map(|view| Arc::clone(&view.connection))
.collect();
Self::shutdown_connections(connections).await;
}
/// Keeps cleanup alive if an interrupt cancels the refresh that requested it.
async fn shutdown_connections(connections: Vec<Arc<McpServerConnection>>) {
let shutdown_task = tokio::spawn(async move {
join_all(connections.iter().map(|connection| connection.shutdown())).await;
});
if let Err(error) = shutdown_task.await {
warn!("MCP client shutdown task failed: {error}");
}
}
pub(crate) fn cancel_startup(&self) {
for view in self.servers.values() {
view.connection.cancel_startup();
}
}
pub(crate) fn retire_shared_connections(&self) {
for view in self.servers.values() {
if let Some(slot) = &view.connection.client.connection_slot {
slot.retire();
}
}
}
pub fn plugin_id_for_server(&self, server_name: &McpServerKey) -> Option<&str> {
self.tool_plugin_context.plugin_id_for_server(server_name)
}
pub fn is_selected_plugin_server(&self, server_name: &McpServerKey) -> bool {
self.tool_plugin_context
.is_selected_plugin_server(server_name)
}
pub async fn wait_for_server_ready(
&self,
server_name: &McpServerKey,
timeout: Duration,
) -> bool {
let Some(view) = self.server(server_name) else {
return false;
};
match tokio::time::timeout(timeout, view.connection.client()).await {
Ok(Ok(_)) => true,
Ok(Err(_)) | Err(_) => false,
}
}
/// Selects the host call's server and validates it before startup or attribution.
/// Returns its exact key so waiting, execution, and attribution use the same server.
pub(crate) fn resolve_tool_call_server(
&self,
server: &str,
tool: &str,
environment_id: Option<&str>,
) -> Result<&McpServerKey> {
let (server_key, view) = match environment_id {
// Hooks name their environment; a thread-wide registration is the only fallback.
Some(id) => self
.servers
.get_key_value(&McpServerKey::in_environment(id, server))
.or_else(|| self.servers.get_key_value(&McpServerKey::unscoped(server))),
// Name-only callers still select the catalog's legacy winner.
None => self
.resolve_server_name(server)
.and_then(|key| self.servers.get_key_value(key)),
}
.ok_or_else(|| anyhow!("unknown MCP server '{server}'"))?;
if let Some(environment_id) = environment_id
&& view.metadata.environment_id != environment_id
{
bail!(
"MCP server `{server}` is running in environment `{}`, expected `{environment_id}`",
view.metadata.environment_id
);
}
if !view.tool_filter.allows(tool) {
return Err(anyhow!(
"tool '{tool}' is disabled for MCP server '{server}'"
));
}
Ok(server_key)
}
/// Invoke the tool indicated by the (server, tool) pair.
#[allow(clippy::too_many_arguments)]
pub async fn call_tool(
&self,
server_key: &McpServerKey,
tool: &str,
arguments: Option<serde_json::Value>,
mut meta: Option<serde_json::Value>,
requested_timeout: Option<Duration>,
wait_for_server: bool,
) -> Result<CallToolResult> {
let view = self
.server(server_key)
.ok_or_else(|| anyhow!("unknown MCP server '{server_key}'"))?;
if !view.tool_filter.allows(tool) {
bail!("tool '{tool}' is disabled for MCP server '{server_key}'");
}
let client = if wait_for_server {
view.connection
.client()
.await
.context("failed to get client")?
} else {
let client = view
.connection
.client
.ready_client()
.ok_or_else(|| anyhow!("MCP server '{server_key}' is not connected"))?;
if client.client.is_closed().await {
bail!("MCP server '{server_key}' is not connected");
}
client
};
let effective_timeout = match (view.tool_timeout, requested_timeout) {
(Some(server_timeout), Some(requested_timeout)) => {
Some(server_timeout.min(requested_timeout))
}
(server_timeout, requested_timeout) => server_timeout.or(requested_timeout),
};
// Direct callers cannot supply host-owned entitlement metadata, even for unlisted tools.
if let Some(serde_json::Value::Object(meta)) = meta.as_mut() {
meta.remove(ENTITLEMENT_CONTEXT_KEY);
}
let result: rmcp::model::CallToolResult = client
.call_tool(
tool.to_string(),
arguments,
self.network_policy_meta_for_call(
&client,
self.config
.as_ref()
.and_then(|config| config.permission_profile_for_server(server_key)),
&view.metadata.environment_id,
meta,
),
effective_timeout,
)
.await
.with_context(|| format!("tool call failed for `{server_key}/{tool}`"))?;
Ok(call_tool_result_from_rmcp(result))
}
/// Capabilities belong to the initialized connection, never a shared tool cache.
pub(crate) fn list_available_server_capabilities(
&self,
include_server: impl Fn(&str) -> bool,
) -> HashMap<String, serde_json::Value> {
self.legacy_servers()
.filter(|(name, _)| include_server(name))
.filter_map(|(name, view)| {
view.connection
.retained_legacy()
.unwrap_or(&view.connection)
.client
.server_capabilities
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone()
.map(|capabilities| (name.to_owned(), capabilities))
})
.collect()
}
/// Returns presentation metadata from the current connection.
/// Codex Apps metadata may come from its existing cache; regular MCP server information is
/// connection-specific, so pending regular clients are awaited.
pub(crate) async fn list_available_server_infos(
&self,
include_server: impl Fn(&str) -> bool,
) -> HashMap<String, McpServerInfo> {
let mut server_infos = HashMap::new();
for (server_name, view) in self
.legacy_servers()
.filter(|(name, _)| include_server(name))
{
let server_name = server_name.to_owned();
let client = &view.connection.client;
if !client.startup_complete.load(Ordering::Acquire)
&& let Some(server_info) = client.cached_server_info.clone()
{
server_infos.insert(server_name.clone(), server_info);
continue;
}
match view.connection.client().await {
Ok(managed_client) => {
server_infos.insert(server_name.clone(), managed_client.server_info);
}
Err(_) => {
if let Some(server_info) = client.cached_server_info.clone() {
server_infos.insert(server_name.clone(), server_info);
}
}
}
}
server_infos
}
}
#[cfg(test)]
#[path = "connection_manager_tests.rs"]
pub(crate) mod tests;
Reference in New Issue View Git Blame Copy Permalink