diff --git a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs index 3756466c330a..75192c054932 100644 --- a/codex-rs/app-server/src/request_processors/apps_processor/installed.rs +++ b/codex-rs/app-server/src/request_processors/apps_processor/installed.rs @@ -70,8 +70,7 @@ impl AppsRequestProcessor { ); let startup_timeout = mcp_servers .get(CODEX_APPS_MCP_SERVER_NAME) - .and_then(|server| server.configured_config()) - .and_then(|config| config.startup_timeout_sec) + .and_then(|server| server.config().startup_timeout_sec) .unwrap_or(CONNECTOR_RUNTIME_REFRESH_TIMEOUT); let runtime_context = McpRuntimeContext::new( self.thread_manager.environment_manager(), diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 7563a80ebd0b..6b154fb57226 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -144,14 +144,12 @@ impl McpRequestProcessor { } }; let effective_servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref()); - let Some(server) = effective_servers - .get(&name) - .and_then(codex_mcp::EffectiveMcpServer::configured_config) - else { + let Some(server) = effective_servers.get(&name) else { return Err(invalid_request(format!( "No MCP server named '{name}' found." ))); }; + let server = server.config(); let (url, http_headers, env_http_headers) = match &server.transport { McpServerTransportConfig::StreamableHttp { diff --git a/codex-rs/codex-mcp/src/binding_tests.rs b/codex-rs/codex-mcp/src/binding_tests.rs index c3be31999d76..c33d5673ffc3 100644 --- a/codex-rs/codex-mcp/src/binding_tests.rs +++ b/codex-rs/codex-mcp/src/binding_tests.rs @@ -26,7 +26,6 @@ use crate::connection_manager::McpConnectionSet; use crate::rmcp_client::ManagedClient; use crate::server::McpServerMetadata; use crate::server::McpServerOrigin; -use crate::tools::ToolFilter; use crate::tools::ToolInfo; const SERVER_NAME: &str = "docs"; @@ -88,7 +87,6 @@ async fn test_step( website_url: None, }, tools: vec![tool.clone()], - tool_filter: ToolFilter::default(), tool_timeout: None, server_instructions: None, server_supports_sandbox_state_meta_capability: supports_sandbox_state_meta, diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 520de0c9cda7..bb71e4eb0ed1 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,46 +7,46 @@ #[path = "connection_manager/required.rs"] mod required; +#[path = "connection_manager/resources.rs"] +mod resources; +#[path = "connection_manager/startup.rs"] +mod startup; #[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 use tool_catalog::tool_is_model_visible; + use std::collections::HashMap; -use std::path::PathBuf; use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Duration; -use crate::binding_clients::McpBindingClients; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; -use crate::elicitation::ElicitationReviewerHandle; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; 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::McpRuntimeContext; -use crate::server::EffectiveMcpServer; +use crate::runtime::McpRuntimeInput; +use crate::server::McpServerConnectionIdentity; use crate::server::McpServerMetadata; -use crate::tool_catalog_cache::McpToolCatalogCache; +use crate::tools::ToolFilter; use crate::tools::ToolInfo; +use crate::tools::filter_tools; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; -use async_channel::Sender; -use codex_api::SharedAuthProvider; -use codex_config::Constrained; -use codex_config::McpServerAuth; -use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; -use codex_config::types::AuthKeyringBackendKind; -use codex_config::types::OAuthCredentialsStoreMode; -use codex_connectors::ConnectorRuntimeContextKey; -use codex_connectors::ConnectorRuntimeManager; -use codex_login::AuthManager; -use codex_login::CodexAuth; use codex_protocol::mcp::CallToolResult; use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; @@ -55,113 +55,176 @@ 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::McpAuthState; -use codex_rmcp_client::McpLoginRequirement; use codex_rmcp_client::determine_streamable_http_auth_status_from_credentials; -use rmcp::model::ElicitationCapability; -use rmcp::model::ListResourcesResult; -use rmcp::model::PaginatedRequestParams; -use rmcp::model::ReadResourceRequestParams; -use rmcp::model::ReadResourceResult; -use serde_json::Value as JsonValue; use tokio::sync::Mutex; use tokio::sync::RwLock; use tokio::task::JoinSet; -use tokio_util::sync::CancellationToken; use tracing::warn; -const MCP_UI_META_KEY: &str = "ui"; -const MCP_UI_VISIBILITY_META_KEY: &str = "visibility"; -const MCP_UI_MODEL_VISIBILITY: &str = "model"; - -/// Returns whether a tool may be included in model-facing tool declarations. -/// -/// Tools without visibility metadata remain visible. -/// Tools with visibility metadata are hidden unless they explicitly include `model`. -/// -/// -pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { - let Some(visibility) = tool - .tool - .meta - .as_deref() - .and_then(|meta| meta.get(MCP_UI_META_KEY)) - .and_then(JsonValue::as_object) - .and_then(|ui| ui.get(MCP_UI_VISIBILITY_META_KEY)) - .and_then(JsonValue::as_array) - else { - return true; - }; - - visibility - .iter() - .any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY)) +#[cfg(test)] +use codex_config::Constrained; +#[cfg(test)] +use codex_protocol::protocol::McpStartupFailureReason; +#[cfg(test)] +use codex_rmcp_client::McpAuthState; +#[cfg(test)] +use codex_rmcp_client::McpLoginRequirement; + +pub(crate) struct McpServerConnection { + identity: Option, + client: AsyncManagedClient, } -/// A thin wrapper around a set of running [`RmcpClient`] instances. +impl McpServerConnection { + async fn reusable_client( + &self, + desired: &McpServerConnectionIdentity, + ) -> Option { + 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()?; + 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 == 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 { + self.client.client().await + } + + async fn shutdown(&self) { + self.client.shutdown().await; + } + + fn cancel_startup(&self) { + if !self.client.startup_complete.load(Ordering::Acquire) { + self.client.cancel_token.cancel(); + } + } +} + +impl Drop for McpServerConnection { + fn drop(&mut self) { + self.client.cancel_token.cancel(); + } +} + +#[derive(Clone)] +struct McpServerView { + connection: Arc, + metadata: McpServerMetadata, + tool_filter: ToolFilter, + tool_timeout: Option, +} + +impl McpServerView { + async fn listed_tools( + &self, + tool_plugin_provenance: &ToolPluginProvenance, + ) -> Option> { + let tools = self.connection.client.listed_tools().await?; + let tools = filter_tools(tools, &self.tool_filter); + Some(if self.connection.client.is_codex_apps_mcp_server { + prepare_codex_apps_tools_for_model(tools, tool_plugin_provenance) + } else { + prepare_regular_mcp_tools_for_model(tools, tool_plugin_provenance) + }) + } +} + +/// A published view over a set of running MCP server connections. pub(crate) struct McpConnectionSet { - clients: HashMap, - server_metadata: HashMap, + servers: HashMap, required_servers: Vec, tool_catalog_revision: Arc>, codex_apps_tools_override: RwLock>>, codex_apps_refresh_lock: Mutex<()>, tool_plugin_provenance: Arc, prefix_mcp_tool_names: bool, - startup_cancellation_token: CancellationToken, + elicitation_requests: ElicitationRequestManager, } impl McpConnectionSet { /// Creates an MCP connection manager. Threadless callers can pass no `tx_event`; startup /// notifications are then skipped and interactive elicitations are declined. - #[allow(clippy::too_many_arguments)] pub async fn new( - mcp_servers: &HashMap, - store_mode: OAuthCredentialsStoreMode, - keyring_backend_kind: AuthKeyringBackendKind, - approval_policy: &Constrained, - submit_id: String, - tx_event: Option>, - startup_cancellation_token: CancellationToken, - initial_permission_profile: PermissionProfile, - runtime_context: McpRuntimeContext, - codex_home: PathBuf, - codex_apps_tools_cache: ConnectorRuntimeManager, - tool_catalog_cache: McpToolCatalogCache, - codex_apps_tools_cache_key: ConnectorRuntimeContextKey, - prefix_mcp_tool_names: bool, - client_elicitation_capability: ElicitationCapability, - supports_openai_form_elicitation: bool, - tool_plugin_provenance: ToolPluginProvenance, - auth: Option<&CodexAuth>, - codex_apps_auth_manager: Option>, - elicitation_reviewer: Option, - elicitation_lifecycle: Option, - elicitation_router: ElicitationRequestRouter, + previous: Option<&Self>, publication_gate: McpPublicationGate, + input: McpRuntimeInput, + elicitation_router: ElicitationRequestRouter, ) -> Self { - let mut clients = HashMap::new(); - let mut server_metadata = HashMap::new(); + let McpRuntimeInput { + config, + 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, + supports_openai_form_elicitation, + auth, + codex_apps_auth_manager, + elicitation_reviewer, + elicitation_lifecycle, + } = input; + let store_mode = config.mcp_oauth_credentials_store_mode; + let keyring_backend_kind = config.auth_keyring_backend_kind; + let approval_policy = &config.approval_policy; + let initial_permission_profile = config.permission_profile.clone(); + let codex_home = config.codex_home.clone(); + let prefix_mcp_tool_names = config.prefix_mcp_tool_names; + let client_elicitation_capability = config.client_elicitation_capability.clone(); + let tool_plugin_provenance = crate::mcp::tool_plugin_provenance(&config); + let auth = auth.as_ref(); + let mut servers = HashMap::new(); let mut required_servers = mcp_servers .iter() .filter(|(_, server)| server.enabled() && server.required()) .map(|(server_name, _)| server_name.clone()) .collect::>(); required_servers.sort(); + let mut reused_ready = Vec::new(); let mut join_set = JoinSet::new(); - let elicitation_requests = ElicitationRequestManager::new( - approval_policy.value(), - initial_permission_profile, - elicitation_reviewer, - elicitation_lifecycle, - elicitation_router, - ); + let reusable_previous = previous.filter(|previous| { + !previous.servers.is_empty() + && previous.elicitation_requests.update( + approval_policy.value(), + initial_permission_profile.clone(), + elicitation_reviewer.clone(), + elicitation_lifecycle.clone(), + ) + }); + let elicitation_requests = if let Some(previous) = reusable_previous { + previous.elicitation_requests.clone() + } else { + ElicitationRequestManager::new( + approval_policy.value(), + initial_permission_profile, + elicitation_reviewer, + elicitation_lifecycle, + elicitation_router, + ) + }; let tool_plugin_provenance = Arc::new(tool_plugin_provenance); let startup_submit_id = submit_id; - let startup_publication_gate = publication_gate.clone(); let static_chatgpt_auth_provider = auth .filter(|auth| auth.uses_codex_backend()) .map(codex_model_provider::auth_provider_from_auth); @@ -170,31 +233,30 @@ impl McpConnectionSet { codex_model_provider::auth_provider_from_auth_manager(auth_manager, auth) }) }); - let mcp_servers = mcp_servers.clone(); for (server_name, server) in mcp_servers .into_iter() .filter(|(_, server)| server.enabled()) { - server_metadata.insert(server_name.clone(), McpServerMetadata::from(&server)); - let cancel_token = startup_cancellation_token.child_token(); - let configured_config = server.configured_config().cloned(); - let resolved_environment = configured_config.as_ref().map_or_else( - || Ok(None), - |config| runtime_context.resolve_server_environment(&server_name, config), + let metadata = McpServerMetadata::from(&server); + let configured_config = server.config().clone(); + let configured_tool_filter = ToolFilter::from_config(&configured_config); + 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 = - configured_config - .as_ref() - .is_some_and(|config| match &config.transport { - McpServerTransportConfig::StreamableHttp { - bearer_token_env_var, - .. - } => bearer_token_env_var.is_some(), - McpServerTransportConfig::Stdio { .. } => false, - }); + let uses_env_bearer_token = match &configured_config.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => bearer_token_env_var.is_some(), + McpServerTransportConfig::Stdio { .. } => false, + }; let shares_codex_apps_tools_cache = should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { @@ -220,14 +282,49 @@ impl McpConnectionSet { } else { chatgpt_auth_provider_for_server(&server, chatgpt_auth_provider) }; + let connection_identity = McpServerConnectionIdentity::new( + &server_name, + &server, + store_mode, + keyring_backend_kind, + &resolved_environment, + &runtime_context, + runtime_auth_provider.as_ref(), + auth, + shares_codex_apps_tools_cache + .then(|| (codex_home.clone(), codex_apps_tools_cache_key.clone())), + client_elicitation_capability.clone(), + supports_openai_form_elicitation, + ); + if let Some(previous_view) = + reusable_previous.and_then(|previous| previous.servers.get(&server_name)) + { + let connection = Arc::clone(&previous_view.connection); + if connection + .reusable_client(&connection_identity) + .await + .is_some() + { + servers.insert( + server_name.clone(), + McpServerView { + connection, + metadata, + tool_filter: configured_tool_filter, + tool_timeout: configured_tool_timeout, + }, + ); + reused_ready.push(server_name); + continue; + } + } + let cancel_token = startup_cancellation_token.child_token(); let tool_catalog_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { None - } else if let Some(config) = configured_config.as_ref() - && let Ok(environment) = resolved_environment.as_ref() - { + } else if let Ok(environment) = resolved_environment.as_ref() { tool_catalog_cache.context( &server_name, - config, + &configured_config, &runtime_context, environment.as_ref(), &client_elicitation_capability, @@ -248,14 +345,24 @@ impl McpConnectionSet { elicitation_requests.clone(), codex_apps_tools_cache_context, tool_catalog_cache_context, - Arc::clone(&tool_plugin_provenance), runtime_context.clone(), resolved_environment, runtime_auth_provider, client_elicitation_capability.clone(), supports_openai_form_elicitation, ); - clients.insert(server_name.clone(), async_managed_client.clone()); + servers.insert( + server_name.clone(), + McpServerView { + connection: Arc::new(McpServerConnection { + identity: Some(connection_identity), + client: async_managed_client.clone(), + }), + metadata, + tool_filter: configured_tool_filter, + tool_timeout: configured_tool_timeout, + }, + ); let tx_event = tx_event.clone(); let submit_id = startup_submit_id.clone(); let publication_gate = publication_gate.clone(); @@ -281,34 +388,33 @@ impl McpConnectionSet { if let Some(tx_event) = tx_event.as_ref() { let auth_state = match &outcome { Err(error) if error.is_authentication_required() && !has_runtime_auth => { - configured_config.as_ref().and_then(|config| { - let McpServerTransportConfig::StreamableHttp { + match &configured_config.transport { + McpServerTransportConfig::StreamableHttp { url, bearer_token_env_var, http_headers, env_http_headers, - } = &config.transport - else { - return None; - }; - match determine_streamable_http_auth_status_from_credentials( - &server_name, - 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 + } => { + match determine_streamable_http_auth_status_from_credentials( + &server_name, + 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, }; @@ -322,7 +428,7 @@ impl McpConnectionSet { let reason = mcp_startup_failure_reason(auth_state, error); let error_str = mcp_init_error_display( server_name.as_str(), - configured_config.as_ref(), + Some(&configured_config), error, ); McpStartupStatus::Failed { @@ -354,23 +460,37 @@ impl McpConnectionSet { }); } let manager = Self { - clients, - server_metadata, + servers, required_servers, tool_catalog_revision: Arc::new(RwLock::new(0)), codex_apps_tools_override: RwLock::new(None), codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance, prefix_mcp_tool_names, - startup_cancellation_token, + elicitation_requests: elicitation_requests.clone(), }; + let summary_publication_gate = publication_gate; tokio::spawn(async move { let outcomes = join_set.join_all().await; - if !startup_publication_gate.wait().await { - return; - } if let Some(tx_event) = tx_event { - let mut summary = McpStartupCompleteEvent::default(); + if !summary_publication_gate.wait().await { + return; + } + let mut summary = McpStartupCompleteEvent { + ready: reused_ready, + ..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 (server_name, outcome) in outcomes { match outcome { Ok(_) => summary.ready.push(server_name), @@ -394,40 +514,68 @@ impl McpConnectionSet { manager } - pub fn empty(prefix_mcp_tool_names: bool) -> Self { + #[cfg(test)] + fn new_uninitialized_with_permission_profile( + approval_policy: &Constrained, + permission_profile: &PermissionProfile, + prefix_mcp_tool_names: bool, + ) -> Self { Self { - clients: HashMap::new(), - server_metadata: HashMap::new(), + servers: HashMap::new(), required_servers: Vec::new(), tool_catalog_revision: Arc::new(RwLock::new(0)), codex_apps_tools_override: RwLock::new(None), codex_apps_refresh_lock: Mutex::new(()), tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), prefix_mcp_tool_names, - startup_cancellation_token: CancellationToken::new(), + elicitation_requests: ElicitationRequestManager::new( + approval_policy.value(), + permission_profile.clone(), + /*reviewer*/ None, + /*lifecycle*/ None, + ElicitationRequestRouter::default(), + ), } } - pub fn has_servers(&self) -> bool { - !self.clients.is_empty() + pub fn empty(prefix_mcp_tool_names: bool) -> Self { + Self { + servers: HashMap::new(), + required_servers: Vec::new(), + tool_catalog_revision: Arc::new(RwLock::new(0)), + codex_apps_tools_override: RwLock::new(None), + codex_apps_refresh_lock: Mutex::new(()), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + prefix_mcp_tool_names, + elicitation_requests: ElicitationRequestManager::new( + AskForApproval::Never, + PermissionProfile::default(), + /*reviewer*/ None, + /*lifecycle*/ None, + ElicitationRequestRouter::default(), + ), + } } - pub(crate) fn cancel_startup(&self) { - self.startup_cancellation_token.cancel(); + pub fn has_servers(&self) -> bool { + !self.servers.is_empty() } pub(crate) fn contains_server(&self, server_name: &str) -> bool { - self.clients.contains_key(server_name) + self.servers.contains_key(server_name) } /// Stop all MCP clients owned by this manager and terminate stdio server processes. pub async fn shutdown(&self) { - self.startup_cancellation_token.cancel(); - let clients = self.clients.values().cloned().collect::>(); + let connections = self + .servers + .values() + .map(|view| Arc::clone(&view.connection)) + .collect::>(); // Keep cleanup alive if an interrupt cancels the refresh that requested it. let shutdown_task = tokio::spawn(async move { - for client in clients { - client.shutdown().await; + for connection in connections { + connection.shutdown().await; } }); if let Err(error) = shutdown_task.await { @@ -435,6 +583,12 @@ impl McpConnectionSet { } } + pub(crate) fn cancel_startup(&self) { + for view in self.servers.values() { + view.connection.cancel_startup(); + } + } + pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> { self.tool_plugin_provenance .plugin_id_for_mcp_server_name(server_name) @@ -446,53 +600,16 @@ impl McpConnectionSet { } pub async fn wait_for_server_ready(&self, server_name: &str, timeout: Duration) -> bool { - let Some(async_managed_client) = self.clients.get(server_name) else { + let Some(view) = self.servers.get(server_name) else { return false; }; - match tokio::time::timeout(timeout, async_managed_client.client()).await { + match tokio::time::timeout(timeout, view.connection.client()).await { Ok(Ok(_)) => true, Ok(Err(_)) | Err(_) => false, } } - pub async fn list_all_resources( - &self, - include_server: impl Fn(&str) -> bool, - ) -> HashMap> { - self.ready_clients_matching(&include_server) - .await - .list_all_resources(|_| true) - .await - } - - pub async fn list_all_resource_templates( - &self, - include_server: impl Fn(&str) -> bool, - ) -> HashMap> { - self.ready_clients_matching(&include_server) - .await - .list_all_resource_templates(|_| true) - .await - } - - async fn ready_clients_matching( - &self, - include_server: &impl Fn(&str) -> bool, - ) -> McpBindingClients { - let mut clients = HashMap::new(); - for (server, client) in self - .clients - .iter() - .filter(|(server, _)| include_server(server)) - { - if let Ok(client) = client.client().await { - clients.insert(server.clone(), Arc::new(client)); - } - } - McpBindingClients::new(clients) - } - /// Invoke the tool indicated by the (server, tool) pair. pub async fn call_tool( &self, @@ -501,16 +618,24 @@ impl McpConnectionSet { arguments: Option, meta: Option, ) -> Result { - let client = self.client_by_name(server).await?; - if !client.tool_filter.allows(tool) { + let view = self + .servers + .get(server) + .ok_or_else(|| anyhow!("unknown MCP server '{server}'"))?; + if !view.tool_filter.allows(tool) { return Err(anyhow!( "tool '{tool}' is disabled for MCP server '{server}'" )); } + let client = view + .connection + .client() + .await + .context("failed to get client")?; let result: rmcp::model::CallToolResult = client .client - .call_tool(tool.to_string(), arguments, meta, client.tool_timeout) + .call_tool(tool.to_string(), arguments, meta, view.tool_timeout) .await .with_context(|| format!("tool call failed for `{server}/{tool}`"))?; @@ -531,45 +656,13 @@ impl McpConnectionSet { }) } - /// List resources from the specified server. - pub async fn list_resources( - &self, - server: &str, - params: Option, - ) -> Result { - let managed = self.client_by_name(server).await?; - let timeout = managed.tool_timeout; - - managed - .client - .list_resources(params, timeout) - .await - .with_context(|| format!("resources/list failed for `{server}`")) - } - - /// Read a resource from the specified server. - pub async fn read_resource( - &self, - server: &str, - params: ReadResourceRequestParams, - ) -> Result { - let managed = self.client_by_name(server).await?; - let client = managed.client.clone(); - let timeout = managed.tool_timeout; - let uri = params.uri.clone(); - - client - .read_resource(params, timeout) - .await - .with_context(|| format!("resources/read failed for `{server}` ({uri})")) - } - /// 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) -> HashMap { let mut server_infos = HashMap::new(); - for (server_name, client) in &self.clients { + for (server_name, view) in &self.servers { + let client = &view.connection.client; if !client.startup_complete.load(Ordering::Acquire) && let Some(server_info) = client.cached_server_info.clone() { @@ -590,136 +683,65 @@ impl McpConnectionSet { server_infos } - async fn client_by_name(&self, name: &str) -> Result { - self.clients - .get(name) - .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))? - .client() - .await - .context("failed to get client") - } - #[cfg(test)] fn new_uninitialized( - _approval_policy: &Constrained, - _permission_profile: &Constrained, + approval_policy: &Constrained, + permission_profile: &Constrained, prefix_mcp_tool_names: bool, ) -> Self { - Self::empty(prefix_mcp_tool_names) - } -} - -impl Drop for McpConnectionSet { - fn drop(&mut self) { - self.startup_cancellation_token.cancel(); - self.clients.clear(); - } -} - -/// Makes ChatGPT authentication available to servers that explicitly opt in. -/// The HTTP transport applies it only when no configured authorization resolves. -fn chatgpt_auth_provider_for_server( - server: &EffectiveMcpServer, - chatgpt_auth_provider: Option, -) -> Option { - if !server - .configured_config() - .is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt)) - { - return None; - } - chatgpt_auth_provider -} - -fn should_share_codex_apps_tools_cache(server_name: &str, uses_env_bearer_token: bool) -> bool { - server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token -} - -async fn emit_update( - submit_id: &str, - tx_event: &Sender, - update: McpStartupUpdateEvent, -) -> Result<(), async_channel::SendError> { - tx_event - .send(Event { - id: submit_id.to_string(), - msg: EventMsg::McpStartupUpdate(update), - }) - .await -} - -fn mcp_startup_failure_reason( - auth_state: Option, - error: &StartupOutcomeError, -) -> Option { - if !error.is_authentication_required() { - return None; + Self::new_uninitialized_with_permission_profile( + approval_policy, + permission_profile.get(), + prefix_mcp_tool_names, + ) } - match auth_state { - Some(McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication)) => { - Some(McpStartupFailureReason::ReauthenticationRequired) - } - Some( - McpAuthState::Unsupported - | McpAuthState::LoggedOut(McpLoginRequirement::Login) - | McpAuthState::BearerToken - | McpAuthState::OAuth, - ) - | None => None, + #[cfg(test)] + fn insert_test_client(&mut self, name: impl Into, client: AsyncManagedClient) { + let name = name.into(); + self.servers.insert( + name, + McpServerView { + tool_filter: ToolFilter::default(), + connection: Arc::new(McpServerConnection { + identity: None, + client, + }), + metadata: McpServerMetadata { + environment_id: String::new(), + pollutes_memory: true, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }, + tool_timeout: None, + }, + ); } -} -fn mcp_init_error_display( - server_name: &str, - config: Option<&McpServerConfig>, - err: &StartupOutcomeError, -) -> String { - if let Some(McpServerTransportConfig::StreamableHttp { - url, - bearer_token_env_var, - http_headers, - .. - }) = config.map(|config| &config.transport) - && url == "https://api.githubcopilot.com/mcp/" - && bearer_token_env_var.is_none() - && http_headers.as_ref().map(HashMap::is_empty).unwrap_or(true) - { - format!( - "GitHub MCP does not support OAuth. Log in by adding a personal access token (https://github.com/settings/personal-access-tokens) to your environment and config.toml:\n[mcp_servers.{server_name}]\nbearer_token_env_var = CODEX_GITHUB_PERSONAL_ACCESS_TOKEN" - ) - } else if is_mcp_client_auth_required_error(err) { - format!( - "The {server_name} MCP server is not logged in. Run `codex mcp login {server_name}`." - ) - } else if is_mcp_client_startup_timeout_error(err) { - let startup_timeout_secs = config - .and_then(|config| config.startup_timeout_sec) - .unwrap_or(DEFAULT_STARTUP_TIMEOUT) - .as_secs(); - format!( - "MCP client for `{server_name}` timed out after {startup_timeout_secs} seconds. Add or adjust `startup_timeout_sec` in your config.toml:\n[mcp_servers.{server_name}]\nstartup_timeout_sec = XX" - ) - } else { - format!("MCP client for `{server_name}` failed to start: {err:#}") + #[cfg(test)] + fn test_client(&self, name: &str) -> &AsyncManagedClient { + &self.servers[name].connection.client } -} -fn is_mcp_client_auth_required_error(error: &StartupOutcomeError) -> bool { - match error { - StartupOutcomeError::Failed { error, .. } => error.contains("Auth required"), - _ => false, + #[cfg(test)] + fn set_test_server_metadata(&mut self, name: &str, metadata: McpServerMetadata) { + self.servers + .get_mut(name) + .expect("test server exists") + .metadata = metadata; } -} -fn is_mcp_client_startup_timeout_error(error: &StartupOutcomeError) -> bool { - match error { - StartupOutcomeError::Failed { error, .. } => { - error.contains("request timed out") - || error.contains("timed out handshaking with MCP server") - || error.contains("MCP client startup timed out") - } - _ => false, + #[cfg(test)] + fn shares_test_connection_with(&self, other: &Self, name: &str) -> bool { + let Some(left) = self.servers.get(name) else { + return false; + }; + let Some(right) = other.servers.get(name) else { + return false; + }; + Arc::ptr_eq(&left.connection, &right.connection) } } diff --git a/codex-rs/codex-mcp/src/connection_manager/required.rs b/codex-rs/codex-mcp/src/connection_manager/required.rs index 82b9a7a36d92..4ba65c53ba32 100644 --- a/codex-rs/codex-mcp/src/connection_manager/required.rs +++ b/codex-rs/codex-mcp/src/connection_manager/required.rs @@ -16,7 +16,7 @@ impl McpConnectionSet { let failures = async { let mut failures = Vec::new(); for server_name in &self.required_servers { - let Some(async_managed_client) = self.clients.get(server_name).cloned() else { + let Some(view) = self.servers.get(server_name) else { failures.push(McpStartupFailure { server: server_name.clone(), error: format!("required MCP server `{server_name}` was not initialized"), @@ -24,7 +24,7 @@ impl McpConnectionSet { continue; }; - match async_managed_client.client().await { + match view.connection.client().await { Ok(_) => {} Err(error) => failures.push(McpStartupFailure { server: server_name.clone(), diff --git a/codex-rs/codex-mcp/src/connection_manager/resources.rs b/codex-rs/codex-mcp/src/connection_manager/resources.rs new file mode 100644 index 000000000000..8a46971acbd2 --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_manager/resources.rs @@ -0,0 +1,183 @@ +use std::collections::HashMap; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use rmcp::model::ListResourcesResult; +use rmcp::model::PaginatedRequestParams; +use rmcp::model::ReadResourceRequestParams; +use rmcp::model::ReadResourceResult; +use rmcp::model::Resource; +use rmcp::model::ResourceTemplate; +use tokio::task::JoinSet; +use tracing::warn; + +use super::McpConnectionSet; +use crate::rmcp_client::ManagedClient; + +impl McpConnectionSet { + /// Returns resources from servers selected by `include_server`. + pub async fn list_all_resources( + &self, + include_server: impl Fn(&str) -> bool, + ) -> HashMap> { + let mut join_set = JoinSet::new(); + for (server_name, view) in self + .servers + .iter() + .filter(|(server_name, _)| include_server(server_name)) + { + let server_name = server_name.clone(); + let Ok(managed_client) = view.connection.client().await else { + continue; + }; + let timeout = view.tool_timeout; + let client = managed_client.client; + join_set.spawn(async move { + let mut resources = Vec::new(); + let mut cursor: Option = None; + loop { + let params = cursor.as_ref().map(|next| { + PaginatedRequestParams::default().with_cursor(Some(next.clone())) + }); + let response = match client.list_resources(params, timeout).await { + Ok(result) => result, + Err(error) => return (server_name, Err(error)), + }; + resources.extend(response.resources); + match response.next_cursor { + Some(next) if cursor.as_ref() == Some(&next) => { + return ( + server_name, + Err(anyhow!("resources/list returned duplicate cursor")), + ); + } + Some(next) => cursor = Some(next), + None => return (server_name, Ok(resources)), + } + } + }); + } + + let mut resources = HashMap::new(); + while let Some(result) = join_set.join_next().await { + match result { + Ok((server_name, Ok(server_resources))) => { + resources.insert(server_name, server_resources); + } + Ok((server_name, Err(error))) => { + warn!("Failed to list resources for MCP server '{server_name}': {error:#}"); + } + Err(error) => { + warn!("Task panic when listing resources for MCP server: {error:#}"); + } + } + } + resources + } + + /// Returns resource templates from servers selected by `include_server`. + pub async fn list_all_resource_templates( + &self, + include_server: impl Fn(&str) -> bool, + ) -> HashMap> { + let mut join_set = JoinSet::new(); + for (server_name, view) in self + .servers + .iter() + .filter(|(server_name, _)| include_server(server_name)) + { + let server_name = server_name.clone(); + let Ok(managed_client) = view.connection.client().await else { + continue; + }; + let timeout = view.tool_timeout; + let client = managed_client.client; + join_set.spawn(async move { + let mut templates = Vec::new(); + let mut cursor: Option = None; + loop { + let params = cursor.as_ref().map(|next| { + PaginatedRequestParams::default().with_cursor(Some(next.clone())) + }); + let response = match client.list_resource_templates(params, timeout).await { + Ok(result) => result, + Err(error) => return (server_name, Err(error)), + }; + templates.extend(response.resource_templates); + match response.next_cursor { + Some(next) if cursor.as_ref() == Some(&next) => { + return ( + server_name, + Err(anyhow!( + "resources/templates/list returned duplicate cursor" + )), + ); + } + Some(next) => cursor = Some(next), + None => return (server_name, Ok(templates)), + } + } + }); + } + + let mut templates = HashMap::new(); + while let Some(result) = join_set.join_next().await { + match result { + Ok((server_name, Ok(server_templates))) => { + templates.insert(server_name, server_templates); + } + Ok((server_name, Err(error))) => { + warn!( + "Failed to list resource templates for MCP server '{server_name}': {error:#}" + ); + } + Err(error) => { + warn!("Task panic when listing resource templates for MCP server: {error:#}"); + } + } + } + templates + } + + pub async fn list_resources( + &self, + server: &str, + params: Option, + ) -> Result { + let (managed, timeout) = self.client_by_name(server).await?; + managed + .client + .list_resources(params, timeout) + .await + .with_context(|| format!("resources/list failed for `{server}`")) + } + + pub async fn read_resource( + &self, + server: &str, + params: ReadResourceRequestParams, + ) -> Result { + let (managed, timeout) = self.client_by_name(server).await?; + let uri = params.uri.clone(); + managed + .client + .read_resource(params, timeout) + .await + .with_context(|| format!("resources/read failed for `{server}` ({uri})")) + } + + async fn client_by_name(&self, name: &str) -> Result<(ManagedClient, Option)> { + let view = self + .servers + .get(name) + .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))?; + let client = view + .connection + .client() + .await + .context("failed to get client")?; + Ok((client, view.tool_timeout)) + } +} diff --git a/codex-rs/codex-mcp/src/connection_manager/startup.rs b/codex-rs/codex-mcp/src/connection_manager/startup.rs new file mode 100644 index 000000000000..ec255e58130c --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_manager/startup.rs @@ -0,0 +1,115 @@ +use std::collections::HashMap; + +use anyhow::Result; +use async_channel::Sender; +use codex_api::SharedAuthProvider; +use codex_config::McpServerAuth; +use codex_config::McpServerConfig; +use codex_config::McpServerTransportConfig; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::McpStartupFailureReason; +use codex_protocol::protocol::McpStartupUpdateEvent; +use codex_rmcp_client::McpAuthState; +use codex_rmcp_client::McpLoginRequirement; + +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; +use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT; +use crate::rmcp_client::StartupOutcomeError; +use crate::server::EffectiveMcpServer; + +/// Makes ChatGPT authentication available to servers that explicitly opt in. +pub(super) fn chatgpt_auth_provider_for_server( + server: &EffectiveMcpServer, + chatgpt_auth_provider: Option, +) -> Option { + if !matches!(&server.config().auth, McpServerAuth::ChatGpt) { + return None; + } + chatgpt_auth_provider +} + +pub(super) fn should_share_codex_apps_tools_cache( + server_name: &str, + uses_env_bearer_token: bool, +) -> bool { + server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token +} + +pub(super) async fn emit_update( + submit_id: &str, + tx_event: &Sender, + update: McpStartupUpdateEvent, +) -> Result<(), async_channel::SendError> { + tx_event + .send(Event { + id: submit_id.to_string(), + msg: EventMsg::McpStartupUpdate(update), + }) + .await +} + +pub(super) fn mcp_startup_failure_reason( + auth_state: Option, + error: &StartupOutcomeError, +) -> Option { + if !error.is_authentication_required() { + return None; + } + match auth_state { + Some(McpAuthState::LoggedOut(McpLoginRequirement::Reauthentication)) => { + Some(McpStartupFailureReason::ReauthenticationRequired) + } + Some( + McpAuthState::Unsupported + | McpAuthState::LoggedOut(McpLoginRequirement::Login) + | McpAuthState::BearerToken + | McpAuthState::OAuth, + ) + | None => None, + } +} + +pub(super) fn mcp_init_error_display( + server_name: &str, + config: Option<&McpServerConfig>, + error: &StartupOutcomeError, +) -> String { + if let Some(McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var, + http_headers, + .. + }) = config.map(|config| &config.transport) + && url == "https://api.githubcopilot.com/mcp/" + && bearer_token_env_var.is_none() + && http_headers.as_ref().map(HashMap::is_empty).unwrap_or(true) + { + format!( + "GitHub MCP does not support OAuth. Log in by adding a personal access token (https://github.com/settings/personal-access-tokens) to your environment and config.toml:\n[mcp_servers.{server_name}]\nbearer_token_env_var = CODEX_GITHUB_PERSONAL_ACCESS_TOKEN" + ) + } else if matches!( + error, + StartupOutcomeError::Failed { error, .. } if error.contains("Auth required") + ) { + format!( + "The {server_name} MCP server is not logged in. Run `codex mcp login {server_name}`." + ) + } else if matches!( + error, + StartupOutcomeError::Failed { error, .. } + if error.contains("request timed out") + || error.contains("timed out handshaking with MCP server") + || error.contains("MCP client startup timed out") + ) { + let startup_timeout_secs = config + .and_then(|config| config.startup_timeout_sec) + .unwrap_or(DEFAULT_STARTUP_TIMEOUT) + .as_secs(); + format!( + "MCP client for `{server_name}` timed out after {startup_timeout_secs} seconds. Add or adjust `startup_timeout_sec` in your config.toml:\n[mcp_servers.{server_name}]\nstartup_timeout_sec = XX" + ) + } else { + format!("MCP client for `{server_name}` failed to start: {error:#}") + } +} diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index c9e7a2549879..f321eb76f073 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -12,31 +12,63 @@ use tracing::trace; use tracing::trace_span; use super::McpConnectionSet; +use super::McpServerMetadata; use crate::binding::McpBinding; use crate::binding::PreparedMcpCall; use crate::binding_clients::McpBindingClients; -use crate::codex_apps::prepare_openai_file_params_for_model; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::rmcp_client::CODEX_APPS_REFRESH_DURATION_METRIC; use crate::rmcp_client::MCP_TOOLS_LIST_DURATION_METRIC; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::list_tools_for_client_uncached; +use crate::rmcp_client::prepare_codex_apps_tools_for_model; use crate::runtime::emit_duration; use crate::tools::ToolInfo; use crate::tools::filter_tools; use crate::tools::normalize_tools_for_model_with_prefix; +const MCP_UI_META_KEY: &str = "ui"; +const MCP_UI_VISIBILITY_META_KEY: &str = "visibility"; +const MCP_UI_MODEL_VISIBILITY: &str = "model"; + +/// Returns whether a tool may be included in model-facing tool declarations. +/// +/// Tools without visibility metadata remain visible. Tools with visibility +/// metadata are hidden unless they explicitly include `model`. +/// +/// +pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { + let Some(visibility) = tool + .tool + .meta + .as_deref() + .and_then(|meta| meta.get(MCP_UI_META_KEY)) + .and_then(serde_json::Value::as_object) + .and_then(|ui| ui.get(MCP_UI_VISIBILITY_META_KEY)) + .and_then(serde_json::Value::as_array) + else { + return true; + }; + visibility + .iter() + .any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY)) +} + impl McpConnectionSet { /// Returns all tools with model-visible names normalized. - #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] + #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.servers.len()))] pub async fn list_all_tools(&self) -> Vec { let mut tools = Vec::new(); let mut available_server_count = 0; let mut unavailable_server_count = 0; - for (server_name, managed_client) in &self.clients { - managed_client.reconnect_failed_startup().await; - let has_cached_tools = managed_client.has_cached_tools(); - let startup_complete = managed_client.startup_complete.load(Ordering::Acquire); + for (server_name, view) in &self.servers { + view.connection.client.reconnect_failed_startup().await; + let has_cached_tools = view.connection.client.has_cached_tools(); + let startup_complete = view + .connection + .client + .startup_complete + .load(Ordering::Acquire); let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { self.codex_apps_tools_override.read().await.clone() } else { @@ -44,8 +76,14 @@ impl McpConnectionSet { }; let Some(server_tools) = async { match catalog_override { - Some(tools) => Some(managed_client.prepare_tools(tools)), - None => managed_client.listed_tools().await, + Some(tools) => { + let tools = filter_tools(tools, &view.tool_filter); + Some(prepare_codex_apps_tools_for_model( + tools, + &self.tool_plugin_provenance, + )) + } + None => view.listed_tools(&self.tool_plugin_provenance).await, } } .instrument(trace_span!( @@ -69,7 +107,7 @@ impl McpConnectionSet { tools.extend( server_tools .into_iter() - .map(|tool| self.with_server_metadata(tool)), + .map(|tool| Self::with_server_metadata(tool, &view.metadata)), ); } let tools = normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names); @@ -94,27 +132,41 @@ impl McpConnectionSet { let revision = self.tool_catalog_revision.read().await; let mut listed_tools = Vec::new(); let mut clients = std::collections::HashMap::new(); - for (server_name, managed_client) in &self.clients { + for (server_name, view) in &self.servers { + if !view + .connection + .client + .startup_complete + .load(Ordering::Acquire) + { + let _ = view.connection.client.client().await; + } + view.connection.client.reconnect_failed_startup().await; + let Ok(mut client) = view.connection.client.client().await else { + trace!(server_name = %server_name, "omitting MCP server without an exact ready client"); + continue; + }; + client.tool_timeout = view.tool_timeout; let catalog_override = if server_name == CODEX_APPS_MCP_SERVER_NAME { self.codex_apps_tools_override.read().await.clone() } else { None }; - let Some((client, server_tools)) = managed_client - .capture_ready_client_and_tools(catalog_override) - .await - else { - trace!( - server_name = %server_name, - "omitting MCP server without an exact ready client" - ); - continue; + let server_tools = catalog_override.unwrap_or_else(|| client.tools.clone()); + let server_tools = filter_tools(server_tools, &view.tool_filter); + let server_tools = if server_name == CODEX_APPS_MCP_SERVER_NAME { + prepare_codex_apps_tools_for_model(server_tools, &self.tool_plugin_provenance) + } else { + crate::rmcp_client::prepare_regular_mcp_tools_for_model( + server_tools, + &self.tool_plugin_provenance, + ) }; - clients.insert(server_name.clone(), client); + clients.insert(server_name.clone(), Arc::new(client)); listed_tools.extend( server_tools .into_iter() - .map(|tool| self.with_server_metadata(tool)), + .map(|tool| Self::with_server_metadata(tool, &view.metadata)), ); } let clients = Arc::new(McpBindingClients::new(clients)); @@ -165,6 +217,7 @@ impl McpConnectionSet { tool_catalog_revision: u64, ) -> Option { let server_name = &tool_info.server_name; + let view = self.servers.get(server_name)?; Some(PreparedMcpCall::new( Arc::clone(self), client, @@ -172,7 +225,7 @@ impl McpConnectionSet { tool_catalog_revision, Arc::clone(&self.tool_catalog_revision), tool_info.clone(), - self.server_metadata.get(server_name)?.clone(), + view.metadata.clone(), self.plugin_id_for_mcp_server_name(server_name) .map(str::to_string), self.is_selected_plugin_mcp_server(server_name), @@ -184,13 +237,15 @@ impl McpConnectionSet { clippy::await_holding_invalid_type, reason = "catalog publication must remain serialized with captured tool calls" )] - pub(crate) async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { + pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { let _refresh = self.codex_apps_refresh_lock.lock().await; let refresh_start = Instant::now(); - let managed_client = self - .clients + let view = self + .servers .get(CODEX_APPS_MCP_SERVER_NAME) - .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))? + .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))?; + let managed_client = view + .connection .client() .await .context("failed to get client")?; @@ -208,7 +263,7 @@ impl McpConnectionSet { /*is_codex_apps_mcp_server*/ true, /*codex_apps_refresh_trigger*/ "explicit", &managed_client.client, - managed_client.tool_timeout, + view.tool_timeout, managed_client.server_instructions.as_deref(), ) .await @@ -237,12 +292,12 @@ impl McpConnectionSet { list_start.elapsed(), &[("cache", "miss")], ); - let tools = filter_tools(tools, &managed_client.tool_filter) - .into_iter() - .map(|mut tool| { - prepare_openai_file_params_for_model(&mut tool); - self.with_server_metadata(tool) - }); + let tools = prepare_codex_apps_tools_for_model( + filter_tools(tools, &view.tool_filter), + &self.tool_plugin_provenance, + ) + .into_iter() + .map(|tool| Self::with_server_metadata(tool, &view.metadata)); let tools = normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names); emit_duration( CODEX_APPS_REFRESH_DURATION_METRIC, @@ -252,13 +307,7 @@ impl McpConnectionSet { Ok(tools) } - fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { - let Some(metadata) = self.server_metadata.get(&tool.server_name) else { - tool.supports_parallel_tool_calls = false; - tool.server_origin = None; - return tool; - }; - + fn with_server_metadata(mut tool: ToolInfo, metadata: &McpServerMetadata) -> ToolInfo { tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; tool.server_origin = metadata .origin diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 15545d266ed1..fdeacb2d6cf4 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -10,9 +10,12 @@ use crate::rmcp_client::CodexAppsStartupReconnect; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::ManagedClientFuture; use crate::rmcp_client::StartupOutcomeError; +use crate::rmcp_client::list_tools_for_client_uncached; +use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; use crate::server::McpServerMetadata; use crate::server::McpServerOrigin; +use crate::tool_catalog_cache::McpToolCatalogCache; use crate::tools::ToolFilter; use crate::tools::ToolInfo; use crate::tools::filter_tools; @@ -29,6 +32,7 @@ use codex_connectors::ConnectorRuntimeContextKey; use codex_connectors::ConnectorRuntimeFetchSource; use codex_connectors::ConnectorRuntimeManager; use codex_exec_server::EnvironmentManager; +use codex_login::CodexAuth; use codex_protocol::ToolName; use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; @@ -40,6 +44,7 @@ use futures::FutureExt; use futures::future::BoxFuture; use pretty_assertions::assert_eq; use rmcp::ErrorData as McpError; +use rmcp::RoleServer; use rmcp::ServerHandler; use rmcp::ServiceExt; use rmcp::model::ClientCapabilities; @@ -56,14 +61,18 @@ use rmcp::model::ProtocolVersion; use rmcp::model::ServerCapabilities; use rmcp::model::ServerInfo; use rmcp::model::Tool; +use rmcp::service::RequestContext; use std::collections::HashMap; use std::collections::HashSet; use std::io; +use std::path::PathBuf; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use tempfile::tempdir; use tokio::io::DuplexStream; use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { @@ -188,6 +197,56 @@ impl InProcessTransportFactory for RefreshTestTransportFactory { } } +#[derive(Clone)] +struct MutableToolsServer { + tools: Arc>>, + block_tool_listing: Arc, +} + +impl ServerHandler for MutableToolsServer { + fn get_info(&self) -> ServerInfo { + ServerInfo::new(ServerCapabilities::builder().enable_tools().build()) + } + + async fn list_tools( + &self, + _request: Option, + _context: RequestContext, + ) -> Result { + if self.block_tool_listing.load(Ordering::Acquire) { + std::future::pending::<()>().await; + } + Ok(ListToolsResult { + tools: self.tools.read().await.clone(), + ..Default::default() + }) + } +} + +struct MutableToolsTransportFactory { + server: MutableToolsServer, +} + +impl InProcessTransportFactory for MutableToolsTransportFactory { + fn open(&self) -> BoxFuture<'static, io::Result> { + let server = self.server.clone(); + async move { + let (client_stream, server_stream) = tokio::io::duplex(4096); + tokio::spawn(async move { + server + .serve(server_stream) + .await + .expect("serve mutable MCP tools") + .waiting() + .await + .expect("mutable MCP tools server completes"); + }); + Ok(client_stream) + } + .boxed() + } +} + async fn create_test_managed_client(tools: Vec) -> ManagedClient { ManagedClient { client: Arc::new( @@ -197,7 +256,6 @@ async fn create_test_managed_client(tools: Vec) -> ManagedClient { ), server_info: create_test_server_info("Ready"), tools, - tool_filter: ToolFilter::default(), tool_timeout: None, server_instructions: None, server_supports_sandbox_state_meta_capability: false, @@ -216,10 +274,8 @@ async fn create_ready_async_managed_client(tools: Vec) -> AsyncManaged cached_server_info: None, codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), } } @@ -255,7 +311,6 @@ async fn create_test_manager_with_ready_apps_client( client, server_info: create_test_server_info("Codex Apps"), tools: vec![tool], - tool_filter: ToolFilter::default(), tool_timeout: Some(Duration::from_secs(5)), server_instructions: None, server_supports_sandbox_state_meta_capability: false, @@ -268,7 +323,7 @@ async fn create_test_manager_with_ready_apps_client( &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: futures::future::ready::>(Ok( @@ -280,15 +335,13 @@ async fn create_test_manager_with_ready_apps_client( cached_server_info: Some(create_test_server_info("Codex Apps")), codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); - manager.server_metadata.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), + manager.set_test_server_metadata( + CODEX_APPS_MCP_SERVER_NAME, McpServerMetadata { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), pollutes_memory: false, @@ -325,7 +378,7 @@ fn create_test_manager_with_failed_apps_startup( &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client, @@ -333,10 +386,8 @@ fn create_test_manager_with_failed_apps_startup( cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), startup_reconnect: Some(Arc::new(CodexAppsStartupReconnect::new(reconnect_factory))), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -477,6 +528,60 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel ); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn concurrent_authority_updates_never_auto_approve_mixed_policy() { + let manager = ElicitationRequestManager::new( + AskForApproval::Never, + PermissionProfile::default(), + /*reviewer*/ None, + /*lifecycle*/ None, + ElicitationRequestRouter::default(), + ); + let updating_manager = manager.clone(); + let updater = tokio::spawn(async move { + for _ in 0..1_000 { + assert!(updating_manager.update( + AskForApproval::OnRequest, + PermissionProfile::Disabled, + /*reviewer*/ None, + /*lifecycle*/ None, + )); + assert!(updating_manager.update( + AskForApproval::Never, + PermissionProfile::default(), + /*reviewer*/ None, + /*lifecycle*/ None, + )); + } + }); + let sender = manager.make_sender("server".to_string(), /*tx_event*/ None); + let elicitation = codex_rmcp_client::Elicitation::Mcp( + CreateElicitationRequestParams::FormElicitationParams { + meta: None, + message: "Confirm?".to_string(), + requested_schema: rmcp::model::ElicitationSchema::builder() + .build() + .expect("schema should build"), + }, + ); + + for _ in 0..1_000 { + let response = sender(NumberOrString::Number(1), elicitation.clone()) + .await + .expect("elicitation should resolve"); + assert_eq!( + response, + ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + } + ); + } + + updater.await.expect("authority updates should finish"); +} + #[tokio::test] async fn shared_elicitation_router_targets_the_exact_pending_request() { struct Registration(Arc); @@ -867,7 +972,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -875,10 +980,8 @@ async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -919,7 +1022,6 @@ async fn capture_binding_uses_the_ready_clients_own_tools() { enabled: None, disabled: HashSet::from(["client_local_blocked".to_string()]), }; - ready_client.tool_filter = tool_filter.clone(); ready_client.codex_apps_tools_cache_context = Some(cache_context.clone()); let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let permission_profile = Constrained::allow_any(PermissionProfile::default()); @@ -928,7 +1030,7 @@ async fn capture_binding_uses_the_ready_clients_own_tools() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: futures::future::ready(Ok(ready_client)).boxed().shared(), @@ -936,15 +1038,18 @@ async fn capture_binding_uses_the_ready_clients_own_tools() { cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter, startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); - manager.server_metadata.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), + manager + .servers + .get_mut(CODEX_APPS_MCP_SERVER_NAME) + .expect("test server exists") + .tool_filter = tool_filter; + manager.set_test_server_metadata( + CODEX_APPS_MCP_SERVER_NAME, McpServerMetadata { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), pollutes_memory: false, @@ -1184,7 +1289,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { /*prefix_mcp_tool_names*/ true, ); let server_info = create_test_server_info("Codex Apps"); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -1192,10 +1297,8 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { cached_server_info: Some(server_info.clone()), codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -1223,7 +1326,7 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { &permission_profile, /*prefix_mcp_tool_names*/ false, ); - manager.clients.insert("rmcp".to_string(), managed_client); + manager.insert_test_client("rmcp", managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -1281,7 +1384,7 @@ async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -1289,15 +1392,13 @@ async fn capture_binding_waits_for_fresh_startup_even_with_cached_tools() { cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete, startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); - manager.server_metadata.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), + manager.set_test_server_metadata( + CODEX_APPS_MCP_SERVER_NAME, McpServerMetadata { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), pollutes_memory: false, @@ -1336,7 +1437,7 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert("rmcp".to_string(), managed_client); + manager.insert_test_client("rmcp", managed_client); let tools = manager.list_all_tools().await; let tool = tools @@ -1368,7 +1469,7 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -1376,10 +1477,8 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tools() { cached_server_info: None, codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -1424,7 +1523,7 @@ async fn shutdown_cancels_pending_tool_listing() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -1432,10 +1531,8 @@ async fn shutdown_cancels_pending_tool_listing() { cached_server_info: None, codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token, }, ); @@ -1472,7 +1569,7 @@ async fn shutdown_continues_after_caller_is_aborted() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: blocking_client, @@ -1480,10 +1577,8 @@ async fn shutdown_continues_after_caller_is_aborted() { cached_server_info: None, codex_apps_tools_cache_context: None, tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -1526,7 +1621,7 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -1534,10 +1629,8 @@ async fn list_all_tools_does_not_block_when_shared_codex_apps_cache_is_empty() { cached_server_info: None, codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(false)), startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -1580,7 +1673,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() /*prefix_mcp_tool_names*/ true, ); let startup_complete = Arc::new(std::sync::atomic::AtomicBool::new(true)); - manager.clients.insert( + manager.insert_test_client( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: failed_client, @@ -1588,10 +1681,8 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails() cached_server_info: Some(server_info.clone()), codex_apps_tools_cache_context: Some(cache_context), tool_catalog_cache_context: None, - tool_filter: ToolFilter::default(), startup_complete, startup_reconnect: None, - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), cancel_token: CancellationToken::new(), }, ); @@ -1638,17 +1729,18 @@ async fn list_all_tools_reconnects_failed_codex_apps_startup_and_reuses_client() .shared() }); let mut manager = create_test_manager_with_failed_apps_startup(Vec::new(), reconnect_factory); - manager.server_metadata.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), - McpServerMetadata { - environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), - pollutes_memory: false, - origin: None, - supports_parallel_tool_calls: false, - default_tools_approval_mode: None, - tool_approval_modes: HashMap::new(), - }, - ); + manager + .servers + .get_mut(CODEX_APPS_MCP_SERVER_NAME) + .expect("test server exists") + .metadata = McpServerMetadata { + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + pollutes_memory: false, + origin: None, + supports_parallel_tool_calls: false, + default_tools_approval_mode: None, + tool_approval_modes: HashMap::new(), + }; let manager = Arc::new(manager); let reconnect_finished_wait = reconnect_finished.notified(); @@ -1829,8 +1921,8 @@ async fn tool_lists_do_not_block_and_share_codex_apps_startup_reconnect() { )], reconnect_factory, ); - manager.server_metadata.insert( - CODEX_APPS_MCP_SERVER_NAME.to_string(), + manager.set_test_server_metadata( + CODEX_APPS_MCP_SERVER_NAME, McpServerMetadata { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), pollutes_memory: false, @@ -1912,8 +2004,9 @@ async fn list_all_tools_adds_server_metadata_to_tools() { &permission_profile, /*prefix_mcp_tool_names*/ true, ); - manager.server_metadata.insert( - server_name.to_string(), + manager.insert_test_client(server_name, managed_client); + manager.set_test_server_metadata( + server_name, McpServerMetadata { environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), pollutes_memory: true, @@ -1925,9 +2018,6 @@ async fn list_all_tools_adds_server_metadata_to_tools() { tool_approval_modes: HashMap::new(), }, ); - manager - .clients - .insert(server_name.to_string(), managed_client); let tools = manager.list_all_tools().await; assert_eq!(tools.len(), 1); @@ -1964,7 +2054,6 @@ fn server_metadata_preserves_tool_approval_policy() { #[tokio::test] async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { - let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); let codex_home = tempdir().expect("tempdir"); let mcp_servers = HashMap::from([ ( @@ -2024,52 +2113,51 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { let cancel_token = CancellationToken::new(); let manager = McpConnectionSet::new( - &mcp_servers, - OAuthCredentialsStoreMode::default(), - AuthKeyringBackendKind::default(), - &approval_policy, - String::new(), - /*tx_event*/ None, - cancel_token.clone(), - PermissionProfile::default(), - McpRuntimeContext::new( - Arc::new(EnvironmentManager::without_environments()), - PathBuf::from("/tmp"), - ), - codex_home.path().to_path_buf(), - ConnectorRuntimeManager::::default(), - McpToolCatalogCache::default(), - ConnectorRuntimeContextKey::personal( - /*account_id*/ None, /*chatgpt_user_id*/ None, - ), - /*prefix_mcp_tool_names*/ true, - ElicitationCapability::default(), - /*supports_openai_form_elicitation*/ false, - ToolPluginProvenance::default(), - /*auth*/ None, - /*codex_apps_auth_manager*/ None, - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, + /*previous*/ None, + McpPublicationGate::already_published(), + McpRuntimeInput { + config: Arc::new(crate::mcp::tests::test_mcp_config( + codex_home.path().to_path_buf(), + )), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers, + submit_id: String::new(), + tx_event: None, + startup_cancellation_token: cancel_token.clone(), + runtime_context: McpRuntimeContext::new( + Arc::new(EnvironmentManager::without_environments()), + PathBuf::from("/tmp"), + ), + codex_apps_tools_cache: ConnectorRuntimeManager::::default(), + tool_catalog_cache: McpToolCatalogCache::default(), + codex_apps_tools_cache_key: ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + supports_openai_form_elicitation: false, + auth: None, + codex_apps_auth_manager: None, + elicitation_reviewer: None, + elicitation_lifecycle: None, + }, ElicitationRequestRouter::default(), - crate::runtime::McpPublicationGate::already_published(), ) .await; - assert!(manager.clients.contains_key("stdio")); - assert!(manager.clients.contains_key("http")); - assert!(manager.clients["http"].tool_catalog_cache_context.is_none()); + assert!(manager.contains_server("stdio")); + assert!(manager.contains_server("http")); + assert!( + manager + .test_client("http") + .tool_catalog_cache_context + .is_none() + ); assert!( !manager .wait_for_server_ready("stdio", Duration::from_millis(10)) .await ); - let error = match manager - .clients - .get("stdio") - .expect("stdio client") - .client() - .await - { + let error = match manager.test_client("stdio").client().await { Ok(_) => panic!("local stdio MCP startup should fail"), Err(error) => error, }; @@ -2250,3 +2338,515 @@ fn mcp_init_error_display_includes_startup_timeout_hint() { ); } } + +fn reusable_server_config(url: &str) -> McpServerConfig { + McpServerConfig { + auth: Default::default(), + transport: McpServerTransportConfig::StreamableHttp { + url: url.to_string(), + bearer_token_env_var: Some("CODEX_MCP_REUSE_TEST_TOKEN".to_string()), + http_headers: None, + env_http_headers: None, + }, + environment_id: codex_config::DEFAULT_MCP_SERVER_ENVIRONMENT_ID.to_string(), + enabled: true, + required: false, + supports_parallel_tool_calls: false, + disabled_reason: None, + startup_timeout_sec: None, + tool_timeout_sec: None, + default_tools_approval_mode: None, + enabled_tools: None, + disabled_tools: None, + scopes: None, + oauth: None, + oauth_resource: None, + tools: HashMap::new(), + } +} + +fn reusable_server_runtime_context() -> McpRuntimeContext { + McpRuntimeContext::new( + Arc::new(EnvironmentManager::without_environments()), + PathBuf::from("/tmp"), + ) +} + +fn reusable_server_identity( + config: &McpServerConfig, + runtime_context: &McpRuntimeContext, +) -> McpServerConnectionIdentity { + let server = EffectiveMcpServer::configured(config.clone()); + McpServerConnectionIdentity::new( + "docs", + &server, + OAuthCredentialsStoreMode::default(), + AuthKeyringBackendKind::default(), + &Ok(None), + runtime_context, + /*runtime_auth_provider*/ None, + /*auth*/ None, + /*codex_apps_cache_identity*/ None, + ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, + ) +} + +async fn manager_with_reusable_ready_server( + config: &McpServerConfig, + runtime_context: &McpRuntimeContext, + tools: Vec, +) -> McpConnectionSet { + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut manager = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + let server = EffectiveMcpServer::configured(config.clone()); + manager.servers.insert( + "docs".to_string(), + McpServerView { + connection: Arc::new(McpServerConnection { + identity: Some(reusable_server_identity(config, runtime_context)), + client: create_ready_async_managed_client(tools).await, + }), + metadata: McpServerMetadata::from(&server), + tool_filter: ToolFilter::from_config(config), + tool_timeout: Some(config.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT)), + }, + ); + manager +} + +async fn reconcile_reusable_server( + previous: &McpConnectionSet, + config: McpServerConfig, + runtime_context: McpRuntimeContext, +) -> McpConnectionSet { + let (tx_event, _rx_event) = async_channel::unbounded(); + let codex_home = tempdir().expect("tempdir"); + McpConnectionSet::new( + Some(previous), + McpPublicationGate::already_published(), + McpRuntimeInput { + config: Arc::new(crate::mcp::tests::test_mcp_config( + codex_home.path().to_path_buf(), + )), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers: HashMap::from([( + "docs".to_string(), + EffectiveMcpServer::configured(config), + )]), + submit_id: "refresh".to_string(), + tx_event: Some(tx_event), + startup_cancellation_token: CancellationToken::new(), + runtime_context, + codex_apps_tools_cache: ConnectorRuntimeManager::default(), + tool_catalog_cache: McpToolCatalogCache::default(), + codex_apps_tools_cache_key: ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ), + supports_openai_form_elicitation: false, + auth: None, + codex_apps_auth_manager: None, + elicitation_reviewer: None, + elicitation_lifecycle: None, + }, + ElicitationRequestRouter::default(), + ) + .await +} + +#[tokio::test] +async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> anyhow::Result<()> { + let tools = Arc::new(tokio::sync::RwLock::new(vec![Tool::new( + "old_search", + "old search", + Arc::new(JsonObject::default()), + )])); + let block_tool_listing = Arc::new(AtomicBool::new(false)); + let client = Arc::new( + RmcpClient::new_in_process_client(Arc::new(MutableToolsTransportFactory { + server: MutableToolsServer { + tools: Arc::clone(&tools), + block_tool_listing: Arc::clone(&block_tool_listing), + }, + })) + .await?, + ); + let initialize = client + .initialize( + InitializeRequestParams::new( + ClientCapabilities::default(), + Implementation::new("codex-test", "0.0.0-test"), + ) + .with_protocol_version(ProtocolVersion::V_2025_06_18), + /*timeout*/ None, + Box::new(|_, _| { + async { + Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + }) + } + .boxed() + }), + ) + .await?; + let initial_tools = list_tools_for_client_uncached( + "docs", + /*is_codex_apps_mcp_server*/ false, + /*codex_apps_refresh_trigger*/ "test", + &client, + /*timeout*/ None, + initialize.instructions.as_deref(), + ) + .await?; + let managed_client = ManagedClient { + client, + server_info: create_test_server_info("Mutable tools"), + tools: initial_tools, + tool_timeout: None, + server_instructions: initialize.instructions, + server_supports_sandbox_state_meta_capability: false, + codex_apps_tools_cache_context: None, + }; + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + let permission_profile = Constrained::allow_any(PermissionProfile::default()); + let mut previous = McpConnectionSet::new_uninitialized( + &approval_policy, + &permission_profile, + /*prefix_mcp_tool_names*/ true, + ); + let server = EffectiveMcpServer::configured(config.clone()); + previous.servers.insert( + "docs".to_string(), + McpServerView { + connection: Arc::new(McpServerConnection { + identity: Some(reusable_server_identity(&config, &runtime_context)), + client: AsyncManagedClient { + client: futures::future::ready(Ok(managed_client)).boxed().shared(), + is_codex_apps_mcp_server: false, + cached_server_info: None, + codex_apps_tools_cache_context: None, + tool_catalog_cache_context: None, + startup_complete: Arc::new(std::sync::atomic::AtomicBool::new(true)), + startup_reconnect: None, + cancel_token: CancellationToken::new(), + }, + }), + metadata: McpServerMetadata::from(&server), + tool_filter: ToolFilter::from_config(&config), + tool_timeout: Some(config.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT)), + }, + ); + let previous = Arc::new(previous); + let old_step = capture_binding(&previous).await; + *tools.write().await = vec![Tool::new( + "new_search", + "new search", + Arc::new(JsonObject::default()), + )]; + block_tool_listing.store(true, Ordering::Release); + + let reconciled = Arc::new( + tokio::time::timeout( + Duration::from_secs(1), + reconcile_reusable_server(&previous, config, runtime_context), + ) + .await + .expect("connection reuse must not wait for a tool-list request"), + ); + let new_step = capture_binding(&reconciled).await; + + assert!(previous.shares_test_connection_with(&reconciled, "docs")); + assert_eq!( + old_step + .tools() + .iter() + .map(|tool| tool.tool.name.to_string()) + .collect::>(), + vec!["old_search".to_string()] + ); + assert_eq!( + new_step + .tools() + .iter() + .map(|tool| tool.tool.name.to_string()) + .collect::>(), + vec!["old_search".to_string()] + ); + Ok(()) +} + +#[tokio::test] +async fn reconciliation_reuses_an_unchanged_ready_server() { + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let previous = manager_with_reusable_ready_server( + &config, + &runtime_context, + vec![create_test_tool("docs", "search")], + ) + .await; + + let reconciled = reconcile_reusable_server(&previous, config, runtime_context.clone()).await; + + assert!(previous.shares_test_connection_with(&reconciled, "docs")); + assert_eq!( + model_tool_names(&reconciled.list_all_tools().await), + HashSet::from([ToolName::namespaced("mcp__docs", "search")]) + ); +} + +#[tokio::test] +async fn reconciliation_updates_elicitation_policy_without_restarting_ready_server() { + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let previous = manager_with_reusable_ready_server( + &config, + &runtime_context, + vec![create_test_tool("docs", "search")], + ) + .await; + { + let mut authority = previous + .elicitation_requests + .authority + .lock() + .expect("elicitation authority lock"); + authority.approval_policy = AskForApproval::Never; + authority.permission_profile = PermissionProfile::Disabled; + } + + let reconciled = reconcile_reusable_server(&previous, config, runtime_context).await; + + assert!(previous.shares_test_connection_with(&reconciled, "docs")); + let authority = reconciled + .elicitation_requests + .authority + .lock() + .expect("elicitation authority lock"); + assert_eq!(authority.approval_policy, AskForApproval::OnRequest); + assert_eq!(authority.permission_profile, PermissionProfile::default()); +} + +#[tokio::test] +async fn reconciliation_reuses_ready_server_when_startup_timeout_changes() { + let runtime_context = reusable_server_runtime_context(); + let mut config = reusable_server_config("http://127.0.0.1:1"); + let previous = manager_with_reusable_ready_server( + &config, + &runtime_context, + vec![create_test_tool("docs", "search")], + ) + .await; + config.startup_timeout_sec = Some(Duration::from_secs(30)); + + let reconciled = reconcile_reusable_server(&previous, config, runtime_context).await; + + assert!(previous.shares_test_connection_with(&reconciled, "docs")); +} + +#[tokio::test] +async fn reconciliation_reconnects_when_connection_identity_changes() { + let runtime_context = reusable_server_runtime_context(); + let previous_config = reusable_server_config("http://127.0.0.1:1"); + let previous = manager_with_reusable_ready_server( + &previous_config, + &runtime_context, + vec![create_test_tool("docs", "search")], + ) + .await; + + let reconciled = reconcile_reusable_server( + &previous, + reusable_server_config("http://127.0.0.1:2"), + runtime_context, + ) + .await; + + assert!(!previous.shares_test_connection_with(&reconciled, "docs")); +} + +#[tokio::test] +async fn connection_identity_distinguishes_accounts_with_the_same_token() -> anyhow::Result<()> { + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let server = EffectiveMcpServer::configured(config); + let access_token = "header.e30.same"; + let previous_auth = CodexAuth::from_external_chatgpt_tokens( + access_token, + "account-a", + /*chatgpt_plan_type*/ None, + )?; + let changed_auth = CodexAuth::from_external_chatgpt_tokens( + access_token, + "account-b", + /*chatgpt_plan_type*/ None, + )?; + let connection_identity = |auth: &CodexAuth| { + let provider = codex_model_provider::auth_provider_from_auth(auth); + McpServerConnectionIdentity::new( + "docs", + &server, + OAuthCredentialsStoreMode::default(), + AuthKeyringBackendKind::default(), + &Ok(None), + &runtime_context, + Some(&provider), + Some(auth), + /*codex_apps_cache_identity*/ None, + ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, + ) + }; + + assert_eq!(previous_auth, changed_auth); + assert_eq!(previous_auth.get_token()?, changed_auth.get_token()?); + assert!( + !connection_identity(&previous_auth) + .has_same_connection_config(&connection_identity(&changed_auth)) + ); + Ok(()) +} + +#[tokio::test] +async fn connection_identity_distinguishes_agent_account_runtime_and_task() -> anyhow::Result<()> { + let runtime_context = reusable_server_runtime_context(); + let config = reusable_server_config("http://127.0.0.1:1"); + let server = EffectiveMcpServer::configured(config); + let record = codex_login::auth::AgentIdentityAuthRecord { + agent_runtime_id: "agent-a".to_string(), + agent_private_key: "MC4CAQAwBQYDK2VwBCIEIJ7kFBaOujmoz1gvBNEC+BeM2IX87FFB0xmISOZ/XO0c" + .to_string(), + account_id: "account-a".to_string(), + chatgpt_user_id: "user-a".to_string(), + email: Some("agent@example.com".to_string()), + plan_type: codex_protocol::account::PlanType::Plus, + chatgpt_account_is_fedramp: false, + task_id: Some("task-a".to_string()), + }; + let auth_route_config = codex_login::test_support::transport_default_auth_route_config(); + let previous_auth = CodexAuth::AgentIdentity( + codex_login::auth::AgentIdentityAuth::from_record( + record.clone(), + "https://auth.openai.com/api/accounts", + &auth_route_config, + ) + .await?, + ); + let connection_identity = |auth: &CodexAuth| { + let provider = codex_model_provider::auth_provider_from_auth(auth); + McpServerConnectionIdentity::new( + CODEX_APPS_MCP_SERVER_NAME, + &server, + OAuthCredentialsStoreMode::default(), + AuthKeyringBackendKind::default(), + &Ok(None), + &runtime_context, + Some(&provider), + Some(auth), + /*codex_apps_cache_identity*/ None, + ElicitationCapability::default(), + /*supports_openai_form_elicitation*/ false, + ) + }; + let previous_identity = connection_identity(&previous_auth); + + for changed_record in [ + codex_login::auth::AgentIdentityAuthRecord { + account_id: "account-b".to_string(), + ..record.clone() + }, + codex_login::auth::AgentIdentityAuthRecord { + chatgpt_user_id: "user-b".to_string(), + ..record.clone() + }, + codex_login::auth::AgentIdentityAuthRecord { + chatgpt_account_is_fedramp: true, + ..record.clone() + }, + codex_login::auth::AgentIdentityAuthRecord { + agent_runtime_id: "agent-b".to_string(), + ..record.clone() + }, + codex_login::auth::AgentIdentityAuthRecord { + task_id: Some("task-b".to_string()), + ..record.clone() + }, + ] { + let changed_auth = CodexAuth::AgentIdentity( + codex_login::auth::AgentIdentityAuth::from_record( + changed_record, + "https://auth.openai.com/api/accounts", + &auth_route_config, + ) + .await?, + ); + assert_eq!(previous_auth, changed_auth); + assert!(!previous_identity.has_same_connection_config(&connection_identity(&changed_auth))); + } + + Ok(()) +} + +#[tokio::test] +async fn view_only_changes_reuse_connection_and_preserve_the_old_step() { + let runtime_context = reusable_server_runtime_context(); + let mut old_config = reusable_server_config("http://127.0.0.1:1"); + old_config.default_tools_approval_mode = Some(AppToolApproval::Prompt); + let previous = Arc::new( + manager_with_reusable_ready_server( + &old_config, + &runtime_context, + vec![ + create_test_tool("docs", "search"), + create_test_tool("docs", "write"), + ], + ) + .await, + ); + let old_step = capture_binding(&previous).await; + let old_call = old_step + .prepare_call("docs", "search") + .expect("old step should prepare search"); + + let mut new_config = old_config; + new_config.enabled_tools = Some(vec!["search".to_string()]); + new_config.default_tools_approval_mode = Some(AppToolApproval::Approve); + let reconciled = + Arc::new(reconcile_reusable_server(previous.as_ref(), new_config, runtime_context).await); + assert!(previous.shares_test_connection_with(&reconciled, "docs")); + + let new_step = capture_binding(&reconciled).await; + let new_call = new_step + .prepare_call("docs", "search") + .expect("new step should prepare search"); + drop(previous); + + assert_eq!( + old_step + .tools() + .iter() + .map(|tool| tool.tool.name.to_string()) + .collect::>(), + HashSet::from(["search".to_string(), "write".to_string()]) + ); + assert_eq!(old_call.tool_approval_mode(), AppToolApproval::Prompt); + assert_eq!( + new_step + .tools() + .iter() + .map(|tool| tool.tool.name.to_string()) + .collect::>(), + vec!["search".to_string()] + ); + assert_eq!(new_call.tool_approval_mode(), AppToolApproval::Approve); +} diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index e2d9ee1eb25e..34d7b314df67 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -7,6 +7,7 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Mutex as StdMutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; @@ -116,14 +117,19 @@ impl ElicitationRequestRouter { } #[derive(Clone)] -pub(crate) struct ElicitationRequestManager { - router: ElicitationRequestRouter, - approval_policy: AskForApproval, - permission_profile: PermissionProfile, +pub(crate) struct ElicitationAuthority { + pub(crate) approval_policy: AskForApproval, + pub(crate) permission_profile: PermissionProfile, reviewer: Option, lifecycle: Option, } +#[derive(Clone)] +pub(crate) struct ElicitationRequestManager { + router: ElicitationRequestRouter, + pub(crate) authority: Arc>, +} + impl ElicitationRequestManager { pub(crate) fn new( approval_policy: AskForApproval, @@ -134,11 +140,32 @@ impl ElicitationRequestManager { ) -> Self { Self { router, + authority: Arc::new(StdMutex::new(ElicitationAuthority { + approval_policy, + permission_profile, + reviewer, + lifecycle, + })), + } + } + + pub(crate) fn update( + &self, + approval_policy: AskForApproval, + permission_profile: PermissionProfile, + reviewer: Option, + lifecycle: Option, + ) -> bool { + let Ok(mut authority) = self.authority.lock() else { + return false; + }; + *authority = ElicitationAuthority { approval_policy, permission_profile, reviewer, lifecycle, - } + }; + true } pub(crate) fn make_sender( @@ -147,18 +174,12 @@ impl ElicitationRequestManager { tx_event: Option>, ) -> SendElicitation { let router = self.router.clone(); - let approval_policy = self.approval_policy; - let permission_profile = self.permission_profile.clone(); - let reviewer = self.reviewer.clone(); - let lifecycle = self.lifecycle.clone(); + let authority = self.authority.clone(); Box::new(move |id, elicitation| { let router = router.clone(); let tx_event = tx_event.clone(); let server_name = server_name.clone(); - let approval_policy = approval_policy; - let permission_profile = permission_profile.clone(); - let reviewer = reviewer.clone(); - let lifecycle = lifecycle.clone(); + let authority = authority.clone(); async move { if router.auto_deny() { return Ok(ElicitationResponse { @@ -168,6 +189,19 @@ impl ElicitationRequestManager { }); } + let Ok(authority) = authority.lock().map(|authority| authority.clone()) else { + return Ok(ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + }); + }; + let ElicitationAuthority { + approval_policy, + permission_profile, + reviewer, + lifecycle, + } = authority; if mcp_permission_prompt_is_auto_approved( approval_policy, &permission_profile, @@ -189,7 +223,7 @@ impl ElicitationRequestManager { }); } - if let Some(reviewer) = reviewer.as_ref() { + if let Some(reviewer) = reviewer { let request = ElicitationReviewRequest { server_name: server_name.clone(), request_id: id.clone(), diff --git a/codex-rs/codex-mcp/src/mcp/auth.rs b/codex-rs/codex-mcp/src/mcp/auth.rs index fcdaefe3a9e0..8b561edae1f1 100644 --- a/codex-rs/codex-mcp/src/mcp/auth.rs +++ b/codex-rs/codex-mcp/src/mcp/auth.rs @@ -192,46 +192,38 @@ where { let futures = servers.into_iter().map(|(name, server)| { let name = name.clone(); - let config = server.configured_config().cloned(); + let config = server.config().clone(); let runtime_context = runtime_context.clone(); - let has_runtime_auth = config - .as_ref() - .is_some_and(|config| matches!(&config.auth, McpServerAuth::ChatGpt)) + let has_runtime_auth = matches!(&config.auth, McpServerAuth::ChatGpt) && auth.is_some_and(CodexAuth::uses_codex_backend) - && config.as_ref().is_some_and(|config| { - matches!( - &config.transport, - McpServerTransportConfig::StreamableHttp { - bearer_token_env_var: None, - .. - } - ) - }); + && matches!( + &config.transport, + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var: None, + .. + } + ); async move { - let auth_state = match config.as_ref() { - Some(config) => { - match compute_auth_status( - &name, - config, - store_mode, - keyring_backend_kind, - has_runtime_auth, - &runtime_context, - ) - .await - { - Ok(status) => status, - Err(error) => { - warn!( - "failed to determine auth status for MCP server `{name}`: {error:?}" - ); - McpAuthState::Unsupported - } - } + let auth_state = match compute_auth_status( + &name, + &config, + store_mode, + keyring_backend_kind, + has_runtime_auth, + &runtime_context, + ) + .await + { + Ok(status) => status, + Err(error) => { + warn!("failed to determine auth status for MCP server `{name}`: {error:?}"); + McpAuthState::Unsupported } - None => McpAuthState::Unsupported, }; - let entry = McpAuthStatusEntry { config, auth_state }; + let entry = McpAuthStatusEntry { + config: Some(config), + auth_state, + }; (name, entry) } }); diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 92321ac3278d..50cc6472d5ed 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -17,6 +17,7 @@ use std::collections::HashMap; use std::collections::HashSet; use std::env; use std::path::PathBuf; +use std::sync::Arc; use std::time::Duration; use codex_config::ConfigLayerStack; @@ -49,7 +50,9 @@ use tokio_util::sync::CancellationToken; use crate::ResolvedMcpCatalog; use crate::connection_manager::McpConnectionSet; +use crate::runtime::McpPublicationGate; use crate::runtime::McpRuntimeContext; +use crate::runtime::McpRuntimeInput; use crate::server::EffectiveMcpServer; use crate::tools::ToolInfo; @@ -319,30 +322,30 @@ pub async fn read_mcp_resource( let mut mcp_servers = effective_mcp_servers(config, auth); mcp_servers.retain(|name, _| name == server); let cancel_token = CancellationToken::new(); + let mut runtime_config = config.clone(); + runtime_config.permission_profile = PermissionProfile::default(); let manager = McpConnectionSet::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind, - &config.approval_policy, - String::new(), - /*tx_event*/ None, - cancel_token.clone(), - PermissionProfile::default(), - runtime_context, - config.codex_home.clone(), - codex_apps_tools_cache, - tool_catalog_cache, - connector_runtime_context_key(auth), - config.prefix_mcp_tool_names, - config.client_elicitation_capability.clone(), - /*supports_openai_form_elicitation*/ false, - tool_plugin_provenance(config), - auth, - /*codex_apps_auth_manager*/ None, - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, + /*previous*/ None, + McpPublicationGate::already_published(), + McpRuntimeInput { + config: Arc::new(runtime_config), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers, + submit_id: String::new(), + tx_event: None, + startup_cancellation_token: cancel_token.clone(), + runtime_context, + codex_apps_tools_cache, + tool_catalog_cache, + codex_apps_tools_cache_key: connector_runtime_context_key(auth), + supports_openai_form_elicitation: false, + auth: auth.cloned(), + codex_apps_auth_manager: None, + elicitation_reviewer: None, + elicitation_lifecycle: None, + }, crate::elicitation::ElicitationRequestRouter::default(), - crate::runtime::McpPublicationGate::already_published(), ) .await; @@ -373,7 +376,6 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let mcp_servers = effective_mcp_servers(config, auth); - let tool_plugin_provenance = tool_plugin_provenance(config); if mcp_servers.is_empty() { return McpServerStatusSnapshot { server_infos: HashMap::new(), @@ -397,35 +399,35 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( let server_names = mcp_servers.keys().cloned().collect(); let cancel_token = CancellationToken::new(); - let mcp_connection_set = McpConnectionSet::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind, - &config.approval_policy, - submit_id, - /*tx_event*/ None, - cancel_token.clone(), - PermissionProfile::default(), - runtime_context, - config.codex_home.clone(), - codex_apps_tools_cache, - tool_catalog_cache, - connector_runtime_context_key(auth), - config.prefix_mcp_tool_names, - config.client_elicitation_capability.clone(), - /*supports_openai_form_elicitation*/ false, - tool_plugin_provenance, - auth, - /*codex_apps_auth_manager*/ None, - /*elicitation_reviewer*/ None, - /*elicitation_lifecycle*/ None, + let mut runtime_config = config.clone(); + runtime_config.permission_profile = PermissionProfile::default(); + let mcp_connection_manager = McpConnectionSet::new( + /*previous*/ None, + McpPublicationGate::already_published(), + McpRuntimeInput { + config: Arc::new(runtime_config), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers, + submit_id, + tx_event: None, + startup_cancellation_token: cancel_token.clone(), + runtime_context, + codex_apps_tools_cache, + tool_catalog_cache, + codex_apps_tools_cache_key: connector_runtime_context_key(auth), + supports_openai_form_elicitation: false, + auth: auth.cloned(), + codex_apps_auth_manager: None, + elicitation_reviewer: None, + elicitation_lifecycle: None, + }, crate::elicitation::ElicitationRequestRouter::default(), - crate::runtime::McpPublicationGate::already_published(), ) .await; let snapshot = collect_mcp_server_status_snapshot_from_manager( - &mcp_connection_set, + &mcp_connection_manager, auth_status_entries, server_names, detail, @@ -652,27 +654,27 @@ fn convert_mcp_resource_templates( } async fn collect_mcp_server_status_snapshot_from_manager( - mcp_connection_set: &McpConnectionSet, + mcp_connection_manager: &McpConnectionSet, auth_status_entries: HashMap, server_names: Vec, detail: McpSnapshotDetail, ) -> McpServerStatusSnapshot { let ((server_infos, tools), resources, resource_templates) = tokio::join!( async { - let server_infos = mcp_connection_set.list_available_server_infos().await; - let tools = mcp_connection_set.list_all_tools().await; + let server_infos = mcp_connection_manager.list_available_server_infos().await; + let tools = mcp_connection_manager.list_all_tools().await; (server_infos, tools) }, async { if detail.include_resources() { - mcp_connection_set.list_all_resources(|_| true).await + mcp_connection_manager.list_all_resources(|_| true).await } else { HashMap::new() } }, async { if detail.include_resources() { - mcp_connection_set + mcp_connection_manager .list_all_resource_templates(|_| true) .await } else { diff --git a/codex-rs/codex-mcp/src/mcp/mod_tests.rs b/codex-rs/codex-mcp/src/mcp/mod_tests.rs index 994c6ea96f31..46927f23075e 100644 --- a/codex-rs/codex-mcp/src/mcp/mod_tests.rs +++ b/codex-rs/codex-mcp/src/mcp/mod_tests.rs @@ -440,15 +440,9 @@ async fn effective_mcp_servers_preserve_runtime_servers() { .get(CODEX_APPS_MCP_SERVER_NAME) .expect("codex apps server should exist"); - let sample = sample - .configured_config() - .expect("configured server should retain transport"); - let docs = docs - .configured_config() - .expect("configured server should retain transport"); - let codex_apps = codex_apps - .configured_config() - .expect("codex apps should use configured transport"); + let sample = sample.config(); + let docs = docs.config(); + let codex_apps = codex_apps.config(); match &sample.transport { McpServerTransportConfig::StreamableHttp { url, .. } => { diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index ee6161add859..32b9e067b2c0 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -29,12 +29,9 @@ use crate::openai_docs_source_attribution::maybe_with_openai_docs_source_attribu use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; use crate::server::EffectiveMcpServer; -use crate::server::McpServerLaunch; use crate::tool_catalog_cache::McpToolCatalogCacheContext; use crate::tool_catalog_cache::McpToolCatalogFetchTicket; -use crate::tools::ToolFilter; use crate::tools::ToolInfo; -use crate::tools::filter_tools; use anyhow::Result; use anyhow::anyhow; use async_channel::Sender; @@ -109,7 +106,6 @@ pub(crate) struct ManagedClient { pub(crate) client: Arc, pub(crate) server_info: McpServerInfo, pub(crate) tools: Vec, - pub(crate) tool_filter: ToolFilter, pub(crate) tool_timeout: Option, pub(crate) server_instructions: Option, pub(crate) server_supports_sandbox_state_meta_capability: bool, @@ -129,7 +125,7 @@ impl ManagedClient { total_start.elapsed(), &[("cache", "hit")], ); - return filter_tools(tools, &self.tool_filter); + return tools; } if self.codex_apps_tools_cache_context.is_some() { @@ -140,7 +136,7 @@ impl ManagedClient { ); } - filter_tools(self.tools.clone(), &self.tool_filter) + self.tools.clone() } } @@ -310,13 +306,9 @@ impl ManagedClientStartup { startup_complete, } = self.clone(); let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; - let tool_filter = server - .configured_config() - .map(ToolFilter::from_config) - .unwrap_or_default(); let startup_timeout = server - .configured_config() - .and_then(|config| config.startup_timeout_sec) + .config() + .startup_timeout_sec .unwrap_or(DEFAULT_STARTUP_TIMEOUT); let cancel_token_for_fut = cancel_token; let tool_catalog_fetch_ticket = tool_catalog_cache_context @@ -356,11 +348,6 @@ impl ManagedClientStartup { StartServerTaskParams { is_codex_apps_mcp_server, startup_timeout: Some(startup_timeout), - tool_timeout: server - .configured_config() - .and_then(|config| config.tool_timeout_sec) - .unwrap_or(DEFAULT_TOOL_TIMEOUT), - tool_filter, tx_event, elicitation_requests, codex_apps_tools_cache_context, @@ -404,10 +391,8 @@ pub(crate) struct AsyncManagedClient { pub(crate) cached_server_info: Option, pub(crate) codex_apps_tools_cache_context: Option>, pub(crate) tool_catalog_cache_context: Option, - pub(crate) tool_filter: ToolFilter, pub(crate) startup_complete: Arc, pub(crate) startup_reconnect: Option>, - pub(crate) tool_plugin_provenance: Arc, pub(crate) cancel_token: CancellationToken, } @@ -427,7 +412,6 @@ impl AsyncManagedClient { elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option>, tool_catalog_cache_context: Option, - tool_plugin_provenance: Arc, runtime_context: McpRuntimeContext, resolved_environment: std::result::Result>, String>, runtime_auth_provider: Option, @@ -437,10 +421,6 @@ impl AsyncManagedClient { let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; let reconnect_server_name = server_name.clone(); let reconnect_tx_event = tx_event.clone(); - let tool_filter = server - .configured_config() - .map(ToolFilter::from_config) - .unwrap_or_default(); let cached_server_info = if is_codex_apps_mcp_server { codex_apps_tools_cache_context .as_ref() @@ -484,10 +464,8 @@ impl AsyncManagedClient { cached_server_info, codex_apps_tools_cache_context, tool_catalog_cache_context, - tool_filter, startup_complete, startup_reconnect, - tool_plugin_provenance, cancel_token, } } @@ -503,50 +481,6 @@ impl AsyncManagedClient { self.client.clone().await } - /// Captures the ready client revision that is current now. - /// - /// A recovered Codex Apps connection replaces the failed startup future for - /// future steps, but cannot reroute a call that was already prepared. - fn ready_client_snapshot(&self) -> Option { - if let Some(client) = self - .startup_reconnect - .as_ref() - .and_then(|reconnect| reconnect.current_client()) - { - return Some(futures::future::ready(Ok(client)).boxed().shared()); - } - match self.client.peek() { - Some(Ok(_)) => Some(self.client.clone()), - Some(Err(_)) | None => None, - } - } - - /// Captures one ready client and derives its model-visible tools. - /// - /// A fresh client waits for initial startup even when metadata is cached. - /// After a failed startup, cached tools remain available to metadata-only - /// callers while recovery runs, but a model step still requires one exact - /// ready client. - pub(crate) async fn capture_ready_client_and_tools( - &self, - catalog_override: Option>, - ) -> Option<(Arc, Vec)> { - if !self.startup_complete.load(Ordering::Acquire) { - let _ = self.client().await; - } - self.reconnect_failed_startup().await; - let client = if self.has_cached_tools() { - self.ready_client_snapshot()? - } else { - self.client().await.ok()?; - self.ready_client_snapshot()? - }; - let managed_client = Arc::new(client.await.ok()?); - let tools = catalog_override.unwrap_or_else(|| managed_client.tools.clone()); - let tools = filter_tools(tools, &managed_client.tool_filter); - Some((Arc::clone(&managed_client), self.prepare_tools(tools))) - } - pub(crate) async fn reconnect_failed_startup(&self) { let Some(startup_reconnect) = self.startup_reconnect.as_ref() else { return; @@ -589,20 +523,11 @@ impl AsyncManagedClient { .as_ref() .and_then(McpToolCatalogCacheContext::current_tools) }) - .map(|tools| filter_tools(tools, &self.tool_filter)) - } - - pub(crate) fn prepare_tools(&self, tools: Vec) -> Vec { - if self.is_codex_apps_mcp_server { - prepare_codex_apps_tools_for_model(tools, &self.tool_plugin_provenance) - } else { - prepare_regular_mcp_tools_for_model(tools, &self.tool_plugin_provenance) - } } pub(crate) async fn listed_tools(&self) -> Option> { // Plugin provenance is resolved per-session rather than stored in shared cache payloads. - let tools = if !self.startup_complete.load(Ordering::Acquire) + if !self.startup_complete.load(Ordering::Acquire) && let Some(startup_tools) = self.cached_tools() { Some(startup_tools) @@ -612,8 +537,7 @@ impl AsyncManagedClient { Err(_) if self.is_codex_apps_mcp_server => self.cached_tools(), Err(_) => None, } - }?; - Some(self.prepare_tools(tools)) + } } } @@ -696,7 +620,7 @@ pub(crate) async fn list_tools_for_client_uncached( /// Presents declared Codex Apps file parameters to the model as local-path inputs and adds plugin /// names to each tool. Plugin membership is resolved by connector ID, falling back to the MCP /// server when absent. -fn prepare_codex_apps_tools_for_model( +pub(crate) fn prepare_codex_apps_tools_for_model( mut tools: Vec, tool_plugin_provenance: &ToolPluginProvenance, ) -> Vec { @@ -750,7 +674,7 @@ fn add_plugin_provenance_to_tool(tool: &mut ToolInfo, plugin_names: &[String]) { } /// Adds server-scoped plugin names to regular MCP tools without changing their input schemas. -fn prepare_regular_mcp_tools_for_model( +pub(crate) fn prepare_regular_mcp_tools_for_model( mut tools: Vec, tool_plugin_provenance: &ToolPluginProvenance, ) -> Vec { @@ -902,8 +826,6 @@ async fn start_server_task( let StartServerTaskParams { is_codex_apps_mcp_server, startup_timeout, - tool_timeout, - tool_filter, tx_event, elicitation_requests, codex_apps_tools_cache_context, @@ -984,8 +906,7 @@ async fn start_server_task( client: Arc::clone(&client), server_info, tools: client_tools, - tool_timeout: Some(tool_timeout), - tool_filter, + tool_timeout: None, server_instructions: initialize_result.instructions, server_supports_sandbox_state_meta_capability, codex_apps_tools_cache_context, @@ -1032,8 +953,6 @@ fn mcp_server_info_from_implementation(server_info: Implementation) -> McpServer struct StartServerTaskParams { is_codex_apps_mcp_server: bool, startup_timeout: Option, // TODO: cancel_token should handle this. - tool_timeout: Duration, - tool_filter: ToolFilter, tx_event: Option>, elicitation_requests: ElicitationRequestManager, codex_apps_tools_cache_context: Option>, @@ -1053,9 +972,7 @@ async fn make_rmcp_client( resolved_environment: std::result::Result>, String>, runtime_auth_provider: Option, ) -> Result { - let config = match server.launch() { - McpServerLaunch::Configured(config) => config.as_ref().clone(), - }; + let config = server.config().clone(); let resolved_environment = resolved_environment.map_err(|err| StartupOutcomeError::from(anyhow!(err)))?; let is_local_environment = config.is_local_environment(); diff --git a/codex-rs/codex-mcp/src/runtime.rs b/codex-rs/codex-mcp/src/runtime.rs index 009ad279f378..1f27fd402041 100644 --- a/codex-rs/codex-mcp/src/runtime.rs +++ b/codex-rs/codex-mcp/src/runtime.rs @@ -138,14 +138,21 @@ impl McpRuntime { runtime } - /// Rebuilds configured servers and publishes their immutable runtime snapshot. + /// Reconciles configured servers and publishes their immutable runtime snapshot. pub async fn replace(&self, input: McpRuntimeInput) { + let current = self.current.load_full(); let (publish, publication_gate) = McpPublicationGate::pending(); let config = Arc::clone(&input.config); let plugins_available = input.plugins_available; let ready_selected_capability_roots = input.ready_selected_capability_roots.clone(); let connections = Arc::new( - Self::materialize(input, self.elicitation_router.clone(), publication_gate).await, + McpConnectionSet::new( + Some(current.connections.as_ref()), + publication_gate, + input, + self.elicitation_router.clone(), + ) + .await, ); self.current.store(Arc::new(PublishedMcpRuntime { connections, @@ -250,57 +257,6 @@ impl McpRuntime { pub async fn shutdown(&self) { self.latest_connections().shutdown().await; } - - async fn materialize( - input: McpRuntimeInput, - elicitation_router: ElicitationRequestRouter, - publication_gate: McpPublicationGate, - ) -> McpConnectionSet { - let McpRuntimeInput { - config, - 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, - supports_openai_form_elicitation, - auth, - codex_apps_auth_manager, - elicitation_reviewer, - elicitation_lifecycle, - } = input; - McpConnectionSet::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - config.auth_keyring_backend_kind, - &config.approval_policy, - submit_id, - tx_event, - startup_cancellation_token, - config.permission_profile.clone(), - runtime_context, - config.codex_home.clone(), - codex_apps_tools_cache, - tool_catalog_cache, - codex_apps_tools_cache_key, - config.prefix_mcp_tool_names, - config.client_elicitation_capability.clone(), - supports_openai_form_elicitation, - crate::mcp::tool_plugin_provenance(&config), - auth.as_ref(), - codex_apps_auth_manager, - elicitation_reviewer, - elicitation_lifecycle, - elicitation_router, - publication_gate, - ) - .await - } } #[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)] diff --git a/codex-rs/codex-mcp/src/server.rs b/codex-rs/codex-mcp/src/server.rs index 39287b21361d..0811d8043580 100644 --- a/codex-rs/codex-mcp/src/server.rs +++ b/codex-rs/codex-mcp/src/server.rs @@ -1,51 +1,213 @@ use std::collections::HashMap; +use std::ffi::OsString; +use std::path::PathBuf; +use std::sync::Arc; +use crate::runtime::McpRuntimeContext; +use codex_api::SharedAuthProvider; use codex_config::AppToolApproval; use codex_config::McpServerConfig; use codex_config::McpServerTransportConfig; - -/// The runtime launch strategy for an effective MCP server. -#[derive(Debug, Clone)] -pub(crate) enum McpServerLaunch { - Configured(Box), -} +use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_connectors::ConnectorRuntimeContextKey; +use codex_exec_server::Environment; +use codex_login::CodexAuth; +use codex_rmcp_client::StoredOAuthTokens; +use codex_rmcp_client::stored_oauth_credentials; +use rmcp::model::ElicitationCapability; +use tracing::warn; /// MCP server after runtime additions have been applied. #[derive(Debug, Clone)] pub struct EffectiveMcpServer { - launch: McpServerLaunch, + config: McpServerConfig, } impl EffectiveMcpServer { pub fn configured(config: McpServerConfig) -> Self { + Self { config } + } + + pub fn config(&self) -> &McpServerConfig { + &self.config + } + + pub fn enabled(&self) -> bool { + self.config.enabled + } + + pub fn required(&self) -> bool { + self.config.required + } +} + +/// Inputs that determine the identity of a live MCP connection. +/// +/// Tool policy and presentation metadata intentionally do not appear here: +/// those belong to a publication and can change without reconnecting. +#[derive(Clone)] +pub(crate) struct McpServerConnectionIdentity { + transport: McpServerTransportConfig, + environment_id: String, + oauth_store: Option<(OAuthCredentialsStoreMode, AuthKeyringBackendKind)>, + oauth_credentials: Result, String>, + resolved_environment: Result>, String>, + local_stdio_fallback_cwd: Option, + referenced_environment_variables: Vec<(String, Option)>, + runtime_auth: Option, + runtime_auth_token: Option, + codex_apps_cache_identity: Option<(PathBuf, ConnectorRuntimeContextKey)>, + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, +} + +impl McpServerConnectionIdentity { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + server_name: &str, + server: &EffectiveMcpServer, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, + resolved_environment: &Result>, String>, + runtime_context: &McpRuntimeContext, + runtime_auth_provider: Option<&SharedAuthProvider>, + auth: Option<&CodexAuth>, + codex_apps_cache_identity: Option<(PathBuf, ConnectorRuntimeContextKey)>, + client_elicitation_capability: ElicitationCapability, + supports_openai_form_elicitation: bool, + ) -> Self { + let config = server.config(); + let stored_oauth_url = if runtime_auth_provider.is_none() { + match &config.transport { + McpServerTransportConfig::StreamableHttp { + url, + bearer_token_env_var: None, + .. + } => Some(url), + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var: Some(_), + .. + } + | McpServerTransportConfig::Stdio { .. } => None, + } + } else { + None + }; + let oauth_credentials = stored_oauth_url.map_or(Ok(None), |url| { + stored_oauth_credentials(server_name, url, store_mode, keyring_backend_kind).map_err( + |error| { + warn!(server_name, %error, "failed to read stored MCP OAuth credentials"); + error.to_string() + }, + ) + }); + let local_stdio_fallback_cwd = (config.is_local_environment() + && matches!( + config.transport, + McpServerTransportConfig::Stdio { cwd: None, .. } + )) + .then(|| runtime_context.local_stdio_fallback_cwd()); + let referenced_environment_variables = referenced_environment_variables(config); + let runtime_auth = runtime_auth_provider.and(auth).cloned(); + let runtime_auth_token = runtime_auth.as_ref().and_then(|auth| auth.get_token().ok()); + Self { - launch: McpServerLaunch::Configured(Box::new(config)), + transport: config.transport.clone(), + environment_id: config.environment_id.clone(), + oauth_store: stored_oauth_url + .is_some() + .then_some((store_mode, keyring_backend_kind)), + oauth_credentials, + resolved_environment: resolved_environment.clone(), + local_stdio_fallback_cwd, + referenced_environment_variables, + runtime_auth, + runtime_auth_token, + codex_apps_cache_identity, + client_elicitation_capability, + supports_openai_form_elicitation, } } - pub(crate) fn launch(&self) -> &McpServerLaunch { - &self.launch + pub(crate) fn has_same_connection_config(&self, other: &Self) -> bool { + let same_runtime_auth = match (&self.runtime_auth, &other.runtime_auth) { + (Some(CodexAuth::AgentIdentity(left)), Some(CodexAuth::AgentIdentity(right))) => { + left.record() == right.record() + } + (Some(left), Some(right)) => { + left == right + && left.get_account_id() == right.get_account_id() + && left.get_chatgpt_user_id() == right.get_chatgpt_user_id() + && left.is_fedramp_account() == right.is_fedramp_account() + } + (None, None) => true, + (Some(_), None) | (None, Some(_)) => false, + }; + self.transport == other.transport + && self.environment_id == other.environment_id + && self.oauth_store == other.oauth_store + && same_resolved_environment(&self.resolved_environment, &other.resolved_environment) + && self.local_stdio_fallback_cwd == other.local_stdio_fallback_cwd + && self.referenced_environment_variables == other.referenced_environment_variables + && same_runtime_auth + && self.runtime_auth_token == other.runtime_auth_token + && self.codex_apps_cache_identity == other.codex_apps_cache_identity + && self.client_elicitation_capability == other.client_elicitation_capability + && self.supports_openai_form_elicitation == other.supports_openai_form_elicitation } - pub fn configured_config(&self) -> Option<&McpServerConfig> { - match &self.launch { - McpServerLaunch::Configured(config) => Some(config.as_ref()), - } + pub(crate) fn oauth_credentials(&self) -> Result<&Option, &String> { + self.oauth_credentials.as_ref() } +} - pub fn enabled(&self) -> bool { - match &self.launch { - McpServerLaunch::Configured(config) => config.enabled, - } +impl PartialEq for McpServerConnectionIdentity { + fn eq(&self, other: &Self) -> bool { + self.has_same_connection_config(other) && self.oauth_credentials == other.oauth_credentials } +} - pub fn required(&self) -> bool { - match &self.launch { - McpServerLaunch::Configured(config) => config.required, - } +fn same_resolved_environment( + left: &Result>, String>, + right: &Result>, String>, +) -> bool { + match (left, right) { + (Ok(Some(left)), Ok(Some(right))) => Arc::ptr_eq(left, right), + (Ok(None), Ok(None)) => true, + (Err(left), Err(right)) => left == right, + (Ok(_), Ok(_)) | (Ok(_), Err(_)) | (Err(_), Ok(_)) => false, } } +fn referenced_environment_variables(config: &McpServerConfig) -> Vec<(String, Option)> { + let mut names = match &config.transport { + McpServerTransportConfig::Stdio { env_vars, .. } => env_vars + .iter() + .filter(|env_var| !env_var.is_remote_source()) + .map(|env_var| env_var.name().to_string()) + .collect::>(), + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + env_http_headers, + .. + } => bearer_token_env_var + .iter() + .chain(env_http_headers.iter().flat_map(|headers| headers.values())) + .cloned() + .collect(), + }; + names.sort(); + names.dedup(); + names + .into_iter() + .map(|name| { + let value = std::env::var_os(&name); + (name, value) + }) + .collect() +} + /// Transport origin retained for metrics and diagnostics after server launch. #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) enum McpServerOrigin { @@ -95,23 +257,22 @@ impl McpServerMetadata { impl From<&EffectiveMcpServer> for McpServerMetadata { fn from(server: &EffectiveMcpServer) -> Self { - match server.launch() { - McpServerLaunch::Configured(config) => Self { - environment_id: config.environment_id.clone(), - pollutes_memory: true, - origin: McpServerOrigin::from_transport(&config.transport), - supports_parallel_tool_calls: config.supports_parallel_tool_calls, - default_tools_approval_mode: config.default_tools_approval_mode, - tool_approval_modes: config - .tools - .iter() - .filter_map(|(name, config)| { - config - .approval_mode - .map(|approval_mode| (name.clone(), approval_mode)) - }) - .collect(), - }, + let config = server.config(); + Self { + environment_id: config.environment_id.clone(), + pollutes_memory: true, + origin: McpServerOrigin::from_transport(&config.transport), + supports_parallel_tool_calls: config.supports_parallel_tool_calls, + default_tools_approval_mode: config.default_tools_approval_mode, + tool_approval_modes: config + .tools + .iter() + .filter_map(|(name, config)| { + config + .approval_mode + .map(|approval_mode| (name.clone(), approval_mode)) + }) + .collect(), } } } diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index f6126824cb11..2b622700519e 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -298,8 +298,8 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( true } else if tools.is_empty() { let timeout = cfg - .configured_config() - .and_then(|config| config.startup_timeout_sec) + .config() + .startup_timeout_sec .unwrap_or(CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS); let ready = mcp_runtime .latest_wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, timeout) diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index aaad252f306c..1c556bd0b61f 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -327,12 +327,19 @@ impl Session { } pub(crate) fn mcp_elicitation_reviewer(self: &Arc) -> ElicitationReviewerHandle { - Arc::new(GuardianMcpElicitationReviewer::new(self)) + Arc::clone( + self.mcp_elicitation_reviewer_handle + .get_or_init(|| Arc::new(GuardianMcpElicitationReviewer::new(self))), + ) } pub(crate) fn mcp_elicitation_lifecycle(&self) -> codex_mcp::ElicitationLifecycle { - let elicitations = self.services.elicitations.clone(); - codex_mcp::ElicitationLifecycle::new(move || elicitations.register()) + self.mcp_elicitation_lifecycle_handle + .get_or_init(|| { + let elicitations = self.services.elicitations.clone(); + codex_mcp::ElicitationLifecycle::new(move || elicitations.register()) + }) + .clone() } #[expect( diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 33c28db6f15a..c7154ab122b6 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -41,6 +41,8 @@ pub(crate) struct Session { pub(super) mcp_refresh_pending: std::sync::atomic::AtomicBool, /// Serializes runtime refreshes without blocking calls that own a snapshot. pub(super) mcp_refresh_lock: Semaphore, + pub(super) mcp_elicitation_reviewer_handle: OnceLock, + pub(super) mcp_elicitation_lifecycle_handle: OnceLock, pub(crate) conversation: Arc, pub(crate) active_turn: Mutex>, pub(crate) input_queue: InputQueue, @@ -1151,6 +1153,8 @@ impl Session { multi_agent_version, mcp_refresh_pending: std::sync::atomic::AtomicBool::new(false), mcp_refresh_lock: Semaphore::new(/*permits*/ 1), + mcp_elicitation_reviewer_handle: OnceLock::new(), + mcp_elicitation_lifecycle_handle: OnceLock::new(), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: InputQueue::new(), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 918d9b502dde..89fdd0568092 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -5542,6 +5542,8 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), mcp_refresh_pending: std::sync::atomic::AtomicBool::new(true), mcp_refresh_lock: Semaphore::new(/*permits*/ 1), + mcp_elicitation_reviewer_handle: OnceLock::new(), + mcp_elicitation_lifecycle_handle: OnceLock::new(), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), @@ -7703,6 +7705,8 @@ where multi_agent_version: OnceLock::from(config.multi_agent_version_from_features()), mcp_refresh_pending: std::sync::atomic::AtomicBool::new(true), mcp_refresh_lock: Semaphore::new(/*permits*/ 1), + mcp_elicitation_reviewer_handle: OnceLock::new(), + mcp_elicitation_lifecycle_handle: OnceLock::new(), conversation: Arc::new(RealtimeConversationManager::new()), active_turn: Mutex::new(None), input_queue: super::input_queue::InputQueue::new(), @@ -7883,6 +7887,18 @@ async fn cancelled_mcp_refresh_remains_pending() { ); } +#[tokio::test] +async fn mcp_elicitation_reviewer_is_reused_across_runtime_refreshes() { + let (session, _turn_context) = make_session_and_context().await; + let session = Arc::new(session); + let previous = session.mcp_elicitation_reviewer(); + + session.mark_mcp_runtime_dirty(); + session.refresh_mcp_if_dirty().await; + + assert!(Arc::ptr_eq(&previous, &session.mcp_elicitation_reviewer())); +} + struct PendingNoiseConnectProvider; impl codex_exec_server::NoiseRendezvousConnectProvider for PendingNoiseConnectProvider { diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 44841b2f35f7..629dc3a95dab 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -381,6 +381,9 @@ async fn later_follow_up_uses_background_recovered_apps_after_mid_thread_startup tokio::fs::remove_dir_all(test.codex_home_path().join("cache/codex_apps_tools")).await?; startup_control.fail_next_initialize_attempts(/*attempts*/ 1); + test.codex + .set_openai_form_elicitation_support(/*supported*/ true) + .await?; test.codex.submit(Op::RefreshMcpServers).await?; test.submit_turn("use Calendar after transient Apps startup failures") .await?; diff --git a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs index 9e5ae1c52ff1..80e0a313c89f 100644 --- a/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs +++ b/codex-rs/ext/mcp/tests/hosted_apps_mcp.rs @@ -33,8 +33,8 @@ async fn contributes_hosted_plugin_runtime_without_an_executor() -> TestResult { let servers = manager.effective_servers(&config, Some(&auth)).await; let server = servers .get(CODEX_APPS_MCP_SERVER_NAME) - .and_then(|server| server.configured_config()) - .ok_or("hosted plugin runtime should be contributed as a configured server")?; + .ok_or("hosted plugin runtime should be contributed as a configured server")? + .config(); let McpServerTransportConfig::StreamableHttp { url, .. } = &server.transport else { panic!("hosted plugin runtime should use streamable HTTP"); }; @@ -94,8 +94,8 @@ async fn default_fallback_overwrites_reserved_config_without_an_extension() -> T let servers = manager.effective_servers(&config, Some(&auth)).await; let server = servers .get(CODEX_APPS_MCP_SERVER_NAME) - .and_then(|server| server.configured_config()) - .ok_or("default Apps MCP should be present")?; + .ok_or("default Apps MCP should be present")? + .config(); let McpServerTransportConfig::StreamableHttp { url, .. } = &server.transport else { panic!("default Apps MCP should use streamable HTTP"); }; diff --git a/codex-rs/rmcp-client/src/lib.rs b/codex-rs/rmcp-client/src/lib.rs index 45799ef4a88e..8f4bb2953288 100644 --- a/codex-rs/rmcp-client/src/lib.rs +++ b/codex-rs/rmcp-client/src/lib.rs @@ -28,6 +28,7 @@ pub use oauth::StoredOAuthTokens; pub use oauth::WrappedOAuthTokenResponse; pub use oauth::delete_oauth_tokens; pub use oauth::save_oauth_tokens; +pub use oauth::stored_oauth_credentials; pub use perform_oauth_login::OAuthProviderError; pub use perform_oauth_login::OauthLoginHandle; pub use perform_oauth_login::perform_oauth_login; diff --git a/codex-rs/rmcp-client/src/oauth.rs b/codex-rs/rmcp-client/src/oauth.rs index b8736168856c..ea6c24172a7d 100644 --- a/codex-rs/rmcp-client/src/oauth.rs +++ b/codex-rs/rmcp-client/src/oauth.rs @@ -127,6 +127,34 @@ pub(crate) fn oauth_token_status( }) } +/// Returns stored OAuth credentials without their derived expiration interval. +pub fn stored_oauth_credentials( + server_name: &str, + url: &str, + store_mode: OAuthCredentialsStoreMode, + keyring_backend_kind: AuthKeyringBackendKind, +) -> Result> { + let Some(resolved) = resolve_oauth_tokens_from_store_policy( + &DefaultKeyringStore, + server_name, + url, + store_mode, + keyring_backend_kind, + )? + else { + return Ok(None); + }; + Ok(normalized_oauth_credentials(Some(&resolved.tokens))) +} + +fn normalized_oauth_credentials(tokens: Option<&StoredOAuthTokens>) -> Option { + tokens.map(|tokens| { + let mut tokens = tokens.clone(); + tokens.token_response.0.set_expires_in(None); + tokens + }) +} + fn oauth_tokens_are_usable(tokens: &StoredOAuthTokens) -> bool { if tokens.client_id.trim().is_empty() { return false; @@ -510,6 +538,11 @@ impl OAuthPersistor { } } + pub(crate) async fn stored_credentials(&self) -> Option { + let credentials = self.inner.last_credentials.lock().await; + normalized_oauth_credentials(credentials.as_ref()) + } + /// Persists RMCP-managed credential changes back to this client's resolved authority. #[expect( clippy::await_holding_invalid_type, @@ -837,6 +870,47 @@ mod tests { use super::test_support::TempCodexHome; + #[test] + fn stored_oauth_credentials_ignore_derived_expiration_and_track_token_changes() -> Result<()> { + let _env = TempCodexHome::new(); + let mut tokens = sample_tokens(); + let credentials = super::normalized_oauth_credentials(Some(&tokens)); + tokens + .token_response + .0 + .set_expires_in(Some(&Duration::from_secs(1))); + assert_eq!( + credentials, + super::normalized_oauth_credentials(Some(&tokens)) + ); + super::save_oauth_tokens_to_file(&tokens)?; + assert_eq!( + credentials, + super::stored_oauth_credentials( + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )? + ); + + tokens + .token_response + .0 + .set_access_token(AccessToken::new("new-access-token".to_string())); + super::save_oauth_tokens_to_file(&tokens)?; + assert_ne!( + credentials, + super::stored_oauth_credentials( + &tokens.server_name, + &tokens.url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )? + ); + Ok(()) + } + #[test] fn resolve_oauth_tokens_from_store_policy_uses_keyring_when_available() -> Result<()> { let _env = TempCodexHome::new(); diff --git a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs index f32c2b691c2b..c6006ff99f0d 100644 --- a/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs +++ b/codex-rs/rmcp-client/src/oauth/tests/persistor_tests.rs @@ -5,6 +5,7 @@ use std::time::Duration; use anyhow::Context; use anyhow::Result; use codex_config::types::AuthKeyringBackendKind; +use codex_config::types::OAuthCredentialsStoreMode; use keyring::Error as KeyringError; use oauth2::AccessToken; use oauth2::TokenResponse; @@ -39,6 +40,7 @@ use crate::oauth::compute_store_key; use crate::oauth::load_oauth_tokens_from_file; use crate::oauth::refresh_lock::RefreshCredentialLock; use crate::oauth::save_oauth_tokens_to_file; +use crate::oauth::stored_oauth_credentials; use crate::startup_error::is_authentication_required_error; const REFRESH_LOCK_CONTENTION_EVENT_TARGET: &str = @@ -114,6 +116,7 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R let first = persistor_for(&initial).await?; let second = persistor_for(&initial).await?; + let initial_credentials = first.stored_credentials().await; let first_task = tokio::spawn({ let first = first.clone(); async move { first.refresh_if_needed().await } @@ -135,6 +138,15 @@ async fn concurrent_refreshes_call_provider_once_and_carry_omitted_fields() -> R first.persist_if_needed().await?; let stored = load_oauth_tokens_from_file(&initial.server_name, &initial.url)? .expect("refreshed credentials should be stored"); + let live_credentials = first.stored_credentials().await; + let disk_credentials = stored_oauth_credentials( + &initial.server_name, + &initial.url, + OAuthCredentialsStoreMode::File, + AuthKeyringBackendKind::Direct, + )?; + assert_eq!(live_credentials, disk_credentials); + assert_ne!(live_credentials, initial_credentials); let mut expected_response = initial.token_response.0.clone(); expected_response.set_access_token(AccessToken::new("refreshed-access-token".to_string())); // File loads derive `expires_in` from stable `expires_at`, so it may tick down before this diff --git a/codex-rs/rmcp-client/src/rmcp_client.rs b/codex-rs/rmcp-client/src/rmcp_client.rs index ed81eefba798..d5627c0ce0ee 100644 --- a/codex-rs/rmcp-client/src/rmcp_client.rs +++ b/codex-rs/rmcp-client/src/rmcp_client.rs @@ -733,6 +733,12 @@ impl RmcpClient { } } + /// Returns `None` when this client does not manage stored OAuth credentials. + pub async fn managed_oauth_credentials(&self) -> Option> { + let persistor = self.oauth_persistor().await?; + Some(persistor.stored_credentials().await) + } + /// Stop the MCP transport and any stdio server process owned by this client. pub async fn shutdown(&self) { let previous_state = {