diff --git a/codex-rs/app-server/src/request_processors.rs b/codex-rs/app-server/src/request_processors.rs index 82dbae7584..1663d2e917 100644 --- a/codex-rs/app-server/src/request_processors.rs +++ b/codex-rs/app-server/src/request_processors.rs @@ -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; diff --git a/codex-rs/app-server/src/request_processors/mcp_processor.rs b/codex-rs/app-server/src/request_processors/mcp_processor.rs index 6b973bccb5..1d5efe9ebe 100644 --- a/codex-rs/app-server/src/request_processors/mcp_processor.rs +++ b/codex-rs/app-server/src/request_processors/mcp_processor.rs @@ -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), + ¶ms, + )?; + (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(), + ¶ms, + )?; + 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::() { - 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 = server_keys[start..end] + let data: Vec = 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, + selected: Option<&McpServerKey>, + params: &ListMcpServerStatusParams, +) -> Result<(Vec, Option), 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 ¶ms.cursor { + Some(cursor) => cursor + .parse::() + .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()), + )) +} diff --git a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs index 19ef0371a9..e904bde011 100644 --- a/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs +++ b/codex-rs/app-server/tests/suite/v2/mcp_server_status.rs @@ -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::>(), + expected_name.into_iter().collect::>() + ); + 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]), diff --git a/codex-rs/codex-mcp/src/mcp/mod.rs b/codex-rs/codex-mcp/src/mcp/mod.rs index 130b84dc8e..3b546cd7c3 100644 --- a/codex-rs/codex-mcp/src/mcp/mod.rs +++ b/codex-rs/codex-mcp/src/mcp/mod.rs @@ -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, pub server_capabilities: HashMap, @@ -553,22 +553,10 @@ pub async fn collect_mcp_server_status_snapshot_with_detail( codex_apps_tools_cache: ConnectorRuntimeManager, tool_catalog_cache: crate::McpToolCatalogCache, detail: McpSnapshotDetail, - server: Option<&McpServerKey>, + mcp_servers: HashMap, ) -> 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(