From 3e3ae088398b4a4ad0833c84f5b1dcaaf5cdabe3 Mon Sep 17 00:00:00 2001 From: jif Date: Wed, 29 Jul 2026 22:24:24 +0000 Subject: [PATCH] Limit MCP catalog pagination (#36039) ## Why MCP catalog discovery must not allow a server to keep pagination running or growing without bounds. ## What changed - Apply the shared pagination collector to tool, resource, and resource-template discovery. - Limit each catalog to 100 pages and 1,024 items, and reject pagination cursors larger than 64 KiB or any repeated cursor. - Bound the entire pagination operation by the configured tool timeout, falling back to 30 seconds when no timeout is configured. ## Testing Add unit coverage for every limit and an end-to-end MCP test that preserves valid multi-page tools while excluding a server that returns an oversized cursor. GitOrigin-RevId: 961a1bfa5fe8a265daf56e342a919c980cd720fa --- codex-rs/codex-mcp/src/binding_clients.rs | 35 +++--- .../src/connection_manager/resources.rs | 62 +++------ codex-rs/codex-mcp/src/pagination.rs | 32 ++++- codex-rs/codex-mcp/src/pagination_tests.rs | 119 ++++++++++++++++++ codex-rs/codex-mcp/src/rmcp_client.rs | 33 +++-- codex-rs/core/tests/suite/rmcp_client.rs | 103 +++++++++++++++ .../rmcp-client/src/bin/test_stdio_server.rs | 22 +++- 7 files changed, 315 insertions(+), 91 deletions(-) diff --git a/codex-rs/codex-mcp/src/binding_clients.rs b/codex-rs/codex-mcp/src/binding_clients.rs index fe0c92e4ff..b2b1c174b8 100644 --- a/codex-rs/codex-mcp/src/binding_clients.rs +++ b/codex-rs/codex-mcp/src/binding_clients.rs @@ -91,15 +91,14 @@ impl McpBindingClients { let client = Arc::clone(&managed.client); let timeout = managed.tool_timeout; join_set.spawn(async move { - let resources = - collect_paginated("resources/list", /*overall_timeout*/ None, |params| { - let client = Arc::clone(&client); - async move { - let response = client.list_resources(params, timeout).await?; - Ok((response.resources, response.next_cursor)) - } - }) - .await; + let resources = collect_paginated("resources/list", timeout, |params| { + let client = Arc::clone(&client); + async move { + let response = client.list_resources(params, timeout).await?; + Ok((response.resources, response.next_cursor)) + } + }) + .await; (server_name, resources) }); } @@ -120,17 +119,13 @@ impl McpBindingClients { let client = Arc::clone(&managed.client); let timeout = managed.tool_timeout; join_set.spawn(async move { - let templates = collect_paginated( - "resources/templates/list", - /*overall_timeout*/ None, - |params| { - let client = Arc::clone(&client); - async move { - let response = client.list_resource_templates(params, timeout).await?; - Ok((response.resource_templates, response.next_cursor)) - } - }, - ) + let templates = collect_paginated("resources/templates/list", timeout, |params| { + let client = Arc::clone(&client); + async move { + let response = client.list_resource_templates(params, timeout).await?; + Ok((response.resource_templates, response.next_cursor)) + } + }) .await; (server_name, templates) }); diff --git a/codex-rs/codex-mcp/src/connection_manager/resources.rs b/codex-rs/codex-mcp/src/connection_manager/resources.rs index 06edde90b4..6762375147 100644 --- a/codex-rs/codex-mcp/src/connection_manager/resources.rs +++ b/codex-rs/codex-mcp/src/connection_manager/resources.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::sync::Arc; use std::time::Duration; use anyhow::Context; @@ -15,6 +16,7 @@ use tokio::task::JoinSet; use tracing::warn; use super::McpConnectionSet; +use crate::pagination::collect_paginated; use crate::rmcp_client::ManagedClient; impl McpConnectionSet { @@ -36,28 +38,15 @@ impl McpConnectionSet { let timeout = view.tool_timeout; let client = managed_client.client; join_set.spawn(async move { - let mut resources = Vec::new(); - let mut cursor: Option = None; - loop { - let params = cursor.as_ref().map(|next| { - PaginatedRequestParams::default().with_cursor(Some(next.clone())) - }); - let response = match client.list_resources(params, timeout).await { - Ok(result) => result, - Err(error) => return (server_name, Err(error)), - }; - resources.extend(response.resources); - match response.next_cursor { - Some(next) if cursor.as_ref() == Some(&next) => { - return ( - server_name, - Err(anyhow!("resources/list returned duplicate cursor")), - ); - } - Some(next) => cursor = Some(next), - None => return (server_name, Ok(resources)), + let resources = collect_paginated("resources/list", timeout, |params| { + let client = Arc::clone(&client); + async move { + let response = client.list_resources(params, timeout).await?; + Ok((response.resources, response.next_cursor)) } - } + }) + .await; + (server_name, resources) }); } @@ -96,30 +85,15 @@ impl McpConnectionSet { let timeout = view.tool_timeout; let client = managed_client.client; join_set.spawn(async move { - let mut templates = Vec::new(); - let mut cursor: Option = None; - loop { - let params = cursor.as_ref().map(|next| { - PaginatedRequestParams::default().with_cursor(Some(next.clone())) - }); - let response = match client.list_resource_templates(params, timeout).await { - Ok(result) => result, - Err(error) => return (server_name, Err(error)), - }; - templates.extend(response.resource_templates); - match response.next_cursor { - Some(next) if cursor.as_ref() == Some(&next) => { - return ( - server_name, - Err(anyhow!( - "resources/templates/list returned duplicate cursor" - )), - ); - } - Some(next) => cursor = Some(next), - None => return (server_name, Ok(templates)), + let templates = collect_paginated("resources/templates/list", timeout, |params| { + let client = Arc::clone(&client); + async move { + let response = client.list_resource_templates(params, timeout).await?; + Ok((response.resource_templates, response.next_cursor)) } - } + }) + .await; + (server_name, templates) }); } diff --git a/codex-rs/codex-mcp/src/pagination.rs b/codex-rs/codex-mcp/src/pagination.rs index 3b699471be..61e3b0658a 100644 --- a/codex-rs/codex-mcp/src/pagination.rs +++ b/codex-rs/codex-mcp/src/pagination.rs @@ -6,6 +6,11 @@ use anyhow::Result; use anyhow::anyhow; use rmcp::model::PaginatedRequestParams; +const MAX_MCP_CATALOG_PAGES: usize = 100; +const MAX_MCP_CATALOG_ITEMS: usize = 1_024; +const MAX_MCP_PAGINATION_CURSOR_BYTES: usize = 64 * 1024; +const DEFAULT_MCP_PAGINATION_TIMEOUT: Duration = Duration::from_secs(30); + pub(crate) async fn collect_paginated( method: &str, overall_timeout: Option, @@ -19,17 +24,34 @@ where let mut collected = Vec::new(); let mut cursor = None; let mut seen_cursors = HashSet::new(); + let mut page_count = 0; loop { + if page_count == MAX_MCP_CATALOG_PAGES { + return Err(anyhow!( + "{method} exceeded the pagination limit of {MAX_MCP_CATALOG_PAGES} pages" + )); + } + page_count += 1; let params = cursor.as_ref().map(|next: &String| { PaginatedRequestParams::default().with_cursor(Some(next.clone())) }); let (items, next_cursor) = fetch(params).await?; + if items.len() > MAX_MCP_CATALOG_ITEMS.saturating_sub(collected.len()) { + return Err(anyhow!( + "{method} exceeded the catalog limit of {MAX_MCP_CATALOG_ITEMS} items" + )); + } collected.extend(items); let Some(next_cursor) = next_cursor else { return Ok(collected); }; + if next_cursor.len() > MAX_MCP_PAGINATION_CURSOR_BYTES { + return Err(anyhow!( + "{method} returned a pagination cursor exceeding {MAX_MCP_PAGINATION_CURSOR_BYTES} bytes" + )); + } if !seen_cursors.insert(next_cursor.clone()) { return Err(anyhow!("{method} returned a repeated pagination cursor")); } @@ -37,12 +59,10 @@ where } }; - match overall_timeout { - Some(timeout) => tokio::time::timeout(timeout, collect) - .await - .map_err(|_| anyhow!("{method} pagination timed out after {timeout:?}"))?, - None => collect.await, - } + let timeout = overall_timeout.unwrap_or(DEFAULT_MCP_PAGINATION_TIMEOUT); + tokio::time::timeout(timeout, collect) + .await + .map_err(|_| anyhow!("{method} pagination timed out after {timeout:?}"))? } #[cfg(test)] diff --git a/codex-rs/codex-mcp/src/pagination_tests.rs b/codex-rs/codex-mcp/src/pagination_tests.rs index 8e58712d9e..709acd299a 100644 --- a/codex-rs/codex-mcp/src/pagination_tests.rs +++ b/codex-rs/codex-mcp/src/pagination_tests.rs @@ -1,10 +1,15 @@ use std::sync::Arc; use std::sync::Mutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::anyhow; use pretty_assertions::assert_eq; +use super::MAX_MCP_CATALOG_ITEMS; +use super::MAX_MCP_CATALOG_PAGES; +use super::MAX_MCP_PAGINATION_CURSOR_BYTES; use super::collect_paginated; #[tokio::test] @@ -60,6 +65,101 @@ async fn rejects_nonconsecutive_repeated_cursors() { ); } +#[tokio::test] +async fn rejects_excessive_pagination_before_fetching_another_page() { + let requests = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&requests); + + let error = collect_paginated( + "tools/list", + /*overall_timeout*/ None, + move |_params| { + let observed = Arc::clone(&observed); + async move { + let page = observed.fetch_add(1, Ordering::Relaxed); + Ok((Vec::<()>::new(), Some(page.to_string()))) + } + }, + ) + .await + .expect_err("unbounded pagination must fail"); + + assert_eq!( + error.to_string(), + format!("tools/list exceeded the pagination limit of {MAX_MCP_CATALOG_PAGES} pages") + ); + assert_eq!(requests.load(Ordering::Relaxed), MAX_MCP_CATALOG_PAGES); +} + +#[tokio::test] +async fn rejects_a_page_exceeding_the_catalog_item_limit() { + let error = collect_paginated( + "tools/list", + /*overall_timeout*/ None, + |_params| async { Ok((vec![(); MAX_MCP_CATALOG_ITEMS + 1], None)) }, + ) + .await + .expect_err("an oversized catalog page must fail"); + + assert_eq!( + error.to_string(), + format!("tools/list exceeded the catalog limit of {MAX_MCP_CATALOG_ITEMS} items") + ); +} + +#[tokio::test] +async fn rejects_a_catalog_exceeding_the_item_limit_across_pages() { + let error = collect_paginated( + "tools/list", + /*overall_timeout*/ None, + |params| async move { + match params.and_then(|params| params.cursor) { + None => Ok((vec![(); MAX_MCP_CATALOG_ITEMS], Some("last".to_string()))), + Some(cursor) if cursor == "last" => Ok((vec![()], None)), + Some(cursor) => Err(anyhow!("unexpected cursor: {cursor}")), + } + }, + ) + .await + .expect_err("catalog items must be bounded across pages"); + + assert_eq!( + error.to_string(), + format!("tools/list exceeded the catalog limit of {MAX_MCP_CATALOG_ITEMS} items") + ); +} + +#[tokio::test] +async fn rejects_oversized_pagination_cursors_before_following_them() { + let requests = Arc::new(AtomicUsize::new(0)); + let observed = Arc::clone(&requests); + + let error = collect_paginated( + "resources/list", + /*overall_timeout*/ None, + move |_params| { + let observed = Arc::clone(&observed); + async move { + observed.fetch_add(1, Ordering::Relaxed); + Ok(( + Vec::<()>::new(), + Some("x".repeat(MAX_MCP_PAGINATION_CURSOR_BYTES + 1)), + )) + } + }, + ) + .await + .expect_err("an oversized pagination cursor must fail"); + + assert_eq!( + error.to_string(), + format!( + "resources/list returned a pagination cursor exceeding {MAX_MCP_PAGINATION_CURSOR_BYTES} bytes" + ) + ); + assert_eq!(requests.load(Ordering::Relaxed), 1); +} + #[tokio::test] async fn forwards_page_failures() { let error = collect_paginated( @@ -73,6 +173,25 @@ async fn forwards_page_failures() { assert_eq!(error.to_string(), "page failed"); } +#[tokio::test(start_paused = true)] +async fn applies_a_default_timeout_when_no_timeout_is_configured() { + let error = collect_paginated( + "resources/list", + /*overall_timeout*/ None, + |_params| async { + tokio::time::sleep(Duration::from_secs(31)).await; + Ok((Vec::<()>::new(), None)) + }, + ) + .await + .expect_err("pagination without a configured timeout must still be bounded"); + + assert_eq!( + error.to_string(), + "resources/list pagination timed out after 30s" + ); +} + #[tokio::test(start_paused = true)] async fn applies_one_timeout_across_individually_timely_pages() { let requests = Arc::new(Mutex::new(Vec::new())); diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index 0e99ce27a1..5a2a8b134e 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -591,26 +591,21 @@ pub(crate) async fn list_tools_for_client_uncached( server_instructions: Option<&str>, ) -> Result> { let fetch_start = Instant::now(); - let tools = match client.protocol_mode() { - McpProtocolMode::Legacy => { - client - .list_tools_with_connector_ids(/*params*/ None, timeout) - .await? - .tools + let protocol_mode = client.protocol_mode(); + let tools = collect_paginated("tools/list", timeout, |params| { + let client = Arc::clone(client); + async move { + let response = client + .list_tools_with_connector_ids(params, timeout) + .await?; + let next_cursor = match protocol_mode { + McpProtocolMode::Legacy => None, + McpProtocolMode::V20260728 => response.next_cursor, + }; + Ok((response.tools, next_cursor)) } - McpProtocolMode::V20260728 => { - collect_paginated("tools/list", timeout, |params| { - let client = Arc::clone(client); - async move { - let response = client - .list_tools_with_connector_ids(params, timeout) - .await?; - Ok((response.tools, response.next_cursor)) - } - }) - .await? - } - } + }) + .await? .into_iter() .map(|tool| { tool_info_from_listed_tool( diff --git a/codex-rs/core/tests/suite/rmcp_client.rs b/codex-rs/core/tests/suite/rmcp_client.rs index 2220730c11..1eb63cf408 100644 --- a/codex-rs/core/tests/suite/rmcp_client.rs +++ b/codex-rs/core/tests/suite/rmcp_client.rs @@ -921,6 +921,109 @@ server_names = ["history", "notes"] Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn modern_mcp_pagination_preserves_valid_tools_and_rejects_oversized_cursors() +-> anyhow::Result<()> { + skip_if_wine_exec!( + Ok(()), + "requires a Windows test_stdio_server in the Wine-exec environment" + ); + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let response = mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_assistant_message("msg-1", "done"), + responses::ev_completed("resp-1"), + ]), + ) + .await; + let command = remote_aware_stdio_server_bin()?; + let fixture = test_codex() + .with_model_info_override("gpt-5.4", |model| model.supports_search_tool = false) + .with_config(move |config| { + config + .features + .enable(Feature::Mcp20260728) + .expect("test config should allow modern MCP"); + for (server_name, pagination) in + [("paginated", "two-pages"), ("rejected", "oversized-cursor")] + { + insert_mcp_server( + config, + server_name, + stdio_transport( + command.clone(), + Some(HashMap::from([ + ( + "CODEX_MCP_PROTOCOL_VERSION".to_string(), + "2026-07-28".to_string(), + ), + ( + "MCP_TEST_TOOL_PAGINATION".to_string(), + pagination.to_string(), + ), + ])), + Vec::new(), + ), + TestMcpServerOptions { + environment_id: remote_aware_environment_id(), + ..Default::default() + }, + ); + } + }) + .build_with_auto_env(&server) + .await?; + + let startup = loop { + let event = fixture.codex.next_event().await?; + if let EventMsg::McpStartupComplete(startup) = event.msg { + break startup; + } + }; + assert!(startup.ready.iter().any(|name| name == "paginated")); + let failure = startup + .failed + .iter() + .find(|failure| failure.server == "rejected") + .expect("oversized cursor should reject only its MCP server"); + assert!( + failure + .error + .contains("tools/list returned a pagination cursor exceeding 65536 bytes"), + "unexpected MCP startup failure: {}", + failure.error + ); + + fixture + .codex + .submit(read_only_user_turn( + &fixture, + "show the paginated MCP tools", + )) + .await?; + wait_for_event(&fixture.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let body = response.single_request().body_json(); + for tool_name in ["echo", "sync"] { + assert!( + responses::namespace_child_tool(&body, "mcp__paginated", tool_name).is_some(), + "expected paginated MCP tool {tool_name} to reach the model" + ); + } + assert!( + responses::namespace_child_tool(&body, "mcp__rejected", "echo").is_none(), + "a rejected MCP catalog must not reach the model" + ); + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn apps_enabled_turn_skips_pending_optional_mcp_without_cached_tools() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs index d95a939dff..d19f48156e 100644 --- a/codex-rs/rmcp-client/src/bin/test_stdio_server.rs +++ b/codex-rs/rmcp-client/src/bin/test_stdio_server.rs @@ -513,7 +513,7 @@ impl ServerHandler for TestToolServer { fn list_tools( &self, - _request: Option, + request: Option, _context: rmcp::service::RequestContext, ) -> impl std::future::Future> + Send + '_ { let tools = self.tools.clone(); @@ -527,7 +527,25 @@ impl ServerHandler for TestToolServer { .get_or_insert_with(MetaObject::new) .insert("ui".to_string(), json!({ "visibility": ["app"] })); } - Ok(ListToolsResult::with_all_items(tools)) + let mut result = ListToolsResult::with_all_items(tools); + match ( + std::env::var("MCP_TEST_TOOL_PAGINATION").as_deref(), + request.and_then(|request| request.cursor).as_deref(), + ) { + (Ok("two-pages"), None) => { + result.tools.retain(|tool| tool.name == "echo"); + result.next_cursor = Some("second".to_string()); + } + (Ok("two-pages"), Some("second")) => { + result.tools.retain(|tool| tool.name == "sync"); + } + (Ok("oversized-cursor"), None) => { + result.tools.retain(|tool| tool.name == "echo"); + result.next_cursor = Some("x".repeat(65_537)); + } + _ => {} + } + Ok(result) } }