From 2c3d906f19b3df31f737c7e3636082a274ff7fca Mon Sep 17 00:00:00 2001 From: Michael Bolin Date: Tue, 12 Aug 2025 11:16:34 -0700 Subject: [PATCH] exploration: new protocol format Here we explore an evolution of the internal protocol that can be mapped to JSON-RPC in a straightforward way. --- codex-rs/Cargo.lock | 12 +++ codex-rs/Cargo.toml | 1 + codex-rs/protocol/Cargo.toml | 21 ++++ codex-rs/protocol/src/lib.rs | 1 + codex-rs/protocol/src/notification.rs | 147 ++++++++++++++++++++++++++ 5 files changed, 182 insertions(+) create mode 100644 codex-rs/protocol/Cargo.toml create mode 100644 codex-rs/protocol/src/lib.rs create mode 100644 codex-rs/protocol/src/notification.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 5ec773adc5..9a90df712f 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -860,6 +860,18 @@ dependencies = [ "wiremock", ] +[[package]] +name = "codex-protocol" +version = "0.0.0" +dependencies = [ + "pretty_assertions", + "serde", + "serde_json", + "strum 0.27.2", + "strum_macros 0.27.2", + "uuid", +] + [[package]] name = "codex-tui" version = "0.0.0" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 0ed8852228..2fb9b9271e 100644 --- a/codex-rs/Cargo.toml +++ b/codex-rs/Cargo.toml @@ -15,6 +15,7 @@ members = [ "mcp-server", "mcp-types", "ollama", + "protocol", "tui", ] resolver = "2" diff --git a/codex-rs/protocol/Cargo.toml b/codex-rs/protocol/Cargo.toml new file mode 100644 index 0000000000..2f122fe3f7 --- /dev/null +++ b/codex-rs/protocol/Cargo.toml @@ -0,0 +1,21 @@ +[package] +edition = "2024" +name = "codex-protocol" +version = { workspace = true } + +[lib] +name = "codex_protocol" +path = "src/lib.rs" + +[lints] +workspace = true + +[dependencies] +serde = { version = "1", features = ["derive"] } +strum = "0.27.2" +strum_macros = "0.27.2" +uuid = { version = "1", features = ["serde", "v4"] } + +[dev-dependencies] +pretty_assertions = "1.4.1" +serde_json = "1" diff --git a/codex-rs/protocol/src/lib.rs b/codex-rs/protocol/src/lib.rs new file mode 100644 index 0000000000..c82dfba556 --- /dev/null +++ b/codex-rs/protocol/src/lib.rs @@ -0,0 +1 @@ +mod notification; diff --git a/codex-rs/protocol/src/notification.rs b/codex-rs/protocol/src/notification.rs new file mode 100644 index 0000000000..f1bbace01a --- /dev/null +++ b/codex-rs/protocol/src/notification.rs @@ -0,0 +1,147 @@ +use serde::Deserialize; +use serde::Serialize; +use strum_macros::Display; +use uuid::Uuid; + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Display)] +#[serde(tag = "method", rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +pub enum NotificationMessage { + Conversation(ConversationNotification), + ShutdownComplete, +} + +/// Notification associated with a conversation. The `conversation_id` is key +/// so clients can dispatch messages to the correct conversation. +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +pub struct ConversationNotification { + conversation_id: Uuid, + #[serde(flatten)] + message: ConversationNotificationMessage, +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Display)] +#[serde(tag = "type", rename_all = "snake_case")] +#[strum(serialize_all = "snake_case")] +enum ConversationNotificationMessage { + Initialized(ConversationInitialized), +} + +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +struct ConversationInitialized { + model: String, +} + +#[allow(clippy::unwrap_used)] +#[cfg(test)] +mod tests { + use super::*; + use pretty_assertions::assert_eq; + use serde::de::Error as _; + use serde_json::json; + + // The idea is that the way we map `NotificationMessage` to an MCP + // notification is to use the tuple type as the `params`. + #[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] + struct JsonrpcNotification { + jsonrpc: String, + method: String, + params: serde_json::Value, + } + + // Example of how JSON-RPC serialization could work. + fn to_jsonrpc_message( + notification: NotificationMessage, + ) -> Result { + let method = notification.to_string(); + let params = match notification { + NotificationMessage::Conversation(notification) => serde_json::to_value(notification)?, + NotificationMessage::ShutdownComplete => serde_json::Value::Null, + }; + + let jsonrpc_notification = JsonrpcNotification { + jsonrpc: "2.0".to_string(), + method, + params, + }; + + serde_json::to_value(jsonrpc_notification) + } + + fn from_jsonrpc_message( + jsonrpc_notification: JsonrpcNotification, + ) -> Result { + let JsonrpcNotification { + jsonrpc: _, + method, + params, + } = jsonrpc_notification; + + match method.as_str() { + "conversation" => { + let conversation_notification: ConversationNotification = + serde_json::from_value(params)?; + Ok(NotificationMessage::Conversation(conversation_notification)) + } + "shutdown_complete" => Ok(NotificationMessage::ShutdownComplete), + _ => Err(serde_json::Error::custom(format!( + "Unknown method: {method}" + ))), + } + } + + #[test] + fn test_serialize_notification_message_conversation() { + let conversation_id = Uuid::parse_str("123e4567-e89b-12d3-a456-426614174000").unwrap(); + let message = NotificationMessage::Conversation(ConversationNotification { + conversation_id, + message: ConversationNotificationMessage::Initialized(ConversationInitialized { + model: "gpt-5".to_string(), + }), + }); + assert_eq!( + json!({ + "method": "conversation", + "conversation_id": conversation_id.to_string(), + "type": "initialized", + "model": "gpt-5", + }), + serde_json::to_value(message.clone()).unwrap() + ); + + let expected_jsonrpc_message = json!({ + "jsonrpc": "2.0", + "method": "conversation", + "params": { + "conversation_id": conversation_id.to_string(), + "type": "initialized", + "model": "gpt-5", + } + }); + assert_eq!( + expected_jsonrpc_message, + to_jsonrpc_message(message.clone()).unwrap() + ); + + let serialized_json_rpc_message = serde_json::to_string(&expected_jsonrpc_message).unwrap(); + let deserialized_json_rpc_message = + serde_json::from_str::(&serialized_json_rpc_message).unwrap(); + assert_eq!( + JsonrpcNotification { + jsonrpc: "2.0".to_string(), + method: "conversation".to_string(), + params: json!({ + "conversation_id": conversation_id.to_string(), + "type": "initialized", + "model": "gpt-5", + }), + }, + deserialized_json_rpc_message + ); + + assert_eq!( + message, + from_jsonrpc_message(deserialized_json_rpc_message).unwrap() + ); + } +}