From 020ecd368efbce7fc7b2dee110969fa6740135df Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 23 Jun 2026 22:16:42 -0700 Subject: [PATCH] Define code mode host wire protocol --- codex-rs/code-mode-protocol/src/host/codec.rs | 101 ++++++++++++ .../src/host/codec_tests.rs | 104 ++++++++++++ .../code-mode-protocol/src/host/host_tests.rs | 154 ++++++++++++++++-- .../code-mode-protocol/src/host/message.rs | 127 ++++++++++++++- codex-rs/code-mode-protocol/src/host/mod.rs | 25 ++- codex-rs/code-mode-protocol/src/host/types.rs | 3 + 6 files changed, 486 insertions(+), 28 deletions(-) create mode 100644 codex-rs/code-mode-protocol/src/host/codec.rs create mode 100644 codex-rs/code-mode-protocol/src/host/codec_tests.rs diff --git a/codex-rs/code-mode-protocol/src/host/codec.rs b/codex-rs/code-mode-protocol/src/host/codec.rs new file mode 100644 index 0000000000..beb50b17e3 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/host/codec.rs @@ -0,0 +1,101 @@ +use std::io; +use std::mem::size_of; + +use serde::Serialize; +use serde::de::DeserializeOwned; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; + +/// Maximum JSON payload size accepted for one IPC frame. +pub const MAX_FRAME_BYTES: usize = 8 * 1024 * 1024; + +/// Decodes JSON messages prefixed by a four-byte little-endian payload length. +pub struct FramedReader { + reader: R, +} + +impl FramedReader +where + R: AsyncRead + Unpin, +{ + pub fn new(reader: R) -> Self { + Self { reader } + } + + /// Reads the next frame, returning `None` only for EOF at a frame boundary. + pub async fn read(&mut self) -> io::Result> + where + T: DeserializeOwned, + { + let mut length_bytes = [0_u8; size_of::()]; + if self.reader.read(&mut length_bytes[..1]).await? == 0 { + return Ok(None); + } + self.reader.read_exact(&mut length_bytes[1..]).await?; + + let length = u32::from_le_bytes(length_bytes) as usize; + if length > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("code-mode IPC frame length {length} exceeds {MAX_FRAME_BYTES} bytes"), + )); + } + + let mut payload = vec![0; length]; + self.reader.read_exact(&mut payload).await?; + serde_json::from_slice(&payload).map(Some).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to decode code-mode IPC frame: {err}"), + ) + }) + } +} + +/// Encodes JSON messages with a four-byte little-endian payload length. +pub struct FramedWriter { + writer: W, +} + +impl FramedWriter +where + W: AsyncWrite + Unpin, +{ + pub fn new(writer: W) -> Self { + Self { writer } + } + + /// Writes and flushes one complete frame. + pub async fn write(&mut self, message: &T) -> io::Result<()> + where + T: Serialize, + { + let payload = serde_json::to_vec(message).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to encode code-mode IPC frame: {err}"), + ) + })?; + if payload.len() > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "code-mode IPC frame length {} exceeds {MAX_FRAME_BYTES} bytes", + payload.len() + ), + )); + } + let length = u32::try_from(payload.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidData, + "code-mode IPC frame length exceeds u32", + ) + })?; + + self.writer.write_all(&length.to_le_bytes()).await?; + self.writer.write_all(&payload).await?; + self.writer.flush().await + } +} diff --git a/codex-rs/code-mode-protocol/src/host/codec_tests.rs b/codex-rs/code-mode-protocol/src/host/codec_tests.rs new file mode 100644 index 0000000000..996bd0679a --- /dev/null +++ b/codex-rs/code-mode-protocol/src/host/codec_tests.rs @@ -0,0 +1,104 @@ +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; + +use super::FramedReader; +use super::FramedWriter; +use super::MAX_FRAME_BYTES; + +#[tokio::test] +async fn frame_wire_format_is_little_endian_length_prefixed_json() { + let (writer, mut reader) = tokio::io::duplex(/*max_buf_size*/ 128); + let write = tokio::spawn(async move { + FramedWriter::new(writer) + .write(&json!({"value": 1})) + .await + .expect("write frame"); + }); + + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await.expect("read bytes"); + write.await.expect("writer task"); + + let payload = br#"{"value":1}"#; + let mut expected = (payload.len() as u32).to_le_bytes().to_vec(); + expected.extend_from_slice(payload); + assert_eq!(bytes, expected); +} + +#[tokio::test] +async fn fragmented_frame_round_trips() { + let value = json!({"type": "session/open", "sessionId": "session-1"}); + let payload = serde_json::to_vec(&value).expect("serialize"); + let mut bytes = (payload.len() as u32).to_le_bytes().to_vec(); + bytes.extend(payload); + + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 128); + let write = tokio::spawn(async move { + for byte in bytes { + writer.write_all(&[byte]).await.expect("write byte"); + tokio::task::yield_now().await; + } + }); + + assert_eq!( + FramedReader::new(reader) + .read::() + .await + .expect("read frame"), + Some(value) + ); + write.await.expect("writer task"); +} + +#[tokio::test] +async fn eof_is_clean_only_at_a_frame_boundary() { + let (writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + drop(writer); + assert_eq!( + FramedReader::new(reader) + .read::() + .await + .expect("clean eof"), + None + ); + + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + writer + .write_all(&[1, 0]) + .await + .expect("write partial header"); + drop(writer); + let err = FramedReader::new(reader) + .read::() + .await + .expect_err("truncated header"); + assert_eq!(err.kind(), std::io::ErrorKind::UnexpectedEof); +} + +#[tokio::test] +async fn oversized_and_malformed_frames_are_rejected() { + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + writer + .write_all(&((MAX_FRAME_BYTES as u32) + 1).to_le_bytes()) + .await + .expect("write oversized header"); + let err = FramedReader::new(reader) + .read::() + .await + .expect_err("oversized frame"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + + let (mut writer, reader) = tokio::io::duplex(/*max_buf_size*/ 16); + writer + .write_all(&(1_u32).to_le_bytes()) + .await + .expect("write length"); + writer.write_all(b"{").await.expect("write malformed json"); + let err = FramedReader::new(reader) + .read::() + .await + .expect_err("malformed frame"); + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); +} diff --git a/codex-rs/code-mode-protocol/src/host/host_tests.rs b/codex-rs/code-mode-protocol/src/host/host_tests.rs index 2d64993440..1d5ca91e41 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -5,12 +5,18 @@ use super::Capability; use super::CapabilitySet; use super::ClientHello; use super::ClientToHost; +use super::DelegateRequest; +use super::DelegateResponse; use super::HandshakeRejectReason; use super::HostHello; +use super::HostRequest; +use super::HostResponse; use super::HostToClient; use super::ProtocolVersion; use super::SessionId; use super::SupportedProtocolVersions; +use super::WireResult; +use crate::CellId; fn session_id() -> SessionId { SessionId::new("session-1").expect("valid session ID") @@ -94,16 +100,36 @@ fn handshake_wire_contract_is_explicit_and_round_trips() { fn session_lifecycle_wire_contract_is_explicit_and_round_trips() { let client_messages = [ ( - ClientToHost::OpenSession { - session_id: session_id(), + ClientToHost::Request { + id: 7, + request: HostRequest::OpenSession { + session_id: session_id(), + }, }, - json!({ "type": "session/open", "sessionId": "session-1" }), + json!({ + "type": "operation/request", + "id": 7, + "request": { + "method": "session/open", + "sessionId": "session-1", + }, + }), ), ( - ClientToHost::CloseSession { - session_id: session_id(), + ClientToHost::Request { + id: 8, + request: HostRequest::ShutdownSession { + session_id: session_id(), + }, }, - json!({ "type": "session/close", "sessionId": "session-1" }), + json!({ + "type": "operation/request", + "id": 8, + "request": { + "method": "session/shutdown", + "sessionId": "session-1", + }, + }), ), ]; for (message, encoded) in client_messages { @@ -116,16 +142,46 @@ fn session_lifecycle_wire_contract_is_explicit_and_round_trips() { let host_messages = [ ( - HostToClient::SessionReady { - session_id: session_id(), + HostToClient::Response { + id: 7, + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session_id(), + }, + }, }, - json!({ "type": "session/ready", "sessionId": "session-1" }), + json!({ + "type": "operation/response", + "id": 7, + "result": { + "status": "ok", + "value": { + "type": "session/ready", + "sessionId": "session-1", + }, + }, + }), ), ( - HostToClient::SessionClosed { - session_id: session_id(), + HostToClient::Response { + id: 8, + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session_id(), + }, + }, }, - json!({ "type": "session/closed", "sessionId": "session-1" }), + json!({ + "type": "operation/response", + "id": 8, + "result": { + "status": "ok", + "value": { + "type": "session/closed", + "sessionId": "session-1", + }, + }, + }), ), ]; for (message, encoded) in host_messages { @@ -137,6 +193,61 @@ fn session_lifecycle_wire_contract_is_explicit_and_round_trips() { } } +#[test] +fn delegate_wire_contract_is_explicit_and_round_trips() { + let request = HostToClient::DelegateRequest { + id: 11, + session_id: session_id(), + request: DelegateRequest::Notify { + call_id: "call-1".to_string(), + cell_id: CellId::new("cell-1".to_string()), + text: "hello".to_string(), + }, + }; + let request_json = json!({ + "type": "delegate/request", + "id": 11, + "sessionId": "session-1", + "request": { + "type": "notification/send", + "callId": "call-1", + "cellId": "cell-1", + "text": "hello", + }, + }); + assert_eq!( + serde_json::to_value(&request).expect("serialize"), + request_json + ); + assert_eq!( + serde_json::from_value::(request_json).expect("deserialize"), + request + ); + + let response = ClientToHost::DelegateResponse { + id: 11, + result: WireResult::Ok { + value: DelegateResponse::NotificationDelivered, + }, + }; + let response_json = json!({ + "type": "delegate/response", + "id": 11, + "result": { + "status": "ok", + "value": { "type": "notification/delivered" }, + }, + }); + assert_eq!( + serde_json::to_value(&response).expect("serialize"), + response_json + ); + assert_eq!( + serde_json::from_value::(response_json).expect("deserialize"), + response + ); +} + #[test] fn invalid_protocol_states_cannot_be_constructed_or_decoded() { assert!(SessionId::new("").is_err()); @@ -168,7 +279,11 @@ fn invalid_protocol_states_cannot_be_constructed_or_decoded() { ); for invalid in [ - json!({ "type": "session/open", "sessionId": "" }), + json!({ + "type": "operation/request", + "id": 1, + "request": { "method": "session/open", "sessionId": "" }, + }), json!({ "type": "connection/hello", "supportedVersions": [], @@ -190,16 +305,21 @@ fn invalid_protocol_states_cannot_be_constructed_or_decoded() { fn unknown_fields_are_rejected() { assert!( serde_json::from_value::(json!({ - "type": "session/open", - "sessionId": "session-1", + "type": "operation/request", + "id": 1, + "request": { "method": "session/open", "sessionId": "session-1" }, "unexpected": true, })) .is_err() ); assert!( serde_json::from_value::(json!({ - "type": "session/ready", - "sessionId": "session-1", + "type": "operation/response", + "id": 1, + "result": { + "status": "ok", + "value": { "type": "session/ready", "sessionId": "session-1" }, + }, "unexpected": true, })) .is_err() diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index ed30ac8f17..307a5802af 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -2,13 +2,22 @@ use std::fmt; use serde::Deserialize; use serde::Serialize; +use serde_json::Value as JsonValue; use super::Capability; use super::CapabilitySet; +use super::DelegateRequestId; use super::HandshakeRejectReason; use super::ProtocolVersion; +use super::RequestId; use super::SessionId; use super::SupportedProtocolVersions; +use crate::CellId; +use crate::CodeModeNestedToolCall; +use crate::ExecuteRequest; +use crate::RuntimeResponse; +use crate::WaitOutcome; +use crate::WaitRequest; #[derive(Clone, Debug, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] @@ -116,27 +125,133 @@ impl HostHello { } /// Messages sent from a client to the code-mode host. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[derive(Debug, Deserialize, PartialEq, Serialize)] #[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] pub enum ClientToHost { #[serde(rename = "connection/hello")] ClientHello(ClientHello), - #[serde(rename = "session/open")] - OpenSession { session_id: SessionId }, - #[serde(rename = "session/close")] - CloseSession { session_id: SessionId }, + #[serde(rename = "operation/request")] + Request { id: RequestId, request: HostRequest }, + #[serde(rename = "delegate/response")] + DelegateResponse { + id: DelegateRequestId, + result: WireResult, + }, } /// Messages sent from the code-mode host to a client. -#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[derive(Debug, Deserialize, PartialEq, Serialize)] #[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] pub enum HostToClient { #[serde(rename = "connection/ready")] HostHello(HostHello), #[serde(rename = "connection/rejected")] HandshakeRejected { reason: HandshakeRejectReason }, + #[serde(rename = "operation/response")] + Response { + id: RequestId, + result: WireResult, + }, + #[serde(rename = "execute/initialResponse")] + InitialResponse { + id: RequestId, + result: WireResult, + }, + #[serde(rename = "delegate/request")] + DelegateRequest { + id: DelegateRequestId, + session_id: SessionId, + request: DelegateRequest, + }, + #[serde(rename = "delegate/cancel")] + CancelDelegateRequest { id: DelegateRequestId }, + #[serde(rename = "cell/closed")] + CellClosed { + session_id: SessionId, + cell_id: CellId, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "method", rename_all_fields = "camelCase")] +pub enum HostRequest { + #[serde(rename = "session/open")] + OpenSession { session_id: SessionId }, + #[serde(rename = "session/execute")] + Execute { + session_id: SessionId, + request: ExecuteRequest, + }, + #[serde(rename = "session/wait")] + Wait { + session_id: SessionId, + request: WaitRequest, + }, + #[serde(rename = "session/terminate")] + Terminate { + session_id: SessionId, + cell_id: CellId, + }, + #[serde(rename = "session/shutdown")] + ShutdownSession { session_id: SessionId }, +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] +pub enum HostResponse { #[serde(rename = "session/ready")] SessionReady { session_id: SessionId }, + #[serde(rename = "execution/started")] + ExecutionStarted { cell_id: CellId }, + #[serde(rename = "wait/completed")] + WaitCompleted { outcome: WaitOutcome }, #[serde(rename = "session/closed")] SessionClosed { session_id: SessionId }, } + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] +pub enum DelegateRequest { + #[serde(rename = "tool/invoke")] + InvokeTool { invocation: CodeModeNestedToolCall }, + #[serde(rename = "notification/send")] + Notify { + call_id: String, + cell_id: CellId, + text: String, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] +pub enum DelegateResponse { + #[serde(rename = "tool/result")] + ToolResult { result: JsonValue }, + #[serde(rename = "notification/delivered")] + NotificationDelivered, +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "status", rename_all_fields = "camelCase")] +pub enum WireResult { + #[serde(rename = "ok")] + Ok { value: T }, + #[serde(rename = "error")] + Err { message: String }, +} + +impl WireResult { + pub fn from_result(result: Result) -> Self { + match result { + Ok(value) => Self::Ok { value }, + Err(message) => Self::Err { message }, + } + } + + pub fn into_result(self) -> Result { + match self { + Self::Ok { value } => Ok(value), + Self::Err { message } => Err(message), + } + } +} diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index f705f80e36..c3e912071b 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -1,29 +1,44 @@ -//! Transport-neutral messages for the callback-only code-mode host boundary. +//! Messages and local IPC framing for the code-mode host boundary. //! -//! Protocol version 1 relies on ordered framing and connection-scoped -//! fail-stop behavior rather than message sequence numbers. It defines no -//! optional capabilities yet; capability names provide an extension point for -//! later versions without weakening the v1 decoder. +//! Protocol version 1 multiplexes session operations and delegate callbacks by +//! request ID over one ordered connection. It defines no optional capabilities +//! yet; capability names provide an extension point for later versions without +//! weakening the v1 decoder. +mod codec; mod error; mod message; mod types; +pub use codec::FramedReader; +pub use codec::FramedWriter; +pub use codec::MAX_FRAME_BYTES; pub use error::HandshakeRejectReason; pub use message::ClientHello; pub use message::ClientHelloError; pub use message::ClientToHost; +pub use message::DelegateRequest; +pub use message::DelegateResponse; pub use message::HostHello; +pub use message::HostRequest; +pub use message::HostResponse; pub use message::HostToClient; +pub use message::WireResult; pub use types::Capability; pub use types::CapabilitySet; +pub use types::DelegateRequestId; pub use types::DuplicateCapability; pub use types::InvalidIdentifier; pub use types::InvalidSupportedProtocolVersions; pub use types::ProtocolVersion; +pub use types::RequestId; pub use types::SessionId; pub use types::SupportedProtocolVersions; #[cfg(test)] #[path = "host_tests.rs"] mod tests; + +#[cfg(test)] +#[path = "codec_tests.rs"] +mod codec_tests; diff --git a/codex-rs/code-mode-protocol/src/host/types.rs b/codex-rs/code-mode-protocol/src/host/types.rs index 69c8fd0881..b51474607a 100644 --- a/codex-rs/code-mode-protocol/src/host/types.rs +++ b/codex-rs/code-mode-protocol/src/host/types.rs @@ -8,6 +8,9 @@ use serde::Serialize; use serde::Serializer; use serde::de::Error as _; +pub type RequestId = u64; +pub type DelegateRequestId = u64; + #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] #[serde(transparent)] pub struct ProtocolVersion(NonZeroU32);