Add ranked tool discovery to JavaScript code mode (#51209)

## What changed

- Add the default-disabled `code_mode_tool_search` feature to expose `await tools.tool_search({query: "...", limit: 8})` in code mode when the model supports tool search.
- Return BM25-ranked deferred tools as `{tools: [{name, description}]}`, including callable names and TypeScript declarations for use in later code cells.
- Filter results through the current step's nested-tool policy before applying the limit, preserving namespace restrictions, normalized identities, description prefixes, and schema budgets.
- Add ranked-search guidance while retaining `ALL_TOOLS` catalog discovery and description overrides.

## Testing

Add coverage for feature and model-capability gating, namespace restrictions, argument validation, schema budgets, and discovery followed by deferred MCP and dynamic tool calls. Extend the scenario snapshot to cover search, invocation, yielding, and completion through `wait`.

GitOrigin-RevId: 092a29dd5f0773c518f53a001b9cafb917422f13

codex-rs/core/src/tools/code_mode/execute_handler.rs

@@ -8,6 +8,7 @@ use codex_prompts::ResolvedModelMessages;
 use codex_tools::IndirectNamespacePrefixes;
 use codex_tools::ToolName;
 use codex_tools::ToolSpec;
+use std::borrow::Cow;
 use std::sync::Arc;
 
 use super::ExecContext;
@@ -15,6 +16,7 @@ use super::PUBLIC_TOOL_NAME;
 use super::handle_runtime_response;
 use super::is_exec_tool_name;
 use super::output::CodeModeToolOutput;
+use super::prepare_code_mode_tool_definitions;
 use super::telemetry::CodeModeToolCallGuard;
 use super::telemetry::trace_id;
 
@@ -55,22 +57,15 @@ impl CodeModeExecuteHandler {
             .tool_input_schema_max_bytes;
         let mut enabled_tools = Vec::with_capacity(self.nested_tool_specs.len());
         for (spec, cached_runtime) in &self.nested_tool_specs {
-            if let Some(cached_definitions) = cached_runtime.as_ref().and_then(|runtime| {
-                runtime.cached_code_mode_definitions(code_mode_input_schema_max_bytes)
-            }) {
-                enabled_tools.extend_from_slice(cached_definitions);
-                continue;
-            }
-
-            let definitions = codex_tools::collect_code_mode_tool_definitions(
-                std::iter::once(spec.as_ref()),
+            let definitions = prepare_code_mode_tool_definitions(
+                cached_runtime.as_deref(),
+                || Cow::Borrowed(spec.as_ref()),
                 code_mode_input_schema_max_bytes,
             );
-            enabled_tools.extend(definitions.into_iter().map(|mut definition| {
-                definition.input_schema = None;
-                definition.output_schema = None;
-                definition
-            }));
+            match definitions {
+                Cow::Borrowed(definitions) => enabled_tools.extend_from_slice(definitions),
+                Cow::Owned(definitions) => enabled_tools.extend(definitions),
+            }
         }
         enabled_tools.sort_by(|left, right| left.name.cmp(&right.name));
         enabled_tools.dedup_by(|left, right| left.name == right.name);

codex-rs/core/src/tools/code_mode/execute_spec.rs

@@ -1,20 +1,8 @@
-use codex_code_mode::ImageDetailVisibility;
-use codex_code_mode::ToolDefinition as CodeModeToolDefinition;
-use codex_protocol::openai_models::CodeModeToolMessages;
 use codex_tools::FreeformTool;
 use codex_tools::FreeformToolFormat;
 use codex_tools::ToolSpec;
-use std::collections::BTreeMap;
 
-pub(crate) fn create_code_mode_tool(
-    enabled_tools: &[CodeModeToolDefinition],
-    deferred_tools: &[CodeModeToolDefinition],
-    namespace_descriptions: &BTreeMap<String, codex_code_mode::ToolNamespaceDescription>,
-    default_exec_yield_time_ms: u64,
-    code_mode_only: bool,
-    image_detail_visibility: ImageDetailVisibility,
-    messages: Option<&CodeModeToolMessages>,
-) -> ToolSpec {
+pub(crate) fn create_code_mode_tool(description: String) -> ToolSpec {
     const CODE_MODE_FREEFORM_GRAMMAR: &str = r#"
 start: pragma_source | plain_source
 pragma_source: PRAGMA_LINE NEWLINE SOURCE
@@ -27,15 +15,7 @@ SOURCE: /[\s\S]+/
 
     ToolSpec::Freeform(FreeformTool {
         name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
-        description: codex_code_mode::build_exec_tool_description(
-            enabled_tools,
-            deferred_tools,
-            namespace_descriptions,
-            default_exec_yield_time_ms,
-            code_mode_only,
-            image_detail_visibility,
-            messages,
-        ),
+        description,
         defer_loading: None,
         format: FreeformToolFormat {
             r#type: "grammar".to_string(),
@@ -48,42 +28,15 @@ SOURCE: /[\s\S]+/
 #[cfg(test)]
 mod tests {
     use super::*;
-    use codex_tools::ToolName;
     use pretty_assertions::assert_eq;
 
     #[test]
     fn create_code_mode_tool_matches_expected_spec() {
-        let enabled_tools = vec![codex_code_mode::ToolDefinition {
-            name: "update_plan".to_string(),
-            tool_name: ToolName::plain("update_plan"),
-            description: "Update the plan".to_string(),
-            kind: codex_code_mode::CodeModeToolKind::Function,
-            input_schema: None,
-            input_schema_max_bytes: None,
-            output_schema: None,
-        }];
-
         assert_eq!(
-            create_code_mode_tool(
-                &enabled_tools,
-                &[],
-                &BTreeMap::new(),
-                codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS,
-                /*code_mode_only*/ true,
-                ImageDetailVisibility::Visible,
-                /*messages*/ None,
-            ),
+            create_code_mode_tool("Run JavaScript code.".to_string()),
             ToolSpec::Freeform(FreeformTool {
                 name: codex_code_mode::PUBLIC_TOOL_NAME.to_string(),
-                description: codex_code_mode::build_exec_tool_description(
-                    &enabled_tools,
-                    &[],
-                    &BTreeMap::new(),
-                    codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS,
-                    /*code_mode_only*/ true,
-                    ImageDetailVisibility::Visible,
-                    /*messages*/ None,
-                ),
+                description: "Run JavaScript code.".to_string(),
                 defer_loading: None,
                 format: FreeformToolFormat {
                     r#type: "grammar".to_string(),

codex-rs/core/src/tools/code_mode/mod.rs

@@ -4,6 +4,7 @@ pub(crate) mod execute_spec;
 mod output;
 mod response_adapter;
 mod telemetry;
+mod tool_definitions;
 mod wait_handler;
 pub(crate) mod wait_spec;
 
@@ -54,6 +55,7 @@ use delegate::CodeModeDispatchWorker;
 pub(crate) use execute_handler::CodeModeExecuteHandler;
 use output::CodeModeToolOutput;
 use response_adapter::into_function_call_output_content_items;
+pub(crate) use tool_definitions::prepare_code_mode_tool_definitions;
 pub(crate) use wait_handler::CodeModeWaitHandler;
 
 pub(crate) const PUBLIC_TOOL_NAME: &str = codex_code_mode::PUBLIC_TOOL_NAME;

codex-rs/core/src/tools/code_mode/tool_definitions.rs

@@ -0,0 +1,30 @@
+//! Prepares nested tool definitions for execution and search.
+//! Cached definitions stay borrowed; rendered schemas are dropped after augmentation.
+
+use crate::tools::registry::CoreToolRuntime;
+use codex_code_mode::ToolDefinition;
+use codex_tools::ToolSpec;
+use std::borrow::Cow;
+
+pub(crate) fn prepare_code_mode_tool_definitions<'a>(
+    runtime: Option<&'a dyn CoreToolRuntime>,
+    spec: impl FnOnce() -> Cow<'a, ToolSpec>,
+    code_mode_input_schema_max_bytes: Option<usize>,
+) -> Cow<'a, [ToolDefinition]> {
+    if let Some(definitions) = runtime
+        .and_then(|runtime| runtime.cached_code_mode_definitions(code_mode_input_schema_max_bytes))
+    {
+        return Cow::Borrowed(definitions);
+    }
+
+    let spec = spec();
+    let mut definitions = codex_tools::collect_code_mode_tool_definitions(
+        std::iter::once(spec.as_ref()),
+        code_mode_input_schema_max_bytes,
+    );
+    for definition in &mut definitions {
+        definition.input_schema = None;
+        definition.output_schema = None;
+    }
+    Cow::Owned(definitions)
+}

codex-rs/core/src/tools/handlers/tool_search.rs

@@ -1,4 +1,6 @@
 use crate::function_tool::FunctionCallError;
+use crate::tools::code_mode::prepare_code_mode_tool_definitions;
+use crate::tools::context::ToolCallSource;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolPayload;
 use crate::tools::context::ToolSearchOutput;
@@ -8,12 +10,14 @@ use crate::tools::handlers::tool_search_spec::create_tool_search_tool;
 use crate::tools::registry::CoreToolRuntime;
 use crate::tools::registry::ToolExecutor;
 use crate::tools::registry::ToolRegistry;
-use bm25::Document;
+use crate::tools::router::ToolRouter;
+use bm25::Embedder;
+use bm25::EmbedderBuilder;
 use bm25::Language;
-use bm25::SearchEngine;
-use bm25::SearchEngineBuilder;
+use bm25::Scorer;
 use codex_prompts::ResolvedModelMessages;
 use codex_tools::IndirectNamespacePrefixes;
+use codex_tools::JsonToolOutput;
 use codex_tools::LoadableToolSpec;
 use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT;
 use codex_tools::TOOL_SEARCH_TOOL_NAME;
@@ -22,16 +26,25 @@ use codex_tools::ToolSearchEntry;
 use codex_tools::ToolSearchInfo;
 use codex_tools::ToolSpec;
 use codex_tools::coalesce_loadable_tool_specs;
+use serde::Serialize;
+use std::borrow::Cow;
 use std::sync::Arc;
 use std::sync::Mutex;
 use std::sync::Weak;
 use tracing::instrument;
 
+#[derive(Serialize)]
+struct CodeModeSearchResult {
+    name: String,
+    description: String,
+}
+
 pub struct ToolSearchHandler {
     search_infos: Vec<ToolSearchInfo>,
     source_listing: ToolSearchSourceListing,
     spec: ToolSpec,
-    search_engine: SearchEngine<usize>,
+    search_embedder: Embedder,
+    search_scorer: Scorer<usize>,
 }
 
 #[derive(Default)]
@@ -150,20 +163,23 @@ impl ToolSearchHandler {
             TOOL_SEARCH_DEFAULT_LIMIT,
             source_listing,
         );
-        let documents: Vec<Document<usize>> = search_infos
+        let corpus = search_infos
             .iter()
-            .map(|search_info| search_info.entry.search_text.clone())
-            .enumerate()
-            .map(|(idx, search_text)| Document::new(idx, search_text))
-            .collect();
-        let search_engine =
-            SearchEngineBuilder::<usize>::with_documents(Language::English, documents).build();
+            .map(|search_info| search_info.entry.search_text.as_str())
+            .collect::<Vec<_>>();
+        let search_embedder =
+            EmbedderBuilder::with_fit_to_corpus(Language::English, &corpus).build();
+        let mut search_scorer = Scorer::new();
+        for (id, text) in corpus.into_iter().enumerate() {
+            search_scorer.upsert(&id, search_embedder.embed(text));
+        }
 
         Self {
             search_infos,
             source_listing,
             spec,
-            search_engine,
+            search_embedder,
+            search_scorer,
         }
     }
 }
@@ -196,12 +212,15 @@ impl ToolSearchHandler {
     ) -> Result<Box<dyn crate::tools::context::ToolOutput>, FunctionCallError> {
         let ToolInvocation {
             payload,
+            source,
             step_context,
             ..
         } = invocation;
+        let nested = matches!(source, ToolCallSource::CodeMode { .. });
 
         let args = match payload {
             ToolPayload::ToolSearch { arguments } => arguments,
+            ToolPayload::Function { arguments } if nested => super::parse_arguments(&arguments)?,
             _ => {
                 return Err(FunctionCallError::Fatal(format!(
                     "{TOOL_SEARCH_TOOL_NAME} handler received unsupported payload"
@@ -223,18 +242,36 @@ impl ToolSearchHandler {
             ));
         }
 
+        let model_messages = ResolvedModelMessages::from_model(&step_context.settings.model_info);
+        let indirect_prefixes = IndirectNamespacePrefixes::new(
+            model_messages.indirect_description_prefixes(),
+            step_context.tool_router.mcp_namespaces(),
+        )
+        .map_err(|error| FunctionCallError::Fatal(error.to_string()))?;
+
+        if nested {
+            let tools = self.search_code_mode(
+                query,
+                limit,
+                &step_context.tool_router,
+                step_context
+                    .turn
+                    .config
+                    .code_mode
+                    .tool_input_schema_max_bytes,
+                &indirect_prefixes,
+            );
+            return Ok(boxed_tool_output(JsonToolOutput::new(serde_json::json!({
+                "tools": tools,
+            }))));
+        }
+
         if self.search_infos.is_empty() {
             return Ok(boxed_tool_output(ToolSearchOutput { tools: Vec::new() }));
         }
 
-        let mut tools = self.search(query, limit)?;
-        let model_messages = ResolvedModelMessages::from_model(&step_context.settings.model_info);
-        IndirectNamespacePrefixes::new(
-            model_messages.indirect_description_prefixes(),
-            step_context.tool_router.mcp_namespaces(),
-        )
-        .map_err(|error| FunctionCallError::Fatal(error.to_string()))?
-        .apply_search(&mut tools);
+        let mut tools = self.search(query, limit);
+        indirect_prefixes.apply_search(&mut tools);
 
         Ok(boxed_tool_output(ToolSearchOutput { tools }))
     }
@@ -243,28 +280,71 @@ impl ToolSearchHandler {
 impl CoreToolRuntime for ToolSearchHandler {}
 
 impl ToolSearchHandler {
-    fn search(
+    fn search_code_mode(
         &self,
         query: &str,
         limit: usize,
-    ) -> Result<Vec<LoadableToolSpec>, FunctionCallError> {
-        let results = self
-            .search_engine
-            .search(query, limit)
-            .into_iter()
-            .map(|result| result.document.id)
-            .filter_map(|id| self.search_infos.get(id))
-            .map(|search_info| &search_info.entry);
-        self.search_output_tools(results)
+        router: &ToolRouter,
+        code_mode_input_schema_max_bytes: Option<usize>,
+        indirect_prefixes: &IndirectNamespacePrefixes<'_>,
+    ) -> Vec<CodeModeSearchResult> {
+        let mut tools = Vec::new();
+        // Apply the exact step's nested-tool policy before limiting ranked results.
+        for entry in self.search_entries(query) {
+            for candidate_name in entry.tool_names() {
+                let candidate_name = candidate_name.with_default_namespace();
+                let name = codex_code_mode::normalize_code_mode_identifier(
+                    &codex_tools::code_mode_name_for_tool_name(&candidate_name),
+                );
+                let Some(tool_name) =
+                    router
+                        .code_mode_tool_names()
+                        .get(&name)
+                        .filter(|tool_name| {
+                            (*tool_name).clone().with_default_namespace() == candidate_name
+                        })
+                else {
+                    continue;
+                };
+                let Some(runtime) = router.tool_runtime(tool_name) else {
+                    continue;
+                };
+                let definitions = prepare_code_mode_tool_definitions(
+                    Some(runtime.as_ref()),
+                    || Cow::Owned(runtime.spec()),
+                    code_mode_input_schema_max_bytes,
+                );
+                let Some(definition) = definitions.iter().find(|definition| {
+                    definition.tool_name.clone().with_default_namespace() == candidate_name
+                }) else {
+                    continue;
+                };
+                let mut description = definition.description.clone();
+                indirect_prefixes
+                    .apply_code_mode_description(&definition.tool_name, &mut description);
+                tools.push(CodeModeSearchResult { name, description });
+                if tools.len() == limit {
+                    return tools;
+                }
+            }
+        }
+        tools
     }
 
-    fn search_output_tools<'a>(
-        &self,
-        results: impl IntoIterator<Item = &'a ToolSearchEntry>,
-    ) -> Result<Vec<LoadableToolSpec>, FunctionCallError> {
-        Ok(coalesce_loadable_tool_specs(
-            results.into_iter().map(ToolSearchEntry::to_loadable_spec),
-        ))
+    fn search(&self, query: &str, limit: usize) -> Vec<LoadableToolSpec> {
+        coalesce_loadable_tool_specs(
+            self.search_entries(query)
+                .take(limit)
+                .map(ToolSearchEntry::to_loadable_spec),
+        )
+    }
+
+    fn search_entries(&self, query: &str) -> impl Iterator<Item = &ToolSearchEntry> + '_ {
+        self.search_scorer
+            .matches(&self.search_embedder.embed(query))
+            .into_iter()
+            .filter_map(|result| self.search_infos.get(result.id))
+            .map(|search_info| &search_info.entry)
     }
 }
 
@@ -413,9 +493,9 @@ mod tests {
             &handler.search_infos[1].entry,
         ];
 
-        let tools = handler
-            .search_output_tools(results)
-            .expect("mixed search output should serialize");
+        let tools = coalesce_loadable_tool_specs(
+            results.into_iter().map(ToolSearchEntry::to_loadable_spec),
+        );
 
         assert_eq!(
             tools,

codex-rs/core/src/tools/registry.rs

@@ -64,7 +64,8 @@ pub(crate) trait CoreToolRuntime: ToolExecutor<ToolInvocation> {
         None
     }
 
-    /// Returns lazily cached Code Mode definitions owned by this runtime.
+    /// Returns lazily cached, augmented Code Mode definitions owned by this runtime.
+    /// Input and output schemas must be cleared after rendering their declarations.
     fn cached_code_mode_definitions(
         &self,
         _code_mode_input_schema_max_bytes: Option<usize>,

codex-rs/core/src/tools/router.rs

@@ -148,8 +148,6 @@ impl ToolRouter {
     }
 
     /// The normalized nested identities chosen after exclusions and collisions.
-    // Consumed by the follow-up cell-origin migration.
-    #[cfg_attr(not(test), allow(dead_code))]
     pub(crate) fn code_mode_tool_names(&self) -> &BTreeMap<String, ToolName> {
         &self.code_mode_tool_names
     }

codex-rs/core/src/tools/spec_plan.rs

@@ -553,9 +553,21 @@ fn apply_direct_model_only_namespace_overrides(
     }
 
     for tool in registry.entries_mut() {
-        let configured = tool
-            .runtime
-            .tool_name()
+        tool.exposure = tool_exposure_with_namespace_override(
+            turn_context,
+            tool.runtime.tool_name(),
+            tool.exposure,
+        );
+    }
+}
+
+fn tool_exposure_with_namespace_override(
+    turn_context: &TurnContext,
+    tool_name: ToolName,
+    exposure: ToolExposure,
+) -> ToolExposure {
+    if exposure.is_available_in_code_mode()
+        && tool_name
             .with_default_namespace()
             .namespace
             .as_ref()
@@ -565,10 +577,11 @@ fn apply_direct_model_only_namespace_overrides(
                     .code_mode
                     .direct_only_tool_namespaces
                     .contains(namespace)
-            });
-        if configured && tool.exposure.is_available_in_code_mode() {
-            tool.exposure = ToolExposure::DirectModelOnly;
-        }
+            })
+    {
+        ToolExposure::DirectModelOnly
+    } else {
+        exposure
     }
 }
 
@@ -892,6 +905,14 @@ fn register_code_mode_executors(
             ToolSpec::Namespace(namespace) if !namespace.tools.is_empty() => {
                 codex_tools::code_mode_name_for_tool_name(&tool_name)
             }
+            ToolSpec::ToolSearch { .. }
+                if turn_context
+                    .config
+                    .features
+                    .enabled(Feature::CodeModeToolSearch) =>
+            {
+                codex_tools::code_mode_name_for_tool_name(&tool_name)
+            }
             ToolSpec::Namespace(_) | ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => {
                 continue;
             }
@@ -930,6 +951,14 @@ fn register_code_mode_executors(
         code_mode_nested_tool_specs.push((spec, cached_runtime));
     }
 
+    let deferred_tool_discovery = if exec_prompt_tool_specs
+        .iter()
+        .any(|spec| matches!(spec, ToolSpec::ToolSearch { .. }))
+    {
+        codex_code_mode::DeferredToolDiscovery::RankedSearch
+    } else {
+        codex_code_mode::DeferredToolDiscovery::Catalog
+    };
     let mut namespace_descriptions = code_mode_namespace_descriptions(&exec_prompt_tool_specs);
     let code_mode_input_schema_max_bytes =
         turn_context.config.code_mode.tool_input_schema_max_bytes;
@@ -947,20 +976,22 @@ fn register_code_mode_executors(
     }
     enabled_tools
         .sort_by(|left, right| compare_code_mode_tools(left, right, &namespace_descriptions));
+    let exec_description = codex_code_mode::build_exec_tool_description(
+        &enabled_tools,
+        &deferred_tools,
+        &namespace_descriptions,
+        turn_context.config.code_mode.default_exec_yield_time_ms,
+        tool_mode == ToolMode::CodeModeOnly,
+        if unified_image_budget_enabled(&turn_context.config.features, model_info) {
+            codex_code_mode::ImageDetailVisibility::Hidden
+        } else {
+            codex_code_mode::ImageDetailVisibility::Visible
+        },
+        deferred_tool_discovery,
+        model_messages.code_mode(),
+    );
     let execute_handler = CodeModeExecuteHandler::new(
-        create_code_mode_tool(
-            &enabled_tools,
-            &deferred_tools,
-            &namespace_descriptions,
-            turn_context.config.code_mode.default_exec_yield_time_ms,
-            tool_mode == ToolMode::CodeModeOnly,
-            if unified_image_budget_enabled(&turn_context.config.features, model_info) {
-                codex_code_mode::ImageDetailVisibility::Hidden
-            } else {
-                codex_code_mode::ImageDetailVisibility::Visible
-            },
-            model_messages.code_mode(),
-        ),
+        create_code_mode_tool(exec_description),
         code_mode_nested_tool_specs,
     );
 
@@ -1521,7 +1552,17 @@ fn append_tool_search_executor(
         ToolSearchSourceListing::Include
     };
     let handler = tool_search_handler_cache.get_or_build(registry, source_listing);
-    registry.register_trusted(handler);
+    // Preserve native search's exposure when nested search is disabled.
+    let exposure = if turn_context
+        .config
+        .features
+        .enabled(Feature::CodeModeToolSearch)
+    {
+        tool_exposure_with_namespace_override(turn_context, handler.tool_name(), handler.exposure())
+    } else {
+        handler.exposure()
+    };
+    registry.register_trusted_with_exposure(handler, exposure);
 }
 
 fn append_extension_tool_executors(

codex-rs/core/src/tools/spec_plan_tests.rs

@@ -1381,6 +1381,226 @@ async fn mcp_and_tool_search_follow_direct_and_deferred_tool_exposure() {
     ));
 }
 
+#[tokio::test]
+async fn code_mode_tool_search_requires_experiment_and_model_capability() {
+    for tool_mode in [ToolMode::Direct, ToolMode::CodeMode, ToolMode::CodeModeOnly] {
+        for experiment_enabled in [false, true] {
+            for supports_search_tool in [false, true] {
+                let plan = probe_with(
+                    |turn| {
+                        set_feature(turn, Feature::Collab, /*enabled*/ false);
+                        set_feature(turn, Feature::CodeModeToolSearch, experiment_enabled);
+                        set_feature(turn, Feature::CodeMode, tool_mode != ToolMode::Direct);
+                        set_feature(
+                            turn,
+                            Feature::CodeModeOnly,
+                            tool_mode == ToolMode::CodeModeOnly,
+                        );
+                        update_turn_settings_for_test(turn, |settings| {
+                            Arc::make_mut(&mut settings.model_info).supports_search_tool =
+                                supports_search_tool;
+                        });
+                    },
+                    ToolPlanInputs {
+                        tool_runtimes: vec![mcp_runtime(
+                            "searchable",
+                            "mcp__searchable",
+                            "lookup",
+                            ToolExposure::Deferred,
+                        )],
+                        ..ToolPlanInputs::default()
+                    },
+                )
+                .await;
+
+                let nested_search_enabled =
+                    experiment_enabled && supports_search_tool && tool_mode != ToolMode::Direct;
+                assert_eq!(
+                    (
+                        plan.tool_mode,
+                        plan.registered_names
+                            .iter()
+                            .any(|name| name == "tool_search"),
+                        plan.visible_names.iter().any(|name| name == "tool_search"),
+                        plan.code_mode_tool_names.get("tool_search").cloned(),
+                    ),
+                    (
+                        tool_mode,
+                        supports_search_tool,
+                        supports_search_tool && tool_mode != ToolMode::CodeModeOnly,
+                        nested_search_enabled.then(|| ToolName::plain("tool_search")),
+                    ),
+                    "tool mode {tool_mode:?}, experiment {experiment_enabled}, model capability {supports_search_tool}",
+                );
+                if tool_mode != ToolMode::Direct {
+                    let ToolSpec::Freeform(exec) =
+                        plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME)
+                    else {
+                        panic!("expected code mode exec tool");
+                    };
+                    assert_eq!(
+                        exec.description.contains("await tools.tool_search("),
+                        nested_search_enabled,
+                    );
+                }
+            }
+        }
+    }
+}
+
+#[tokio::test]
+async fn code_mode_tool_search_respects_direct_only_functions() {
+    for (code_mode_only, nested_search_enabled) in
+        [(false, false), (false, true), (true, false), (true, true)]
+    {
+        let plan = probe_with(
+            |turn| {
+                set_feature(turn, Feature::Collab, /*enabled*/ false);
+                set_feature(turn, Feature::CodeMode, /*enabled*/ true);
+                set_feature(turn, Feature::CodeModeOnly, code_mode_only);
+                set_feature(turn, Feature::CodeModeToolSearch, nested_search_enabled);
+                update_turn_settings_for_test(turn, |settings| {
+                    Arc::make_mut(&mut settings.model_info).supports_search_tool = true;
+                });
+                update_config(turn, |config| {
+                    config.code_mode.direct_only_tool_namespaces = vec!["functions".to_string()];
+                });
+            },
+            ToolPlanInputs {
+                tool_runtimes: vec![mcp_runtime(
+                    "searchable",
+                    "mcp__searchable",
+                    "lookup",
+                    ToolExposure::Deferred,
+                )],
+                ..ToolPlanInputs::default()
+            },
+        )
+        .await;
+
+        assert_eq!(
+            plan.visible_names.iter().any(|name| name == "tool_search"),
+            !code_mode_only || nested_search_enabled,
+        );
+        assert_eq!(
+            plan.exposure("tool_search"),
+            if nested_search_enabled {
+                ToolExposure::DirectModelOnly
+            } else {
+                ToolExposure::Direct
+            },
+        );
+        assert_eq!(plan.code_mode_tool_names.get("tool_search"), None);
+        let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else {
+            panic!("expected code mode exec tool");
+        };
+        assert!(!exec.description.contains("await tools.tool_search("));
+    }
+}
+
+#[tokio::test]
+async fn dynamic_tool_search_does_not_enable_ranked_search_guidance() {
+    for (experiment_enabled, supports_search_tool, has_deferred_tools) in [
+        (false, false, true),
+        (true, false, true),
+        (true, true, false),
+    ] {
+        let plan = probe_with(
+            |turn| {
+                set_features(turn, &[Feature::CodeMode, Feature::CodeModeOnly]);
+                set_feature(turn, Feature::Collab, /*enabled*/ false);
+                set_feature(turn, Feature::CodeModeToolSearch, experiment_enabled);
+                update_turn_settings_for_test(turn, |settings| {
+                    Arc::make_mut(&mut settings.model_info).supports_search_tool =
+                        supports_search_tool;
+                });
+            },
+            ToolPlanInputs {
+                dynamic_tools: vec![dynamic_tool(
+                    /*namespace*/ None,
+                    "tool_search",
+                    /*defer_loading*/ false,
+                )],
+                tool_runtimes: if has_deferred_tools {
+                    vec![mcp_runtime(
+                        "searchable",
+                        "mcp__searchable",
+                        "lookup",
+                        ToolExposure::Deferred,
+                    )]
+                } else {
+                    Vec::new()
+                },
+                ..ToolPlanInputs::default()
+            },
+        )
+        .await;
+
+        assert_eq!(
+            plan.code_mode_tool_names.get("tool_search"),
+            Some(&ToolName::plain("tool_search")),
+        );
+        let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else {
+            panic!("expected code mode exec tool");
+        };
+        assert!(exec.description.contains("tool_search dynamic tool"));
+        assert!(!exec.description.contains("await tools.tool_search("));
+    }
+}
+
+#[tokio::test]
+async fn code_mode_tool_search_respects_exclusions_and_requires_deferred_tools() {
+    for (exclude_functions, has_deferred_tools) in [(true, true), (false, false)] {
+        let plan = probe_with(
+            |turn| {
+                set_features(
+                    turn,
+                    &[
+                        Feature::CodeMode,
+                        Feature::CodeModeOnly,
+                        Feature::CodeModeToolSearch,
+                    ],
+                );
+                set_feature(turn, Feature::Collab, /*enabled*/ false);
+                update_turn_settings_for_test(turn, |settings| {
+                    Arc::make_mut(&mut settings.model_info).supports_search_tool = true;
+                });
+                if exclude_functions {
+                    update_config(turn, |config| {
+                        config.code_mode.excluded_tool_namespaces = vec!["functions".to_string()];
+                    });
+                }
+            },
+            ToolPlanInputs {
+                tool_runtimes: if has_deferred_tools {
+                    vec![mcp_runtime(
+                        "searchable",
+                        "mcp__searchable",
+                        "lookup",
+                        ToolExposure::Deferred,
+                    )]
+                } else {
+                    Vec::new()
+                },
+                ..ToolPlanInputs::default()
+            },
+        )
+        .await;
+
+        assert_eq!(plan.code_mode_tool_names.get("tool_search"), None);
+        assert_eq!(
+            plan.registered_names
+                .iter()
+                .any(|name| name == "tool_search"),
+            has_deferred_tools,
+        );
+        let ToolSpec::Freeform(exec) = plan.visible_spec(codex_code_mode::PUBLIC_TOOL_NAME) else {
+            panic!("expected code mode exec tool");
+        };
+        assert!(!exec.description.contains("await tools.tool_search("));
+    }
+}
+
 #[tokio::test]
 async fn tool_namespaces_info_is_opt_in_and_tracks_mcp_exposure() {
     for (enabled, use_responses_lite) in [(false, true), (true, false), (true, true)] {