diff --git a/codex-rs/codex-mcp/src/connection_generation.rs b/codex-rs/codex-mcp/src/connection_generation.rs new file mode 100644 index 0000000000..cd37746e54 --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_generation.rs @@ -0,0 +1,327 @@ +//! One atomically publishable generation of MCP connections. + +use std::collections::HashMap; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use crate::McpAuthStatusEntry; +use crate::codex_apps::CodexAppsToolsCacheContext; +use crate::codex_apps::CodexAppsToolsCacheKey; +use crate::connection_manager::emit_update; +use crate::connection_manager::mcp_init_error_display; +use crate::elicitation::ElicitationRequestManager; +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::StartupOutcomeError; +use crate::runtime::McpRuntimeContext; +use crate::server::EffectiveMcpServer; +use crate::server::McpServerMetadata; +use async_channel::Sender; +use codex_config::McpServerTransportConfig; +use codex_config::types::OAuthCredentialsStoreMode; +use codex_login::CodexAuth; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::Event; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::McpStartupCompleteEvent; +use codex_protocol::protocol::McpStartupFailure; +use codex_protocol::protocol::McpStartupStatus; +use codex_protocol::protocol::McpStartupUpdateEvent; +use rmcp::model::ElicitationCapability; +use tokio::sync::Notify; +use tokio::task::JoinHandle; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; + +/// Everything needed to start one complete generation of MCP connections. +pub struct McpConnectionStartParams { + pub mcp_servers: HashMap, + pub store_mode: OAuthCredentialsStoreMode, + pub auth_entries: HashMap, + pub approval_policy: AskForApproval, + pub submit_id: String, + pub tx_event: Sender, + pub permission_profile: PermissionProfile, + pub runtime_context: McpRuntimeContext, + pub codex_home: PathBuf, + pub codex_apps_tools_cache_key: CodexAppsToolsCacheKey, + pub host_owned_codex_apps_enabled: bool, + pub prefix_mcp_tool_names: bool, + pub client_elicitation_capability: ElicitationCapability, + pub tool_plugin_provenance: ToolPluginProvenance, + pub auth: Option, + pub elicitation_reviewer: Option, +} + +pub(crate) struct ConnectionGeneration { + pub(crate) clients: HashMap, + pub(crate) server_metadata: HashMap, + pub(crate) tool_plugin_provenance: Arc, + pub(crate) host_owned_codex_apps_enabled: bool, + pub(crate) prefix_mcp_tool_names: bool, + startup_cancellation_token: CancellationToken, + startup_completion: Arc>>>, + active_operations: AtomicUsize, + idle: Notify, +} + +impl ConnectionGeneration { + pub(crate) fn empty(prefix_mcp_tool_names: bool) -> Self { + Self { + clients: HashMap::new(), + server_metadata: HashMap::new(), + tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), + host_owned_codex_apps_enabled: false, + prefix_mcp_tool_names, + startup_cancellation_token: CancellationToken::new(), + startup_completion: Arc::new(StdMutex::new(None)), + active_operations: AtomicUsize::new(0), + idle: Notify::new(), + } + } + + pub(crate) async fn start( + params: McpConnectionStartParams, + elicitation_requests: ElicitationRequestManager, + startup_cancellation_token: CancellationToken, + ) -> Self { + let McpConnectionStartParams { + mcp_servers, + store_mode, + auth_entries, + approval_policy: _, + submit_id, + tx_event, + permission_profile: _, + runtime_context, + codex_home, + codex_apps_tools_cache_key, + host_owned_codex_apps_enabled, + prefix_mcp_tool_names, + client_elicitation_capability, + tool_plugin_provenance, + auth, + elicitation_reviewer: _, + } = params; + let mut clients = HashMap::new(); + let mut server_metadata = HashMap::new(); + let mut join_set = JoinSet::new(); + let tool_plugin_provenance = Arc::new(tool_plugin_provenance); + let startup_submit_id = submit_id.clone(); + let codex_apps_auth_provider = auth + .as_ref() + .filter(|auth| auth.uses_codex_backend()) + .map(codex_model_provider::auth_provider_from_auth); + + 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 _ = emit_update( + startup_submit_id.as_str(), + &tx_event, + McpStartupUpdateEvent { + server: server_name.clone(), + status: McpStartupStatus::Starting, + }, + ) + .await; + let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { + Some(CodexAppsToolsCacheContext { + codex_home: codex_home.clone(), + user_key: codex_apps_tools_cache_key.clone(), + }) + } else { + None + }; + let uses_env_bearer_token = + server + .configured_config() + .is_some_and(|config| match &config.transport { + McpServerTransportConfig::StreamableHttp { + bearer_token_env_var, + .. + } => bearer_token_env_var.is_some(), + McpServerTransportConfig::Stdio { .. } => false, + }); + let runtime_auth_provider = + if server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token { + codex_apps_auth_provider.clone() + } else { + None + }; + let async_managed_client = AsyncManagedClient::new( + server_name.clone(), + server, + store_mode, + cancel_token.clone(), + tx_event.clone(), + elicitation_requests.clone(), + codex_apps_tools_cache_context, + Arc::clone(&tool_plugin_provenance), + runtime_context.clone(), + runtime_auth_provider, + client_elicitation_capability.clone(), + ); + clients.insert(server_name.clone(), async_managed_client.clone()); + let tx_event = tx_event.clone(); + let submit_id = startup_submit_id.clone(); + let auth_entry = auth_entries.get(&server_name).cloned(); + join_set.spawn(async move { + let mut outcome = async_managed_client.client().await; + if cancel_token.is_cancelled() { + outcome = Err(StartupOutcomeError::Cancelled); + } + let status = match &outcome { + Ok(_) => McpStartupStatus::Ready, + Err(StartupOutcomeError::Cancelled) => McpStartupStatus::Cancelled, + Err(error) => McpStartupStatus::Failed { + error: mcp_init_error_display( + server_name.as_str(), + auth_entry.as_ref(), + error, + ), + }, + }; + + let _ = emit_update( + submit_id.as_str(), + &tx_event, + McpStartupUpdateEvent { + server: server_name.clone(), + status, + }, + ) + .await; + + (server_name, outcome) + }); + } + + let startup_completion = tokio::spawn(async move { + let outcomes = join_set.join_all().await; + let mut summary = McpStartupCompleteEvent::default(); + for (server_name, outcome) in outcomes { + match outcome { + Ok(_) => summary.ready.push(server_name), + Err(StartupOutcomeError::Cancelled) => summary.cancelled.push(server_name), + Err(StartupOutcomeError::Failed { error }) => { + summary.failed.push(McpStartupFailure { + server: server_name, + error, + }); + } + } + } + let _ = tx_event + .send(Event { + id: startup_submit_id, + msg: EventMsg::McpStartupComplete(summary), + }) + .await; + }); + + Self { + clients, + server_metadata, + tool_plugin_provenance, + host_owned_codex_apps_enabled, + prefix_mcp_tool_names, + startup_cancellation_token, + startup_completion: Arc::new(StdMutex::new(Some(startup_completion))), + active_operations: AtomicUsize::new(0), + idle: Notify::new(), + } + } + + pub(crate) fn acquire(self: &Arc) -> ConnectionGenerationLease { + self.active_operations.fetch_add(1, Ordering::AcqRel); + ConnectionGenerationLease { + generation: Arc::clone(self), + } + } + + pub(crate) async fn wait_until_idle(&self) { + loop { + let idle = self.idle.notified(); + if self.active_operations.load(Ordering::Acquire) == 0 { + return; + } + idle.await; + } + } + + pub(crate) fn cancel_startup(&self) { + self.startup_cancellation_token.cancel(); + } + + pub(crate) async fn shutdown(&self) { + self.cancel_startup(); + for client in self.clients.values() { + client.shutdown().await; + } + let startup_completion = self + .startup_completion + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + if let Some(startup_completion) = startup_completion { + let _ = startup_completion.await; + } + } +} + +impl Clone for ConnectionGeneration { + fn clone(&self) -> Self { + Self { + clients: self.clients.clone(), + server_metadata: self.server_metadata.clone(), + tool_plugin_provenance: Arc::clone(&self.tool_plugin_provenance), + host_owned_codex_apps_enabled: self.host_owned_codex_apps_enabled, + prefix_mcp_tool_names: self.prefix_mcp_tool_names, + startup_cancellation_token: self.startup_cancellation_token.clone(), + startup_completion: Arc::clone(&self.startup_completion), + active_operations: AtomicUsize::new(0), + idle: Notify::new(), + } + } +} + +pub(crate) struct ConnectionGenerationLease { + generation: Arc, +} + +impl std::ops::Deref for ConnectionGenerationLease { + type Target = ConnectionGeneration; + + fn deref(&self) -> &Self::Target { + &self.generation + } +} + +impl Drop for ConnectionGenerationLease { + fn drop(&mut self) { + if self + .generation + .active_operations + .fetch_sub(1, Ordering::AcqRel) + == 1 + { + self.generation.idle.notify_waiters(); + } + } +} + +impl Drop for ConnectionGeneration { + fn drop(&mut self) { + self.startup_cancellation_token.cancel(); + } +} diff --git a/codex-rs/codex-mcp/src/connection_lifecycle.rs b/codex-rs/codex-mcp/src/connection_lifecycle.rs new file mode 100644 index 0000000000..95c51a04b8 --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_lifecycle.rs @@ -0,0 +1,233 @@ +//! Lifecycle ownership for a thread's MCP connections. +//! +//! A manager remains stable while complete connection generations are prepared +//! and published atomically. Startup cancellation and shutdown therefore stay +//! with the object that owns the live clients instead of being coordinated by +//! callers. + +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::RwLock as StdRwLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use crate::connection_generation::ConnectionGeneration; +use crate::connection_generation::ConnectionGenerationLease; +use crate::connection_generation::McpConnectionStartParams; +use crate::elicitation::ElicitationRequestManager; +#[cfg(test)] +use crate::rmcp_client::AsyncManagedClient; +#[cfg(test)] +use crate::server::McpServerMetadata; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use tokio::sync::Mutex; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; + +struct McpConnectionLifecycle { + active: StdRwLock>, + pending_startup: StdMutex>, + mutation: Mutex<()>, + retirements: StdMutex>, + elicitation_requests: ElicitationRequestManager, + closed: AtomicBool, +} + +impl Drop for McpConnectionLifecycle { + fn drop(&mut self) { + if let Ok(pending_startup) = self.pending_startup.get_mut() + && let Some(token) = pending_startup.take() + { + token.cancel(); + } + if let Ok(active) = self.active.get_mut() { + active.cancel_startup(); + } + } +} + +/// Owns and atomically replaces all live MCP connections for one thread. +#[derive(Clone)] +pub struct McpConnectionManager { + lifecycle: Arc, +} + +impl McpConnectionManager { + /// Creates a valid manager with no live servers. + pub fn empty( + approval_policy: AskForApproval, + permission_profile: PermissionProfile, + prefix_mcp_tool_names: bool, + ) -> Self { + Self { + lifecycle: Arc::new(McpConnectionLifecycle { + active: StdRwLock::new(Arc::new(ConnectionGeneration::empty( + prefix_mcp_tool_names, + ))), + pending_startup: StdMutex::new(None), + mutation: Mutex::new(()), + retirements: StdMutex::new(JoinSet::new()), + elicitation_requests: ElicitationRequestManager::new( + approval_policy, + permission_profile, + /*reviewer*/ None, + ), + closed: AtomicBool::new(false), + }), + } + } + + /// Creates a manager and publishes its first connection generation. + pub async fn start(params: McpConnectionStartParams) -> Self { + let manager = Self::empty( + params.approval_policy, + params.permission_profile.clone(), + params.prefix_mcp_tool_names, + ); + manager.initialize(params).await; + manager + } + + /// Publishes the initial connection generation for an empty manager. + pub async fn initialize(&self, params: McpConnectionStartParams) { + self.replace_generation(params).await; + } + + /// Prepares and atomically publishes a replacement connection generation. + pub async fn reconfigure(&self, params: McpConnectionStartParams) { + self.replace_generation(params).await; + } + + async fn replace_generation(&self, params: McpConnectionStartParams) { + let _mutation_guard = self.lifecycle.mutation.lock().await; + if self.lifecycle.closed.load(Ordering::Acquire) { + return; + } + self.set_approval_policy(params.approval_policy); + self.set_permission_profile(params.permission_profile.clone()); + let elicitation_requests = self + .lifecycle + .elicitation_requests + .for_generation(params.elicitation_reviewer.clone()); + + let startup_cancellation_token = CancellationToken::new(); + let previous_pending = self + .lifecycle + .pending_startup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .replace(startup_cancellation_token.clone()); + if let Some(previous) = previous_pending { + previous.cancel(); + } + let generation = Arc::new( + ConnectionGeneration::start(params, elicitation_requests, startup_cancellation_token) + .await, + ); + let previous = match self.lifecycle.active.write() { + Ok(mut active) => std::mem::replace(&mut *active, generation), + Err(poisoned) => { + let mut active = poisoned.into_inner(); + std::mem::replace(&mut *active, generation) + } + }; + self.lifecycle + .pending_startup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take(); + previous.cancel_startup(); + let mut retirements = self + .lifecycle + .retirements + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + while retirements.try_join_next().is_some() {} + retirements.spawn(async move { + previous.wait_until_idle().await; + previous.shutdown().await; + }); + } + + pub fn cancel_startup(&self) { + let pending_startup = self + .lifecycle + .pending_startup + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + if let Some(pending_startup) = pending_startup { + pending_startup.cancel(); + } else { + self.generation().cancel_startup(); + } + } + + /// Atomically removes the active generation and stops all owned clients. + pub async fn shutdown(&self) { + self.lifecycle.closed.store(true, Ordering::Release); + self.cancel_startup(); + let _mutation_guard = self.lifecycle.mutation.lock().await; + let prefix_mcp_tool_names = self.generation().prefix_mcp_tool_names; + let empty = Arc::new(ConnectionGeneration::empty(prefix_mcp_tool_names)); + let previous = match self.lifecycle.active.write() { + Ok(mut active) => std::mem::replace(&mut *active, empty), + Err(poisoned) => { + let mut active = poisoned.into_inner(); + std::mem::replace(&mut *active, empty) + } + }; + previous.wait_until_idle().await; + previous.shutdown().await; + let mut retirements = std::mem::take( + &mut *self + .lifecycle + .retirements + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + ); + while retirements.join_next().await.is_some() {} + } + + pub(crate) fn generation(&self) -> ConnectionGenerationLease { + let active = self + .lifecycle + .active + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner); + active.acquire() + } + + pub(crate) fn elicitation_requests(&self) -> &ElicitationRequestManager { + &self.lifecycle.elicitation_requests + } + + #[cfg(test)] + pub(crate) fn insert_client_for_test(&self, server_name: String, client: AsyncManagedClient) { + let mut active = self + .lifecycle + .active + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Arc::make_mut(&mut active) + .clients + .insert(server_name, client); + } + + #[cfg(test)] + pub(crate) fn insert_server_metadata_for_test( + &self, + server_name: String, + metadata: McpServerMetadata, + ) { + let mut active = self + .lifecycle + .active + .write() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Arc::make_mut(&mut active) + .server_metadata + .insert(server_name, metadata); + } +} diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index d15a2e1500..201fb7c287 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -7,31 +7,24 @@ //! `codex-core`. use std::collections::HashMap; -use std::path::PathBuf; -use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; use crate::McpAuthStatusEntry; -use crate::codex_apps::CodexAppsToolsCacheContext; -use crate::codex_apps::CodexAppsToolsCacheKey; use crate::codex_apps::write_cached_codex_apps_tools_if_needed; -use crate::elicitation::ElicitationRequestManager; -use crate::elicitation::ElicitationReviewerHandle; +use crate::connection_generation::ConnectionGeneration; +use crate::connection_generation::ConnectionGenerationLease; +use crate::connection_lifecycle::McpConnectionManager; 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::MCP_TOOLS_FETCH_UNCACHED_DURATION_METRIC; use crate::rmcp_client::MCP_TOOLS_LIST_DURATION_METRIC; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::StartupOutcomeError; use crate::rmcp_client::list_tools_for_client_uncached; -use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; -use crate::server::EffectiveMcpServer; -use crate::server::McpServerMetadata; +use crate::server::McpServerOrigin; use crate::tools::ToolInfo; use crate::tools::filter_tools; use crate::tools::normalize_tools_for_model_with_prefix; @@ -40,22 +33,15 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use async_channel::Sender; -use codex_config::Constrained; use codex_config::McpServerTransportConfig; -use codex_config::types::OAuthCredentialsStoreMode; -use codex_login::CodexAuth; use codex_protocol::mcp::CallToolResult; use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; -use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::Event; use codex_protocol::protocol::EventMsg; -use codex_protocol::protocol::McpStartupCompleteEvent; use codex_protocol::protocol::McpStartupFailure; -use codex_protocol::protocol::McpStartupStatus; use codex_protocol::protocol::McpStartupUpdateEvent; use codex_rmcp_client::ElicitationResponse; -use rmcp::model::ElicitationCapability; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; @@ -66,7 +52,6 @@ use rmcp::model::Resource; use rmcp::model::ResourceTemplate; use serde_json::Value as JsonValue; use tokio::task::JoinSet; -use tokio_util::sync::CancellationToken; use tracing::Instrument; use tracing::instrument; use tracing::trace; @@ -101,267 +86,62 @@ pub fn tool_is_model_visible(tool: &ToolInfo) -> bool { .any(|target| target.as_str() == Some(MCP_UI_MODEL_VISIBILITY)) } -/// A thin wrapper around a set of running [`RmcpClient`] instances. -pub struct McpConnectionManager { - clients: HashMap, - server_metadata: HashMap, - tool_plugin_provenance: Arc, - host_owned_codex_apps_enabled: bool, - prefix_mcp_tool_names: bool, - elicitation_requests: ElicitationRequestManager, - startup_cancellation_token: CancellationToken, -} - impl McpConnectionManager { - pub fn new_uninitialized( - approval_policy: &Constrained, - permission_profile: &Constrained, - prefix_mcp_tool_names: bool, - ) -> Self { - Self::new_uninitialized_with_permission_profile( - approval_policy, - permission_profile.get(), - prefix_mcp_tool_names, - ) - } - - pub 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(), - tool_plugin_provenance: Arc::new(ToolPluginProvenance::default()), - host_owned_codex_apps_enabled: false, - prefix_mcp_tool_names, - elicitation_requests: ElicitationRequestManager::new( - approval_policy.value(), - permission_profile.clone(), - /*reviewer*/ None, - ), - startup_cancellation_token: CancellationToken::new(), - } - } - pub fn has_servers(&self) -> bool { - !self.clients.is_empty() + !self.generation().clients.is_empty() } - /// Drain all MCP clients from this manager and return a future that stops - /// them and terminates their stdio server processes. - pub fn begin_shutdown(&mut self) -> impl std::future::Future + Send + 'static { - self.startup_cancellation_token.cancel(); - let clients = std::mem::take(&mut self.clients); - self.server_metadata.clear(); - async move { - for client in clients.into_values() { - client.shutdown().await; - } - } + /// Returns a future that atomically removes and stops all MCP clients. + pub fn begin_shutdown(&self) -> impl std::future::Future + Send + 'static { + let manager = self.clone(); + async move { manager.shutdown().await } } - /// Stop all MCP clients owned by this manager and terminate stdio server processes. - pub async fn shutdown(&mut self) { - self.begin_shutdown().await; - } - - pub fn server_origin(&self, server_name: &str) -> Option<&str> { - self.server_metadata + pub fn server_origin(&self, server_name: &str) -> Option { + self.generation() + .server_metadata .get(server_name) .and_then(|metadata| metadata.origin.as_ref()) - .map(super::server::McpServerOrigin::as_str) + .map(McpServerOrigin::as_str) + .map(ToOwned::to_owned) } pub fn server_pollutes_memory(&self, server_name: &str) -> bool { - self.server_metadata + self.generation() + .server_metadata .get(server_name) .is_none_or(|metadata| metadata.pollutes_memory) } - pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option<&str> { - self.tool_plugin_provenance + pub fn plugin_id_for_mcp_server_name(&self, server_name: &str) -> Option { + self.generation() + .tool_plugin_provenance .plugin_id_for_mcp_server_name(server_name) + .map(ToOwned::to_owned) } pub fn is_host_owned_codex_apps_server(&self, server_name: &str) -> bool { - self.host_owned_codex_apps_enabled && server_name == CODEX_APPS_MCP_SERVER_NAME + self.generation().host_owned_codex_apps_enabled && server_name == CODEX_APPS_MCP_SERVER_NAME } - pub fn set_approval_policy(&self, approval_policy: &Constrained) { - if let Ok(mut policy) = self.elicitation_requests.approval_policy.lock() { - *policy = approval_policy.value(); + pub fn set_approval_policy(&self, approval_policy: codex_protocol::protocol::AskForApproval) { + if let Ok(mut policy) = self.elicitation_requests().approval_policy.lock() { + *policy = approval_policy; } } pub fn set_permission_profile(&self, permission_profile: PermissionProfile) { - if let Ok(mut profile) = self.elicitation_requests.permission_profile.lock() { + if let Ok(mut profile) = self.elicitation_requests().permission_profile.lock() { *profile = permission_profile; } } pub fn elicitations_auto_deny(&self) -> bool { - self.elicitation_requests.auto_deny() + self.elicitation_requests().auto_deny() } pub fn set_elicitations_auto_deny(&self, auto_deny: bool) { - self.elicitation_requests.set_auto_deny(auto_deny); - } - - #[allow(clippy::new_ret_no_self, clippy::too_many_arguments)] - pub async fn new( - mcp_servers: &HashMap, - store_mode: OAuthCredentialsStoreMode, - auth_entries: HashMap, - approval_policy: &Constrained, - submit_id: String, - tx_event: Sender, - initial_permission_profile: PermissionProfile, - runtime_context: McpRuntimeContext, - codex_home: PathBuf, - codex_apps_tools_cache_key: CodexAppsToolsCacheKey, - host_owned_codex_apps_enabled: bool, - prefix_mcp_tool_names: bool, - client_elicitation_capability: ElicitationCapability, - tool_plugin_provenance: ToolPluginProvenance, - auth: Option<&CodexAuth>, - elicitation_reviewer: Option, - ) -> (Self, CancellationToken) { - let cancel_token = CancellationToken::new(); - let mut clients = HashMap::new(); - let mut server_metadata = HashMap::new(); - let mut join_set = JoinSet::new(); - let elicitation_requests = ElicitationRequestManager::new( - approval_policy.value(), - initial_permission_profile, - elicitation_reviewer, - ); - let tool_plugin_provenance = Arc::new(tool_plugin_provenance); - let startup_submit_id = submit_id.clone(); - let codex_apps_auth_provider = auth - .filter(|auth| auth.uses_codex_backend()) - .map(codex_model_provider::auth_provider_from_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 = cancel_token.child_token(); - let _ = emit_update( - startup_submit_id.as_str(), - &tx_event, - McpStartupUpdateEvent { - server: server_name.clone(), - status: McpStartupStatus::Starting, - }, - ) - .await; - let codex_apps_tools_cache_context = if server_name == CODEX_APPS_MCP_SERVER_NAME { - Some(CodexAppsToolsCacheContext { - codex_home: codex_home.clone(), - user_key: codex_apps_tools_cache_key.clone(), - }) - } else { - None - }; - let uses_env_bearer_token = - server - .configured_config() - .is_some_and(|config| match &config.transport { - McpServerTransportConfig::StreamableHttp { - bearer_token_env_var, - .. - } => bearer_token_env_var.is_some(), - McpServerTransportConfig::Stdio { .. } => false, - }); - let runtime_auth_provider = - if server_name == CODEX_APPS_MCP_SERVER_NAME && !uses_env_bearer_token { - codex_apps_auth_provider.clone() - } else { - None - }; - let async_managed_client = AsyncManagedClient::new( - server_name.clone(), - server, - store_mode, - cancel_token.clone(), - tx_event.clone(), - elicitation_requests.clone(), - codex_apps_tools_cache_context, - Arc::clone(&tool_plugin_provenance), - runtime_context.clone(), - runtime_auth_provider, - client_elicitation_capability.clone(), - ); - clients.insert(server_name.clone(), async_managed_client.clone()); - let tx_event = tx_event.clone(); - let submit_id = startup_submit_id.clone(); - let auth_entry = auth_entries.get(&server_name).cloned(); - join_set.spawn(async move { - let mut outcome = async_managed_client.client().await; - if cancel_token.is_cancelled() { - outcome = Err(StartupOutcomeError::Cancelled); - } - let status = match &outcome { - Ok(_) => McpStartupStatus::Ready, - Err(StartupOutcomeError::Cancelled) => McpStartupStatus::Cancelled, - Err(error) => { - let error_str = mcp_init_error_display( - server_name.as_str(), - auth_entry.as_ref(), - error, - ); - McpStartupStatus::Failed { error: error_str } - } - }; - - let _ = emit_update( - submit_id.as_str(), - &tx_event, - McpStartupUpdateEvent { - server: server_name.clone(), - status, - }, - ) - .await; - - (server_name, outcome) - }); - } - let manager = Self { - clients, - server_metadata, - tool_plugin_provenance, - host_owned_codex_apps_enabled, - prefix_mcp_tool_names, - elicitation_requests: elicitation_requests.clone(), - startup_cancellation_token: cancel_token.clone(), - }; - tokio::spawn(async move { - let outcomes = join_set.join_all().await; - let mut summary = McpStartupCompleteEvent::default(); - for (server_name, outcome) in outcomes { - match outcome { - Ok(_) => summary.ready.push(server_name), - Err(StartupOutcomeError::Cancelled) => summary.cancelled.push(server_name), - Err(StartupOutcomeError::Failed { error }) => { - summary.failed.push(McpStartupFailure { - server: server_name, - error, - }) - } - } - } - let _ = tx_event - .send(Event { - id: startup_submit_id, - msg: EventMsg::McpStartupComplete(summary), - }) - .await; - }); - (manager, cancel_token) + self.elicitation_requests().set_auto_deny(auto_deny); } pub async fn resolve_elicitation( @@ -370,13 +150,14 @@ impl McpConnectionManager { id: RequestId, response: ElicitationResponse, ) -> Result<()> { - self.elicitation_requests + self.elicitation_requests() .resolve(server_name, id, response) .await } 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 generation = self.generation(); + let Some(async_managed_client) = generation.clients.get(server_name) else { return false; }; @@ -390,9 +171,10 @@ impl McpConnectionManager { &self, required_servers: &[String], ) -> Vec { + let generation = self.generation(); let mut failures = Vec::new(); for server_name in required_servers { - let Some(async_managed_client) = self.clients.get(server_name).cloned() else { + let Some(async_managed_client) = generation.clients.get(server_name).cloned() else { failures.push(McpStartupFailure { server: server_name.clone(), error: format!("required MCP server `{server_name}` was not initialized"), @@ -412,10 +194,11 @@ impl McpConnectionManager { } /// 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)] pub async fn list_all_tools(&self) -> Vec { + let generation = self.generation(); let mut tools = Vec::new(); - for (server_name, managed_client) in &self.clients { + for (server_name, managed_client) in &generation.clients { let has_cached_tool_info_snapshot = managed_client.cached_tool_info_snapshot.is_some(); let startup_complete = managed_client .startup_complete @@ -446,17 +229,18 @@ impl McpConnectionManager { tools.extend( server_tools .into_iter() - .map(|tool| self.with_server_metadata(tool)), + .map(|tool| Self::with_server_metadata(&generation, tool)), ); } - normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names) + normalize_tools_for_model_with_prefix(tools, generation.prefix_mcp_tool_names) } /// Returns presentation metadata without waiting for uncached clients still initializing. /// Cached values will be used if available and the server is still starting up. pub async fn list_available_server_infos(&self) -> HashMap { + let generation = self.generation(); let mut server_infos = HashMap::new(); - for (server_name, client) in &self.clients { + for (server_name, client) in &generation.clients { if !client.startup_complete.load(Ordering::Acquire) { if let Some(server_info) = client.cached_server_info.clone() { server_infos.insert(server_name.clone(), server_info); @@ -483,7 +267,8 @@ impl McpConnectionManager { /// latest filtered tools are returned directly to the caller. On /// failure, the existing cache remains unchanged. pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { - let managed_client = self + let generation = self.generation(); + let managed_client = generation .clients .get(CODEX_APPS_MCP_SERVER_NAME) .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))? @@ -524,16 +309,16 @@ impl McpConnectionManager { .into_iter() .map(|mut tool| { tool.tool = tool_with_model_visible_input_schema(&tool.tool); - self.with_server_metadata(tool) + Self::with_server_metadata(&generation, tool) }); Ok(normalize_tools_for_model_with_prefix( tools, - self.prefix_mcp_tool_names, + generation.prefix_mcp_tool_names, )) } - fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { - let Some(metadata) = self.server_metadata.get(&tool.server_name) else { + fn with_server_metadata(generation: &ConnectionGeneration, mut tool: ToolInfo) -> ToolInfo { + let Some(metadata) = generation.server_metadata.get(&tool.server_name) else { tool.supports_parallel_tool_calls = false; tool.server_origin = None; return tool; @@ -550,11 +335,10 @@ impl McpConnectionManager { /// Returns a single map that contains all resources. Each key is the /// server name and the value is a vector of resources. pub async fn list_all_resources(&self) -> HashMap> { + let generation = self.generation(); let mut join_set = JoinSet::new(); - let clients_snapshot = &self.clients; - - for (server_name, async_managed_client) in clients_snapshot { + for (server_name, async_managed_client) in &generation.clients { let server_name = server_name.clone(); let Ok(managed_client) = async_managed_client.client().await else { continue; @@ -615,11 +399,10 @@ impl McpConnectionManager { /// Returns a single map that contains all resource templates. Each key is the /// server name and the value is a vector of resource templates. pub async fn list_all_resource_templates(&self) -> HashMap> { + let generation = self.generation(); let mut join_set = JoinSet::new(); - let clients_snapshot = &self.clients; - - for (server_name, async_managed_client) in clients_snapshot { + for (server_name, async_managed_client) in &generation.clients { let server_name_cloned = server_name.clone(); let Ok(managed_client) = async_managed_client.client().await else { continue; @@ -689,7 +472,7 @@ impl McpConnectionManager { arguments: Option, meta: Option, ) -> Result { - let client = self.client_by_name(server).await?; + let (_generation, client) = self.client_by_name(server).await?; if !client.tool_filter.allows(tool) { return Err(anyhow!( "tool '{tool}' is disabled for MCP server '{server}'" @@ -723,10 +506,8 @@ impl McpConnectionManager { &self, server: &str, ) -> Result { - Ok(self - .client_by_name(server) - .await? - .server_supports_sandbox_state_meta_capability) + let (_generation, client) = self.client_by_name(server).await?; + Ok(client.server_supports_sandbox_state_meta_capability) } /// List resources from the specified server. @@ -735,7 +516,7 @@ impl McpConnectionManager { server: &str, params: Option, ) -> Result { - let managed = self.client_by_name(server).await?; + let (_generation, managed) = self.client_by_name(server).await?; let timeout = managed.tool_timeout; managed @@ -751,7 +532,7 @@ impl McpConnectionManager { server: &str, params: Option, ) -> Result { - let managed = self.client_by_name(server).await?; + let (_generation, managed) = self.client_by_name(server).await?; let client = managed.client.clone(); let timeout = managed.tool_timeout; @@ -767,7 +548,7 @@ impl McpConnectionManager { server: &str, params: ReadResourceRequestParams, ) -> Result { - let managed = self.client_by_name(server).await?; + let (_generation, managed) = self.client_by_name(server).await?; let client = managed.client.clone(); let timeout = managed.tool_timeout; let uri = params.uri.clone(); @@ -778,24 +559,23 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } - async fn client_by_name(&self, name: &str) -> Result { - self.clients + async fn client_by_name( + &self, + name: &str, + ) -> Result<(ConnectionGenerationLease, ManagedClient)> { + let generation = self.generation(); + let client = generation + .clients .get(name) .ok_or_else(|| anyhow!("unknown MCP server '{name}'"))? .client() .await - .context("failed to get client") + .context("failed to get client")?; + Ok((generation, client)) } } -impl Drop for McpConnectionManager { - fn drop(&mut self) { - self.startup_cancellation_token.cancel(); - self.clients.clear(); - } -} - -async fn emit_update( +pub(crate) async fn emit_update( submit_id: &str, tx_event: &Sender, update: McpStartupUpdateEvent, @@ -808,7 +588,7 @@ async fn emit_update( .await } -fn mcp_init_error_display( +pub(crate) fn mcp_init_error_display( server_name: &str, entry: Option<&McpAuthStatusEntry>, err: &StartupOutcomeError, diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 60a0026eeb..406ba0ff53 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,4 +1,9 @@ use super::*; +use crate::CodexAppsToolsCacheKey; +use crate::EffectiveMcpServer; +use crate::McpConnectionStartParams; +use crate::McpRuntimeContext; +use crate::ToolPluginProvenance; use crate::codex_apps::CODEX_APPS_TOOLS_CACHE_SCHEMA_VERSION; use crate::codex_apps::CodexAppsToolsCacheContext; use crate::codex_apps::load_startup_cached_codex_apps_server_info; @@ -12,6 +17,7 @@ use crate::elicitation::elicitation_is_rejected_by_policy; use crate::rmcp_client::AsyncManagedClient; use crate::rmcp_client::ManagedClient; use crate::rmcp_client::StartupOutcomeError; +use crate::server::McpServerMetadata; use crate::server::McpServerOrigin; use crate::tools::ToolFilter; use crate::tools::ToolInfo; @@ -20,10 +26,12 @@ use crate::tools::normalize_tools_for_model_with_prefix; use crate::tools::tool_with_model_visible_input_schema; use codex_config::Constrained; use codex_config::McpServerConfig; +use codex_config::types::OAuthCredentialsStoreMode; use codex_exec_server::EnvironmentManager; use codex_protocol::ToolName; use codex_protocol::mcp::McpServerInfo; use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::GranularApprovalConfig; use codex_protocol::protocol::McpAuthStatus; use futures::FutureExt; @@ -36,8 +44,10 @@ use rmcp::model::Meta; use rmcp::model::NumberOrString; use rmcp::model::Tool; use std::collections::HashSet; +use std::path::PathBuf; use std::sync::Arc; use tempfile::tempdir; +use tokio_util::sync::CancellationToken; fn create_test_tool(server_name: &str, tool_name: &str) -> ToolInfo { ToolInfo { @@ -304,6 +314,69 @@ async fn disabled_permissions_do_not_auto_accept_elicitation_with_requested_fiel ); } +#[tokio::test] +async fn elicitations_with_the_same_id_resolve_in_request_order() -> anyhow::Result<()> { + let manager = ElicitationRequestManager::new( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*reviewer*/ None, + ); + let first_generation = manager.for_generation(/*reviewer*/ None); + let second_generation = manager.for_generation(/*reviewer*/ None); + let (tx_event, rx_event) = async_channel::bounded(2); + let first_sender = first_generation.make_sender("server".to_string(), tx_event.clone()); + let second_sender = second_generation.make_sender("server".to_string(), tx_event); + let request_id = NumberOrString::Number(1); + let elicitation = CreateElicitationRequestParams::FormElicitationParams { + meta: None, + message: "What should I say?".to_string(), + requested_schema: rmcp::model::ElicitationSchema::builder() + .required_property( + "message", + rmcp::model::PrimitiveSchema::String(rmcp::model::StringSchema::new()), + ) + .build() + .expect("schema should build"), + }; + + let first = tokio::spawn(first_sender(request_id.clone(), elicitation.clone())); + rx_event.recv().await?; + let second = tokio::spawn(second_sender(request_id.clone(), elicitation)); + rx_event.recv().await?; + + let first_response = ElicitationResponse { + action: ElicitationAction::Accept, + content: Some(serde_json::json!({"message": "first"})), + meta: None, + }; + let second_response = ElicitationResponse { + action: ElicitationAction::Decline, + content: None, + meta: None, + }; + manager + .resolve( + "server".to_string(), + request_id.clone(), + first_response.clone(), + ) + .await?; + manager + .resolve("server".to_string(), request_id, second_response.clone()) + .await?; + + assert_eq!( + ( + first.await.expect("first elicitation task should finish")?, + second + .await + .expect("second elicitation task should finish")?, + ), + (first_response, second_response), + ); + Ok(()) +} + #[test] fn test_normalize_tools_short_non_duplicated_names() { let tools = vec![ @@ -799,12 +872,12 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_while_client_is_pending() .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_client_for_test( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -835,13 +908,13 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() { .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); let server_info = create_test_server_info("Codex Apps"); - manager.clients.insert( + manager.insert_client_for_test( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -873,12 +946,12 @@ async fn list_all_tools_accepts_canonical_namespaced_tool_names() { .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ false, ); - manager.clients.insert( + manager.insert_client_for_test( "rmcp".to_string(), AsyncManagedClient { client: pending_client, @@ -916,12 +989,12 @@ async fn list_all_tools_applies_legacy_mcp_prefix_by_default() { .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_client_for_test( "rmcp".to_string(), AsyncManagedClient { client: pending_client, @@ -958,12 +1031,12 @@ async fn list_all_tools_blocks_while_client_is_pending_without_cached_tool_info_ .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_client_for_test( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -987,12 +1060,12 @@ async fn list_all_tools_does_not_block_when_cached_tool_info_snapshot_is_empty() .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); - manager.clients.insert( + manager.insert_client_for_test( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: pending_client, @@ -1026,13 +1099,13 @@ async fn list_all_tools_uses_cached_tool_info_snapshot_when_client_startup_fails .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); let startup_complete = Arc::new(std::sync::atomic::AtomicBool::new(true)); - manager.clients.insert( + manager.insert_client_for_test( CODEX_APPS_MCP_SERVER_NAME.to_string(), AsyncManagedClient { client: failed_client, @@ -1072,12 +1145,12 @@ async fn list_all_tools_adds_server_metadata_to_cached_tools() { .shared(); let approval_policy = Constrained::allow_any(AskForApproval::OnFailure); let permission_profile = Constrained::allow_any(PermissionProfile::default()); - let mut manager = McpConnectionManager::new_uninitialized( - &approval_policy, - &permission_profile, + let manager = McpConnectionManager::empty( + approval_policy.value(), + permission_profile.get().clone(), /*prefix_mcp_tool_names*/ true, ); - manager.server_metadata.insert( + manager.insert_server_metadata_for_test( server_name.to_string(), McpServerMetadata { pollutes_memory: true, @@ -1087,7 +1160,7 @@ async fn list_all_tools_adds_server_metadata_to_cached_tools() { supports_parallel_tool_calls: true, }, ); - manager.clients.insert( + manager.insert_client_for_test( server_name.to_string(), AsyncManagedClient { client: pending_client, @@ -1167,35 +1240,42 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { ), ]); - let (manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - OAuthCredentialsStoreMode::default(), - HashMap::new(), - &approval_policy, - String::new(), + let manager = McpConnectionManager::start(McpConnectionStartParams { + mcp_servers, + store_mode: OAuthCredentialsStoreMode::default(), + auth_entries: HashMap::new(), + approval_policy: approval_policy.value(), + submit_id: String::new(), tx_event, - PermissionProfile::default(), - McpRuntimeContext::new( + permission_profile: PermissionProfile::default(), + runtime_context: McpRuntimeContext::new( Arc::new(EnvironmentManager::without_environments()), PathBuf::from("/tmp"), ), - codex_home.path().to_path_buf(), - CodexAppsToolsCacheKey { + codex_home: codex_home.path().to_path_buf(), + codex_apps_tools_cache_key: CodexAppsToolsCacheKey { account_id: None, chatgpt_user_id: None, is_workspace_account: false, }, - /*host_owned_codex_apps_enabled*/ false, - /*prefix_mcp_tool_names*/ true, - ElicitationCapability::default(), - ToolPluginProvenance::default(), - /*auth*/ None, - /*elicitation_reviewer*/ None, - ) + host_owned_codex_apps_enabled: false, + prefix_mcp_tool_names: true, + client_elicitation_capability: ElicitationCapability::default(), + tool_plugin_provenance: ToolPluginProvenance::default(), + auth: None, + elicitation_reviewer: None, + }) .await; - assert!(manager.clients.contains_key("stdio")); - assert!(manager.clients.contains_key("http")); + assert_eq!( + manager + .generation() + .clients + .keys() + .cloned() + .collect::>(), + HashSet::from(["http".to_string(), "stdio".to_string()]) + ); assert!( !manager .wait_for_server_ready("stdio", Duration::from_millis(10)) @@ -1210,7 +1290,49 @@ async fn no_local_runtime_fails_local_stdio_but_keeps_local_http_server() { failures[0].error, "local stdio MCP server `stdio` requires a local environment" ); - cancel_token.cancel(); + manager.cancel_startup(); +} + +#[tokio::test] +async fn shutdown_prevents_reconfiguration_from_publishing_connections() { + let manager = McpConnectionManager::empty( + AskForApproval::OnRequest, + PermissionProfile::default(), + /*prefix_mcp_tool_names*/ true, + ); + manager.shutdown().await; + let (tx_event, rx_event) = async_channel::bounded(1); + drop(rx_event); + + manager + .reconfigure(McpConnectionStartParams { + mcp_servers: HashMap::new(), + store_mode: OAuthCredentialsStoreMode::default(), + auth_entries: HashMap::new(), + approval_policy: AskForApproval::OnRequest, + submit_id: String::new(), + tx_event, + permission_profile: PermissionProfile::default(), + runtime_context: McpRuntimeContext::new( + Arc::new(EnvironmentManager::without_environments()), + PathBuf::from("/tmp"), + ), + codex_home: PathBuf::from("/tmp"), + codex_apps_tools_cache_key: CodexAppsToolsCacheKey { + account_id: None, + chatgpt_user_id: None, + is_workspace_account: false, + }, + host_owned_codex_apps_enabled: true, + prefix_mcp_tool_names: false, + client_elicitation_capability: ElicitationCapability::default(), + tool_plugin_provenance: ToolPluginProvenance::default(), + auth: None, + elicitation_reviewer: None, + }) + .await; + + assert!(!manager.is_host_owned_codex_apps_server(CODEX_APPS_MCP_SERVER_NAME)); } #[test] diff --git a/codex-rs/codex-mcp/src/elicitation.rs b/codex-rs/codex-mcp/src/elicitation.rs index a51cd7c623..d59b034f7c 100644 --- a/codex-rs/codex-mcp/src/elicitation.rs +++ b/codex-rs/codex-mcp/src/elicitation.rs @@ -6,6 +6,7 @@ //! and later resolved through the stored responder. use std::collections::HashMap; +use std::collections::VecDeque; use std::sync::Arc; use std::sync::Mutex as StdMutex; @@ -54,7 +55,7 @@ pub(crate) struct ElicitationRequestManager { pub(crate) approval_policy: Arc>, pub(crate) permission_profile: Arc>, auto_deny: Arc>, - reviewer: Option, + reviewer: Arc>>, } impl ElicitationRequestManager { @@ -68,7 +69,7 @@ impl ElicitationRequestManager { approval_policy: Arc::new(StdMutex::new(approval_policy)), permission_profile: Arc::new(StdMutex::new(permission_profile)), auto_deny: Arc::new(StdMutex::new(false)), - reviewer, + reviewer: Arc::new(StdMutex::new(reviewer)), } } @@ -85,19 +86,38 @@ impl ElicitationRequestManager { } } + pub(crate) fn for_generation(&self, reviewer: Option) -> Self { + Self { + requests: Arc::clone(&self.requests), + approval_policy: Arc::clone(&self.approval_policy), + permission_profile: Arc::clone(&self.permission_profile), + auto_deny: Arc::clone(&self.auto_deny), + reviewer: Arc::new(StdMutex::new(reviewer)), + } + } + pub(crate) async fn resolve( &self, server_name: String, id: RequestId, - response: ElicitationResponse, + mut response: ElicitationResponse, ) -> Result<()> { - self.requests - .lock() - .await - .remove(&(server_name, id)) - .ok_or_else(|| anyhow!("elicitation request not found"))? - .send(response) - .map_err(|e| anyhow!("failed to send elicitation response: {e:?}")) + let key = (server_name, id); + loop { + let responder = { + let mut requests = self.requests.lock().await; + let responder = requests.get_mut(&key).and_then(VecDeque::pop_front); + if requests.get(&key).is_some_and(VecDeque::is_empty) { + requests.remove(&key); + } + responder + } + .ok_or_else(|| anyhow!("elicitation request not found"))?; + match responder.send(response) { + Ok(()) => return Ok(()), + Err(returned_response) => response = returned_response, + } + } } pub(crate) fn make_sender( @@ -117,7 +137,7 @@ impl ElicitationRequestManager { let approval_policy = approval_policy.clone(); let permission_profile = permission_profile.clone(); let auto_deny = auto_deny.clone(); - let reviewer = reviewer.clone(); + let reviewer = reviewer.lock().ok().and_then(|reviewer| reviewer.clone()); async move { let auto_deny = auto_deny .lock() @@ -203,7 +223,9 @@ impl ElicitationRequestManager { let (tx, rx) = oneshot::channel(); { let mut lock = elicitation_requests.lock().await; - lock.insert((server_name.clone(), id.clone()), tx); + lock.entry((server_name.clone(), id.clone())) + .or_default() + .push_back(tx); } let _ = tx_event .send(Event { @@ -241,7 +263,7 @@ pub(crate) fn elicitation_is_rejected_by_policy(approval_policy: AskForApproval) } } -type ResponderMap = HashMap<(String, RequestId), oneshot::Sender>; +type ResponderMap = HashMap<(String, RequestId), VecDeque>>; fn can_auto_accept_elicitation(elicitation: &CreateElicitationRequestParams) -> bool { match elicitation { diff --git a/codex-rs/codex-mcp/src/lib.rs b/codex-rs/codex-mcp/src/lib.rs index a56365277a..c3bddf550a 100644 --- a/codex-rs/codex-mcp/src/lib.rs +++ b/codex-rs/codex-mcp/src/lib.rs @@ -1,4 +1,5 @@ -pub use connection_manager::McpConnectionManager; +pub use connection_generation::McpConnectionStartParams; +pub use connection_lifecycle::McpConnectionManager; pub use connection_manager::tool_is_model_visible; pub use elicitation::ElicitationReviewRequest; pub use elicitation::ElicitationReviewer; @@ -55,6 +56,8 @@ pub use tools::declared_openai_file_input_param_names; pub(crate) mod auth_elicitation; pub(crate) mod codex_apps; +pub(crate) mod connection_generation; +pub(crate) mod connection_lifecycle; pub(crate) mod connection_manager; pub(crate) mod elicitation; pub(crate) mod mcp; diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index a74f95bae3..fd641f865b 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -37,7 +37,8 @@ use rmcp::model::ReadResourceResult; use serde_json::Value; use crate::codex_apps::codex_apps_tools_cache_key; -use crate::connection_manager::McpConnectionManager; +use crate::connection_generation::McpConnectionStartParams; +use crate::connection_lifecycle::McpConnectionManager; use crate::runtime::McpRuntimeContext; use crate::server::EffectiveMcpServer; @@ -293,30 +294,30 @@ pub async fn read_mcp_resource( .await; let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - auth_statuses, - &config.approval_policy, - String::new(), + let manager = McpConnectionManager::start(McpConnectionStartParams { + mcp_servers, + store_mode: config.mcp_oauth_credentials_store_mode, + auth_entries: auth_statuses, + approval_policy: config.approval_policy.value(), + submit_id: String::new(), tx_event, - PermissionProfile::default(), + permission_profile: PermissionProfile::default(), runtime_context, - config.codex_home.clone(), - codex_apps_tools_cache_key(auth), + codex_home: config.codex_home.clone(), + codex_apps_tools_cache_key: codex_apps_tools_cache_key(auth), host_owned_codex_apps_enabled, - config.prefix_mcp_tool_names, - config.client_elicitation_capability.clone(), - tool_plugin_provenance(config), - auth, - /*elicitation_reviewer*/ None, - ) + prefix_mcp_tool_names: config.prefix_mcp_tool_names, + client_elicitation_capability: config.client_elicitation_capability.clone(), + tool_plugin_provenance: tool_plugin_provenance(config), + auth: auth.cloned(), + elicitation_reviewer: None, + }) .await; let result = manager .read_resource(server, ReadResourceRequestParams::new(uri)) .await; - cancel_token.cancel(); + manager.shutdown().await; result } @@ -363,24 +364,24 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - auth_status_entries.clone(), - &config.approval_policy, + let mcp_connection_manager = McpConnectionManager::start(McpConnectionStartParams { + mcp_servers, + store_mode: config.mcp_oauth_credentials_store_mode, + auth_entries: auth_status_entries.clone(), + approval_policy: config.approval_policy.value(), submit_id, tx_event, - PermissionProfile::default(), + permission_profile: PermissionProfile::default(), runtime_context, - config.codex_home.clone(), - codex_apps_tools_cache_key(auth), + codex_home: config.codex_home.clone(), + codex_apps_tools_cache_key: codex_apps_tools_cache_key(auth), host_owned_codex_apps_enabled, - config.prefix_mcp_tool_names, - config.client_elicitation_capability.clone(), + prefix_mcp_tool_names: config.prefix_mcp_tool_names, + client_elicitation_capability: config.client_elicitation_capability.clone(), tool_plugin_provenance, - auth, - /*elicitation_reviewer*/ None, - ) + auth: auth.cloned(), + elicitation_reviewer: None, + }) .await; let snapshot = collect_mcp_server_status_snapshot_from_manager( @@ -391,7 +392,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( ) .await; - cancel_token.cancel(); + mcp_connection_manager.shutdown().await; snapshot } diff --git a/codex-rs/core/src/codex_delegate_tests.rs b/codex-rs/core/src/codex_delegate_tests.rs index 7dabb09d33..3ad695ed30 100644 --- a/codex-rs/core/src/codex_delegate_tests.rs +++ b/codex-rs/core/src/codex_delegate_tests.rs @@ -2,7 +2,6 @@ use super::*; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_DECLINE_SYNTHETIC; use crate::mcp_tool_call::MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX; use async_channel::bounded; -use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::models::NetworkPermissions; use codex_protocol::models::ResponseItem; @@ -450,72 +449,3 @@ async fn delegated_mcp_guardian_abort_returns_synthetic_decline_answer() { }) ); } - -#[tokio::test] -async fn delegated_mcp_user_reviewer_waits_for_metadata_lookup() { - let (parent_session, parent_ctx, _rx_events) = - crate::session::tests::make_session_and_context_with_rx().await; - let pending_mcp_invocations = Arc::new(Mutex::new(HashMap::from([( - "call-1".to_string(), - McpInvocation { - server: CODEX_APPS_MCP_SERVER_NAME.to_string(), - tool: "dangerous_tool".to_string(), - arguments: None, - }, - )]))); - let cancel_token = CancellationToken::new(); - let manager = Arc::clone(&parent_session.services.mcp_connection_manager); - let (manager_locked_tx, manager_locked_rx) = std::sync::mpsc::sync_channel(0); - let (release_manager_tx, release_manager_rx) = std::sync::mpsc::sync_channel(0); - let manager_lock = tokio::task::spawn_blocking(move || { - let _manager_guard = manager.blocking_write(); - manager_locked_tx - .send(()) - .expect("manager lock receiver should remain open"); - release_manager_rx - .recv() - .expect("manager lock release sender should remain open"); - }); - manager_locked_rx - .recv_timeout(Duration::from_secs(1)) - .expect("manager write lock should be acquired"); - - let event = RequestUserInputEvent { - call_id: "call-1".to_string(), - turn_id: "child-turn-1".to_string(), - questions: vec![RequestUserInputQuestion { - id: format!("{MCP_TOOL_APPROVAL_QUESTION_ID_PREFIX}_call-1"), - header: "Approve app tool call?".to_string(), - question: "Allow this app tool?".to_string(), - is_other: false, - is_secret: false, - options: None, - }], - }; - let response = maybe_auto_review_mcp_request_user_input( - &parent_session, - &parent_ctx, - &pending_mcp_invocations, - &event, - &cancel_token, - ); - tokio::pin!(response); - assert!( - timeout(Duration::from_millis(100), &mut response) - .await - .is_err(), - "manual reviewer should wait for MCP metadata" - ); - release_manager_tx - .send(()) - .expect("manager lock holder should remain open"); - manager_lock - .await - .expect("manager lock task should not panic"); - assert_eq!( - timeout(Duration::from_secs(1), response) - .await - .expect("manual reviewer should finish after MCP metadata lookup"), - None - ); -} diff --git a/codex-rs/core/src/connectors.rs b/codex-rs/core/src/connectors.rs index 623cb44216..4c49454bdd 100644 --- a/codex-rs/core/src/connectors.rs +++ b/codex-rs/core/src/connectors.rs @@ -35,13 +35,11 @@ use codex_login::CodexAuth; use codex_login::default_client::originator; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::McpConnectionManager; +use codex_mcp::McpConnectionStartParams; use codex_mcp::McpRuntimeContext; use codex_mcp::ToolInfo; use codex_mcp::ToolPluginProvenance; use codex_mcp::codex_apps_tools_cache_key; -use codex_mcp::compute_auth_statuses; -use codex_mcp::effective_mcp_servers; -use codex_mcp::host_owned_codex_apps_enabled; const CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS: Duration = Duration::from_secs(30); @@ -264,47 +262,45 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( }); } - let mcp_config = mcp_manager.runtime_config(config).await; - let mut mcp_servers = effective_mcp_servers(&mcp_config, auth.as_ref()); - mcp_servers.retain(|name, _| name == CODEX_APPS_MCP_SERVER_NAME); - let host_owned_codex_apps_enabled = host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()); - if mcp_servers.is_empty() { + let mut resolved = mcp_manager.resolve_connections(config, auth.as_ref()).await; + resolved + .servers + .retain(|name, _| name == CODEX_APPS_MCP_SERVER_NAME); + if resolved.servers.is_empty() { return Ok(AccessibleConnectorsStatus { connectors: Vec::new(), codex_apps_ready: true, }); } - - let auth_status_entries = compute_auth_statuses( - mcp_servers.iter(), - config.mcp_oauth_credentials_store_mode, - auth.as_ref(), - ) - .await; + let codex_apps_startup_timeout = resolved + .servers + .get(CODEX_APPS_MCP_SERVER_NAME) + .and_then(|server| server.configured_config()) + .and_then(|config| config.startup_timeout_sec); let (tx_event, rx_event) = unbounded(); drop(rx_event); - let (mut mcp_connection_manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - auth_status_entries, - &config.permissions.approval_policy, - INITIAL_SUBMIT_ID.to_owned(), + let mcp_connection_manager = McpConnectionManager::start(McpConnectionStartParams { + mcp_servers: resolved.servers, + store_mode: resolved.store_mode, + auth_entries: resolved.auth_statuses, + approval_policy: config.permissions.approval_policy.value(), + submit_id: INITIAL_SUBMIT_ID.to_owned(), tx_event, - PermissionProfile::default(), + permission_profile: PermissionProfile::default(), // Connector discovery is threadless. Use an actually configured env if // one exists, but do not reintroduce the old hidden-local fallback. - McpRuntimeContext::new(environment_manager, config.cwd.to_path_buf()), - config.codex_home.to_path_buf(), - codex_apps_tools_cache_key(auth.as_ref()), - host_owned_codex_apps_enabled, - mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability, - ToolPluginProvenance::default(), - auth.as_ref(), - /*elicitation_reviewer*/ None, - ) + runtime_context: McpRuntimeContext::new(environment_manager, config.cwd.to_path_buf()), + codex_home: config.codex_home.to_path_buf(), + codex_apps_tools_cache_key: codex_apps_tools_cache_key(auth.as_ref()), + host_owned_codex_apps_enabled: resolved.host_owned_codex_apps_enabled, + prefix_mcp_tool_names: resolved.prefix_mcp_tool_names, + client_elicitation_capability: resolved.client_elicitation_capability, + tool_plugin_provenance: resolved.tool_plugin_provenance, + auth: auth.clone(), + elicitation_reviewer: None, + }) .await; let refreshed_tools = if force_refetch { @@ -333,17 +329,15 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( let mut should_reload_tools = false; let codex_apps_ready = if refreshed_tools_succeeded { true - } else if let Some(cfg) = mcp_servers.get(CODEX_APPS_MCP_SERVER_NAME) { + } else { let immediate_ready = mcp_connection_manager .wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, Duration::ZERO) .await; if immediate_ready { true } else if tools.is_empty() { - let timeout = cfg - .configured_config() - .and_then(|config| config.startup_timeout_sec) - .unwrap_or(CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS); + let timeout = + codex_apps_startup_timeout.unwrap_or(CONNECTORS_READY_TIMEOUT_ON_EMPTY_TOOLS); let ready = mcp_connection_manager .wait_for_server_ready(CODEX_APPS_MCP_SERVER_NAME, timeout) .await; @@ -352,14 +346,12 @@ pub async fn list_accessible_connectors_from_mcp_tools_with_mcp_manager( } else { false } - } else { - false }; if should_reload_tools { tools = mcp_connection_manager.list_all_tools().await; } if codex_apps_ready { - cancel_token.cancel(); + mcp_connection_manager.cancel_startup(); } let accessible_connectors = codex_connectors::filter::filter_disallowed_connectors( diff --git a/codex-rs/core/src/mcp.rs b/codex-rs/core/src/mcp.rs index 1b693d70fd..4601001a9a 100644 --- a/codex-rs/core/src/mcp.rs +++ b/codex-rs/core/src/mcp.rs @@ -3,17 +3,33 @@ use std::sync::Arc; use crate::config::Config; use codex_config::McpServerConfig; +use codex_config::types::OAuthCredentialsStoreMode; use codex_core_plugins::PluginsManager; use codex_extension_api::ExtensionRegistry; use codex_extension_api::McpServerContribution; use codex_login::CodexAuth; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::EffectiveMcpServer; +use codex_mcp::McpAuthStatusEntry; use codex_mcp::McpConfig; use codex_mcp::ToolPluginProvenance; +use codex_mcp::compute_auth_statuses; use codex_mcp::configured_mcp_servers; use codex_mcp::effective_mcp_servers; +use codex_mcp::effective_mcp_servers_from_configured; +use codex_mcp::host_owned_codex_apps_enabled; use codex_mcp::tool_plugin_provenance as collect_tool_plugin_provenance; +use rmcp::model::ElicitationCapability; + +pub(crate) struct ResolvedMcpConnections { + pub(crate) servers: HashMap, + pub(crate) store_mode: OAuthCredentialsStoreMode, + pub(crate) auth_statuses: HashMap, + pub(crate) host_owned_codex_apps_enabled: bool, + pub(crate) prefix_mcp_tool_names: bool, + pub(crate) client_elicitation_capability: ElicitationCapability, + pub(crate) tool_plugin_provenance: ToolPluginProvenance, +} #[derive(Clone)] pub struct McpManager { @@ -76,6 +92,54 @@ impl McpManager { effective_mcp_servers(&mcp_config, auth) } + pub(crate) async fn resolve_connections( + &self, + config: &Config, + auth: Option<&CodexAuth>, + ) -> ResolvedMcpConnections { + let mcp_config = self.runtime_config(config).await; + let servers = effective_mcp_servers(&mcp_config, auth); + self.resolve_connections_from_servers( + servers, + mcp_config.mcp_oauth_credentials_store_mode, + &mcp_config, + auth, + ) + .await + } + + pub(crate) async fn resolve_refreshed_connections( + &self, + config: &Config, + configured_servers: HashMap, + store_mode: OAuthCredentialsStoreMode, + auth: Option<&CodexAuth>, + ) -> ResolvedMcpConnections { + let mcp_config = self.runtime_config(config).await; + let servers = effective_mcp_servers_from_configured(configured_servers, &mcp_config, auth); + self.resolve_connections_from_servers(servers, store_mode, &mcp_config, auth) + .await + } + + async fn resolve_connections_from_servers( + &self, + servers: HashMap, + store_mode: OAuthCredentialsStoreMode, + mcp_config: &McpConfig, + auth: Option<&CodexAuth>, + ) -> ResolvedMcpConnections { + let auth_statuses = compute_auth_statuses(servers.iter(), store_mode, auth).await; + ResolvedMcpConnections { + host_owned_codex_apps_enabled: host_owned_codex_apps_enabled(mcp_config, auth), + prefix_mcp_tool_names: mcp_config.prefix_mcp_tool_names, + client_elicitation_capability: mcp_config.client_elicitation_capability.clone(), + tool_plugin_provenance: collect_tool_plugin_provenance(mcp_config), + servers, + store_mode, + auth_statuses, + } + } + /// Returns provenance for plugin-owned servers in the configured view. pub async fn tool_plugin_provenance(&self, config: &Config) -> ToolPluginProvenance { let mcp_config = config.to_mcp_config(self.plugins_manager.as_ref()).await; diff --git a/codex-rs/core/src/mcp_tool_call.rs b/codex-rs/core/src/mcp_tool_call.rs index b9a8d4ff08..3ccadac517 100644 --- a/codex-rs/core/src/mcp_tool_call.rs +++ b/codex-rs/core/src/mcp_tool_call.rs @@ -316,13 +316,7 @@ async fn handle_approved_mcp_tool_call( let arguments_value = invocation.arguments.clone(); let connector_id = metadata.and_then(|metadata| metadata.connector_id.as_deref()); let connector_name = metadata.and_then(|metadata| metadata.connector_name.as_deref()); - let server_origin = sess - .services - .mcp_connection_manager - .read() - .await - .server_origin(&server) - .map(str::to_string); + let server_origin = sess.services.mcp_connection_manager.server_origin(&server); let start = Instant::now(); let rewrite = rewrite_mcp_tool_arguments_for_openai_files( @@ -606,8 +600,6 @@ async fn maybe_request_codex_apps_auth_elicitation( if !sess .services .mcp_connection_manager - .read() - .await .is_host_owned_codex_apps_server(server) { return result; @@ -669,15 +661,12 @@ async fn maybe_request_codex_apps_auth_elicitation( auth_elicitation_completed_result(&plan.auth_failure, result.meta) } -#[expect( - clippy::await_holding_invalid_type, - reason = "Codex Apps cache refresh reads through the session-owned manager guard" -)] async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: &TurnContext) { - let mcp_tools_result = { - let manager = sess.services.mcp_connection_manager.read().await; - manager.hard_refresh_codex_apps_tools_cache().await - }; + let mcp_tools_result = sess + .services + .mcp_connection_manager + .hard_refresh_codex_apps_tools_cache() + .await; match mcp_tools_result { Ok(mcp_tools) => { @@ -694,10 +683,6 @@ async fn refresh_codex_apps_after_connector_auth(sess: &Session, turn_context: & } } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP sandbox metadata reads through the session-owned manager guard" -)] async fn augment_mcp_tool_request_meta_with_sandbox_state( sess: &Session, turn_context: &TurnContext, @@ -707,8 +692,6 @@ async fn augment_mcp_tool_request_meta_with_sandbox_state( let supports_sandbox_state_meta = sess .services .mcp_connection_manager - .read() - .await .server_supports_sandbox_state_meta_capability(server) .await .unwrap_or(false); @@ -757,8 +740,6 @@ async fn maybe_mark_thread_memory_mode_polluted( let pollutes_memory = sess .services .mcp_connection_manager - .read() - .await .server_pollutes_memory(server); if !pollutes_memory { return; @@ -1414,20 +1395,14 @@ async fn mcp_tool_approval_decision_from_guardian( } } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP approval metadata reads through the session-owned manager guard" -)] pub(crate) async fn lookup_mcp_tool_metadata( sess: &Session, turn_context: &TurnContext, server: &str, tool_name: &str, ) -> Option { - let manager = sess.services.mcp_connection_manager.read().await; - let plugin_id = manager - .plugin_id_for_mcp_server_name(server) - .map(str::to_string); + let manager = &sess.services.mcp_connection_manager; + let plugin_id = manager.plugin_id_for_mcp_server_name(server); let tools = manager.list_all_tools().await; let tool_info = tools .into_iter() @@ -1509,22 +1484,12 @@ fn get_mcp_app_resource_uri( }) } -#[expect( - clippy::await_holding_invalid_type, - reason = "MCP app metadata reads through the session-owned manager guard" -)] async fn lookup_mcp_app_usage_metadata( sess: &Session, server: &str, tool_name: &str, ) -> Option { - let tools = sess - .services - .mcp_connection_manager - .read() - .await - .list_all_tools() - .await; + let tools = sess.services.mcp_connection_manager.list_all_tools().await; tools.into_iter().find_map(|tool_info| { if tool_info.server_name == server && tool_info.tool.name == tool_name { diff --git a/codex-rs/core/src/mcp_tool_call_tests.rs b/codex-rs/core/src/mcp_tool_call_tests.rs index bd71c3a054..d1467a117b 100644 --- a/codex-rs/core/src/mcp_tool_call_tests.rs +++ b/codex-rs/core/src/mcp_tool_call_tests.rs @@ -1256,29 +1256,35 @@ fn codex_apps_auth_failure_metadata() -> McpToolApprovalMetadata { async fn install_host_owned_codex_apps_manager(session: &Session, turn_context: &TurnContext) { let auth = session.services.auth_manager.auth().await; - let (manager, _cancel_token) = codex_mcp::McpConnectionManager::new( - &HashMap::new(), - turn_context.config.mcp_oauth_credentials_store_mode, - HashMap::new(), - &turn_context.approval_policy, - turn_context.sub_id.clone(), - session.get_tx_event(), - turn_context.permission_profile(), - codex_mcp::McpRuntimeContext::new(Arc::clone(&session.services.environment_manager), { - #[allow(deprecated)] - turn_context.cwd.to_path_buf() - }), - turn_context.config.codex_home.to_path_buf(), - codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), - /*host_owned_codex_apps_enabled*/ true, - turn_context.config.prefix_mcp_tool_names(), - rmcp::model::ElicitationCapability::default(), - codex_mcp::ToolPluginProvenance::default(), - auth.as_ref(), - /*elicitation_reviewer*/ None, - ) - .await; - *session.services.mcp_connection_manager.write().await = manager; + let params = codex_mcp::McpConnectionStartParams { + mcp_servers: HashMap::new(), + store_mode: turn_context.config.mcp_oauth_credentials_store_mode, + auth_entries: HashMap::new(), + approval_policy: turn_context.approval_policy.value(), + submit_id: turn_context.sub_id.clone(), + tx_event: session.get_tx_event(), + permission_profile: turn_context.permission_profile(), + runtime_context: codex_mcp::McpRuntimeContext::new( + Arc::clone(&session.services.environment_manager), + { + #[allow(deprecated)] + turn_context.cwd.to_path_buf() + }, + ), + codex_home: turn_context.config.codex_home.to_path_buf(), + codex_apps_tools_cache_key: codex_mcp::codex_apps_tools_cache_key(auth.as_ref()), + host_owned_codex_apps_enabled: true, + prefix_mcp_tool_names: turn_context.config.prefix_mcp_tool_names(), + client_elicitation_capability: rmcp::model::ElicitationCapability::default(), + tool_plugin_provenance: codex_mcp::ToolPluginProvenance::default(), + auth, + elicitation_reviewer: None, + }; + session + .services + .mcp_connection_manager + .reconfigure(params) + .await; } #[tokio::test] diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index c0f62ec664..1ad8b64698 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -632,11 +632,7 @@ async fn shutdown_session_runtime(sess: &Arc) { if let Err(err) = sess.services.code_mode_service.shutdown().await { warn!("failed to shutdown code mode session: {err}"); } - let mcp_shutdown = { - let mut manager = sess.services.mcp_connection_manager.write().await; - manager.begin_shutdown() - }; - mcp_shutdown.await; + sess.services.mcp_connection_manager.shutdown().await; sess.guardian_review_session.shutdown().await; } diff --git a/codex-rs/core/src/session/mcp.rs b/codex-rs/core/src/session/mcp.rs index f2643d0e5c..668df78d41 100644 --- a/codex-rs/core/src/session/mcp.rs +++ b/codex-rs/core/src/session/mcp.rs @@ -1,7 +1,9 @@ use super::*; +use crate::mcp::ResolvedMcpConnections; use codex_mcp::ElicitationReviewRequest; use codex_mcp::ElicitationReviewer; use codex_mcp::ElicitationReviewerHandle; +use codex_mcp::McpConnectionStartParams; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_KEY as MCP_ELICITATION_APPROVAL_KIND_KEY; use codex_protocol::mcp_approval_meta::APPROVAL_KIND_MCP_TOOL_CALL as MCP_ELICITATION_APPROVAL_KIND_MCP_TOOL_CALL; @@ -78,6 +80,63 @@ impl Session { Arc::new(GuardianMcpElicitationReviewer::new(self)) } + pub(crate) async fn initialize_mcp_connections( + &self, + resolved: ResolvedMcpConnections, + auth: Option, + session_configuration: &SessionConfiguration, + runtime_context: McpRuntimeContext, + codex_home: PathBuf, + elicitation_reviewer: Option, + ) -> anyhow::Result<()> { + let mut required_servers = resolved + .servers + .iter() + .filter(|(_, server)| server.enabled() && server.required()) + .map(|(name, _)| name.clone()) + .collect::>(); + required_servers.sort(); + let required_server_count = required_servers.len(); + let params = mcp_connection_start_params( + resolved, + auth, + session_configuration.approval_policy.value(), + INITIAL_SUBMIT_ID.to_owned(), + self.get_tx_event(), + session_configuration.permission_profile(), + runtime_context, + codex_home, + elicitation_reviewer, + ); + self.services + .mcp_connection_manager + .initialize(params) + .await; + + if required_servers.is_empty() { + return Ok(()); + } + let failures = self + .services + .mcp_connection_manager + .required_startup_failures(&required_servers) + .instrument(info_span!( + "session_init.required_mcp_wait", + otel.name = "session_init.required_mcp_wait", + session_init.required_mcp_server_count = required_server_count, + )) + .await; + if failures.is_empty() { + return Ok(()); + } + let details = failures + .iter() + .map(|failure| format!("{}: {}", failure.server, failure.error)) + .collect::>() + .join("; "); + anyhow::bail!("required MCP servers failed to initialize: {details}") + } + #[expect( clippy::await_holding_invalid_type, reason = "active turn checks and turn state updates must remain atomic" @@ -91,8 +150,6 @@ impl Session { if self .services .mcp_connection_manager - .read() - .await .elicitations_auto_deny() { return McpServerElicitationOutcome { @@ -226,16 +283,10 @@ impl Session { self.services .mcp_connection_manager - .read() - .await .resolve_elicitation(server_name, id, response) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource calls are serialized through the session-owned manager guard" - )] pub async fn list_resources( &self, server: &str, @@ -243,16 +294,10 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await .list_resources(server, params) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource calls are serialized through the session-owned manager guard" - )] pub async fn list_resource_templates( &self, server: &str, @@ -260,16 +305,10 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await .list_resource_templates(server, params) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP resource calls are serialized through the session-owned manager guard" - )] pub async fn read_resource( &self, server: &str, @@ -277,16 +316,10 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await .read_resource(server, params) .await } - #[expect( - clippy::await_holding_invalid_type, - reason = "MCP tool calls are serialized through the session-owned manager guard" - )] pub async fn call_tool( &self, server: &str, @@ -296,8 +329,6 @@ impl Session { ) -> anyhow::Result { self.services .mcp_connection_manager - .read() - .await .call_tool(server, tool, arguments, meta) .await } @@ -311,22 +342,11 @@ impl Session { ) { let auth = self.services.auth_manager.auth().await; let config = self.get_config().await; - let mcp_config = self + let resolved = self .services .mcp_manager - .runtime_config(config.as_ref()) + .resolve_refreshed_connections(config.as_ref(), mcp_servers, store_mode, auth.as_ref()) .await; - let tool_plugin_provenance = self - .services - .mcp_manager - .tool_plugin_provenance(config.as_ref()) - .await; - let mcp_servers = - effective_mcp_servers_from_configured(mcp_servers, &mcp_config, auth.as_ref()); - let host_owned_codex_apps_enabled = - host_owned_codex_apps_enabled(&mcp_config, auth.as_ref()); - let auth_statuses = - compute_auth_statuses(mcp_servers.iter(), store_mode, auth.as_ref()).await; let mcp_runtime_context = match turn_context.environments.primary() { Some(turn_environment) => McpRuntimeContext::new( Arc::clone(&self.services.environment_manager), @@ -338,47 +358,21 @@ impl Session { turn_context.cwd.to_path_buf(), ), }; - { - let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - guard.cancel(); - *guard = CancellationToken::new(); - } - let (refreshed_manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - store_mode, - auth_statuses, - &turn_context.approval_policy, + let params = mcp_connection_start_params( + resolved, + auth, + turn_context.approval_policy.value(), turn_context.sub_id.clone(), self.get_tx_event(), turn_context.permission_profile(), mcp_runtime_context, config.codex_home.to_path_buf(), - codex_apps_tools_cache_key(auth.as_ref()), - host_owned_codex_apps_enabled, - mcp_config.prefix_mcp_tool_names, - mcp_config.client_elicitation_capability, - tool_plugin_provenance, - auth.as_ref(), elicitation_reviewer, - ) - .await; - { - let current_manager = self.services.mcp_connection_manager.read().await; - refreshed_manager.set_elicitations_auto_deny(current_manager.elicitations_auto_deny()); - } - { - let mut guard = self.services.mcp_startup_cancellation_token.lock().await; - if guard.is_cancelled() { - cancel_token.cancel(); - } - *guard = cancel_token; - } - - let mut old_manager = { - let mut manager = self.services.mcp_connection_manager.write().await; - std::mem::replace(&mut *manager, refreshed_manager) - }; - old_manager.shutdown().await; + ); + self.services + .mcp_connection_manager + .reconfigure(params) + .await; } pub(crate) async fn refresh_mcp_servers_if_requested( @@ -429,21 +423,39 @@ impl Session { .await; } - #[cfg(test)] - pub(crate) async fn mcp_startup_cancellation_token(&self) -> CancellationToken { - self.services - .mcp_startup_cancellation_token - .lock() - .await - .clone() - } - pub(crate) async fn cancel_mcp_startup(&self) { - self.services - .mcp_startup_cancellation_token - .lock() - .await - .cancel(); + self.services.mcp_connection_manager.cancel_startup(); + } +} + +fn mcp_connection_start_params( + resolved: ResolvedMcpConnections, + auth: Option, + approval_policy: AskForApproval, + submit_id: String, + tx_event: Sender, + permission_profile: PermissionProfile, + runtime_context: McpRuntimeContext, + codex_home: PathBuf, + elicitation_reviewer: Option, +) -> McpConnectionStartParams { + McpConnectionStartParams { + mcp_servers: resolved.servers, + store_mode: resolved.store_mode, + auth_entries: resolved.auth_statuses, + approval_policy, + submit_id, + tx_event, + permission_profile, + runtime_context, + codex_home, + codex_apps_tools_cache_key: codex_apps_tools_cache_key(auth.as_ref()), + host_owned_codex_apps_enabled: resolved.host_owned_codex_apps_enabled, + prefix_mcp_tool_names: resolved.prefix_mcp_tool_names, + client_elicitation_capability: resolved.client_elicitation_capability, + tool_plugin_provenance: resolved.tool_plugin_provenance, + auth, + elicitation_reviewer, } } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 23dbe1f45e..808296b783 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -146,15 +146,12 @@ use codex_utils_output_truncation::TruncationPolicy; use futures::future::BoxFuture; use futures::future::Shared; use futures::prelude::*; -use rmcp::model::ElicitationCapability; -use rmcp::model::FormElicitationCapability; use rmcp::model::ListResourceTemplatesResult; use rmcp::model::ListResourcesResult; use rmcp::model::PaginatedRequestParams; use rmcp::model::ReadResourceRequestParams; use rmcp::model::ReadResourceResult; use rmcp::model::RequestId; -use rmcp::model::UrlElicitationCapability; use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::RwLock; @@ -316,9 +313,6 @@ use crate::unified_exec::UnifiedExecProcessManager; use crate::windows_sandbox::WindowsSandboxLevelExt; use codex_core_plugins::PluginsManager; use codex_git_utils::get_git_repo_root; -use codex_mcp::compute_auth_statuses; -use codex_mcp::effective_mcp_servers_from_configured; -use codex_mcp::host_owned_codex_apps_enabled; use codex_otel::SessionTelemetry; use codex_otel::THREAD_STARTED_METRIC; use codex_otel::TelemetryAuthMode; @@ -793,8 +787,10 @@ impl Codex { ..Default::default() }) .await?; - let mcp_connection_manager = self.session.services.mcp_connection_manager.read().await; - mcp_connection_manager.set_elicitations_auto_deny(mcp_elicitations_auto_deny); + self.session + .services + .mcp_connection_manager + .set_elicitations_auto_deny(mcp_elicitations_auto_deny); Ok(()) } @@ -2845,10 +2841,9 @@ impl Session { } } if turn_context.config.include_apps_instructions && turn_context.apps_enabled() { - let mcp_connection_manager = self.services.mcp_connection_manager.read().await; let accessible_and_enabled_connectors = connectors::list_accessible_and_enabled_connectors_from_manager( - &mcp_connection_manager, + self.services.mcp_connection_manager.as_ref(), &turn_context.config, ) .await; diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index 50ff984231..b913b5d1f5 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -620,16 +620,10 @@ impl Session { let mcp_manager_for_mcp = Arc::clone(&mcp_manager); let auth_and_mcp_fut = async move { let auth = auth_manager_clone.auth().await; - let mcp_servers = mcp_manager_for_mcp - .effective_servers(&config_for_mcp, auth.as_ref()) + let resolved_mcp_connections = mcp_manager_for_mcp + .resolve_connections(&config_for_mcp, auth.as_ref()) .await; - let auth_statuses = compute_auth_statuses( - mcp_servers.iter(), - config_for_mcp.mcp_oauth_credentials_store_mode, - auth.as_ref(), - ) - .await; - (auth, mcp_servers, auth_statuses) + (auth, resolved_mcp_connections) } .instrument(info_span!( "session_init.auth_mcp", @@ -652,7 +646,7 @@ impl Session { let ( thread_persistence_result, state_db_ctx, - (auth, mcp_servers, auth_statuses), + (auth, resolved_mcp_connections), plugin_skill_errors, ) = tokio::join!( thread_persistence_fut, @@ -817,7 +811,11 @@ impl Session { config .permissions .legacy_sandbox_policy(session_configuration.cwd().as_path()), - mcp_servers.keys().map(String::as_str).collect(), + resolved_mcp_connections + .servers + .keys() + .map(String::as_str) + .collect(), ); let use_zsh_fork_shell = config.features.enabled(Feature::ShellZshFork); @@ -978,21 +976,11 @@ impl Session { } let services = SessionServices { - // Initialize the MCP connection manager with an uninitialized - // instance. It will be replaced with one created via - // McpConnectionManager::new() once all its constructor args are - // available. This also ensures `SessionConfigured` is emitted - // before any MCP-related events. It is reasonable to consider - // changing this to use Option or OnceCell, though the current - // setup is straightforward enough and performs well. - mcp_connection_manager: Arc::new(RwLock::new( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), + mcp_connection_manager: Arc::new(McpConnectionManager::empty( + config.permissions.approval_policy.value(), + config.permissions.permission_profile().clone(), + config.prefix_mcp_tool_names(), )), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -1116,32 +1104,16 @@ impl Session { sess.send_event_raw(event).await; } - let mut required_mcp_servers: Vec = mcp_servers + let enabled_mcp_server_count = resolved_mcp_connections + .servers + .values() + .filter(|server| server.enabled()) + .count(); + let required_mcp_server_count = resolved_mcp_connections + .servers .iter() .filter(|(_, server)| server.enabled() && server.required()) - .map(|(name, _)| name.clone()) - .collect(); - required_mcp_servers.sort(); - let enabled_mcp_server_count = - mcp_servers.values().filter(|server| server.enabled()).count(); - let required_mcp_server_count = required_mcp_servers.len(); - let tool_plugin_provenance = mcp_manager.tool_plugin_provenance(config.as_ref()).await; - let host_owned_codex_apps_enabled = config - .features - .apps_enabled_for_auth(auth.as_ref().is_some_and(|auth| auth.uses_codex_backend())); - let client_elicitation_capability = if config.features.enabled(Feature::AuthElicitation) { - ElicitationCapability { - form: Some(FormElicitationCapability::default()), - url: Some(UrlElicitationCapability::default()), - } - } else { - ElicitationCapability::default() - }; - { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - cancel_guard.cancel(); - *cancel_guard = CancellationToken::new(); - } + .count(); let turn_environment = crate::environment_selection::resolve_environment_selections( sess.services.environment_manager.as_ref(), session_configuration.environment_selections(), @@ -1164,22 +1136,12 @@ impl Session { session_configuration.cwd().to_path_buf(), ), }; - let (mcp_connection_manager, cancel_token) = McpConnectionManager::new( - &mcp_servers, - config.mcp_oauth_credentials_store_mode, - auth_statuses.clone(), - &session_configuration.approval_policy, - INITIAL_SUBMIT_ID.to_owned(), - tx_event.clone(), - session_configuration.permission_profile(), + sess.initialize_mcp_connections( + resolved_mcp_connections, + auth.cloned(), + &session_configuration, mcp_runtime_context, config.codex_home.to_path_buf(), - codex_apps_tools_cache_key(auth), - host_owned_codex_apps_enabled, - config.prefix_mcp_tool_names(), - client_elicitation_capability, - tool_plugin_provenance, - auth, Some(sess.mcp_elicitation_reviewer()), ) .instrument(info_span!( @@ -1188,40 +1150,7 @@ impl Session { session_init.enabled_mcp_server_count = enabled_mcp_server_count, session_init.required_mcp_server_count = required_mcp_server_count, )) - .await; - { - let mut manager_guard = sess.services.mcp_connection_manager.write().await; - *manager_guard = mcp_connection_manager; - } - { - let mut cancel_guard = sess.services.mcp_startup_cancellation_token.lock().await; - if cancel_guard.is_cancelled() { - cancel_token.cancel(); - } - *cancel_guard = cancel_token; - } - if !required_mcp_servers.is_empty() { - let failures = sess - .services - .mcp_connection_manager - .read() - .await - .required_startup_failures(&required_mcp_servers) - .instrument(info_span!( - "session_init.required_mcp_wait", - otel.name = "session_init.required_mcp_wait", - session_init.required_mcp_server_count = required_mcp_server_count, - )) - .await; - if !failures.is_empty() { - let details = failures - .iter() - .map(|failure| format!("{}: {}", failure.server, failure.error)) - .collect::>() - .join("; "); - anyhow::bail!("required MCP servers failed to initialize: {details}"); - } - } + .await?; sess.schedule_startup_prewarm(session_configuration.base_instructions.clone()) .await; let session_start_source = match &initial_history { diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index f6166260a0..6fcf6189f6 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -319,8 +319,6 @@ async fn request_mcp_server_elicitation_auto_accepts_when_auto_deny_is_enabled() session .services .mcp_connection_manager - .read() - .await .set_elicitations_auto_deny(/*auto_deny*/ true); let requested_schema: McpElicitationSchema = serde_json::from_value(json!({ @@ -4801,14 +4799,11 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { ); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), + mcp_connection_manager: Arc::new(McpConnectionManager::empty( + config.permissions.approval_policy.value(), + config.permissions.permission_profile().clone(), + config.prefix_mcp_tool_names(), )), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -6867,14 +6862,11 @@ where ); let services = SessionServices { - mcp_connection_manager: Arc::new(RwLock::new( - McpConnectionManager::new_uninitialized_with_permission_profile( - &config.permissions.approval_policy, - config.permissions.permission_profile(), - config.prefix_mcp_tool_names(), - ), + mcp_connection_manager: Arc::new(McpConnectionManager::empty( + config.permissions.approval_policy.value(), + config.permissions.permission_profile().clone(), + config.prefix_mcp_tool_names(), )), - mcp_startup_cancellation_token: Mutex::new(CancellationToken::new()), unified_exec_manager: UnifiedExecProcessManager::new( config.background_terminal_max_timeout, ), @@ -7025,8 +7017,6 @@ pub(crate) async fn make_session_and_context_with_rx() -> ( #[tokio::test] async fn refresh_mcp_servers_is_deferred_until_next_turn() { let (session, turn_context) = make_session_and_context().await; - let old_token = session.mcp_startup_cancellation_token().await; - assert!(!old_token.is_cancelled()); let mcp_oauth_credentials_store_mode = serde_json::to_value(OAuthCredentialsStoreMode::Auto).expect("serialize store mode"); @@ -7039,7 +7029,6 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { *guard = Some(refresh_config); } - assert!(!old_token.is_cancelled()); assert!( session .pending_mcp_server_refresh_config @@ -7052,7 +7041,6 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { .refresh_mcp_servers_if_requested(&turn_context, /*elicitation_reviewer*/ None) .await; - assert!(old_token.is_cancelled()); assert!( session .pending_mcp_server_refresh_config @@ -7060,8 +7048,6 @@ async fn refresh_mcp_servers_is_deferred_until_next_turn() { .await .is_none() ); - let new_token = session.mcp_startup_cancellation_token().await; - assert!(!new_token.is_cancelled()); } #[tokio::test] @@ -9405,18 +9391,12 @@ async fn abort_review_task_emits_exited_then_aborted_and_records_history() { } #[tokio::test] -#[expect( - clippy::await_holding_invalid_type, - reason = "test builds a router from session-owned MCP manager state" -)] async fn fatal_tool_error_stops_turn_and_reports_error() { let (session, turn_context, _rx) = make_session_and_context_with_rx().await; let tools = { session .services .mcp_connection_manager - .read() - .await .list_all_tools() .await }; diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index 959e28834a..c6a528b26e 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -464,8 +464,6 @@ async fn build_skills_and_plugins( match sess .services .mcp_connection_manager - .read() - .await .list_all_tools() .or_cancel(cancellation_token) .await @@ -1079,18 +1077,12 @@ pub(crate) async fn built_tools( turn_context: &TurnContext, cancellation_token: &CancellationToken, ) -> CodexResult> { - let mcp_connection_manager = sess - .services - .mcp_connection_manager - .read() - .instrument(trace_span!("read_mcp_connection_manager")) - .await; + let mcp_connection_manager = &sess.services.mcp_connection_manager; let has_mcp_servers = mcp_connection_manager.has_servers(); let all_mcp_tools = mcp_connection_manager .list_all_tools() .or_cancel(cancellation_token) .await?; - drop(mcp_connection_manager); let loaded_plugins = sess .services .plugins_manager diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index e3bd01081a..fc9e149280 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -723,12 +723,12 @@ impl Session { .map(|turn_environment| turn_environment.cwd.clone()) .unwrap_or_else(|| session_configuration.cwd().clone()); let per_turn_config = Self::build_per_turn_config(&session_configuration, cwd.clone()); - { - let mcp_connection_manager = self.services.mcp_connection_manager.read().await; - mcp_connection_manager.set_approval_policy(&session_configuration.approval_policy); - mcp_connection_manager - .set_permission_profile(session_configuration.permission_profile()); - } + self.services + .mcp_connection_manager + .set_approval_policy(session_configuration.approval_policy.value()); + self.services + .mcp_connection_manager + .set_permission_profile(session_configuration.permission_profile()); let model_info = self .services diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 59def9c7ab..af82a7c607 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -34,13 +34,10 @@ use codex_thread_store::ThreadStore; use std::path::PathBuf; use tokio::runtime::Handle; use tokio::sync::Mutex; -use tokio::sync::RwLock; use tokio::sync::watch; -use tokio_util::sync::CancellationToken; pub(crate) struct SessionServices { - pub(crate) mcp_connection_manager: Arc>, - pub(crate) mcp_startup_cancellation_token: Mutex, + pub(crate) mcp_connection_manager: Arc, pub(crate) unified_exec_manager: UnifiedExecProcessManager, #[cfg_attr(not(unix), allow(dead_code))] pub(crate) shell_zsh_path: Option, diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs index b1e82aa8a6..d15ee52f59 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resource_templates.rs @@ -107,8 +107,6 @@ impl ToolExecutor for ListMcpResourceTemplatesHandler { let templates = session .services .mcp_connection_manager - .read() - .await .list_all_resource_templates() .await; Ok(ListResourceTemplatesPayload::from_all_servers(templates)) diff --git a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs index c2cdec8b15..d11c065ad7 100644 --- a/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs +++ b/codex-rs/core/src/tools/handlers/mcp_resource/list_mcp_resources.rs @@ -105,8 +105,6 @@ impl ToolExecutor for ListMcpResourcesHandler { let resources = session .services .mcp_connection_manager - .read() - .await .list_all_resources() .await; Ok(ListResourcesPayload::from_all_servers(resources)) diff --git a/codex-rs/core/src/tools/handlers/request_plugin_install.rs b/codex-rs/core/src/tools/handlers/request_plugin_install.rs index 3f7bd92d25..464b7de8ff 100644 --- a/codex-rs/core/src/tools/handlers/request_plugin_install.rs +++ b/codex-rs/core/src/tools/handlers/request_plugin_install.rs @@ -304,10 +304,6 @@ fn is_remote_plugin_install_suggestion(plugin_id: &str) -> bool { .is_some_and(|(_, marketplace_name)| marketplace_name == REMOTE_GLOBAL_MARKETPLACE_NAME) } -#[expect( - clippy::await_holding_invalid_type, - reason = "connector cache refresh reads through the session-owned manager guard" -)] async fn refresh_missing_requested_connectors( session: &crate::session::session::Session, turn: &crate::session::turn_context::TurnContext, @@ -319,7 +315,7 @@ async fn refresh_missing_requested_connectors( return Some(Vec::new()); } - let manager = session.services.mcp_connection_manager.read().await; + let manager = &session.services.mcp_connection_manager; let mcp_tools = manager.list_all_tools().await; let accessible_connectors = connectors::with_app_enabled_state( connectors::accessible_connectors_from_mcp_tools(&mcp_tools), diff --git a/codex-rs/core/src/tools/router_tests.rs b/codex-rs/core/src/tools/router_tests.rs index cece7fc171..d6a935036e 100644 --- a/codex-rs/core/src/tools/router_tests.rs +++ b/codex-rs/core/src/tools/router_tests.rs @@ -95,17 +95,11 @@ fn extension_tool_test_registry() -> Arc> { } #[tokio::test] -#[expect( - clippy::await_holding_invalid_type, - reason = "test builds a router from session-owned MCP manager state" -)] async fn parallel_support_does_not_match_namespaced_local_tool_names() -> anyhow::Result<()> { let (session, turn) = make_session_and_context().await; let mcp_tools = session .services .mcp_connection_manager - .read() - .await .list_all_tools() .await; let router = ToolRouter::from_turn_context(