From 04fc208b6dae3b76ba535db54de716aa818bb70d Mon Sep 17 00:00:00 2001 From: sayan-oai Date: Thu, 9 Apr 2026 18:15:10 -0700 Subject: [PATCH] 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 --- codex-rs/tools/src/tool_discovery.rs | 27 +++++----- codex-rs/tools/src/tool_discovery_tests.rs | 57 +++++++++++++++------- 2 files changed, 53 insertions(+), 31 deletions(-) diff --git a/codex-rs/tools/src/tool_discovery.rs b/codex-rs/tools/src/tool_discovery.rs index 594d4011b..f5d84a952 100644 --- a/codex-rs/tools/src/tool_discovery.rs +++ b/codex-rs/tools/src/tool_discovery.rs @@ -204,16 +204,17 @@ pub fn create_tool_search_tool( pub fn collect_tool_search_output_tools<'a>( tool_sources: impl IntoIterator>, ) -> Result, serde_json::Error> { - let grouped = tool_sources.into_iter().fold( - BTreeMap::)>>::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>)> = 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::, _>>()?; diff --git a/codex-rs/tools/src/tool_discovery_tests.rs b/codex-rs/tools/src/tool_discovery_tests.rs index dbbf56c3c..10e108fe0 100644 --- a/codex-rs/tools/src/tool_discovery_tests.rs +++ b/codex-rs/tools/src/tool_discovery_tests.rs @@ -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(),