From 208395eaa1d0fea56ae28ef52d463ea307b33837 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 16 Jun 2026 01:02:28 +0000 Subject: [PATCH] code-mode: define stdio wire protocol --- codex-rs/Cargo.lock | 1 + codex-rs/code-mode-protocol/Cargo.toml | 2 + codex-rs/code-mode-protocol/src/lib.rs | 1 + codex-rs/code-mode-protocol/src/wire.rs | 290 ++++++++++++++++++ codex-rs/code-mode-protocol/src/wire_tests.rs | 171 +++++++++++ 5 files changed, 465 insertions(+) create mode 100644 codex-rs/code-mode-protocol/src/wire.rs create mode 100644 codex-rs/code-mode-protocol/src/wire_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 0b3295100e..ea4e8357a5 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2519,6 +2519,7 @@ dependencies = [ "pretty_assertions", "serde", "serde_json", + "tokio", "tokio-util", ] diff --git a/codex-rs/code-mode-protocol/Cargo.toml b/codex-rs/code-mode-protocol/Cargo.toml index f3ef06fd32..f6280fd091 100644 --- a/codex-rs/code-mode-protocol/Cargo.toml +++ b/codex-rs/code-mode-protocol/Cargo.toml @@ -16,7 +16,9 @@ workspace = true codex-protocol = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } +tokio = { workspace = true, features = ["io-util", "sync"] } tokio-util = { workspace = true, features = ["rt"] } [dev-dependencies] pretty_assertions = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/code-mode-protocol/src/lib.rs b/codex-rs/code-mode-protocol/src/lib.rs index ec7e03c50b..f3b7f44422 100644 --- a/codex-rs/code-mode-protocol/src/lib.rs +++ b/codex-rs/code-mode-protocol/src/lib.rs @@ -2,6 +2,7 @@ mod description; mod response; mod runtime; mod session; +pub mod wire; pub use description::CODE_MODE_PRAGMA_PREFIX; pub use description::CodeModeToolKind; diff --git a/codex-rs/code-mode-protocol/src/wire.rs b/codex-rs/code-mode-protocol/src/wire.rs new file mode 100644 index 0000000000..7e67e5979a --- /dev/null +++ b/codex-rs/code-mode-protocol/src/wire.rs @@ -0,0 +1,290 @@ +use std::io; + +use serde::Deserialize; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value as JsonValue; +use tokio::io::AsyncRead; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWrite; +use tokio::io::AsyncWriteExt; + +pub const MAX_FRAME_BYTES: usize = 128 * 1024 * 1024; + +pub type RequestId = u64; +pub type SessionId = u64; +pub type CallbackId = u64; + +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(transparent)] +pub struct CellId(String); + +impl CellId { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ClientMessage { + Request { + id: RequestId, + request: HostRequest, + }, + CancelRequest { + id: RequestId, + }, + CallbackResponse { + id: CallbackId, + response: CallbackResponse, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum HostMessage { + Response { + id: RequestId, + result: WireResult, + }, + CallbackRequest { + id: CallbackId, + session_id: SessionId, + request: CallbackRequest, + }, + CancelCallback { + id: CallbackId, + }, + CellClosed { + session_id: SessionId, + cell_id: CellId, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "method", rename_all = "snake_case")] +pub enum HostRequest { + CreateSession, + ShutdownSession { + session_id: SessionId, + }, + CreateCell { + session_id: SessionId, + request: CreateCellRequest, + }, + Observe { + session_id: SessionId, + cell_id: CellId, + mode: ObserveMode, + }, + Terminate { + session_id: SessionId, + cell_id: CellId, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum HostResponse { + SessionCreated { session_id: SessionId }, + SessionShutdown, + CellCreated { cell_id: CellId }, + Observed { event: CellEvent }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "status", rename_all = "snake_case")] +pub enum WireResult { + Ok { value: T }, + Err { error: Error }, +} + +impl WireResult { + pub fn from_result(result: Result) -> Self { + match result { + Ok(value) => Self::Ok { value }, + Err(error) => Self::Err { error }, + } + } +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(tag = "code", rename_all = "snake_case")] +pub enum Error { + MissingSession { session_id: SessionId }, + ShuttingDown, + DuplicateCell { cell_id: CellId }, + MissingCell { cell_id: CellId }, + ClosedCell { cell_id: CellId }, + BusyObserver { cell_id: CellId }, + AlreadyTerminating { cell_id: CellId }, + Cancelled, + Runtime { message: String }, + InvalidRequest { message: String }, + CallbackFailed { message: String }, + Internal { message: String }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct CreateCellRequest { + pub tool_call_id: String, + pub enabled_tools: Vec, + pub source: String, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct ToolDefinition { + pub name: String, + pub tool_name: ToolName, + pub description: String, + pub kind: ToolKind, +} + +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +pub struct ToolName { + pub name: String, + pub namespace: Option, +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolKind { + Function, + Freeform, +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ObserveMode { + YieldAfter { duration_ms: u64 }, + PendingFrontier, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum CellEvent { + Yielded { + content_items: Vec, + }, + Pending { + content_items: Vec, + pending_tool_call_ids: Vec, + }, + Completed { + content_items: Vec, + error_text: Option, + }, + Terminated { + content_items: Vec, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum OutputItem { + Text { + text: String, + }, + Image { + image_url: String, + detail: Option, + }, +} + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ImageDetail { + Auto, + Low, + High, + Original, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum CallbackRequest { + InvokeTool { + invocation: NestedToolCall, + }, + Notify { + call_id: String, + cell_id: CellId, + text: String, + }, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +pub struct NestedToolCall { + pub cell_id: CellId, + pub runtime_tool_call_id: String, + pub tool_name: ToolName, + pub tool_kind: ToolKind, + pub input: Option, +} + +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum CallbackResponse { + ToolResult { result: JsonValue }, + ToolError { error_text: String }, + NotificationDelivered, + NotificationError { error_text: String }, +} + +pub async fn read_frame(reader: &mut R) -> io::Result> +where + R: AsyncRead + Unpin, + T: DeserializeOwned, +{ + let first_length_byte = match reader.read_u8().await { + Ok(byte) => byte, + Err(err) if err.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), + Err(err) => return Err(err), + }; + let mut length_bytes = [first_length_byte, 0, 0, 0]; + reader.read_exact(&mut length_bytes[1..]).await?; + let length = u32::from_be_bytes(length_bytes) as usize; + if length > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!("code-mode IPC frame exceeds {MAX_FRAME_BYTES} bytes"), + )); + } + let mut payload = vec![0; length]; + reader.read_exact(&mut payload).await?; + serde_json::from_slice(&payload) + .map(Some) + .map_err(io::Error::other) +} + +pub async fn write_frame(writer: &mut W, message: &T) -> io::Result<()> +where + W: AsyncWrite + Unpin, + T: Serialize, +{ + let payload = serde_json::to_vec(message).map_err(io::Error::other)?; + if payload.len() > MAX_FRAME_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + format!("code-mode IPC frame exceeds {MAX_FRAME_BYTES} bytes"), + )); + } + let length = u32::try_from(payload.len()).map_err(|_| { + io::Error::new( + io::ErrorKind::InvalidInput, + "code-mode IPC frame length exceeds u32", + ) + })?; + writer.write_all(&length.to_be_bytes()).await?; + writer.write_all(&payload).await?; + writer.flush().await +} + +#[cfg(test)] +#[path = "wire_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-protocol/src/wire_tests.rs b/codex-rs/code-mode-protocol/src/wire_tests.rs new file mode 100644 index 0000000000..23d13455c2 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/wire_tests.rs @@ -0,0 +1,171 @@ +use std::io; + +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::io::AsyncWriteExt; +use tokio::io::duplex; + +use super::*; + +#[test] +fn observe_request_has_a_stable_tagged_shape() { + let message = ClientMessage::Request { + id: 7, + request: HostRequest::Observe { + session_id: 3, + cell_id: CellId::new("cell-9"), + mode: ObserveMode::YieldAfter { duration_ms: 250 }, + }, + }; + + assert_eq!( + serde_json::to_value(message).unwrap(), + json!({ + "type": "request", + "id": 7, + "request": { + "method": "observe", + "session_id": 3, + "cell_id": "cell-9", + "mode": { + "type": "yield_after", + "duration_ms": 250, + }, + }, + }) + ); +} + +#[test] +fn cell_closed_notification_has_a_stable_tagged_shape() { + let message = HostMessage::CellClosed { + session_id: 3, + cell_id: CellId::new("cell-9"), + }; + + assert_eq!( + serde_json::to_value(message).unwrap(), + json!({ + "type": "cell_closed", + "session_id": 3, + "cell_id": "cell-9", + }) + ); +} + +#[test] +fn busy_observer_error_has_a_stable_tagged_shape() { + let message = HostMessage::Response { + id: 11, + result: WireResult::Err { + error: Error::BusyObserver { + cell_id: CellId::new("cell-2"), + }, + }, + }; + + assert_eq!( + serde_json::to_value(message).unwrap(), + json!({ + "type": "response", + "id": 11, + "result": { + "status": "err", + "error": { + "code": "busy_observer", + "cell_id": "cell-2", + }, + }, + }) + ); +} + +#[test] +fn callback_result_and_error_round_trip() { + for response in [ + CallbackResponse::ToolResult { + result: json!({"value": 42}), + }, + CallbackResponse::ToolError { + error_text: "tool failed".to_string(), + }, + CallbackResponse::NotificationDelivered, + CallbackResponse::NotificationError { + error_text: "notify failed".to_string(), + }, + ] { + let message = ClientMessage::CallbackResponse { id: 19, response }; + let encoded = serde_json::to_vec(&message).unwrap(); + + assert_eq!( + serde_json::from_slice::(&encoded).unwrap(), + message + ); + } +} + +#[tokio::test] +async fn frame_round_trip_preserves_message() { + let (mut client, mut server) = duplex(1024); + let message = ClientMessage::Request { + id: 7, + request: HostRequest::CreateSession, + }; + + write_frame(&mut client, &message).await.unwrap(); + + assert_eq!(read_frame(&mut server).await.unwrap(), Some(message)); +} + +#[tokio::test] +async fn clean_eof_returns_none() { + let (client, mut server) = duplex(16); + drop(client); + + assert_eq!( + read_frame::<_, ClientMessage>(&mut server).await.unwrap(), + None + ); +} + +#[tokio::test] +async fn partial_frame_header_is_rejected() { + let (mut client, mut server) = duplex(16); + client.write_all(&[0, 1]).await.unwrap(); + client.shutdown().await.unwrap(); + + let error = read_frame::<_, ClientMessage>(&mut server) + .await + .unwrap_err(); + + assert_eq!(error.kind(), io::ErrorKind::UnexpectedEof); +} + +#[tokio::test] +async fn oversized_frame_is_rejected_before_allocation() { + let (mut client, mut server) = duplex(16); + let oversized_length = u32::try_from(MAX_FRAME_BYTES + 1).unwrap(); + client + .write_all(&oversized_length.to_be_bytes()) + .await + .unwrap(); + + let error = read_frame::<_, ClientMessage>(&mut server) + .await + .unwrap_err(); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); +} + +#[tokio::test] +async fn malformed_json_frame_is_rejected() { + let (mut client, mut server) = duplex(16); + client.write_all(&1_u32.to_be_bytes()).await.unwrap(); + client.write_all(b"{").await.unwrap(); + + let error = read_frame::<_, ClientMessage>(&mut server) + .await + .unwrap_err(); + + assert_eq!(error.kind(), io::ErrorKind::Other); +}