From fd429c488c34ef82ff9a3aaf762f17af5a7cc961 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 5 May 2025 16:19:12 -0700 Subject: [PATCH] feat: support mcp_servers in config.toml --- codex-rs/Cargo.lock | 3 +- codex-rs/README.md | 4 + codex-rs/core/Cargo.toml | 2 + codex-rs/core/src/client.rs | 45 ++++- codex-rs/core/src/codex.rs | 93 +++++++++- codex-rs/core/src/config.rs | 10 + codex-rs/core/src/lib.rs | 2 + codex-rs/core/src/mcp_connection_manager.rs | 194 ++++++++++++++++++++ codex-rs/core/src/mcp_server_config.rs | 14 ++ codex-rs/mcp-client/Cargo.toml | 1 - codex-rs/mcp-client/src/mcp_client.rs | 18 +- 11 files changed, 371 insertions(+), 15 deletions(-) create mode 100644 codex-rs/core/src/mcp_connection_manager.rs create mode 100644 codex-rs/core/src/mcp_server_config.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 4b73372fb6..3e68b7ed70 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -493,6 +493,7 @@ dependencies = [ "bytes", "clap", "codex-apply-patch", + "codex-mcp-client", "dirs", "env-flags", "eventsource-stream", @@ -500,6 +501,7 @@ dependencies = [ "futures", "landlock", "libc", + "mcp-types", "mime_guess", "openssl-sys", "patch", @@ -561,7 +563,6 @@ name = "codex-mcp-client" version = "0.1.0" dependencies = [ "anyhow", - "codex-core", "mcp-types", "pretty_assertions", "serde", diff --git a/codex-rs/README.md b/codex-rs/README.md index 3c42ceff4a..e00f59a565 100644 --- a/codex-rs/README.md +++ b/codex-rs/README.md @@ -79,6 +79,10 @@ sandbox_permissions = [ ] ``` +### mcp_servers + +FIXME: document this part of the config + ### disable_response_storage Currently, customers whose accounts are set to use Zero Data Retention (ZDR) must set `disable_response_storage` to `true` so that Codex uses an alternative to the Responses API that works with ZDR: diff --git a/codex-rs/core/Cargo.toml b/codex-rs/core/Cargo.toml index 0ed550f9a8..0d880eed4a 100644 --- a/codex-rs/core/Cargo.toml +++ b/codex-rs/core/Cargo.toml @@ -19,6 +19,7 @@ env-flags = "0.1.1" eventsource-stream = "0.2.3" fs-err = "3.1.0" futures = "0.3" +codex-mcp-client = { path = "../mcp-client" } mime_guess = "2.0" patch = "0.7" path-absolutize = "3.1.1" @@ -39,6 +40,7 @@ toml = "0.8.20" tracing = { version = "0.1.41", features = ["log"] } tree-sitter = "0.25.3" tree-sitter-bash = "0.23.3" +mcp-types = { path = "../mcp-types" } [target.'cfg(target_os = "linux")'.dependencies] libc = "0.2.172" diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 10ec0b9780..47ccfa580b 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1,4 +1,5 @@ use std::collections::BTreeMap; +use std::collections::HashMap; use std::io::BufRead; use std::path::Path; use std::pin::Pin; @@ -13,6 +14,7 @@ use futures::prelude::*; use reqwest::StatusCode; use serde::Deserialize; use serde::Serialize; +use serde_json::json; use serde_json::Value; use tokio::sync::mpsc; use tokio::time::timeout; @@ -42,6 +44,12 @@ pub struct Prompt { pub instructions: Option, /// Whether to store response on server side (disable_response_storage = !store). pub store: bool, + + /// Additional tools sourced from external MCP servers. Filled in only for + /// the first turn. Note the key is the "fully qualified" tool name + /// (i.e., prefixed with the server name), which should be reported to the + /// model in place of Tool::name. + pub extra_tools: HashMap, } #[derive(Debug)] @@ -59,7 +67,7 @@ struct Payload<'a> { // we code defensively to avoid this case, but perhaps we should use a // separate enum for serialization. input: &'a Vec, - tools: &'a [Tool], + tools: &'a [serde_json::Value], tool_choice: &'static str, parallel_tool_calls: bool, reasoning: Option, @@ -78,7 +86,7 @@ struct Reasoning { } #[derive(Debug, Serialize)] -struct Tool { +struct ToolInternal { name: &'static str, #[serde(rename = "type")] kind: &'static str, // "function" @@ -105,7 +113,7 @@ enum JsonSchema { } /// Tool usage specification -static TOOLS: LazyLock> = LazyLock::new(|| { +static TOOLS_INTERNAL: LazyLock> = LazyLock::new(|| { let mut properties = BTreeMap::new(); properties.insert( "command".to_string(), @@ -116,7 +124,7 @@ static TOOLS: LazyLock> = LazyLock::new(|| { properties.insert("workdir".to_string(), JsonSchema::String); properties.insert("timeout".to_string(), JsonSchema::Number); - vec![Tool { + vec![ToolInternal { name: "shell", kind: "function", description: "Runs a shell command, and returns its output.", @@ -149,11 +157,26 @@ impl ModelClient { return stream_from_fixture(path).await; } + // Assemble tool list: built-in tools + any extra tools from the prompt. + let mut tools_json: Vec = TOOLS_INTERNAL + .iter() + .map(|t| serde_json::to_value(t).expect("serialize builtin tool")) + .collect(); + tools_json.extend( + prompt + .extra_tools + .clone() + .into_iter() + .map(|(name, tool)| mcp_tool_to_openai_tool(name, tool)), + ); + + debug!("tools_json: {}", serde_json::to_string_pretty(&tools_json)?); + let payload = Payload { model: &self.model, instructions: prompt.instructions.as_ref(), input: &prompt.input, - tools: &TOOLS, + tools: &tools_json, tool_choice: "auto", parallel_tool_calls: false, reasoning: Some(Reasoning { @@ -235,6 +258,18 @@ impl ModelClient { } } +fn mcp_tool_to_openai_tool( + fully_qualified_name: String, + tool: mcp_types::Tool, +) -> serde_json::Value { + json!({ + "name": fully_qualified_name, + "description": tool.description, + "parameters": tool.input_schema, + "type": "function", + }) +} + #[derive(Debug, Deserialize, Serialize)] struct SseEvent { #[serde(rename = "type")] diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index c74d0079ee..52fe3c178b 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -38,6 +38,7 @@ use crate::exec::ExecParams; use crate::exec::ExecToolCallOutput; use crate::exec::SandboxType; use crate::flags::OPENAI_STREAM_MAX_RETRIES; +use crate::mcp_connection_manager::try_parse_fully_qualified_tool_name; use crate::models::ContentItem; use crate::models::FunctionCallOutputPayload; use crate::models::ResponseInputItem; @@ -202,6 +203,9 @@ struct Session { sandbox_policy: SandboxPolicy, writable_roots: Mutex>, + /// Manager for external MCP servers/tools. + mcp: crate::mcp_connection_manager::McpConnectionManager, + /// External notifier command (will be passed as args to exec()). When /// `None` this feature is disabled. notify: Option>, @@ -554,6 +558,34 @@ async fn submission_loop( }; let writable_roots = Mutex::new(get_writable_roots(&cwd)); + + // Load config to initialise the MCP connection manager. + let config = match crate::config::Config::load_with_overrides( + crate::config::ConfigOverrides::default(), + ) { + Ok(cfg) => cfg, + Err(e) => { + error!("Failed to load config for MCP servers: {e:#}"); + // Fall back to empty server map so the session can still proceed. + crate::config::Config::load_default_config_for_test() + } + }; + + let mcp = match crate::mcp_connection_manager::create_mcp_connection_manager( + config.mcp_servers.clone(), + ) + .await + { + Ok(mgr) => mgr, + Err(e) => { + error!("Failed to create MCP connection manager: {e:#}"); + // Use an empty manager so we can still continue. + crate::mcp_connection_manager::McpConnectionManager::new(HashMap::new()) + .await + .expect("empty manager should never fail") + } + }; + sess = Some(Arc::new(Session { client, tx_event: tx_event.clone(), @@ -565,6 +597,7 @@ async fn submission_loop( writable_roots, notify, state: Mutex::new(state), + mcp, })); // ack @@ -753,11 +786,22 @@ async fn run_turn( } else { None }; + + // FIXME: cache the list of tool calls + let extra_tools = match sess.mcp.list_all_tools().await { + Ok(v) => v, + Err(e) => { + tracing::warn!("failed to list tools from MCP servers: {e:#}"); + HashMap::new() + } + }; + let prompt = Prompt { input, prev_id, instructions, store, + extra_tools, }; let mut retries = 0; @@ -1141,13 +1185,48 @@ async fn handle_function_call( } } _ => { - // Unknown function: reply with structured failure so the model can adapt. - ResponseInputItem::FunctionCallOutput { - call_id, - output: crate::models::FunctionCallOutputPayload { - content: format!("unsupported call: {}", name), - success: None, - }, + match try_parse_fully_qualified_tool_name(&name) { + Some((server, tool_name)) => { + // Attempt to route to external MCP server. + let arguments_value: Option = + serde_json::from_str(&arguments).ok(); + + match sess + .mcp + .call_tool(&server, &tool_name, arguments_value) + .await + { + Ok(result) => { + let success = !result.is_error.unwrap_or(false); + let content = serde_json::to_string(&result) + .unwrap_or_else(|_| "".to_string()); + ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content, + success: Some(success), + }, + } + } + Err(e) => ResponseInputItem::FunctionCallOutput { + call_id, + output: FunctionCallOutputPayload { + content: format!("tool call error: {e}"), + success: Some(false), + }, + }, + } + } + None => { + // Unknown function: reply with structured failure so the model can adapt. + ResponseInputItem::FunctionCallOutput { + call_id, + output: crate::models::FunctionCallOutputPayload { + content: format!("unsupported call: {}", name), + success: None, + }, + } + } } } } diff --git a/codex-rs/core/src/config.rs b/codex-rs/core/src/config.rs index 1557ce2752..fbb9d81d3b 100644 --- a/codex-rs/core/src/config.rs +++ b/codex-rs/core/src/config.rs @@ -1,10 +1,12 @@ use crate::approval_mode_cli_arg::parse_sandbox_permission_with_base_path; use crate::flags::OPENAI_DEFAULT_MODEL; +use crate::mcp_server_config::McpServerConfig; use crate::protocol::AskForApproval; use crate::protocol::SandboxPermission; use crate::protocol::SandboxPolicy; use dirs::home_dir; use serde::Deserialize; +use std::collections::HashMap; use std::path::PathBuf; /// Embedded fallback instructions that mirror the TypeScript CLI’s default @@ -57,6 +59,9 @@ pub struct Config { /// for the session. All relative paths inside the business-logic layer are /// resolved against this path. pub cwd: PathBuf, + + /// Definition for MCP servers that Codex can reach out to for tool calls. + pub mcp_servers: HashMap, } /// Base config deserialized from ~/.codex/config.toml. @@ -85,6 +90,10 @@ pub struct ConfigToml { /// System instructions. pub instructions: Option, + + /// Definition for MCP servers that Codex can reach out to for tool calls. + #[serde(default)] + pub mcp_servers: HashMap, } impl ConfigToml { @@ -213,6 +222,7 @@ impl Config { .unwrap_or(false), notify: cfg.notify, instructions, + mcp_servers: cfg.mcp_servers, } } diff --git a/codex-rs/core/src/lib.rs b/codex-rs/core/src/lib.rs index a5909ed63d..a79cd49403 100644 --- a/codex-rs/core/src/lib.rs +++ b/codex-rs/core/src/lib.rs @@ -15,6 +15,8 @@ mod flags; mod is_safe_command; #[cfg(target_os = "linux")] pub mod linux; +mod mcp_connection_manager; +pub mod mcp_server_config; mod models; pub mod protocol; mod safety; diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs new file mode 100644 index 0000000000..711aef0cc9 --- /dev/null +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -0,0 +1,194 @@ +//! Connection manager for Model Context Protocol (MCP) servers. +//! +//! The [`McpConnectionManager`] owns one [`codex_mcp_client::McpClient`] per +//! configured server (keyed by the *server name*). It offers convenience +//! helpers to query the available tools across *all* servers and returns them +//! in a single aggregated map using the fully-qualified tool name +//! `""` as the key. + +use std::collections::HashMap; + +use anyhow::anyhow; +use anyhow::Result; +use codex_mcp_client::McpClient; +use mcp_types::Tool; +use tokio::task::JoinSet; +use tracing::info; +use tracing::warn; + +use crate::mcp_server_config::McpServerConfig; + +/// Delimiter used to separate the server name from the tool name in a fully +/// qualified tool name. +/// +/// OpenAI requires tool names to conform to `^[a-zA-Z0-9_-]+$`, so we must +/// choose a delimiter from this character set. +const MCP_TOOL_NAME_DELIMITER: &str = "__OAI_CODEX_MCP__"; + +fn fully_qualified_tool_name(server: &str, tool: &str) -> String { + format!("{server}{MCP_TOOL_NAME_DELIMITER}{tool}") +} + +pub(crate) fn try_parse_fully_qualified_tool_name(fq_name: &str) -> Option<(String, String)> { + let (server, tool) = fq_name.split_once(MCP_TOOL_NAME_DELIMITER)?; + if server.is_empty() || tool.is_empty() { + return None; + } + Some((server.to_string(), tool.to_string())) +} + +/// A thin wrapper around a set of running [`McpClient`] instances. +/// +/// The struct is intentionally lightweight – cloning just clones the internal +/// `HashMap` of clients which in turn clones the `Arc`s of each client. +#[derive(Clone)] +pub(crate) struct McpConnectionManager { + /// Server-name → client instance. + /// + /// The server name originates from the keys of the `mcp_servers` map in + /// the user configuration. + clients: HashMap>, // Arc to cheaply clone +} + +impl McpConnectionManager { + /// Spawn a [`McpClient`] for each configured server. + /// + /// * `mcp_servers` – Map loaded from the user configuration where *keys* + /// are human-readable server identifiers and *values* are the spawn + /// instructions. + pub async fn new(mcp_servers: HashMap) -> Result { + // Early exit if no servers are configured. + if mcp_servers.is_empty() { + return Ok(Self { + clients: HashMap::new(), + }); + } + + // Spin up all servers concurrently. + let mut join_set = JoinSet::new(); + + // Spawn tasks to launch each server. + for (server_name, cfg) in mcp_servers { + // Perform slash validation up-front so we can return early without + // spawning any tasks when the name is invalid. + if server_name.contains('/') { + return Err(anyhow!( + "MCP server name '{server_name}' must not contain a forward slash (/)" + )); + } + + join_set.spawn(async move { + // Build argv vector: first element is the command itself followed + // by the optional additional args from the config. + let mut argv = vec![cfg.command.clone()]; + argv.extend(cfg.args.clone()); + + // FIXME: take cfg.env into account when spawning the command. + let client_res = McpClient::new_stdio_client(argv).await; + + (server_name, client_res) + }); + } + + // Collect results. + let mut clients: HashMap> = HashMap::new(); + + while let Some(res) = join_set.join_next().await { + let (server_name, client_res) = res?; // propagate JoinError + + let client = client_res + .map_err(|e| anyhow!("failed to spawn MCP server '{server_name}': {e}"))?; + + clients.insert(server_name, std::sync::Arc::new(client)); + } + + Ok(Self { clients }) + } + + /// Return a reference to the internal client for the given server. + #[allow(dead_code)] + pub fn client_for_server(&self, server_name: &str) -> Option> { + self.clients.get(server_name).cloned() + } + + /// Query every server for its available tools and return a single map that + /// contains **all** tools. The key is the fully-qualified name + /// `/`. + pub async fn list_all_tools(&self) -> Result> { + let mut join_set = JoinSet::new(); + + // Spawn one task per server so we can query them concurrently. This + // keeps the overall latency roughly at the slowest server instead of + // the cumulative latency. + for (server_name, client) in self.clients.clone() { + let server_name_cloned = server_name.clone(); + let client_clone = client.clone(); + join_set.spawn(async move { + let res = client_clone.list_tools(None).await; + (server_name_cloned, res) + }); + } + + let mut aggregated: HashMap = HashMap::new(); + + while let Some(join_res) = join_set.join_next().await { + let (server_name, list_result) = join_res?; // propagate JoinError + + let list_result = list_result?; + + for tool in list_result.tools { + if tool.name.contains('/') { + warn!( + server = %server_name, + tool_name = %tool.name, + "tool name contains '/' – skipping to avoid ambiguity" + ); + continue; + } + + let fq_name = fully_qualified_tool_name(&server_name, &tool.name); + + if aggregated.insert(fq_name.clone(), tool).is_some() { + warn!("tool name collision for '{fq_name}' – overwriting previous entry"); + } + } + } + + info!( + "aggregated {} tools from {} servers", + aggregated.len(), + self.clients.len() + ); + + Ok(aggregated) + } + + /// Route a fully-qualified tool call to the matching server. + pub async fn call_tool( + &self, + server: &str, + tool: &str, + arguments: Option, + ) -> Result { + let client = self + .clients + .get(server) + .ok_or_else(|| anyhow!("unknown MCP server '{server}'"))? + .clone(); + + client + .call_tool(tool.to_string(), arguments) + .await + .map_err(|e| anyhow!("tool call failed for '{server}/{tool}': {e}")) + } +} + +/// Convenience helper that mirrors the previous `create_mcp_connection_manager` +/// free-standing function but returns `Result` and is **async**. Existing +/// call-sites can continue to call the function while new code can use the +/// `McpConnectionManager::new` associated function directly. +pub(crate) async fn create_mcp_connection_manager( + mcp_servers: HashMap, +) -> Result { + McpConnectionManager::new(mcp_servers).await +} diff --git a/codex-rs/core/src/mcp_server_config.rs b/codex-rs/core/src/mcp_server_config.rs new file mode 100644 index 0000000000..261a75d13e --- /dev/null +++ b/codex-rs/core/src/mcp_server_config.rs @@ -0,0 +1,14 @@ +use std::collections::HashMap; + +use serde::Deserialize; + +#[derive(Deserialize, Debug, Clone)] +pub struct McpServerConfig { + pub command: String, + + #[serde(default)] + pub args: Vec, + + #[serde(default)] + pub env: Option>, +} diff --git a/codex-rs/mcp-client/Cargo.toml b/codex-rs/mcp-client/Cargo.toml index 2101a1e697..b3792922cc 100644 --- a/codex-rs/mcp-client/Cargo.toml +++ b/codex-rs/mcp-client/Cargo.toml @@ -5,7 +5,6 @@ edition = "2021" [dependencies] anyhow = "1" -codex-core = { path = "../core", features = ["cli"] } mcp-types = { path = "../mcp-types" } serde = { version = "1", features = ["derive"] } serde_json = "1" diff --git a/codex-rs/mcp-client/src/mcp_client.rs b/codex-rs/mcp-client/src/mcp_client.rs index ccab93dc7d..1a892c7d5b 100644 --- a/codex-rs/mcp-client/src/mcp_client.rs +++ b/codex-rs/mcp-client/src/mcp_client.rs @@ -18,6 +18,8 @@ use std::sync::Arc; use anyhow::anyhow; use anyhow::Result; +use mcp_types::CallToolRequest; +use mcp_types::CallToolRequestParams; use mcp_types::JSONRPCMessage; use mcp_types::JSONRPCNotification; use mcp_types::JSONRPCRequest; @@ -37,6 +39,7 @@ use tokio::process::Command; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio::sync::Mutex; +use tracing::debug; use tracing::error; use tracing::info; use tracing::warn; @@ -122,6 +125,7 @@ impl McpClient { while let Some(msg) = outgoing_rx.recv().await { match serde_json::to_string(&msg) { Ok(json) => { + debug!("MCP message to server: {json}"); if stdin.write_all(json.as_bytes()).await.is_err() { error!("failed to write message to child stdin"); break; @@ -149,6 +153,7 @@ impl McpClient { tokio::spawn(async move { while let Ok(Some(line)) = lines.next_line().await { + debug!("MCP message from server: {line}"); match serde_json::from_str::(&line) { Ok(JSONRPCMessage::Response(resp)) => { Self::dispatch_response(resp, &pending).await; @@ -229,7 +234,7 @@ impl McpClient { // Send to writer task. if self.outgoing_tx.send(message).await.is_err() { return Err(anyhow!( - "failed to send message to writer task – channel closed" + "failed to send message to writer task - channel closed" )); } @@ -262,6 +267,17 @@ impl McpClient { self.send_request::(params).await } + /// Convenience wrapper around `tools/call`. + pub async fn call_tool( + &self, + name: String, + arguments: Option, + ) -> Result { + let params = CallToolRequestParams { name, arguments }; + debug!("MCP tool call: {params:?}"); + self.send_request::(params).await + } + /// Internal helper: route a JSON-RPC *response* object to the pending map. async fn dispatch_response( resp: JSONRPCResponse,