From 020ecd368efbce7fc7b2dee110969fa6740135df Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 23 Jun 2026 22:16:42 -0700 Subject: [PATCH 1/6] 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 000000000000..beb50b17e3ac --- /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 000000000000..996bd0679a86 --- /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 2d64993440de..1d5ca91e411b 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 ed30ac8f17d1..307a5802af7b 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 f705f80e3687..c3e912071b76 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 69c8fd088194..b51474607a2e 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); From 884e77a94dc85f8c4f6e4b714cdf1f7d58af752e Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Wed, 24 Jun 2026 10:52:38 -0700 Subject: [PATCH 2/6] Pin code mode host V1 payloads --- codex-rs/code-mode-protocol/src/host/codec.rs | 2 +- .../code-mode-protocol/src/host/host_tests.rs | 711 ++++++++++++++---- .../code-mode-protocol/src/host/message.rs | 30 +- codex-rs/code-mode-protocol/src/host/mod.rs | 12 + .../code-mode-protocol/src/host/payload.rs | 406 ++++++++++ 5 files changed, 983 insertions(+), 178 deletions(-) create mode 100644 codex-rs/code-mode-protocol/src/host/payload.rs diff --git a/codex-rs/code-mode-protocol/src/host/codec.rs b/codex-rs/code-mode-protocol/src/host/codec.rs index beb50b17e3ac..6bf83d367da2 100644 --- a/codex-rs/code-mode-protocol/src/host/codec.rs +++ b/codex-rs/code-mode-protocol/src/host/codec.rs @@ -9,7 +9,7 @@ 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; +pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024; /// Decodes JSON messages prefixed by a four-byte little-endian payload length. pub struct FramedReader { 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 1d5ca91e411b..77a9cf71634e 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -1,4 +1,9 @@ +use std::fmt::Debug; + use pretty_assertions::assert_eq; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::Value; use serde_json::json; use super::Capability; @@ -15,13 +20,27 @@ use super::HostToClient; use super::ProtocolVersion; use super::SessionId; use super::SupportedProtocolVersions; +use super::WireCellId; +use super::WireContentItem; +use super::WireExecuteRequest; +use super::WireImageDetail; +use super::WireNestedToolCall; use super::WireResult; -use crate::CellId; +use super::WireRuntimeResponse; +use super::WireToolDefinition; +use super::WireToolKind; +use super::WireToolName; +use super::WireWaitOutcome; +use super::WireWaitRequest; fn session_id() -> SessionId { SessionId::new("session-1").expect("valid session ID") } +fn cell_id(value: &str) -> WireCellId { + WireCellId::new(value) +} + fn capability(value: &str) -> Capability { Capability::new(value).expect("valid capability") } @@ -31,220 +50,512 @@ fn supported_versions() -> SupportedProtocolVersions { .expect("nonempty unique protocol versions") } -#[test] -fn handshake_wire_contract_is_explicit_and_round_trips() { - let client_hello = ClientToHost::ClientHello( - ClientHello::new( - supported_versions(), - CapabilitySet::try_new([capability("required")]).expect("valid required set"), - CapabilitySet::try_new([capability("optional")]).expect("valid optional set"), - ) - .expect("disjoint capabilities"), - ); - let client_hello_json = json!({ - "type": "connection/hello", - "supportedVersions": [1], - "requiredCapabilities": ["required"], - "optionalCapabilities": ["optional"], - }); - assert_eq!( - serde_json::to_value(&client_hello).expect("serialize"), - client_hello_json - ); +fn assert_wire_round_trip(message: T, encoded: Value) +where + T: Debug + DeserializeOwned + PartialEq + Serialize, +{ + assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded); assert_eq!( - serde_json::from_value::(client_hello_json).expect("deserialize"), - client_hello + serde_json::from_value::(encoded).expect("deserialize"), + message ); +} - let host_hello = HostToClient::HostHello(HostHello::new( - ProtocolVersion::V1, - CapabilitySet::try_new([capability("required")]).expect("valid capabilities"), - )); - let host_hello_json = json!({ - "type": "connection/ready", - "selectedVersion": 1, - "capabilities": ["required"], - }); - assert_eq!( - serde_json::to_value(&host_hello).expect("serialize"), - host_hello_json - ); - assert_eq!( - serde_json::from_value::(host_hello_json).expect("deserialize"), - host_hello - ); +fn execute_request() -> WireExecuteRequest { + WireExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: vec![ + WireToolDefinition { + name: "function_tool".to_string(), + tool_name: WireToolName { + name: "function_tool".to_string(), + namespace: None, + }, + description: "function tool".to_string(), + kind: WireToolKind::Function, + input_schema: Some(json!({ "type": "object" })), + output_schema: None, + }, + WireToolDefinition { + name: "freeform_tool".to_string(), + tool_name: WireToolName { + name: "freeform_tool".to_string(), + namespace: Some("mcp__sample__".to_string()), + }, + description: "freeform tool".to_string(), + kind: WireToolKind::Freeform, + input_schema: None, + output_schema: Some(json!({ "type": "string" })), + }, + ], + source: "text('hello');".to_string(), + yield_time_ms: Some(25), + max_output_tokens: Some(100), + } +} - let rejected = HostToClient::HandshakeRejected { - reason: HandshakeRejectReason::NoCompatibleVersion { - supported_versions: supported_versions(), +fn content_items() -> Vec { + vec![ + WireContentItem::InputText { + text: "hello".to_string(), }, - }; - let rejected_json = json!({ - "type": "connection/rejected", - "reason": { - "type": "noCompatibleVersion", - "supportedVersions": [1], + WireContentItem::InputImage { + image_url: "data:image/png;base64,none".to_string(), + detail: None, }, - }); - assert_eq!( - serde_json::to_value(&rejected).expect("serialize"), - rejected_json + WireContentItem::InputImage { + image_url: "data:image/png;base64,auto".to_string(), + detail: Some(WireImageDetail::Auto), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,low".to_string(), + detail: Some(WireImageDetail::Low), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,high".to_string(), + detail: Some(WireImageDetail::High), + }, + WireContentItem::InputImage { + image_url: "data:image/png;base64,original".to_string(), + detail: Some(WireImageDetail::Original), + }, + ] +} + +fn content_items_json() -> Value { + json!([ + { "type": "input_text", "text": "hello" }, + { "type": "input_image", "image_url": "data:image/png;base64,none" }, + { + "type": "input_image", + "image_url": "data:image/png;base64,auto", + "detail": "auto", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,low", + "detail": "low", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,high", + "detail": "high", + }, + { + "type": "input_image", + "image_url": "data:image/png;base64,original", + "detail": "original", + }, + ]) +} + +#[test] +fn handshake_v1_variants_are_pinned() { + assert_wire_round_trip( + ClientToHost::ClientHello( + ClientHello::new( + supported_versions(), + CapabilitySet::try_new([capability("required")]).expect("valid required set"), + CapabilitySet::try_new([capability("optional")]).expect("valid optional set"), + ) + .expect("disjoint capabilities"), + ), + json!({ + "type": "connection/hello", + "supportedVersions": [1], + "requiredCapabilities": ["required"], + "optionalCapabilities": ["optional"], + }), ); - assert_eq!( - serde_json::from_value::(rejected_json).expect("deserialize"), - rejected + assert_wire_round_trip( + HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::try_new([capability("required")]).expect("valid capabilities"), + )), + json!({ + "type": "connection/ready", + "selectedVersion": 1, + "capabilities": ["required"], + }), ); + for (reason, encoded) in [ + ( + HandshakeRejectReason::NoCompatibleVersion { + supported_versions: supported_versions(), + }, + json!({ + "type": "connection/rejected", + "reason": { + "type": "noCompatibleVersion", + "supportedVersions": [1], + }, + }), + ), + ( + HandshakeRejectReason::MissingRequiredCapability { + capability: capability("required"), + }, + json!({ + "type": "connection/rejected", + "reason": { + "type": "missingRequiredCapability", + "capability": "required", + }, + }), + ), + ( + HandshakeRejectReason::InvalidHello { + message: "invalid hello".to_string(), + }, + json!({ + "type": "connection/rejected", + "reason": { + "type": "invalidHello", + "message": "invalid hello", + }, + }), + ), + ] { + assert_wire_round_trip(HostToClient::HandshakeRejected { reason }, encoded); + } } #[test] -fn session_lifecycle_wire_contract_is_explicit_and_round_trips() { - let client_messages = [ +fn client_to_host_v1_variants_are_pinned() { + let execute_request = execute_request(); + for (id, request, encoded_request) in [ ( - ClientToHost::Request { - id: 7, - request: HostRequest::OpenSession { - session_id: session_id(), - }, + 1, + HostRequest::OpenSession { + session_id: session_id(), + }, + json!({ "method": "session/open", "sessionId": "session-1" }), + ), + ( + 2, + HostRequest::Execute { + session_id: session_id(), + request: execute_request, }, json!({ - "type": "operation/request", - "id": 7, + "method": "session/execute", + "sessionId": "session-1", "request": { - "method": "session/open", - "sessionId": "session-1", + "tool_call_id": "call-1", + "enabled_tools": [ + { + "name": "function_tool", + "tool_name": { "name": "function_tool", "namespace": null }, + "description": "function tool", + "kind": "function", + "input_schema": { "type": "object" }, + "output_schema": null, + }, + { + "name": "freeform_tool", + "tool_name": { + "name": "freeform_tool", + "namespace": "mcp__sample__", + }, + "description": "freeform tool", + "kind": "freeform", + "input_schema": null, + "output_schema": { "type": "string" }, + }, + ], + "source": "text('hello');", + "yield_time_ms": 25, + "max_output_tokens": 100, }, }), ), ( - ClientToHost::Request { - id: 8, - request: HostRequest::ShutdownSession { - session_id: session_id(), + 3, + HostRequest::Wait { + session_id: session_id(), + request: WireWaitRequest { + cell_id: cell_id("cell-1"), + yield_time_ms: 50, }, }, + json!({ + "method": "session/wait", + "sessionId": "session-1", + "request": { "cell_id": "cell-1", "yield_time_ms": 50 }, + }), + ), + ( + 4, + HostRequest::Terminate { + session_id: session_id(), + cell_id: cell_id("cell-1"), + }, + json!({ + "method": "session/terminate", + "sessionId": "session-1", + "cellId": "cell-1", + }), + ), + ( + 5, + HostRequest::ShutdownSession { + session_id: session_id(), + }, + json!({ "method": "session/shutdown", "sessionId": "session-1" }), + ), + ] { + assert_wire_round_trip( + ClientToHost::Request { id, request }, json!({ "type": "operation/request", - "id": 8, - "request": { - "method": "session/shutdown", - "sessionId": "session-1", + "id": id, + "request": encoded_request, + }), + ); + } + + for (id, result, encoded_result) in [ + ( + 6, + WireResult::Ok { + value: DelegateResponse::ToolResult { + result: json!({ "answer": 42 }), }, + }, + json!({ + "status": "ok", + "value": { "type": "tool/result", "result": { "answer": 42 } }, + }), + ), + ( + 7, + WireResult::Ok { + value: DelegateResponse::NotificationDelivered, + }, + json!({ + "status": "ok", + "value": { "type": "notification/delivered" }, }), ), - ]; - for (message, encoded) in client_messages { - assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded); - assert_eq!( - serde_json::from_value::(encoded).expect("deserialize"), - message + ( + 8, + WireResult::Err { + message: "delegate failed".to_string(), + }, + json!({ "status": "error", "message": "delegate failed" }), + ), + ] { + assert_wire_round_trip( + ClientToHost::DelegateResponse { id, result }, + json!({ + "type": "delegate/response", + "id": id, + "result": encoded_result, + }), ); } +} - let host_messages = [ +#[test] +fn host_to_client_v1_variants_are_pinned() { + for (id, response, encoded_response) in [ ( - HostToClient::Response { - id: 7, - result: WireResult::Ok { - value: HostResponse::SessionReady { - session_id: session_id(), + 1, + HostResponse::SessionReady { + session_id: session_id(), + }, + json!({ "type": "session/ready", "sessionId": "session-1" }), + ), + ( + 2, + HostResponse::ExecutionStarted { + cell_id: cell_id("cell-1"), + }, + json!({ "type": "execution/started", "cellId": "cell-1" }), + ), + ( + 3, + HostResponse::WaitCompleted { + outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { + cell_id: cell_id("cell-1"), + content_items: content_items(), + }), + }, + json!({ + "type": "wait/completed", + "outcome": { + "LiveCell": { + "Yielded": { + "cell_id": "cell-1", + "content_items": content_items_json(), + }, }, }, + }), + ), + ( + 4, + HostResponse::WaitCompleted { + outcome: WireWaitOutcome::MissingCell(WireRuntimeResponse::Result { + cell_id: cell_id("missing-cell"), + content_items: Vec::new(), + error_text: Some("cell not found".to_string()), + }), }, json!({ - "type": "operation/response", - "id": 7, - "result": { - "status": "ok", - "value": { - "type": "session/ready", - "sessionId": "session-1", + "type": "wait/completed", + "outcome": { + "MissingCell": { + "Result": { + "cell_id": "missing-cell", + "content_items": [], + "error_text": "cell not found", + }, }, }, }), ), ( + 5, + HostResponse::SessionClosed { + session_id: session_id(), + }, + json!({ "type": "session/closed", "sessionId": "session-1" }), + ), + ] { + assert_wire_round_trip( HostToClient::Response { - id: 8, - result: WireResult::Ok { - value: HostResponse::SessionClosed { - session_id: session_id(), - }, - }, + id, + result: WireResult::Ok { value: response }, }, json!({ "type": "operation/response", - "id": 8, - "result": { - "status": "ok", - "value": { - "type": "session/closed", - "sessionId": "session-1", - }, - }, + "id": id, + "result": { "status": "ok", "value": encoded_response }, }), - ), - ]; - for (message, encoded) in host_messages { - assert_eq!(serde_json::to_value(&message).expect("serialize"), encoded); - assert_eq!( - serde_json::from_value::(encoded).expect("deserialize"), - message ); } -} - -#[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(), + assert_wire_round_trip( + HostToClient::Response { + id: 6, + result: WireResult::Err { + message: "operation failed".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", + json!({ + "type": "operation/response", + "id": 6, + "result": { "status": "error", "message": "operation failed" }, + }), + ); + + assert_wire_round_trip( + HostToClient::InitialResponse { + id: 7, + result: WireResult::Ok { + value: WireRuntimeResponse::Terminated { + cell_id: cell_id("cell-1"), + content_items: Vec::new(), + }, + }, }, - }); - assert_eq!( - serde_json::to_value(&request).expect("serialize"), - request_json + json!({ + "type": "execute/initialResponse", + "id": 7, + "result": { + "status": "ok", + "value": { + "Terminated": { "cell_id": "cell-1", "content_items": [] }, + }, + }, + }), ); - assert_eq!( - serde_json::from_value::(request_json).expect("deserialize"), - request + assert_wire_round_trip( + HostToClient::InitialResponse { + id: 8, + result: WireResult::Err { + message: "execution failed".to_string(), + }, + }, + json!({ + "type": "execute/initialResponse", + "id": 8, + "result": { "status": "error", "message": "execution failed" }, + }), ); - let response = ClientToHost::DelegateResponse { - id: 11, - result: WireResult::Ok { - value: DelegateResponse::NotificationDelivered, + assert_wire_round_trip( + HostToClient::DelegateRequest { + id: 9, + session_id: session_id(), + request: DelegateRequest::InvokeTool { + invocation: WireNestedToolCall { + cell_id: cell_id("cell-1"), + runtime_tool_call_id: "runtime-call-1".to_string(), + tool_name: WireToolName { + name: "freeform_tool".to_string(), + namespace: Some("mcp__sample__".to_string()), + }, + tool_kind: WireToolKind::Freeform, + input: Some(json!({ "value": 1 })), + }, + }, }, - }; - let response_json = json!({ - "type": "delegate/response", - "id": 11, - "result": { - "status": "ok", - "value": { "type": "notification/delivered" }, + json!({ + "type": "delegate/request", + "id": 9, + "sessionId": "session-1", + "request": { + "type": "tool/invoke", + "invocation": { + "cell_id": "cell-1", + "runtime_tool_call_id": "runtime-call-1", + "tool_name": { + "name": "freeform_tool", + "namespace": "mcp__sample__", + }, + "tool_kind": "freeform", + "input": { "value": 1 }, + }, + }, + }), + ); + assert_wire_round_trip( + HostToClient::DelegateRequest { + id: 10, + session_id: session_id(), + request: DelegateRequest::Notify { + call_id: "call-1".to_string(), + cell_id: cell_id("cell-1"), + text: "important".to_string(), + }, }, - }); - assert_eq!( - serde_json::to_value(&response).expect("serialize"), - response_json + json!({ + "type": "delegate/request", + "id": 10, + "sessionId": "session-1", + "request": { + "type": "notification/send", + "callId": "call-1", + "cellId": "cell-1", + "text": "important", + }, + }), ); - assert_eq!( - serde_json::from_value::(response_json).expect("deserialize"), - response + assert_wire_round_trip( + HostToClient::CancelDelegateRequest { id: 11 }, + json!({ "type": "delegate/cancel", "id": 11 }), + ); + assert_wire_round_trip( + HostToClient::CellClosed { + session_id: session_id(), + cell_id: cell_id("cell-1"), + }, + json!({ + "type": "cell/closed", + "sessionId": "session-1", + "cellId": "cell-1", + }), ); } @@ -302,7 +613,7 @@ fn invalid_protocol_states_cannot_be_constructed_or_decoded() { } #[test] -fn unknown_fields_are_rejected() { +fn every_nested_v1_object_rejects_unknown_fields() { assert!( serde_json::from_value::(json!({ "type": "operation/request", @@ -312,6 +623,82 @@ fn unknown_fields_are_rejected() { })) .is_err() ); + assert!( + serde_json::from_value::(json!({ + "method": "session/open", + "sessionId": "session-1", + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "tool_call_id": "call-1", + "enabled_tools": [], + "source": "text('hello');", + "yield_time_ms": null, + "max_output_tokens": null, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "name": "tool", + "tool_name": { "name": "tool", "namespace": null }, + "description": "tool", + "kind": "function", + "input_schema": null, + "output_schema": null, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "name": "tool", + "namespace": null, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "cell_id": "cell-1", + "yield_time_ms": 50, + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "Yielded": { + "cell_id": "cell-1", + "content_items": [], + "unexpected": true, + }, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "type": "input_text", + "text": "hello", + "unexpected": true, + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "cell_id": "cell-1", + "runtime_tool_call_id": "runtime-call-1", + "tool_name": { "name": "tool", "namespace": null }, + "tool_kind": "function", + "input": null, + "unexpected": true, + })) + .is_err() + ); assert!( serde_json::from_value::(json!({ "type": "operation/response", diff --git a/codex-rs/code-mode-protocol/src/host/message.rs b/codex-rs/code-mode-protocol/src/host/message.rs index 307a5802af7b..78f0961d518f 100644 --- a/codex-rs/code-mode-protocol/src/host/message.rs +++ b/codex-rs/code-mode-protocol/src/host/message.rs @@ -12,12 +12,12 @@ 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; +use super::WireCellId; +use super::WireExecuteRequest; +use super::WireNestedToolCall; +use super::WireRuntimeResponse; +use super::WireWaitOutcome; +use super::WireWaitRequest; #[derive(Clone, Debug, PartialEq, Serialize)] #[serde(rename_all = "camelCase")] @@ -155,7 +155,7 @@ pub enum HostToClient { #[serde(rename = "execute/initialResponse")] InitialResponse { id: RequestId, - result: WireResult, + result: WireResult, }, #[serde(rename = "delegate/request")] DelegateRequest { @@ -168,7 +168,7 @@ pub enum HostToClient { #[serde(rename = "cell/closed")] CellClosed { session_id: SessionId, - cell_id: CellId, + cell_id: WireCellId, }, } @@ -180,17 +180,17 @@ pub enum HostRequest { #[serde(rename = "session/execute")] Execute { session_id: SessionId, - request: ExecuteRequest, + request: WireExecuteRequest, }, #[serde(rename = "session/wait")] Wait { session_id: SessionId, - request: WaitRequest, + request: WireWaitRequest, }, #[serde(rename = "session/terminate")] Terminate { session_id: SessionId, - cell_id: CellId, + cell_id: WireCellId, }, #[serde(rename = "session/shutdown")] ShutdownSession { session_id: SessionId }, @@ -202,9 +202,9 @@ pub enum HostResponse { #[serde(rename = "session/ready")] SessionReady { session_id: SessionId }, #[serde(rename = "execution/started")] - ExecutionStarted { cell_id: CellId }, + ExecutionStarted { cell_id: WireCellId }, #[serde(rename = "wait/completed")] - WaitCompleted { outcome: WaitOutcome }, + WaitCompleted { outcome: WireWaitOutcome }, #[serde(rename = "session/closed")] SessionClosed { session_id: SessionId }, } @@ -213,11 +213,11 @@ pub enum HostResponse { #[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")] pub enum DelegateRequest { #[serde(rename = "tool/invoke")] - InvokeTool { invocation: CodeModeNestedToolCall }, + InvokeTool { invocation: WireNestedToolCall }, #[serde(rename = "notification/send")] Notify { call_id: String, - cell_id: CellId, + cell_id: WireCellId, text: String, }, } diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index c3e912071b76..7e78768ab6d6 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -8,6 +8,7 @@ mod codec; mod error; mod message; +mod payload; mod types; pub use codec::FramedReader; @@ -24,6 +25,17 @@ pub use message::HostRequest; pub use message::HostResponse; pub use message::HostToClient; pub use message::WireResult; +pub use payload::WireCellId; +pub use payload::WireContentItem; +pub use payload::WireExecuteRequest; +pub use payload::WireImageDetail; +pub use payload::WireNestedToolCall; +pub use payload::WireRuntimeResponse; +pub use payload::WireToolDefinition; +pub use payload::WireToolKind; +pub use payload::WireToolName; +pub use payload::WireWaitOutcome; +pub use payload::WireWaitRequest; pub use types::Capability; pub use types::CapabilitySet; pub use types::DelegateRequestId; diff --git a/codex-rs/code-mode-protocol/src/host/payload.rs b/codex-rs/code-mode-protocol/src/host/payload.rs new file mode 100644 index 000000000000..f6d5169a83c3 --- /dev/null +++ b/codex-rs/code-mode-protocol/src/host/payload.rs @@ -0,0 +1,406 @@ +use codex_protocol::ToolName; +use serde::Deserialize; +use serde::Serialize; +use serde_json::Value as JsonValue; + +use crate::CellId; +use crate::CodeModeNestedToolCall; +use crate::CodeModeToolKind; +use crate::ExecuteRequest; +use crate::FunctionCallOutputContentItem; +use crate::ImageDetail; +use crate::RuntimeResponse; +use crate::ToolDefinition; +use crate::WaitOutcome; +use crate::WaitRequest; + +/// A cell identifier with a wire representation owned by protocol V1. +#[derive(Clone, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(transparent)] +pub struct WireCellId(String); + +impl WireCellId { + pub fn new(value: impl Into) -> Self { + Self(value.into()) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl From for WireCellId { + fn from(value: CellId) -> Self { + Self(value.as_str().to_string()) + } +} + +impl From<&CellId> for WireCellId { + fn from(value: &CellId) -> Self { + Self(value.as_str().to_string()) + } +} + +impl From for CellId { + fn from(value: WireCellId) -> Self { + Self::new(value.0) + } +} + +/// The V1 wire representation of a tool's stable name. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireToolName { + pub name: String, + pub namespace: Option, +} + +impl From for WireToolName { + fn from(value: ToolName) -> Self { + Self { + name: value.name, + namespace: value.namespace, + } + } +} + +impl From for ToolName { + fn from(value: WireToolName) -> Self { + Self::new(value.namespace, value.name) + } +} + +/// The tool invocation shape supported by protocol V1. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum WireToolKind { + Function, + Freeform, +} + +impl From for WireToolKind { + fn from(value: CodeModeToolKind) -> Self { + match value { + CodeModeToolKind::Function => Self::Function, + CodeModeToolKind::Freeform => Self::Freeform, + } + } +} + +impl From for CodeModeToolKind { + fn from(value: WireToolKind) -> Self { + match value { + WireToolKind::Function => Self::Function, + WireToolKind::Freeform => Self::Freeform, + } + } +} + +/// A V1 tool definition embedded in an execute request. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireToolDefinition { + pub name: String, + pub tool_name: WireToolName, + pub description: String, + pub kind: WireToolKind, + pub input_schema: Option, + pub output_schema: Option, +} + +impl From for WireToolDefinition { + fn from(value: ToolDefinition) -> Self { + Self { + name: value.name, + tool_name: value.tool_name.into(), + description: value.description, + kind: value.kind.into(), + input_schema: value.input_schema, + output_schema: value.output_schema, + } + } +} + +impl From for ToolDefinition { + fn from(value: WireToolDefinition) -> Self { + Self { + name: value.name, + tool_name: value.tool_name.into(), + description: value.description, + kind: value.kind.into(), + input_schema: value.input_schema, + output_schema: value.output_schema, + } + } +} + +/// The complete execute request shape supported by protocol V1. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireExecuteRequest { + pub tool_call_id: String, + pub enabled_tools: Vec, + pub source: String, + pub yield_time_ms: Option, + pub max_output_tokens: Option, +} + +impl From for WireExecuteRequest { + fn from(value: ExecuteRequest) -> Self { + Self { + tool_call_id: value.tool_call_id, + enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(), + source: value.source, + yield_time_ms: value.yield_time_ms, + max_output_tokens: value.max_output_tokens, + } + } +} + +impl From for ExecuteRequest { + fn from(value: WireExecuteRequest) -> Self { + Self { + tool_call_id: value.tool_call_id, + enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(), + source: value.source, + yield_time_ms: value.yield_time_ms, + max_output_tokens: value.max_output_tokens, + } + } +} + +/// The complete wait request shape supported by protocol V1. +#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireWaitRequest { + pub cell_id: WireCellId, + pub yield_time_ms: u64, +} + +impl From for WireWaitRequest { + fn from(value: WaitRequest) -> Self { + Self { + cell_id: value.cell_id.into(), + yield_time_ms: value.yield_time_ms, + } + } +} + +impl From for WaitRequest { + fn from(value: WireWaitRequest) -> Self { + Self { + cell_id: value.cell_id.into(), + yield_time_ms: value.yield_time_ms, + } + } +} + +/// Image detail values accepted in a V1 runtime response. +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum WireImageDetail { + Auto, + Low, + High, + Original, +} + +impl From for WireImageDetail { + fn from(value: ImageDetail) -> Self { + match value { + ImageDetail::Auto => Self::Auto, + ImageDetail::Low => Self::Low, + ImageDetail::High => Self::High, + ImageDetail::Original => Self::Original, + } + } +} + +impl From for ImageDetail { + fn from(value: WireImageDetail) -> Self { + match value { + WireImageDetail::Auto => Self::Auto, + WireImageDetail::Low => Self::Low, + WireImageDetail::High => Self::High, + WireImageDetail::Original => Self::Original, + } + } +} + +/// One output item emitted by a V1 runtime response. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields, tag = "type", rename_all = "snake_case")] +pub enum WireContentItem { + InputText { + text: String, + }, + InputImage { + image_url: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + detail: Option, + }, +} + +impl From for WireContentItem { + fn from(value: FunctionCallOutputContentItem) -> Self { + match value { + FunctionCallOutputContentItem::InputText { text } => Self::InputText { text }, + FunctionCallOutputContentItem::InputImage { image_url, detail } => Self::InputImage { + image_url, + detail: detail.map(Into::into), + }, + } + } +} + +impl From for FunctionCallOutputContentItem { + fn from(value: WireContentItem) -> Self { + match value { + WireContentItem::InputText { text } => Self::InputText { text }, + WireContentItem::InputImage { image_url, detail } => Self::InputImage { + image_url, + detail: detail.map(Into::into), + }, + } + } +} + +/// Runtime output returned over the V1 host connection. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub enum WireRuntimeResponse { + Yielded { + cell_id: WireCellId, + content_items: Vec, + }, + Terminated { + cell_id: WireCellId, + content_items: Vec, + }, + Result { + cell_id: WireCellId, + content_items: Vec, + error_text: Option, + }, +} + +impl From for WireRuntimeResponse { + fn from(value: RuntimeResponse) -> Self { + match value { + RuntimeResponse::Yielded { + cell_id, + content_items, + } => Self::Yielded { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + RuntimeResponse::Terminated { + cell_id, + content_items, + } => Self::Terminated { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + RuntimeResponse::Result { + cell_id, + content_items, + error_text, + } => Self::Result { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + error_text, + }, + } + } +} + +impl From for RuntimeResponse { + fn from(value: WireRuntimeResponse) -> Self { + match value { + WireRuntimeResponse::Yielded { + cell_id, + content_items, + } => Self::Yielded { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + WireRuntimeResponse::Terminated { + cell_id, + content_items, + } => Self::Terminated { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + }, + WireRuntimeResponse::Result { + cell_id, + content_items, + error_text, + } => Self::Result { + cell_id: cell_id.into(), + content_items: content_items.into_iter().map(Into::into).collect(), + error_text, + }, + } + } +} + +/// Whether a waited-for cell remained live in protocol V1. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub enum WireWaitOutcome { + LiveCell(WireRuntimeResponse), + MissingCell(WireRuntimeResponse), +} + +impl From for WireWaitOutcome { + fn from(value: WaitOutcome) -> Self { + match value { + WaitOutcome::LiveCell(response) => Self::LiveCell(response.into()), + WaitOutcome::MissingCell(response) => Self::MissingCell(response.into()), + } + } +} + +impl From for WaitOutcome { + fn from(value: WireWaitOutcome) -> Self { + match value { + WireWaitOutcome::LiveCell(response) => Self::LiveCell(response.into()), + WireWaitOutcome::MissingCell(response) => Self::MissingCell(response.into()), + } + } +} + +/// A nested tool invocation sent over the V1 host connection. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(deny_unknown_fields)] +pub struct WireNestedToolCall { + pub cell_id: WireCellId, + pub runtime_tool_call_id: String, + pub tool_name: WireToolName, + pub tool_kind: WireToolKind, + pub input: Option, +} + +impl From for WireNestedToolCall { + fn from(value: CodeModeNestedToolCall) -> Self { + Self { + cell_id: value.cell_id.into(), + runtime_tool_call_id: value.runtime_tool_call_id, + tool_name: value.tool_name.into(), + tool_kind: value.tool_kind.into(), + input: value.input, + } + } +} + +impl From for CodeModeNestedToolCall { + fn from(value: WireNestedToolCall) -> Self { + Self { + cell_id: value.cell_id.into(), + runtime_tool_call_id: value.runtime_tool_call_id, + tool_name: value.tool_name.into(), + tool_kind: value.tool_kind.into(), + input: value.input, + } + } +} From 077f8199de79658661e22a7b1f97cb00f3ac0836 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Wed, 24 Jun 2026 14:43:18 -0700 Subject: [PATCH 3/6] Strengthen code mode wire integer types --- .../code-mode-protocol/src/host/host_tests.rs | 75 ++++++++++++++----- .../code-mode-protocol/src/host/payload.rs | 28 ++++--- codex-rs/code-mode-protocol/src/host/types.rs | 23 +++++- 3 files changed, 94 insertions(+), 32 deletions(-) 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 77a9cf71634e..0ccf8230c7cb 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -11,6 +11,7 @@ use super::CapabilitySet; use super::ClientHello; use super::ClientToHost; use super::DelegateRequest; +use super::DelegateRequestId; use super::DelegateResponse; use super::HandshakeRejectReason; use super::HostHello; @@ -18,6 +19,7 @@ use super::HostRequest; use super::HostResponse; use super::HostToClient; use super::ProtocolVersion; +use super::RequestId; use super::SessionId; use super::SupportedProtocolVersions; use super::WireCellId; @@ -32,6 +34,7 @@ use super::WireToolKind; use super::WireToolName; use super::WireWaitOutcome; use super::WireWaitRequest; +use crate::ExecuteRequest; fn session_id() -> SessionId { SessionId::new("session-1").expect("valid session ID") @@ -41,6 +44,14 @@ fn cell_id(value: &str) -> WireCellId { WireCellId::new(value) } +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + +fn delegate_request_id(value: i64) -> DelegateRequestId { + DelegateRequestId::new(value) +} + fn capability(value: &str) -> Capability { Capability::new(value).expect("valid capability") } @@ -225,14 +236,14 @@ fn client_to_host_v1_variants_are_pinned() { let execute_request = execute_request(); for (id, request, encoded_request) in [ ( - 1, + request_id(1), HostRequest::OpenSession { session_id: session_id(), }, json!({ "method": "session/open", "sessionId": "session-1" }), ), ( - 2, + request_id(2), HostRequest::Execute { session_id: session_id(), request: execute_request, @@ -270,7 +281,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - 3, + request_id(3), HostRequest::Wait { session_id: session_id(), request: WireWaitRequest { @@ -285,7 +296,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - 4, + request_id(4), HostRequest::Terminate { session_id: session_id(), cell_id: cell_id("cell-1"), @@ -297,7 +308,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - 5, + request_id(5), HostRequest::ShutdownSession { session_id: session_id(), }, @@ -316,7 +327,7 @@ fn client_to_host_v1_variants_are_pinned() { for (id, result, encoded_result) in [ ( - 6, + delegate_request_id(6), WireResult::Ok { value: DelegateResponse::ToolResult { result: json!({ "answer": 42 }), @@ -328,7 +339,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - 7, + delegate_request_id(7), WireResult::Ok { value: DelegateResponse::NotificationDelivered, }, @@ -338,7 +349,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - 8, + delegate_request_id(8), WireResult::Err { message: "delegate failed".to_string(), }, @@ -360,21 +371,21 @@ fn client_to_host_v1_variants_are_pinned() { fn host_to_client_v1_variants_are_pinned() { for (id, response, encoded_response) in [ ( - 1, + request_id(1), HostResponse::SessionReady { session_id: session_id(), }, json!({ "type": "session/ready", "sessionId": "session-1" }), ), ( - 2, + request_id(2), HostResponse::ExecutionStarted { cell_id: cell_id("cell-1"), }, json!({ "type": "execution/started", "cellId": "cell-1" }), ), ( - 3, + request_id(3), HostResponse::WaitCompleted { outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { cell_id: cell_id("cell-1"), @@ -394,7 +405,7 @@ fn host_to_client_v1_variants_are_pinned() { }), ), ( - 4, + request_id(4), HostResponse::WaitCompleted { outcome: WireWaitOutcome::MissingCell(WireRuntimeResponse::Result { cell_id: cell_id("missing-cell"), @@ -416,7 +427,7 @@ fn host_to_client_v1_variants_are_pinned() { }), ), ( - 5, + request_id(5), HostResponse::SessionClosed { session_id: session_id(), }, @@ -437,7 +448,7 @@ fn host_to_client_v1_variants_are_pinned() { } assert_wire_round_trip( HostToClient::Response { - id: 6, + id: request_id(6), result: WireResult::Err { message: "operation failed".to_string(), }, @@ -451,7 +462,7 @@ fn host_to_client_v1_variants_are_pinned() { assert_wire_round_trip( HostToClient::InitialResponse { - id: 7, + id: request_id(7), result: WireResult::Ok { value: WireRuntimeResponse::Terminated { cell_id: cell_id("cell-1"), @@ -472,7 +483,7 @@ fn host_to_client_v1_variants_are_pinned() { ); assert_wire_round_trip( HostToClient::InitialResponse { - id: 8, + id: request_id(8), result: WireResult::Err { message: "execution failed".to_string(), }, @@ -486,7 +497,7 @@ fn host_to_client_v1_variants_are_pinned() { assert_wire_round_trip( HostToClient::DelegateRequest { - id: 9, + id: delegate_request_id(9), session_id: session_id(), request: DelegateRequest::InvokeTool { invocation: WireNestedToolCall { @@ -522,7 +533,7 @@ fn host_to_client_v1_variants_are_pinned() { ); assert_wire_round_trip( HostToClient::DelegateRequest { - id: 10, + id: delegate_request_id(10), session_id: session_id(), request: DelegateRequest::Notify { call_id: "call-1".to_string(), @@ -543,7 +554,9 @@ fn host_to_client_v1_variants_are_pinned() { }), ); assert_wire_round_trip( - HostToClient::CancelDelegateRequest { id: 11 }, + HostToClient::CancelDelegateRequest { + id: delegate_request_id(11), + }, json!({ "type": "delegate/cancel", "id": 11 }), ); assert_wire_round_trip( @@ -559,6 +572,30 @@ fn host_to_client_v1_variants_are_pinned() { ); } +#[test] +fn execute_request_integer_bounds_are_enforced() { + let wire_request = execute_request(); + let domain_request = ExecuteRequest::try_from(wire_request.clone()) + .expect("valid wire request converts to the domain"); + assert_eq!( + WireExecuteRequest::try_from(domain_request.clone()) + .expect("valid domain request converts to the wire"), + wire_request + ); + + let too_large = ExecuteRequest { + max_output_tokens: Some(usize::try_from(i32::MAX).expect("i32::MAX fits usize") + 1), + ..domain_request + }; + assert!(WireExecuteRequest::try_from(too_large).is_err()); + + let negative = WireExecuteRequest { + max_output_tokens: Some(-1), + ..wire_request + }; + assert!(ExecuteRequest::try_from(negative).is_err()); +} + #[test] fn invalid_protocol_states_cannot_be_constructed_or_decoded() { assert!(SessionId::new("").is_err()); diff --git a/codex-rs/code-mode-protocol/src/host/payload.rs b/codex-rs/code-mode-protocol/src/host/payload.rs index f6d5169a83c3..a5ba41071105 100644 --- a/codex-rs/code-mode-protocol/src/host/payload.rs +++ b/codex-rs/code-mode-protocol/src/host/payload.rs @@ -1,3 +1,5 @@ +use std::num::TryFromIntError; + use codex_protocol::ToolName; use serde::Deserialize; use serde::Serialize; @@ -142,30 +144,34 @@ pub struct WireExecuteRequest { pub enabled_tools: Vec, pub source: String, pub yield_time_ms: Option, - pub max_output_tokens: Option, + pub max_output_tokens: Option, } -impl From for WireExecuteRequest { - fn from(value: ExecuteRequest) -> Self { - Self { +impl TryFrom for WireExecuteRequest { + type Error = TryFromIntError; + + fn try_from(value: ExecuteRequest) -> Result { + Ok(Self { tool_call_id: value.tool_call_id, enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(), source: value.source, yield_time_ms: value.yield_time_ms, - max_output_tokens: value.max_output_tokens, - } + max_output_tokens: value.max_output_tokens.map(i32::try_from).transpose()?, + }) } } -impl From for ExecuteRequest { - fn from(value: WireExecuteRequest) -> Self { - Self { +impl TryFrom for ExecuteRequest { + type Error = TryFromIntError; + + fn try_from(value: WireExecuteRequest) -> Result { + Ok(Self { tool_call_id: value.tool_call_id, enabled_tools: value.enabled_tools.into_iter().map(Into::into).collect(), source: value.source, yield_time_ms: value.yield_time_ms, - max_output_tokens: value.max_output_tokens, - } + max_output_tokens: value.max_output_tokens.map(usize::try_from).transpose()?, + }) } } diff --git a/codex-rs/code-mode-protocol/src/host/types.rs b/codex-rs/code-mode-protocol/src/host/types.rs index b51474607a2e..40c69df90b90 100644 --- a/codex-rs/code-mode-protocol/src/host/types.rs +++ b/codex-rs/code-mode-protocol/src/host/types.rs @@ -8,8 +8,27 @@ use serde::Serialize; use serde::Serializer; use serde::de::Error as _; -pub type RequestId = u64; -pub type DelegateRequestId = u64; +/// Correlates one client operation request with the host's response. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(transparent)] +pub struct RequestId(i64); + +impl RequestId { + pub const fn new(value: i64) -> Self { + Self(value) + } +} + +/// Correlates one host delegate request with the client's response. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] +#[serde(transparent)] +pub struct DelegateRequestId(i64); + +impl DelegateRequestId { + pub const fn new(value: i64) -> Self { + Self(value) + } +} #[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd, Serialize)] #[serde(transparent)] From c52c04ba7b4af03d0599b166dec4035c6e65d9b3 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 01:00:31 +0000 Subject: [PATCH 4/6] code-mode: satisfy argument comment lint --- .../code-mode-protocol/src/host/host_tests.rs | 38 +++++++++---------- 1 file changed, 19 insertions(+), 19 deletions(-) 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 0ccf8230c7cb..3aa2a18307c5 100644 --- a/codex-rs/code-mode-protocol/src/host/host_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/host_tests.rs @@ -236,14 +236,14 @@ fn client_to_host_v1_variants_are_pinned() { let execute_request = execute_request(); for (id, request, encoded_request) in [ ( - request_id(1), + request_id(/*value*/ 1), HostRequest::OpenSession { session_id: session_id(), }, json!({ "method": "session/open", "sessionId": "session-1" }), ), ( - request_id(2), + request_id(/*value*/ 2), HostRequest::Execute { session_id: session_id(), request: execute_request, @@ -281,7 +281,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - request_id(3), + request_id(/*value*/ 3), HostRequest::Wait { session_id: session_id(), request: WireWaitRequest { @@ -296,7 +296,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - request_id(4), + request_id(/*value*/ 4), HostRequest::Terminate { session_id: session_id(), cell_id: cell_id("cell-1"), @@ -308,7 +308,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - request_id(5), + request_id(/*value*/ 5), HostRequest::ShutdownSession { session_id: session_id(), }, @@ -327,7 +327,7 @@ fn client_to_host_v1_variants_are_pinned() { for (id, result, encoded_result) in [ ( - delegate_request_id(6), + delegate_request_id(/*value*/ 6), WireResult::Ok { value: DelegateResponse::ToolResult { result: json!({ "answer": 42 }), @@ -339,7 +339,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - delegate_request_id(7), + delegate_request_id(/*value*/ 7), WireResult::Ok { value: DelegateResponse::NotificationDelivered, }, @@ -349,7 +349,7 @@ fn client_to_host_v1_variants_are_pinned() { }), ), ( - delegate_request_id(8), + delegate_request_id(/*value*/ 8), WireResult::Err { message: "delegate failed".to_string(), }, @@ -371,21 +371,21 @@ fn client_to_host_v1_variants_are_pinned() { fn host_to_client_v1_variants_are_pinned() { for (id, response, encoded_response) in [ ( - request_id(1), + request_id(/*value*/ 1), HostResponse::SessionReady { session_id: session_id(), }, json!({ "type": "session/ready", "sessionId": "session-1" }), ), ( - request_id(2), + request_id(/*value*/ 2), HostResponse::ExecutionStarted { cell_id: cell_id("cell-1"), }, json!({ "type": "execution/started", "cellId": "cell-1" }), ), ( - request_id(3), + request_id(/*value*/ 3), HostResponse::WaitCompleted { outcome: WireWaitOutcome::LiveCell(WireRuntimeResponse::Yielded { cell_id: cell_id("cell-1"), @@ -405,7 +405,7 @@ fn host_to_client_v1_variants_are_pinned() { }), ), ( - request_id(4), + request_id(/*value*/ 4), HostResponse::WaitCompleted { outcome: WireWaitOutcome::MissingCell(WireRuntimeResponse::Result { cell_id: cell_id("missing-cell"), @@ -427,7 +427,7 @@ fn host_to_client_v1_variants_are_pinned() { }), ), ( - request_id(5), + request_id(/*value*/ 5), HostResponse::SessionClosed { session_id: session_id(), }, @@ -448,7 +448,7 @@ fn host_to_client_v1_variants_are_pinned() { } assert_wire_round_trip( HostToClient::Response { - id: request_id(6), + id: request_id(/*value*/ 6), result: WireResult::Err { message: "operation failed".to_string(), }, @@ -462,7 +462,7 @@ fn host_to_client_v1_variants_are_pinned() { assert_wire_round_trip( HostToClient::InitialResponse { - id: request_id(7), + id: request_id(/*value*/ 7), result: WireResult::Ok { value: WireRuntimeResponse::Terminated { cell_id: cell_id("cell-1"), @@ -483,7 +483,7 @@ fn host_to_client_v1_variants_are_pinned() { ); assert_wire_round_trip( HostToClient::InitialResponse { - id: request_id(8), + id: request_id(/*value*/ 8), result: WireResult::Err { message: "execution failed".to_string(), }, @@ -497,7 +497,7 @@ fn host_to_client_v1_variants_are_pinned() { assert_wire_round_trip( HostToClient::DelegateRequest { - id: delegate_request_id(9), + id: delegate_request_id(/*value*/ 9), session_id: session_id(), request: DelegateRequest::InvokeTool { invocation: WireNestedToolCall { @@ -533,7 +533,7 @@ fn host_to_client_v1_variants_are_pinned() { ); assert_wire_round_trip( HostToClient::DelegateRequest { - id: delegate_request_id(10), + id: delegate_request_id(/*value*/ 10), session_id: session_id(), request: DelegateRequest::Notify { call_id: "call-1".to_string(), @@ -555,7 +555,7 @@ fn host_to_client_v1_variants_are_pinned() { ); assert_wire_round_trip( HostToClient::CancelDelegateRequest { - id: delegate_request_id(11), + id: delegate_request_id(/*value*/ 11), }, json!({ "type": "delegate/cancel", "id": 11 }), ); From 464273e64d63df4abaf3226d6393a792dfa9bb0c Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Tue, 23 Jun 2026 22:18:58 -0700 Subject: [PATCH 5/6] Add process-owned code mode session --- codex-rs/Cargo.lock | 11 + codex-rs/code-mode-host/Cargo.toml | 18 + codex-rs/code-mode-host/src/delegate.rs | 88 ++++ codex-rs/code-mode-host/src/host_tests.rs | 257 +++++++++++ codex-rs/code-mode-host/src/lib.rs | 345 +++++++++++++++ codex-rs/code-mode-host/src/main.rs | 5 +- codex-rs/code-mode-host/src/peer.rs | 124 ++++++ codex-rs/code-mode-host/tests/stdio.rs | 192 +++++++++ codex-rs/code-mode/Cargo.toml | 2 +- codex-rs/code-mode/src/lib.rs | 3 + codex-rs/code-mode/src/remote_session.rs | 297 +++++++++++++ .../src/remote_session/connection.rs | 405 ++++++++++++++++++ .../src/remote_session/connection/reader.rs | 131 ++++++ .../code-mode/src/remote_session_tests.rs | 52 +++ justfile | 14 + 15 files changed, 1942 insertions(+), 2 deletions(-) create mode 100644 codex-rs/code-mode-host/src/delegate.rs create mode 100644 codex-rs/code-mode-host/src/host_tests.rs create mode 100644 codex-rs/code-mode-host/src/lib.rs create mode 100644 codex-rs/code-mode-host/src/peer.rs create mode 100644 codex-rs/code-mode-host/tests/stdio.rs create mode 100644 codex-rs/code-mode/src/remote_session.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection.rs create mode 100644 codex-rs/code-mode/src/remote_session/connection/reader.rs create mode 100644 codex-rs/code-mode/src/remote_session_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 2c50f268f0b8..dfe4a7fa67ee 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2514,6 +2514,17 @@ dependencies = [ [[package]] name = "codex-code-mode-host" version = "0.0.0" +dependencies = [ + "anyhow", + "codex-code-mode", + "codex-code-mode-protocol", + "codex-protocol", + "codex-utils-cargo-bin", + "pretty_assertions", + "serde_json", + "tokio", + "tokio-util", +] [[package]] name = "codex-code-mode-protocol" diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index a2c384d010a6..2be47134100d 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -8,5 +8,23 @@ license.workspace = true name = "codex-code-mode-host" path = "src/main.rs" +[lib] +doctest = false +name = "codex_code_mode_host" +path = "src/lib.rs" + [lints] workspace = true + +[dependencies] +anyhow = { workspace = true } +codex-code-mode = { workspace = true } +codex-code-mode-protocol = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt"] } +tokio-util = { workspace = true, features = ["rt"] } + +[dev-dependencies] +codex-protocol = { workspace = true } +codex-utils-cargo-bin = { workspace = true } +pretty_assertions = { workspace = true } +serde_json = { workspace = true } diff --git a/codex-rs/code-mode-host/src/delegate.rs b/codex-rs/code-mode-host/src/delegate.rs new file mode 100644 index 000000000000..b4fee25c0ae4 --- /dev/null +++ b/codex-rs/code-mode-host/src/delegate.rs @@ -0,0 +1,88 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeNestedToolCall; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::NotificationFuture; +use codex_code_mode_protocol::ToolInvocationFuture; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::SessionId; +use tokio_util::sync::CancellationToken; + +use crate::peer::HostPeer; + +pub(super) struct RemoteDelegate { + session_id: SessionId, + peer: Arc, +} + +impl RemoteDelegate { + pub(super) fn new(session_id: SessionId, peer: Arc) -> Self { + Self { session_id, peer } + } +} + +impl CodeModeSessionDelegate for RemoteDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::InvokeTool { + invocation: invocation.into(), + }, + cancellation_token, + ) + .await? + { + DelegateResponse::ToolResult { result } => Ok(result), + DelegateResponse::NotificationDelivered => { + Err("code-mode client returned an invalid tool result".to_string()) + } + } + }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + Box::pin(async move { + match self + .peer + .call( + self.session_id.clone(), + DelegateRequest::Notify { + call_id, + cell_id: cell_id.into(), + text, + }, + cancellation_token, + ) + .await? + { + DelegateResponse::NotificationDelivered => Ok(()), + DelegateResponse::ToolResult { .. } => { + Err("code-mode client returned an invalid notification result".to_string()) + } + } + }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.peer.send(HostToClient::CellClosed { + session_id: self.session_id.clone(), + cell_id: cell_id.into(), + }); + } +} diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs new file mode 100644 index 000000000000..4521c8c3ad25 --- /dev/null +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -0,0 +1,257 @@ +use codex_code_mode_protocol::host::Capability; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use codex_code_mode_protocol::host::WireResult; +use pretty_assertions::assert_eq; + +use super::run; + +fn client_hello( + versions: impl IntoIterator, + required_capabilities: CapabilitySet, +) -> ClientToHost { + ClientToHost::ClientHello( + ClientHello::new( + SupportedProtocolVersions::try_new(versions).expect("supported versions"), + required_capabilities, + CapabilitySet::empty(), + ) + .expect("client hello"), + ) +} + +fn session_id(value: &str) -> SessionId { + SessionId::new(value).expect("session ID") +} + +fn request_id(value: i64) -> RequestId { + RequestId::new(value) +} + +#[tokio::test] +async fn handshake_and_multiple_session_lifecycles_are_ordered() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 4096); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("read hello"), + Some(HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + CapabilitySet::empty(), + ))) + ); + + for (request_id, id) in [(request_id(1), "session-1"), (request_id(2), "session-2")] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::OpenSession { + session_id: session_id(id), + }, + }) + .await + .expect("open session"); + assert_eq!( + reader.read::().await.expect("session ready"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionReady { + session_id: session_id(id), + }, + }, + }) + ); + } + + for (request_id, id) in [(request_id(3), "session-1"), (request_id(4), "session-2")] { + writer + .write(&ClientToHost::Request { + id: request_id, + request: HostRequest::ShutdownSession { + session_id: session_id(id), + }, + }) + .await + .expect("shutdown session"); + assert_eq!( + reader.read::().await.expect("session closed"), + Some(HostToClient::Response { + id: request_id, + result: WireResult::Ok { + value: HostResponse::SessionClosed { + session_id: session_id(id), + }, + }, + }) + ); + } + + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn incompatible_or_invalid_handshake_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let version_two = ProtocolVersion::new(/*value*/ 2).expect("protocol version"); + + writer + .write(&client_hello([version_two], CapabilitySet::empty())) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { + supported_versions: SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .expect("host versions"), + }, + }) + ); + host.await.expect("host task").expect("host connection"); + + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&ClientToHost::Request { + id: request_id(1), + request: HostRequest::OpenSession { + session_id: session_id("session-1"), + }, + }) + .await + .expect("write invalid first message"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn unsupported_required_capability_is_rejected() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + let capability = Capability::new("required").expect("capability"); + + writer + .write(&client_hello( + [ProtocolVersion::V1], + CapabilitySet::try_new([capability.clone()]).expect("capabilities"), + )) + .await + .expect("write hello"); + assert_eq!( + reader.read::().await.expect("rejection"), + Some(HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { capability }, + }) + ); + host.await.expect("host task").expect("host connection"); +} + +#[tokio::test] +async fn session_id_cannot_be_reused_after_shutdown() { + let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 2048); + let (host_reader, host_writer) = tokio::io::split(host_stream); + let (client_reader, client_writer) = tokio::io::split(client_stream); + let host = tokio::spawn(run(host_reader, host_writer)); + let mut reader = FramedReader::new(client_reader); + let mut writer = FramedWriter::new(client_writer); + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write hello"); + reader + .read::() + .await + .expect("read hello") + .expect("host hello"); + + let id = session_id("session-1"); + for (request_id, request) in [ + ( + request_id(1), + HostRequest::OpenSession { + session_id: id.clone(), + }, + ), + ( + request_id(2), + HostRequest::ShutdownSession { + session_id: id.clone(), + }, + ), + ] { + writer + .write(&ClientToHost::Request { + id: request_id, + request, + }) + .await + .expect("session request"); + reader + .read::() + .await + .expect("session response") + .expect("session response message"); + } + writer + .write(&ClientToHost::Request { + id: request_id(3), + request: HostRequest::OpenSession { session_id: id }, + }) + .await + .expect("reuse session ID"); + assert_eq!( + reader.read::().await.expect("reuse response"), + Some(HostToClient::Response { + id: request_id(3), + result: WireResult::Err { + message: "code-mode session ID `session-1` was reused".to_string(), + }, + }) + ); + drop(writer); + drop(reader); + host.await.expect("host task").expect("host connection"); +} diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs new file mode 100644 index 000000000000..93f1299605b7 --- /dev/null +++ b/codex-rs/code-mode-host/src/lib.rs @@ -0,0 +1,345 @@ +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + +use anyhow::Context; +use anyhow::Result; +use codex_code_mode::InProcessCodeModeSession; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HandshakeRejectReason; +use codex_code_mode_protocol::host::HostHello; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use codex_code_mode_protocol::host::WireResult; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::sync::mpsc; +use tokio_util::task::TaskTracker; + +use self::delegate::RemoteDelegate; +use self::peer::HostPeer; + +mod delegate; +mod peer; + +/// Runs one code-mode host connection over the process standard streams. +pub async fn run_stdio() -> Result<()> { + run(tokio::io::stdin(), tokio::io::stdout()).await +} + +/// Runs one code-mode host connection over an ordered input/output pair. +async fn run(reader: R, writer: W) -> Result<()> +where + R: AsyncRead + Send + Unpin + 'static, + W: AsyncWrite + Send + Unpin + 'static, +{ + let mut reader = FramedReader::new(reader); + let mut writer = FramedWriter::new(writer); + if !negotiate(&mut reader, &mut writer).await? { + return Ok(()); + } + + let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel(); + let peer = Arc::new(HostPeer::new(outgoing_tx)); + let state = Arc::new(HostState { + sessions: Mutex::new(HashMap::new()), + seen_session_ids: Mutex::new(HashSet::new()), + request_tasks: TaskTracker::new(), + closing: AtomicBool::new(false), + peer: Arc::clone(&peer), + }); + let writer_task = tokio::spawn(async move { + while let Some(message) = outgoing_rx.recv().await { + writer + .write(&message) + .await + .context("failed to write code-mode host message")?; + } + Ok::<(), anyhow::Error>(()) + }); + + let input_result = async { + while let Some(message) = reader + .read::() + .await + .context("failed to read code-mode client message")? + { + match message { + ClientToHost::ClientHello(_) => { + anyhow::bail!("received a second code-mode client hello"); + } + ClientToHost::Request { id, request } => { + state.spawn_request(id, request); + } + ClientToHost::DelegateResponse { id, result } => { + peer.complete(id, result.into_result()).await; + } + } + } + Ok::<(), anyhow::Error>(()) + } + .await; + + peer.disconnect(); + state.disconnect().await; + drop(state); + drop(peer); + let writer_result = writer_task.await.context("code-mode writer task failed")?; + input_result?; + writer_result +} + +async fn negotiate(reader: &mut FramedReader, writer: &mut FramedWriter) -> Result +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + let Some(first_message) = reader + .read::() + .await + .context("failed to read code-mode client hello")? + else { + return Ok(false); + }; + let ClientToHost::ClientHello(client_hello) = first_message else { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::InvalidHello { + message: "first message must be connection/hello".to_string(), + }, + }) + .await + .context("failed to reject invalid code-mode client hello")?; + return Ok(false); + }; + + let supported_versions = SupportedProtocolVersions::try_new([ProtocolVersion::V1])?; + if !client_hello + .supported_versions() + .contains(ProtocolVersion::V1) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::NoCompatibleVersion { supported_versions }, + }) + .await + .context("failed to reject incompatible code-mode client")?; + return Ok(false); + } + + let host_capabilities = CapabilitySet::empty(); + if let Some(capability) = client_hello + .required_capabilities() + .iter() + .find(|capability| !host_capabilities.contains(capability)) + { + writer + .write(&HostToClient::HandshakeRejected { + reason: HandshakeRejectReason::MissingRequiredCapability { + capability: capability.clone(), + }, + }) + .await + .context("failed to reject unsupported code-mode capability")?; + return Ok(false); + } + + writer + .write(&HostToClient::HostHello(HostHello::new( + ProtocolVersion::V1, + host_capabilities, + ))) + .await + .context("failed to write code-mode host hello")?; + Ok(true) +} + +struct HostState { + sessions: Mutex>>, + seen_session_ids: Mutex>, + request_tasks: TaskTracker, + closing: AtomicBool, + peer: Arc, +} + +impl HostState { + fn spawn_request(self: &Arc, request_id: RequestId, request: HostRequest) { + let state = Arc::clone(self); + self.request_tasks.spawn(async move { + state.handle_request(request_id, request).await; + }); + } + + async fn handle_request(&self, request_id: RequestId, request: HostRequest) { + if self.closing.load(Ordering::Acquire) { + self.respond( + request_id, + Err("code-mode host is shutting down".to_string()), + ); + return; + } + match request { + HostRequest::OpenSession { session_id } => { + let result = self + .open_session(session_id.clone()) + .map(|()| HostResponse::SessionReady { session_id }); + self.respond(request_id, result); + } + HostRequest::Execute { + session_id, + request, + } => { + let request = match request.try_into() { + Ok(request) => request, + Err(err) => { + self.respond( + request_id, + Err(format!("invalid code-mode execute request: {err}")), + ); + return; + } + }; + let result = match self.session(&session_id) { + Ok(session) => session.execute(request).await, + Err(err) => { + self.respond(request_id, Err(err)); + return; + } + }; + match result { + Ok(started) => { + let cell_id = started.cell_id.clone(); + self.respond( + request_id, + Ok(HostResponse::ExecutionStarted { + cell_id: cell_id.into(), + }), + ); + self.peer.send(HostToClient::InitialResponse { + id: request_id, + result: WireResult::from_result( + started.initial_response().await.map(Into::into), + ), + }); + } + Err(err) => self.respond(request_id, Err(err)), + } + } + HostRequest::Wait { + session_id, + request, + } => { + let result = match self.session(&session_id) { + Ok(session) => session.wait(request.into()).await.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::Terminate { + session_id, + cell_id, + } => { + let result = match self.session(&session_id) { + Ok(session) => session.terminate(cell_id.into()).await.map(|outcome| { + HostResponse::WaitCompleted { + outcome: outcome.into(), + } + }), + Err(err) => Err(err), + }; + self.respond(request_id, result); + } + HostRequest::ShutdownSession { session_id } => { + let session = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .remove(&session_id); + let result = match session { + Some(session) => session + .shutdown() + .await + .map(|()| HostResponse::SessionClosed { session_id }), + None => Err(format!("unknown code-mode session {session_id}")), + }; + self.respond(request_id, result); + } + } + } + + fn open_session(&self, session_id: SessionId) -> Result<(), String> { + let mut sessions = self.sessions.lock().unwrap_or_else(PoisonError::into_inner); + if self.closing.load(Ordering::Acquire) { + return Err("code-mode host is shutting down".to_string()); + } + if !self + .seen_session_ids + .lock() + .unwrap_or_else(PoisonError::into_inner) + .insert(session_id.clone()) + { + return Err(format!("code-mode session ID `{session_id}` was reused")); + } + let delegate = Arc::new(RemoteDelegate::new( + session_id.clone(), + Arc::clone(&self.peer), + )); + sessions.insert( + session_id, + Arc::new(InProcessCodeModeSession::with_delegate(delegate)), + ); + Ok(()) + } + + fn session(&self, session_id: &SessionId) -> Result, String> { + self.sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .get(session_id) + .cloned() + .ok_or_else(|| format!("unknown code-mode session {session_id}")) + } + + fn respond(&self, id: RequestId, result: Result) { + self.peer.send(HostToClient::Response { + id, + result: WireResult::from_result(result), + }); + } + + async fn disconnect(&self) { + self.closing.store(true, Ordering::Release); + let sessions = self + .sessions + .lock() + .unwrap_or_else(PoisonError::into_inner) + .drain() + .map(|(_, session)| session) + .collect::>(); + for session in sessions { + let _ = session.shutdown().await; + } + self.request_tasks.close(); + self.request_tasks.wait().await; + } +} + +#[cfg(test)] +#[path = "host_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index f328e4d9d04c..b215869c8ab4 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1 +1,4 @@ -fn main() {} +#[tokio::main(flavor = "current_thread")] +async fn main() -> anyhow::Result<()> { + codex_code_mode_host::run_stdio().await +} diff --git a/codex-rs/code-mode-host/src/peer.rs b/codex-rs/code-mode-host/src/peer.rs new file mode 100644 index 000000000000..582e01a518b5 --- /dev/null +++ b/codex-rs/code-mode-host/src/peer.rs @@ -0,0 +1,124 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::SessionId; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +pub(super) struct HostPeer { + outgoing_tx: mpsc::UnboundedSender, + pending: Mutex>>>, + next_request_id: AtomicI64, + disconnected: CancellationToken, +} + +impl HostPeer { + pub(super) fn new(outgoing_tx: mpsc::UnboundedSender) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + next_request_id: AtomicI64::new(1), + disconnected: CancellationToken::new(), + } + } + + pub(super) fn send(&self, message: HostToClient) -> bool { + self.outgoing_tx.send(message).is_ok() + } + + pub(super) async fn call( + self: &Arc, + session_id: SessionId, + request: DelegateRequest, + cancellation_token: CancellationToken, + ) -> Result { + if self.disconnected.is_cancelled() { + return Err("code-mode client connection closed".to_string()); + } + let id = DelegateRequestId::new(self.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + self.pending.lock().await.insert(id, response_tx); + let mut pending = PendingDelegateRequest::new(Arc::clone(self), id); + if !self.send(HostToClient::DelegateRequest { + id, + session_id, + request, + }) { + self.pending.lock().await.remove(&id); + pending.disarm(); + return Err("code-mode client connection closed".to_string()); + } + + tokio::select! { + response = response_rx => { + pending.disarm(); + response.map_err(|_| { + "code-mode client closed before returning delegate output".to_string() + })? + } + _ = cancellation_token.cancelled() => { + if self.pending.lock().await.remove(&id).is_some() { + self.send(HostToClient::CancelDelegateRequest { id }); + } + pending.disarm(); + Err("code mode delegate request cancelled".to_string()) + } + _ = self.disconnected.cancelled() => { + self.pending.lock().await.remove(&id); + pending.disarm(); + Err("code-mode client connection closed".to_string()) + } + } + } + + pub(super) async fn complete( + &self, + id: DelegateRequestId, + response: Result, + ) { + if let Some(sender) = self.pending.lock().await.remove(&id) { + let _ = sender.send(response); + } + } + + pub(super) fn disconnect(&self) { + self.disconnected.cancel(); + } +} + +struct PendingDelegateRequest { + peer: Arc, + id: Option, +} + +impl PendingDelegateRequest { + fn new(peer: Arc, id: DelegateRequestId) -> Self { + Self { peer, id: Some(id) } + } + + fn disarm(&mut self) { + self.id = None; + } +} + +impl Drop for PendingDelegateRequest { + fn drop(&mut self) { + let Some(id) = self.id.take() else { + return; + }; + let peer = Arc::clone(&self.peer); + tokio::spawn(async move { + if peer.pending.lock().await.remove(&id).is_some() { + peer.send(HostToClient::CancelDelegateRequest { id }); + } + }); + } +} diff --git a/codex-rs/code-mode-host/tests/stdio.rs b/codex-rs/code-mode-host/tests/stdio.rs new file mode 100644 index 000000000000..4c2c5d17185a --- /dev/null +++ b/codex-rs/code-mode-host/tests/stdio.rs @@ -0,0 +1,192 @@ +#![allow(clippy::expect_used)] + +use std::sync::Arc; +use std::sync::Mutex; + +use codex_code_mode::CellId; +use codex_code_mode::CodeModeNestedToolCall; +use codex_code_mode::CodeModeSession; +use codex_code_mode::CodeModeSessionDelegate; +use codex_code_mode::CodeModeSessionProvider; +use codex_code_mode::CodeModeToolKind; +use codex_code_mode::ExecuteRequest; +use codex_code_mode::FunctionCallOutputContentItem; +use codex_code_mode::NotificationFuture; +use codex_code_mode::ProcessOwnedCodeModeSessionProvider; +use codex_code_mode::RuntimeResponse; +use codex_code_mode::ToolDefinition; +use codex_code_mode::ToolInvocationFuture; +use codex_code_mode::WaitOutcome; +use codex_code_mode::WaitRequest; +use codex_protocol::ToolName; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio_util::sync::CancellationToken; + +#[derive(Default)] +struct RecordingDelegate { + invocations: Mutex>, + notifications: Mutex>, + closed_cells: Mutex>, +} + +impl CodeModeSessionDelegate for RecordingDelegate { + fn invoke_tool<'a>( + &'a self, + invocation: CodeModeNestedToolCall, + _cancellation_token: CancellationToken, + ) -> ToolInvocationFuture<'a> { + self.invocations + .lock() + .expect("invocations lock") + .push(invocation); + Box::pin(async { Ok(json!({ "value": "output" })) }) + } + + fn notify<'a>( + &'a self, + call_id: String, + cell_id: CellId, + text: String, + _cancellation_token: CancellationToken, + ) -> NotificationFuture<'a> { + self.notifications + .lock() + .expect("notifications lock") + .push((call_id, cell_id, text)); + Box::pin(async { Ok(()) }) + } + + fn cell_closed(&self, cell_id: &CellId) { + self.closed_cells + .lock() + .expect("closed cells lock") + .push(cell_id.clone()); + } +} + +fn cell_id(value: &str) -> CellId { + CellId::new(value.to_string()) +} + +fn execute_request(source: &str) -> ExecuteRequest { + ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: source.to_string(), + yield_time_ms: None, + max_output_tokens: None, + } +} + +async fn execute(session: &Arc, request: ExecuteRequest) -> RuntimeResponse { + session + .execute(request) + .await + .expect("start execution") + .initial_response() + .await + .expect("initial response") +} + +#[tokio::test] +async fn remote_session_persists_values_forwards_delegates_and_controls_cells() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").expect("host binary"), + ); + let delegate = Arc::new(RecordingDelegate::default()); + let session = provider + .create_session(delegate.clone()) + .await + .expect("create remote session"); + + assert_eq!( + execute(&session, execute_request(r#"store("key", "persisted");"#),).await, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: None, + } + ); + + let mut callback_request = execute_request( + r#" +const result = await tools.echo({ value: String(load("key")) }); +notify("notice"); +text(result.value); +"#, + ); + callback_request.tool_call_id = "call-2".to_string(); + callback_request.enabled_tools = vec![ToolDefinition { + name: "echo".to_string(), + tool_name: ToolName::plain("echo"), + description: String::new(), + kind: CodeModeToolKind::Function, + input_schema: None, + output_schema: None, + }]; + assert_eq!( + execute(&session, callback_request).await, + RuntimeResponse::Result { + cell_id: cell_id("2"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "output".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + *delegate.invocations.lock().expect("invocations lock"), + vec![CodeModeNestedToolCall { + cell_id: cell_id("2"), + runtime_tool_call_id: "tool-1".to_string(), + tool_name: ToolName::plain("echo"), + tool_kind: CodeModeToolKind::Function, + input: Some(json!({ "value": "persisted" })), + }] + ); + assert_eq!( + *delegate.notifications.lock().expect("notifications lock"), + vec![("call-2".to_string(), cell_id("2"), "notice".to_string())] + ); + + let mut pending_request = execute_request("await new Promise(() => {});"); + pending_request.tool_call_id = "call-3".to_string(); + pending_request.yield_time_ms = Some(1); + assert_eq!( + execute(&session, pending_request).await, + RuntimeResponse::Yielded { + cell_id: cell_id("3"), + content_items: Vec::new(), + } + ); + assert_eq!( + session + .wait(WaitRequest { + cell_id: cell_id("3"), + yield_time_ms: 1, + }) + .await + .expect("wait for cell"), + WaitOutcome::LiveCell(RuntimeResponse::Yielded { + cell_id: cell_id("3"), + content_items: Vec::new(), + }) + ); + assert_eq!( + session + .terminate(cell_id("3")) + .await + .expect("terminate cell"), + WaitOutcome::LiveCell(RuntimeResponse::Terminated { + cell_id: cell_id("3"), + content_items: Vec::new(), + }) + ); + + session.shutdown().await.expect("shutdown remote session"); + assert_eq!( + *delegate.closed_cells.lock().expect("closed cells lock"), + vec![cell_id("1"), cell_id("2"), cell_id("3")] + ); +} diff --git a/codex-rs/code-mode/Cargo.toml b/codex-rs/code-mode/Cargo.toml index 09f1d7380dc4..513628a50772 100644 --- a/codex-rs/code-mode/Cargo.toml +++ b/codex-rs/code-mode/Cargo.toml @@ -20,7 +20,7 @@ codex-code-mode-protocol = { workspace = true } codex-protocol = { workspace = true } deno_core_icudata = { workspace = true } serde_json = { workspace = true } -tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } +tokio = { workspace = true, features = ["io-util", "macros", "process", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } v8 = { workspace = true } diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 089bb6f78da8..b186bd00e67d 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,9 +1,12 @@ mod cell_actor; +mod remote_session; mod runtime; mod service; mod session_runtime; pub use codex_code_mode_protocol::*; +pub use remote_session::ProcessOwnedCodeModeSession; +pub use remote_session::ProcessOwnedCodeModeSessionProvider; pub use service::InProcessCodeModeSession; pub use service::InProcessCodeModeSessionProvider; pub use service::NoopCodeModeSessionDelegate; diff --git a/codex-rs/code-mode/src/remote_session.rs b/codex-rs/code-mode/src/remote_session.rs new file mode 100644 index 000000000000..f5eb443809e4 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session.rs @@ -0,0 +1,297 @@ +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSession; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::CodeModeSessionProvider; +use codex_code_mode_protocol::CodeModeSessionProviderFuture; +use codex_code_mode_protocol::CodeModeSessionResultFuture; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::SessionId; +use tokio::sync::Semaphore; + +use self::connection::Connection; +use crate::NoopCodeModeSessionDelegate; + +mod connection; + +const CODE_MODE_HOST_PATH_ENV: &str = "CODEX_CODE_MODE_HOST_PATH"; + +/// Creates code-mode sessions backed by one lazily spawned process host. +pub struct ProcessOwnedCodeModeSessionProvider { + host_program: PathBuf, + process_host: StdMutex>>, +} + +impl ProcessOwnedCodeModeSessionProvider { + pub fn with_host_program(host_program: PathBuf) -> Self { + Self { + host_program, + process_host: StdMutex::new(None), + } + } + + fn process_host(&self) -> Arc { + let mut process_host = self + .process_host + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(process_host) = process_host.as_ref() { + return Arc::clone(process_host); + } + + let new_process_host = Arc::new(OwnedProcessHost::new(self.host_program.clone())); + *process_host = Some(Arc::clone(&new_process_host)); + new_process_host + } +} + +impl Default for ProcessOwnedCodeModeSessionProvider { + fn default() -> Self { + Self::with_host_program(default_host_program()) + } +} + +impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider { + fn create_session<'a>( + &'a self, + delegate: Arc, + ) -> CodeModeSessionProviderFuture<'a> { + let session = ProcessOwnedCodeModeSession::with_process_host(delegate, self.process_host()); + Box::pin(async move { + session.connection().await?; + let session: Arc = Arc::new(session); + Ok(session) + }) + } +} + +struct OwnedProcessHost { + host_program: PathBuf, + connection: StdMutex>>, + spawn_permit: Semaphore, + next_session_id: AtomicU64, +} + +impl OwnedProcessHost { + fn new(host_program: PathBuf) -> Self { + Self { + host_program, + connection: StdMutex::new(None), + spawn_permit: Semaphore::new(/*permits*/ 1), + next_session_id: AtomicU64::new(1), + } + } + + async fn connection(&self) -> Result, String> { + if let Some(connection) = self.live_connection() { + return Ok(connection); + } + + let _spawn_permit = self + .spawn_permit + .acquire() + .await + .map_err(|_| "code-mode host spawn coordinator closed".to_string())?; + if let Some(connection) = self.live_connection() { + return Ok(connection); + } + let new_connection = Arc::new(Connection::spawn(&self.host_program).await?); + *self + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(Arc::clone(&new_connection)); + Ok(new_connection) + } + + fn live_connection(&self) -> Option> { + self.connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .filter(|connection| connection.is_alive()) + .cloned() + } + + fn allocate_session_id(&self) -> SessionId { + let value = self.next_session_id.fetch_add(1, Ordering::Relaxed); + match SessionId::new(format!("session-{value}")) { + Ok(session_id) => session_id, + Err(_) => unreachable!("a generated code-mode session ID is nonempty"), + } + } +} + +enum SessionState { + New, + Open(Arc), + Shutdown, +} + +/// A logical code-mode session assigned to a process-owned host. +pub struct ProcessOwnedCodeModeSession { + process_host: Arc, + session_id: SessionId, + delegate: Arc, + state: StdMutex, + transition_permit: Semaphore, +} + +impl ProcessOwnedCodeModeSession { + pub fn new() -> Self { + Self::with_process_host( + Arc::new(NoopCodeModeSessionDelegate), + Arc::new(OwnedProcessHost::new(default_host_program())), + ) + } + + fn with_process_host( + delegate: Arc, + process_host: Arc, + ) -> Self { + let session_id = process_host.allocate_session_id(); + Self { + process_host, + session_id, + delegate, + state: StdMutex::new(SessionState::New), + transition_permit: Semaphore::new(/*permits*/ 1), + } + } + + async fn connection(&self) -> Result, String> { + if let Some(connection) = self.current_connection()? { + return Ok(connection); + } + + let _transition_permit = self + .transition_permit + .acquire() + .await + .map_err(|_| "code-mode session transition coordinator closed".to_string())?; + if let Some(connection) = self.current_connection()? { + return Ok(connection); + } + let connection = self.process_host.connection().await?; + connection + .open_session(self.session_id.clone(), Arc::clone(&self.delegate)) + .await?; + *self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + SessionState::Open(Arc::clone(&connection)); + Ok(connection) + } + + fn current_connection(&self) -> Result>, String> { + match &*self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + { + SessionState::New => Ok(None), + SessionState::Open(connection) => Ok(Some(Arc::clone(connection))), + SessionState::Shutdown => Err("code mode session is shutting down".to_string()), + } + } + + pub async fn execute(&self, request: ExecuteRequest) -> Result { + self.connection() + .await? + .execute(self.session_id.clone(), request) + .await + } + + pub async fn wait(&self, request: WaitRequest) -> Result { + self.connection() + .await? + .wait(self.session_id.clone(), request) + .await + } + + pub async fn terminate(&self, cell_id: CellId) -> Result { + self.connection() + .await? + .terminate(self.session_id.clone(), cell_id) + .await + } + + pub async fn shutdown(&self) -> Result<(), String> { + let transition_permit = self + .transition_permit + .acquire() + .await + .map_err(|_| "code-mode session transition coordinator closed".to_string())?; + let connection = { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match std::mem::replace(&mut *state, SessionState::Shutdown) { + SessionState::Open(connection) => connection, + SessionState::New | SessionState::Shutdown => return Ok(()), + } + }; + drop(transition_permit); + connection.shutdown_session(self.session_id.clone()).await + } +} + +impl Default for ProcessOwnedCodeModeSession { + fn default() -> Self { + Self::new() + } +} + +impl CodeModeSession for ProcessOwnedCodeModeSession { + fn execute<'a>( + &'a self, + request: ExecuteRequest, + ) -> CodeModeSessionResultFuture<'a, StartedCell> { + Box::pin(ProcessOwnedCodeModeSession::execute(self, request)) + } + + fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(ProcessOwnedCodeModeSession::wait(self, request)) + } + + fn terminate<'a>(&'a self, cell_id: CellId) -> CodeModeSessionResultFuture<'a, WaitOutcome> { + Box::pin(ProcessOwnedCodeModeSession::terminate(self, cell_id)) + } + + fn shutdown<'a>(&'a self) -> CodeModeSessionResultFuture<'a, ()> { + Box::pin(ProcessOwnedCodeModeSession::shutdown(self)) + } +} + +fn default_host_program() -> PathBuf { + if let Some(path) = std::env::var_os(CODE_MODE_HOST_PATH_ENV) { + return PathBuf::from(path); + } + let executable_name = if cfg!(windows) { + "codex-code-mode-host.exe" + } else { + "codex-code-mode-host" + }; + if let Ok(current_exe) = std::env::current_exe() + && let Some(parent) = current_exe.parent() + { + let sibling = parent.join(executable_name); + if sibling.is_file() { + return sibling; + } + } + PathBuf::from(executable_name) +} + +#[cfg(test)] +#[path = "remote_session_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode/src/remote_session/connection.rs b/codex-rs/code-mode/src/remote_session/connection.rs new file mode 100644 index 000000000000..413f25fc21f4 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection.rs @@ -0,0 +1,405 @@ +use std::collections::HashMap; +use std::path::Path; +use std::process::Stdio; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicI64; +use std::sync::atomic::Ordering; + +use codex_code_mode_protocol::CellId; +use codex_code_mode_protocol::CodeModeSessionDelegate; +use codex_code_mode_protocol::ExecuteRequest; +use codex_code_mode_protocol::RuntimeResponse; +use codex_code_mode_protocol::StartedCell; +use codex_code_mode_protocol::WaitOutcome; +use codex_code_mode_protocol::WaitRequest; +use codex_code_mode_protocol::host::CapabilitySet; +use codex_code_mode_protocol::host::ClientHello; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::DelegateRequestId; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HostRequest; +use codex_code_mode_protocol::host::HostResponse; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::ProtocolVersion; +use codex_code_mode_protocol::host::RequestId; +use codex_code_mode_protocol::host::SessionId; +use codex_code_mode_protocol::host::SupportedProtocolVersions; +use tokio::io::AsyncBufReadExt; +use tokio::io::BufReader; +use tokio::process::Command; +use tokio::sync::Mutex; +use tokio::sync::mpsc; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; +use tracing::debug; +use tracing::warn; + +use self::reader::drive_reader; + +mod reader; + +const IPC_CHANNEL_CAPACITY: usize = 128; + +pub(super) struct Connection { + state: Arc, + cancellation: CancellationToken, +} + +impl Connection { + pub(super) async fn spawn(host_program: &Path) -> Result { + let mut command = Command::new(host_program); + #[cfg(unix)] + command.process_group(0); + let mut child = command + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|err| { + format!( + "failed to spawn code-mode host {}: {err}", + host_program.display() + ) + })?; + let stdin = child + .stdin + .take() + .ok_or_else(|| "spawned code-mode host has no stdin".to_string())?; + let stdout = child + .stdout + .take() + .ok_or_else(|| "spawned code-mode host has no stdout".to_string())?; + let stderr = child.stderr.take(); + let mut reader = FramedReader::new(stdout); + let mut writer = FramedWriter::new(stdin); + + let hello = ClientHello::new( + SupportedProtocolVersions::try_new([ProtocolVersion::V1]) + .map_err(|err| err.to_string())?, + CapabilitySet::empty(), + CapabilitySet::empty(), + ) + .map_err(|err| err.to_string())?; + writer + .write(&ClientToHost::ClientHello(hello)) + .await + .map_err(|err| format!("failed to write code-mode host hello: {err}"))?; + match reader + .read::() + .await + .map_err(|err| format!("failed to read code-mode host hello: {err}"))? + { + Some(HostToClient::HostHello(hello)) + if hello.selected_version() == ProtocolVersion::V1 => {} + Some(HostToClient::HandshakeRejected { reason }) => { + return Err(format!("code-mode host rejected the handshake: {reason:?}")); + } + Some(message) => { + return Err(format!( + "code-mode host returned an invalid handshake response: {message:?}" + )); + } + None => return Err("code-mode host exited during handshake".to_string()), + } + + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(IPC_CHANNEL_CAPACITY); + let cancellation = CancellationToken::new(); + let state = Arc::new(ConnectionState::new(outgoing_tx, cancellation.clone())); + + let writer_state = Arc::downgrade(&state); + let writer_cancellation = cancellation.clone(); + tokio::spawn(async move { + loop { + tokio::select! { + _ = writer_cancellation.cancelled() => break, + message = outgoing_rx.recv() => { + let Some(message) = message else { + break; + }; + if let Err(err) = writer.write(&message).await { + if let Some(state) = writer_state.upgrade() { + state.fail(format!( + "failed to write code-mode host message: {err}" + )).await; + } + break; + } + } + } + } + }); + + let reader_state = Arc::downgrade(&state); + let reader_cancellation = cancellation.clone(); + tokio::spawn(async move { + drive_reader(reader, reader_state, reader_cancellation).await; + }); + + if let Some(stderr) = stderr { + tokio::spawn(async move { + let mut lines = BufReader::new(stderr).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => debug!("code-mode host stderr: {line}"), + Ok(None) => break, + Err(err) => { + warn!("failed to read code-mode host stderr: {err}"); + break; + } + } + } + }); + } + + let supervisor_state = Arc::downgrade(&state); + let supervisor_cancellation = cancellation.clone(); + tokio::spawn(async move { + tokio::select! { + result = child.wait() => { + let reason = match result { + Ok(status) => format!("code-mode host exited with status {status}"), + Err(err) => format!("failed waiting for code-mode host: {err}"), + }; + if let Some(state) = supervisor_state.upgrade() { + state.fail(reason).await; + } + } + _ = supervisor_cancellation.cancelled() => { + let _ = child.start_kill(); + let _ = child.wait().await; + } + } + }); + + Ok(Self { + state, + cancellation, + }) + } + + pub(super) fn is_alive(&self) -> bool { + self.state.alive.load(Ordering::Acquire) + } + + pub(super) async fn open_session( + &self, + session_id: SessionId, + delegate: Arc, + ) -> Result<(), String> { + self.state + .delegates + .lock() + .await + .insert(session_id.clone(), delegate); + let response = self + .request(HostRequest::OpenSession { + session_id: session_id.clone(), + }) + .await; + match response { + Ok(HostResponse::SessionReady { + session_id: ready_session_id, + }) if ready_session_id == session_id => Ok(()), + Ok(_) => { + self.state.delegates.lock().await.remove(&session_id); + Err("code-mode host returned an invalid open-session response".to_string()) + } + Err(err) => { + self.state.delegates.lock().await.remove(&session_id); + Err(err) + } + } + } + + pub(super) async fn execute( + &self, + session_id: SessionId, + request: ExecuteRequest, + ) -> Result { + let request = request + .try_into() + .map_err(|err| format!("failed to encode code-mode execute request: {err}"))?; + let id = RequestId::new(self.state.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + let (initial_tx, initial_rx) = oneshot::channel(); + self.state.pending.lock().await.insert(id, response_tx); + self.state + .initial_responses + .lock() + .await + .insert(id, initial_tx); + if let Err(err) = self + .state + .send(ClientToHost::Request { + id, + request: HostRequest::Execute { + session_id, + request, + }, + }) + .await + { + self.state.pending.lock().await.remove(&id); + self.state.initial_responses.lock().await.remove(&id); + return Err(err); + } + let response = match response_rx.await { + Ok(Ok(response)) => response, + Ok(Err(err)) => { + self.state.initial_responses.lock().await.remove(&id); + return Err(err); + } + Err(_) => { + self.state.initial_responses.lock().await.remove(&id); + return Err(self.state.failure_message()); + } + }; + match response { + HostResponse::ExecutionStarted { cell_id } => Ok(StartedCell::from_result_receiver( + cell_id.into(), + initial_rx, + )), + _ => { + self.state.initial_responses.lock().await.remove(&id); + Err("code-mode host returned an invalid execute response".to_string()) + } + } + } + + pub(super) async fn wait( + &self, + session_id: SessionId, + request: WaitRequest, + ) -> Result { + match self + .request(HostRequest::Wait { + session_id, + request: request.into(), + }) + .await? + { + HostResponse::WaitCompleted { outcome } => Ok(outcome.into()), + _ => Err("code-mode host returned an invalid wait response".to_string()), + } + } + + pub(super) async fn terminate( + &self, + session_id: SessionId, + cell_id: CellId, + ) -> Result { + match self + .request(HostRequest::Terminate { + session_id, + cell_id: cell_id.into(), + }) + .await? + { + HostResponse::WaitCompleted { outcome } => Ok(outcome.into()), + _ => Err("code-mode host returned an invalid terminate response".to_string()), + } + } + + pub(super) async fn shutdown_session(&self, session_id: SessionId) -> Result<(), String> { + let response = self + .request(HostRequest::ShutdownSession { + session_id: session_id.clone(), + }) + .await; + self.state.delegates.lock().await.remove(&session_id); + match response? { + HostResponse::SessionClosed { + session_id: closed_session_id, + } if closed_session_id == session_id => Ok(()), + _ => Err("code-mode host returned an invalid shutdown response".to_string()), + } + } + + async fn request(&self, request: HostRequest) -> Result { + let id = RequestId::new(self.state.next_request_id.fetch_add(1, Ordering::Relaxed)); + let (response_tx, response_rx) = oneshot::channel(); + self.state.pending.lock().await.insert(id, response_tx); + if let Err(err) = self.state.send(ClientToHost::Request { id, request }).await { + self.state.pending.lock().await.remove(&id); + return Err(err); + } + response_rx + .await + .map_err(|_| self.state.failure_message())? + } +} + +impl Drop for Connection { + fn drop(&mut self) { + self.cancellation.cancel(); + } +} + +struct ConnectionState { + outgoing_tx: mpsc::Sender, + pending: Mutex>>>, + initial_responses: Mutex>>>, + delegates: Mutex>>, + delegate_cancellations: Mutex>, + next_request_id: AtomicI64, + alive: AtomicBool, + failure: std::sync::Mutex>, + cancellation: CancellationToken, +} + +impl ConnectionState { + fn new(outgoing_tx: mpsc::Sender, cancellation: CancellationToken) -> Self { + Self { + outgoing_tx, + pending: Mutex::new(HashMap::new()), + initial_responses: Mutex::new(HashMap::new()), + delegates: Mutex::new(HashMap::new()), + delegate_cancellations: Mutex::new(HashMap::new()), + next_request_id: AtomicI64::new(1), + alive: AtomicBool::new(true), + failure: std::sync::Mutex::new(None), + cancellation, + } + } + + async fn send(&self, message: ClientToHost) -> Result<(), String> { + if !self.alive.load(Ordering::Acquire) { + return Err(self.failure_message()); + } + self.outgoing_tx + .send(message) + .await + .map_err(|_| self.failure_message()) + } + + fn failure_message(&self) -> String { + self.failure + .lock() + .ok() + .and_then(|failure| failure.clone()) + .unwrap_or_else(|| "code-mode host connection closed".to_string()) + } + + async fn fail(&self, reason: String) { + if !self.alive.swap(false, Ordering::AcqRel) { + return; + } + if let Ok(mut failure) = self.failure.lock() { + *failure = Some(reason.clone()); + } + self.cancellation.cancel(); + for (_, sender) in self.pending.lock().await.drain() { + let _ = sender.send(Err(reason.clone())); + } + for (_, sender) in self.initial_responses.lock().await.drain() { + let _ = sender.send(Err(reason.clone())); + } + self.delegates.lock().await.clear(); + for (_, cancellation) in self.delegate_cancellations.lock().await.drain() { + cancellation.cancel(); + } + } +} diff --git a/codex-rs/code-mode/src/remote_session/connection/reader.rs b/codex-rs/code-mode/src/remote_session/connection/reader.rs new file mode 100644 index 000000000000..277a2930243b --- /dev/null +++ b/codex-rs/code-mode/src/remote_session/connection/reader.rs @@ -0,0 +1,131 @@ +use std::sync::Arc; +use std::sync::Weak; + +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::WireResult; +use tokio::process::ChildStdout; +use tokio_util::sync::CancellationToken; + +use super::ConnectionState; + +pub(super) async fn drive_reader( + mut reader: FramedReader, + state: Weak, + cancellation: CancellationToken, +) { + loop { + let message = tokio::select! { + _ = cancellation.cancelled() => return, + result = reader.read::() => result, + }; + match message { + Ok(Some(message)) => { + let Some(state) = state.upgrade() else { + return; + }; + handle_host_message(state, message).await; + } + Ok(None) => { + if let Some(state) = state.upgrade() { + state + .fail("code-mode host closed its stdout".to_string()) + .await; + } + return; + } + Err(err) => { + if let Some(state) = state.upgrade() { + state + .fail(format!("failed to read code-mode host message: {err}")) + .await; + } + return; + } + } + } +} + +async fn handle_host_message(state: Arc, message: HostToClient) { + match message { + HostToClient::Response { id, result } => { + if let Some(sender) = state.pending.lock().await.remove(&id) { + let _ = sender.send(result.into_result()); + } + } + HostToClient::InitialResponse { id, result } => { + if let Some(sender) = state.initial_responses.lock().await.remove(&id) { + let _ = sender.send(result.into_result().map(Into::into)); + } + } + HostToClient::DelegateRequest { + id, + session_id, + request, + } => { + let delegate = state.delegates.lock().await.get(&session_id).cloned(); + let Some(delegate) = delegate else { + let _ = state + .send(ClientToHost::DelegateResponse { + id, + result: WireResult::Err { + message: format!("unknown code-mode session {session_id}"), + }, + }) + .await; + return; + }; + let cancellation = CancellationToken::new(); + state + .delegate_cancellations + .lock() + .await + .insert(id, cancellation.clone()); + tokio::spawn(async move { + let result = match request { + DelegateRequest::InvokeTool { invocation } => delegate + .invoke_tool(invocation.into(), cancellation) + .await + .map(|result| DelegateResponse::ToolResult { result }), + DelegateRequest::Notify { + call_id, + cell_id, + text, + } => delegate + .notify(call_id, cell_id.into(), text, cancellation) + .await + .map(|()| DelegateResponse::NotificationDelivered), + }; + state.delegate_cancellations.lock().await.remove(&id); + let _ = state + .send(ClientToHost::DelegateResponse { + id, + result: WireResult::from_result(result), + }) + .await; + }); + } + HostToClient::CancelDelegateRequest { id } => { + if let Some(cancellation) = state.delegate_cancellations.lock().await.remove(&id) { + cancellation.cancel(); + } + } + HostToClient::CellClosed { + session_id, + cell_id, + } => { + if let Some(delegate) = state.delegates.lock().await.get(&session_id).cloned() { + let cell_id = cell_id.into(); + delegate.cell_closed(&cell_id); + } + } + HostToClient::HostHello(_) | HostToClient::HandshakeRejected { .. } => { + state + .fail("code-mode host sent a second handshake response".to_string()) + .await; + } + } +} diff --git a/codex-rs/code-mode/src/remote_session_tests.rs b/codex-rs/code-mode/src/remote_session_tests.rs new file mode 100644 index 000000000000..d97120755361 --- /dev/null +++ b/codex-rs/code-mode/src/remote_session_tests.rs @@ -0,0 +1,52 @@ +use std::sync::Arc; + +use codex_code_mode_protocol::CodeModeSessionProvider; + +use super::ProcessOwnedCodeModeSession; +use super::ProcessOwnedCodeModeSessionProvider; +use crate::NoopCodeModeSessionDelegate; + +#[test] +fn provider_reuses_its_live_process_host() { + let provider = ProcessOwnedCodeModeSessionProvider::default(); + + let first = provider.process_host(); + let second = provider.process_host(); + + assert!(Arc::ptr_eq(&first, &second)); +} + +#[tokio::test] +async fn provider_reports_host_spawn_failure() { + let provider = ProcessOwnedCodeModeSessionProvider::with_host_program( + "codex-code-mode-host-does-not-exist".into(), + ); + + let error = provider + .create_session(Arc::new(NoopCodeModeSessionDelegate)) + .await + .err() + .expect("session creation should fail"); + + assert!(error.contains("failed to spawn code-mode host")); +} + +#[tokio::test] +async fn shutdown_before_open_does_not_spawn_the_host() { + let session = ProcessOwnedCodeModeSession::new(); + + session.shutdown().await.expect("shutdown session"); + let error = session + .execute(codex_code_mode_protocol::ExecuteRequest { + tool_call_id: "call-1".to_string(), + enabled_tools: Vec::new(), + source: "text('unreachable')".to_string(), + yield_time_ms: None, + max_output_tokens: None, + }) + .await + .err() + .expect("shutdown session should reject execution"); + + assert_eq!(error, "code mode session is shutting down"); +} diff --git a/justfile b/justfile index 03f87391ac4d..d4a0e6566c67 100644 --- a/justfile +++ b/justfile @@ -31,6 +31,10 @@ tui-with-exec-server *args: file-search *args: cargo run --bin codex-file-search -- {args} +# Run the standalone code-mode host from source. +code-mode-host *args: + cargo run --bin codex-code-mode-host -- {args} + # Build the CLI and run the app-server test client app-server-test-client *args: cargo build -p codex-cli @@ -107,6 +111,16 @@ bazel-codex *args: bazel-codex *args: bazel run //codex-rs/cli:codex --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) +# Build and run the standalone code-mode host from source using Bazel. +[no-cd] +[unix] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under="cd $PWD &&" -- "$@" + +[windows] +bazel-code-mode-host *args: + bazel run //codex-rs/code-mode-host:codex-code-mode-host --run_under='cd /d "{{ invocation_directory_native() }}" &&' -- @($args | Select-Object -Skip 1) + [no-cd] bazel-lock-update: bazel mod deps --lockfile_mode=update From 1886ad634427d3f28eead4f5a4ec87ce6df129d1 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Thu, 25 Jun 2026 01:03:58 +0000 Subject: [PATCH 6/6] code-mode: annotate host request ids --- codex-rs/code-mode-host/src/host_tests.rs | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index 4521c8c3ad25..1dde4d4b1713 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -61,7 +61,10 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() { ))) ); - for (request_id, id) in [(request_id(1), "session-1"), (request_id(2), "session-2")] { + for (request_id, id) in [ + (request_id(/*value*/ 1), "session-1"), + (request_id(/*value*/ 2), "session-2"), + ] { writer .write(&ClientToHost::Request { id: request_id, @@ -84,7 +87,10 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() { ); } - for (request_id, id) in [(request_id(3), "session-1"), (request_id(4), "session-2")] { + for (request_id, id) in [ + (request_id(/*value*/ 3), "session-1"), + (request_id(/*value*/ 4), "session-2"), + ] { writer .write(&ClientToHost::Request { id: request_id, @@ -145,7 +151,7 @@ async fn incompatible_or_invalid_handshake_is_rejected() { let mut writer = FramedWriter::new(client_writer); writer .write(&ClientToHost::Request { - id: request_id(1), + id: request_id(/*value*/ 1), request: HostRequest::OpenSession { session_id: session_id("session-1"), }, @@ -210,13 +216,13 @@ async fn session_id_cannot_be_reused_after_shutdown() { let id = session_id("session-1"); for (request_id, request) in [ ( - request_id(1), + request_id(/*value*/ 1), HostRequest::OpenSession { session_id: id.clone(), }, ), ( - request_id(2), + request_id(/*value*/ 2), HostRequest::ShutdownSession { session_id: id.clone(), }, @@ -237,7 +243,7 @@ async fn session_id_cannot_be_reused_after_shutdown() { } writer .write(&ClientToHost::Request { - id: request_id(3), + id: request_id(/*value*/ 3), request: HostRequest::OpenSession { session_id: id }, }) .await @@ -245,7 +251,7 @@ async fn session_id_cannot_be_reused_after_shutdown() { assert_eq!( reader.read::().await.expect("reuse response"), Some(HostToClient::Response { - id: request_id(3), + id: request_id(/*value*/ 3), result: WireResult::Err { message: "code-mode session ID `session-1` was reused".to_string(), },