From 54af11932681ab178e6307ba9ec16468fa9d0b3d Mon Sep 17 00:00:00 2001 From: pash Date: Tue, 14 Apr 2026 15:26:37 -0700 Subject: [PATCH] Make deferred dynamic tools searchable --- codex-rs/app-server/README.md | 2 +- .../core/src/tools/handlers/tool_search.rs | 234 +++++++++++++++--- codex-rs/core/src/tools/spec.rs | 12 +- codex-rs/core/tests/suite/search_tool.rs | 2 +- codex-rs/tools/src/tool_discovery.rs | 4 +- codex-rs/tools/src/tool_discovery_tests.rs | 4 +- codex-rs/tools/src/tool_registry_plan.rs | 43 +++- .../tools/src/tool_registry_plan_tests.rs | 51 ++++ 8 files changed, 292 insertions(+), 60 deletions(-) diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 5cfe052b50..ed500062d7 100644 --- a/codex-rs/app-server/README.md +++ b/codex-rs/app-server/README.md @@ -1167,7 +1167,7 @@ If the session approval policy uses `Granular` with `request_permissions: false` `dynamicTools` on `thread/start` and the corresponding `item/tool/call` request/response flow are experimental APIs. To enable them, set `initialize.params.capabilities.experimentalApi = true`. -Each dynamic tool may set `deferLoading`. When omitted, it defaults to `false`. Set it to `true` to keep the tool registered and callable by runtime features such as `js_repl`, while excluding it from the model-facing tool list sent on ordinary turns. +Each dynamic tool may set `deferLoading`. When omitted, it defaults to `false`. Set it to `true` to keep the tool registered and callable by runtime features such as `js_repl`, while excluding it from the model-facing tool list sent on ordinary turns. When `tool_search` is available, deferred dynamic tools are searchable and can be exposed by a matching search result. When a dynamic tool is invoked during a turn, the server sends an `item/tool/call` JSON-RPC request to the client: diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 4bab746ac2..bc43e7b575 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -9,28 +9,49 @@ use bm25::Language; use bm25::SearchEngine; use bm25::SearchEngineBuilder; use codex_mcp::ToolInfo; +use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT; use codex_tools::TOOL_SEARCH_TOOL_NAME; +use codex_tools::ToolSearchOutputTool; use codex_tools::ToolSearchResultSource; use codex_tools::collect_tool_search_output_tools; +use codex_tools::dynamic_tool_to_responses_api_tool; const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use"; const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20; pub struct ToolSearchHandler { - entries: Vec<(String, ToolInfo)>, + entries: Vec, search_engine: SearchEngine, } impl ToolSearchHandler { - pub fn new(tools: std::collections::HashMap) -> Self { - let mut entries: Vec<(String, ToolInfo)> = tools.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); + pub fn new( + mcp_tools: std::collections::HashMap, + dynamic_tools: Vec, + ) -> Self { + let mut mcp_entries: Vec = mcp_tools + .into_iter() + .map(|(name, info)| ToolSearchEntry::Mcp { + name, + info: Box::new(info), + }) + .collect(); + mcp_entries.sort_by(|a, b| a.sort_key().cmp(b.sort_key())); + + let mut dynamic_entries: Vec = dynamic_tools + .into_iter() + .map(|tool| ToolSearchEntry::Dynamic { tool }) + .collect(); + dynamic_entries.sort_by(|a, b| a.sort_key().cmp(b.sort_key())); + + let mut entries = mcp_entries; + entries.extend(dynamic_entries); let documents: Vec> = entries .iter() .enumerate() - .map(|(idx, (name, info))| Document::new(idx, build_search_text(name, info))) + .map(|(idx, entry)| Document::new(idx, entry.search_text())) .collect(); let search_engine = SearchEngineBuilder::::with_documents(Language::English, documents).build(); @@ -83,35 +104,61 @@ impl ToolHandler for ToolSearchHandler { return Ok(ToolSearchOutput { tools: Vec::new() }); } - let results = self.search_result_entries(query, limit, requested_limit.is_none()); - - let tools = collect_tool_search_output_tools(results.into_iter().map(|(_, tool)| { - ToolSearchResultSource { - server_name: tool.server_name.as_str(), - tool_namespace: tool.callable_namespace.as_str(), - tool_name: tool.callable_name.as_str(), - tool: &tool.tool, - connector_name: tool.connector_name.as_deref(), - connector_description: tool.connector_description.as_deref(), - } - })) - .map_err(|err| { - FunctionCallError::Fatal(format!( - "failed to encode {TOOL_SEARCH_TOOL_NAME} output: {err}" - )) - })?; + let tools = self.search(query, limit, requested_limit.is_none())?; Ok(ToolSearchOutput { tools }) } } impl ToolSearchHandler { + fn search( + &self, + query: &str, + limit: usize, + use_default_limit: bool, + ) -> Result, FunctionCallError> { + let results = self.search_result_entries(query, limit, use_default_limit); + let mut tools = Vec::new(); + let mut pending_mcp_sources = Vec::new(); + + for entry in results { + match entry { + ToolSearchEntry::Mcp { info, .. } => { + pending_mcp_sources.push(ToolSearchResultSource { + server_name: info.server_name.as_str(), + tool_namespace: info.callable_namespace.as_str(), + tool_name: info.callable_name.as_str(), + tool: &info.tool, + connector_name: info.connector_name.as_deref(), + connector_description: info.connector_description.as_deref(), + }); + } + ToolSearchEntry::Dynamic { tool } => { + tools.extend( + collect_tool_search_output_tools(pending_mcp_sources.drain(..)) + .map_err(tool_search_output_error)?, + ); + tools.push(ToolSearchOutputTool::Function( + dynamic_tool_to_responses_api_tool(tool) + .map_err(tool_search_output_error)?, + )); + } + } + } + + tools.extend( + collect_tool_search_output_tools(pending_mcp_sources) + .map_err(tool_search_output_error)?, + ); + Ok(tools) + } + fn search_result_entries( &self, query: &str, limit: usize, use_default_limit: bool, - ) -> Vec<&(String, ToolInfo)> { + ) -> Vec<&ToolSearchEntry> { let mut results = self .search_engine .search(query, limit) @@ -124,7 +171,7 @@ impl ToolSearchHandler { if results .iter() - .any(|(_, tool)| tool.server_name == COMPUTER_USE_MCP_SERVER_NAME) + .any(|entry| entry.mcp_server_name() == Some(COMPUTER_USE_MCP_SERVER_NAME)) { results = self .search_engine @@ -137,15 +184,17 @@ impl ToolSearchHandler { } } -fn limit_results_per_server(results: Vec<&(String, ToolInfo)>) -> Vec<&(String, ToolInfo)> { +fn limit_results_per_server(results: Vec<&ToolSearchEntry>) -> Vec<&ToolSearchEntry> { results .into_iter() .scan( std::collections::HashMap::<&str, usize>::new(), |counts, entry| { - let tool = &entry.1; - let count = counts.entry(tool.server_name.as_str()).or_default(); - if *count >= default_limit_for_server(tool.server_name.as_str()) { + let Some(server_name) = entry.mcp_server_name() else { + return Some(Some(entry)); + }; + let count = counts.entry(server_name).or_default(); + if *count >= default_limit_for_server(server_name) { Some(None) } else { *count += 1; @@ -165,7 +214,41 @@ fn default_limit_for_server(server_name: &str) -> usize { } } -fn build_search_text(name: &str, info: &ToolInfo) -> String { +enum ToolSearchEntry { + Mcp { name: String, info: Box }, + Dynamic { tool: DynamicToolSpec }, +} + +impl ToolSearchEntry { + fn sort_key(&self) -> &str { + match self { + Self::Mcp { name, .. } => name.as_str(), + Self::Dynamic { tool } => tool.name.as_str(), + } + } + + fn search_text(&self) -> String { + match self { + Self::Mcp { name, info } => build_mcp_search_text(name, info), + Self::Dynamic { tool } => build_dynamic_search_text(tool), + } + } + + fn mcp_server_name(&self) -> Option<&str> { + match self { + Self::Mcp { info, .. } => Some(info.server_name.as_str()), + Self::Dynamic { .. } => None, + } + } +} + +fn tool_search_output_error(err: serde_json::Error) -> FunctionCallError { + FunctionCallError::Fatal(format!( + "failed to encode {TOOL_SEARCH_TOOL_NAME} output: {err}" + )) +} + +fn build_mcp_search_text(name: &str, info: &ToolInfo) -> String { let mut parts = vec![ name.to_string(), info.callable_name.clone(), @@ -218,20 +301,91 @@ fn build_search_text(name: &str, info: &ToolInfo) -> String { parts.join(" ") } +fn build_dynamic_search_text(tool: &DynamicToolSpec) -> String { + let mut parts = vec![ + tool.name.clone(), + tool.name.replace('_', " "), + tool.description.clone(), + ]; + + parts.extend( + tool.input_schema + .get("properties") + .and_then(serde_json::Value::as_object) + .map(|map| map.keys().cloned().collect::>()) + .unwrap_or_default(), + ); + + parts.join(" ") +} + #[cfg(test)] mod tests { use super::*; + use codex_tools::JsonSchema; + use codex_tools::ResponsesApiTool; use pretty_assertions::assert_eq; use rmcp::model::Tool; use std::sync::Arc; + #[test] + fn search_returns_deferred_dynamic_tools() { + let handler = ToolSearchHandler::new( + std::collections::HashMap::new(), + vec![DynamicToolSpec { + name: "automation_update".to_string(), + description: "Create, update, view, or delete recurring automations.".to_string(), + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "mode": { "type": "string" }, + "kind": { "type": "string" } + }, + "required": ["mode"], + "additionalProperties": false, + }), + defer_loading: true, + }], + ); + + let tools = handler + .search( + "automation_update", + TOOL_SEARCH_DEFAULT_LIMIT, + /*use_default_limit*/ true, + ) + .expect("search deferred dynamic tools"); + + assert_eq!( + tools, + vec![ToolSearchOutputTool::Function(ResponsesApiTool { + name: "automation_update".to_string(), + description: "Create, update, view, or delete recurring automations.".to_string(), + strict: false, + defer_loading: Some(true), + parameters: JsonSchema::object( + std::collections::BTreeMap::from([ + ("kind".to_string(), JsonSchema::string(None)), + ("mode".to_string(), JsonSchema::string(None)), + ]), + Some(vec!["mode".to_string()]), + Some(false.into()), + ), + output_schema: None, + })], + ); + } + #[test] fn computer_use_tool_search_uses_larger_limit() { - let handler = ToolSearchHandler::new(numbered_tools( - COMPUTER_USE_MCP_SERVER_NAME, - "computer use", - /*count*/ 100, - )); + let handler = ToolSearchHandler::new( + numbered_tools( + COMPUTER_USE_MCP_SERVER_NAME, + "computer use", + /*count*/ 100, + ), + Vec::new(), + ); let results = handler.search_result_entries( "computer use", @@ -243,7 +397,7 @@ mod tests { assert!( results .iter() - .all(|(_, tool)| tool.server_name == COMPUTER_USE_MCP_SERVER_NAME) + .all(|entry| entry.mcp_server_name() == Some(COMPUTER_USE_MCP_SERVER_NAME)) ); let explicit_results = handler.search_result_entries( @@ -267,7 +421,7 @@ mod tests { "calendar", /*count*/ 100, )); - let handler = ToolSearchHandler::new(tools); + let handler = ToolSearchHandler::new(tools, Vec::new()); let results = handler.search_result_entries( "calendar", @@ -279,7 +433,7 @@ mod tests { assert!( results .iter() - .all(|(_, tool)| tool.server_name == "other-server") + .all(|entry| entry.mcp_server_name() == Some("other-server")) ); let explicit_results = handler.search_result_entries( @@ -301,7 +455,7 @@ mod tests { "computer use", /*count*/ 100, )); - let handler = ToolSearchHandler::new(tools); + let handler = ToolSearchHandler::new(tools, Vec::new()); let results = handler.search_result_entries( "computer use", @@ -360,10 +514,10 @@ mod tests { } } - fn count_results_for_server(results: &[&(String, ToolInfo)], server_name: &str) -> usize { + fn count_results_for_server(results: &[&ToolSearchEntry], server_name: &str) -> usize { results .iter() - .filter(|(_, tool)| tool.server_name == server_name) + .filter(|entry| entry.mcp_server_name() == Some(server_name)) .count() } } diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 1a1a9f2e5f..de0184b7de 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -157,6 +157,11 @@ pub(crate) fn build_specs_with_discoverable_tools( let request_user_input_handler = Arc::new(RequestUserInputHandler { default_mode_request_user_input: config.default_mode_request_user_input, }); + let deferred_dynamic_tools = dynamic_tools + .iter() + .filter(|tool| tool.defer_loading) + .cloned() + .collect::>(); let mut tool_search_handler = None; let tool_suggest_handler = Arc::new(ToolSuggestHandler); let code_mode_handler = Arc::new(CodeModeExecuteHandler); @@ -259,9 +264,10 @@ pub(crate) fn build_specs_with_discoverable_tools( } ToolHandlerKind::ToolSearch => { if tool_search_handler.is_none() { - tool_search_handler = deferred_mcp_tools - .as_ref() - .map(|tools| Arc::new(ToolSearchHandler::new(tools.clone()))); + tool_search_handler = Some(Arc::new(ToolSearchHandler::new( + deferred_mcp_tools.clone().unwrap_or_default(), + deferred_dynamic_tools.clone(), + ))); } if let Some(tool_search_handler) = tool_search_handler.as_ref() { builder.register_handler(handler.name, tool_search_handler.clone()); diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index eff7bfb427..5fe2863cfa 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -39,7 +39,7 @@ use std::collections::HashMap; use std::time::Duration; const SEARCH_TOOL_DESCRIPTION_SNIPPETS: [&str; 2] = [ - "You have access to tools from the following MCP servers/connectors", + "You have access to tools from the following sources", "- Calendar: Plan events and manage your calendar.", ]; const TOOL_SEARCH_TOOL_NAME: &str = "tool_search"; diff --git a/codex-rs/tools/src/tool_discovery.rs b/codex-rs/tools/src/tool_discovery.rs index b773623edf..272d2576c7 100644 --- a/codex-rs/tools/src/tool_discovery.rs +++ b/codex-rs/tools/src/tool_discovery.rs @@ -153,7 +153,7 @@ pub fn create_tool_search_tool( let properties = BTreeMap::from([ ( "query".to_string(), - JsonSchema::string(Some("Search query for MCP tools.".to_string())), + JsonSchema::string(Some("Search query for deferred tools.".to_string())), ), ( "limit".to_string(), @@ -189,7 +189,7 @@ pub fn create_tool_search_tool( }; let description = format!( - "# MCP tool discovery\n\nSearches over MCP tool metadata with BM25 and exposes matching tools for the next model call.\n\nYou have access to tools from the following MCP servers/connectors:\n{source_descriptions}\nSome of the tools may not have been provided to you upfront, and you should use this tool (`{TOOL_SEARCH_TOOL_NAME}`) to search for the required MCP tools. For MCP tool discovery, always use `{TOOL_SEARCH_TOOL_NAME}` instead of `list_mcp_resources` or `list_mcp_resource_templates`." + "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.\n\nYou have access to tools from the following sources:\n{source_descriptions}\nSome of the tools may not have been provided to you upfront, and you should use this tool (`{TOOL_SEARCH_TOOL_NAME}`) to search for the required tools. For MCP tool discovery, always use `{TOOL_SEARCH_TOOL_NAME}` instead of `list_mcp_resources` or `list_mcp_resource_templates`." ); ToolSpec::ToolSearch { diff --git a/codex-rs/tools/src/tool_discovery_tests.rs b/codex-rs/tools/src/tool_discovery_tests.rs index b745f756cc..bc0f84b0b4 100644 --- a/codex-rs/tools/src/tool_discovery_tests.rs +++ b/codex-rs/tools/src/tool_discovery_tests.rs @@ -50,7 +50,7 @@ fn create_tool_search_tool_deduplicates_and_renders_enabled_sources() { ), ToolSpec::ToolSearch { execution: "client".to_string(), - description: "# MCP tool discovery\n\nSearches over MCP tool metadata with BM25 and exposes matching tools for the next model call.\n\nYou have access to tools from the following MCP servers/connectors:\n- Google Drive: Use Google Drive as the single entrypoint for Drive, Docs, Sheets, and Slides work.\n- docs\nSome of the tools may not have been provided to you upfront, and you should use this tool (`tool_search`) to search for the required MCP tools. For MCP tool discovery, always use `tool_search` instead of `list_mcp_resources` or `list_mcp_resource_templates`.".to_string(), + description: "# Tool discovery\n\nSearches over deferred tool metadata with BM25 and exposes matching tools for the next model call.\n\nYou have access to tools from the following sources:\n- Google Drive: Use Google Drive as the single entrypoint for Drive, Docs, Sheets, and Slides work.\n- docs\nSome of the tools may not have been provided to you upfront, and you should use this tool (`tool_search`) to search for the required tools. For MCP tool discovery, always use `tool_search` instead of `list_mcp_resources` or `list_mcp_resource_templates`.".to_string(), parameters: JsonSchema::object(BTreeMap::from([ ( "limit".to_string(), @@ -61,7 +61,7 @@ fn create_tool_search_tool_deduplicates_and_renders_enabled_sources() { ), ( "query".to_string(), - JsonSchema::string(Some("Search query for MCP tools.".to_string()),), + JsonSchema::string(Some("Search query for deferred tools.".to_string()),), ), ]), Some(vec!["query".to_string()]), Some(false.into())), } diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index ac11a2c4be..79ec8bb350 100644 --- a/codex-rs/tools/src/tool_registry_plan.rs +++ b/codex-rs/tools/src/tool_registry_plan.rs @@ -11,6 +11,7 @@ use crate::ToolHandlerKind; use crate::ToolRegistryPlan; use crate::ToolRegistryPlanParams; use crate::ToolSearchSource; +use crate::ToolSearchSourceInfo; use crate::ToolSpec; use crate::ToolsConfig; use crate::ViewImageToolOptions; @@ -251,17 +252,35 @@ pub fn build_tool_registry_plan( plan.register_handler("request_permissions", ToolHandlerKind::RequestPermissions); } + let deferred_dynamic_tools = params + .dynamic_tools + .iter() + .filter(|tool| tool.defer_loading) + .collect::>(); + if config.search_tool - && let Some(deferred_mcp_tools) = params.deferred_mcp_tools + && (params.deferred_mcp_tools.is_some() || !deferred_dynamic_tools.is_empty()) { - let search_source_infos = - collect_tool_search_source_infos(deferred_mcp_tools.iter().map(|tool| { - ToolSearchSource { - server_name: tool.server_name, - connector_name: tool.connector_name, - connector_description: tool.connector_description, - } - })); + let mut search_source_infos = params + .deferred_mcp_tools + .map(|deferred_mcp_tools| { + collect_tool_search_source_infos(deferred_mcp_tools.iter().map(|tool| { + ToolSearchSource { + server_name: tool.server_name, + connector_name: tool.connector_name, + connector_description: tool.connector_description, + } + })) + }) + .unwrap_or_default(); + + if !deferred_dynamic_tools.is_empty() { + search_source_infos.push(ToolSearchSourceInfo { + name: "Dynamic tools".to_string(), + description: Some("Tools provided by the current Codex thread.".to_string()), + }); + } + plan.push_spec( create_tool_search_tool(&search_source_infos, TOOL_SEARCH_DEFAULT_LIMIT), /*supports_parallel_tool_calls*/ true, @@ -269,8 +288,10 @@ pub fn build_tool_registry_plan( ); plan.register_handler(TOOL_SEARCH_TOOL_NAME, ToolHandlerKind::ToolSearch); - for tool in deferred_mcp_tools { - plan.register_handler(tool.name.clone(), ToolHandlerKind::Mcp); + if let Some(deferred_mcp_tools) = params.deferred_mcp_tools { + for tool in deferred_mcp_tools { + plan.register_handler(tool.name.clone(), ToolHandlerKind::Mcp); + } } } diff --git a/codex-rs/tools/src/tool_registry_plan_tests.rs b/codex-rs/tools/src/tool_registry_plan_tests.rs index cc8c05dd9a..798f7ec85f 100644 --- a/codex-rs/tools/src/tool_registry_plan_tests.rs +++ b/codex-rs/tools/src/tool_registry_plan_tests.rs @@ -1407,6 +1407,57 @@ fn search_tool_requires_model_capability_and_enabled_feature() { assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME]); } +#[test] +fn search_tool_registers_for_deferred_dynamic_tools() { + let model_info = search_capable_model_info(); + let mut features = Features::with_defaults(); + features.enable(Feature::ToolSearch); + let available_models = Vec::new(); + let tools_config = ToolsConfig::new(&ToolsConfigParams { + model_info: &model_info, + available_models: &available_models, + features: &features, + image_generation_tool_auth_allowed: true, + web_search_mode: Some(WebSearchMode::Cached), + session_source: SessionSource::Cli, + sandbox_policy: &SandboxPolicy::DangerFullAccess, + windows_sandbox_level: WindowsSandboxLevel::Disabled, + }); + let dynamic_tool = DynamicToolSpec { + name: "automation_update".to_string(), + description: "Create, update, view, or delete recurring automations.".to_string(), + input_schema: json!({ + "type": "object", + "properties": { + "mode": { "type": "string" }, + }, + }), + defer_loading: true, + }; + + let (tools, handlers) = build_specs( + &tools_config, + /*mcp_tools*/ None, + /*deferred_mcp_tools*/ None, + &[dynamic_tool], + ); + + let search_tool = find_tool(&tools, TOOL_SEARCH_TOOL_NAME); + let ToolSpec::ToolSearch { description, .. } = &search_tool.spec else { + panic!("expected tool_search tool"); + }; + assert!(description.contains("- Dynamic tools: Tools provided by the current Codex thread.")); + assert_contains_tool_names(&tools, &[TOOL_SEARCH_TOOL_NAME, "automation_update"]); + assert!(handlers.contains(&ToolHandlerSpec { + name: ToolName::plain(TOOL_SEARCH_TOOL_NAME), + kind: ToolHandlerKind::ToolSearch, + })); + assert!(handlers.contains(&ToolHandlerSpec { + name: ToolName::plain("automation_update"), + kind: ToolHandlerKind::DynamicTool, + })); +} + #[test] fn tool_suggest_is_not_registered_without_feature_flag() { let model_info = search_capable_model_info();