From 51b3cd51f6f94488c0e05564cbcad9512f73e3db Mon Sep 17 00:00:00 2001 From: Ahmed Ibrahim Date: Tue, 9 Jun 2026 16:07:34 -0700 Subject: [PATCH] [codex] Tighten MCP connection manager API visibility and order (#27257) ## Summary - order `McpConnectionManager` methods by visibility, with the primary constructor and public API first - restrict `list_available_server_infos` to `codex-mcp` - make `new_uninitialized` a private test-only helper ## Why The manager exposed methods that are only used inside `codex-mcp` or its unit tests. Tightening those methods keeps the exported API intentional, while the new ordering makes the supported surface easier to scan. ## Validation - `just fmt` - `git diff --check` - local tests not run; relying on CI --- codex-rs/codex-mcp/src/connection_manager.rs | 273 ++++++++++--------- 1 file changed, 137 insertions(+), 136 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index d15a2e150..5558916e5 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -113,102 +113,6 @@ pub struct McpConnectionManager { } 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() - } - - /// 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; - } - } - } - - /// 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 - .get(server_name) - .and_then(|metadata| metadata.origin.as_ref()) - .map(super::server::McpServerOrigin::as_str) - } - - pub fn server_pollutes_memory(&self, server_name: &str) -> bool { - self.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 - .plugin_id_for_mcp_server_name(server_name) - } - - 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 - } - - 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_permission_profile(&self, permission_profile: PermissionProfile) { - 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() - } - - 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, @@ -364,6 +268,90 @@ impl McpConnectionManager { (manager, cancel_token) } + 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() + } + + /// 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; + } + } + } + + /// 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 + .get(server_name) + .and_then(|metadata| metadata.origin.as_ref()) + .map(super::server::McpServerOrigin::as_str) + } + + pub fn server_pollutes_memory(&self, server_name: &str) -> bool { + self.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 + .plugin_id_for_mcp_server_name(server_name) + } + + 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 + } + + 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_permission_profile(&self, permission_profile: PermissionProfile) { + 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() + } + + pub fn set_elicitations_auto_deny(&self, auto_deny: bool) { + self.elicitation_requests.set_auto_deny(auto_deny); + } + pub async fn resolve_elicitation( &self, server_name: String, @@ -452,31 +440,6 @@ impl McpConnectionManager { normalize_tools_for_model_with_prefix(tools, self.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 mut server_infos = HashMap::new(); - for (server_name, client) in &self.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); - } - continue; - } - match client.client().await { - Ok(managed_client) => { - server_infos.insert(server_name.clone(), managed_client.server_info); - } - Err(_) => { - if let Some(server_info) = client.cached_server_info.clone() { - server_infos.insert(server_name.clone(), server_info); - } - } - } - } - server_infos - } - /// Force-refresh codex apps tools by bypassing the in-process cache. /// /// On success, the refreshed tools replace the cache contents and the @@ -532,21 +495,6 @@ impl McpConnectionManager { )) } - fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { - let Some(metadata) = self.server_metadata.get(&tool.server_name) else { - tool.supports_parallel_tool_calls = false; - tool.server_origin = None; - return tool; - }; - - tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; - tool.server_origin = metadata - .origin - .as_ref() - .map(|origin| origin.as_str().to_string()); - tool - } - /// 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> { @@ -778,6 +726,46 @@ impl McpConnectionManager { .with_context(|| format!("resources/read failed for `{server}` ({uri})")) } + /// 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(crate) async fn list_available_server_infos(&self) -> HashMap { + let mut server_infos = HashMap::new(); + for (server_name, client) in &self.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); + } + continue; + } + match client.client().await { + Ok(managed_client) => { + server_infos.insert(server_name.clone(), managed_client.server_info); + } + Err(_) => { + if let Some(server_info) = client.cached_server_info.clone() { + server_infos.insert(server_name.clone(), server_info); + } + } + } + } + server_infos + } + + fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { + let Some(metadata) = self.server_metadata.get(&tool.server_name) else { + tool.supports_parallel_tool_calls = false; + tool.server_origin = None; + return tool; + }; + + tool.supports_parallel_tool_calls = metadata.supports_parallel_tool_calls; + tool.server_origin = metadata + .origin + .as_ref() + .map(|origin| origin.as_str().to_string()); + tool + } + async fn client_by_name(&self, name: &str) -> Result { self.clients .get(name) @@ -786,6 +774,19 @@ impl McpConnectionManager { .await .context("failed to get client") } + + #[cfg(test)] + 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, + ) + } } impl Drop for McpConnectionManager {