From e4e040881acab7e0059775ec58843c43ac6b882f Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 3 Aug 2026 16:57:35 +0000 Subject: [PATCH] Raise the host-owned Codex Apps catalog limit (#36772) ## Why Codex Apps tool catalogs can exceed the standard MCP limit of 2,048 items and still need to remain available through tool search. ## What changed - Allow up to 8,192 catalog items for host-owned `codex_apps` registrations. - Keep the 2,048-item limit for other MCP servers, including extensions that register the `codex_apps` name. - Restrict the shared Codex Apps tool cache to host-owned registrations. ## Testing - Cover large catalogs for built-in and hosted-plugin registrations while verifying that another extension retains the standard limit. - Verify that a same-named extension cannot read the host-owned Apps cache. GitOrigin-RevId: 2265693a460e9be63c4c849169cba7324def4f49 --- codex-rs/codex-mcp/src/connection_manager.rs | 31 ++++- .../src/connection_manager/tool_catalog.rs | 1 + .../codex-mcp/src/connection_manager_tests.rs | 80 +++++++++++++ codex-rs/codex-mcp/src/pagination.rs | 20 +++- codex-rs/codex-mcp/src/rmcp_client.rs | 13 +- .../core/tests/suite/mcp_tool_exposure.rs | 112 +++++++++++++++++- 6 files changed, 247 insertions(+), 10 deletions(-) diff --git a/codex-rs/codex-mcp/src/connection_manager.rs b/codex-rs/codex-mcp/src/connection_manager.rs index 90288db8ff..810ab7181f 100644 --- a/codex-rs/codex-mcp/src/connection_manager.rs +++ b/codex-rs/codex-mcp/src/connection_manager.rs @@ -27,10 +27,13 @@ use std::sync::OnceLock; use std::sync::atomic::Ordering; use std::time::Duration; +use crate::McpServerSource; use crate::elicitation::ElicitationRequestManager; use crate::elicitation::ElicitationRequestRouter; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; +use crate::pagination::MAX_CODEX_APPS_TOOL_CATALOG_ITEMS; +use crate::pagination::MAX_MCP_CATALOG_ITEMS; use crate::rmcp_client::AsyncManagedClient; use crate::rmcp_client::DEFAULT_TOOL_TIMEOUT; use crate::rmcp_client::ManagedClient; @@ -124,6 +127,7 @@ struct McpServerView { metadata: McpServerMetadata, tool_filter: ToolFilter, tool_timeout: Option, + catalog_item_limit: usize, } impl McpServerView { @@ -236,6 +240,21 @@ impl McpConnectionSet { .into_iter() .filter(|(_, server)| server.enabled()) { + let is_host_owned_codex_apps = server_name == CODEX_APPS_MCP_SERVER_NAME + && config.mcp_server_catalog.server(&server_name).is_some_and( + |server| match server.source() { + McpServerSource::Compatibility { .. } => true, + McpServerSource::Extension { id } => id == "hosted_plugin_runtime", + McpServerSource::Plugin(_) + | McpServerSource::SelectedPlugin(_) + | McpServerSource::Config => false, + }, + ); + let catalog_item_limit = if is_host_owned_codex_apps { + MAX_CODEX_APPS_TOOL_CATALOG_ITEMS + } else { + MAX_MCP_CATALOG_ITEMS + }; let metadata = McpServerMetadata::from(&server); let configured_config = server.config().clone(); let configured_tool_filter = ToolFilter::from_config(&configured_config); @@ -256,8 +275,8 @@ impl McpConnectionSet { } => bearer_token_env_var.is_some(), McpServerTransportConfig::Stdio { .. } => false, }; - let shares_codex_apps_tools_cache = - should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); + let shares_codex_apps_tools_cache = is_host_owned_codex_apps + && should_share_codex_apps_tools_cache(&server_name, uses_env_bearer_token); let codex_apps_tools_cache_context = shares_codex_apps_tools_cache.then(|| { codex_apps_tools_cache .context(codex_home.clone(), codex_apps_tools_cache_key.clone()) @@ -323,8 +342,9 @@ impl McpConnectionSet { .reusable_client(&connection_identity) .await .is_some_and(|client| { - expected_protocol_mode - .is_some_and(|expected| client.client.protocol_mode() == expected) + previous_view.catalog_item_limit == catalog_item_limit + && expected_protocol_mode + .is_some_and(|expected| client.client.protocol_mode() == expected) }) { servers.insert( @@ -334,6 +354,7 @@ impl McpConnectionSet { metadata, tool_filter: configured_tool_filter, tool_timeout: configured_tool_timeout, + catalog_item_limit, }, ); reused_ready.push(server_name); @@ -373,6 +394,7 @@ impl McpConnectionSet { client_elicitation_capability.clone(), supports_openai_form_elicitation, protocol_mode, + catalog_item_limit, ); servers.insert( server_name.clone(), @@ -384,6 +406,7 @@ impl McpConnectionSet { metadata, tool_filter: configured_tool_filter, tool_timeout: configured_tool_timeout, + catalog_item_limit, }, ); let tx_event = tx_event.clone(); diff --git a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs index 3e4f7721d6..b5c27c903e 100644 --- a/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs +++ b/codex-rs/codex-mcp/src/connection_manager/tool_catalog.rs @@ -347,6 +347,7 @@ impl McpConnectionSet { /*codex_apps_refresh_trigger*/ "explicit", &managed_client.client, view.tool_timeout, + view.catalog_item_limit, managed_client.server_instructions.as_deref(), ) .await diff --git a/codex-rs/codex-mcp/src/connection_manager_tests.rs b/codex-rs/codex-mcp/src/connection_manager_tests.rs index c7a4c1f602..461991c866 100644 --- a/codex-rs/codex-mcp/src/connection_manager_tests.rs +++ b/codex-rs/codex-mcp/src/connection_manager_tests.rs @@ -124,6 +124,7 @@ impl McpConnectionSet { tool_approval_modes: HashMap::new(), }, tool_timeout: None, + catalog_item_limit: crate::pagination::MAX_MCP_CATALOG_ITEMS, }, ); } @@ -387,6 +388,7 @@ async fn legacy_tool_catalog_does_not_follow_pagination_cursor() -> anyhow::Resu "test", &client, Some(Duration::from_secs(5)), + crate::pagination::MAX_MCP_CATALOG_ITEMS, /*server_instructions*/ None, ) .await?; @@ -1287,6 +1289,81 @@ fn codex_apps_env_bearer_token_bypasses_shared_tools_cache() { )); } +#[tokio::test] +async fn codex_apps_extension_does_not_share_host_owned_tools_cache() -> anyhow::Result<()> { + let codex_home = tempdir()?; + let cache_key = ConnectorRuntimeContextKey::personal( + /*account_id*/ None, /*chatgpt_user_id*/ None, + ); + let codex_apps_tools_cache = ConnectorRuntimeManager::::default(); + let cache_context = + codex_apps_tools_cache.context(codex_home.path().to_path_buf(), cache_key.clone()); + store_current_tools( + &cache_context, + vec![create_test_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + )], + ); + + let server_config: McpServerConfig = + serde_json::from_value(serde_json::json!({ "url": "http://127.0.0.1:1" }))?; + let mut config = crate::mcp::tests::test_mcp_config(codex_home.path().to_path_buf()); + let mut catalog = crate::ResolvedMcpCatalog::builder(); + catalog.register(crate::McpServerRegistration::from_extension( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + "test-extension", + /*contribution_order*/ 0, + server_config.clone(), + )); + config.mcp_server_catalog = catalog.build(); + + let startup_cancellation_token = CancellationToken::new(); + startup_cancellation_token.cancel(); + let manager = McpConnectionSet::new( + /*previous*/ None, + McpPublicationGate::already_published(), + McpRuntimeInput { + config: Arc::new(config), + plugins_available: false, + ready_selected_capability_roots: Vec::new(), + mcp_servers: HashMap::from([( + CODEX_APPS_MCP_SERVER_NAME.to_string(), + EffectiveMcpServer::configured(server_config), + )]), + submit_id: "cache-ownership-test".to_string(), + tx_event: None, + startup_cancellation_token, + runtime_context: McpRuntimeContext::new( + Arc::new(environment_manager_without_environments()), + codex_home.path().to_path_buf(), + ), + codex_apps_tools_cache, + tool_catalog_cache: McpToolCatalogCache::default(), + codex_apps_tools_cache_key: cache_key, + supports_openai_form_elicitation: false, + auth: None, + codex_apps_auth_manager: None, + elicitation_reviewer: None, + elicitation_lifecycle: None, + }, + ElicitationRequestRouter::default(), + ) + .await; + + let client = manager.test_client(CODEX_APPS_MCP_SERVER_NAME); + assert!( + client.codex_apps_tools_cache_context.is_none(), + "an extension must not receive the host-owned Apps cache" + ); + assert!( + !client.has_cached_tools(), + "an extension must not expose cached host-owned Apps tools" + ); + + Ok(()) +} + #[tokio::test] async fn list_all_tools_uses_shared_codex_apps_cache_while_client_is_pending() { let codex_home = tempdir().expect("tempdir"); @@ -3367,6 +3444,7 @@ async fn manager_with_reusable_ready_server( metadata: McpServerMetadata::from(&server), tool_filter: ToolFilter::from_config(config), tool_timeout: Some(config.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT)), + catalog_item_limit: crate::pagination::MAX_MCP_CATALOG_ITEMS, }, ); manager @@ -3455,6 +3533,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a /*codex_apps_refresh_trigger*/ "test", &client, /*timeout*/ None, + crate::pagination::MAX_MCP_CATALOG_ITEMS, initialize.instructions.as_deref(), ) .await?; @@ -3496,6 +3575,7 @@ async fn reconciliation_reuses_connection_without_relisting_regular_tools() -> a metadata: McpServerMetadata::from(&server), tool_filter: ToolFilter::from_config(&config), tool_timeout: Some(config.tool_timeout_sec.unwrap_or(DEFAULT_TOOL_TIMEOUT)), + catalog_item_limit: crate::pagination::MAX_MCP_CATALOG_ITEMS, }, ); let previous = Arc::new(previous); diff --git a/codex-rs/codex-mcp/src/pagination.rs b/codex-rs/codex-mcp/src/pagination.rs index 9e9e618c83..9b22b3453e 100644 --- a/codex-rs/codex-mcp/src/pagination.rs +++ b/codex-rs/codex-mcp/src/pagination.rs @@ -7,13 +7,27 @@ use anyhow::anyhow; use rmcp::model::PaginatedRequestParams; const MAX_MCP_CATALOG_PAGES: usize = 100; -const MAX_MCP_CATALOG_ITEMS: usize = 2_048; +pub(crate) const MAX_MCP_CATALOG_ITEMS: usize = 2_048; +pub(crate) const MAX_CODEX_APPS_TOOL_CATALOG_ITEMS: usize = 8_192; 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, + fetch: F, +) -> Result> +where + F: FnMut(Option) -> Fut, + Fut: Future, Option)>>, +{ + collect_paginated_with_limit(method, overall_timeout, MAX_MCP_CATALOG_ITEMS, fetch).await +} + +pub(crate) async fn collect_paginated_with_limit( + method: &str, + overall_timeout: Option, + max_items: usize, mut fetch: F, ) -> Result> where @@ -37,9 +51,9 @@ where 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()) { + if items.len() > max_items.saturating_sub(collected.len()) { return Err(anyhow!( - "{method} exceeded the catalog limit of {MAX_MCP_CATALOG_ITEMS} items" + "{method} exceeded the catalog limit of {max_items} items" )); } collected.extend(items); diff --git a/codex-rs/codex-mcp/src/rmcp_client.rs b/codex-rs/codex-mcp/src/rmcp_client.rs index eac72c0e72..3a0e81984e 100644 --- a/codex-rs/codex-mcp/src/rmcp_client.rs +++ b/codex-rs/codex-mcp/src/rmcp_client.rs @@ -26,7 +26,7 @@ use crate::elicitation::ElicitationRequestManager; use crate::mcp::CODEX_APPS_MCP_SERVER_NAME; use crate::mcp::ToolPluginProvenance; use crate::openai_docs_source_attribution::maybe_with_openai_docs_source_attribution; -use crate::pagination::collect_paginated; +use crate::pagination::collect_paginated_with_limit; use crate::runtime::McpRuntimeContext; use crate::runtime::emit_duration; use crate::server::EffectiveMcpServer; @@ -285,6 +285,7 @@ struct ManagedClientStartup { client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, protocol_mode: McpProtocolMode, + catalog_item_limit: usize, cancel_token: CancellationToken, startup_complete: Arc, } @@ -306,6 +307,7 @@ impl ManagedClientStartup { client_elicitation_capability, supports_openai_form_elicitation, protocol_mode, + catalog_item_limit, cancel_token, startup_complete, } = self.clone(); @@ -360,6 +362,7 @@ impl ManagedClientStartup { tool_catalog_fetch_ticket, client_elicitation_capability, supports_openai_form_elicitation, + catalog_item_limit, }, ) .await @@ -423,6 +426,7 @@ impl AsyncManagedClient { client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, protocol_mode: McpProtocolMode, + catalog_item_limit: usize, ) -> Self { let is_codex_apps_mcp_server = server_name == CODEX_APPS_MCP_SERVER_NAME; let reconnect_server_name = server_name.clone(); @@ -450,6 +454,7 @@ impl AsyncManagedClient { client_elicitation_capability, supports_openai_form_elicitation, protocol_mode, + catalog_item_limit, cancel_token: cancel_token.clone(), startup_complete: Arc::clone(&startup_complete), }); @@ -590,11 +595,12 @@ pub(crate) async fn list_tools_for_client_uncached( codex_apps_refresh_trigger: &'static str, client: &Arc, timeout: Option, + catalog_item_limit: usize, server_instructions: Option<&str>, ) -> Result> { let fetch_start = Instant::now(); let protocol_mode = client.protocol_mode(); - let tools = collect_paginated("tools/list", timeout, |params| { + let tools = collect_paginated_with_limit("tools/list", timeout, catalog_item_limit, |params| { let client = Arc::clone(client); async move { let response = client @@ -850,6 +856,7 @@ async fn start_server_task( tool_catalog_fetch_ticket, client_elicitation_capability, supports_openai_form_elicitation, + catalog_item_limit, } = params; let params = mcp_initialize_request_params( client_elicitation_capability, @@ -891,6 +898,7 @@ async fn start_server_task( /*codex_apps_refresh_trigger*/ "initial", &client, startup_timeout, + catalog_item_limit, initialize_result.instructions.as_deref(), ) .await @@ -982,6 +990,7 @@ struct StartServerTaskParams { tool_catalog_fetch_ticket: Option, client_elicitation_capability: ElicitationCapability, supports_openai_form_elicitation: bool, + catalog_item_limit: usize, } #[allow(clippy::too_many_arguments)] diff --git a/codex-rs/core/tests/suite/mcp_tool_exposure.rs b/codex-rs/core/tests/suite/mcp_tool_exposure.rs index 627d1abd71..38fa96e775 100644 --- a/codex-rs/core/tests/suite/mcp_tool_exposure.rs +++ b/codex-rs/core/tests/suite/mcp_tool_exposure.rs @@ -49,7 +49,14 @@ use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; use std::time::Duration; +use tempfile::TempDir; use tokio::sync::Semaphore; +use wiremock::Mock; +use wiremock::Request; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::method; +use wiremock::matchers::path_regex; struct McpResourceClientCapture { client: Arc>>, @@ -63,6 +70,7 @@ struct CoalescingMcpContributor { } struct AppsMcpServerContributor { + id: &'static str, url: String, } @@ -116,7 +124,7 @@ impl McpServerContributor for CoalescingMcpContributor { impl McpServerContributor for AppsMcpServerContributor { fn id(&self) -> &'static str { - "deferred_apps_recovery_test" + self.id } fn contribute<'a>( @@ -330,6 +338,107 @@ startup_timeout_sec = 0.1 Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn elevated_apps_catalog_limit_requires_host_owned_registration() -> Result<()> { + skip_if_no_network!(Ok(())); + + let codex_home = Arc::new(TempDir::new()?); + for extension_id in [None, Some("hosted_plugin_runtime"), Some("test-extension")] { + let server = responses::start_mock_server().await; + let apps_server = AppsTestServer::mount_searchable(&server).await?; + let tools = Arc::new( + (0..2_603) + .map(|index| { + json!({ + "name": format!("calendar_catalog_tool_{index}"), + "description": format!("Read calendar catalog entry {index}."), + "inputSchema": { + "type": "object", + "properties": {}, + "additionalProperties": false, + }, + "_meta": { + "connector_id": "calendar", + "connector_name": "Calendar", + "connector_description": "Plan events and manage your calendar.", + }, + }) + }) + .collect::>(), + ); + Mock::given(method("POST")) + .and(path_regex("^/api/codex/ps/mcp/?$")) + .and(body_partial_json(json!({ "method": "tools/list" }))) + .respond_with(move |request: &Request| { + let body: Value = serde_json::from_slice(&request.body) + .expect("Apps tools/list should be a valid JSON-RPC request"); + ResponseTemplate::new(200).set_body_json(json!({ + "jsonrpc": "2.0", + "id": body.get("id").cloned().unwrap_or(Value::Null), + "result": { + "tools": tools.as_ref(), + }, + })) + }) + .with_priority(1) + .mount(&server) + .await; + + let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone()) + .with_home(Arc::clone(&codex_home)); + if let Some(id) = extension_id { + let mut extensions = ExtensionRegistryBuilder::new(); + extensions.mcp_server_contributor(Arc::new(AppsMcpServerContributor { + id, + url: format!("{}/api/codex/ps/mcp", apps_server.chatgpt_base_url), + })); + builder = builder.with_extensions(Arc::new(extensions.build())); + } + let test = builder.build_with_auto_env(&server).await?; + let startup = wait_for_mcp_server(&test.codex, CODEX_APPS_MCP_SERVER_NAME).await; + + if extension_id == Some("test-extension") { + let error = startup.expect_err("an extension must retain the standard catalog limit"); + assert!( + error.to_string().contains("catalog limit of 2048 items"), + "an extension named codex_apps must not inherit the trusted Apps limit: {error}" + ); + continue; + } + + startup?; + let response = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + test.submit_turn("inspect the large Apps tool catalog") + .await?; + let body = response.single_request().body_json(); + let description = body["tools"] + .as_array() + .and_then(|tools| { + tools.iter().find_map(|tool| { + (tool.get("type").and_then(Value::as_str) == Some("tool_search")) + .then(|| tool.get("description").and_then(Value::as_str)) + .flatten() + }) + }) + .expect("large Apps catalogs should remain discoverable through tool_search"); + assert!( + description.contains("Calendar"), + "the accepted Apps catalog should remain model-discoverable: {description}" + ); + test.codex.shutdown_and_wait().await?; + } + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_only_exposes_direct_model_only_mcp_namespaces() -> Result<()> { skip_if_no_network!(Ok(())); @@ -618,6 +727,7 @@ async fn apps_guidance_and_deferred_namespace_appear_after_recovery_within_a_tur let release_apps_recovery = startup_control.hold_next_successful_initialize(); let mut extensions = ExtensionRegistryBuilder::new(); extensions.mcp_server_contributor(Arc::new(AppsMcpServerContributor { + id: "deferred_apps_recovery_test", url: format!("{}/api/codex/ps/mcp", gated_apps_server.chatgpt_base_url), })); let call_id = "pause-for-apps";