preserve search results order in tool_search_output (#17263)

we used to alpha-sort tool search results because we were using
`BTreeMap`, which threw away the actual search result ordering.

Now we use a vec to preserve it.

### Tests
Updated tests
This commit is contained in:
sayan-oai
2026-04-09 18:15:10 -07:00
committed by GitHub
Unverified
parent b976e701a8
commit 04fc208b6d
2 changed files with 53 additions and 31 deletions
+13 -14
View File
@@ -204,16 +204,17 @@ pub fn create_tool_search_tool(
pub fn collect_tool_search_output_tools<'a>(
tool_sources: impl IntoIterator<Item = ToolSearchResultSource<'a>>,
) -> Result<Vec<ToolSearchOutputTool>, serde_json::Error> {
let grouped = tool_sources.into_iter().fold(
BTreeMap::<String, Vec<(String, ToolSearchResultSource<'a>)>>::new(),
|mut grouped, tool| {
grouped
.entry(tool.tool_namespace.to_string())
.or_default()
.push((tool.tool_name.to_string(), tool));
grouped
},
);
let mut grouped: Vec<(&'a str, Vec<ToolSearchResultSource<'a>>)> = 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 {
@@ -222,12 +223,10 @@ pub fn collect_tool_search_output_tools<'a>(
};
let description = first_tool
.1
.connector_description
.map(str::to_string)
.or_else(|| {
first_tool
.1
.connector_name
.map(str::trim)
.filter(|connector_name| !connector_name.is_empty())
@@ -236,14 +235,14 @@ pub fn collect_tool_search_output_tools<'a>(
.or_else(|| {
Some(format!(
"Tools from the {} MCP server.",
first_tool.1.server_name
first_tool.server_name
))
});
let tools = tools
.iter()
.map(|tool| {
mcp_tool_to_deferred_responses_api_tool(tool.0.clone(), tool.1.tool)
mcp_tool_to_deferred_responses_api_tool(tool.tool_name.to_string(), tool.tool)
.map(ResponsesApiNamespaceTool::Function)
})
.collect::<Result<Vec<_>, _>>()?;
+40 -17
View File
@@ -137,13 +137,22 @@ fn create_tool_suggest_tool_uses_plugin_summary_fallback() {
}
#[test]
fn collect_tool_search_output_tools_groups_results_by_namespace() {
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",
@@ -182,6 +191,36 @@ fn collect_tool_search_output_tools_groups_results_by_namespace() {
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(),
@@ -212,22 +251,6 @@ fn collect_tool_search_output_tools_groups_results_by_namespace() {
}),
],
}),
ToolSearchOutputTool::Namespace(ResponsesApiNamespace {
name: "mcp__codex_apps__gmail".to_string(),
description: "Read mail".to_string(),
tools: vec![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__docs__".to_string(),
description: "Tools from the docs MCP server.".to_string(),