openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-13 21:30:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Limit MCP status discovery to the requested page (#61365)

Browse Source 
## Why

Paginated MCP server status requests previously discovered all configured servers before selecting the response page, causing unnecessary connections to servers outside that page.

## What changed

Select and deduplicate server identities before discovery, then collect status only for servers on the requested page. Include thread runtime statuses when selecting pages and preserve separate entries for same-named servers in different environments. Empty pages skip discovery.

## Testing

Add coverage that counts MCP initialization requests across successive pages and an empty final page. Extend environment inventory coverage to verify that pagination preserves each registration's identity, authentication status, and tools.

GitOrigin-RevId: 0c9c9ca6c8016147eb9a12b53ffe533ac3334444
This commit is contained in:
Charlie Marsh
2026-11-13 19:50:06 +00:00
committed by copyberry
parent 4413deac57
commit 6d99161bce
4 changed files with 180 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/app-server/src/request_processors.rs
-1
View File
@@ -421,7 +421,6 @@ use codex_login::oauth_client_id;
use codex_login::request_device_code;
use codex_login::run_login_server;
use codex_mcp::McpRuntimeContext;
use codex_mcp::McpServerStatusSnapshot;
use codex_mcp::McpSnapshotDetail;
use codex_mcp::collect_mcp_server_status_snapshot_with_detail;
use codex_mcp::discover_supported_scopes;
codex-rs/app-server/src/request_processors/mcp_processor.rs
+85 -74
View File
@@ -5,6 +5,7 @@ use super::*;
use codex_config::types::McpServerAuth;
use codex_core::McpRefreshReason;
use codex_extension_api::McpServerContributionContext;
use codex_mcp::McpServerKey;
use codex_mcp::McpServerSource;
use codex_mcp::ReadResourceRequestParams;
use codex_mcp::ema_auth_scope;
@@ -414,10 +415,12 @@ impl McpRequestProcessor {
McpServerStatusDetail::ToolsAndAuthOnly => McpSnapshotDetail::ToolsAndAuthOnly,
};
let (mcp_config, snapshot) = if let (Some(thread), Some(server)) =
let runtime_statuses;
let (keys, next_cursor);
let (mcp_config, mut snapshot) = if let (Some(thread), Some(server)) =
(thread.as_ref(), params.server_name.as_deref())
{
thread
let (config, snapshot) = thread
.mcp_server_status_snapshot(
server,
params.environment_id.as_deref(),
@@ -427,7 +430,17 @@ impl McpRequestProcessor {
.await
.map_err(|error| {
internal_error(format!("failed to read MCP server status: {error:#}"))
})?
})?;
runtime_statuses = thread.mcp_connection_statuses(&config).await;
let selected = config
.mcp_server_catalog
.resolve_server_key(server, params.environment_id.as_deref());
(keys, next_cursor) = status_page(
snapshot.server_keys.iter().chain(runtime_statuses.keys()),
Some(&selected),
&params,
)?;
(config, snapshot)
} else {
if params.environment_id.is_some() {
return Err(invalid_request(
@@ -468,11 +481,23 @@ impl McpRequestProcessor {
),
),
};
let selected_server = params.server_name.as_deref().map(|name| {
let page_statuses = match thread.as_ref() {
Some(thread) => thread.mcp_connection_statuses(&mcp_config).await,
None => HashMap::new(),
};
let mut servers = codex_mcp::effective_mcp_servers(&mcp_config, auth.as_ref());
let selected = params.server_name.as_deref().map(|name| {
mcp_config
.mcp_server_catalog
.resolve_server_key(name, /*environment_id*/ None)
.resolve_server_key(name, params.environment_id.as_deref())
});
(keys, next_cursor) = status_page(
servers.keys().chain(page_statuses.keys()),
selected.as_ref(),
&params,
)?;
let page: HashSet<_> = keys.iter().collect();
servers.retain(|key, _| page.contains(key));
let mut snapshot = collect_mcp_server_status_snapshot_with_detail(
&mcp_config,
auth.as_ref(),
@@ -482,7 +507,7 @@ impl McpRequestProcessor {
mcp_manager.codex_apps_tools_cache(),
mcp_manager.tool_catalog_cache(),
detail,
selected_server.as_ref(),
servers,
)
.await;
if mcp_config.xaa_enabled
@@ -491,7 +516,9 @@ impl McpRequestProcessor {
// EMA can only use its admitted thread connection. Keep ordinary
// status discovery unchanged and replace the threadless EMA results.
for (key, server) in mcp_config.mcp_server_catalog.servers() {
if !matches!(server.config().auth, McpServerAuth::EmaAuth) {
if !page.contains(key)
|| !matches!(server.config().auth, McpServerAuth::EmaAuth)
{
continue;
}
let registration = Some(server);
@@ -531,62 +558,17 @@ impl McpRequestProcessor {
snapshot.auth_statuses.extend(selected.auth_statuses);
}
}
// Discovery can finish a thread connection's startup after we select the page.
runtime_statuses = match thread.as_ref() {
Some(thread) if !keys.is_empty() => {
thread.mcp_connection_statuses(&mcp_config).await
}
_ => page_statuses,
};
(Arc::new(mcp_config), snapshot)
};
let runtime_statuses = match thread {
Some(thread) => thread.mcp_connection_statuses(&mcp_config).await,
None => HashMap::new(),
};
let McpServerStatusSnapshot {
mut server_infos,
mut server_capabilities,
mut tool_list_meta,
mut tools_by_server,
mut tools_errors,
mut resources,
mut resource_templates,
mut auth_statuses,
mut server_keys,
} = snapshot;
server_keys.extend(runtime_statuses.keys().cloned());
server_keys.extend(
auth_statuses
.keys()
.cloned()
.chain(resources.keys().cloned())
.chain(resource_templates.keys().cloned()),
);
server_keys.sort();
server_keys.dedup();
server_keys.retain(|key| {
params.server_name.as_deref().is_none_or(|server| {
*key == mcp_config
.mcp_server_catalog
.resolve_server_key(server, params.environment_id.as_deref())
})
});
let total = server_keys.len();
let limit = params.limit.unwrap_or(total as u32).max(1) as usize;
let effective_limit = limit.min(total);
let start = match params.cursor {
Some(cursor) => match cursor.parse::<usize>() {
Ok(idx) => idx,
Err(_) => return Err(invalid_request(format!("invalid cursor: {cursor}"))),
},
None => 0,
};
if start > total {
return Err(invalid_request(format!(
"cursor {start} exceeds total MCP servers {total}"
)));
}
let end = start.saturating_add(effective_limit).min(total);
let data: Vec<McpServerStatus> = server_keys[start..end]
let data: Vec<McpServerStatus> = keys
.iter()
.map(|key| McpServerStatus {
name: key.name().to_owned(),
@@ -615,26 +597,24 @@ impl McpRequestProcessor {
}
McpServerTransportConfig::Stdio { .. } => None,
}),
server_info: server_infos.remove(key),
server_capabilities: server_capabilities.remove(key),
tool_list_meta: tool_list_meta.remove(key).map(Arc::unwrap_or_clone),
tools: tools_by_server.remove(key).unwrap_or_default(),
tools_error: tools_errors.remove(key),
resources: resources.remove(key).unwrap_or_default(),
resource_templates: resource_templates.remove(key).unwrap_or_default(),
auth_status: auth_statuses
server_info: snapshot.server_infos.remove(key),
server_capabilities: snapshot.server_capabilities.remove(key),
tool_list_meta: snapshot
.tool_list_meta
.remove(key)
.map(Arc::unwrap_or_clone),
tools: snapshot.tools_by_server.remove(key).unwrap_or_default(),
tools_error: snapshot.tools_errors.remove(key),
resources: snapshot.resources.remove(key).unwrap_or_default(),
resource_templates: snapshot.resource_templates.remove(key).unwrap_or_default(),
auth_status: snapshot
.auth_statuses
.remove(key)
.unwrap_or(CoreMcpAuthStatus::Unsupported)
.into(),
})
.collect();
let next_cursor = if end < total {
Some(end.to_string())
} else {
None
};
Ok(ListMcpServerStatusResponse { data, next_cursor })
}
@@ -892,3 +872,34 @@ fn with_mcp_tool_call_thread_id_meta(
other => other,
}
}
/// Selects complete server identities before starting discovery for a page.
fn status_page<'a>(
keys: impl Iterator<Item = &'a McpServerKey>,
selected: Option<&McpServerKey>,
params: &ListMcpServerStatusParams,
) -> Result<(Vec<McpServerKey>, Option<String>), JSONRPCErrorError> {
let mut keys: Vec<_> = keys
.filter(|key| selected.is_none_or(|selected| *key == selected))
.collect();
keys.sort_unstable();
keys.dedup();
let total = keys.len();
let start = match &params.cursor {
Some(cursor) => cursor
.parse::<usize>()
.map_err(|_| invalid_request(format!("invalid cursor: {cursor}")))?,
None => 0,
};
if start > total {
return Err(invalid_request(format!(
"cursor {start} exceeds total MCP servers {total}"
)));
}
let limit = params.limit.map_or(total, |limit| limit.max(1) as usize);
let end = start.saturating_add(limit).min(total);
Ok((
keys[start..end].iter().copied().cloned().collect(),
(end < total).then(|| end.to_string()),
))
}
codex-rs/app-server/tests/suite/v2/mcp_server_status.rs
+92 -4
View File
@@ -86,9 +86,7 @@ use wiremock::matchers::path;
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
/// Selected status reads return the thread's tools without repeating MCP discovery.
#[tokio::test]
async fn mcp_server_status_list_reuses_thread_connection() -> Result<()> {
async fn recording_widget_server() -> (MockServer, serde_json::Value) {
let server = MockServer::start().await;
let tool = json!({
"name": "open_widget",
@@ -100,7 +98,8 @@ async fn mcp_server_status_list_reuses_thread_connection() -> Result<()> {
Mock::given(method("POST"))
.and(path("/mcp"))
.respond_with(move |request: &wiremock::Request| {
let request: serde_json::Value = serde_json::from_slice(&request.body).unwrap();
let request: serde_json::Value =
serde_json::from_slice(&request.body).expect("valid MCP request");
let result = match request["method"].as_str() {
Some("initialize") => json!({
"protocolVersion": "2025-11-25",
@@ -118,6 +117,13 @@ async fn mcp_server_status_list_reuses_thread_connection() -> Result<()> {
})
.mount(&server)
.await;
(server, tool)
}
/// Selected status reads return the thread's tools without repeating MCP discovery.
#[tokio::test]
async fn mcp_server_status_list_reuses_thread_connection() -> Result<()> {
let (server, tool) = recording_widget_server().await;
let codex_home = TempDir::new()?;
mock_responses_config(&server.uri())
.with_extra_config(&format!(
@@ -184,6 +190,63 @@ async fn mcp_server_status_list_reuses_thread_connection() -> Result<()> {
Ok(())
}
#[tokio::test]
async fn mcp_server_status_list_discovers_only_the_requested_page() -> Result<()> {
let (server, _) = recording_widget_server().await;
let codex_home = TempDir::new()?;
mock_responses_config(&server.uri())
.with_extra_config(&format!(
"[mcp_servers.first]\nurl = \"{0}/mcp\"\n[mcp_servers.last]\nurl = \"{0}/mcp\"",
server.uri()
))
.write(codex_home.path())?;
let mut app = TestAppServer::builder()
.with_codex_home(codex_home.path())
.build_initialized()
.await?;
for (cursor, expected_name, next, initialize_count) in [
(None, Some("first"), Some("1"), 1),
(Some("1"), Some("last"), None, 2),
(Some("2"), None, None, 2),
] {
let response: ListMcpServerStatusResponse = app
.request(|request_id| ClientRequest::McpServerStatusList {
request_id,
params: ListMcpServerStatusParams {
environment_id: None,
thread_id: None,
server_name: None,
cursor: cursor.map(str::to_owned),
limit: Some(1),
detail: Some(McpServerStatusDetail::Full),
},
})
.await?;
assert_eq!(
response
.data
.iter()
.map(|row| row.name.as_str())
.collect::<Vec<_>>(),
expected_name.into_iter().collect::<Vec<_>>()
);
assert_eq!(response.next_cursor.as_deref(), next);
let initialize = body_partial_json(json!({"method": "initialize"}));
assert_eq!(
server
.received_requests()
.await
.unwrap()
.iter()
.filter(|request| initialize.matches(request))
.count(),
initialize_count
);
}
Ok(())
}
#[test_case(false; "standalone")]
#[test_case(true; "thread runtime")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
@@ -1677,6 +1740,31 @@ async fn mcp_server_status_keeps_same_named_environment_inventory_separate() ->
),
],
);
for (index, next_cursor) in [(0, Some("1")), (1, None)] {
let page: ListMcpServerStatusResponse = app
.request(|request_id| ClientRequest::McpServerStatusList {
request_id,
params: ListMcpServerStatusParams {
environment_id: None,
thread_id: Some(thread.id.clone()),
server_name: None,
detail: Some(McpServerStatusDetail::Full),
cursor: (index != 0).then(|| index.to_string()),
limit: Some(1),
},
})
.await?;
let [row] = page.data.as_slice() else {
anyhow::bail!("expected one MCP registration on page {index}");
};
let expected = &response.data[index];
assert_eq!(row.name, expected.name);
assert_eq!(row.environment_id, expected.environment_id);
assert_eq!(row.http_origin, expected.http_origin);
assert_eq!(row.auth_status, expected.auth_status);
assert_eq!(row.tools, expected.tools);
assert_eq!(page.next_cursor.as_deref(), next_cursor);
}
for (environment_id, expected) in [
(Some("env-a"), &response.data[0]),
(Some("env-b"), &response.data[1]),
codex-rs/codex-mcp/src/mcp/mod.rs
+3 -15
View File
@@ -527,7 +527,7 @@ pub async fn read_mcp_resource(
result
}
#[derive(Debug, Clone)]
#[derive(Debug, Clone, Default)]
pub struct McpServerStatusSnapshot {
pub server_infos: HashMap<McpServerKey, McpServerInfo>,
pub server_capabilities: HashMap<McpServerKey, serde_json::Value>,
@@ -553,22 +553,10 @@ pub async fn collect_mcp_server_status_snapshot_with_detail(
codex_apps_tools_cache: ConnectorRuntimeManager<ToolInfo>,
tool_catalog_cache: crate::McpToolCatalogCache,
detail: McpSnapshotDetail,
server: Option<&McpServerKey>,
mcp_servers: HashMap<McpServerKey, EffectiveMcpServer>,
) -> McpServerStatusSnapshot {
let mut mcp_servers = effective_mcp_servers(config, auth);
mcp_servers.retain(|key, _| server.is_none_or(|selected| key == selected));
if mcp_servers.is_empty() {
return McpServerStatusSnapshot {
server_infos: HashMap::new(),
server_capabilities: HashMap::new(),
tool_list_meta: HashMap::new(),
tools_by_server: HashMap::new(),
tools_errors: HashMap::new(),
resources: HashMap::new(),
resource_templates: HashMap::new(),
auth_statuses: HashMap::new(),
server_keys: Vec::new(),
};
return McpServerStatusSnapshot::default();
}
let auth_status_entries = compute_auth_statuses(