From 2b72d05c5ef574c695fdf062171ce37e0f38905f Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 5 May 2025 07:16:19 -0700 Subject: [PATCH 1/2] feat: make Codex available as a tool when running it as an MCP server (#811) This PR replaces the placeholder `"echo"` tool call in the MCP server with a `"codex"` tool that calls Codex. Events such as `ExecApprovalRequest` and `ApplyPatchApprovalRequest` are not handled properly yet, but I have `approval_policy = "never"` set in my `~/.codex/config.toml` such that those codepaths are not exercised. The schema for this MPC tool is defined by a new `CodexToolCallParam` struct introduced in this PR. It is fairly similar to `ConfigOverrides`, as the param is used to help create the `Config` used to start the Codex session, though it also includes the `prompt` used to kick off the session. This PR also introduces the use of the third-party `schemars` crate to generate the JSON schema, which is verified in the `verify_codex_tool_json_schema()` unit test. Events that are dispatched during the Codex session are sent back to the MCP client as MCP notifications. This gives the client a way to monitor progress as the tool call itself may take minutes to complete depending on the complexity of the task requested by the user. In the video below, I launched the server via: ```shell mcp-server$ RUST_LOG=debug npx @modelcontextprotocol/inspector cargo run -- ``` In the video, you can see the flow of: * requesting the list of tools * choosing the **codex** tool * entering a value for **prompt** and then making the tool call Note that I left the other fields blank because when unspecified, the values in my `~/.codex/config.toml` were used: https://github.com/user-attachments/assets/1975058c-b004-43ef-8c8d-800a953b8192 Note that while using the inspector, I did run into https://github.com/modelcontextprotocol/inspector/issues/293, though the tip about ensuring I had only one instance of the **MCP Inspector** tab open in my browser seemed to fix things. --- codex-rs/Cargo.lock | 43 ++++ codex-rs/mcp-server/Cargo.toml | 15 +- codex-rs/mcp-server/src/codex_tool_config.rs | 244 +++++++++++++++++++ codex-rs/mcp-server/src/codex_tool_runner.rs | 181 ++++++++++++++ codex-rs/mcp-server/src/main.rs | 4 + codex-rs/mcp-server/src/message_processor.rs | 102 +++++--- 6 files changed, 548 insertions(+), 41 deletions(-) create mode 100644 codex-rs/mcp-server/src/codex_tool_config.rs create mode 100644 codex-rs/mcp-server/src/codex_tool_runner.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index f2f865b02b..0a4d879746 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -562,6 +562,8 @@ version = "0.1.0" dependencies = [ "codex-core", "mcp-types", + "pretty_assertions", + "schemars", "serde", "serde_json", "tokio", @@ -934,6 +936,12 @@ dependencies = [ "syn 2.0.100", ] +[[package]] +name = "dyn-clone" +version = "1.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c7a8fb8a9fbf66c1f703fe16184d10ca0ee9d23be5b4436400408ba54a95005" + [[package]] name = "either" version = "1.15.0" @@ -2824,6 +2832,30 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "schemars" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615" +dependencies = [ + "dyn-clone", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn 2.0.100", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -2882,6 +2914,17 @@ dependencies = [ "syn 2.0.100", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.100", +] + [[package]] name = "serde_json" version = "1.0.140" diff --git a/codex-rs/mcp-server/Cargo.toml b/codex-rs/mcp-server/Cargo.toml index 258a37aace..fdd2a304cd 100644 --- a/codex-rs/mcp-server/Cargo.toml +++ b/codex-rs/mcp-server/Cargo.toml @@ -4,19 +4,9 @@ version = "0.1.0" edition = "2021" [dependencies] -# -# codex-core contains optional functionality that is gated behind the "cli" -# feature. Unfortunately there is an unconditional reference to a module that -# is only compiled when the feature is enabled, which breaks the build when -# the default (no-feature) variant is used. -# -# We therefore explicitly enable the "cli" feature when codex-mcp-server pulls -# in codex-core so that the required symbols are present. This does _not_ -# change the public API of codex-core – it merely opts into compiling the -# extra, feature-gated source files so the build succeeds. -# codex-core = { path = "../core", features = ["cli"] } mcp-types = { path = "../mcp-types" } +schemars = "0.8.22" serde = { version = "1", features = ["derive"] } serde_json = "1" tracing = { version = "0.1.41", features = ["log"] } @@ -28,3 +18,6 @@ tokio = { version = "1", features = [ "rt-multi-thread", "signal", ] } + +[dev-dependencies] +pretty_assertions = "1.4.1" diff --git a/codex-rs/mcp-server/src/codex_tool_config.rs b/codex-rs/mcp-server/src/codex_tool_config.rs new file mode 100644 index 0000000000..aa1a620dc0 --- /dev/null +++ b/codex-rs/mcp-server/src/codex_tool_config.rs @@ -0,0 +1,244 @@ +//! Configuration object accepted by the `codex` MCP tool-call. + +use std::path::PathBuf; + +use mcp_types::Tool; +use mcp_types::ToolInputSchema; +use schemars::r#gen::SchemaSettings; +use schemars::JsonSchema; +use serde::Deserialize; + +use codex_core::protocol::AskForApproval; +use codex_core::protocol::SandboxPolicy; + +/// Client-supplied configuration for a `codex` tool-call. +#[derive(Debug, Clone, Deserialize, JsonSchema)] +#[serde(rename_all = "kebab-case")] +pub(crate) struct CodexToolCallParam { + /// The *initial user prompt* to start the Codex conversation. + pub prompt: String, + + /// Optional override for the model name (e.g. "o3", "o4-mini") + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// Working directory for the session. If relative, it is resolved against + /// the server process's current working directory. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub cwd: Option, + + /// Execution approval policy expressed as the kebab-case variant name + /// (`unless-allow-listed`, `auto-edit`, `on-failure`, `never`). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub approval_policy: Option, + + /// Sandbox permissions using the same string values accepted by the CLI + /// (e.g. "disk-write-cwd", "network-full-access"). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sandbox_permissions: Option>, + + /// Disable server-side response storage. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub disable_response_storage: Option, + // Custom system instructions. + // #[serde(default, skip_serializing_if = "Option::is_none")] + // pub instructions: Option, +} + +// Create custom enums for use with `CodexToolCallApprovalPolicy` where we +// intentionally exclude docstrings from the generated schema because they +// introduce anyOf in the the generated JSON schema, which makes it more complex +// without adding any real value since we aspire to use self-descriptive names. + +#[derive(Debug, Clone, Deserialize, JsonSchema)] +#[serde(rename_all = "kebab-case")] +pub(crate) enum CodexToolCallApprovalPolicy { + AutoEdit, + UnlessAllowListed, + OnFailure, + Never, +} + +impl From for AskForApproval { + fn from(value: CodexToolCallApprovalPolicy) -> Self { + match value { + CodexToolCallApprovalPolicy::AutoEdit => AskForApproval::AutoEdit, + CodexToolCallApprovalPolicy::UnlessAllowListed => AskForApproval::UnlessAllowListed, + CodexToolCallApprovalPolicy::OnFailure => AskForApproval::OnFailure, + CodexToolCallApprovalPolicy::Never => AskForApproval::Never, + } + } +} + +// TODO: Support additional writable folders via a separate property on +// CodexToolCallParam. + +#[derive(Debug, Clone, Deserialize, JsonSchema)] +#[serde(rename_all = "kebab-case")] +pub(crate) enum CodexToolCallSandboxPermission { + DiskFullReadAccess, + DiskWriteCwd, + DiskWritePlatformUserTempFolder, + DiskWritePlatformGlobalTempFolder, + DiskFullWriteAccess, + NetworkFullAccess, +} + +impl From for codex_core::protocol::SandboxPermission { + fn from(value: CodexToolCallSandboxPermission) -> Self { + match value { + CodexToolCallSandboxPermission::DiskFullReadAccess => { + codex_core::protocol::SandboxPermission::DiskFullReadAccess + } + CodexToolCallSandboxPermission::DiskWriteCwd => { + codex_core::protocol::SandboxPermission::DiskWriteCwd + } + CodexToolCallSandboxPermission::DiskWritePlatformUserTempFolder => { + codex_core::protocol::SandboxPermission::DiskWritePlatformUserTempFolder + } + CodexToolCallSandboxPermission::DiskWritePlatformGlobalTempFolder => { + codex_core::protocol::SandboxPermission::DiskWritePlatformGlobalTempFolder + } + CodexToolCallSandboxPermission::DiskFullWriteAccess => { + codex_core::protocol::SandboxPermission::DiskFullWriteAccess + } + CodexToolCallSandboxPermission::NetworkFullAccess => { + codex_core::protocol::SandboxPermission::NetworkFullAccess + } + } + } +} + +pub(crate) fn create_tool_for_codex_tool_call_param() -> Tool { + let schema = SchemaSettings::draft2019_09() + .with(|s| { + s.inline_subschemas = true; + s.option_add_null_type = false + }) + .into_generator() + .into_root_schema_for::(); + let schema_value = + serde_json::to_value(&schema).expect("Codex tool schema should serialise to JSON"); + + let tool_input_schema = + serde_json::from_value::(schema_value).unwrap_or_else(|e| { + panic!("failed to create Tool from schema: {e}"); + }); + Tool { + name: "codex".to_string(), + input_schema: tool_input_schema, + description: Some( + "Run a Codex session. Accepts configuration parameters matching the Codex Config struct." + .to_string(), + ), + annotations: None, + } +} + +impl CodexToolCallParam { + /// Returns the initial user prompt to start the Codex conversation and the + /// Config. + pub fn into_config(self) -> std::io::Result<(String, codex_core::config::Config)> { + let Self { + prompt, + model, + cwd, + approval_policy, + sandbox_permissions, + disable_response_storage, + } = self; + let sandbox_policy = sandbox_permissions.map(|perms| { + SandboxPolicy::from(perms.into_iter().map(Into::into).collect::>()) + }); + + // Build ConfigOverrides recognised by codex-core. + let overrides = codex_core::config::ConfigOverrides { + model, + cwd: cwd.map(PathBuf::from), + approval_policy: approval_policy.map(Into::into), + sandbox_policy, + disable_response_storage, + }; + + let cfg = codex_core::config::Config::load_with_overrides(overrides)?; + + Ok((prompt, cfg)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + + /// We include a test to verify the exact JSON schema as "executable + /// documentation" for the schema. When can track changes to this test as a + /// way to audit changes to the generated schema. + /// + /// Seeing the fully expanded schema makes it easier to casually verify that + /// the generated JSON for enum types such as "approval-policy" is compact. + /// Ideally, modelcontextprotocol/inspector would provide a simpler UI for + /// enum fields versus open string fields to take advantage of this. + /// + /// As of 2025-05-04, there is an open PR for this: + /// https://github.com/modelcontextprotocol/inspector/pull/196 + #[test] + fn verify_codex_tool_json_schema() { + let tool = create_tool_for_codex_tool_call_param(); + let tool_json = serde_json::to_value(&tool).expect("tool serializes"); + let expected_tool_json = serde_json::json!({ + "name": "codex", + "description": "Run a Codex session. Accepts configuration parameters matching the Codex Config struct.", + "inputSchema": { + "type": "object", + "properties": { + "approval-policy": { + "description": "Execution approval policy expressed as the kebab-case variant name (`unless-allow-listed`, `auto-edit`, `on-failure`, `never`).", + "enum": [ + "auto-edit", + "unless-allow-listed", + "on-failure", + "never" + ], + "type": "string" + }, + "cwd": { + "description": "Working directory for the session. If relative, it is resolved against the server process's current working directory.", + "type": "string" + }, + "disable-response-storage": { + "description": "Disable server-side response storage.", + "type": "boolean" + }, + "model": { + "description": "Optional override for the model name (e.g. \"o3\", \"o4-mini\")", + "type": "string" + }, + "prompt": { + "description": "The *initial user prompt* to start the Codex conversation.", + "type": "string" + }, + "sandbox-permissions": { + "description": "Sandbox permissions using the same string values accepted by the CLI (e.g. \"disk-write-cwd\", \"network-full-access\").", + "items": { + "enum": [ + "disk-full-read-access", + "disk-write-cwd", + "disk-write-platform-user-temp-folder", + "disk-write-platform-global-temp-folder", + "disk-full-write-access", + "network-full-access" + ], + "type": "string" + }, + "type": "array" + } + }, + "required": [ + "prompt" + ] + } + }); + assert_eq!(expected_tool_json, tool_json); + } +} diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs new file mode 100644 index 0000000000..c35b855c49 --- /dev/null +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -0,0 +1,181 @@ +//! Asynchronous worker that executes a **Codex** tool-call inside a spawned +//! Tokio task. Separated from `message_processor.rs` to keep that file small +//! and to make future feature-growth easier to manage. + +use codex_core::codex_wrapper::init_codex; +use codex_core::config::Config as CodexConfig; +use codex_core::protocol::Event; +use codex_core::protocol::EventMsg; +use codex_core::protocol::InputItem; +use codex_core::protocol::Op; +use mcp_types::CallToolResult; +use mcp_types::CallToolResultContent; +use mcp_types::JSONRPCMessage; +use mcp_types::JSONRPCResponse; +use mcp_types::RequestId; +use mcp_types::TextContent; +use mcp_types::JSONRPC_VERSION; +use tokio::sync::mpsc::Sender; + +/// Convert a Codex [`Event`] to an MCP notification. +fn codex_event_to_notification(event: &Event) -> JSONRPCMessage { + JSONRPCMessage::Notification(mcp_types::JSONRPCNotification { + jsonrpc: JSONRPC_VERSION.into(), + method: "codex/event".into(), + params: Some(serde_json::to_value(event).expect("Event must serialize")), + }) +} + +/// Run a complete Codex session and stream events back to the client. +/// +/// On completion (success or error) the function sends the appropriate +/// `tools/call` response so the LLM can continue the conversation. +pub async fn run_codex_tool_session( + id: RequestId, + initial_prompt: String, + config: CodexConfig, + outgoing: Sender, +) { + let (codex, first_event, _ctrl_c) = match init_codex(config).await { + Ok(res) => res, + Err(e) => { + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: format!("Failed to start Codex session: {e}"), + annotations: None, + })], + is_error: Some(true), + }; + let _ = outgoing + .send(JSONRPCMessage::Response(JSONRPCResponse { + jsonrpc: JSONRPC_VERSION.into(), + id, + result: result.into(), + })) + .await; + return; + } + }; + + // Send initial SessionConfigured event. + let _ = outgoing + .send(codex_event_to_notification(&first_event)) + .await; + + if let Err(e) = codex + .submit(Op::UserInput { + items: vec![InputItem::Text { + text: initial_prompt.clone(), + }], + }) + .await + { + tracing::error!("Failed to submit initial prompt: {e}"); + } + + let mut last_agent_message: Option = None; + + // Stream events until the task needs to pause for user interaction or + // completes. + loop { + match codex.next_event().await { + Ok(event) => { + let _ = outgoing.send(codex_event_to_notification(&event)).await; + + match &event.msg { + EventMsg::AgentMessage { message } => { + last_agent_message = Some(message.clone()); + } + EventMsg::ExecApprovalRequest { .. } => { + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: "EXEC_APPROVAL_REQUIRED".to_string(), + annotations: None, + })], + is_error: None, + }; + let _ = outgoing + .send(JSONRPCMessage::Response(JSONRPCResponse { + jsonrpc: JSONRPC_VERSION.into(), + id: id.clone(), + result: result.into(), + })) + .await; + break; + } + EventMsg::ApplyPatchApprovalRequest { .. } => { + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: "PATCH_APPROVAL_REQUIRED".to_string(), + annotations: None, + })], + is_error: None, + }; + let _ = outgoing + .send(JSONRPCMessage::Response(JSONRPCResponse { + jsonrpc: JSONRPC_VERSION.into(), + id: id.clone(), + result: result.into(), + })) + .await; + break; + } + EventMsg::TaskComplete => { + let result = if let Some(msg) = last_agent_message { + CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: msg, + annotations: None, + })], + is_error: None, + } + } else { + CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: String::new(), + annotations: None, + })], + is_error: None, + } + }; + let _ = outgoing + .send(JSONRPCMessage::Response(JSONRPCResponse { + jsonrpc: JSONRPC_VERSION.into(), + id: id.clone(), + result: result.into(), + })) + .await; + break; + } + EventMsg::SessionConfigured { .. } => { + tracing::error!("unexpected SessionConfigured event"); + } + _ => {} + } + } + Err(e) => { + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: format!("Codex runtime error: {e}"), + annotations: None, + })], + is_error: Some(true), + }; + let _ = outgoing + .send(JSONRPCMessage::Response(JSONRPCResponse { + jsonrpc: JSONRPC_VERSION.into(), + id: id.clone(), + result: result.into(), + })) + .await; + break; + } + } + } +} diff --git a/codex-rs/mcp-server/src/main.rs b/codex-rs/mcp-server/src/main.rs index b0fb7fece5..87e8d7bbe2 100644 --- a/codex-rs/mcp-server/src/main.rs +++ b/codex-rs/mcp-server/src/main.rs @@ -1,4 +1,5 @@ //! Prototype MCP server. +#![deny(clippy::print_stdout, clippy::print_stderr)] use std::io::Result as IoResult; @@ -12,7 +13,10 @@ use tracing::debug; use tracing::error; use tracing::info; +mod codex_tool_config; +mod codex_tool_runner; mod message_processor; + use crate::message_processor::MessageProcessor; /// Size of the bounded channels used to communicate between tasks. The value diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index 6fcdc75dd5..5fa2085a15 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -1,6 +1,9 @@ -//! Very small proof-of-concept request router for the MCP prototype server. +use crate::codex_tool_config::create_tool_for_codex_tool_call_param; +use crate::codex_tool_config::CodexToolCallParam; +use codex_core::config::Config as CodexConfig; use mcp_types::CallToolRequestParams; +use mcp_types::CallToolResult; use mcp_types::CallToolResultContent; use mcp_types::ClientRequest; use mcp_types::JSONRPCBatchRequest; @@ -17,11 +20,10 @@ use mcp_types::RequestId; use mcp_types::ServerCapabilitiesTools; use mcp_types::ServerNotification; use mcp_types::TextContent; -use mcp_types::Tool; -use mcp_types::ToolInputSchema; use mcp_types::JSONRPC_VERSION; use serde_json::json; use tokio::sync::mpsc; +use tokio::task; pub(crate) struct MessageProcessor { outgoing: mpsc::Sender, @@ -303,21 +305,7 @@ impl MessageProcessor { ) { tracing::trace!("tools/list -> {params:?}"); let result = ListToolsResult { - tools: vec![Tool { - name: "echo".to_string(), - input_schema: ToolInputSchema { - r#type: "object".to_string(), - properties: Some(json!({ - "input": { - "type": "string", - "description": "The input to echo back" - } - })), - required: Some(vec!["input".to_string()]), - }, - description: Some("Echoes the request back".to_string()), - annotations: None, - }], + tools: vec![create_tool_for_codex_tool_call_param()], next_cursor: None, }; @@ -331,26 +319,80 @@ impl MessageProcessor { ) { tracing::info!("tools/call -> params: {:?}", params); let CallToolRequestParams { name, arguments } = params; - match name.as_str() { - "echo" => { - let result = mcp_types::CallToolResult { + + // We only support the "codex" tool for now. + if name != "codex" { + // Tool not found – return error result so the LLM can react. + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_string(), + text: format!("Unknown tool '{name}'"), + annotations: None, + })], + is_error: Some(true), + }; + self.send_response::(id, result); + return; + } + + let (initial_prompt, config): (String, CodexConfig) = match arguments { + Some(json_val) => match serde_json::from_value::(json_val) { + Ok(tool_cfg) => match tool_cfg.into_config() { + Ok(cfg) => cfg, + Err(e) => { + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_owned(), + text: format!( + "Failed to load Codex configuration from overrides: {e}" + ), + annotations: None, + })], + is_error: Some(true), + }; + self.send_response::(id, result); + return; + } + }, + Err(e) => { + let result = CallToolResult { + content: vec![CallToolResultContent::TextContent(TextContent { + r#type: "text".to_owned(), + text: format!("Failed to parse configuration for Codex tool: {e}"), + annotations: None, + })], + is_error: Some(true), + }; + self.send_response::(id, result); + return; + } + }, + None => { + let result = CallToolResult { content: vec![CallToolResultContent::TextContent(TextContent { r#type: "text".to_string(), - text: format!("Echo: {arguments:?}"), + text: + "Missing arguments for codex tool-call; the `prompt` field is required." + .to_string(), annotations: None, })], - is_error: None, - }; - self.send_response::(id, result); - } - _ => { - let result = mcp_types::CallToolResult { - content: vec![], is_error: Some(true), }; self.send_response::(id, result); + return; } - } + }; + + // Clone outgoing sender to move into async task. + let outgoing = self.outgoing.clone(); + + // Spawn an async task to handle the Codex session so that we do not + // block the synchronous message-processing loop. + task::spawn(async move { + // Run the Codex session and stream events back to the client. + crate::codex_tool_runner::run_codex_tool_session(id, initial_prompt, config, outgoing) + .await; + }); } fn handle_set_level( From 306d0b561882dd681313ad98a663cee55180a72b Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Mon, 5 May 2025 07:26:50 -0700 Subject: [PATCH 2/2] feat: mcp-client --- codex-rs/Cargo.lock | 15 ++ codex-rs/Cargo.toml | 1 + codex-rs/mcp-client/Cargo.toml | 23 +++ codex-rs/mcp-client/src/lib.rs | 301 ++++++++++++++++++++++++++++++++ codex-rs/mcp-client/src/main.rs | 43 +++++ 5 files changed, 383 insertions(+) create mode 100644 codex-rs/mcp-client/Cargo.toml create mode 100644 codex-rs/mcp-client/src/lib.rs create mode 100644 codex-rs/mcp-client/src/main.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0a4d879746..4b73372fb6 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -556,6 +556,21 @@ dependencies = [ "tempfile", ] +[[package]] +name = "codex-mcp-client" +version = "0.1.0" +dependencies = [ + "anyhow", + "codex-core", + "mcp-types", + "pretty_assertions", + "serde", + "serde_json", + "tokio", + "tracing", + "tracing-subscriber", +] + [[package]] name = "codex-mcp-server" version = "0.1.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 55aab2101b..9afcc11f4c 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -7,6 +7,7 @@ members = [ "core", "exec", "execpolicy", + "mcp-client", "mcp-server", "mcp-types", "tui", diff --git a/codex-rs/mcp-client/Cargo.toml b/codex-rs/mcp-client/Cargo.toml new file mode 100644 index 0000000000..2101a1e697 --- /dev/null +++ b/codex-rs/mcp-client/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "codex-mcp-client" +version = "0.1.0" +edition = "2021" + +[dependencies] +anyhow = "1" +codex-core = { path = "../core", features = ["cli"] } +mcp-types = { path = "../mcp-types" } +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tracing = { version = "0.1.41", features = ["log"] } +tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } +tokio = { version = "1", features = [ + "io-std", + "macros", + "process", + "rt-multi-thread", + "signal", +] } + +[dev-dependencies] +pretty_assertions = "1.4.1" diff --git a/codex-rs/mcp-client/src/lib.rs b/codex-rs/mcp-client/src/lib.rs new file mode 100644 index 0000000000..be41dfaa09 --- /dev/null +++ b/codex-rs/mcp-client/src/lib.rs @@ -0,0 +1,301 @@ +//! A minimal async client for the Model Context Protocol (MCP). +//! +//! The client is intentionally lightweight – it is only capable of: +//! 1. Spawning a subprocess (typically `codex-mcp-server`) whose STDIN/STDOUT +//! transports newline-delimited JSON-RPC messages. +//! 2. Sending MCP requests and pairing them with their corresponding +//! responses. +//! 3. Offering a convenience helper for the common `tools/list` request. +//! +//! The crate hides all JSON‐RPC framing details behind a typed API. Users +//! interact with the [`ModelContextProtocolRequest`] trait from `mcp-types` to +//! issue requests and receive strongly-typed results. + +use std::collections::HashMap; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; +use std::sync::Arc; + +use anyhow::anyhow; +use anyhow::Result; +use mcp_types::JSONRPCMessage; +use mcp_types::JSONRPCNotification; +use mcp_types::JSONRPCRequest; +use mcp_types::JSONRPCResponse; +use mcp_types::ListToolsRequest; +use mcp_types::ListToolsRequestParams; +use mcp_types::ListToolsResult; +use mcp_types::ModelContextProtocolRequest; +use mcp_types::RequestId; +use mcp_types::JSONRPC_VERSION; +use serde::de::DeserializeOwned; +use serde::Serialize; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +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; + +/// Capacity of the bounded channels used for transporting messages between the +/// client API and the IO tasks. +const CHANNEL_CAPACITY: usize = 128; + +/// Internal representation of a pending request sender. +type PendingSender = oneshot::Sender; + +/// A running MCP client instance. +pub struct McpClient { + /// Channel for sending JSON-RPC messages *to* the background writer task. + outgoing_tx: mpsc::Sender, + + /// Map of `request.id -> oneshot::Sender` used to dispatch responses back + /// to the originating caller. + pending: Arc>>, + + /// Monotonically increasing counter used to generate request IDs. + id_counter: AtomicI64, +} + +impl McpClient { + /// Spawn the given command and establish an MCP session over its STDIO. + /// + /// `args` follows the Unix convention where the first element is the + /// executable path and the rest are arguments. For example: + /// + /// ```no_run + /// # use codex_mcp_client::McpClient; + /// # async fn run() -> anyhow::Result<()> { + /// let client = McpClient::new_stdio_client(vec![ + /// "codex-mcp-server".to_string(), + /// ]).await?; + /// # Ok(()) } + /// ``` + pub async fn new_stdio_client(args: Vec) -> std::io::Result { + if args.is_empty() { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "expected at least one element in `args` - the program to spawn", + )); + } + + let program = &args[0]; + let mut command = Command::new(program); + if args.len() > 1 { + command.args(&args[1..]); + } + + command.stdin(std::process::Stdio::piped()); + command.stdout(std::process::Stdio::piped()); + // As noted in the `kill_on_drop` documentation, the Tokio runtime makes + // a "best effort" to reap-after-exit to avoid zombie processes, but it + // is not a guarantee. + command.kill_on_drop(true); + let mut child = command.spawn()?; + + let stdin = child.stdin.take().ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::Other, "failed to capture child stdin") + })?; + let stdout = child.stdout.take().ok_or_else(|| { + std::io::Error::new(std::io::ErrorKind::Other, "failed to capture child stdout") + })?; + + // Because we have invoked take() on both stdin and stdout, calling + // `child.wait()` will not close the pipes now owned by tokio tasks. + // We invoke `wait()` proactively to ensure the child process is reaped. + tokio::spawn(async move { + let _ = child.wait().await; + }); + + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); + let pending: Arc>> = Arc::new(Mutex::new(HashMap::new())); + + // Spawn writer task. It listens on the `outgoing_rx` channel and + // writes messages to the child's STDIN. + let writer_handle = { + let mut stdin = stdin; + tokio::spawn(async move { + while let Some(msg) = outgoing_rx.recv().await { + match serde_json::to_string(&msg) { + Ok(json) => { + if stdin.write_all(json.as_bytes()).await.is_err() { + error!("failed to write message to child stdin"); + break; + } + if stdin.write_all(b"\n").await.is_err() { + error!("failed to write newline to child stdin"); + break; + } + if stdin.flush().await.is_err() { + error!("failed to flush child stdin"); + break; + } + } + Err(e) => error!("failed to serialize JSONRPCMessage: {e}"), + } + } + }) + }; + + // Spawn reader task. It reads line-delimited JSON from the child's + // STDOUT and dispatches responses to the pending map. + let reader_handle = { + let pending = pending.clone(); + let mut lines = BufReader::new(stdout).lines(); + + tokio::spawn(async move { + while let Ok(Some(line)) = lines.next_line().await { + match serde_json::from_str::(&line) { + Ok(JSONRPCMessage::Response(resp)) => { + Self::dispatch_response(resp, &pending).await; + } + Ok(JSONRPCMessage::Error(err)) => { + Self::dispatch_error(err, &pending).await; + } + Ok(JSONRPCMessage::Notification(JSONRPCNotification { .. })) => { + // For now we only log server-initiated notifications. + info!("<- notification: {}", line); + } + Ok(other) => { + // Batch responses and requests are currently not + // expected from the server – log and ignore. + info!("<- unhandled message: {:?}", other); + } + Err(e) => { + error!("failed to deserialize JSONRPCMessage: {e}; line = {}", line) + } + } + } + }) + }; + + // We intentionally *detach* the tasks. They will keep running in the + // background as long as their respective resources (channels/stdin/ + // stdout) are alive. Dropping `McpClient` cancels the tasks due to + // dropped resources. + let _ = (writer_handle, reader_handle); + + Ok(Self { + outgoing_tx, + pending, + id_counter: AtomicI64::new(0), + }) + } + + /// Send an arbitrary MCP request and await the typed result. + pub async fn send_request(&self, params: R::Params) -> Result + where + R: ModelContextProtocolRequest, + R::Params: Serialize, + R::Result: DeserializeOwned, + { + // Create a new unique ID. + let id = self.id_counter.fetch_add(1, Ordering::SeqCst); + let request_id = RequestId::Integer(id); + + // Serialize params -> JSON. For many request types `Params` is + // `Option` and `None` should be encoded as *absence* of the field. + let params_json = serde_json::to_value(¶ms)?; + let params_field = if params_json.is_null() { + None + } else { + Some(params_json) + }; + + let jsonrpc_request = JSONRPCRequest { + id: request_id.clone(), + jsonrpc: JSONRPC_VERSION.to_string(), + method: R::METHOD.to_string(), + params: params_field, + }; + + let message = JSONRPCMessage::Request(jsonrpc_request); + + // oneshot channel for the response. + let (tx, rx) = oneshot::channel(); + + // Register in pending map *before* sending the message so a race where + // the response arrives immediately cannot be lost. + { + let mut guard = self.pending.lock().await; + guard.insert(id, tx); + } + + // 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" + )); + } + + // Await the response. + let msg = rx + .await + .map_err(|_| anyhow!("response channel closed before a reply was received"))?; + + match msg { + JSONRPCMessage::Response(JSONRPCResponse { result, .. }) => { + let typed: R::Result = serde_json::from_value(result)?; + Ok(typed) + } + JSONRPCMessage::Error(err) => Err(anyhow!(format!( + "server returned JSON-RPC error: code = {}, message = {}", + err.error.code, err.error.message + ))), + other => Err(anyhow!(format!( + "unexpected message variant received in reply path: {:?}", + other + ))), + } + } + + /// Convenience wrapper around `tools/list`. + pub async fn list_tools( + &self, + params: Option, + ) -> Result { + self.send_request::(params).await + } + + /// Internal helper: route a JSON-RPC *response* object to the pending map. + async fn dispatch_response( + resp: JSONRPCResponse, + pending: &Arc>>, + ) { + let id = match resp.id { + RequestId::Integer(i) => i, + RequestId::String(_) => { + // We only ever generate integer IDs. Receiving a string here + // means we will not find a matching entry in `pending`. + debug!("response with string ID - no matching pending request"); + return; + } + }; + + if let Some(tx) = pending.lock().await.remove(&id) { + // Ignore send errors – the receiver might have been dropped. + let _ = tx.send(JSONRPCMessage::Response(resp)); + } else { + debug!(id, "no pending request found for response"); + } + } + + /// Internal helper: route a JSON-RPC *error* object to the pending map. + async fn dispatch_error( + err: mcp_types::JSONRPCError, + pending: &Arc>>, + ) { + let id = match err.id { + RequestId::Integer(i) => i, + RequestId::String(_) => return, // see comment above + }; + + if let Some(tx) = pending.lock().await.remove(&id) { + let _ = tx.send(JSONRPCMessage::Error(err)); + } + } +} diff --git a/codex-rs/mcp-client/src/main.rs b/codex-rs/mcp-client/src/main.rs new file mode 100644 index 0000000000..fe8c0f6600 --- /dev/null +++ b/codex-rs/mcp-client/src/main.rs @@ -0,0 +1,43 @@ +//! Simple command-line utility to exercise `McpClient`. +//! +//! Example usage: +//! +//! ```bash +//! cargo run -p codex-mcp-client -- `codex-mcp-server` +//! ``` +//! +//! Any additional arguments after the first one are forwarded to the spawned +//! program. The utility connects, issues a `tools/list` request and prints the +//! server's response as pretty JSON. + +use anyhow::Context; +use anyhow::Result; +use codex_mcp_client::McpClient; +use mcp_types::ListToolsRequestParams; + +#[tokio::main] +async fn main() -> Result<()> { + // Collect command-line arguments excluding the program name itself. + let cmd_args: Vec = std::env::args().skip(1).collect(); + + if cmd_args.is_empty() || cmd_args[0] == "--help" || cmd_args[0] == "-h" { + eprintln!("Usage: mcp-client [args..]\n\nExample: mcp-client codex-mcp-server"); + std::process::exit(1); + } + + // Spawn the subprocess and connect the client. + let client = McpClient::new_stdio_client(cmd_args.clone()) + .await + .with_context(|| format!("failed to spawn subprocess: {:?}", cmd_args))?; + + // Issue `tools/list` request (no params). + let tools = client + .list_tools(None::) + .await + .context("tools/list request failed")?; + + // Print the result in a human readable form. + println!("{}", serde_json::to_string_pretty(&tools)?); + + Ok(()) +}