From 11fd3123beb2f09371e4ab0f6568673236130457 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Sat, 19 Jul 2025 00:30:56 -0400 Subject: [PATCH 1/2] chore: introduce OutgoingMessageSender (#1622) Previous to this change, `MessageProcessor` had a `tokio::sync::mpsc::Sender` as an abstraction for server code to send a message down to the MCP client. Because `Sender` is cheap to `clone()`, it was straightforward to make it available to tasks scheduled with `tokio::task::spawn()`. This worked well when we were only sending notifications or responses back down to the client, but we want to add support for sending elicitations in #1623, which means that we need to be able to send _requests_ to the client, and now we need a bit of centralization to ensure all request ids are unique. To that end, this PR introduces `OutgoingMessageSender`, which houses the existing `Sender` as well as an `AtomicI64` to mint out new, unique request ids. It has methods like `send_request()` and `send_response()` so that callers do not have to deal with `JSONRPCMessage` directly, as having to set the `jsonrpc` for each message was a bit tedious (this cleans up `codex_tool_runner.rs` quite a bit). We do not have `OutgoingMessageSender` implement `Clone` because it is important that the `AtomicI64` is shared across all users of `OutgoingMessageSender`. As such, `Arc` must be used instead, as it is frequently shared with new tokio tasks. As part of this change, we update `message_processor.rs` to embrace `await`, though we must be careful that no individual handler blocks the main loop and prevents other messages from being handled. --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/openai/codex/pull/1622). * #1623 * __->__ #1622 * #1621 * #1620 --- codex-rs/mcp-server/src/codex_tool_runner.rs | 65 ++-------- codex-rs/mcp-server/src/lib.rs | 13 +- codex-rs/mcp-server/src/message_processor.rs | 82 ++++++------ codex-rs/mcp-server/src/outgoing_message.rs | 129 +++++++++++++++++++ 4 files changed, 186 insertions(+), 103 deletions(-) create mode 100644 codex-rs/mcp-server/src/outgoing_message.rs diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index 00cadcf0d8..a20566d61c 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -2,10 +2,11 @@ //! Tokio task. Separated from `message_processor.rs` to keep that file small //! and to make future feature-growth easier to manage. +use std::sync::Arc; + use codex_core::codex_wrapper::init_codex; use codex_core::config::Config as CodexConfig; use codex_core::protocol::AgentMessageEvent; -use codex_core::protocol::Event; use codex_core::protocol::EventMsg; use codex_core::protocol::InputItem; use codex_core::protocol::Op; @@ -13,22 +14,10 @@ use codex_core::protocol::Submission; use codex_core::protocol::TaskCompleteEvent; use mcp_types::CallToolResult; use mcp_types::ContentBlock; -use mcp_types::JSONRPC_VERSION; -use mcp_types::JSONRPCMessage; -use mcp_types::JSONRPCResponse; use mcp_types::RequestId; use mcp_types::TextContent; -use tokio::sync::mpsc::Sender; -/// Convert a Codex [`Event`] to an MCP notification. -fn codex_event_to_notification(event: &Event) -> JSONRPCMessage { - #[expect(clippy::expect_used)] - JSONRPCMessage::Notification(mcp_types::JSONRPCNotification { - jsonrpc: JSONRPC_VERSION.into(), - method: "codex/event".into(), - params: Some(serde_json::to_value(event).expect("Event must serialize")), - }) -} +use crate::outgoing_message::OutgoingMessageSender; /// Run a complete Codex session and stream events back to the client. /// @@ -38,7 +27,7 @@ pub async fn run_codex_tool_session( id: RequestId, initial_prompt: String, config: CodexConfig, - outgoing: Sender, + outgoing: Arc, ) { let (codex, first_event, _ctrl_c) = match init_codex(config).await { Ok(res) => res, @@ -52,21 +41,13 @@ pub async fn run_codex_tool_session( is_error: Some(true), structured_content: None, }; - let _ = outgoing - .send(JSONRPCMessage::Response(JSONRPCResponse { - jsonrpc: JSONRPC_VERSION.into(), - id, - result: result.into(), - })) - .await; + outgoing.send_response(id.clone(), result.into()).await; return; } }; // Send initial SessionConfigured event. - let _ = outgoing - .send(codex_event_to_notification(&first_event)) - .await; + outgoing.send_event_as_notification(&first_event).await; // Use the original MCP request ID as the `sub_id` for the Codex submission so that // any events emitted for this tool-call can be correlated with the @@ -94,7 +75,7 @@ pub async fn run_codex_tool_session( loop { match codex.next_event().await { Ok(event) => { - let _ = outgoing.send(codex_event_to_notification(&event)).await; + outgoing.send_event_as_notification(&event).await; match &event.msg { EventMsg::ExecApprovalRequest(_) => { @@ -107,13 +88,7 @@ pub async fn run_codex_tool_session( is_error: None, structured_content: None, }; - let _ = outgoing - .send(JSONRPCMessage::Response(JSONRPCResponse { - jsonrpc: JSONRPC_VERSION.into(), - id: id.clone(), - result: result.into(), - })) - .await; + outgoing.send_response(id.clone(), result.into()).await; break; } EventMsg::ApplyPatchApprovalRequest(_) => { @@ -126,13 +101,7 @@ pub async fn run_codex_tool_session( is_error: None, structured_content: None, }; - let _ = outgoing - .send(JSONRPCMessage::Response(JSONRPCResponse { - jsonrpc: JSONRPC_VERSION.into(), - id: id.clone(), - result: result.into(), - })) - .await; + outgoing.send_response(id.clone(), result.into()).await; break; } EventMsg::TaskComplete(TaskCompleteEvent { last_agent_message }) => { @@ -149,13 +118,7 @@ pub async fn run_codex_tool_session( is_error: None, structured_content: None, }; - let _ = outgoing - .send(JSONRPCMessage::Response(JSONRPCResponse { - jsonrpc: JSONRPC_VERSION.into(), - id: id.clone(), - result: result.into(), - })) - .await; + outgoing.send_response(id.clone(), result.into()).await; break; } EventMsg::SessionConfigured(_) => { @@ -203,13 +166,7 @@ pub async fn run_codex_tool_session( // structured way. structured_content: None, }; - let _ = outgoing - .send(JSONRPCMessage::Response(JSONRPCResponse { - jsonrpc: JSONRPC_VERSION.into(), - id: id.clone(), - result: result.into(), - })) - .await; + outgoing.send_response(id.clone(), result.into()).await; break; } } diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index db41013ab6..b968b4976e 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -18,8 +18,11 @@ mod codex_tool_config; mod codex_tool_runner; mod json_to_toml; mod message_processor; +mod outgoing_message; use crate::message_processor::MessageProcessor; +use crate::outgoing_message::OutgoingMessage; +use crate::outgoing_message::OutgoingMessageSender; /// Size of the bounded channels used to communicate between tasks. The value /// is a balance between throughput and memory usage – 128 messages should be @@ -35,7 +38,7 @@ pub async fn run_main(codex_linux_sandbox_exe: Option) -> IoResult<()> // Set up channels. let (incoming_tx, mut incoming_rx) = mpsc::channel::(CHANNEL_CAPACITY); - let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(CHANNEL_CAPACITY); // Task: read from stdin, push to `incoming_tx`. let stdin_reader_handle = tokio::spawn({ @@ -63,11 +66,12 @@ pub async fn run_main(codex_linux_sandbox_exe: Option) -> IoResult<()> // Task: process incoming messages. let processor_handle = tokio::spawn({ - let mut processor = MessageProcessor::new(outgoing_tx.clone(), codex_linux_sandbox_exe); + let outgoing_message_sender = OutgoingMessageSender::new(outgoing_tx); + let mut processor = MessageProcessor::new(outgoing_message_sender, codex_linux_sandbox_exe); async move { while let Some(msg) = incoming_rx.recv().await { match msg { - JSONRPCMessage::Request(r) => processor.process_request(r), + JSONRPCMessage::Request(r) => processor.process_request(r).await, JSONRPCMessage::Response(r) => processor.process_response(r), JSONRPCMessage::Notification(n) => processor.process_notification(n), JSONRPCMessage::Error(e) => processor.process_error(e), @@ -81,7 +85,8 @@ pub async fn run_main(codex_linux_sandbox_exe: Option) -> IoResult<()> // Task: write outgoing messages to stdout. let stdout_writer_handle = tokio::spawn(async move { let mut stdout = io::stdout(); - while let Some(msg) = outgoing_rx.recv().await { + while let Some(outgoing_message) = outgoing_rx.recv().await { + let msg: JSONRPCMessage = outgoing_message.into(); match serde_json::to_string(&msg) { Ok(json) => { if let Err(e) = stdout.write_all(json.as_bytes()).await { diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index dcc6ae62f9..aad7f21132 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -1,17 +1,17 @@ use std::path::PathBuf; +use std::sync::Arc; use crate::codex_tool_config::CodexToolCallParam; use crate::codex_tool_config::create_tool_for_codex_tool_call_param; +use crate::outgoing_message::OutgoingMessageSender; use codex_core::config::Config as CodexConfig; use mcp_types::CallToolRequestParams; use mcp_types::CallToolResult; use mcp_types::ClientRequest; use mcp_types::ContentBlock; -use mcp_types::JSONRPC_VERSION; use mcp_types::JSONRPCError; use mcp_types::JSONRPCErrorError; -use mcp_types::JSONRPCMessage; use mcp_types::JSONRPCNotification; use mcp_types::JSONRPCRequest; use mcp_types::JSONRPCResponse; @@ -22,11 +22,10 @@ use mcp_types::ServerCapabilitiesTools; use mcp_types::ServerNotification; use mcp_types::TextContent; use serde_json::json; -use tokio::sync::mpsc; use tokio::task; pub(crate) struct MessageProcessor { - outgoing: mpsc::Sender, + outgoing: Arc, initialized: bool, codex_linux_sandbox_exe: Option, } @@ -35,17 +34,17 @@ impl MessageProcessor { /// Create a new `MessageProcessor`, retaining a handle to the outgoing /// `Sender` so handlers can enqueue messages to be written to stdout. pub(crate) fn new( - outgoing: mpsc::Sender, + outgoing: OutgoingMessageSender, codex_linux_sandbox_exe: Option, ) -> Self { Self { - outgoing, + outgoing: Arc::new(outgoing), initialized: false, codex_linux_sandbox_exe, } } - pub(crate) fn process_request(&mut self, request: JSONRPCRequest) { + pub(crate) async fn process_request(&mut self, request: JSONRPCRequest) { // Hold on to the ID so we can respond. let request_id = request.id.clone(); @@ -60,10 +59,10 @@ impl MessageProcessor { // Dispatch to a dedicated handler for each request type. match client_request { ClientRequest::InitializeRequest(params) => { - self.handle_initialize(request_id, params); + self.handle_initialize(request_id, params).await; } ClientRequest::PingRequest(params) => { - self.handle_ping(request_id, params); + self.handle_ping(request_id, params).await; } ClientRequest::ListResourcesRequest(params) => { self.handle_list_resources(params); @@ -87,10 +86,10 @@ impl MessageProcessor { self.handle_get_prompt(params); } ClientRequest::ListToolsRequest(params) => { - self.handle_list_tools(request_id, params); + self.handle_list_tools(request_id, params).await; } ClientRequest::CallToolRequest(params) => { - self.handle_call_tool(request_id, params); + self.handle_call_tool(request_id, params).await; } ClientRequest::SetLevelRequest(params) => { self.handle_set_level(params); @@ -148,7 +147,7 @@ impl MessageProcessor { tracing::error!("<- error: {:?}", err); } - fn handle_initialize( + async fn handle_initialize( &mut self, id: RequestId, params: ::Params, @@ -157,19 +156,12 @@ impl MessageProcessor { if self.initialized { // Already initialised: send JSON-RPC error response. - let error_msg = JSONRPCMessage::Error(JSONRPCError { - jsonrpc: JSONRPC_VERSION.into(), - id, - error: JSONRPCErrorError { - code: -32600, // Invalid Request - message: "initialize called more than once".to_string(), - data: None, - }, - }); - - if let Err(e) = self.outgoing.try_send(error_msg) { - tracing::error!("Failed to send initialization error: {e}"); - } + let error = JSONRPCErrorError { + code: -32600, // Invalid Request + message: "initialize called more than once".to_string(), + data: None, + }; + self.outgoing.send_error(id, error).await; return; } @@ -196,34 +188,29 @@ impl MessageProcessor { }, }; - self.send_response::(id, result); + self.send_response::(id, result) + .await; } - fn send_response(&self, id: RequestId, result: T::Result) + async fn send_response(&self, id: RequestId, result: T::Result) where T: ModelContextProtocolRequest, { // result has `Serialized` instance so should never fail #[expect(clippy::unwrap_used)] - let response = JSONRPCMessage::Response(JSONRPCResponse { - jsonrpc: JSONRPC_VERSION.into(), - id, - result: serde_json::to_value(result).unwrap(), - }); - - if let Err(e) = self.outgoing.try_send(response) { - tracing::error!("Failed to send response: {e}"); - } + let result = serde_json::to_value(result).unwrap(); + self.outgoing.send_response(id, result).await; } - fn handle_ping( + async fn handle_ping( &self, id: RequestId, params: ::Params, ) { tracing::info!("ping -> params: {:?}", params); let result = json!({}); - self.send_response::(id, result); + self.send_response::(id, result) + .await; } fn handle_list_resources( @@ -276,7 +263,7 @@ impl MessageProcessor { tracing::info!("prompts/get -> params: {:?}", params); } - fn handle_list_tools( + async fn handle_list_tools( &self, id: RequestId, params: ::Params, @@ -287,10 +274,11 @@ impl MessageProcessor { next_cursor: None, }; - self.send_response::(id, result); + self.send_response::(id, result) + .await; } - fn handle_call_tool( + async fn handle_call_tool( &self, id: RequestId, params: ::Params, @@ -310,7 +298,8 @@ impl MessageProcessor { is_error: Some(true), structured_content: None, }; - self.send_response::(id, result); + self.send_response::(id, result) + .await; return; } @@ -330,7 +319,8 @@ impl MessageProcessor { is_error: Some(true), structured_content: None, }; - self.send_response::(id, result); + self.send_response::(id, result) + .await; return; } }, @@ -344,7 +334,8 @@ impl MessageProcessor { is_error: Some(true), structured_content: None, }; - self.send_response::(id, result); + self.send_response::(id, result) + .await; return; } }, @@ -360,7 +351,8 @@ impl MessageProcessor { is_error: Some(true), structured_content: None, }; - self.send_response::(id, result); + self.send_response::(id, result) + .await; return; } }; diff --git a/codex-rs/mcp-server/src/outgoing_message.rs b/codex-rs/mcp-server/src/outgoing_message.rs new file mode 100644 index 0000000000..93a760d36b --- /dev/null +++ b/codex-rs/mcp-server/src/outgoing_message.rs @@ -0,0 +1,129 @@ +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_core::protocol::Event; +use mcp_types::JSONRPC_VERSION; +use mcp_types::JSONRPCError; +use mcp_types::JSONRPCErrorError; +use mcp_types::JSONRPCMessage; +use mcp_types::JSONRPCNotification; +use mcp_types::JSONRPCRequest; +use mcp_types::JSONRPCResponse; +use mcp_types::RequestId; +use mcp_types::Result; +use serde::Serialize; +use tokio::sync::mpsc; + +pub(crate) struct OutgoingMessageSender { + next_request_id: AtomicI64, + sender: mpsc::Sender, +} + +impl OutgoingMessageSender { + pub(crate) fn new(sender: mpsc::Sender) -> Self { + Self { + next_request_id: AtomicI64::new(0), + sender, + } + } + + #[allow(dead_code)] + pub(crate) async fn send_request(&self, method: &str, params: Option) { + let outgoing_message = OutgoingMessage::Request(OutgoingRequest { + id: RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::Relaxed)), + method: method.to_string(), + params, + }); + let _ = self.sender.send(outgoing_message).await; + } + + pub(crate) async fn send_response(&self, id: RequestId, result: Result) { + let outgoing_message = OutgoingMessage::Response(OutgoingResponse { id, result }); + let _ = self.sender.send(outgoing_message).await; + } + + pub(crate) async fn send_event_as_notification(&self, event: &Event) { + #[expect(clippy::expect_used)] + let params = Some(serde_json::to_value(event).expect("Event must serialize")); + let outgoing_message = OutgoingMessage::Notification(OutgoingNotification { + method: "codex/event".to_string(), + params, + }); + let _ = self.sender.send(outgoing_message).await; + } + + pub(crate) async fn send_error(&self, id: RequestId, error: JSONRPCErrorError) { + let outgoing_message = OutgoingMessage::Error(OutgoingError { id, error }); + let _ = self.sender.send(outgoing_message).await; + } +} + +/// Outgoing message from the server to the client. +pub(crate) enum OutgoingMessage { + Request(OutgoingRequest), + Notification(OutgoingNotification), + Response(OutgoingResponse), + Error(OutgoingError), +} + +impl From for JSONRPCMessage { + fn from(val: OutgoingMessage) -> Self { + use OutgoingMessage::*; + match val { + Request(OutgoingRequest { id, method, params }) => { + JSONRPCMessage::Request(JSONRPCRequest { + jsonrpc: JSONRPC_VERSION.into(), + id, + method, + params, + }) + } + Notification(OutgoingNotification { method, params }) => { + JSONRPCMessage::Notification(JSONRPCNotification { + jsonrpc: JSONRPC_VERSION.into(), + method, + params, + }) + } + Response(OutgoingResponse { id, result }) => { + JSONRPCMessage::Response(JSONRPCResponse { + jsonrpc: JSONRPC_VERSION.into(), + id, + result, + }) + } + Error(OutgoingError { id, error }) => JSONRPCMessage::Error(JSONRPCError { + jsonrpc: JSONRPC_VERSION.into(), + id, + error, + }), + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(crate) struct OutgoingRequest { + pub id: RequestId, + pub method: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub params: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(crate) struct OutgoingNotification { + pub method: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub params: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(crate) struct OutgoingResponse { + pub id: RequestId, + pub result: Result, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub(crate) struct OutgoingError { + pub error: JSONRPCErrorError, + pub id: RequestId, +} From 7d7d3084bccf927cb855efdd150f1966dbbdab87 Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Sat, 19 Jul 2025 00:31:04 -0400 Subject: [PATCH 2/2] feat: leverage elicitations in the MCP server --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/mcp_connection_manager.rs | 5 +- codex-rs/mcp-server/Cargo.toml | 1 + codex-rs/mcp-server/src/codex_tool_runner.rs | 79 ++++++++++++++++---- codex-rs/mcp-server/src/lib.rs | 2 +- codex-rs/mcp-server/src/message_processor.rs | 4 +- codex-rs/mcp-server/src/outgoing_message.rs | 42 ++++++++++- codex-rs/mcp-types/generate_mcp_types.py | 5 +- codex-rs/mcp-types/src/lib.rs | 4 +- 9 files changed, 121 insertions(+), 22 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index a25e0f8be0..9171369ae2 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -799,6 +799,7 @@ dependencies = [ "schemars 0.8.22", "serde", "serde_json", + "shlex", "tokio", "toml 0.9.1", "tracing", diff --git a/codex-rs/core/src/mcp_connection_manager.rs b/codex-rs/core/src/mcp_connection_manager.rs index cb91bc6127..1322306243 100644 --- a/codex-rs/core/src/mcp_connection_manager.rs +++ b/codex-rs/core/src/mcp_connection_manager.rs @@ -18,6 +18,7 @@ use mcp_types::ClientCapabilities; use mcp_types::Implementation; use mcp_types::Tool; +use serde_json::json; use sha1::Digest; use sha1::Sha1; use tokio::task::JoinSet; @@ -135,7 +136,9 @@ impl McpConnectionManager { experimental: None, roots: None, sampling: None, - elicitation: None, + // TODO(mbolin): Research the expected structure + // for this field. + elicitation: Some(json!({})), }, client_info: Implementation { name: "codex-mcp-client".to_owned(), diff --git a/codex-rs/mcp-server/Cargo.toml b/codex-rs/mcp-server/Cargo.toml index f91a3dc8f8..640a999317 100644 --- a/codex-rs/mcp-server/Cargo.toml +++ b/codex-rs/mcp-server/Cargo.toml @@ -22,6 +22,7 @@ mcp-types = { path = "../mcp-types" } schemars = "0.8.22" serde = { version = "1", features = ["derive"] } serde_json = "1" +shlex = "1.3.0" toml = "0.9" tracing = { version = "0.1.41", features = ["log"] } tracing-subscriber = { version = "0.3", features = ["fmt", "env-filter"] } diff --git a/codex-rs/mcp-server/src/codex_tool_runner.rs b/codex-rs/mcp-server/src/codex_tool_runner.rs index a20566d61c..9342a93ef0 100644 --- a/codex-rs/mcp-server/src/codex_tool_runner.rs +++ b/codex-rs/mcp-server/src/codex_tool_runner.rs @@ -4,18 +4,25 @@ use std::sync::Arc; +use codex_core::Codex; use codex_core::codex_wrapper::init_codex; use codex_core::config::Config as CodexConfig; use codex_core::protocol::AgentMessageEvent; use codex_core::protocol::EventMsg; +use codex_core::protocol::ExecApprovalRequestEvent; use codex_core::protocol::InputItem; use codex_core::protocol::Op; use codex_core::protocol::Submission; use codex_core::protocol::TaskCompleteEvent; use mcp_types::CallToolResult; use mcp_types::ContentBlock; +use mcp_types::ElicitRequest; +use mcp_types::ElicitRequestParamsRequestedSchema; +use mcp_types::ModelContextProtocolRequest; use mcp_types::RequestId; use mcp_types::TextContent; +use serde_json::json; +use tracing::error; use crate::outgoing_message::OutgoingMessageSender; @@ -45,6 +52,7 @@ pub async fn run_codex_tool_session( return; } }; + let codex = Arc::new(codex); // Send initial SessionConfigured event. outgoing.send_event_as_notification(&first_event).await; @@ -58,7 +66,7 @@ pub async fn run_codex_tool_session( }; let submission = Submission { - id: sub_id, + id: sub_id.clone(), op: Op::UserInput { items: vec![InputItem::Text { text: initial_prompt.clone(), @@ -77,18 +85,46 @@ pub async fn run_codex_tool_session( Ok(event) => { outgoing.send_event_as_notification(&event).await; - match &event.msg { - EventMsg::ExecApprovalRequest(_) => { - let result = CallToolResult { - content: vec![ContentBlock::TextContent(TextContent { - r#type: "text".to_string(), - text: "EXEC_APPROVAL_REQUIRED".to_string(), - annotations: None, - })], - is_error: None, - structured_content: None, - }; - outgoing.send_response(id.clone(), result.into()).await; + match event.msg { + EventMsg::ExecApprovalRequest(ExecApprovalRequestEvent { + command, + cwd, + reason: _, + }) => { + let escaped_command = shlex::try_join(command.iter().map(|s| s.as_str())) + .unwrap_or_else(|_| command.join(" ")); + let message = format!("Allow Codex to run `{escaped_command}` in {cwd:?}?"); + + // This `params` conforms to ElicitRequestParams, but + // contains additional metadata fields. + let params = json!({ + "message": message, + "requestedSchema": ElicitRequestParamsRequestedSchema { + properties: json!({}), + required: None, + r#type: "object".to_string(), + }, + + "codex-tool-call-id": sub_id, + "codex-event-id": event.id, + "codex-elicitation": "exec-approval", + "codex-command": command, + "codex-cwd": cwd.to_string_lossy().to_string() + }); + let on_response = outgoing + .send_request(ElicitRequest::METHOD, Some(params)) + .await; + + // Listen for the response on a separate task so we do + // not block the main loop of this function. + { + let outgoing = outgoing.clone(); + let codex = codex.clone(); + tokio::spawn(async move { + on_exec_approval_response(on_response, outgoing, codex) + }); + } + break; } EventMsg::ApplyPatchApprovalRequest(_) => { @@ -172,3 +208,20 @@ pub async fn run_codex_tool_session( } } } + +async fn on_exec_approval_response( + receiver: tokio::sync::oneshot::Receiver, + outgoing: Arc, + codex: Arc, +) { + let response = receiver.await; + let value = match response { + Ok(value) => value, + Err(err) => { + error!("request failed: {err:?}"); + return; + } + }; + + // Try to deserialize `value` and then make the appropriate call to `codex`. +} diff --git a/codex-rs/mcp-server/src/lib.rs b/codex-rs/mcp-server/src/lib.rs index b968b4976e..3b984ecf13 100644 --- a/codex-rs/mcp-server/src/lib.rs +++ b/codex-rs/mcp-server/src/lib.rs @@ -72,7 +72,7 @@ pub async fn run_main(codex_linux_sandbox_exe: Option) -> IoResult<()> while let Some(msg) = incoming_rx.recv().await { match msg { JSONRPCMessage::Request(r) => processor.process_request(r).await, - JSONRPCMessage::Response(r) => processor.process_response(r), + JSONRPCMessage::Response(r) => processor.process_response(r).await, JSONRPCMessage::Notification(n) => processor.process_notification(n), JSONRPCMessage::Error(e) => processor.process_error(e), } diff --git a/codex-rs/mcp-server/src/message_processor.rs b/codex-rs/mcp-server/src/message_processor.rs index aad7f21132..d994d8a7a2 100644 --- a/codex-rs/mcp-server/src/message_processor.rs +++ b/codex-rs/mcp-server/src/message_processor.rs @@ -101,8 +101,10 @@ impl MessageProcessor { } /// Handle a standalone JSON-RPC response originating from the peer. - pub(crate) fn process_response(&mut self, response: JSONRPCResponse) { + pub(crate) async fn process_response(&mut self, response: JSONRPCResponse) { tracing::info!("<- response: {:?}", response); + let JSONRPCResponse { id, result, .. } = response; + self.outgoing.notify_client_response(id, result).await } /// Handle a fire-and-forget JSON-RPC notification. diff --git a/codex-rs/mcp-server/src/outgoing_message.rs b/codex-rs/mcp-server/src/outgoing_message.rs index 93a760d36b..a1eea65f25 100644 --- a/codex-rs/mcp-server/src/outgoing_message.rs +++ b/codex-rs/mcp-server/src/outgoing_message.rs @@ -1,3 +1,4 @@ +use std::collections::HashMap; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; @@ -12,11 +13,15 @@ use mcp_types::JSONRPCResponse; use mcp_types::RequestId; use mcp_types::Result; use serde::Serialize; +use tokio::sync::Mutex; use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tracing::warn; pub(crate) struct OutgoingMessageSender { next_request_id: AtomicI64, sender: mpsc::Sender, + request_id_to_callback: Mutex>>, } impl OutgoingMessageSender { @@ -24,17 +29,48 @@ impl OutgoingMessageSender { Self { next_request_id: AtomicI64::new(0), sender, + request_id_to_callback: Mutex::new(HashMap::new()), } } - #[allow(dead_code)] - pub(crate) async fn send_request(&self, method: &str, params: Option) { + pub(crate) async fn send_request( + &self, + method: &str, + params: Option, + ) -> oneshot::Receiver { + let id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::Relaxed)); + let outgoing_message_id = id.clone(); + let (tx_approve, rx_approve) = oneshot::channel(); + { + let mut request_id_to_callback = self.request_id_to_callback.lock().await; + request_id_to_callback.insert(id, tx_approve); + } + let outgoing_message = OutgoingMessage::Request(OutgoingRequest { - id: RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::Relaxed)), + id: outgoing_message_id, method: method.to_string(), params, }); let _ = self.sender.send(outgoing_message).await; + rx_approve + } + + pub(crate) async fn notify_client_response(&self, id: RequestId, result: Result) { + let entry = { + let mut request_id_to_callback = self.request_id_to_callback.lock().await; + request_id_to_callback.remove_entry(&id) + }; + + match entry { + Some((id, sender)) => { + if let Err(err) = sender.send(result) { + warn!("could not notify callback for {id:?} due to: {err:?}"); + } + } + None => { + warn!("could not find callback for {id:?}"); + } + } } pub(crate) async fn send_response(&self, id: RequestId, result: Result) { diff --git a/codex-rs/mcp-types/generate_mcp_types.py b/codex-rs/mcp-types/generate_mcp_types.py index 224e04c0a5..38f57e9a1b 100755 --- a/codex-rs/mcp-types/generate_mcp_types.py +++ b/codex-rs/mcp-types/generate_mcp_types.py @@ -18,6 +18,9 @@ SCHEMA_VERSION = "2025-06-18" JSONRPC_VERSION = "2.0" STANDARD_DERIVE = "#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]\n" +STANDARD_HASHABLE_DERIVE = ( + "#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Hash, Eq)]\n" +) # Will be populated with the schema's `definitions` map in `main()` so that # helper functions (for example `define_any_of`) can perform look-ups while @@ -391,7 +394,7 @@ def define_string_enum( def define_untagged_enum(name: str, type_list: list[str], out: list[str]) -> None: - out.append(STANDARD_DERIVE) + out.append(STANDARD_HASHABLE_DERIVE) out.append("#[serde(untagged)]\n") out.append(f"pub enum {name} {{\n") for simple_type in type_list: diff --git a/codex-rs/mcp-types/src/lib.rs b/codex-rs/mcp-types/src/lib.rs index 6341fb62b4..cf09d67e35 100644 --- a/codex-rs/mcp-types/src/lib.rs +++ b/codex-rs/mcp-types/src/lib.rs @@ -931,7 +931,7 @@ pub struct ProgressNotificationParams { pub total: Option, } -#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Hash, Eq)] #[serde(untagged)] pub enum ProgressToken { String(String), @@ -1031,7 +1031,7 @@ pub struct Request { pub params: Option, } -#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)] +#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, Hash, Eq)] #[serde(untagged)] pub enum RequestId { String(String),