agent/codex
2
0
Fork 0
mirror of https://github.com/pchuan98/codex.git synced 2026-07-01 00:31:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Adjust default tool search result caps (#17684)

Browse Source 
## Summary

- Allows selected MCP results to return a larger default result set.
- Keeps the existing default cap for other MCP results.
- Applies the cap consistently when higher explicit limits are
requested.

## Testing

- `cargo test -p codex-core tool_search`
- Ran a local CLI smoke test with two stdio MCP servers exposing 100
tools each; the selected-server query returned 20 tools and the
regular-server query returned 8.
This commit is contained in:
malone hedges
2026-04-14 14:57:19 -07:00
committed by GitHub
parent 8b7d0e9201
commit 78835d7e63
1 changed files with 226 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/tools/handlers/tool_search.rs
+226 -15
View File
@@ -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<usize>,
@@ -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::<Vec<_>>();
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<String, ToolInfo> {
(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()
}
}