[mcp-tools] Use IndexMap for tool ordering

This commit is contained in:
Dylan Hurd
2025-08-25 07:31:34 -07:00
parent 295ca27e98
commit e573e9d435
9 changed files with 63 additions and 52 deletions

23
codex-rs/Cargo.lock generated
View File

@@ -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",

View File

@@ -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"

View File

@@ -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<String, McpServerConfig>,
pub mcp_servers: IndexMap<String, McpServerConfig>,
/// Combined provider map (defaults merged with user-defined overrides).
pub model_providers: HashMap<String, ModelProviderInfo>,
@@ -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<String, McpServerConfig>,
pub mcp_servers: IndexMap<String, McpServerConfig>,
/// 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(),

View File

@@ -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<String, anyhow::Error>;
fn qualify_tools(tools: Vec<ToolInfo>) -> HashMap<String, ToolInfo> {
fn qualify_tools(tools: Vec<ToolInfo>) -> IndexMap<String, ToolInfo> {
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<String, std::sync::Arc<McpClient>>,
clients: IndexMap<String, std::sync::Arc<McpClient>>,
/// Fully qualified tool name -> tool instance.
tools: HashMap<String, ToolInfo>,
tools: IndexMap<String, ToolInfo>,
}
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<String, McpServerConfig>,
mcp_servers: IndexMap<String, McpServerConfig>,
) -> 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<String, std::sync::Arc<McpClient>> =
HashMap::with_capacity(join_set.len());
let mut clients: IndexMap<String, std::sync::Arc<McpClient>> =
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<String, Tool> {
pub fn list_all_tools(&self) -> IndexMap<String, Tool> {
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<String, std::sync::Arc<McpClient>>,
clients: &IndexMap<String, std::sync::Arc<McpClient>>,
) -> Result<Vec<ToolInfo>> {
let mut join_set = JoinSet::new();
@@ -248,27 +249,34 @@ async fn list_all_tools(
});
}
let mut aggregated: Vec<ToolInfo> = 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)
}

View File

@@ -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<HashMap<String, mcp_types::Tool>>,
mcp_tools: Option<IndexMap<String, mcp_types::Tool>>,
) -> Vec<OpenAiTool> {
let mut tools: Vec<OpenAiTool> = 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<String, mcp_types::Tool> = HashMap::from([
let tools_map: IndexMap<String, mcp_types::Tool> = 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(),

View File

@@ -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"] }

View File

@@ -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<String, McpTool>,
pub tools: IndexMap<String, McpTool>,
}
#[derive(Debug, Default, Clone, Deserialize, Serialize)]

View File

@@ -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"

View File

@@ -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<String, mcp_types::Tool>,
tools: indexmap::IndexMap<String, mcp_types::Tool>,
) -> PlainHistoryCell {
let mut lines: Vec<Line<'static>> = vec![
Line::from("/mcp".magenta()),