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/code-mode-protocol/src/description.rs

@@ -17,6 +17,7 @@ const DEFERRED_NESTED_TOOLS_GUIDANCE: &str = r#"Some deferred nested tools may b
 To find one, filter `ALL_TOOLS` by `name` and `description`.
 
 Tool availability can change between calls."#;
+pub const TOOL_SEARCH_GUIDANCE: &str = r#"Use `await tools.tool_search({query: "...", limit: 8})` for BM25-ranked discovery of deferred tools (up to `limit` matches; 8 by default). It returns `{tools: [{name, description}]}` with callable names and full TypeScript declarations. Print results with `text()` and inspect unfamiliar declarations before calling `tools[result.name](arguments)` in a later execution."#;
 const LEGACY_IMAGE_HELPER_DESCRIPTION: &str = r#"`image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null } | ImageContent, detail?: "auto" | "low" | "high" | "original" | null)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`. MCP image blocks may request detail with `_meta: { "codex/imageDetail": "original" }`. When provided, the second `detail` argument overrides any detail embedded in the first argument."#;
 const UNIFIED_IMAGE_HELPER_DESCRIPTION: &str = r#"`image(imageUrlOrItem: string | { image_url: string } | ImageContent)`: Appends an image item. `image_url` should be a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`."#;
 const EXEC_DESCRIPTION_TEMPLATE: &str = r#"Run JavaScript code to orchestrate/compose tool calls
@@ -268,6 +269,13 @@ pub enum ImageDetailVisibility {
     Hidden,
 }
 
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum DeferredToolDiscovery {
+    Catalog,
+    RankedSearch,
+}
+
+#[allow(clippy::too_many_arguments)]
 pub fn build_exec_tool_description(
     enabled_tools: &[ToolDefinition],
     _deferred_tools: &[ToolDefinition],
@@ -275,6 +283,7 @@ pub fn build_exec_tool_description(
     default_exec_yield_time_ms: u64,
     code_mode_only: bool,
     image_detail_visibility: ImageDetailVisibility,
+    deferred_tool_discovery: DeferredToolDiscovery,
     messages: Option<&CodeModeToolMessages>,
 ) -> String {
     let mut sections = Vec::new();
@@ -294,9 +303,18 @@ pub fn build_exec_tool_description(
     if !description.is_empty() {
         sections.push(description);
     }
+    let default_guidance = match deferred_tool_discovery {
+        // In Code Mode Only, the enabled search tool's description is rendered below.
+        DeferredToolDiscovery::RankedSearch if !code_mode_only => {
+            format!("{DEFERRED_NESTED_TOOLS_GUIDANCE}\n\n{TOOL_SEARCH_GUIDANCE}")
+        }
+        DeferredToolDiscovery::Catalog | DeferredToolDiscovery::RankedSearch => {
+            DEFERRED_NESTED_TOOLS_GUIDANCE.to_string()
+        }
+    };
     let guidance = messages
         .and_then(|messages| messages.deferred_nested_tools_guidance.as_deref())
-        .unwrap_or(DEFERRED_NESTED_TOOLS_GUIDANCE);
+        .unwrap_or(&default_guidance);
     if !guidance.is_empty() {
         sections.push(guidance.to_string());
     }
@@ -529,6 +547,7 @@ mod description_override_tests;
 #[cfg(test)]
 mod tests {
     use super::CodeModeToolKind;
+    use super::DeferredToolDiscovery;
     use super::ImageDetailVisibility;
     use super::ParsedExecSource;
     use super::ToolDefinition;
@@ -737,6 +756,7 @@ mod tests {
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ true,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             /*messages*/ None,
         );
         assert!(description.contains(
@@ -755,6 +775,7 @@ bar"
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ false,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             /*messages*/ None,
         );
         assert!(description.contains("`audio(audioUrlOrItem:"));
@@ -813,6 +834,7 @@ bar"
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ true,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             /*messages*/ None,
         );
         assert_eq!(description.matches("## mcp__sample").count(), 1);
@@ -857,6 +879,7 @@ bar"
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ true,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             /*messages*/ None,
         );
 
@@ -963,6 +986,7 @@ bar"
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ true,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             /*messages*/ None,
         );
 
@@ -1002,6 +1026,7 @@ bar"
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ true,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             /*messages*/ None,
         );
 
@@ -1011,27 +1036,63 @@ bar"
     }
 
     #[test]
-    fn exec_description_mentions_deferred_nested_tools_when_available() {
-        let description = build_exec_tool_description(
-            &[],
-            &[ToolDefinition {
-                name: "deferred_tool".to_string(),
-                tool_name: ToolName::plain("deferred_tool"),
-                description: "Deferred tool".to_string(),
-                kind: CodeModeToolKind::Function,
-                input_schema: None,
-                input_schema_max_bytes: None,
-                output_schema: None,
-            }],
-            &BTreeMap::new(),
-            crate::DEFAULT_EXEC_YIELD_TIME_MS,
-            /*code_mode_only*/ false,
-            ImageDetailVisibility::Visible,
-            /*messages*/ None,
-        );
+    fn exec_description_preserves_discovery_guidance_across_catalog_changes() {
+        let deferred_tools = [ToolDefinition {
+            name: "deferred_tool".to_string(),
+            tool_name: ToolName::plain("deferred_tool"),
+            description: "Deferred tool".to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: None,
+            input_schema_max_bytes: None,
+            output_schema: None,
+        }];
+        let search_tool = ToolDefinition {
+            name: "tool_search".to_string(),
+            tool_name: ToolName::plain("tool_search"),
+            description: super::TOOL_SEARCH_GUIDANCE.to_string(),
+            kind: CodeModeToolKind::Function,
+            input_schema: None,
+            input_schema_max_bytes: None,
+            output_schema: None,
+        };
+        for code_mode_only in [false, true] {
+            for discovery in [
+                DeferredToolDiscovery::Catalog,
+                DeferredToolDiscovery::RankedSearch,
+            ] {
+                let enabled_tools = match discovery {
+                    DeferredToolDiscovery::Catalog => Vec::new(),
+                    DeferredToolDiscovery::RankedSearch => vec![search_tool.clone()],
+                };
+                let description = build_exec_tool_description(
+                    &enabled_tools,
+                    &deferred_tools,
+                    &BTreeMap::new(),
+                    crate::DEFAULT_EXEC_YIELD_TIME_MS,
+                    code_mode_only,
+                    ImageDetailVisibility::Visible,
+                    discovery,
+                    /*messages*/ None,
+                );
+                assert!(description.contains("Some deferred nested tools may be omitted"));
+                assert!(description.contains("filter `ALL_TOOLS` by `name` and `description`"));
+                assert_eq!(
+                    description.matches("await tools.tool_search(").count(),
+                    usize::from(discovery == DeferredToolDiscovery::RankedSearch),
+                );
 
-        assert!(description.contains("Some deferred nested tools may be omitted"));
-        assert!(description.contains("filter `ALL_TOOLS` by `name` and `description`"));
-        assert!(!description.contains("do not print the full `ALL_TOOLS` array"));
+                let empty_catalog_description = build_exec_tool_description(
+                    &enabled_tools,
+                    &[],
+                    &BTreeMap::new(),
+                    crate::DEFAULT_EXEC_YIELD_TIME_MS,
+                    code_mode_only,
+                    ImageDetailVisibility::Visible,
+                    discovery,
+                    /*messages*/ None,
+                );
+                assert_eq!(description, empty_catalog_description);
+            }
+        }
     }
 }

codex-rs/code-mode-protocol/src/description_override_tests.rs

@@ -2,9 +2,11 @@
 
 use super::CodeModeToolKind;
 use super::DEFERRED_NESTED_TOOLS_GUIDANCE;
+use super::DeferredToolDiscovery;
 use super::ImageDetailVisibility;
 use super::LEGACY_IMAGE_HELPER_DESCRIPTION;
 use super::MCP_TYPESCRIPT_PREAMBLE;
+use super::TOOL_SEARCH_GUIDANCE;
 use super::ToolDefinition;
 use super::UNIFIED_IMAGE_HELPER_DESCRIPTION;
 use super::build_exec_tool_description;
@@ -34,6 +36,7 @@ fn exec_override_renders_only_known_literal_placeholders() {
             /*default_exec_yield_time_ms*/ 4567,
             /*code_mode_only*/ false,
             image_detail_visibility,
+            DeferredToolDiscovery::Catalog,
             Some(&CodeModeToolMessages {
                 exec: Some(ToolMessage {
                     description: Some(" \nDefaults to 10000 ms. {{ default_exec_yield_time_ms }} ms.\n{{ image_helper }}\n{{ unknown }} {{default_exec_yield_time_ms}}\t ".to_string()),
@@ -63,6 +66,7 @@ fn exec_override_preserves_empty_and_whitespace_only_text() {
                 crate::DEFAULT_EXEC_YIELD_TIME_MS,
                 /*code_mode_only*/ true,
                 ImageDetailVisibility::Visible,
+                DeferredToolDiscovery::Catalog,
                 Some(&CodeModeToolMessages {
                     exec: Some(ToolMessage {
                         description: Some(description_override.to_string()),
@@ -153,6 +157,7 @@ fn exec_override_preserves_runtime_sections() {
                 crate::DEFAULT_EXEC_YIELD_TIME_MS,
                 code_mode_only,
                 ImageDetailVisibility::Visible,
+                DeferredToolDiscovery::Catalog,
                 Some(&CodeModeToolMessages {
                     exec: Some(ToolMessage {
                         description: Some(String::new()),
@@ -168,6 +173,53 @@ fn exec_override_preserves_runtime_sections() {
     }
 }
 
+#[test]
+fn ranked_search_guidance_honors_overrides() {
+    let deferred_tools = [ToolDefinition {
+        name: "deferred_tool".to_string(),
+        tool_name: ToolName::plain("deferred_tool"),
+        description: "Deferred tool".to_string(),
+        kind: CodeModeToolKind::Function,
+        input_schema: None,
+        input_schema_max_bytes: None,
+        output_schema: None,
+    }];
+    for code_mode_only in [false, true] {
+        for (guidance, expected) in [
+            (
+                None,
+                if code_mode_only {
+                    DEFERRED_NESTED_TOOLS_GUIDANCE.to_string()
+                } else {
+                    format!("{DEFERRED_NESTED_TOOLS_GUIDANCE}\n\n{TOOL_SEARCH_GUIDANCE}")
+                },
+            ),
+            (Some(""), String::new()),
+            (Some("Catalog discovery."), "Catalog discovery.".to_string()),
+        ] {
+            let description = build_exec_tool_description(
+                &[],
+                &deferred_tools,
+                &BTreeMap::new(),
+                crate::DEFAULT_EXEC_YIELD_TIME_MS,
+                code_mode_only,
+                ImageDetailVisibility::Visible,
+                DeferredToolDiscovery::RankedSearch,
+                Some(&CodeModeToolMessages {
+                    exec: Some(ToolMessage {
+                        description: Some(String::new()),
+                        ..Default::default()
+                    }),
+                    deferred_nested_tools_guidance: guidance.map(str::to_string),
+                    mcp_typescript_preamble: Some(String::new()),
+                    ..Default::default()
+                }),
+            );
+            assert_eq!(description, expected);
+        }
+    }
+}
+
 #[test]
 fn mcp_types_stay_stable_when_a_deferred_tool_changes_its_output_schema() {
     let mut tool = ToolDefinition {
@@ -187,6 +239,7 @@ fn mcp_types_stay_stable_when_a_deferred_tool_changes_its_output_schema() {
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             /*code_mode_only*/ true,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             Some(&CodeModeToolMessages {
                 mcp_typescript_preamble: preamble.map(str::to_string),
                 ..Default::default()
@@ -229,6 +282,7 @@ fn nested_guidance_survives_catalog_changes_and_respects_overrides() {
             crate::DEFAULT_EXEC_YIELD_TIME_MS,
             code_mode_only,
             ImageDetailVisibility::Visible,
+            DeferredToolDiscovery::Catalog,
             Some(&CodeModeToolMessages {
                 exec: Some(ToolMessage {
                     description: Some(String::new()),

codex-rs/code-mode-protocol/src/lib.rs

@@ -8,8 +8,10 @@ mod session;
 
 pub use description::CODE_MODE_PRAGMA_PREFIX;
 pub use description::CodeModeToolKind;
+pub use description::DeferredToolDiscovery;
 pub use description::EnabledToolMetadata;
 pub use description::ImageDetailVisibility;
+pub use description::TOOL_SEARCH_GUIDANCE;
 pub use description::ToolDefinition;
 pub use description::ToolNamespaceDescription;
 pub use description::augment_tool_definition;