diff --git a/codex-rs/code-mode-protocol/src/description.rs b/codex-rs/code-mode-protocol/src/description.rs index 87c918055f..3e0d1affde 100644 --- a/codex-rs/code-mode-protocol/src/description.rs +++ b/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); + } + } } } diff --git a/codex-rs/code-mode-protocol/src/description_override_tests.rs b/codex-rs/code-mode-protocol/src/description_override_tests.rs index 57c8b9f0cb..294aa1011e 100644 --- a/codex-rs/code-mode-protocol/src/description_override_tests.rs +++ b/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()), diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs index aab810d954..608b871b34 100644 --- a/codex-rs/code-mode-protocol/src/lib.rs +++ b/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; diff --git a/codex-rs/core/config.schema.json b/codex-rs/core/config.schema.json index 7322488bad..8931d9c789 100644 --- a/codex-rs/core/config.schema.json +++ b/codex-rs/core/config.schema.json @@ -815,6 +815,9 @@ "code_mode_prewarm": { "type": "boolean" }, + "code_mode_tool_search": { + "type": "boolean" + }, "codex_apps_mcp_2026_07_28": { "type": "boolean" }, @@ -6988,6 +6991,9 @@ "code_mode_prewarm": { "type": "boolean" }, + "code_mode_tool_search": { + "type": "boolean" + }, "codex_apps_mcp_2026_07_28": { "type": "boolean" }, diff --git a/codex-rs/core/src/tools/code_mode/execute_handler.rs b/codex-rs/core/src/tools/code_mode/execute_handler.rs index 2215209053..75733b9dfa 100644 --- a/codex-rs/core/src/tools/code_mode/execute_handler.rs +++ b/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); diff --git a/codex-rs/core/src/tools/code_mode/execute_spec.rs b/codex-rs/core/src/tools/code_mode/execute_spec.rs index 8068614d78..784341246b 100644 --- a/codex-rs/core/src/tools/code_mode/execute_spec.rs +++ b/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, - 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(), diff --git a/codex-rs/core/src/tools/code_mode/mod.rs b/codex-rs/core/src/tools/code_mode/mod.rs index 79a95b2ac2..2b241ffc7d 100644 --- a/codex-rs/core/src/tools/code_mode/mod.rs +++ b/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; diff --git a/codex-rs/core/src/tools/code_mode/tool_definitions.rs b/codex-rs/core/src/tools/code_mode/tool_definitions.rs new file mode 100644 index 0000000000..84b262319a --- /dev/null +++ b/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, +) -> 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) +} diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 67b782927e..570b6610a2 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/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, source_listing: ToolSearchSourceListing, spec: ToolSpec, - search_engine: SearchEngine, + search_embedder: Embedder, + search_scorer: Scorer, } #[derive(Default)] @@ -150,20 +163,23 @@ impl ToolSearchHandler { TOOL_SEARCH_DEFAULT_LIMIT, source_listing, ); - let documents: Vec> = 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::::with_documents(Language::English, documents).build(); + .map(|search_info| search_info.entry.search_text.as_str()) + .collect::>(); + 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, 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, 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, + indirect_prefixes: &IndirectNamespacePrefixes<'_>, + ) -> Vec { + 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, - ) -> Result, FunctionCallError> { - Ok(coalesce_loadable_tool_specs( - results.into_iter().map(ToolSearchEntry::to_loadable_spec), - )) + fn search(&self, query: &str, limit: usize) -> Vec { + coalesce_loadable_tool_specs( + self.search_entries(query) + .take(limit) + .map(ToolSearchEntry::to_loadable_spec), + ) + } + + fn search_entries(&self, query: &str) -> impl Iterator + '_ { + 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, diff --git a/codex-rs/core/src/tools/registry.rs b/codex-rs/core/src/tools/registry.rs index fb2fad72bd..e0acae262e 100644 --- a/codex-rs/core/src/tools/registry.rs +++ b/codex-rs/core/src/tools/registry.rs @@ -64,7 +64,8 @@ pub(crate) trait CoreToolRuntime: ToolExecutor { 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, diff --git a/codex-rs/core/src/tools/router.rs b/codex-rs/core/src/tools/router.rs index fbf0e773be..baf39f082a 100644 --- a/codex-rs/core/src/tools/router.rs +++ b/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 { &self.code_mode_tool_names } diff --git a/codex-rs/core/src/tools/spec_plan.rs b/codex-rs/core/src/tools/spec_plan.rs index e5ab1500ee..7e2213456d 100644 --- a/codex-rs/core/src/tools/spec_plan.rs +++ b/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( diff --git a/codex-rs/core/src/tools/spec_plan_tests.rs b/codex-rs/core/src/tools/spec_plan_tests.rs index 952d1a7866..395cc95e5d 100644 --- a/codex-rs/core/src/tools/spec_plan_tests.rs +++ b/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)] { diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index fed86b0c66..95c7bdaaad 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -3611,46 +3611,109 @@ async fn code_mode_complete_call_survives_unrelated_truncation() -> Result<()> { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_only_guides_all_tools_search_and_calls_deferred_app_tools() -> Result<()> { +#[test_case(false, None; "control")] +#[test_case(true, None; "ranked_search")] +#[test_case(true, Some(30_000); "ranked_search_expanded_schema")] +async fn code_mode_only_searches_and_calls_deferred_app_tools( + ranked_search: bool, + schema_max_bytes: Option, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let apps_server = AppsTestServer::mount_searchable(&server).await?; + // Repeated refs keep the MCP schema small while its declaration exceeds the default budget. + let properties = (0..8) + .map(|index| { + ( + format!("option{index}"), + serde_json::json!({"$ref": "#/$defs/option"}), + ) + }) + .chain(std::iter::once(( + "timezone".to_string(), + serde_json::json!({"type": "string"}), + ))) + .collect::>(); + let input_schema = serde_json::json!({ + "type": "object", + "properties": properties, + "$defs": {"option": {"type": "object", "properties": {"term": {"type": "string", "description": format!("{}schema_budget_marker", "reference guidance ".repeat(120))}}}}, + }); + let tools = (0..100).map(|index| serde_json::json!({ + "name": format!("calendar_timezone_option_{index}"), + "description": format!("Read timezone option {index}."), + "annotations": {"readOnlyHint": true}, + "inputSchema": if index == 99 { input_schema.clone() } else { serde_json::json!({"type": "object"}) }, + "_meta": {"connector_id": "calendar", "connector_name": "Calendar", "connector_description": "Plan events and manage your calendar."}, + })).collect(); + let apps_server = + AppsTestServer::mount_with_tools(&server, Arc::new(Mutex::new(tools))).await?; let originating_item_id = "ctc_code_mode_origin"; - let mut exec_call = ev_custom_tool_call( - "call-1", - "exec", + let lookup = if ranked_search { + r#" +const { tools: matches } = await tools.tool_search({ + query: "calendar_timezone_option_99", limit: 1, +}); +const tool = matches[0]; +"# + } else { r#" const tool = ALL_TOOLS.find( ({ name }) => name === "mcp__codex_apps__calendar_timezone_option_99" ); +"# + }; + let call = r#" if (!tool) { text(JSON.stringify({ found: false })); } else { const result = await tools[tool.name]({ timezone: "UTC" }); text(JSON.stringify({ found: true, + rankedSearchAvailable: typeof tools.tool_search === "function", + name: tool.name, + hasDeclaration: tool.description.includes("mcp__codex_apps__calendar_timezone_option_99(args:"), + hasSchemaDescription: tool.description.includes("schema_budget_marker"), + hasNamespacePrefix: tool.description.startsWith("Calendar search context.\n\n"), + matchesCatalog: tool.description === ALL_TOOLS.find(({ name }) => name === tool.name)?.description, isError: Boolean(result.isError), text: result.content?.[0]?.text ?? "", })); } -"#, +"#; + let mut exec_call = ev_custom_tool_call( + "call-1", + "exec", + &format!("const tool = load(\"deferredAppTool\");\n{call}"), ); exec_call["item"]["id"] = serde_json::json!(originating_item_id); let resp_mock = responses::mount_sse_once( &server, sse(vec![ ev_response_created("resp-1"), - exec_call, + ev_custom_tool_call( + "search-1", + "exec", + &format!("{lookup}\nstore(\"deferredAppTool\", tool ?? null);"), + ), ev_completed("resp-1"), ]), ) .await; + let call_mock = responses::mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-2"), + exec_call, + ev_completed("resp-2"), + ]), + ) + .await; let follow_up_mock = responses::mount_sse_once( &server, sse(vec![ ev_assistant_message("msg-1", "done"), - ev_completed("resp-2"), + ev_completed("resp-3"), ]), ) .await; @@ -3671,6 +3734,13 @@ if (!tool) { .features .enable(Feature::CodeModeOnly) .expect("test config should allow feature update"); + config.code_mode.tool_input_schema_max_bytes = schema_max_bytes; + if ranked_search { + config + .features + .enable(Feature::CodeModeToolSearch) + .expect("ranked Code Mode search should be enabled"); + } let mut model_catalog = bundled_models_response().expect("bundled models.json should parse"); let model = model_catalog @@ -3681,9 +3751,21 @@ if (!tool) { config.chatgpt_base_url = apps_base_url; config.model = Some("gpt-5.5".to_string()); model.supports_search_tool = true; + model + .model_messages + .as_mut() + .expect("model messages") + .tools + .get_or_insert_with(Default::default) + .indirect_description_prefixes = Some( + serde_json::from_value(serde_json::json!({ + "mcp_servers": {"codex_apps": " Calendar search context. "}, + })) + .expect("indirect namespace prefixes"), + ); config.model_catalog = Some(model_catalog); }); - let test = builder.build(&server).await?; + let test = builder.build_with_auto_env(&server).await?; test.submit_turn("inspect tools in code mode only").await?; let first_body = resp_mock.single_request().body_json(); @@ -3716,6 +3798,8 @@ if (!tool) { }) .expect("exec description should be present"); assert!(exec_description.contains("filter `ALL_TOOLS` by `name` and `description`")); + let has_ranked_search = exec_description.contains("await tools.tool_search("); + assert_eq!(has_ranked_search, ranked_search); assert!(exec_description.contains("Shared MCP Types:")); assert!(!exec_description.contains("calendar_timezone_option_99")); @@ -3726,15 +3810,26 @@ if (!tool) { Some(false), "code_mode_only deferred app tool call failed unexpectedly: {output}" ); - let parsed: Value = serde_json::from_str(&output)?; + let parsed: Value = serde_json::from_str(&output).unwrap_or_else(|error| { + panic!("deferred app call-1 output was not JSON: {error}\n{output}") + }); assert_eq!( parsed, serde_json::json!({ "found": true, + "rankedSearchAvailable": ranked_search, + "name": "mcp__codex_apps__calendar_timezone_option_99", + "hasDeclaration": true, + "hasSchemaDescription": schema_max_bytes.is_some(), + "hasNamespacePrefix": true, + "matchesCatalog": true, "isError": false, "text": "called calendar_timezone_option_99 for at with ", }) ); + assert!(request.inputs_of_type("tool_search_output").is_empty()); + assert!(request.inputs_of_type("tool_search_call").is_empty()); + let call_body = call_mock.single_request().body_json(); let apps_tool_calls = recorded_apps_tool_calls(&server).await; assert!( apps_tool_calls.iter().any(|call| { @@ -3742,7 +3837,7 @@ if (!tool) { && call.pointer("/params/_meta/sessionId") == Some(&serde_json::json!(test.session_configured.session_id)) && call.pointer("/params/_meta/windowId") - == Some(&first_body["client_metadata"]["x-codex-window-id"]) + == Some(&call_body["client_metadata"]["x-codex-window-id"]) }), "the nested MCP call should inherit its code cell's originating Responses item" ); @@ -3761,6 +3856,7 @@ async fn app_only_tools_are_not_visible_or_runnable_by_code_mode_model() -> Resu r#" const visibleTool = ALL_TOOLS.find(({{ name }}) => name === {visible_tool_name:?}); const tool = ALL_TOOLS.find(({{ name }}) => name === {tool_name:?}); +const search = await tools.tool_search({{ query: {tool_name:?} }}); let error = null; try {{ await tools[{tool_name:?}]({{}}); @@ -3770,6 +3866,7 @@ try {{ text(JSON.stringify({{ visibleListed: visibleTool !== undefined, listed: tool !== undefined, + searchListed: search.tools.some(({{ name }}) => name === {tool_name:?}), callable: typeof tools[{tool_name:?}] === "function", error, }})); @@ -3798,6 +3895,10 @@ text(JSON.stringify({{ let mut builder = search_capable_apps_builder(apps_server.chatgpt_base_url.clone()) .with_config(|config| { + config + .features + .enable(Feature::CodeModeToolSearch) + .expect("ranked Code Mode search should be enabled"); config .features .enable(Feature::CodeMode) @@ -3821,6 +3922,7 @@ text(JSON.stringify({{ let parsed: Value = serde_json::from_str(&output)?; assert_eq!(parsed["visibleListed"], true); assert_eq!(parsed["listed"], false); + assert_eq!(parsed["searchListed"], false); assert_eq!(parsed["callable"], false); assert!( parsed["error"] @@ -8250,36 +8352,57 @@ async fn code_mode_renders_local_refs_in_outbound_exec_description() -> Result<( } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_can_call_hidden_dynamic_tools() -> Result<()> { +async fn code_mode_search_normalizes_and_calls_deferred_dynamic_tools() -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; - let mut builder = test_codex().with_config(move |config| { - let _ = config.features.enable(Feature::CodeMode); - }); - let base_test = builder.build(&server).await?; + let mut builder = test_codex() + .with_model_info_override("gpt-5.5", |model_info| { + model_info.supports_search_tool = true; + model_info.tool_mode = Some(ToolMode::CodeMode); + }) + .with_config(|config| { + let _ = config.features.enable(Feature::CodeMode); + let _ = config.features.enable(Feature::CodeModeToolSearch); + config.code_mode.excluded_tool_namespaces = vec!["blocked".to_string()]; + }); + let base_test = builder.build_with_auto_env(&server).await?; let new_thread = base_test .thread_manager .start_thread(StartThreadOptions { - dynamic_tools: vec![DynamicToolSpec::Namespace(DynamicToolNamespaceSpec { - name: "codex_app".to_string(), - description: "Codex app tools.".to_string(), - tools: vec![DynamicToolNamespaceTool::Function( - DynamicToolFunctionSpec { - name: "hidden_dynamic_tool".to_string(), - description: "A hidden dynamic tool.".to_string(), - input_schema: serde_json::json!({ - "type": "object", - "properties": { - "city": { "type": "string" } - }, - "required": ["city"], - "additionalProperties": false, - }), - defer_loading: true, - }, - )], - })], + dynamic_tools: vec![ + DynamicToolSpec::Namespace(DynamicToolNamespaceSpec { + name: "blocked".to_string(), + description: "Blocked hidden dynamic tools.".to_string(), + tools: vec![DynamicToolNamespaceTool::Function( + DynamicToolFunctionSpec { + name: "hidden-dynamic-tool".to_string(), + description: "A blocked hidden dynamic tool.".to_string(), + input_schema: serde_json::json!({ "type": "object" }), + defer_loading: true, + }, + )], + }), + DynamicToolSpec::Namespace(DynamicToolNamespaceSpec { + name: "codex_app".to_string(), + description: "Codex app tools.".to_string(), + tools: vec![DynamicToolNamespaceTool::Function( + DynamicToolFunctionSpec { + name: "hidden-dynamic-tool".to_string(), + description: "A hidden dynamic tool.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "city": { "type": "string" } + }, + "required": ["city"], + "additionalProperties": false, + }), + defer_loading: true, + }, + )], + }), + ], ..StartThreadOptions::new(base_test.config.clone()) }) .await?; @@ -8288,13 +8411,21 @@ async fn code_mode_can_call_hidden_dynamic_tools() -> Result<()> { test.session_configured = new_thread.session_configured; let code = r#" -const tool = ALL_TOOLS.find(({ name }) => name === "codex_app__hidden_dynamic_tool"); -const out = await tools.codex_app__hidden_dynamic_tool({ city: "Paris" }); +const errors = []; +for (const args of [ + { query: " " }, { query: "hidden", limit: 0 }, +]) { + try { await tools.tool_search(args); } catch (error) { errors.push(String(error)); } +} +const { tools: matches } = await tools.tool_search({ query: "blocked hidden-dynamic-tool", limit: 1 }); +const tool = matches[0]; +const out = await tools[tool.name]({ city: "Paris" }); text( JSON.stringify({ name: tool?.name ?? null, description: tool?.description ?? null, out, + recoveredFromInvalidArguments: errors.length === 2, }) ); "#; @@ -8360,7 +8491,7 @@ text( }) .await; assert_eq!(request.namespace.as_deref(), Some("codex_app")); - assert_eq!(request.tool, "hidden_dynamic_tool"); + assert_eq!(request.tool, "hidden-dynamic-tool"); assert_eq!(request.arguments, serde_json::json!({ "city": "Paris" })); test.codex .submit(Op::DynamicToolResponse { @@ -8399,6 +8530,7 @@ text( parsed.get("out"), Some(&Value::String("hidden-ok".to_string())) ); + assert_eq!(parsed["recoveredFromInvalidArguments"], true); assert!( parsed .get("description") diff --git a/codex-rs/core/tests/suite/scenarios.rs b/codex-rs/core/tests/suite/scenarios.rs index 765b75f83f..aff8dbb68b 100644 --- a/codex-rs/core/tests/suite/scenarios.rs +++ b/codex-rs/core/tests/suite/scenarios.rs @@ -1415,8 +1415,10 @@ async fn astra_continues_after_input_yields_a_code_mode_cell() -> Result<()> { Ok(()) } +#[test_case::test_case(false; "catalog")] +#[test_case::test_case(true; "ranked_search")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn code_mode_catalog_messages() -> Result<()> { +async fn code_mode_catalog_messages(ranked_search: bool) -> Result<()> { skip_if_no_network!(Ok(())); let server = start_mock_server().await; let apps = @@ -1428,6 +1430,8 @@ async fn code_mode_catalog_messages() -> Result<()> { configure_scenario_catalog(config); config.chatgpt_base_url = apps.chatgpt_base_url; config.features.enable(Feature::Apps).expect("enable apps"); + config.features.set_enabled(Feature::CodeModeToolSearch, ranked_search) + .expect("set ranked tool search"); config.workspace_roots = vec![config.cwd.clone()]; config.code_mode.disable_in_process_fallback = true; config.code_mode.default_exec_yield_time_ms = 1234; @@ -1464,44 +1468,86 @@ async fn code_mode_catalog_messages() -> Result<()> { }) .build_with_auto_env(&server) .await?; - let mock = mount_sse_sequence( - &server, - vec![ + let exec_code = if ranked_search { + r#"const name = load("calendarToolName"); +text((await tools[name]({ timezone: "UTC" })).content[0].text); +text('started'); yield_control(); text('finished');"# + } else { + "text('started'); yield_control(); text('finished');" + }; + let mut events = vec![ + sse(vec![ + ev_response_created("exec-response"), + ev_custom_tool_call("exec-call", "exec", exec_code), + ev_completed("exec-response"), + ]), + sse(vec![ + ev_response_created("wait-response"), + ev_function_call_with_namespace( + "wait-call", + "functions", + "wait", + if ranked_search { + r#"{"cell_id":"2"}"# + } else { + r#"{"cell_id":"1"}"# + }, + ), + ev_completed("wait-response"), + ]), + sse(vec![ + ev_assistant_message("final", "The cell started, yielded, and finished."), + ev_completed("final-response"), + ]), + ]; + if ranked_search { + events.insert( + 0, sse(vec![ - ev_response_created("exec-response"), + ev_response_created("search-response"), ev_custom_tool_call( - "exec-call", + "search-call", "exec", - "text('started'); yield_control(); text('finished');", + r#"const result = await tools.tool_search({query: "calendar_timezone_option_99", limit: 1}); +store("calendarToolName", result.tools[0].name); +text(result.tools[0].name); +text(result.tools[0].description);"#, ), - ev_completed("exec-response"), + ev_completed("search-response"), ]), - sse(vec![ - ev_response_created("wait-response"), - ev_function_call_with_namespace( - "wait-call", - "functions", - "wait", - r#"{"cell_id":"1"}"#, - ), - ev_completed("wait-response"), - ]), - sse(vec![ - ev_assistant_message("final", "The cell started, yielded, and finished."), - ev_completed("final-response"), - ]), - ], - ) - .await; - test.submit_turn("Start a code cell, yield its initial output, then wait for its completion.") - .await?; + ); + } + let mock = mount_sse_sequence(&server, events).await; + let prompt = if ranked_search { + "Find calendar timezone option 99, call it with UTC, then yield and wait for completion." + } else { + "Start a code cell, yield its initial output, then wait for its completion." + }; + test.submit_turn(prompt).await?; let requests = mock.requests(); + assert_eq!(requests.len(), if ranked_search { 4 } else { 3 }); + if ranked_search { + assert!(requests[0].body_contains_text("await tools.tool_search(")); + assert!(requests[1].body_contains_text("calendar_timezone_option_99(args:")); + assert!(requests[2].body_contains_text("called calendar_timezone_option_99")); + } assert!(requests[0].body_contains_text("Catalog discovery: find nested tools")); assert!(requests[0].body_contains_text("type CallToolResult")); - insta::assert_snapshot!( - "code_mode_catalog_messages", - context_snapshot::format_request_history_snapshot( + let (snapshot_name, scenario) = if ranked_search { + ( + "code_mode_catalog_messages_ranked_search", + "Ranked search prints a deferred tool declaration before a later code cell invokes it, yields, and completes through wait.", + ) + } else { + ( + "code_mode_catalog_messages", "Catalog Code Mode instructions and wait parameter descriptions accompany a yielded cell through completion.", + ) + }; + insta::assert_snapshot!( + snapshot_name, + context_snapshot::format_request_history_snapshot( + scenario, &requests, &ContextSnapshotOptions::default().include_request_settings(), ) diff --git a/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_catalog_messages_ranked_search.snap b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_catalog_messages_ranked_search.snap new file mode 100644 index 0000000000..db8001d90b --- /dev/null +++ b/codex-rs/core/tests/suite/snapshots/all__suite__scenarios__code_mode_catalog_messages_ranked_search.snap @@ -0,0 +1,132 @@ +--- +source: core/tests/suite/scenarios.rs +expression: "context_snapshot::format_request_history_snapshot(scenario, &requests,\n&ContextSnapshotOptions::default().include_request_settings(),)" +--- +Scenario: Ranked search prints a deferred tool declaration before a later code cell invokes it, yields, and completes through wait. + +## Window 1 +Settings: + include: ["reasoning.encrypted_content"] + model: "gpt-6-astra" + parallel_tool_calls: true + prompt_cache_key: "" + reasoning: {"effort":"low"} + store: false + stream: true + text: {"verbosity":"low"} + tool_choice: "auto" + tools (7; hash=b6a71dafc21b6db7): + - custom/exec: Catalog exec instructions. Default yield: 1234 ms. `image(imageUrlOrItem: string | { image_url: s...; hash=1b186fc5b7e14af3 + - function/wait: Catalog wait instructions: resume the yielded cell.; args=[cell_id, max_tokens, terminate, yield_time_ms]; hash=67bc0f04da3505b3 + - function/request_user_input: Request user input for one to three short questions and wait for the response. This tool is only ...; args=[questions]; hash=64a729b25c424fc6 + - function/request_user_input_async: Ask the user one or more questions during ongoing work. Use this tool only to request missing inf...; args=[questions]; hash=d6d74f33db330162 + - namespace/clock: Tools for reading and waiting on time.; hash=f402e8c5e9b5e317 + - function/sleep + - namespace/collaboration: Tools for spawning and managing sub-agents.; hash=b263e3c84ad07c37 + - function/followup_task + - function/interrupt_agent + - function/list_agents + - function/send_message + - function/spawn_agent + - function/wait_agent + - web_search; hash=e214d4db1bc49e89 +-- request 1 (turn) -- +00:message/developer: + You are Codex, an agent based on GPT-6. You and the user share one workspace, and your job is to collabo...them until their intended goal is completely handled. [hash=efa50f074efa5803] + + # When to ask the user for permission + + Use your best judgement given task context for when you really need user permission, like a competent co...ork without ending the turn to clarify with the user. [hash=ca7367e658b51388] + + User authorization and preferences persist across turns. Do not request permission again when the user h... any guidelines provided in skills or external files. [hash=911f91c92c72882c] + + + + - Skill naming: If a plugin contributes skills, those skill entries are prefixed with plugin_name: in the Skills list. + - MCP naming: Plugin-provided MCP tools keep standard MCP identifiers such as mcp__server__tool; use tool provenance to tell which plugin they come from. + - Trigger rules: If the user explicitly names a plugin, prefer capabilities associated with that plugin for that turn. + - Relationship to capabilities: Plugins are not invoked directly. Use their underlying skills, MCP tools, and app tools to help solve the task. + - Relevance: Determine what a plugin can help with from explicit user mention or from the plugin-associa..., MCP tools, and apps exposed elsewhere in this turn. [hash=620459a4e7b590f0] + - Missing/blocked: If the user requests a plugin that does not have relevant callable capabilities for t..., say so briefly and continue with the best fallback. [hash=746b0469a8f6cb3e] + +01:message/developer[2]: + [01] + Filesystem sandboxing defines which files can be read or written. `sandbox_mode` is `danger-full-access`...ll commands are permitted. Network access is enabled. [hash=c134061cebb94a1b] + Approval policy is currently never. Do not provide the `sandbox_permissions` for any reason, commands will be rejected. + + [02] # Collaboration Mode: Default + + You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active. + + Your active mode changes only when new developer instructions with a different `...<...by themselves. Known mode names are Default and Plan. [hash=e75f6b12fc52d812] + + ## request_user_input availability + + Use the `request_user_input` tool only when it is listed in the available tools for this turn. + + Use the `request_user_input` tool only for optional questions where the answer would materially improve the quality of the work. + + If `request_user_input` returns no answers, continue with best judgment instead of asking again or treating the turn as blocked. + + Never use the `request_user_input` tool for permission requests or permission-related escalations. + +02:message/developer: + You are `/root`, the primary agent in a team of agents collaborating to fulfill the user's goals. + + At the start of your turn, you are the active agent. + You can spawn sub-agents to handle subtasks, and those sub-agents can spawn their own sub-agents. + All agents in the team, including the agents that you can assign tasks to, are equally intelligent and capable, and have access to the same set of tools. + + You can use `spawn_agent` to create a new agent, `followup_task` to give an existing agent a new task an...message to a running agent without triggering a turn. [hash=7d40ff8b58977202] + `send_message` calls may be read by a human, so ensure they are legible. Always put proper spaces between words and/or numbers. + + - All agents use the same current working directory. + - As a result, edits made by one agent are immediately visible to all other agents. + + When calling `wait_agent`, prefer longer waits (minutes) to avoid busy polling. + + There are 4 available concurrency slots, meaning that up to 4 agents can be active at once, including you. + + Full-history forks (`fork_turns` omitted or `"all"`) inherit the parent model and reasoning effort and d...ne"` or a positive integer string. [hash=c7f52e90688299ce] +03:message/developer: + Any earlier instruction enabling proactive multi-agent delegation no longer applies. D...elegation, or parallel agent work. [hash=2b6e037067b4390c] +04:message/user: + + + + + + <...e="unrestricted" /> [hash=d40b2f45cb028888] + +05:message/user: + Find calendar timezone option 99, call it with UTC, then yield and wait for completion. +-- request 2 (turn) -- +06:custom_tool_call/exec:const result = await tools.tool_search({query: "calendar_timezone_option_99", limit: 1}); + store("calendarToolName", result.tools[0].name); + text(result.tools[0].name); + text(result.tools[0].description); +07:custom_tool_call_output:Script completed + Wall time seconds + Output: + | mcp__codex_apps__calendar_timezone_option_99 | Plan events and manage your calendar. + + Read timezone option 99. + + exec tool declaration: + ```ts + declare const tools: { mcp__codex_apps__calendar_timezone_option_99(args: { timezone?: string; }): Promise; }; + ``` +-- request 3 (turn) -- +08:custom_tool_call/exec:const name = load("calendarToolName"); + text((await tools[name]({ timezone: "UTC" })).content[0].text); + text('started'); yield_control(); text('finished'); +09:custom_tool_call_output:Script running with cell ID 2 + Wall time seconds + Output: + | called calendar_timezone_option_99 for at with | started +-- request 4 (turn) -- +10:function_call/functions.wait:{"cell_id":"2"} +11:function_call_output:Script completed + Wall time seconds + Output: + | finished diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index d1713bd384..5d3804662f 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -127,6 +127,8 @@ pub enum Feature { ExecutedToolCallMetadata, /// Enable JavaScript code mode backed by the standalone host process. CodeMode, + /// Expose ranked tool discovery inside JavaScript code mode. + CodeModeToolSearch, /// Removed compatibility flag for the configurable code-mode exec yield timeout. CodeModeBufferedExec, /// Run JavaScript code mode in the standalone host process. @@ -1112,6 +1114,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::UnderDevelopment, default_enabled: false, }, + FeatureSpec { + id: Feature::CodeModeToolSearch, + key: "code_mode_tool_search", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::CodeModeBufferedExec, key: "code_mode_buffered_exec", diff --git a/codex-rs/tools/src/code_mode.rs b/codex-rs/tools/src/code_mode.rs index 7deb8c810a..8194387921 100644 --- a/codex-rs/tools/src/code_mode.rs +++ b/codex-rs/tools/src/code_mode.rs @@ -223,7 +223,38 @@ fn code_mode_tool_definitions_for_spec( } }) .collect(), - ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => Vec::new(), + ToolSpec::ToolSearch { parameters, .. } => { + let name = spec.name().to_string(); + vec![CodeModeToolDefinition { + tool_name: ToolName::plain(name.clone()), + name, + description: codex_code_mode::TOOL_SEARCH_GUIDANCE.to_string(), + kind: CodeModeToolKind::Function, + input_schema: serde_json::to_value(parameters).ok(), + input_schema_max_bytes: Some(effective_input_schema_max_bytes( + parameters.mcp_input_schema_max_bytes, + code_mode_input_schema_max_bytes, + )), + output_schema: Some(serde_json::json!({ + "type": "object", + "properties": { + "tools": { + "type": "array", + "items": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "description": { "type": "string" } + }, + "required": ["name", "description"] + } + } + }, + "required": ["tools"] + })), + }] + } + ToolSpec::WebSearch { .. } => Vec::new(), } } diff --git a/codex-rs/tools/src/code_mode_tests.rs b/codex-rs/tools/src/code_mode_tests.rs index 135694bc95..0486372a5c 100644 --- a/codex-rs/tools/src/code_mode_tests.rs +++ b/codex-rs/tools/src/code_mode_tests.rs @@ -238,17 +238,50 @@ declare const tools: { editor__apply_patch(input: string): Promise; }; } #[test] -fn tool_spec_to_code_mode_tool_definition_skips_unsupported_variants() { +fn tool_search_spec_converts_to_a_callable_code_mode_definition() { + let definition = tool_spec_to_code_mode_tool_definition(&ToolSpec::ToolSearch { + execution: "client".to_string(), + description: "Native tool discovery".to_string(), + parameters: JsonSchema::object( + BTreeMap::from([( + "query".to_string(), + JsonSchema::string(/*description*/ None), + )]), + Some(vec!["query".to_string()]), + Some(AdditionalProperties::Boolean(false)), + ), + }) + .expect("tool search should be callable from Code Mode"); assert_eq!( - tool_spec_to_code_mode_tool_definition(&ToolSpec::ToolSearch { - execution: "sync".to_string(), - description: "Search".to_string(), - parameters: JsonSchema::object( - BTreeMap::new(), - /*required*/ None, - /*additional_properties*/ None - ), - }), - None + ( + &definition.name, + &definition.tool_name, + &definition.input_schema, + definition.input_schema_max_bytes, + ), + ( + &"tool_search".to_string(), + &ToolName::plain("tool_search"), + &Some(json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + }, + "required": ["query"], + "additionalProperties": false, + })), + Some(codex_code_mode::DEFAULT_INPUT_SCHEMA_MAX_BYTES), + ), ); + assert!( + definition + .description + .contains("tool_search(args: { query: string; })") + ); + assert!( + definition + .description + .contains("tools[result.name](arguments)") + ); + assert!(!definition.description.contains("Native tool discovery")); } diff --git a/codex-rs/tools/src/indirect_namespace_prefixes.rs b/codex-rs/tools/src/indirect_namespace_prefixes.rs index 4261520627..48d4e43d6c 100644 --- a/codex-rs/tools/src/indirect_namespace_prefixes.rs +++ b/codex-rs/tools/src/indirect_namespace_prefixes.rs @@ -83,9 +83,14 @@ impl<'a> IndirectNamespacePrefixes<'a> { return; } for tool in tools { - if let Some(prefix) = self.prefix(namespace_name(&tool.tool_name)) { - prepend(prefix, &mut tool.description); - } + self.apply_code_mode_description(&tool.tool_name, &mut tool.description); + } + } + + /// Adds guidance to one owned description without copying the tool's schemas. + pub fn apply_code_mode_description(&self, tool_name: &ToolName, description: &mut String) { + if let Some(prefix) = self.prefix(namespace_name(tool_name)) { + prepend(prefix, description); } } diff --git a/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs b/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs index fed2443ec3..6f2df660a1 100644 --- a/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs +++ b/codex-rs/tools/src/indirect_namespace_prefixes_tests.rs @@ -38,6 +38,7 @@ fn plain_tools_join_an_existing_functions_group_only_with_a_prefix() { codex_code_mode::DEFAULT_EXEC_YIELD_TIME_MS, /*code_mode_only*/ true, codex_code_mode::ImageDetailVisibility::Visible, + codex_code_mode::DeferredToolDiscovery::Catalog, /*messages*/ None, ) }; diff --git a/codex-rs/tools/src/tool_search.rs b/codex-rs/tools/src/tool_search.rs index 8cd45d9005..829910ea8f 100644 --- a/codex-rs/tools/src/tool_search.rs +++ b/codex-rs/tools/src/tool_search.rs @@ -3,6 +3,7 @@ use crate::LoadableToolSpec; use crate::ResponsesApiNamespace; use crate::ResponsesApiNamespaceTool; use crate::ResponsesApiTool; +use crate::ToolName; use crate::ToolSearchSourceInfo; use crate::ToolSpec; use crate::default_namespace_description; @@ -16,6 +17,27 @@ pub struct ToolSearchEntry { } impl ToolSearchEntry { + /// Iterate candidate identities without materializing their tool schemas. + pub fn tool_names(&self) -> impl Iterator + '_ { + let (namespace, name, tools): (&str, Option<&str>, &[ResponsesApiNamespaceTool]) = + match self.spec.as_ref() { + ToolSpec::Function(tool) => (DEFAULT_FUNCTION_NAMESPACE, Some(&tool.name), &[]), + ToolSpec::Freeform(tool) => (DEFAULT_FUNCTION_NAMESPACE, Some(&tool.name), &[]), + ToolSpec::Namespace(namespace) => { + (&namespace.name, None, namespace.tools.as_slice()) + } + ToolSpec::ToolSearch { .. } | ToolSpec::WebSearch { .. } => { + unreachable!("search entries contain only loadable tools"); + } + }; + name.into_iter() + .chain(tools.iter().map(|tool| match tool { + ResponsesApiNamespaceTool::Function(tool) => tool.name.as_str(), + ResponsesApiNamespaceTool::Custom(tool) => tool.name.as_str(), + })) + .map(move |name| ToolName::namespaced(namespace, name)) + } + /// Materialize only selected results; output schemas remain shared until discarded. pub fn to_loadable_spec(&self) -> LoadableToolSpec { let Some(output) = normalize_search_spec(self.spec.as_ref().clone()) else { diff --git a/codex-rs/tools/src/tool_search_tests.rs b/codex-rs/tools/src/tool_search_tests.rs index 106c617868..9e34967e59 100644 --- a/codex-rs/tools/src/tool_search_tests.rs +++ b/codex-rs/tools/src/tool_search_tests.rs @@ -26,23 +26,46 @@ fn shared_search_specs_preserve_results_and_release_the_source() { definition: "start: \"patch\"".to_string(), }, }; - for spec in [ - ToolSpec::Function(function.clone()), - ToolSpec::Freeform(custom.clone()), - ToolSpec::Namespace(ResponsesApiNamespace { - name: "example".to_string(), - description: String::new(), - tools: vec![ - ResponsesApiNamespaceTool::Function(function), - ResponsesApiNamespaceTool::Custom(custom), + for (spec, expected_names) in [ + ( + ToolSpec::Function(function.clone()), + vec![ToolName::namespaced("functions", "lookup")], + ), + ( + ToolSpec::Freeform(custom.clone()), + vec![ToolName::namespaced("functions", "patch")], + ), + ( + ToolSpec::Namespace(ResponsesApiNamespace { + name: "example".to_string(), + description: String::new(), + tools: vec![ + ResponsesApiNamespaceTool::Function(function), + ResponsesApiNamespaceTool::Custom(custom), + ], + }), + vec![ + ToolName::namespaced("example", "lookup"), + ToolName::namespaced("example", "patch"), ], - }), + ), + ( + ToolSpec::Namespace(ResponsesApiNamespace { + name: "empty".to_string(), + description: String::new(), + tools: Vec::new(), + }), + Vec::new(), + ), ] { - let expected = + let normalized = ToolSearchInfo::from_spec("query".to_string(), spec.clone(), /*source_info*/ None) - .unwrap() - .entry - .to_loadable_spec(); + .unwrap(); + assert_eq!( + normalized.entry.tool_names().collect::>(), + expected_names + ); + let expected = normalized.entry.to_loadable_spec(); let spec = Arc::new(spec); let weak = Arc::downgrade(&spec); let info = ToolSearchInfo::from_shared_spec( @@ -57,6 +80,7 @@ fn shared_search_specs_preserve_results_and_release_the_source() { "search retains the original spec rather than a deep copy" ); assert_eq!(info.entry.to_loadable_spec(), expected); + assert_eq!(info.entry.tool_names().collect::>(), expected_names); drop(info); assert!( weak.upgrade().is_none(),