valknar/llmx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: introduce OutgoingMessageSender (#1622)

Browse Source 
Previous to this change, `MessageProcessor` had a
`tokio::sync::mpsc::Sender<JSONRPCMessage>` 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<OutgoingMessage>` 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<OutgoingMessageSender>` 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
This commit is contained in:
Michael Bolin
2025-07-19 00:30:56 -04:00
committed by GitHub
parent e78ec00e73
commit 11fd3123be
4 changed files with 186 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
codex-rs/mcp-server/src/codex_tool_runner.rs
View File

@@ -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<JSONRPCMessage>,
outgoing: Arc<OutgoingMessageSender>,
) {
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;
}
}