mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
Add single-server MCP status discovery with thread connection reuse (#48783)
## Why Inspecting one MCP server should not require full-inventory discovery or a separate connection when a thread already has one. ## What changed - Add optional `serverName` to `mcpServerStatus/list` and update generated schemas and bindings. - With `threadId`, apply any pending runtime refresh and reuse the thread's current connection and tool catalog, waiting only for the selected server during discovery. - Without `threadId`, create a discovery connection for only the selected server. Unknown names return an empty page; omitting `serverName` preserves full-inventory discovery. ## Testing Add a regression test covering both status detail modes, checking that tool metadata is preserved and that status reads do not repeat `initialize` or `tools/list` requests. GitOrigin-RevId: 1b361b67a993ab45cae493e88c146f50cfaf753a
This commit is contained in:
@@ -1049,9 +1049,12 @@ impl McpConnectionSet {
|
||||
/// Returns presentation metadata from the current connection.
|
||||
/// Codex Apps metadata may come from its existing cache; regular MCP server information is
|
||||
/// connection-specific, so pending regular clients are awaited.
|
||||
pub(crate) async fn list_available_server_infos(&self) -> HashMap<String, McpServerInfo> {
|
||||
pub(crate) async fn list_available_server_infos(
|
||||
&self,
|
||||
include_server: impl Fn(&str) -> bool,
|
||||
) -> HashMap<String, McpServerInfo> {
|
||||
let mut server_infos = HashMap::new();
|
||||
for (server_name, view) in &self.servers {
|
||||
for (server_name, view) in self.servers.iter().filter(|(name, _)| include_server(name)) {
|
||||
let client = &view.connection.client;
|
||||
if !client.startup_complete.load(Ordering::Acquire)
|
||||
&& let Some(server_info) = client.cached_server_info.clone()
|
||||
|
||||
@@ -179,49 +179,57 @@ impl McpConnectionSet {
|
||||
|
||||
/// Returns all tools with model-visible names normalized.
|
||||
pub async fn list_all_tools(&self) -> Vec<ToolInfo> {
|
||||
self.list_tools_with_errors().await.0
|
||||
Box::pin(self.list_tools_with_errors(|_| true)).await.0
|
||||
}
|
||||
|
||||
#[instrument(level = "trace", skip_all, fields(mcp_server_count = self.servers.len()))]
|
||||
pub(crate) async fn list_tools_with_errors(&self) -> (Vec<ToolInfo>, HashMap<String, String>) {
|
||||
pub(crate) async fn list_tools_with_errors(
|
||||
&self,
|
||||
include_server: impl Fn(&str) -> bool,
|
||||
) -> (Vec<ToolInfo>, HashMap<String, String>) {
|
||||
let mut tools = Vec::new();
|
||||
let mut errors = HashMap::new();
|
||||
let mut available_server_count = 0;
|
||||
let mut unavailable_server_count = 0;
|
||||
let server_results = join_all(self.servers.iter().map(|(server_name, view)| async move {
|
||||
view.connection.client.reconnect_failed_startup().await;
|
||||
let has_cached_tools = view.connection.client.has_cached_tools();
|
||||
let startup_complete = view
|
||||
.connection
|
||||
.client
|
||||
.startup_complete
|
||||
.load(Ordering::Acquire);
|
||||
let server_tools = view
|
||||
.listed_tools(&self.tool_plugin_context)
|
||||
.instrument(trace_span!(
|
||||
"list_tools_for_server",
|
||||
server_name = %server_name,
|
||||
has_cached_tools,
|
||||
startup_complete
|
||||
))
|
||||
.await;
|
||||
let result = match server_tools {
|
||||
Ok(server_tools) => Ok(server_tools
|
||||
.into_iter()
|
||||
.map(|tool| Self::with_server_metadata(tool, &view.metadata))
|
||||
.collect::<Vec<_>>()),
|
||||
Err(error) => {
|
||||
trace!(
|
||||
server_name = %server_name,
|
||||
has_cached_tools,
|
||||
startup_complete,
|
||||
"MCP server tools unavailable while building tool list"
|
||||
);
|
||||
Err(error)
|
||||
}
|
||||
};
|
||||
(server_name, result)
|
||||
}))
|
||||
let server_results = join_all(
|
||||
self.servers
|
||||
.iter()
|
||||
.filter(|(name, _)| include_server(name))
|
||||
.map(|(server_name, view)| async move {
|
||||
view.connection.client.reconnect_failed_startup().await;
|
||||
let has_cached_tools = view.connection.client.has_cached_tools();
|
||||
let startup_complete = view
|
||||
.connection
|
||||
.client
|
||||
.startup_complete
|
||||
.load(Ordering::Acquire);
|
||||
let server_tools = view
|
||||
.listed_tools(&self.tool_plugin_context)
|
||||
.instrument(trace_span!(
|
||||
"list_tools_for_server",
|
||||
server_name = %server_name,
|
||||
has_cached_tools,
|
||||
startup_complete
|
||||
))
|
||||
.await;
|
||||
let result = match server_tools {
|
||||
Ok(server_tools) => Ok(server_tools
|
||||
.into_iter()
|
||||
.map(|tool| Self::with_server_metadata(tool, &view.metadata))
|
||||
.collect::<Vec<_>>()),
|
||||
Err(error) => {
|
||||
trace!(
|
||||
server_name = %server_name,
|
||||
has_cached_tools,
|
||||
startup_complete,
|
||||
"MCP server tools unavailable while building tool list"
|
||||
);
|
||||
Err(error)
|
||||
}
|
||||
};
|
||||
(server_name, result)
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
for (server_name, server_tools) in server_results {
|
||||
match server_tools {
|
||||
|
||||
@@ -2718,7 +2718,7 @@ async fn list_available_server_infos_uses_cache_while_client_is_pending() {
|
||||
|
||||
let timeout_result = tokio::time::timeout(
|
||||
Duration::from_millis(10),
|
||||
manager.list_available_server_infos(),
|
||||
manager.list_available_server_infos(|_| true),
|
||||
)
|
||||
.await;
|
||||
let server_infos = timeout_result.expect("server info lookup should not block on startup");
|
||||
@@ -4024,7 +4024,7 @@ async fn list_all_tools_uses_shared_codex_apps_cache_when_client_startup_fails()
|
||||
assert_eq!(tool.callable_name, "calendar_create_event");
|
||||
assert_eq!(
|
||||
manager
|
||||
.list_available_server_infos()
|
||||
.list_available_server_infos(|_| true)
|
||||
.await
|
||||
.get(CODEX_APPS_MCP_SERVER_NAME),
|
||||
Some(&server_info)
|
||||
|
||||
@@ -497,13 +497,14 @@ pub struct McpServerStatusSnapshot {
|
||||
pub async fn collect_mcp_server_status_snapshot_with_detail(
|
||||
config: &McpConfig,
|
||||
auth: Option<&CodexAuth>,
|
||||
submit_id: String,
|
||||
runtime_context: McpRuntimeContext,
|
||||
codex_apps_tools_cache: ConnectorRuntimeManager<ToolInfo>,
|
||||
tool_catalog_cache: crate::McpToolCatalogCache,
|
||||
detail: McpSnapshotDetail,
|
||||
server_name: Option<&str>,
|
||||
) -> McpServerStatusSnapshot {
|
||||
let mcp_servers = effective_mcp_servers(config, auth);
|
||||
let mut mcp_servers = effective_mcp_servers(config, auth);
|
||||
mcp_servers.retain(|name, _| server_name.is_none_or(|selected| name == selected));
|
||||
if mcp_servers.is_empty() {
|
||||
return McpServerStatusSnapshot {
|
||||
server_infos: HashMap::new(),
|
||||
@@ -539,7 +540,7 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
|
||||
plugins_available: false,
|
||||
ready_selected_capability_roots: Vec::new(),
|
||||
mcp_servers,
|
||||
submit_id,
|
||||
submit_id: String::new(),
|
||||
tx_event: None,
|
||||
startup_cancellation_token: cancel_token.clone(),
|
||||
runtime_context,
|
||||
@@ -788,21 +789,29 @@ fn convert_mcp_resource_templates(
|
||||
.collect::<HashMap<_, _>>()
|
||||
}
|
||||
|
||||
async fn collect_mcp_server_status_snapshot_from_manager(
|
||||
pub(crate) async fn collect_mcp_server_status_snapshot_from_manager(
|
||||
mcp_connection_manager: &McpConnectionSet,
|
||||
auth_status_entries: HashMap<String, crate::mcp::auth::McpAuthStatusEntry>,
|
||||
server_names: Vec<String>,
|
||||
detail: McpSnapshotDetail,
|
||||
) -> McpServerStatusSnapshot {
|
||||
let selected_servers: HashSet<&str> = server_names.iter().map(String::as_str).collect();
|
||||
let include_server = |name: &str| selected_servers.contains(name);
|
||||
let ((server_infos, (tools, tools_errors)), resources, resource_templates) = tokio::join!(
|
||||
async {
|
||||
let server_infos = mcp_connection_manager.list_available_server_infos().await;
|
||||
let tools = mcp_connection_manager.list_tools_with_errors().await;
|
||||
let server_infos = mcp_connection_manager
|
||||
.list_available_server_infos(include_server)
|
||||
.await;
|
||||
let tools = mcp_connection_manager
|
||||
.list_tools_with_errors(include_server)
|
||||
.await;
|
||||
(server_infos, tools)
|
||||
},
|
||||
async {
|
||||
if detail.include_resources() {
|
||||
mcp_connection_manager.list_all_resources(|_| true).await
|
||||
mcp_connection_manager
|
||||
.list_all_resources(include_server)
|
||||
.await
|
||||
} else {
|
||||
HashMap::new()
|
||||
}
|
||||
@@ -810,7 +819,7 @@ async fn collect_mcp_server_status_snapshot_from_manager(
|
||||
async {
|
||||
if detail.include_resources() {
|
||||
mcp_connection_manager
|
||||
.list_all_resource_templates(|_| true)
|
||||
.list_all_resource_templates(include_server)
|
||||
.await
|
||||
} else {
|
||||
HashMap::new()
|
||||
|
||||
@@ -5,6 +5,8 @@
|
||||
//! [`crate::rmcp_client`] and connection-set behavior lives in
|
||||
//! [`crate::connection_manager`].
|
||||
|
||||
mod status;
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use std::path::PathBuf;
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
//! Read selected-server metadata from one published thread runtime.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::McpRuntime;
|
||||
use super::McpRuntimeContext;
|
||||
use crate::McpConfig;
|
||||
use crate::McpServerStatusSnapshot;
|
||||
use crate::McpSnapshotDetail;
|
||||
use crate::mcp::collect_mcp_server_status_snapshot_from_manager;
|
||||
use crate::mcp::compute_auth_statuses;
|
||||
use crate::mcp::effective_mcp_servers;
|
||||
|
||||
impl McpRuntime {
|
||||
/// Reuses the selected connection and catalog without starting unrelated servers.
|
||||
/// Configuration and metadata belong to the same captured runtime publication.
|
||||
pub async fn server_status_snapshot(
|
||||
&self,
|
||||
server: &str,
|
||||
detail: McpSnapshotDetail,
|
||||
runtime_context: &McpRuntimeContext,
|
||||
) -> anyhow::Result<(Arc<McpConfig>, McpServerStatusSnapshot)> {
|
||||
let current = self.current.load_full();
|
||||
let runtime_context = runtime_context.clone().with_selected_environments(
|
||||
Arc::clone(¤t.environment_selections),
|
||||
current.ready_environments.clone(),
|
||||
);
|
||||
let config = Arc::clone(
|
||||
current
|
||||
.config
|
||||
.as_ref()
|
||||
.ok_or_else(|| anyhow::anyhow!("MCP runtime is not configured"))?,
|
||||
);
|
||||
let mut servers = effective_mcp_servers(&config, current.auth.as_ref());
|
||||
servers.retain(|name, _| name == server);
|
||||
let auth_statuses = compute_auth_statuses(
|
||||
servers.iter(),
|
||||
config.mcp_oauth_credentials_store_mode,
|
||||
config.auth_keyring_backend_kind,
|
||||
current.auth.as_ref(),
|
||||
&runtime_context,
|
||||
)
|
||||
.await;
|
||||
let snapshot = collect_mcp_server_status_snapshot_from_manager(
|
||||
¤t.connections,
|
||||
auth_statuses,
|
||||
servers.into_keys().collect(),
|
||||
detail,
|
||||
)
|
||||
.await;
|
||||
Ok((config, snapshot))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user