diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 58e1d3b23d..19ad27a1b3 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -6,34 +6,31 @@ //! calls to the right client, and exposes the public manager API used by //! `codex-core`. +#[path = "connection_manager/required.rs"] +mod required; +#[path = "connection_manager/tool_catalog.rs"] +mod tool_catalog; + 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::codex_apps::prepare_openai_file_params_for_model; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; use crate::elicitation::ElicitationReviewerHandle; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; use crate::rmcp_client::AsyncManagedClient; -use crate::rmcp_client::CODEX_APPS_REFRESH_DURATION_METRIC; use crate::rmcp_client::DEFAULT_STARTUP_TIMEOUT; -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::tool_catalog_cache::McpToolCatalogCache; use crate::tools::ToolInfo; -use crate::tools::filter_tools; -use crate::tools::normalize_tools_for_model_with_prefix; use anyhow::Context; use anyhow::Result; use anyhow::anyhow; @@ -46,7 +43,6 @@ use codex_config::McpServerTransportConfig; use codex_config::types::AuthKeyringBackendKind; use codex_config::types::OAuthCredentialsStoreMode; use codex_connectors::ConnectorRuntimeContextKey; -use codex_connectors::ConnectorRuntimeFetchSource; use codex_connectors::ConnectorRuntimeManager; use codex_login::AuthManager; use codex_login::CodexAuth; @@ -77,11 +73,6 @@ use rmcp::model::ResourceTemplate; use serde_json::Value as JsonValue; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; -use tracing::Instrument; -use tracing::info_span; -use tracing::instrument; -use tracing::trace; -use tracing::trace_span; use tracing::warn; const MCP_UI_META_KEY: &str = "ui"; @@ -392,52 +383,6 @@ impl McpConnectionManager { manager } - /// Waits for every required server and reports their startup failures together. - /// - /// Callers must make the manager reachable to request handlers before awaiting this method, - /// because server initialization may require client elicitation. - pub async fn validate_required_servers(&self) -> Result<()> { - let failures = async { - let mut failures = Vec::new(); - for server_name in &self.required_servers { - let Some(async_managed_client) = self.clients.get(server_name).cloned() else { - failures.push(McpStartupFailure { - server: server_name.clone(), - error: format!("required MCP server `{server_name}` was not initialized"), - }); - continue; - }; - - match async_managed_client.client().await { - Ok(_) => {} - Err(error) => failures.push(McpStartupFailure { - server: server_name.clone(), - error: startup_outcome_error_message(error), - }), - } - } - failures - } - .instrument(info_span!( - "session_init.required_mcp_wait", - otel.name = "session_init.required_mcp_wait", - session_init.required_mcp_server_count = self.required_servers.len(), - )) - .await; - if failures.is_empty() { - return Ok(()); - } - - let details = failures - .iter() - .map(|failure| format!("{}: {}", failure.server, failure.error)) - .collect::>() - .join("; "); - Err(anyhow!( - "required MCP servers failed to initialize: {details}" - )) - } - pub fn new_uninitialized_with_permission_profile( approval_policy: &Constrained, permission_profile: &PermissionProfile, @@ -573,131 +518,6 @@ impl McpConnectionManager { } } - /// Returns all tools with model-visible names normalized. - #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] - pub async fn list_all_tools(&self) -> Vec { - let mut tools = Vec::new(); - let mut available_server_count = 0; - let mut unavailable_server_count = 0; - for (server_name, managed_client) in &self.clients { - managed_client.reconnect_failed_startup().await; - let has_cached_tools = managed_client.has_cached_tools(); - let startup_complete = managed_client - .startup_complete - .load(std::sync::atomic::Ordering::Acquire); - let Some(server_tools) = managed_client - .listed_tools() - .instrument(trace_span!( - "list_tools_for_server", - server_name = %server_name, - has_cached_tools, - startup_complete - )) - .await - else { - unavailable_server_count += 1; - trace!( - server_name = %server_name, - has_cached_tools, - startup_complete, - "MCP server tools unavailable while building tool list" - ); - continue; - }; - available_server_count += 1; - tools.extend( - server_tools - .into_iter() - .map(|tool| self.with_server_metadata(tool)), - ); - } - let tools = normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names); - trace!( - available_server_count, - unavailable_server_count, - tool_count = tools.len(), - "built MCP tool list" - ); - tools - } - - /// Returns one tool from the current live connection. - pub async fn tool_info(&self, server: &str, tool: &str) -> Option { - let client = self.clients.get(server)?; - let managed_client = client.client().await.ok()?; - let tool = client - .prepare_tools(managed_client.listed_tools()) - .into_iter() - .find(|tool_info| tool_info.tool.name == tool)?; - Some(self.with_server_metadata(tool)) - } - - /// Force-refresh codex apps tools by bypassing the in-process cache. - /// - /// On success, the refreshed tools replace shared cache contents when the - /// cache is enabled and the latest filtered tools are returned directly to - /// the caller. On failure, existing shared cache contents remain unchanged. - pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { - let refresh_start = Instant::now(); - let managed_client = self - .clients - .get(CODEX_APPS_MCP_SERVER_NAME) - .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))? - .client() - .await - .context("failed to get client")?; - - let list_start = Instant::now(); - let fetch_ticket = - managed_client - .codex_apps_tools_cache_context - .as_ref() - .map(|cache_context| { - cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh) - }); - let tools = list_tools_for_client_uncached( - CODEX_APPS_MCP_SERVER_NAME, - /*is_codex_apps_mcp_server*/ true, - /*codex_apps_refresh_trigger*/ "explicit", - &managed_client.client, - managed_client.tool_timeout, - managed_client.server_instructions.as_deref(), - ) - .await - .with_context(|| { - format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'") - })?; - - let tools = - match ( - managed_client.codex_apps_tools_cache_context.as_ref(), - fetch_ticket, - ) { - (Some(cache_context), Some(fetch_ticket)) => cache_context - .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), - (None, None) => tools, - _ => unreachable!("Codex Apps fetch ticket requires cache context"), - }; - emit_duration( - MCP_TOOLS_LIST_DURATION_METRIC, - list_start.elapsed(), - &[("cache", "miss")], - ); - let tools = filter_tools(tools, &managed_client.tool_filter) - .into_iter() - .map(|mut tool| { - prepare_openai_file_params_for_model(&mut tool); - self.with_server_metadata(tool) - }); - let tools = normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names); - emit_duration( - CODEX_APPS_REFRESH_DURATION_METRIC, - refresh_start.elapsed(), - &[("path", "legacy"), ("trigger", "explicit")], - ); - Ok(tools) - } - /// Returns resources from servers selected by `include_server`. Each key /// is the server name and the value is a vector of resources. pub async fn list_all_resources( @@ -967,21 +787,6 @@ impl McpConnectionManager { 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) @@ -1101,13 +906,6 @@ fn mcp_init_error_display( } } -fn startup_outcome_error_message(error: StartupOutcomeError) -> String { - match error { - StartupOutcomeError::Cancelled => "MCP startup cancelled".to_string(), - StartupOutcomeError::Failed { error, .. } => error, - } -} - fn is_mcp_client_auth_required_error(error: &StartupOutcomeError) -> bool { match error { StartupOutcomeError::Failed { error, .. } => error.contains("Auth required"), diff --git a/codex-rs/codex-mcp/src/connection_manager/required.rs b/codex-rs/codex-mcp/src/connection_manager/required.rs new file mode 100644 index 0000000000..73d38e4f5a --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_manager/required.rs @@ -0,0 +1,63 @@ +use anyhow::Result; +use anyhow::anyhow; +use codex_protocol::protocol::McpStartupFailure; +use tracing::Instrument; +use tracing::info_span; + +use super::McpConnectionManager; +use crate::rmcp_client::StartupOutcomeError; + +impl McpConnectionManager { + /// Waits for every required server and reports their startup failures together. + /// + /// Callers must make the manager reachable to request handlers before awaiting this method, + /// because server initialization may require client elicitation. + pub async fn validate_required_servers(&self) -> Result<()> { + let failures = async { + let mut failures = Vec::new(); + for server_name in &self.required_servers { + let Some(async_managed_client) = self.clients.get(server_name).cloned() else { + failures.push(McpStartupFailure { + server: server_name.clone(), + error: format!("required MCP server `{server_name}` was not initialized"), + }); + continue; + }; + + match async_managed_client.client().await { + Ok(_) => {} + Err(error) => failures.push(McpStartupFailure { + server: server_name.clone(), + error: startup_outcome_error_message(error), + }), + } + } + failures + } + .instrument(info_span!( + "session_init.required_mcp_wait", + otel.name = "session_init.required_mcp_wait", + session_init.required_mcp_server_count = self.required_servers.len(), + )) + .await; + if failures.is_empty() { + return Ok(()); + } + + let details = failures + .iter() + .map(|failure| format!("{}: {}", failure.server, failure.error)) + .collect::>() + .join("; "); + Err(anyhow!( + "required MCP servers failed to initialize: {details}" + )) + } +} + +pub(super) fn startup_outcome_error_message(error: StartupOutcomeError) -> String { + match error { + StartupOutcomeError::Cancelled => "MCP startup cancelled".to_string(), + StartupOutcomeError::Failed { error, .. } => error, + } +} diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs new file mode 100644 index 0000000000..7d79765ad0 --- /dev/null +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -0,0 +1,163 @@ +use std::time::Instant; + +use anyhow::Context; +use anyhow::Result; +use anyhow::anyhow; +use codex_connectors::ConnectorRuntimeFetchSource; +use tracing::Instrument; +use tracing::instrument; +use tracing::trace; +use tracing::trace_span; + +use super::McpConnectionManager; +use crate::codex_apps::prepare_openai_file_params_for_model; +use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; +use crate::rmcp_client::CODEX_APPS_REFRESH_DURATION_METRIC; +use crate::rmcp_client::MCP_TOOLS_LIST_DURATION_METRIC; +use crate::rmcp_client::list_tools_for_client_uncached; +use crate::runtime::emit_duration; +use crate::tools::ToolInfo; +use crate::tools::filter_tools; +use crate::tools::normalize_tools_for_model_with_prefix; + +impl McpConnectionManager { + /// Returns all tools with model-visible names normalized. + #[instrument(level = "trace", skip_all, fields(mcp_server_count = self.clients.len()))] + pub async fn list_all_tools(&self) -> Vec { + let mut tools = Vec::new(); + let mut available_server_count = 0; + let mut unavailable_server_count = 0; + for (server_name, managed_client) in &self.clients { + managed_client.reconnect_failed_startup().await; + let has_cached_tools = managed_client.has_cached_tools(); + let startup_complete = managed_client + .startup_complete + .load(std::sync::atomic::Ordering::Acquire); + let Some(server_tools) = managed_client + .listed_tools() + .instrument(trace_span!( + "list_tools_for_server", + server_name = %server_name, + has_cached_tools, + startup_complete + )) + .await + else { + unavailable_server_count += 1; + trace!( + server_name = %server_name, + has_cached_tools, + startup_complete, + "MCP server tools unavailable while building tool list" + ); + continue; + }; + available_server_count += 1; + tools.extend( + server_tools + .into_iter() + .map(|tool| self.with_server_metadata(tool)), + ); + } + let tools = normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names); + trace!( + available_server_count, + unavailable_server_count, + tool_count = tools.len(), + "built MCP tool list" + ); + tools + } + + /// Returns one tool from the current live connection. + pub async fn tool_info(&self, server: &str, tool: &str) -> Option { + let client = self.clients.get(server)?; + let managed_client = client.client().await.ok()?; + let tool = client + .prepare_tools(managed_client.listed_tools()) + .into_iter() + .find(|tool_info| tool_info.tool.name == tool)?; + Some(self.with_server_metadata(tool)) + } + + /// Force-refresh codex apps tools by bypassing the in-process cache. + /// + /// On success, the refreshed tools replace shared cache contents when the + /// cache is enabled and the latest filtered tools are returned directly to + /// the caller. On failure, existing shared cache contents remain unchanged. + pub async fn hard_refresh_codex_apps_tools_cache(&self) -> Result> { + let refresh_start = Instant::now(); + let managed_client = self + .clients + .get(CODEX_APPS_MCP_SERVER_NAME) + .ok_or_else(|| anyhow!("unknown MCP server '{CODEX_APPS_MCP_SERVER_NAME}'"))? + .client() + .await + .context("failed to get client")?; + + let list_start = Instant::now(); + let fetch_ticket = + managed_client + .codex_apps_tools_cache_context + .as_ref() + .map(|cache_context| { + cache_context.begin_fetch(ConnectorRuntimeFetchSource::HardRefresh) + }); + let tools = list_tools_for_client_uncached( + CODEX_APPS_MCP_SERVER_NAME, + /*is_codex_apps_mcp_server*/ true, + /*codex_apps_refresh_trigger*/ "explicit", + &managed_client.client, + managed_client.tool_timeout, + managed_client.server_instructions.as_deref(), + ) + .await + .with_context(|| { + format!("failed to refresh tools for MCP server '{CODEX_APPS_MCP_SERVER_NAME}'") + })?; + + let tools = + match ( + managed_client.codex_apps_tools_cache_context.as_ref(), + fetch_ticket, + ) { + (Some(cache_context), Some(fetch_ticket)) => cache_context + .publish_if_newest_accepted(fetch_ticket, &managed_client.server_info, tools), + (None, None) => tools, + _ => unreachable!("Codex Apps fetch ticket requires cache context"), + }; + emit_duration( + MCP_TOOLS_LIST_DURATION_METRIC, + list_start.elapsed(), + &[("cache", "miss")], + ); + let tools = filter_tools(tools, &managed_client.tool_filter) + .into_iter() + .map(|mut tool| { + prepare_openai_file_params_for_model(&mut tool); + self.with_server_metadata(tool) + }); + let tools = normalize_tools_for_model_with_prefix(tools, self.prefix_mcp_tool_names); + emit_duration( + CODEX_APPS_REFRESH_DURATION_METRIC, + refresh_start.elapsed(), + &[("path", "legacy"), ("trigger", "explicit")], + ); + Ok(tools) + } + + fn with_server_metadata(&self, mut tool: ToolInfo) -> ToolInfo { + let Some(metadata) = self.server_metadata.get(&tool.server_name) else { + tool.supports_parallel_tool_calls = false; + tool.server_origin = None; + return tool; + }; + + 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 + } +} diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index 18e8afe86e..88297a5592 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -1,3 +1,4 @@ +use super::required::startup_outcome_error_message; use super::*; use crate::elicitation::ElicitationLifecycle; use crate::elicitation::ElicitationRequestManager;