From e573e9d435b83ebb9329a8d36c0ecd6fd2735ae6 Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Mon, 25 Aug 2025 07:31:34 -0700 Subject: [PATCH] [mcp-tools] Use IndexMap for tool ordering --- codex-rs/Cargo.lock | 23 ++++++----- codex-rs/core/Cargo.toml | 1 + codex-rs/core/src/config.rs | 11 +++--- codex-rs/core/src/mcp_connection_manager.rs | 42 ++++++++++++--------- codex-rs/core/src/openai_tools.rs | 31 +++++++-------- codex-rs/protocol/Cargo.toml | 1 + codex-rs/protocol/src/protocol.rs | 3 +- codex-rs/tui/Cargo.toml | 1 + codex-rs/tui/src/history_cell.rs | 2 +- 9 files changed, 63 insertions(+), 52 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 43cea8f9f2..b261331830 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -726,6 +726,7 @@ dependencies = [ "env-flags", "eventsource-stream", "futures", + "indexmap 2.11.0", "landlock", "libc", "maplit", @@ -927,6 +928,7 @@ name = "codex-protocol" version = "0.0.0" dependencies = [ "base64 0.22.1", + "indexmap 2.11.0", "mcp-types", "mime_guess", "pretty_assertions", @@ -972,6 +974,7 @@ dependencies = [ "crossterm", "diffy", "image", + "indexmap 2.11.0", "insta", "lazy_static", "libc", @@ -2008,7 +2011,7 @@ dependencies = [ "futures-core", "futures-sink", "http", - "indexmap 2.10.0", + "indexmap 2.11.0", "slab", "tokio", "tokio-util", @@ -2421,9 +2424,9 @@ dependencies = [ [[package]] name = "indexmap" -version = "2.10.0" +version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fe4cd85333e22411419a0bcae1297d25e58c9443848b11dc6a86fefe8c78a661" +checksum = "f2481980430f9f78649238835720ddccc57e52df14ffce1c6f37391d61b563e9" dependencies = [ "equivalent", "hashbrown 0.15.4", @@ -3389,7 +3392,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4c5cc86750666a3ed20bdaf5ca2a0344f9c67674cae0515bec2da16fbaa47db" dependencies = [ "fixedbitset", - "indexmap 2.10.0", + "indexmap 2.11.0", ] [[package]] @@ -3426,7 +3429,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3af6b589e163c5a788fab00ce0c0366f6efbb9959c2f9874b224936af7fce7e1" dependencies = [ "base64 0.22.1", - "indexmap 2.10.0", + "indexmap 2.11.0", "quick-xml", "serde", "time", @@ -4328,7 +4331,7 @@ version = "1.0.143" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d401abef1d108fbd9cbaebc3e46611f4b1021f714a0597a71f41ee463f5f4a5a" dependencies = [ - "indexmap 2.10.0", + "indexmap 2.11.0", "itoa", "memchr", "ryu", @@ -4386,7 +4389,7 @@ dependencies = [ "chrono", "hex", "indexmap 1.9.3", - "indexmap 2.10.0", + "indexmap 2.11.0", "schemars 0.9.0", "schemars 1.0.4", "serde", @@ -5153,7 +5156,7 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "75129e1dc5000bfbaa9fee9d1b21f974f9fbad9daec557a521ee6e080825f6e8" dependencies = [ - "indexmap 2.10.0", + "indexmap 2.11.0", "serde", "serde_spanned 1.0.0", "toml_datetime 0.7.0", @@ -5186,7 +5189,7 @@ version = "0.22.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ - "indexmap 2.10.0", + "indexmap 2.11.0", "serde", "serde_spanned 0.6.9", "toml_datetime 0.6.11", @@ -5199,7 +5202,7 @@ version = "0.23.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "17d3b47e6b7a040216ae5302712c94d1cf88c95b47efa80e2c59ce96c878267e" dependencies = [ - "indexmap 2.10.0", + "indexmap 2.11.0", "toml_datetime 0.7.0", "toml_parser", "toml_writer", diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index aa5747df6d..c9d41ddd65 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -25,6 +25,7 @@ dirs = "6" env-flags = "0.1.1" eventsource-stream = "0.2.3" futures = "0.3" +indexmap = "2.11.0" libc = "0.2.175" mcp-types = { path = "../mcp-types" } mime_guess = "2.0" diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index 98a8fde135..704cf9bea0 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -20,6 +20,7 @@ use codex_protocol::config_types::ReasoningEffort; use codex_protocol::config_types::ReasoningSummary; use codex_protocol::config_types::SandboxMode; use dirs::home_dir; +use indexmap::IndexMap; use serde::Deserialize; use std::collections::HashMap; use std::path::Path; @@ -114,7 +115,7 @@ pub struct Config { pub cwd: PathBuf, /// Definition for MCP servers that Codex can reach out to for tool calls. - pub mcp_servers: HashMap, + pub mcp_servers: IndexMap, /// Combined provider map (defaults merged with user-defined overrides). pub model_providers: HashMap, @@ -421,7 +422,7 @@ pub struct ConfigToml { /// Definition for MCP servers that Codex can reach out to for tool calls. #[serde(default)] - pub mcp_servers: HashMap, + pub mcp_servers: IndexMap, /// User-defined provider entries that extend/override the built-in list. #[serde(default)] @@ -1130,7 +1131,7 @@ disable_response_storage = true user_instructions: None, notify: None, cwd: fixture.cwd(), - mcp_servers: HashMap::new(), + mcp_servers: IndexMap::new(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, codex_home: fixture.codex_home(), @@ -1186,7 +1187,7 @@ disable_response_storage = true user_instructions: None, notify: None, cwd: fixture.cwd(), - mcp_servers: HashMap::new(), + mcp_servers: IndexMap::new(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, codex_home: fixture.codex_home(), @@ -1257,7 +1258,7 @@ disable_response_storage = true user_instructions: None, notify: None, cwd: fixture.cwd(), - mcp_servers: HashMap::new(), + mcp_servers: IndexMap::new(), model_providers: fixture.model_provider_map.clone(), project_doc_max_bytes: PROJECT_DOC_MAX_BYTES, codex_home: fixture.codex_home(), diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index b5813a0462..31b92df2a1 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -15,6 +15,7 @@ use anyhow::Context; use anyhow::Result; use anyhow::anyhow; use codex_mcp_client::McpClient; +use indexmap::IndexMap; use mcp_types::ClientCapabilities; use mcp_types::Implementation; use mcp_types::Tool; @@ -43,9 +44,9 @@ const LIST_TOOLS_TIMEOUT: Duration = Duration::from_secs(10); /// spawned successfully. pub type ClientStartErrors = HashMap; -fn qualify_tools(tools: Vec) -> HashMap { +fn qualify_tools(tools: Vec) -> IndexMap { let mut used_names = HashSet::new(); - let mut qualified_tools = HashMap::new(); + let mut qualified_tools = IndexMap::new(); for tool in tools { let mut qualified_name = format!( "{}{}{}", @@ -88,10 +89,10 @@ pub(crate) struct McpConnectionManager { /// /// The server name originates from the keys of the `mcp_servers` map in /// the user configuration. - clients: HashMap>, + clients: IndexMap>, /// Fully qualified tool name -> tool instance. - tools: HashMap, + tools: IndexMap, } impl McpConnectionManager { @@ -104,7 +105,7 @@ impl McpConnectionManager { /// Servers that fail to start are reported in `ClientStartErrors`: the /// user should be informed about these errors. pub async fn new( - mcp_servers: HashMap, + mcp_servers: IndexMap, ) -> Result<(Self, ClientStartErrors)> { // Early exit if no servers are configured. if mcp_servers.is_empty() { @@ -168,8 +169,8 @@ impl McpConnectionManager { }); } - let mut clients: HashMap> = - HashMap::with_capacity(join_set.len()); + let mut clients: IndexMap> = + IndexMap::with_capacity(join_set.len()); while let Some(res) = join_set.join_next().await { let (server_name, client_res) = res?; // JoinError propagation @@ -193,7 +194,7 @@ impl McpConnectionManager { /// Returns a single map that contains **all** tools. Each key is the /// fully-qualified name for the tool. - pub fn list_all_tools(&self) -> HashMap { + pub fn list_all_tools(&self) -> IndexMap { self.tools .iter() .map(|(name, tool)| (name.clone(), tool.tool.clone())) @@ -230,7 +231,7 @@ impl McpConnectionManager { /// Query every server for its available tools and return a single map that /// contains **all** tools. Each key is the fully-qualified name for the tool. async fn list_all_tools( - clients: &HashMap>, + clients: &IndexMap>, ) -> Result> { let mut join_set = JoinSet::new(); @@ -248,27 +249,34 @@ async fn list_all_tools( }); } - let mut aggregated: Vec = Vec::with_capacity(join_set.len()); - + let mut response_map = HashMap::with_capacity(join_set.len()); while let Some(join_res) = join_set.join_next().await { let (server_name, list_result) = join_res?; let list_result = list_result?; + // sort tools by name to ensure stable ordering + let mut tools = Vec::with_capacity(list_result.tools.len()); for tool in list_result.tools { let tool_info = ToolInfo { server_name: server_name.clone(), tool_name: tool.name.clone(), tool, }; - aggregated.push(tool_info); + tools.push(tool_info); } + tools.sort_by_key(|tool| tool.tool_name.clone()); + response_map.insert(server_name, tools); } - info!( - "aggregated {} tools from {} servers", - aggregated.len(), - clients.len() - ); + // Our mcp tool list must be stable across requests to ensure prompt-caching, + // so we preserve the cfg ordering + let mut aggregated = Vec::with_capacity(response_map.len()); + for server_name in clients.keys() { + if let Some(tools) = response_map.remove(server_name) { + info!("loaded {} with {} tools", &server_name, tools.len()); + aggregated.extend(tools); + } + } Ok(aggregated) } diff --git a/codex-rs/core/src/openai_tools.rs b/codex-rs/core/src/openai_tools.rs index a9fdb4f0e4..bd2c5f4da8 100644 --- a/codex-rs/core/src/openai_tools.rs +++ b/codex-rs/core/src/openai_tools.rs @@ -1,9 +1,9 @@ +use indexmap::IndexMap; use serde::Deserialize; use serde::Serialize; use serde_json::Value as JsonValue; use serde_json::json; use std::collections::BTreeMap; -use std::collections::HashMap; use crate::model_family::ModelFamily; use crate::plan_tool::PLAN_TOOL; @@ -498,7 +498,7 @@ fn sanitize_json_schema(value: &mut JsonValue) { /// [`McpConnectionManager`] for more details. pub(crate) fn get_openai_tools( config: &ToolsConfig, - mcp_tools: Option>, + mcp_tools: Option>, ) -> Vec { let mut tools: Vec = Vec::new(); @@ -542,12 +542,7 @@ pub(crate) fn get_openai_tools( } if let Some(mcp_tools) = mcp_tools { - // Ensure deterministic ordering to maximize prompt cache hits. - // HashMap iteration order is non-deterministic, so sort by fully-qualified tool name. - let mut entries: Vec<(String, mcp_types::Tool)> = mcp_tools.into_iter().collect(); - entries.sort_by(|a, b| a.0.cmp(&b.0)); - - for (name, tool) in entries.into_iter() { + for (name, tool) in mcp_tools.into_iter() { match mcp_tool_to_openai_tool(name.clone(), tool.clone()) { Ok(converted_tool) => tools.push(OpenAiTool::Function(converted_tool)), Err(e) => { @@ -605,7 +600,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, }); - let tools = get_openai_tools(&config, Some(HashMap::new())); + let tools = get_openai_tools(&config, Some(IndexMap::new())); assert_eq_tool_names(&tools, &["local_shell", "update_plan", "web_search"]); } @@ -622,7 +617,7 @@ mod tests { include_web_search_request: true, use_streamable_shell_tool: false, }); - let tools = get_openai_tools(&config, Some(HashMap::new())); + let tools = get_openai_tools(&config, Some(IndexMap::new())); assert_eq_tool_names(&tools, &["shell", "update_plan", "web_search"]); } @@ -641,7 +636,7 @@ mod tests { }); let tools = get_openai_tools( &config, - Some(HashMap::from([( + Some(IndexMap::from([( "test_server/do_something_cool".to_string(), mcp_types::Tool { name: "do_something_cool".to_string(), @@ -727,7 +722,7 @@ mod tests { } #[test] - fn test_get_openai_tools_mcp_tools_sorted_by_name() { + fn test_get_openai_tools_mcp_tools_sorted_by_insertion_order() { let model_family = find_family_for_model("o3").expect("o3 should be a valid model family"); let config = ToolsConfig::new(&ToolsConfigParams { model_family: &model_family, @@ -740,7 +735,7 @@ mod tests { }); // Intentionally construct a map with keys that would sort alphabetically. - let tools_map: HashMap = HashMap::from([ + let tools_map: IndexMap = IndexMap::from([ ( "test_server/do".to_string(), mcp_types::Tool { @@ -794,8 +789,8 @@ mod tests { &tools, &[ "shell", - "test_server/cool", "test_server/do", + "test_server/cool", "test_server/something", ], ); @@ -816,7 +811,7 @@ mod tests { let tools = get_openai_tools( &config, - Some(HashMap::from([( + Some(IndexMap::from([( "dash/search".to_string(), mcp_types::Tool { name: "search".to_string(), @@ -874,7 +869,7 @@ mod tests { let tools = get_openai_tools( &config, - Some(HashMap::from([( + Some(IndexMap::from([( "dash/paginate".to_string(), mcp_types::Tool { name: "paginate".to_string(), @@ -927,7 +922,7 @@ mod tests { let tools = get_openai_tools( &config, - Some(HashMap::from([( + Some(IndexMap::from([( "dash/tags".to_string(), mcp_types::Tool { name: "tags".to_string(), @@ -983,7 +978,7 @@ mod tests { let tools = get_openai_tools( &config, - Some(HashMap::from([( + Some(IndexMap::from([( "dash/value".to_string(), mcp_types::Tool { name: "value".to_string(), diff --git a/codex-rs/protocol/Cargo.toml b/codex-rs/protocol/Cargo.toml index aac2444014..89f1d649a9 100644 --- a/codex-rs/protocol/Cargo.toml +++ b/codex-rs/protocol/Cargo.toml @@ -12,6 +12,7 @@ workspace = true [dependencies] base64 = "0.22.1" +indexmap = { version = "2.11.0", features = ["serde"] } mcp-types = { path = "../mcp-types" } mime_guess = "2.0.5" serde = { version = "1", features = ["derive"] } diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 71c1538197..3a7505a878 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -10,6 +10,7 @@ use std::path::PathBuf; use std::str::FromStr; use std::time::Duration; +use indexmap::IndexMap; use mcp_types::CallToolResult; use mcp_types::Tool as McpTool; use serde::Deserialize; @@ -798,7 +799,7 @@ pub struct GetHistoryEntryResponseEvent { #[derive(Debug, Clone, Deserialize, Serialize)] pub struct McpListToolsResponseEvent { /// Fully qualified tool name -> tool definition. - pub tools: std::collections::HashMap, + pub tools: IndexMap, } #[derive(Debug, Default, Clone, Deserialize, Serialize)] diff --git a/codex-rs/tui/Cargo.toml b/codex-rs/tui/Cargo.toml index 6d69e97e73..aeffa3a99a 100644 --- a/codex-rs/tui/Cargo.toml +++ b/codex-rs/tui/Cargo.toml @@ -46,6 +46,7 @@ image = { version = "^0.25.6", default-features = false, features = [ "jpeg", "png", ] } +indexmap = "2.11.0" lazy_static = "1" mcp-types = { path = "../mcp-types" } once_cell = "1" diff --git a/codex-rs/tui/src/history_cell.rs b/codex-rs/tui/src/history_cell.rs index 0b2af7a100..e3f684e0dd 100644 --- a/codex-rs/tui/src/history_cell.rs +++ b/codex-rs/tui/src/history_cell.rs @@ -782,7 +782,7 @@ pub(crate) fn empty_mcp_output() -> PlainHistoryCell { /// Render MCP tools grouped by connection using the fully-qualified tool names. pub(crate) fn new_mcp_tools_output( config: &Config, - tools: std::collections::HashMap, + tools: indexmap::IndexMap, ) -> PlainHistoryCell { let mut lines: Vec> = vec![ Line::from("/mcp".magenta()),