From 54af11932681ab178e6307ba9ec16468fa9d0b3d Mon Sep 17 00:00:00 2001 From: pash Date: Tue, 14 Apr 2026 15:26:37 -0700 Subject: [PATCH 1/7] 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 5cfe052b508e..ed500062d702 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 4bab746ac28a..bc43e7b57565 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 1a1a9f2e5ffb..de0184b7de53 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 eff7bfb4275a..5fe2863cfabc 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 b773623edfe2..272d2576c76c 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 b745f756cca4..bc0f84b0b484 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 ac11a2c4beec..79ec8bb350a3 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 cc8c05dd9a5a..798f7ec85f61 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(); From f1e0ee2de505ffe74249343b9c54d71202de5d59 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Thu, 16 Apr 2026 23:44:07 +0800 Subject: [PATCH 2/7] Support tool search for dynamic tools --- .../tests/suite/v2/dynamic_tools.rs | 220 ++++++++++++++++++ .../core/src/tools/handlers/tool_search.rs | 219 ++++++++++------- codex-rs/core/tests/suite/search_tool.rs | 2 +- 3 files changed, 361 insertions(+), 80 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs index 0a3315a07a29..115c7e0de63c 100644 --- a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs +++ b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs @@ -4,6 +4,7 @@ use app_test_support::McpProcess; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; +use app_test_support::write_models_cache_with_models; use codex_app_server_protocol::DynamicToolCallOutputContentItem; use codex_app_server_protocol::DynamicToolCallParams; use codex_app_server_protocol::DynamicToolCallResponse; @@ -21,6 +22,7 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; +use codex_models_manager::bundled_models_response; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -28,6 +30,7 @@ use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; +use std::io::Write; use std::path::Path; use std::time::Duration; use tempfile::TempDir; @@ -196,6 +199,178 @@ async fn thread_start_keeps_hidden_dynamic_tools_out_of_model_requests() -> Resu Ok(()) } +/// Exercises deferred dynamic tool discovery, the follow-up tool call, and the tool response. +#[tokio::test] +async fn deferred_dynamic_tool_can_be_discovered_and_called_through_tool_search() -> Result<()> { + 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 responses = vec![ + responses::sse(vec![ + responses::ev_response_created("resp-1"), + responses::ev_tool_search_call( + search_call_id, + &json!({ + "query": "recurring automations", + "limit": 8, + }), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_response_created("resp-2"), + responses::ev_function_call(dynamic_call_id, tool_name, &tool_call_arguments), + responses::ev_completed("resp-2"), + ]), + create_final_assistant_message_sse_response("Done")?, + ]; + let server = create_mock_responses_server_sequence_unchecked(responses).await; + + let codex_home = TempDir::new()?; + write_search_capable_models_cache(codex_home.path())?; + create_config_toml(codex_home.path(), &server.uri())?; + enable_tool_search_feature(codex_home.path())?; + + let mut mcp = McpProcess::new(codex_home.path()).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).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 thread_req = mcp + .send_thread_start_request(ThreadStartParams { + dynamic_tools: Some(vec![dynamic_tool]), + ..Default::default() + }) + .await?; + let thread_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), + ) + .await??; + let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; + let thread_id = thread.id.clone(); + + let turn_req = mcp + .send_turn_start_request(TurnStartParams { + thread_id: thread_id.clone(), + input: vec![V2UserInput::Text { + text: "Use the automation tool".to_string(), + text_elements: Vec::new(), + }], + ..Default::default() + }) + .await?; + let turn_resp: JSONRPCResponse = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), + ) + .await??; + let TurnStartResponse { turn } = to_response::(turn_resp)?; + let turn_id = turn.id.clone(); + + let started = wait_for_dynamic_tool_started(&mut mcp, dynamic_call_id).await?; + assert_eq!(started.thread_id, thread_id.clone()); + assert_eq!(started.turn_id, turn_id.clone()); + + let request = timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_request_message(), + ) + .await??; + let (request_id, params) = match request { + ServerRequest::DynamicToolCall { request_id, params } => (request_id, params), + other => panic!("expected DynamicToolCall request, got {other:?}"), + }; + assert_eq!( + params, + DynamicToolCallParams { + thread_id: thread_id.clone(), + turn_id: turn_id.clone(), + call_id: dynamic_call_id.to_string(), + tool: tool_name.to_string(), + arguments: tool_args.clone(), + } + ); + + mcp.send_response( + request_id, + serde_json::to_value(DynamicToolCallResponse { + content_items: vec![DynamicToolCallOutputContentItem::InputText { + text: "dynamic-search-ok".to_string(), + }], + success: true, + })?, + ) + .await?; + + let completed = wait_for_dynamic_tool_completed(&mut mcp, dynamic_call_id).await?; + assert_eq!(completed.thread_id, thread_id); + assert_eq!(completed.turn_id, turn_id); + + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + + let bodies = responses_bodies(&server).await?; + let first_request = bodies + .first() + .context("expected an initial responses request")?; + assert!( + find_tool_by_type(first_request, "tool_search").is_some(), + "initial request should advertise tool_search: {first_request:?}" + ); + assert!( + find_tool(first_request, tool_name).is_none(), + "deferred dynamic tool should not be directly advertised before search" + ); + + let search_tools = bodies + .iter() + .find_map(|body| tool_search_output_tools(body, search_call_id)) + .context("expected tool_search_output in follow-up request")?; + assert_eq!( + search_tools, + vec![json!({ + "type": "function", + "name": tool_name, + "description": tool_description, + "strict": false, + "defer_loading": true, + "parameters": input_schema, + })] + ); + + let payload = bodies + .iter() + .find_map(|body| function_call_output_payload(body, dynamic_call_id)) + .context("expected function_call_output in post-tool follow-up request")?; + assert_eq!( + payload, + FunctionCallOutputPayload::from_text("dynamic-search-ok".to_string()) + ); + + Ok(()) +} + /// Exercises the full dynamic tool call path (server request, client response, model output). #[tokio::test] async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Result<()> { @@ -583,6 +758,30 @@ fn find_tool<'a>(body: &'a Value, name: &str) -> Option<&'a Value> { }) } +fn find_tool_by_type<'a>(body: &'a Value, tool_type: &str) -> Option<&'a Value> { + body.get("tools") + .and_then(Value::as_array) + .and_then(|tools| { + tools + .iter() + .find(|tool| tool.get("type").and_then(Value::as_str) == Some(tool_type)) + }) +} + +fn tool_search_output_tools(body: &Value, call_id: &str) -> Option> { + body.get("input") + .and_then(Value::as_array) + .and_then(|items| { + items.iter().find(|item| { + item.get("type").and_then(Value::as_str) == Some("tool_search_output") + && item.get("call_id").and_then(Value::as_str) == Some(call_id) + }) + }) + .and_then(|item| item.get("tools")) + .and_then(Value::as_array) + .cloned() +} + fn function_call_output_payload(body: &Value, call_id: &str) -> Option { function_call_output_raw_output(body, call_id) .and_then(|output| serde_json::from_value(output).ok()) @@ -663,3 +862,24 @@ stream_max_retries = 0 ), ) } + +fn enable_tool_search_feature(codex_home: &Path) -> std::io::Result<()> { + let mut config_toml = std::fs::OpenOptions::new() + .append(true) + .open(codex_home.join("config.toml"))?; + config_toml.write_all(b"\n[features]\ntool_search = true\n") +} + +fn write_search_capable_models_cache(codex_home: &Path) -> Result<()> { + let mut model = bundled_models_response() + .context("bundled models should parse")? + .models + .into_iter() + .find(|model| model.slug == "gpt-5.4") + .context("expected bundled gpt-5.4 model")?; + model.slug = "mock-model".to_string(); + model.display_name = "mock-model".to_string(); + model.supports_search_tool = true; + write_models_cache_with_models(codex_home, vec![model])?; + Ok(()) +} diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index bc43e7b57565..9f920ae2b782 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -10,12 +10,13 @@ use bm25::SearchEngine; use bm25::SearchEngineBuilder; use codex_mcp::ToolInfo; use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_tools::ResponsesApiNamespace; +use codex_tools::ResponsesApiNamespaceTool; 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; +use codex_tools::mcp_tool_to_deferred_responses_api_tool; const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use"; const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20; @@ -31,9 +32,9 @@ impl ToolSearchHandler { dynamic_tools: Vec, ) -> Self { let mut mcp_entries: Vec = mcp_tools - .into_iter() - .map(|(name, info)| ToolSearchEntry::Mcp { - name, + .into_values() + .map(|info| ToolSearchEntry::Mcp { + name: info.canonical_tool_name().display(), info: Box::new(info), }) .collect(); @@ -118,39 +119,7 @@ impl ToolSearchHandler { 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) + search_output_tools(results) } fn search_result_entries( @@ -184,6 +153,60 @@ impl ToolSearchHandler { } } +fn search_output_tools<'a>( + results: impl IntoIterator, +) -> Result, FunctionCallError> { + let mut tools = Vec::new(); + // Preserve search order: group MCP tools under namespaces, emit dynamic tools directly. + for entry in results { + match entry { + ToolSearchEntry::Mcp { info, .. } => { + let tool_name = info.canonical_tool_name(); + let namespace = info.callable_namespace.as_str(); + let namespace_tool = + mcp_tool_to_deferred_responses_api_tool(&tool_name, &info.tool) + .map(ResponsesApiNamespaceTool::Function) + .map_err(tool_search_output_error)?; + + if let Some(output) = tools.iter_mut().find_map(|tool| match tool { + ToolSearchOutputTool::Namespace(output) if output.name == namespace => { + Some(output) + } + ToolSearchOutputTool::Namespace(_) | ToolSearchOutputTool::Function(_) => None, + }) { + output.tools.push(namespace_tool); + } else { + tools.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { + name: namespace.to_string(), + description: mcp_namespace_description(info), + tools: vec![namespace_tool], + })); + } + } + ToolSearchEntry::Dynamic { tool } => { + tools.push(ToolSearchOutputTool::Function( + dynamic_tool_to_responses_api_tool(tool).map_err(tool_search_output_error)?, + )); + } + } + } + + Ok(tools) +} + +fn mcp_namespace_description(info: &ToolInfo) -> String { + info.connector_description + .clone() + .or_else(|| { + info.connector_name + .as_deref() + .map(str::trim) + .filter(|connector_name| !connector_name.is_empty()) + .map(|connector_name| format!("Tools for working with {connector_name}.")) + }) + .unwrap_or_else(|| format!("Tools from the {} MCP server.", info.server_name)) +} + fn limit_results_per_server(results: Vec<&ToolSearchEntry>) -> Vec<&ToolSearchEntry> { results .into_iter() @@ -322,57 +345,95 @@ fn build_dynamic_search_text(tool: &DynamicToolSpec) -> String { #[cfg(test)] mod tests { use super::*; - use codex_tools::JsonSchema; + 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 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, - }], - ); + fn mixed_search_results_coalesce_mcp_namespaces() { + let entries = [ + ToolSearchEntry::Mcp { + name: "mcp__calendar__create_event".to_string(), + info: Box::new(tool_info("calendar", "create_event", "Create events")), + }, + ToolSearchEntry::Dynamic { + tool: 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, + }, + }, + ToolSearchEntry::Mcp { + name: "mcp__calendar__list_events".to_string(), + info: Box::new(tool_info("calendar", "list_events", "List events")), + }, + ]; - let tools = handler - .search( - "automation_update", - TOOL_SEARCH_DEFAULT_LIMIT, - /*use_default_limit*/ true, - ) - .expect("search deferred dynamic tools"); + let tools = + search_output_tools(entries.iter()).expect("mixed search output should serialize"); 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, - })], + vec![ + ToolSearchOutputTool::Namespace(ResponsesApiNamespace { + name: "mcp__calendar__".to_string(), + description: "Tools from the calendar MCP server.".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(None), + )]), + Some(vec!["mode".to_string()]), + Some(false.into()), + ), + output_schema: None, + }), + ], ); } diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 5fe2863cfabc..23f2e8c474ad 100644 --- a/codex-rs/core/tests/suite/search_tool.rs +++ b/codex-rs/core/tests/suite/search_tool.rs @@ -177,7 +177,7 @@ async fn search_tool_enabled_by_default_adds_tool_search() -> 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"], From 108e7b5875299e1031bfa8687d7f527860a6ee2f Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 17 Apr 2026 09:34:00 +0800 Subject: [PATCH 3/7] rm ToolSearchEntry --- .../core/src/tools/handlers/tool_search.rs | 225 ++++++++---------- 1 file changed, 99 insertions(+), 126 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 9f920ae2b782..b3a96818273e 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -22,7 +22,8 @@ const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use"; const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20; pub struct ToolSearchHandler { - entries: Vec, + mcp_tools: Vec, + dynamic_tools: Vec, search_engine: SearchEngine, } @@ -31,34 +32,25 @@ impl ToolSearchHandler { mcp_tools: std::collections::HashMap, dynamic_tools: Vec, ) -> Self { - let mut mcp_entries: Vec = mcp_tools - .into_values() - .map(|info| ToolSearchEntry::Mcp { - name: info.canonical_tool_name().display(), - 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 mcp_tools: Vec = mcp_tools.into_values().collect(); + mcp_tools.sort_by_key(|info| info.canonical_tool_name().display()); - let mut entries = mcp_entries; - entries.extend(dynamic_entries); + let mut dynamic_tools = dynamic_tools; + dynamic_tools.sort_by(|a, b| a.name.cmp(&b.name)); - let documents: Vec> = entries + let documents: Vec> = mcp_tools .iter() + .map(build_mcp_search_text) + .chain(dynamic_tools.iter().map(build_dynamic_search_text)) .enumerate() - .map(|(idx, entry)| Document::new(idx, entry.search_text())) + .map(|(idx, search_text)| Document::new(idx, search_text)) .collect(); let search_engine = SearchEngineBuilder::::with_documents(Language::English, documents).build(); Self { - entries, + mcp_tools, + dynamic_tools, search_engine, } } @@ -101,7 +93,7 @@ impl ToolHandler for ToolSearchHandler { )); } - if self.entries.is_empty() { + if self.mcp_tools.is_empty() && self.dynamic_tools.is_empty() { return Ok(ToolSearchOutput { tools: Vec::new() }); } @@ -118,49 +110,44 @@ impl ToolSearchHandler { limit: usize, use_default_limit: bool, ) -> Result, FunctionCallError> { - let results = self.search_result_entries(query, limit, use_default_limit); - search_output_tools(results) + let result_ids = self.search_result_ids(query, limit, use_default_limit); + self.search_output_tools(result_ids) } - fn search_result_entries( - &self, - query: &str, - limit: usize, - use_default_limit: bool, - ) -> Vec<&ToolSearchEntry> { + fn search_result_ids(&self, query: &str, limit: usize, use_default_limit: bool) -> Vec { 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) .collect::>(); if !use_default_limit { return results; } - if results - .iter() - .any(|entry| entry.mcp_server_name() == Some(COMPUTER_USE_MCP_SERVER_NAME)) - { + if results.iter().any(|&id| { + self.mcp_tools + .get(id) + .is_some_and(|tool| tool.server_name == 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) .collect(); } - limit_results_per_server(results) + limit_results_per_server(&self.mcp_tools, results) } -} -fn search_output_tools<'a>( - results: impl IntoIterator, -) -> Result, FunctionCallError> { - let mut tools = Vec::new(); - // Preserve search order: group MCP tools under namespaces, emit dynamic tools directly. - for entry in results { - match entry { - ToolSearchEntry::Mcp { info, .. } => { + fn search_output_tools( + &self, + result_ids: impl IntoIterator, + ) -> Result, FunctionCallError> { + let mut tools = Vec::new(); + // Preserve search order: group MCP tools under namespaces, emit dynamic tools directly. + for result_id in result_ids { + if let Some(info) = self.mcp_tools.get(result_id) { let tool_name = info.canonical_tool_name(); let namespace = info.callable_namespace.as_str(); let namespace_tool = @@ -182,16 +169,22 @@ fn search_output_tools<'a>( tools: vec![namespace_tool], })); } + continue; } - ToolSearchEntry::Dynamic { tool } => { - tools.push(ToolSearchOutputTool::Function( - dynamic_tool_to_responses_api_tool(tool).map_err(tool_search_output_error)?, - )); - } + + let Some(dynamic_tool_index) = result_id.checked_sub(self.mcp_tools.len()) else { + continue; + }; + let Some(tool) = self.dynamic_tools.get(dynamic_tool_index) else { + continue; + }; + tools.push(ToolSearchOutputTool::Function( + dynamic_tool_to_responses_api_tool(tool).map_err(tool_search_output_error)?, + )); } - } - Ok(tools) + Ok(tools) + } } fn mcp_namespace_description(info: &ToolInfo) -> String { @@ -207,21 +200,22 @@ fn mcp_namespace_description(info: &ToolInfo) -> String { .unwrap_or_else(|| format!("Tools from the {} MCP server.", info.server_name)) } -fn limit_results_per_server(results: Vec<&ToolSearchEntry>) -> Vec<&ToolSearchEntry> { +fn limit_results_per_server(mcp_tools: &[ToolInfo], results: Vec) -> Vec { results .into_iter() .scan( std::collections::HashMap::<&str, usize>::new(), - |counts, entry| { - let Some(server_name) = entry.mcp_server_name() else { - return Some(Some(entry)); + |counts, result_id| { + let Some(tool) = mcp_tools.get(result_id) else { + return Some(Some(result_id)); }; + let server_name = tool.server_name.as_str(); let count = counts.entry(server_name).or_default(); if *count >= default_limit_for_server(server_name) { Some(None) } else { *count += 1; - Some(Some(entry)) + Some(Some(result_id)) } }, ) @@ -237,43 +231,15 @@ fn default_limit_for_server(server_name: &str) -> usize { } } -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 { +fn build_mcp_search_text(info: &ToolInfo) -> String { let mut parts = vec![ - name.to_string(), + info.canonical_tool_name().display(), info.callable_name.clone(), info.tool.name.to_string(), info.server_name.clone(), @@ -354,35 +320,35 @@ mod tests { #[test] fn mixed_search_results_coalesce_mcp_namespaces() { - let entries = [ - ToolSearchEntry::Mcp { - name: "mcp__calendar__create_event".to_string(), - info: Box::new(tool_info("calendar", "create_event", "Create events")), - }, - ToolSearchEntry::Dynamic { - tool: 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, - }, - }, - ToolSearchEntry::Mcp { - name: "mcp__calendar__list_events".to_string(), - info: Box::new(tool_info("calendar", "list_events", "List events")), - }, - ]; + let handler = ToolSearchHandler::new( + 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"), + ), + ]), + 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 tools = - search_output_tools(entries.iter()).expect("mixed search output should serialize"); + let tools = handler + .search_output_tools([0, 2, 1]) + .expect("mixed search output should serialize"); assert_eq!( tools, @@ -448,7 +414,7 @@ mod tests { Vec::new(), ); - let results = handler.search_result_entries( + let results = handler.search_result_ids( "computer use", TOOL_SEARCH_DEFAULT_LIMIT, /*use_default_limit*/ true, @@ -458,10 +424,10 @@ mod tests { assert!( results .iter() - .all(|entry| entry.mcp_server_name() == Some(COMPUTER_USE_MCP_SERVER_NAME)) + .all(|&id| handler.mcp_tools[id].server_name == COMPUTER_USE_MCP_SERVER_NAME) ); - let explicit_results = handler.search_result_entries( + let explicit_results = handler.search_result_ids( "computer use", /*limit*/ 100, /*use_default_limit*/ false, @@ -484,7 +450,7 @@ mod tests { )); let handler = ToolSearchHandler::new(tools, Vec::new()); - let results = handler.search_result_entries( + let results = handler.search_result_ids( "calendar", TOOL_SEARCH_DEFAULT_LIMIT, /*use_default_limit*/ true, @@ -494,10 +460,10 @@ mod tests { assert!( results .iter() - .all(|entry| entry.mcp_server_name() == Some("other-server")) + .all(|&id| handler.mcp_tools[id].server_name == "other-server") ); - let explicit_results = handler.search_result_entries( + let explicit_results = handler.search_result_ids( "calendar", /*limit*/ 100, /*use_default_limit*/ false, ); @@ -518,17 +484,20 @@ mod tests { )); let handler = ToolSearchHandler::new(tools, Vec::new()); - let results = handler.search_result_entries( + let results = handler.search_result_ids( "computer use", TOOL_SEARCH_DEFAULT_LIMIT, /*use_default_limit*/ true, ); assert!( - count_results_for_server(&results, COMPUTER_USE_MCP_SERVER_NAME) + count_results_for_server(&handler, &results, COMPUTER_USE_MCP_SERVER_NAME) <= COMPUTER_USE_TOOL_SEARCH_LIMIT ); - assert!(count_results_for_server(&results, "other-server") <= TOOL_SEARCH_DEFAULT_LIMIT); + assert!( + count_results_for_server(&handler, &results, "other-server") + <= TOOL_SEARCH_DEFAULT_LIMIT + ); } fn numbered_tools( @@ -575,10 +544,14 @@ mod tests { } } - fn count_results_for_server(results: &[&ToolSearchEntry], server_name: &str) -> usize { + fn count_results_for_server( + handler: &ToolSearchHandler, + results: &[usize], + server_name: &str, + ) -> usize { results .iter() - .filter(|entry| entry.mcp_server_name() == Some(server_name)) + .filter(|&&id| handler.mcp_tools[id].server_name == server_name) .count() } } From 46cc87f951db838c1f1fb64fabeea2b9c80c9784 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 17 Apr 2026 12:21:33 +0800 Subject: [PATCH 4/7] Clean up tool search entries --- .../core/src/tools/handlers/tool_search.rs | 292 ++++++------------ codex-rs/core/src/tools/mod.rs | 1 + codex-rs/core/src/tools/spec.rs | 10 +- codex-rs/core/src/tools/tool_search_entry.rs | 169 ++++++++++ 4 files changed, 273 insertions(+), 199 deletions(-) create mode 100644 codex-rs/core/src/tools/tool_search_entry.rs diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index b3a96818273e..c1b9177530e5 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -4,44 +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 crate::tools::tool_search_entry::ToolSearchEntryOutput; use bm25::Document; use bm25::Language; use bm25::SearchEngine; use bm25::SearchEngineBuilder; -use codex_mcp::ToolInfo; -use codex_protocol::dynamic_tools::DynamicToolSpec; use codex_tools::ResponsesApiNamespace; -use codex_tools::ResponsesApiNamespaceTool; use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT; use codex_tools::TOOL_SEARCH_TOOL_NAME; use codex_tools::ToolSearchOutputTool; -use codex_tools::dynamic_tool_to_responses_api_tool; -use codex_tools::mcp_tool_to_deferred_responses_api_tool; +use std::collections::HashMap; const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use"; const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20; pub struct ToolSearchHandler { - mcp_tools: Vec, - dynamic_tools: Vec, + entries: Vec, search_engine: SearchEngine, } impl ToolSearchHandler { - pub fn new( - mcp_tools: std::collections::HashMap, - dynamic_tools: Vec, - ) -> Self { - let mut mcp_tools: Vec = mcp_tools.into_values().collect(); - mcp_tools.sort_by_key(|info| info.canonical_tool_name().display()); - - let mut dynamic_tools = dynamic_tools; - dynamic_tools.sort_by(|a, b| a.name.cmp(&b.name)); - - let documents: Vec> = mcp_tools + pub(crate) fn new(entries: Vec) -> Self { + let documents: Vec> = entries .iter() - .map(build_mcp_search_text) - .chain(dynamic_tools.iter().map(build_dynamic_search_text)) + .map(|entry| entry.search_text.clone()) .enumerate() .map(|(idx, search_text)| Document::new(idx, search_text)) .collect(); @@ -49,8 +36,7 @@ impl ToolSearchHandler { SearchEngineBuilder::::with_documents(Language::English, documents).build(); Self { - mcp_tools, - dynamic_tools, + entries, search_engine, } } @@ -93,7 +79,7 @@ impl ToolHandler for ToolSearchHandler { )); } - if self.mcp_tools.is_empty() && self.dynamic_tools.is_empty() { + if self.entries.is_empty() { return Ok(ToolSearchOutput { tools: Vec::new() }); } @@ -126,9 +112,10 @@ impl ToolSearchHandler { } if results.iter().any(|&id| { - self.mcp_tools + self.entries .get(id) - .is_some_and(|tool| tool.server_name == COMPUTER_USE_MCP_SERVER_NAME) + .and_then(|entry| entry.limit_bucket.as_deref()) + .is_some_and(|bucket| bucket == COMPUTER_USE_MCP_SERVER_NAME) }) { results = self .search_engine @@ -137,7 +124,7 @@ impl ToolSearchHandler { .map(|result| result.document.id) .collect(); } - limit_results_per_server(&self.mcp_tools, results) + limit_results_by_bucket(&self.entries, results) } fn search_output_tools( @@ -145,172 +132,80 @@ impl ToolSearchHandler { result_ids: impl IntoIterator, ) -> Result, FunctionCallError> { let mut tools = Vec::new(); - // Preserve search order: group MCP tools under namespaces, emit dynamic tools directly. + // Preserve search order: group namespace children, emit standalone tools directly. for result_id in result_ids { - if let Some(info) = self.mcp_tools.get(result_id) { - let tool_name = info.canonical_tool_name(); - let namespace = info.callable_namespace.as_str(); - let namespace_tool = - mcp_tool_to_deferred_responses_api_tool(&tool_name, &info.tool) - .map(ResponsesApiNamespaceTool::Function) - .map_err(tool_search_output_error)?; - - if let Some(output) = tools.iter_mut().find_map(|tool| match tool { - ToolSearchOutputTool::Namespace(output) if output.name == namespace => { - Some(output) + let Some(entry) = self.entries.get(result_id) else { + continue; + }; + match &entry.output { + ToolSearchEntryOutput::Function(tool) => { + tools.push(ToolSearchOutputTool::Function(tool.clone())); + } + ToolSearchEntryOutput::NamespacedFunction { + namespace, + namespace_description, + tool: namespace_tool, + } => { + if let Some(output) = tools.iter_mut().find_map(|tool| match tool { + ToolSearchOutputTool::Namespace(output) if output.name == *namespace => { + Some(output) + } + ToolSearchOutputTool::Namespace(_) | ToolSearchOutputTool::Function(_) => { + None + } + }) { + output.tools.push(namespace_tool.clone()); + } else { + tools.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { + name: namespace.clone(), + description: namespace_description.clone(), + tools: vec![namespace_tool.clone()], + })); } - ToolSearchOutputTool::Namespace(_) | ToolSearchOutputTool::Function(_) => None, - }) { - output.tools.push(namespace_tool); - } else { - tools.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: namespace.to_string(), - description: mcp_namespace_description(info), - tools: vec![namespace_tool], - })); } - continue; } - - let Some(dynamic_tool_index) = result_id.checked_sub(self.mcp_tools.len()) else { - continue; - }; - let Some(tool) = self.dynamic_tools.get(dynamic_tool_index) else { - continue; - }; - tools.push(ToolSearchOutputTool::Function( - dynamic_tool_to_responses_api_tool(tool).map_err(tool_search_output_error)?, - )); } Ok(tools) } } -fn mcp_namespace_description(info: &ToolInfo) -> String { - info.connector_description - .clone() - .or_else(|| { - info.connector_name - .as_deref() - .map(str::trim) - .filter(|connector_name| !connector_name.is_empty()) - .map(|connector_name| format!("Tools for working with {connector_name}.")) - }) - .unwrap_or_else(|| format!("Tools from the {} MCP server.", info.server_name)) -} - -fn limit_results_per_server(mcp_tools: &[ToolInfo], results: Vec) -> Vec { +fn limit_results_by_bucket(entries: &[ToolSearchEntry], results: Vec) -> Vec { results .into_iter() - .scan( - std::collections::HashMap::<&str, usize>::new(), - |counts, result_id| { - let Some(tool) = mcp_tools.get(result_id) else { - return Some(Some(result_id)); - }; - let server_name = tool.server_name.as_str(); - let count = counts.entry(server_name).or_default(); - if *count >= default_limit_for_server(server_name) { - Some(None) - } else { - *count += 1; - Some(Some(result_id)) - } - }, - ) + .scan(HashMap::<&str, usize>::new(), |counts, result_id| { + let Some(bucket) = entries + .get(result_id) + .and_then(|entry| entry.limit_bucket.as_deref()) + else { + return Some(Some(result_id)); + }; + let count = counts.entry(bucket).or_default(); + if *count >= default_limit_for_bucket(bucket) { + Some(None) + } else { + *count += 1; + Some(Some(result_id)) + } + }) .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 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(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(" ") -} - #[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; @@ -320,8 +215,21 @@ mod tests { #[test] fn mixed_search_results_coalesce_mcp_namespaces() { - let handler = ToolSearchHandler::new( - std::collections::HashMap::from([ + 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"), @@ -330,20 +238,8 @@ mod tests { "mcp__calendar__list_events".to_string(), tool_info("calendar", "list_events", "List events"), ), - ]), - 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, - }], + ])), + &dynamic_tools, ); let tools = handler @@ -405,14 +301,12 @@ mod tests { #[test] fn computer_use_tool_search_uses_larger_limit() { - let handler = ToolSearchHandler::new( - numbered_tools( - COMPUTER_USE_MCP_SERVER_NAME, - "computer use", - /*count*/ 100, - ), - Vec::new(), + 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_ids( "computer use", @@ -424,7 +318,8 @@ mod tests { assert!( results .iter() - .all(|&id| handler.mcp_tools[id].server_name == COMPUTER_USE_MCP_SERVER_NAME) + .all(|&id| handler.entries[id].limit_bucket.as_deref() + == Some(COMPUTER_USE_MCP_SERVER_NAME)) ); let explicit_results = handler.search_result_ids( @@ -448,7 +343,7 @@ mod tests { "calendar", /*count*/ 100, )); - let handler = ToolSearchHandler::new(tools, Vec::new()); + let handler = handler_from_tools(Some(&tools), &[]); let results = handler.search_result_ids( "calendar", @@ -460,7 +355,7 @@ mod tests { assert!( results .iter() - .all(|&id| handler.mcp_tools[id].server_name == "other-server") + .all(|&id| handler.entries[id].limit_bucket.as_deref() == Some("other-server")) ); let explicit_results = handler.search_result_ids( @@ -482,7 +377,7 @@ mod tests { "computer use", /*count*/ 100, )); - let handler = ToolSearchHandler::new(tools, Vec::new()); + let handler = handler_from_tools(Some(&tools), &[]); let results = handler.search_result_ids( "computer use", @@ -551,7 +446,14 @@ mod tests { ) -> usize { results .iter() - .filter(|&&id| handler.mcp_tools[id].server_name == server_name) + .filter(|&&id| handler.entries[id].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 06c20d6922a1..ee790d7d2ac9 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 de0184b7de53..41dc422e3d0b 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); @@ -264,10 +265,11 @@ pub(crate) fn build_specs_with_discoverable_tools( } ToolHandlerKind::ToolSearch => { if tool_search_handler.is_none() { - tool_search_handler = Some(Arc::new(ToolSearchHandler::new( - deferred_mcp_tools.clone().unwrap_or_default(), - deferred_dynamic_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 000000000000..77771ec8dc0e --- /dev/null +++ b/codex-rs/core/src/tools/tool_search_entry.rs @@ -0,0 +1,169 @@ +use codex_mcp::ToolInfo; +use codex_protocol::dynamic_tools::DynamicToolSpec; +use codex_tools::ResponsesApiNamespaceTool; +use codex_tools::ResponsesApiTool; +use codex_tools::dynamic_tool_to_responses_api_tool; +use codex_tools::mcp_tool_to_deferred_responses_api_tool; +use std::collections::HashMap; + +#[derive(Clone)] +pub(crate) struct ToolSearchEntry { + pub(crate) search_text: String, + pub(crate) output: ToolSearchEntryOutput, + pub(crate) limit_bucket: Option, +} + +#[derive(Clone)] +pub(crate) enum ToolSearchEntryOutput { + Function(ResponsesApiTool), + NamespacedFunction { + namespace: String, + namespace_description: String, + tool: ResponsesApiNamespaceTool, + }, +} + +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 { + let tool_name = info.canonical_tool_name(); + let tool = mcp_tool_to_deferred_responses_api_tool(&tool_name, &info.tool) + .map(ResponsesApiNamespaceTool::Function)?; + let namespace_description = info + .connector_description + .clone() + .or_else(|| { + info.connector_name + .as_deref() + .map(str::trim) + .filter(|connector_name| !connector_name.is_empty()) + .map(|connector_name| format!("Tools for working with {connector_name}.")) + }) + .unwrap_or_else(|| format!("Tools from the {} MCP server.", info.server_name)); + + Ok(ToolSearchEntry { + search_text: build_mcp_search_text(info), + output: ToolSearchEntryOutput::NamespacedFunction { + namespace: info.callable_namespace.clone(), + namespace_description, + tool, + }, + 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: ToolSearchEntryOutput::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(" ") +} From c955ab8c106e1425c81b7289ec6bb7bc0226ec9a Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 17 Apr 2026 18:55:11 +0800 Subject: [PATCH 5/7] reuse tool search logic --- .../core/src/tools/handlers/tool_search.rs | 26 +++------ codex-rs/core/src/tools/tool_search_entry.rs | 48 +++++----------- codex-rs/tools/src/lib.rs | 1 + codex-rs/tools/src/tool_discovery.rs | 56 ++++++++++++------- 4 files changed, 59 insertions(+), 72 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index c1b9177530e5..85613029e5e9 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -5,12 +5,10 @@ use crate::tools::context::ToolSearchOutput; use crate::tools::registry::ToolHandler; use crate::tools::registry::ToolKind; use crate::tools::tool_search_entry::ToolSearchEntry; -use crate::tools::tool_search_entry::ToolSearchEntryOutput; use bm25::Document; use bm25::Language; use bm25::SearchEngine; use bm25::SearchEngineBuilder; -use codex_tools::ResponsesApiNamespace; use codex_tools::TOOL_SEARCH_DEFAULT_LIMIT; use codex_tools::TOOL_SEARCH_TOOL_NAME; use codex_tools::ToolSearchOutputTool; @@ -138,29 +136,23 @@ impl ToolSearchHandler { continue; }; match &entry.output { - ToolSearchEntryOutput::Function(tool) => { + ToolSearchOutputTool::Function(tool) => { tools.push(ToolSearchOutputTool::Function(tool.clone())); } - ToolSearchEntryOutput::NamespacedFunction { - namespace, - namespace_description, - tool: namespace_tool, - } => { + ToolSearchOutputTool::Namespace(namespace) => { if let Some(output) = tools.iter_mut().find_map(|tool| match tool { - ToolSearchOutputTool::Namespace(output) if output.name == *namespace => { + ToolSearchOutputTool::Namespace(output) + if output.name == namespace.name => + { Some(output) } ToolSearchOutputTool::Namespace(_) | ToolSearchOutputTool::Function(_) => { None } }) { - output.tools.push(namespace_tool.clone()); + output.tools.extend(namespace.tools.clone()); } else { - tools.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: namespace.clone(), - description: namespace_description.clone(), - tools: vec![namespace_tool.clone()], - })); + tools.push(ToolSearchOutputTool::Namespace(namespace.clone())); } } } @@ -251,7 +243,7 @@ mod tests { vec![ ToolSearchOutputTool::Namespace(ResponsesApiNamespace { name: "mcp__calendar__".to_string(), - description: "Tools from the calendar MCP server.".to_string(), + description: "Tools in the mcp__calendar__ namespace.".to_string(), tools: vec![ ResponsesApiNamespaceTool::Function(ResponsesApiTool { name: "create_event".to_string(), @@ -288,7 +280,7 @@ mod tests { parameters: codex_tools::JsonSchema::object( std::collections::BTreeMap::from([( "mode".to_string(), - codex_tools::JsonSchema::string(None), + codex_tools::JsonSchema::string(/*description*/ None), )]), Some(vec!["mode".to_string()]), Some(false.into()), diff --git a/codex-rs/core/src/tools/tool_search_entry.rs b/codex-rs/core/src/tools/tool_search_entry.rs index 77771ec8dc0e..6fa493132a71 100644 --- a/codex-rs/core/src/tools/tool_search_entry.rs +++ b/codex-rs/core/src/tools/tool_search_entry.rs @@ -1,28 +1,18 @@ use codex_mcp::ToolInfo; use codex_protocol::dynamic_tools::DynamicToolSpec; -use codex_tools::ResponsesApiNamespaceTool; -use codex_tools::ResponsesApiTool; +use codex_tools::ToolSearchOutputTool; +use codex_tools::ToolSearchResultSource; use codex_tools::dynamic_tool_to_responses_api_tool; -use codex_tools::mcp_tool_to_deferred_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: ToolSearchEntryOutput, + pub(crate) output: ToolSearchOutputTool, pub(crate) limit_bucket: Option, } -#[derive(Clone)] -pub(crate) enum ToolSearchEntryOutput { - Function(ResponsesApiTool), - NamespacedFunction { - namespace: String, - namespace_description: String, - tool: ResponsesApiNamespaceTool, - }, -} - pub(crate) fn build_tool_search_entries( mcp_tools: Option<&HashMap>, dynamic_tools: &[DynamicToolSpec], @@ -63,28 +53,16 @@ pub(crate) fn build_tool_search_entries( } fn mcp_tool_search_entry(info: &ToolInfo) -> Result { - let tool_name = info.canonical_tool_name(); - let tool = mcp_tool_to_deferred_responses_api_tool(&tool_name, &info.tool) - .map(ResponsesApiNamespaceTool::Function)?; - let namespace_description = info - .connector_description - .clone() - .or_else(|| { - info.connector_name - .as_deref() - .map(str::trim) - .filter(|connector_name| !connector_name.is_empty()) - .map(|connector_name| format!("Tools for working with {connector_name}.")) - }) - .unwrap_or_else(|| format!("Tools from the {} MCP server.", info.server_name)); - Ok(ToolSearchEntry { search_text: build_mcp_search_text(info), - output: ToolSearchEntryOutput::NamespacedFunction { - namespace: info.callable_namespace.clone(), - namespace_description, - tool, - }, + 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()), }) } @@ -92,7 +70,7 @@ fn mcp_tool_search_entry(info: &ToolInfo) -> Result Result { Ok(ToolSearchEntry { search_text: build_dynamic_search_text(tool), - output: ToolSearchEntryOutput::Function(dynamic_tool_to_responses_api_tool(tool)?), + output: ToolSearchOutputTool::Function(dynamic_tool_to_responses_api_tool(tool)?), limit_bucket: None, }) } diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index b7c6a57d2b03..d2e460816c3b 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -115,6 +115,7 @@ 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 272d2576c76c..518a29581f37 100644 --- a/codex-rs/tools/src/tool_discovery.rs +++ b/codex-rs/tools/src/tool_discovery.rs @@ -224,32 +224,14 @@ pub fn collect_tool_search_output_tools<'a>( continue; }; - 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) - }) + .map(|tool| tool_search_result_source_to_namespace_tool(*tool)) .collect::, _>>()?; results.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { name: tool_namespace.to_string(), - description: description - .unwrap_or_else(|| default_namespace_description(tool_namespace)), + description: tool_search_result_source_namespace_description(*first_tool), tools, })); } @@ -257,6 +239,40 @@ pub fn collect_tool_search_output_tools<'a>( Ok(results) } +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)?], + })) +} + +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)) +} + +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>( searchable_tools: impl IntoIterator>, ) -> Vec { From 4cc404087d309177a933334a68b14eea54f00af6 Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Fri, 17 Apr 2026 20:41:55 +0800 Subject: [PATCH 6/7] move to test_codex integration test rather than app-server --- .../tests/suite/v2/dynamic_tools.rs | 220 ------------------ codex-rs/core/tests/suite/search_tool.rs | 189 ++++++++++++++- 2 files changed, 183 insertions(+), 226 deletions(-) diff --git a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs index 115c7e0de63c..0a3315a07a29 100644 --- a/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs +++ b/codex-rs/app-server/tests/suite/v2/dynamic_tools.rs @@ -4,7 +4,6 @@ use app_test_support::McpProcess; use app_test_support::create_final_assistant_message_sse_response; use app_test_support::create_mock_responses_server_sequence_unchecked; use app_test_support::to_response; -use app_test_support::write_models_cache_with_models; use codex_app_server_protocol::DynamicToolCallOutputContentItem; use codex_app_server_protocol::DynamicToolCallParams; use codex_app_server_protocol::DynamicToolCallResponse; @@ -22,7 +21,6 @@ use codex_app_server_protocol::ThreadStartResponse; use codex_app_server_protocol::TurnStartParams; use codex_app_server_protocol::TurnStartResponse; use codex_app_server_protocol::UserInput as V2UserInput; -use codex_models_manager::bundled_models_response; use codex_protocol::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -30,7 +28,6 @@ use core_test_support::responses; use pretty_assertions::assert_eq; use serde_json::Value; use serde_json::json; -use std::io::Write; use std::path::Path; use std::time::Duration; use tempfile::TempDir; @@ -199,178 +196,6 @@ async fn thread_start_keeps_hidden_dynamic_tools_out_of_model_requests() -> Resu Ok(()) } -/// Exercises deferred dynamic tool discovery, the follow-up tool call, and the tool response. -#[tokio::test] -async fn deferred_dynamic_tool_can_be_discovered_and_called_through_tool_search() -> Result<()> { - 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 responses = vec![ - responses::sse(vec![ - responses::ev_response_created("resp-1"), - responses::ev_tool_search_call( - search_call_id, - &json!({ - "query": "recurring automations", - "limit": 8, - }), - ), - responses::ev_completed("resp-1"), - ]), - responses::sse(vec![ - responses::ev_response_created("resp-2"), - responses::ev_function_call(dynamic_call_id, tool_name, &tool_call_arguments), - responses::ev_completed("resp-2"), - ]), - create_final_assistant_message_sse_response("Done")?, - ]; - let server = create_mock_responses_server_sequence_unchecked(responses).await; - - let codex_home = TempDir::new()?; - write_search_capable_models_cache(codex_home.path())?; - create_config_toml(codex_home.path(), &server.uri())?; - enable_tool_search_feature(codex_home.path())?; - - let mut mcp = McpProcess::new(codex_home.path()).await?; - timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).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 thread_req = mcp - .send_thread_start_request(ThreadStartParams { - dynamic_tools: Some(vec![dynamic_tool]), - ..Default::default() - }) - .await?; - let thread_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(thread_req)), - ) - .await??; - let ThreadStartResponse { thread, .. } = to_response::(thread_resp)?; - let thread_id = thread.id.clone(); - - let turn_req = mcp - .send_turn_start_request(TurnStartParams { - thread_id: thread_id.clone(), - input: vec![V2UserInput::Text { - text: "Use the automation tool".to_string(), - text_elements: Vec::new(), - }], - ..Default::default() - }) - .await?; - let turn_resp: JSONRPCResponse = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_response_message(RequestId::Integer(turn_req)), - ) - .await??; - let TurnStartResponse { turn } = to_response::(turn_resp)?; - let turn_id = turn.id.clone(); - - let started = wait_for_dynamic_tool_started(&mut mcp, dynamic_call_id).await?; - assert_eq!(started.thread_id, thread_id.clone()); - assert_eq!(started.turn_id, turn_id.clone()); - - let request = timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_request_message(), - ) - .await??; - let (request_id, params) = match request { - ServerRequest::DynamicToolCall { request_id, params } => (request_id, params), - other => panic!("expected DynamicToolCall request, got {other:?}"), - }; - assert_eq!( - params, - DynamicToolCallParams { - thread_id: thread_id.clone(), - turn_id: turn_id.clone(), - call_id: dynamic_call_id.to_string(), - tool: tool_name.to_string(), - arguments: tool_args.clone(), - } - ); - - mcp.send_response( - request_id, - serde_json::to_value(DynamicToolCallResponse { - content_items: vec![DynamicToolCallOutputContentItem::InputText { - text: "dynamic-search-ok".to_string(), - }], - success: true, - })?, - ) - .await?; - - let completed = wait_for_dynamic_tool_completed(&mut mcp, dynamic_call_id).await?; - assert_eq!(completed.thread_id, thread_id); - assert_eq!(completed.turn_id, turn_id); - - timeout( - DEFAULT_READ_TIMEOUT, - mcp.read_stream_until_notification_message("turn/completed"), - ) - .await??; - - let bodies = responses_bodies(&server).await?; - let first_request = bodies - .first() - .context("expected an initial responses request")?; - assert!( - find_tool_by_type(first_request, "tool_search").is_some(), - "initial request should advertise tool_search: {first_request:?}" - ); - assert!( - find_tool(first_request, tool_name).is_none(), - "deferred dynamic tool should not be directly advertised before search" - ); - - let search_tools = bodies - .iter() - .find_map(|body| tool_search_output_tools(body, search_call_id)) - .context("expected tool_search_output in follow-up request")?; - assert_eq!( - search_tools, - vec![json!({ - "type": "function", - "name": tool_name, - "description": tool_description, - "strict": false, - "defer_loading": true, - "parameters": input_schema, - })] - ); - - let payload = bodies - .iter() - .find_map(|body| function_call_output_payload(body, dynamic_call_id)) - .context("expected function_call_output in post-tool follow-up request")?; - assert_eq!( - payload, - FunctionCallOutputPayload::from_text("dynamic-search-ok".to_string()) - ); - - Ok(()) -} - /// Exercises the full dynamic tool call path (server request, client response, model output). #[tokio::test] async fn dynamic_tool_call_round_trip_sends_text_content_items_to_model() -> Result<()> { @@ -758,30 +583,6 @@ fn find_tool<'a>(body: &'a Value, name: &str) -> Option<&'a Value> { }) } -fn find_tool_by_type<'a>(body: &'a Value, tool_type: &str) -> Option<&'a Value> { - body.get("tools") - .and_then(Value::as_array) - .and_then(|tools| { - tools - .iter() - .find(|tool| tool.get("type").and_then(Value::as_str) == Some(tool_type)) - }) -} - -fn tool_search_output_tools(body: &Value, call_id: &str) -> Option> { - body.get("input") - .and_then(Value::as_array) - .and_then(|items| { - items.iter().find(|item| { - item.get("type").and_then(Value::as_str) == Some("tool_search_output") - && item.get("call_id").and_then(Value::as_str) == Some(call_id) - }) - }) - .and_then(|item| item.get("tools")) - .and_then(Value::as_array) - .cloned() -} - fn function_call_output_payload(body: &Value, call_id: &str) -> Option { function_call_output_raw_output(body, call_id) .and_then(|output| serde_json::from_value(output).ok()) @@ -862,24 +663,3 @@ stream_max_retries = 0 ), ) } - -fn enable_tool_search_feature(codex_home: &Path) -> std::io::Result<()> { - let mut config_toml = std::fs::OpenOptions::new() - .append(true) - .open(codex_home.join("config.toml"))?; - config_toml.write_all(b"\n[features]\ntool_search = true\n") -} - -fn write_search_capable_models_cache(codex_home: &Path) -> Result<()> { - let mut model = bundled_models_response() - .context("bundled models should parse")? - .models - .into_iter() - .find(|model| model.slug == "gpt-5.4") - .context("expected bundled gpt-5.4 model")?; - model.slug = "mock-model".to_string(); - model.display_name = "mock-model".to_string(); - model.supports_search_tool = true; - write_models_cache_with_models(codex_home, vec![model])?; - Ok(()) -} diff --git a/codex-rs/core/tests/suite/search_tool.rs b/codex-rs/core/tests/suite/search_tool.rs index 23f2e8c474ad..2e9c8aa541ff 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; @@ -94,12 +98,7 @@ fn tool_search_output_tools(request: &ResponsesRequest, call_id: &str) -> Vec 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(())); From 1aacfe4661690673be802def49c525551d248ece Mon Sep 17 00:00:00 2001 From: Sayan Sisodiya Date: Sat, 18 Apr 2026 01:39:58 +0800 Subject: [PATCH 7/7] pass around ToolSearchEntry objects and cleanup dead fn --- .../core/src/tools/handlers/tool_search.rs | 82 ++++---- codex-rs/tools/src/lib.rs | 1 - codex-rs/tools/src/tool_discovery.rs | 36 ---- codex-rs/tools/src/tool_discovery_tests.rs | 190 ------------------ 4 files changed, 40 insertions(+), 269 deletions(-) diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 85613029e5e9..7103023d9d5f 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -94,25 +94,31 @@ impl ToolSearchHandler { limit: usize, use_default_limit: bool, ) -> Result, FunctionCallError> { - let result_ids = self.search_result_ids(query, limit, use_default_limit); - self.search_output_tools(result_ids) + let results = self.search_result_entries(query, limit, use_default_limit); + self.search_output_tools(results) } - fn search_result_ids(&self, query: &str, limit: usize, use_default_limit: bool) -> Vec { + fn search_result_entries( + &self, + query: &str, + limit: usize, + use_default_limit: bool, + ) -> Vec<&ToolSearchEntry> { let mut results = self .search_engine .search(query, limit) .into_iter() .map(|result| result.document.id) + .filter_map(|id| self.entries.get(id)) .collect::>(); if !use_default_limit { return results; } - if results.iter().any(|&id| { - self.entries - .get(id) - .and_then(|entry| entry.limit_bucket.as_deref()) + if results.iter().any(|entry| { + entry + .limit_bucket + .as_deref() .is_some_and(|bucket| bucket == COMPUTER_USE_MCP_SERVER_NAME) }) { results = self @@ -120,21 +126,19 @@ impl ToolSearchHandler { .search(query, COMPUTER_USE_TOOL_SEARCH_LIMIT) .into_iter() .map(|result| result.document.id) + .filter_map(|id| self.entries.get(id)) .collect(); } - limit_results_by_bucket(&self.entries, results) + limit_results_by_bucket(results) } - fn search_output_tools( + fn search_output_tools<'a>( &self, - result_ids: impl IntoIterator, + results: impl IntoIterator, ) -> Result, FunctionCallError> { let mut tools = Vec::new(); // Preserve search order: group namespace children, emit standalone tools directly. - for result_id in result_ids { - let Some(entry) = self.entries.get(result_id) else { - continue; - }; + for entry in results { match &entry.output { ToolSearchOutputTool::Function(tool) => { tools.push(ToolSearchOutputTool::Function(tool.clone())); @@ -162,22 +166,19 @@ impl ToolSearchHandler { } } -fn limit_results_by_bucket(entries: &[ToolSearchEntry], results: Vec) -> Vec { +fn limit_results_by_bucket(results: Vec<&ToolSearchEntry>) -> Vec<&ToolSearchEntry> { results .into_iter() - .scan(HashMap::<&str, usize>::new(), |counts, result_id| { - let Some(bucket) = entries - .get(result_id) - .and_then(|entry| entry.limit_bucket.as_deref()) - else { - return Some(Some(result_id)); + .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_id)) + Some(Some(result)) } }) .flatten() @@ -233,9 +234,14 @@ mod tests { ])), &dynamic_tools, ); + let results = [ + &handler.entries[0], + &handler.entries[2], + &handler.entries[1], + ]; let tools = handler - .search_output_tools([0, 2, 1]) + .search_output_tools(results) .expect("mixed search output should serialize"); assert_eq!( @@ -300,7 +306,7 @@ mod tests { ); let handler = handler_from_tools(Some(&tools), &[]); - let results = handler.search_result_ids( + let results = handler.search_result_entries( "computer use", TOOL_SEARCH_DEFAULT_LIMIT, /*use_default_limit*/ true, @@ -310,11 +316,10 @@ mod tests { assert!( results .iter() - .all(|&id| handler.entries[id].limit_bucket.as_deref() - == Some(COMPUTER_USE_MCP_SERVER_NAME)) + .all(|entry| entry.limit_bucket.as_deref() == Some(COMPUTER_USE_MCP_SERVER_NAME)) ); - let explicit_results = handler.search_result_ids( + let explicit_results = handler.search_result_entries( "computer use", /*limit*/ 100, /*use_default_limit*/ false, @@ -337,7 +342,7 @@ mod tests { )); let handler = handler_from_tools(Some(&tools), &[]); - let results = handler.search_result_ids( + let results = handler.search_result_entries( "calendar", TOOL_SEARCH_DEFAULT_LIMIT, /*use_default_limit*/ true, @@ -347,10 +352,10 @@ mod tests { assert!( results .iter() - .all(|&id| handler.entries[id].limit_bucket.as_deref() == Some("other-server")) + .all(|entry| entry.limit_bucket.as_deref() == Some("other-server")) ); - let explicit_results = handler.search_result_ids( + let explicit_results = handler.search_result_entries( "calendar", /*limit*/ 100, /*use_default_limit*/ false, ); @@ -371,20 +376,17 @@ mod tests { )); let handler = handler_from_tools(Some(&tools), &[]); - let results = handler.search_result_ids( + let results = handler.search_result_entries( "computer use", TOOL_SEARCH_DEFAULT_LIMIT, /*use_default_limit*/ true, ); assert!( - count_results_for_server(&handler, &results, COMPUTER_USE_MCP_SERVER_NAME) + count_results_for_server(&results, COMPUTER_USE_MCP_SERVER_NAME) <= COMPUTER_USE_TOOL_SEARCH_LIMIT ); - assert!( - count_results_for_server(&handler, &results, "other-server") - <= TOOL_SEARCH_DEFAULT_LIMIT - ); + assert!(count_results_for_server(&results, "other-server") <= TOOL_SEARCH_DEFAULT_LIMIT); } fn numbered_tools( @@ -431,14 +433,10 @@ mod tests { } } - fn count_results_for_server( - handler: &ToolSearchHandler, - results: &[usize], - server_name: &str, - ) -> usize { + fn count_results_for_server(results: &[&ToolSearchEntry], server_name: &str) -> usize { results .iter() - .filter(|&&id| handler.entries[id].limit_bucket.as_deref() == Some(server_name)) + .filter(|entry| entry.limit_bucket.as_deref() == Some(server_name)) .count() } diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index d2e460816c3b..07e98ab5a9d3 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -109,7 +109,6 @@ 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; diff --git a/codex-rs/tools/src/tool_discovery.rs b/codex-rs/tools/src/tool_discovery.rs index 518a29581f37..359b1d792f99 100644 --- a/codex-rs/tools/src/tool_discovery.rs +++ b/codex-rs/tools/src/tool_discovery.rs @@ -203,42 +203,6 @@ 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])); - } - } - - let mut results = Vec::with_capacity(grouped.len()); - for (tool_namespace, tools) in grouped { - let Some(first_tool) = tools.first() else { - continue; - }; - - let tools = tools - .iter() - .map(|tool| tool_search_result_source_to_namespace_tool(*tool)) - .collect::, _>>()?; - - results.push(ToolSearchOutputTool::Namespace(ResponsesApiNamespace { - name: tool_namespace.to_string(), - description: tool_search_result_source_namespace_description(*first_tool), - tools, - })); - } - - Ok(results) -} - pub fn tool_search_result_source_to_output_tool( source: ToolSearchResultSource<'_>, ) -> Result { diff --git a/codex-rs/tools/src/tool_discovery_tests.rs b/codex-rs/tools/src/tool_discovery_tests.rs index bc0f84b0b484..9afd32231e1d 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() { @@ -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!(