From d0047de7cbd0d8ffd03804d3736478594f0742c8 Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Fri, 17 Apr 2026 18:34:06 +0800 Subject: [PATCH] add token-based tool deferral behind feature flag (#18097) add new `tool_search_always_defer_mcp_tools` feature flag that always defers all mcp tools rather than deferring once > 100 deferrable tools. add new tests, also move `mcp_exposure` tests into dedicated file rather than polluting `codex_tests`. --- codex-rs/core/config.schema.json | 6 + codex-rs/core/src/codex_tests.rs | 177 +------------ codex-rs/core/src/mcp_tool_exposure.rs | 13 +- codex-rs/core/src/mcp_tool_exposure_test.rs | 259 ++++++++++++++++++++ codex-rs/core/tests/suite/search_tool.rs | 46 ++++ codex-rs/features/src/lib.rs | 8 + 6 files changed, 332 insertions(+), 177 deletions(-) create mode 100644 codex-rs/core/src/mcp_tool_exposure_test.rs diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index eb44f5c8f..2eb8cf727 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -497,6 +497,9 @@ "tool_search": { "type": "boolean" }, + "tool_search_always_defer_mcp_tools": { + "type": "boolean" + }, "tool_suggest": { "type": "boolean" }, @@ -2390,6 +2393,9 @@ "tool_search": { "type": "boolean" }, + "tool_search_always_defer_mcp_tools": { + "type": "boolean" + }, "tool_suggest": { "type": "boolean" }, diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index aa5583405..2a05ccca0 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -10,15 +10,12 @@ use crate::config_loader::RequirementSource; use crate::config_loader::Sourced; use crate::exec::ExecCapturePolicy; use crate::function_tool::FunctionCallError; -use crate::mcp_tool_exposure::DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD; -use crate::mcp_tool_exposure::build_mcp_tool_exposure; use crate::shell::default_user_shell; use crate::tools::format_exec_output_str; +use codex_features::Feature; use codex_features::Features; use codex_login::CodexAuth; -use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; -use codex_mcp::ToolInfo; use codex_model_provider_info::ModelProviderInfo; use codex_models_manager::bundled_models_response; use codex_models_manager::model_info; @@ -117,8 +114,6 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; use codex_protocol::mcp::CallToolResult as McpCallToolResult; use pretty_assertions::assert_eq; -use rmcp::model::JsonObject; -use rmcp::model::Tool; use serde::Deserialize; use serde_json::json; use std::path::PathBuf; @@ -405,80 +400,6 @@ fn assistant_message_stream_parsers_seed_plan_parser_across_added_and_delta_boun assert!(tail.plan_segments.is_empty()); } -fn make_mcp_tool( - server_name: &str, - tool_name: &str, - connector_id: Option<&str>, - connector_name: Option<&str>, -) -> ToolInfo { - let tool_namespace = if server_name == CODEX_APPS_MCP_SERVER_NAME { - connector_name - .map(codex_connectors::metadata::sanitize_name) - .map(|connector_name| format!("mcp__{server_name}__{connector_name}")) - .unwrap_or_else(|| server_name.to_string()) - } else { - format!("mcp__{server_name}__") - }; - - ToolInfo { - server_name: server_name.to_string(), - callable_name: tool_name.to_string(), - callable_namespace: tool_namespace, - server_instructions: None, - tool: Tool { - name: tool_name.to_string().into(), - title: None, - description: Some(format!("Test tool: {tool_name}").into()), - input_schema: Arc::new(JsonObject::default()), - output_schema: None, - annotations: None, - execution: None, - icons: None, - meta: None, - }, - connector_id: connector_id.map(str::to_string), - connector_name: connector_name.map(str::to_string), - plugin_display_names: Vec::new(), - connector_description: None, - } -} - -fn numbered_mcp_tools(count: usize) -> HashMap { - (0..count) - .map(|index| { - let tool_name = format!("tool_{index}"); - ( - format!("mcp__rmcp__{tool_name}"), - make_mcp_tool( - "rmcp", &tool_name, /*connector_id*/ None, /*connector_name*/ None, - ), - ) - }) - .collect() -} - -async fn tools_config_for_mcp_tool_exposure(search_tool: bool) -> ToolsConfig { - let config = test_config().await; - let model_info = ModelsManager::construct_model_info_offline_for_tests( - "gpt-5-codex", - &config.to_models_manager_config(), - ); - let features = Features::with_defaults(); - let available_models = Vec::new(); - let mut tools_config = ToolsConfig::new(&ToolsConfigParams { - model_info: &model_info, - available_models: &available_models, - features: &features, - image_generation_tool_auth_allowed: true, - web_search_mode: Some(WebSearchMode::Cached), - session_source: SessionSource::Cli, - sandbox_policy: &SandboxPolicy::DangerFullAccess, - windows_sandbox_level: WindowsSandboxLevel::Disabled, - }); - tools_config.search_tool = search_tool; - tools_config -} - #[test] fn validated_network_policy_amendment_host_allows_normalized_match() { let amendment = NetworkPolicyAmendment { @@ -1146,102 +1067,6 @@ fn collect_explicit_app_ids_from_skill_items_skips_plain_mentions_with_skill_con assert_eq!(connector_ids, HashSet::::new()); } -#[tokio::test] -async fn mcp_tool_exposure_directly_exposes_small_effective_tool_sets() { - let config = test_config().await; - let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; - let mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1); - - let exposure = build_mcp_tool_exposure( - &mcp_tools, - /*connectors*/ None, - &[], - &config, - &tools_config, - ); - - let mut direct_tool_names: Vec<_> = exposure.direct_tools.keys().cloned().collect(); - direct_tool_names.sort(); - let mut expected_tool_names: Vec<_> = mcp_tools.keys().cloned().collect(); - expected_tool_names.sort(); - assert_eq!(direct_tool_names, expected_tool_names); - assert!(exposure.deferred_tools.is_none()); -} - -#[tokio::test] -async fn mcp_tool_exposure_searches_large_effective_tool_sets() { - let config = test_config().await; - let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; - let mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD); - - let exposure = build_mcp_tool_exposure( - &mcp_tools, - /*connectors*/ None, - &[], - &config, - &tools_config, - ); - - assert!(exposure.direct_tools.is_empty()); - let deferred_tools = exposure - .deferred_tools - .as_ref() - .expect("large tool sets should be discoverable through tool_search"); - let mut deferred_tool_names: Vec<_> = deferred_tools.keys().cloned().collect(); - deferred_tool_names.sort(); - let mut expected_tool_names: Vec<_> = mcp_tools.keys().cloned().collect(); - expected_tool_names.sort(); - assert_eq!(deferred_tool_names, expected_tool_names); -} - -#[tokio::test] -async fn mcp_tool_exposure_directly_exposes_explicit_apps_without_deferred_overlap() { - let config = test_config().await; - let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; - let mut mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1); - mcp_tools.extend([( - "mcp__codex_apps__calendar_create_event".to_string(), - make_mcp_tool( - CODEX_APPS_MCP_SERVER_NAME, - "calendar_create_event", - Some("calendar"), - Some("Calendar"), - ), - )]); - let connectors = vec![make_connector("calendar", "Calendar")]; - - let exposure = build_mcp_tool_exposure( - &mcp_tools, - Some(connectors.as_slice()), - connectors.as_slice(), - &config, - &tools_config, - ); - - let mut tool_names: Vec = exposure.direct_tools.into_keys().collect(); - tool_names.sort(); - assert_eq!( - tool_names, - vec!["mcp__codex_apps__calendar_create_event".to_string()] - ); - assert_eq!( - exposure.deferred_tools.as_ref().map(HashMap::len), - Some(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1) - ); - let deferred_tools = exposure - .deferred_tools - .as_ref() - .expect("large tool sets should be discoverable through tool_search"); - assert!( - tool_names - .iter() - .all(|direct_tool_name| !deferred_tools.contains_key(direct_tool_name)), - "direct tools should not also be deferred: {tool_names:?}" - ); - assert!(!deferred_tools.contains_key("mcp__codex_apps__calendar_create_event")); - assert!(deferred_tools.contains_key("mcp__rmcp__tool_0")); -} - #[tokio::test] async fn reconstruct_history_matches_live_compactions() { let (session, turn_context) = make_session_and_context().await; diff --git a/codex-rs/core/src/mcp_tool_exposure.rs b/codex-rs/core/src/mcp_tool_exposure.rs index 16aa50cfd..3917a1d5e 100644 --- a/codex-rs/core/src/mcp_tool_exposure.rs +++ b/codex-rs/core/src/mcp_tool_exposure.rs @@ -1,6 +1,7 @@ use std::collections::HashMap; use std::collections::HashSet; +use codex_features::Feature; use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; use codex_mcp::ToolInfo as McpToolInfo; use codex_mcp::filter_non_codex_apps_mcp_tools_only; @@ -32,7 +33,13 @@ pub(crate) fn build_mcp_tool_exposure( )); } - if !tools_config.search_tool || deferred_tools.len() < DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD { + let should_defer = tools_config.search_tool + && (config + .features + .enabled(Feature::ToolSearchAlwaysDeferMcpTools) + || deferred_tools.len() >= DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD); + + if !should_defer { return McpToolExposure { direct_tools: deferred_tools, deferred_tools: None, @@ -75,3 +82,7 @@ fn filter_codex_apps_mcp_tools( .map(|(name, tool)| (name.clone(), tool.clone())) .collect() } + +#[cfg(test)] +#[path = "mcp_tool_exposure_test.rs"] +mod tests; diff --git a/codex-rs/core/src/mcp_tool_exposure_test.rs b/codex-rs/core/src/mcp_tool_exposure_test.rs new file mode 100644 index 000000000..3cf463497 --- /dev/null +++ b/codex-rs/core/src/mcp_tool_exposure_test.rs @@ -0,0 +1,259 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use codex_connectors::metadata::sanitize_name; +use codex_features::Feature; +use codex_features::Features; +use codex_mcp::CODEX_APPS_MCP_SERVER_NAME; +use codex_mcp::ToolInfo; +use codex_models_manager::manager::ModelsManager; +use codex_protocol::config_types::WebSearchMode; +use codex_protocol::config_types::WindowsSandboxLevel; +use codex_protocol::protocol::SandboxPolicy; +use codex_protocol::protocol::SessionSource; +use codex_tools::ToolsConfig; +use codex_tools::ToolsConfigParams; +use pretty_assertions::assert_eq; +use rmcp::model::JsonObject; +use rmcp::model::Tool; + +use super::*; +use crate::config::test_config; +use crate::connectors::AppInfo; + +fn make_connector(id: &str, name: &str) -> AppInfo { + AppInfo { + id: id.to_string(), + name: name.to_string(), + description: None, + logo_url: None, + logo_url_dark: None, + distribution_channel: None, + branding: None, + app_metadata: None, + labels: None, + install_url: None, + is_accessible: true, + is_enabled: true, + plugin_display_names: Vec::new(), + } +} + +fn make_mcp_tool( + server_name: &str, + tool_name: &str, + connector_id: Option<&str>, + connector_name: Option<&str>, +) -> ToolInfo { + let tool_namespace = if server_name == CODEX_APPS_MCP_SERVER_NAME { + connector_name + .map(sanitize_name) + .map(|connector_name| format!("mcp__{server_name}__{connector_name}")) + .unwrap_or_else(|| server_name.to_string()) + } else { + format!("mcp__{server_name}__") + }; + + ToolInfo { + server_name: server_name.to_string(), + callable_name: tool_name.to_string(), + callable_namespace: tool_namespace, + server_instructions: None, + tool: Tool { + name: tool_name.to_string().into(), + title: None, + description: Some(format!("Test tool: {tool_name}").into()), + input_schema: Arc::new(JsonObject::default()), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: connector_id.map(str::to_string), + connector_name: connector_name.map(str::to_string), + plugin_display_names: Vec::new(), + connector_description: None, + } +} + +fn numbered_mcp_tools(count: usize) -> HashMap { + (0..count) + .map(|index| { + let tool_name = format!("tool_{index}"); + ( + format!("mcp__rmcp__{tool_name}"), + make_mcp_tool( + "rmcp", &tool_name, /*connector_id*/ None, /*connector_name*/ None, + ), + ) + }) + .collect() +} + +async fn tools_config_for_mcp_tool_exposure(search_tool: bool) -> ToolsConfig { + let config = test_config().await; + let model_info = ModelsManager::construct_model_info_offline_for_tests( + "gpt-5-codex", + &config.to_models_manager_config(), + ); + let features = Features::with_defaults(); + let available_models = Vec::new(); + let mut tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + sandbox_policy: &SandboxPolicy::DangerFullAccess, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + tools_config.search_tool = search_tool; + tools_config +} + +#[tokio::test] +async fn directly_exposes_small_effective_tool_sets() { + let config = test_config().await; + let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; + let mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1); + + let exposure = build_mcp_tool_exposure( + &mcp_tools, + /*connectors*/ None, + &[], + &config, + &tools_config, + ); + + let mut direct_tool_names: Vec<_> = exposure.direct_tools.keys().cloned().collect(); + direct_tool_names.sort(); + let mut expected_tool_names: Vec<_> = mcp_tools.keys().cloned().collect(); + expected_tool_names.sort(); + assert_eq!(direct_tool_names, expected_tool_names); + assert!(exposure.deferred_tools.is_none()); +} + +#[tokio::test] +async fn searches_large_effective_tool_sets() { + let config = test_config().await; + let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; + let mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD); + + let exposure = build_mcp_tool_exposure( + &mcp_tools, + /*connectors*/ None, + &[], + &config, + &tools_config, + ); + + assert!(exposure.direct_tools.is_empty()); + let deferred_tools = exposure + .deferred_tools + .as_ref() + .expect("large tool sets should be discoverable through tool_search"); + let mut deferred_tool_names: Vec<_> = deferred_tools.keys().cloned().collect(); + deferred_tool_names.sort(); + let mut expected_tool_names: Vec<_> = mcp_tools.keys().cloned().collect(); + expected_tool_names.sort(); + assert_eq!(deferred_tool_names, expected_tool_names); +} + +#[tokio::test] +async fn directly_exposes_explicit_apps_without_deferred_overlap() { + let config = test_config().await; + let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; + let mut mcp_tools = numbered_mcp_tools(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1); + mcp_tools.extend([( + "mcp__codex_apps__calendar_create_event".to_string(), + make_mcp_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + Some("calendar"), + Some("Calendar"), + ), + )]); + let connectors = vec![make_connector("calendar", "Calendar")]; + + let exposure = build_mcp_tool_exposure( + &mcp_tools, + Some(connectors.as_slice()), + connectors.as_slice(), + &config, + &tools_config, + ); + + let mut tool_names: Vec = exposure.direct_tools.into_keys().collect(); + tool_names.sort(); + assert_eq!( + tool_names, + vec!["mcp__codex_apps__calendar_create_event".to_string()] + ); + assert_eq!( + exposure.deferred_tools.as_ref().map(HashMap::len), + Some(DIRECT_MCP_TOOL_EXPOSURE_THRESHOLD - 1) + ); + let deferred_tools = exposure + .deferred_tools + .as_ref() + .expect("large tool sets should be discoverable through tool_search"); + assert!( + tool_names + .iter() + .all(|direct_tool_name| !deferred_tools.contains_key(direct_tool_name)), + "direct tools should not also be deferred: {tool_names:?}" + ); + assert!(!deferred_tools.contains_key("mcp__codex_apps__calendar_create_event")); + assert!(deferred_tools.contains_key("mcp__rmcp__tool_0")); +} + +#[tokio::test] +async fn always_defer_feature_preserves_explicit_apps() { + let mut config = test_config().await; + config + .features + .enable(Feature::ToolSearchAlwaysDeferMcpTools) + .expect("test config should allow feature update"); + let tools_config = tools_config_for_mcp_tool_exposure(/*search_tool*/ true).await; + let mcp_tools = HashMap::from([ + ( + "mcp__rmcp__tool".to_string(), + make_mcp_tool( + "rmcp", "tool", /*connector_id*/ None, /*connector_name*/ None, + ), + ), + ( + "mcp__codex_apps__calendar_create_event".to_string(), + make_mcp_tool( + CODEX_APPS_MCP_SERVER_NAME, + "calendar_create_event", + Some("calendar"), + Some("Calendar"), + ), + ), + ]); + let connectors = vec![make_connector("calendar", "Calendar")]; + + let exposure = build_mcp_tool_exposure( + &mcp_tools, + Some(connectors.as_slice()), + connectors.as_slice(), + &config, + &tools_config, + ); + + let mut direct_tool_names: Vec = exposure.direct_tools.into_keys().collect(); + direct_tool_names.sort(); + assert_eq!( + direct_tool_names, + vec!["mcp__codex_apps__calendar_create_event".to_string()] + ); + let deferred_tools = exposure + .deferred_tools + .as_ref() + .expect("MCP tools should be discoverable through tool_search"); + assert!(deferred_tools.contains_key("mcp__rmcp__tool")); + assert!(!deferred_tools.contains_key("mcp__codex_apps__calendar_create_event")); +} diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 41f8e623c..cfcec7992 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -189,6 +189,52 @@ async fn search_tool_enabled_by_default_adds_tool_search() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn always_defer_feature_hides_small_app_tool_sets() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let apps_server = AppsTestServer::mount(&server).await?; + let mock = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-1"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-1"), + ]), + ) + .await; + + let mut builder = + configured_builder(apps_server.chatgpt_base_url.clone()).with_config(|config| { + config + .features + .enable(Feature::ToolSearchAlwaysDeferMcpTools) + .expect("test config should allow feature update"); + }); + let test = builder.build(&server).await?; + + test.submit_turn_with_policies( + "list tools", + AskForApproval::Never, + SandboxPolicy::DangerFullAccess, + ) + .await?; + + let body = mock.single_request().body_json(); + let tools = tool_names(&body); + assert!( + tools.iter().any(|name| name == TOOL_SEARCH_TOOL_NAME), + "small app tool sets should be deferred behind tool_search: {tools:?}" + ); + assert!( + tools.iter().all(|name| !name.starts_with("mcp__")), + "MCP tools should not be directly exposed: {tools:?}" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tool_search_disabled_exposes_apps_tools_directly() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index 888f0a409..cef303502 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -148,6 +148,8 @@ pub enum Feature { Apps, /// Enable the tool_search tool for apps. ToolSearch, + /// Always defer MCP tools behind tool_search instead of exposing small sets directly. + ToolSearchAlwaysDeferMcpTools, /// Expose placeholder tools for unavailable historical tool calls. UnavailableDummyTools, /// Enable discoverable tool suggestions for apps. @@ -818,6 +820,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::ToolSearchAlwaysDeferMcpTools, + key: "tool_search_always_defer_mcp_tools", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::UnavailableDummyTools, key: "unavailable_dummy_tools",