diff --git a/codex-rs/core/src/tools/handlers/tool_search.rs b/codex-rs/core/src/tools/handlers/tool_search.rs index 1c715e43f..4bab746ac 100644 --- a/codex-rs/core/src/tools/handlers/tool_search.rs +++ b/codex-rs/core/src/tools/handlers/tool_search.rs @@ -14,6 +14,9 @@ use codex_tools::TOOL_SEARCH_TOOL_NAME; use codex_tools::ToolSearchResultSource; use codex_tools::collect_tool_search_output_tools; +const COMPUTER_USE_MCP_SERVER_NAME: &str = "computer-use"; +const COMPUTER_USE_TOOL_SEARCH_LIMIT: usize = 20; + pub struct ToolSearchHandler { entries: Vec<(String, ToolInfo)>, search_engine: SearchEngine, @@ -67,7 +70,8 @@ impl ToolHandler for ToolSearchHandler { "query must not be empty".to_string(), )); } - let limit = args.limit.unwrap_or(TOOL_SEARCH_DEFAULT_LIMIT); + let requested_limit = args.limit; + let limit = requested_limit.unwrap_or(TOOL_SEARCH_DEFAULT_LIMIT); if limit == 0 { return Err(FunctionCallError::RespondToModel( @@ -79,21 +83,18 @@ impl ToolHandler for ToolSearchHandler { return Ok(ToolSearchOutput { tools: Vec::new() }); } - let results = self.search_engine.search(query, limit); + let results = self.search_result_entries(query, limit, requested_limit.is_none()); - let tools = collect_tool_search_output_tools( - results - .into_iter() - .filter_map(|result| self.entries.get(result.document.id)) - .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(), - }), - ) + 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}" @@ -104,6 +105,66 @@ impl ToolHandler for ToolSearchHandler { } } +impl ToolSearchHandler { + fn search_result_entries( + &self, + query: &str, + limit: usize, + use_default_limit: bool, + ) -> Vec<&(String, ToolInfo)> { + let mut results = self + .search_engine + .search(query, limit) + .into_iter() + .filter_map(|result| self.entries.get(result.document.id)) + .collect::>(); + if !use_default_limit { + return results; + } + + if results + .iter() + .any(|(_, 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)) + .collect(); + } + limit_results_per_server(results) + } +} + +fn limit_results_per_server(results: Vec<&(String, ToolInfo)>) -> Vec<&(String, ToolInfo)> { + results + .into_iter() + .scan( + std::collections::HashMap::<&str, usize>::new(), + |counts, entry| { + let tool = &entry.1; + let count = counts.entry(tool.server_name.as_str()).or_default(); + if *count >= default_limit_for_server(tool.server_name.as_str()) { + Some(None) + } else { + *count += 1; + Some(Some(entry)) + } + }, + ) + .flatten() + .collect() +} + +fn default_limit_for_server(server_name: &str) -> usize { + if server_name == COMPUTER_USE_MCP_SERVER_NAME { + COMPUTER_USE_TOOL_SEARCH_LIMIT + } else { + TOOL_SEARCH_DEFAULT_LIMIT + } +} + fn build_search_text(name: &str, info: &ToolInfo) -> String { let mut parts = vec![ name.to_string(), @@ -156,3 +217,153 @@ fn build_search_text(name: &str, info: &ToolInfo) -> String { parts.join(" ") } + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use rmcp::model::Tool; + use std::sync::Arc; + + #[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 results = handler.search_result_entries( + "computer use", + TOOL_SEARCH_DEFAULT_LIMIT, + /*use_default_limit*/ true, + ); + + assert_eq!(results.len(), COMPUTER_USE_TOOL_SEARCH_LIMIT); + assert!( + results + .iter() + .all(|(_, tool)| tool.server_name == COMPUTER_USE_MCP_SERVER_NAME) + ); + + let explicit_results = handler.search_result_entries( + "computer use", + /*limit*/ 100, + /*use_default_limit*/ false, + ); + + assert_eq!(explicit_results.len(), 100); + } + + #[test] + fn non_computer_use_query_keeps_default_limit_with_computer_use_tools_installed() { + let mut tools = numbered_tools( + COMPUTER_USE_MCP_SERVER_NAME, + "computer use", + /*count*/ 100, + ); + tools.extend(numbered_tools( + "other-server", + "calendar", + /*count*/ 100, + )); + let handler = ToolSearchHandler::new(tools); + + let results = handler.search_result_entries( + "calendar", + TOOL_SEARCH_DEFAULT_LIMIT, + /*use_default_limit*/ true, + ); + + assert_eq!(results.len(), TOOL_SEARCH_DEFAULT_LIMIT); + assert!( + results + .iter() + .all(|(_, tool)| tool.server_name == "other-server") + ); + + let explicit_results = handler.search_result_entries( + "calendar", /*limit*/ 100, /*use_default_limit*/ false, + ); + + assert_eq!(explicit_results.len(), 100); + } + + #[test] + fn expanded_search_keeps_non_computer_use_servers_at_default_limit() { + let mut tools = numbered_tools( + COMPUTER_USE_MCP_SERVER_NAME, + "computer use", + /*count*/ 100, + ); + tools.extend(numbered_tools( + "other-server", + "computer use", + /*count*/ 100, + )); + let handler = ToolSearchHandler::new(tools); + + let results = handler.search_result_entries( + "computer use", + TOOL_SEARCH_DEFAULT_LIMIT, + /*use_default_limit*/ true, + ); + + assert!( + count_results_for_server(&results, COMPUTER_USE_MCP_SERVER_NAME) + <= COMPUTER_USE_TOOL_SEARCH_LIMIT + ); + assert!(count_results_for_server(&results, "other-server") <= TOOL_SEARCH_DEFAULT_LIMIT); + } + + fn numbered_tools( + server_name: &str, + description_prefix: &str, + count: usize, + ) -> std::collections::HashMap { + (0..count) + .map(|index| { + let tool_name = format!("tool_{index:03}"); + ( + format!("mcp__{server_name}__{tool_name}"), + tool_info(server_name, &tool_name, description_prefix), + ) + }) + .collect() + } + + fn tool_info(server_name: &str, tool_name: &str, description_prefix: &str) -> ToolInfo { + ToolInfo { + server_name: server_name.to_string(), + callable_name: tool_name.to_string(), + callable_namespace: format!("mcp__{server_name}__"), + server_instructions: None, + tool: Tool { + name: tool_name.to_string().into(), + title: None, + description: Some(format!("{description_prefix} desktop tool").into()), + input_schema: Arc::new(rmcp::model::object(serde_json::json!({ + "type": "object", + "properties": {}, + "additionalProperties": false, + }))), + output_schema: None, + annotations: None, + execution: None, + icons: None, + meta: None, + }, + connector_id: None, + connector_name: None, + plugin_display_names: Vec::new(), + connector_description: None, + } + } + + fn count_results_for_server(results: &[&(String, ToolInfo)], server_name: &str) -> usize { + results + .iter() + .filter(|(_, tool)| tool.server_name == server_name) + .count() + } +}