openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-19 17:28:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Split MCP connection manager into focused modules (#34522)

Browse Source 
## What changed

- Move required-server startup validation into `connection_manager/required.rs`.
- Move tool listing, lookup, metadata attachment, and Codex Apps cache refresh into `connection_manager/tool_catalog.rs`.
- Keep the existing connection-manager API and behavior unchanged.

GitOrigin-RevId: 7249896699f756bdbb6e00bc2bbc4a98098741a9
This commit is contained in:
jif
2026-07-21 11:45:36 +00:00
committed by copyberry
parent b9800de486
commit 2d85e6d3a6
4 changed files with 232 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-mcp/src/connection_manager.rs
+5 -207
View File
@@ -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::<Vec<_>>()
.join("; ");
Err(anyhow!(
"required MCP servers failed to initialize: {details}"
))
}
pub fn new_uninitialized_with_permission_profile(
approval_policy: &Constrained<AskForApproval>,
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<ToolInfo> {
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<ToolInfo> {
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<Vec<ToolInfo>> {
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<ManagedClient> {
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"),
codex-rs/codex-mcp/src/connection_manager/required.rs
+63
View File
@@ -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::<Vec<_>>()
.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,
}
}
codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs
+163
View File
@@ -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<ToolInfo> {
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<ToolInfo> {
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<Vec<ToolInfo>> {
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
}
}
codex-rs/codex-mcp/src/connection_manager_tests.rs
+1
View File
@@ -1,3 +1,4 @@
use super::required::startup_outcome_error_message;
use super::*;
use crate::elicitation::ElicitationLifecycle;
use crate::elicitation::ElicitationRequestManager;