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

@@ -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"
         },

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<String, ToolInfo> {
-    (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::<String>::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<String> = 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;

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;

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<String, ToolInfo> {
+    (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<String> = 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<String> = 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"));
+}

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(()));

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",