From 6991be7eadda6bc37f15723b250709229c2be94b Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Sat, 18 Apr 2026 02:07:59 +0800 Subject: [PATCH] enable tool search over dynamic tools (#18263) ## Summary - Normalize deferred MCP and dynamic tools into `ToolSearchEntry` values before constructing `ToolSearchHandler`. - Move the tool-search entry adapter out of `tools/handlers` and into `tools/tool_search_entry.rs` so the handlers directory stays focused on handlers. - Keep `ToolSearchHandler` operating over one generic entry list for BM25 search, namespace grouping, and per-bucket default limits. ## Why Follow-up cleanup for #17849. The dynamic tool-search support made the handler juggle source-specific MCP and dynamic tool lists, index arithmetic, output conversion, and namespace emission. This keeps source adaptation outside the handler so the search loop itself is smaller and source-agnostic. ## Validation - `just fmt` - `cargo test -p codex-core tools::handlers::tool_search::tests` - `git diff --check` - `cargo test -p codex-core` currently fails in unrelated `plugins::manager::tests::list_marketplaces_ignores_installed_roots_missing_from_config`; rerunning that single test fails the same way at `core/src/plugins/manager_tests.rs:1692`. --------- Co-authored-by: pash --- codex-rs/app-server/README.md | 2 +- .../core/src/tools/handlers/tool_search.rs | 302 +++++++++++------- codex-rs/core/src/tools/mod.rs | 1 + codex-rs/core/src/tools/spec.rs | 14 +- codex-rs/core/src/tools/tool_search_entry.rs | 147 +++++++++ codex-rs/core/tests/suite/search_tool.rs | 193 ++++++++++- codex-rs/tools/src/lib.rs | 2 +- codex-rs/tools/src/tool_discovery.rs | 84 ++--- codex-rs/tools/src/tool_discovery_tests.rs | 194 +---------- codex-rs/tools/src/tool_registry_plan.rs | 43 ++- .../tools/src/tool_registry_plan_tests.rs | 51 +++ 11 files changed, 654 insertions(+), 379 deletions(-) create mode 100644 codex-rs/core/src/tools/tool_search_entry.rs diff --git a/codex-rs/app-server/README.md b/codex-rs/app-server/README.md index 1b1c1e920..8c1be6083 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 4bab746ac..7103023d9 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -4,33 +4,31 @@ use crate::tools::context::ToolPayload; use crate::tools::context::ToolSearchOutput; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; +use crate::tools::tool_search_entry::ToolSearchEntry; use bm25::Document; use bm25::Language; use bm25::SearchEngine; use bm25::SearchEngineBuilder; -use codex_mcp::ToolInfo; use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT; use codex_tools::TOOL_SEARCH_TOOL_NAME; -use codex_tools::ToolSearchResultSource; -use codex_tools::collect_tool_search_output_tools; +use codex_tools::ToolSearchOutputTool; +use std::collections::HashMap; 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(crate) fn new(entries: Vec) -> Self { let documents: Vec> = entries .iter() + .map(|entry| entry.search_text.clone()) .enumerate() - .map(|(idx, (name, info))| Document::new(idx, build_search_text(name, info))) + .map(|(idx, search_text)| Document::new(idx, search_text)) .collect(); let search_engine = SearchEngineBuilder::::with_documents(Language::English, documents).build(); @@ -83,155 +81,230 @@ 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); + self.search_output_tools(results) + } + 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) .into_iter() - .filter_map(|result| self.entries.get(result.document.id)) + .map(|result| result.document.id) + .filter_map(|id| self.entries.get(id)) .collect::>(); if !use_default_limit { return results; } - if results - .iter() - .any(|(_, tool)| tool.server_name == COMPUTER_USE_MCP_SERVER_NAME) - { + if results.iter().any(|entry| { + entry + .limit_bucket + .as_deref() + .is_some_and(|bucket| bucket == COMPUTER_USE_MCP_SERVER_NAME) + }) { results = self .search_engine .search(query, COMPUTER_USE_TOOL_SEARCH_LIMIT) .into_iter() - .filter_map(|result| self.entries.get(result.document.id)) + .map(|result| result.document.id) + .filter_map(|id| self.entries.get(id)) .collect(); } - limit_results_per_server(results) + limit_results_by_bucket(results) + } + + fn search_output_tools<'a>( + &self, + results: impl IntoIterator, + ) -> Result, FunctionCallError> { + let mut tools = Vec::new(); + // Preserve search order: group namespace children, emit standalone tools directly. + for entry in results { + match &entry.output { + ToolSearchOutputTool::Function(tool) => { + tools.push(ToolSearchOutputTool::Function(tool.clone())); + } + ToolSearchOutputTool::Namespace(namespace) => { + if let Some(output) = tools.iter_mut().find_map(|tool| match tool { + ToolSearchOutputTool::Namespace(output) + if output.name == namespace.name => + { + Some(output) + } + ToolSearchOutputTool::Namespace(_) | ToolSearchOutputTool::Function(_) => { + None + } + }) { + output.tools.extend(namespace.tools.clone()); + } else { + tools.push(ToolSearchOutputTool::Namespace(namespace.clone())); + } + } + } + } + + Ok(tools) } } -fn limit_results_per_server(results: Vec<&(String, ToolInfo)>) -> Vec<&(String, ToolInfo)> { +fn limit_results_by_bucket(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()) { - Some(None) - } else { - *count += 1; - Some(Some(entry)) - } - }, - ) + .scan(HashMap::<&str, usize>::new(), |counts, result| { + let Some(bucket) = result.limit_bucket.as_deref() else { + return Some(Some(result)); + }; + let count = counts.entry(bucket).or_default(); + if *count >= default_limit_for_bucket(bucket) { + Some(None) + } else { + *count += 1; + Some(Some(result)) + } + }) .flatten() .collect() } -fn default_limit_for_server(server_name: &str) -> usize { - if server_name == COMPUTER_USE_MCP_SERVER_NAME { +fn default_limit_for_bucket(bucket: &str) -> usize { + if bucket == COMPUTER_USE_MCP_SERVER_NAME { COMPUTER_USE_TOOL_SEARCH_LIMIT } else { TOOL_SEARCH_DEFAULT_LIMIT } } -fn build_search_text(name: &str, info: &ToolInfo) -> String { - let mut parts = vec![ - name.to_string(), - info.callable_name.clone(), - info.tool.name.to_string(), - info.server_name.clone(), - ]; - - if let Some(title) = info.tool.title.as_deref() - && !title.trim().is_empty() - { - parts.push(title.to_string()); - } - - if let Some(description) = info.tool.description.as_deref() - && !description.trim().is_empty() - { - parts.push(description.to_string()); - } - - if let Some(connector_name) = info.connector_name.as_deref() - && !connector_name.trim().is_empty() - { - parts.push(connector_name.to_string()); - } - - if let Some(connector_description) = info.connector_description.as_deref() - && !connector_description.trim().is_empty() - { - parts.push(connector_description.to_string()); - } - - parts.extend( - info.plugin_display_names - .iter() - .map(String::as_str) - .map(str::trim) - .filter(|name| !name.is_empty()) - .map(str::to_string), - ); - - parts.extend( - info.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 crate::tools::tool_search_entry::build_tool_search_entries; + use codex_mcp::ToolInfo; + use codex_protocol::dynamic_tools::DynamicToolSpec; + use codex_tools::ResponsesApiNamespace; + use codex_tools::ResponsesApiNamespaceTool; + use codex_tools::ResponsesApiTool; use pretty_assertions::assert_eq; use rmcp::model::Tool; use std::sync::Arc; + #[test] + fn mixed_search_results_coalesce_mcp_namespaces() { + let dynamic_tools = 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" }, + }, + "required": ["mode"], + "additionalProperties": false, + }), + defer_loading: true, + }]; + let handler = handler_from_tools( + Some(&std::collections::HashMap::from([ + ( + "mcp__calendar__create_event".to_string(), + tool_info("calendar", "create_event", "Create events"), + ), + ( + "mcp__calendar__list_events".to_string(), + tool_info("calendar", "list_events", "List events"), + ), + ])), + &dynamic_tools, + ); + let results = [ + &handler.entries[0], + &handler.entries[2], + &handler.entries[1], + ]; + + let tools = handler + .search_output_tools(results) + .expect("mixed search output should serialize"); + + assert_eq!( + tools, + vec![ + ToolSearchOutputTool::Namespace(ResponsesApiNamespace { + name: "mcp__calendar__".to_string(), + description: "Tools in the mcp__calendar__ namespace.".to_string(), + tools: vec![ + ResponsesApiNamespaceTool::Function(ResponsesApiTool { + name: "create_event".to_string(), + description: "Create events desktop tool".to_string(), + strict: false, + defer_loading: Some(true), + parameters: codex_tools::JsonSchema::object( + Default::default(), + /*required*/ None, + Some(false.into()), + ), + output_schema: None, + }), + ResponsesApiNamespaceTool::Function(ResponsesApiTool { + name: "list_events".to_string(), + description: "List events desktop tool".to_string(), + strict: false, + defer_loading: Some(true), + parameters: codex_tools::JsonSchema::object( + Default::default(), + /*required*/ None, + Some(false.into()), + ), + output_schema: None, + }), + ], + }), + 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: codex_tools::JsonSchema::object( + std::collections::BTreeMap::from([( + "mode".to_string(), + codex_tools::JsonSchema::string(/*description*/ 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( + let tools = numbered_tools( COMPUTER_USE_MCP_SERVER_NAME, "computer use", /*count*/ 100, - )); + ); + let handler = handler_from_tools(Some(&tools), &[]); let results = handler.search_result_entries( "computer use", @@ -243,7 +316,7 @@ mod tests { assert!( results .iter() - .all(|(_, tool)| tool.server_name == COMPUTER_USE_MCP_SERVER_NAME) + .all(|entry| entry.limit_bucket.as_deref() == Some(COMPUTER_USE_MCP_SERVER_NAME)) ); let explicit_results = handler.search_result_entries( @@ -267,7 +340,7 @@ mod tests { "calendar", /*count*/ 100, )); - let handler = ToolSearchHandler::new(tools); + let handler = handler_from_tools(Some(&tools), &[]); let results = handler.search_result_entries( "calendar", @@ -279,7 +352,7 @@ mod tests { assert!( results .iter() - .all(|(_, tool)| tool.server_name == "other-server") + .all(|entry| entry.limit_bucket.as_deref() == Some("other-server")) ); let explicit_results = handler.search_result_entries( @@ -301,7 +374,7 @@ mod tests { "computer use", /*count*/ 100, )); - let handler = ToolSearchHandler::new(tools); + let handler = handler_from_tools(Some(&tools), &[]); let results = handler.search_result_entries( "computer use", @@ -360,10 +433,17 @@ 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.limit_bucket.as_deref() == Some(server_name)) .count() } + + fn handler_from_tools( + mcp_tools: Option<&std::collections::HashMap>, + dynamic_tools: &[DynamicToolSpec], + ) -> ToolSearchHandler { + ToolSearchHandler::new(build_tool_search_entries(mcp_tools, dynamic_tools)) + } } diff --git a/codex-rs/core/src/tools/mod.rs b/codex-rs/core/src/tools/mod.rs index 06c20d692..ee790d7d2 100644 --- a/codex-rs/core/src/tools/mod.rs +++ b/codex-rs/core/src/tools/mod.rs @@ -11,6 +11,7 @@ pub(crate) mod router; pub(crate) mod runtimes; pub(crate) mod sandboxing; pub(crate) mod spec; +pub(crate) mod tool_search_entry; use codex_protocol::exec_output::ExecToolCallOutput; use codex_utils_output_truncation::TruncationPolicy; diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 1a1a9f2e5..41dc422e3 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -108,6 +108,7 @@ pub(crate) fn build_specs_with_discoverable_tools( use crate::tools::handlers::multi_agents_v2::SpawnAgentHandler as SpawnAgentHandlerV2; use crate::tools::handlers::multi_agents_v2::WaitAgentHandler as WaitAgentHandlerV2; use crate::tools::handlers::unavailable_tool_message; + use crate::tools::tool_search_entry::build_tool_search_entries; let mut builder = ToolRegistryBuilder::new(); let mcp_tool_plan_inputs = mcp_tools.as_ref().map(map_mcp_tools_for_plan); @@ -157,6 +158,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 +265,11 @@ 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()))); + let entries = build_tool_search_entries( + deferred_mcp_tools.as_ref(), + &deferred_dynamic_tools, + ); + tool_search_handler = Some(Arc::new(ToolSearchHandler::new(entries))); } 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/src/tools/tool_search_entry.rs b/codex-rs/core/src/tools/tool_search_entry.rs new file mode 100644 index 000000000..6fa493132 --- /dev/null +++ b/codex-rs/core/src/tools/tool_search_entry.rs @@ -0,0 +1,147 @@ +use codex_mcp::ToolInfo; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_tools::ToolSearchOutputTool; +use codex_tools::ToolSearchResultSource; +use codex_tools::dynamic_tool_to_responses_api_tool; +use codex_tools::tool_search_result_source_to_output_tool; +use std::collections::HashMap; + +#[derive(Clone)] +pub(crate) struct ToolSearchEntry { + pub(crate) search_text: String, + pub(crate) output: ToolSearchOutputTool, + pub(crate) limit_bucket: Option, +} + +pub(crate) fn build_tool_search_entries( + mcp_tools: Option<&HashMap>, + dynamic_tools: &[DynamicToolSpec], +) -> Vec { + let mut entries = Vec::new(); + + let mut mcp_tools = mcp_tools + .map(|tools| tools.values().collect::>()) + .unwrap_or_default(); + mcp_tools.sort_by_key(|info| info.canonical_tool_name().display()); + for info in mcp_tools { + match mcp_tool_search_entry(info) { + Ok(entry) => entries.push(entry), + Err(error) => { + let tool_name = info.canonical_tool_name(); + tracing::error!( + "Failed to convert deferred MCP tool `{tool_name}` to OpenAI tool: {error:?}" + ); + } + } + } + + let mut dynamic_tools = dynamic_tools.iter().collect::>(); + dynamic_tools.sort_by(|a, b| a.name.cmp(&b.name)); + for tool in dynamic_tools { + match dynamic_tool_search_entry(tool) { + Ok(entry) => entries.push(entry), + Err(error) => { + tracing::error!( + "Failed to convert deferred dynamic tool {:?} to OpenAI tool: {error:?}", + tool.name + ); + } + } + } + + entries +} + +fn mcp_tool_search_entry(info: &ToolInfo) -> Result { + Ok(ToolSearchEntry { + search_text: build_mcp_search_text(info), + output: tool_search_result_source_to_output_tool(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(), + })?, + limit_bucket: Some(info.server_name.clone()), + }) +} + +fn dynamic_tool_search_entry(tool: &DynamicToolSpec) -> Result { + Ok(ToolSearchEntry { + search_text: build_dynamic_search_text(tool), + output: ToolSearchOutputTool::Function(dynamic_tool_to_responses_api_tool(tool)?), + limit_bucket: None, + }) +} + +fn build_mcp_search_text(info: &ToolInfo) -> String { + let mut parts = vec![ + info.canonical_tool_name().display(), + info.callable_name.clone(), + info.tool.name.to_string(), + info.server_name.clone(), + ]; + + if let Some(title) = info.tool.title.as_deref() + && !title.trim().is_empty() + { + parts.push(title.to_string()); + } + + if let Some(description) = info.tool.description.as_deref() + && !description.trim().is_empty() + { + parts.push(description.to_string()); + } + + if let Some(connector_name) = info.connector_name.as_deref() + && !connector_name.trim().is_empty() + { + parts.push(connector_name.to_string()); + } + + if let Some(connector_description) = info.connector_description.as_deref() + && !connector_description.trim().is_empty() + { + parts.push(connector_description.to_string()); + } + + parts.extend( + info.plugin_display_names + .iter() + .map(String::as_str) + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(str::to_string), + ); + + parts.extend( + info.tool + .input_schema + .get("properties") + .and_then(serde_json::Value::as_object) + .map(|map| map.keys().cloned().collect::>()) + .unwrap_or_default(), + ); + + 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(" ") +} diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index cfcec7992..2cd51c4cf 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -8,6 +8,10 @@ use codex_core::config::Config; use codex_features::Feature; use codex_login::CodexAuth; use codex_models_manager::bundled_models_response; +use codex_protocol::dynamic_tools::DynamicToolCallOutputContentItem; +use codex_protocol::dynamic_tools::DynamicToolResponse; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; use codex_protocol::protocol::McpInvocation; @@ -39,7 +43,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"; @@ -94,12 +98,7 @@ fn tool_search_output_tools(request: &ResponsesRequest, call_id: &str) -> Vec Result<()> { "parameters": { "type": "object", "properties": { - "query": {"type": "string", "description": "Search query for MCP tools."}, + "query": {"type": "string", "description": "Search query for deferred tools."}, "limit": {"type": "number", "description": "Maximum number of tools to return (defaults to 8)."}, }, "required": ["query"], @@ -691,6 +699,175 @@ async fn tool_search_returns_deferred_tools_without_follow_up_tool_injection() - Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn tool_search_returns_deferred_dynamic_tool_and_routes_follow_up_call() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let search_call_id = "tool-search-1"; + let dynamic_call_id = "dyn-search-call-1"; + let tool_name = "automation_update"; + let tool_description = "Create, update, view, or delete recurring automations."; + let tool_args = json!({ "mode": "create" }); + let tool_call_arguments = serde_json::to_string(&tool_args)?; + let mock = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_tool_search_call( + search_call_id, + &json!({ + "query": "recurring automations", + "limit": 8, + }), + ), + ev_completed("resp-1"), + ]), + sse(vec![ + ev_response_created("resp-2"), + json!({ + "type": "response.output_item.done", + "item": { + "type": "function_call", + "call_id": dynamic_call_id, + "name": tool_name, + "arguments": tool_call_arguments, + } + }), + ev_completed("resp-2"), + ]), + sse(vec![ + ev_response_created("resp-3"), + ev_assistant_message("msg-1", "done"), + ev_completed("resp-3"), + ]), + ], + ) + .await; + + let input_schema = json!({ + "type": "object", + "properties": { + "mode": { "type": "string" }, + }, + "required": ["mode"], + "additionalProperties": false, + }); + let dynamic_tool = DynamicToolSpec { + name: tool_name.to_string(), + description: tool_description.to_string(), + input_schema: input_schema.clone(), + defer_loading: true, + }; + + let mut builder = test_codex().with_config(configure_search_capable_model); + let base_test = builder.build(&server).await?; + let new_thread = base_test + .thread_manager + .start_thread_with_tools( + base_test.config.clone(), + vec![dynamic_tool], + /*persist_extended_history*/ false, + ) + .await?; + let mut test = base_test; + test.codex = new_thread.thread; + test.session_configured = new_thread.session_configured; + + test.codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "Use the automation tool".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + }) + .await?; + + let EventMsg::DynamicToolCallRequest(request) = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::DynamicToolCallRequest(_)) + }) + .await + else { + unreachable!("event guard guarantees DynamicToolCallRequest"); + }; + assert_eq!(request.call_id, dynamic_call_id); + assert_eq!(request.tool, tool_name); + assert_eq!(request.arguments, tool_args); + + test.codex + .submit(Op::DynamicToolResponse { + id: request.call_id, + response: DynamicToolResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "dynamic-search-ok".to_string(), + }], + success: true, + }, + }) + .await?; + + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnComplete(_)) + }) + .await; + + let requests = mock.requests(); + assert_eq!(requests.len(), 3); + + let first_request_tools = tool_names(&requests[0].body_json()); + assert!( + first_request_tools + .iter() + .any(|name| name == TOOL_SEARCH_TOOL_NAME), + "first request should advertise tool_search: {first_request_tools:?}" + ); + assert!( + !first_request_tools.iter().any(|name| name == tool_name), + "deferred dynamic tool should be hidden before search: {first_request_tools:?}" + ); + + let tools = tool_search_output_tools(&requests[1], search_call_id); + assert_eq!( + tools, + vec![json!({ + "type": "function", + "name": tool_name, + "description": tool_description, + "strict": false, + "defer_loading": true, + "parameters": input_schema, + })] + ); + + let second_request_tools = tool_names(&requests[1].body_json()); + assert!( + !second_request_tools.iter().any(|name| name == tool_name), + "follow-up request should rely on tool_search_output history, not tool injection: {second_request_tools:?}" + ); + + let output = requests[2] + .function_call_output(dynamic_call_id) + .get("output") + .cloned() + .expect("dynamic tool output should be present"); + let payload: FunctionCallOutputPayload = serde_json::from_value(output)?; + assert_eq!( + payload, + FunctionCallOutputPayload::from_text("dynamic-search-ok".to_string()) + ); + + let third_request_tools = tool_names(&requests[2].body_json()); + assert!( + !third_request_tools.iter().any(|name| name == tool_name), + "post-tool follow-up should rely on tool_search_output history, not tool injection: {third_request_tools:?}" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn tool_search_indexes_only_enabled_non_app_mcp_tools() -> Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index b7c6a57d2..07e98ab5a 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -109,12 +109,12 @@ pub use tool_discovery::ToolSearchResultSource; pub use tool_discovery::ToolSearchSource; pub use tool_discovery::ToolSearchSourceInfo; pub use tool_discovery::ToolSuggestEntry; -pub use tool_discovery::collect_tool_search_output_tools; pub use tool_discovery::collect_tool_search_source_infos; pub use tool_discovery::collect_tool_suggest_entries; pub use tool_discovery::create_tool_search_tool; pub use tool_discovery::create_tool_suggest_tool; pub use tool_discovery::filter_tool_suggest_discoverable_tools_for_client; +pub use tool_discovery::tool_search_result_source_to_output_tool; pub use tool_registry_plan::build_tool_registry_plan; pub use tool_registry_plan_types::ToolHandlerKind; pub use tool_registry_plan_types::ToolHandlerSpec; diff --git a/codex-rs/tools/src/tool_discovery.rs b/codex-rs/tools/src/tool_discovery.rs index b773623ed..359b1d792 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 { @@ -203,58 +203,38 @@ pub fn create_tool_search_tool( } } -pub fn collect_tool_search_output_tools<'a>( - tool_sources: impl IntoIterator>, -) -> Result, serde_json::Error> { - let mut grouped: Vec<(&'a str, Vec>)> = Vec::new(); - for tool in tool_sources { - if let Some((_, tools)) = grouped - .iter_mut() - .find(|(tool_namespace, _)| *tool_namespace == tool.tool_namespace) - { - tools.push(tool); - } else { - grouped.push((tool.tool_namespace, vec![tool])); - } - } +pub fn tool_search_result_source_to_output_tool( + source: ToolSearchResultSource<'_>, +) -> Result { + Ok(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { + name: source.tool_namespace.to_string(), + description: tool_search_result_source_namespace_description(source), + tools: vec![tool_search_result_source_to_namespace_tool(source)?], + })) +} - let mut results = Vec::with_capacity(grouped.len()); - for (tool_namespace, tools) in grouped { - let Some(first_tool) = tools.first() else { - continue; - }; +fn tool_search_result_source_namespace_description(source: ToolSearchResultSource<'_>) -> String { + source + .connector_description + .map(str::trim) + .filter(|description| !description.is_empty()) + .map(str::to_string) + .or_else(|| { + source + .connector_name + .map(str::trim) + .filter(|connector_name| !connector_name.is_empty()) + .map(|connector_name| format!("Tools for working with {connector_name}.")) + }) + .unwrap_or_else(|| default_namespace_description(source.tool_namespace)) +} - let description = first_tool - .connector_description - .map(str::trim) - .filter(|description| !description.is_empty()) - .map(str::to_string) - .or_else(|| { - first_tool - .connector_name - .map(str::trim) - .filter(|connector_name| !connector_name.is_empty()) - .map(|connector_name| format!("Tools for working with {connector_name}.")) - }); - - let tools = tools - .iter() - .map(|tool| { - let tool_name = ToolName::namespaced(tool.tool_namespace, tool.tool_name); - mcp_tool_to_deferred_responses_api_tool(&tool_name, tool.tool) - .map(ResponsesApiNamespaceTool::Function) - }) - .collect::, _>>()?; - - results.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: tool_namespace.to_string(), - description: description - .unwrap_or_else(|| default_namespace_description(tool_namespace)), - tools, - })); - } - - Ok(results) +fn tool_search_result_source_to_namespace_tool( + source: ToolSearchResultSource<'_>, +) -> Result { + let tool_name = ToolName::namespaced(source.tool_namespace, source.tool_name); + mcp_tool_to_deferred_responses_api_tool(&tool_name, source.tool) + .map(ResponsesApiNamespaceTool::Function) } pub fn collect_tool_search_source_infos<'a>( diff --git a/codex-rs/tools/src/tool_discovery_tests.rs b/codex-rs/tools/src/tool_discovery_tests.rs index b745f756c..9afd32231 100644 --- a/codex-rs/tools/src/tool_discovery_tests.rs +++ b/codex-rs/tools/src/tool_discovery_tests.rs @@ -2,28 +2,8 @@ use super::*; use crate::JsonSchema; use codex_app_server_protocol::AppInfo; use pretty_assertions::assert_eq; -use rmcp::model::JsonObject; -use rmcp::model::Tool; use serde_json::json; use std::collections::BTreeMap; -use std::sync::Arc; - -fn mcp_tool(name: &str, description: &str) -> Tool { - Tool { - name: name.to_string().into(), - title: None, - description: Some(description.to_string().into()), - input_schema: Arc::new(JsonObject::from_iter([( - "type".to_string(), - json!("object"), - )])), - output_schema: None, - annotations: None, - execution: None, - icons: None, - meta: None, - } -} #[test] fn create_tool_search_tool_deduplicates_and_renders_enabled_sources() { @@ -50,7 +30,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 +41,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())), } @@ -136,176 +116,6 @@ fn create_tool_suggest_tool_uses_plugin_summary_fallback() { ); } -#[test] -fn collect_tool_search_output_tools_preserves_search_order_while_grouping_by_namespace() { - let calendar_create_event = mcp_tool("calendar-create-event", "Create a calendar event."); - let gmail_read_email = mcp_tool("gmail-read-email", "Read an email."); - let gmail_send_email = mcp_tool("gmail-send-email", "Send an email."); - let calendar_list_events = mcp_tool("calendar-list-events", "List calendar events."); - let docs_search = mcp_tool("search", "Search docs."); - - let tools = collect_tool_search_output_tools([ - ToolSearchResultSource { - server_name: "codex_apps", - tool_namespace: "mcp__codex_apps__gmail", - tool_name: "_send_email", - tool: &gmail_send_email, - connector_name: Some("Gmail"), - connector_description: Some("Read mail"), - }, - ToolSearchResultSource { - server_name: "codex_apps", - tool_namespace: "mcp__codex_apps__calendar", - tool_name: "_create_event", - tool: &calendar_create_event, - connector_name: Some("Calendar"), - connector_description: Some("Plan events"), - }, - ToolSearchResultSource { - server_name: "codex_apps", - tool_namespace: "mcp__codex_apps__gmail", - tool_name: "_read_email", - tool: &gmail_read_email, - connector_name: Some("Gmail"), - connector_description: Some("Read mail"), - }, - ToolSearchResultSource { - server_name: "codex_apps", - tool_namespace: "mcp__codex_apps__calendar", - tool_name: "_list_events", - tool: &calendar_list_events, - connector_name: Some("Calendar"), - connector_description: Some("Plan events"), - }, - ToolSearchResultSource { - server_name: "docs", - tool_namespace: "mcp__docs__", - tool_name: "search", - tool: &docs_search, - connector_name: None, - connector_description: None, - }, - ]) - .expect("collect tool search output tools"); - - assert_eq!( - tools, - vec![ - ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: "mcp__codex_apps__gmail".to_string(), - description: "Read mail".to_string(), - tools: vec![ - ResponsesApiNamespaceTool::Function(ResponsesApiTool { - name: "_send_email".to_string(), - description: "Send an email.".to_string(), - strict: false, - defer_loading: Some(true), - parameters: JsonSchema::object( - Default::default(), - /*required*/ None, - /*additional_properties*/ None - ), - output_schema: None, - }), - ResponsesApiNamespaceTool::Function(ResponsesApiTool { - name: "_read_email".to_string(), - description: "Read an email.".to_string(), - strict: false, - defer_loading: Some(true), - parameters: JsonSchema::object( - Default::default(), - /*required*/ None, - /*additional_properties*/ None - ), - output_schema: None, - }), - ], - }), - ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: "mcp__codex_apps__calendar".to_string(), - description: "Plan events".to_string(), - tools: vec![ - ResponsesApiNamespaceTool::Function(ResponsesApiTool { - name: "_create_event".to_string(), - description: "Create a calendar event.".to_string(), - strict: false, - defer_loading: Some(true), - parameters: JsonSchema::object( - Default::default(), - /*required*/ None, - /*additional_properties*/ None - ), - output_schema: None, - }), - ResponsesApiNamespaceTool::Function(ResponsesApiTool { - name: "_list_events".to_string(), - description: "List calendar events.".to_string(), - strict: false, - defer_loading: Some(true), - parameters: JsonSchema::object( - Default::default(), - /*required*/ None, - /*additional_properties*/ None - ), - output_schema: None, - }), - ], - }), - ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: "mcp__docs__".to_string(), - description: "Tools in the mcp__docs__ namespace.".to_string(), - tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool { - name: "search".to_string(), - description: "Search docs.".to_string(), - strict: false, - defer_loading: Some(true), - parameters: JsonSchema::object( - Default::default(), - /*required*/ None, - /*additional_properties*/ None - ), - output_schema: None, - })], - }), - ], - ); -} - -#[test] -fn collect_tool_search_output_tools_ignores_blank_connector_description() { - let gmail_batch_read_email = mcp_tool("gmail-batch-read-email", "Read multiple emails."); - - let tools = collect_tool_search_output_tools([ToolSearchResultSource { - server_name: "codex_apps", - tool_namespace: "mcp__codex_apps__gmail", - tool_name: "_batch_read_email", - tool: &gmail_batch_read_email, - connector_name: Some("Gmail"), - connector_description: Some(" "), - }]) - .expect("collect tool search output tools"); - - assert_eq!( - tools, - vec![ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: "mcp__codex_apps__gmail".to_string(), - description: "Tools for working with Gmail.".to_string(), - tools: vec![ResponsesApiNamespaceTool::Function(ResponsesApiTool { - name: "_batch_read_email".to_string(), - description: "Read multiple emails.".to_string(), - strict: false, - defer_loading: Some(true), - parameters: JsonSchema::object( - Default::default(), - /*required*/ None, - /*additional_properties*/ None - ), - output_schema: None, - })], - })], - ); -} - #[test] fn discoverable_tool_enums_use_expected_wire_names() { assert_eq!( diff --git a/codex-rs/tools/src/tool_registry_plan.rs b/codex-rs/tools/src/tool_registry_plan.rs index ac11a2c4b..79ec8bb35 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 cc8c05dd9..798f7ec85 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();