From 0dfa778dae6a94b2ff2c69176cbaf063a3bf18a1 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Fri, 24 Jul 2026 02:35:19 +0000 Subject: [PATCH] Add WebSocket transport to the code-mode host (#35078) ## What changed - Add a `--listen` option that accepts `stdio`, `stdio://`, or a `ws://IP:PORT` endpoint, while retaining stdio as the default. - Serve the existing length-prefixed protocol in binary WebSocket messages, with isolated connections, shared host limits, and a `/readyz` endpoint. - Reject browser-origin handshakes and contain malformed frames to the affected connection. ## Testing - Cover listen URL parsing and complete-frame encoding and decoding. - Exercise readiness, cell execution, tool callbacks, large frames, concurrent connections, malformed frames, and origin rejection through the WebSocket listener. GitOrigin-RevId: 01c8be4c6256b8ce4a3a0002440dcb3294e5f887 --- codex-rs/Cargo.lock | 6 + codex-rs/code-mode-host/Cargo.toml | 8 +- codex-rs/code-mode-host/src/host_tests.rs | 97 ++++- codex-rs/code-mode-host/src/lib.rs | 61 ++- codex-rs/code-mode-host/src/main.rs | 21 +- codex-rs/code-mode-host/src/transport.rs | 233 ++++++++++ .../code-mode-host/src/transport_tests.rs | 53 +++ codex-rs/code-mode-host/tests/websocket.rs | 405 ++++++++++++++++++ codex-rs/code-mode-protocol/src/host/codec.rs | 51 ++- .../src/host/codec_tests.rs | 33 ++ codex-rs/code-mode-protocol/src/host/mod.rs | 2 +- 11 files changed, 946 insertions(+), 24 deletions(-) create mode 100644 codex-rs/code-mode-host/src/transport.rs create mode 100644 codex-rs/code-mode-host/src/transport_tests.rs create mode 100644 codex-rs/code-mode-host/tests/websocket.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 43fc9f2e17..bf1d6f84ce 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2502,15 +2502,21 @@ name = "codex-code-mode-host" version = "0.0.0" dependencies = [ "anyhow", + "axum", + "clap", "codex-code-mode", "codex-code-mode-protocol", "codex-protocol", "codex-utils-cargo-bin", + "futures", "pretty_assertions", "serde_json", "tempfile", "tokio", + "tokio-tungstenite", "tokio-util", + "tracing", + "tracing-subscriber", ] [[package]] diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 86a36d2a27..0d9ee48b02 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -18,10 +18,15 @@ workspace = true [dependencies] anyhow = { workspace = true } +axum = { workspace = true, features = ["http1", "tokio", "ws"] } +clap = { workspace = true, features = ["derive"] } codex-code-mode = { workspace = true } codex-code-mode-protocol = { workspace = true } -tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt", "sync", "time"] } +futures = { workspace = true } +tokio = { workspace = true, features = ["io-std", "io-util", "macros", "net", "process", "rt", "sync", "time"] } tokio-util = { workspace = true, features = ["rt"] } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } [dev-dependencies] codex-protocol = { workspace = true } @@ -29,3 +34,4 @@ codex-utils-cargo-bin = { workspace = true } pretty_assertions = { workspace = true } serde_json = { workspace = true } tempfile = { workspace = true } +tokio-tungstenite = { workspace = true } diff --git a/codex-rs/code-mode-host/src/host_tests.rs b/codex-rs/code-mode-host/src/host_tests.rs index b189b89e09..b3639e7cc0 100644 --- a/codex-rs/code-mode-host/src/host_tests.rs +++ b/codex-rs/code-mode-host/src/host_tests.rs @@ -1,3 +1,13 @@ +use std::collections::HashMap; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::Mutex; +use std::sync::PoisonError; +use std::sync::atomic::AtomicBool; +use std::task::Context; +use std::task::Poll; +use std::time::Duration; + use codex_code_mode_protocol::host::Capability; use codex_code_mode_protocol::host::CapabilitySet; use codex_code_mode_protocol::host::ClientHello; @@ -17,8 +27,10 @@ use codex_code_mode_protocol::host::SupportedProtocolVersions; use codex_code_mode_protocol::host::WireExecuteRequest; use codex_code_mode_protocol::host::WireResult; use pretty_assertions::assert_eq; +use tokio::io::AsyncWrite; use tokio::sync::Semaphore; use tokio::sync::mpsc; +use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; @@ -159,6 +171,85 @@ async fn handshake_and_multiple_session_lifecycles_are_ordered() { host.await.expect("host task").expect("host connection"); } +#[tokio::test] +async fn disconnect_cancels_a_backpressured_host_writer() { + let (host_reader, client_writer) = tokio::io::duplex(/*max_buf_size*/ 4096); + let (blocked_tx, blocked_rx) = oneshot::channel(); + let host = tokio::spawn(run( + host_reader, + BlockingWriter { + blocked_tx: Some(blocked_tx), + handshake_flushed: false, + }, + )); + let mut writer = FramedWriter::new(client_writer); + + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write client hello"); + writer + .write(&ClientToHost::Request { + id: request_id(/*value*/ 1), + request: HostRequest::OpenSession { + session_id: session_id("backpressured-session"), + }, + }) + .await + .expect("write session-open request"); + + tokio::time::timeout(Duration::from_secs(1), blocked_rx) + .await + .expect("host writer should reach backpressure") + .expect("host writer should report backpressure"); + + writer + .write(&client_hello([ProtocolVersion::V1], CapabilitySet::empty())) + .await + .expect("write invalid second client hello"); + + let error = tokio::time::timeout(Duration::from_secs(1), host) + .await + .expect("disconnect should cancel the backpressured writer") + .expect("host task should finish") + .expect_err("a second client hello should fail the connection"); + assert_eq!( + error.to_string(), + "received a second code-mode client hello" + ); +} + +struct BlockingWriter { + blocked_tx: Option>, + handshake_flushed: bool, +} + +impl AsyncWrite for BlockingWriter { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + bytes: &[u8], + ) -> Poll> { + if self.handshake_flushed { + if let Some(blocked_tx) = self.blocked_tx.take() { + let _ = blocked_tx.send(()); + } + Poll::Pending + } else { + Poll::Ready(Ok(bytes.len())) + } + } + + fn poll_flush(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.handshake_flushed = true; + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} + #[tokio::test] async fn incompatible_or_invalid_handshake_is_rejected() { let (host_stream, client_stream) = tokio::io::duplex(/*max_buf_size*/ 1024); @@ -478,9 +569,3 @@ async fn cell_forwarding_panic_disconnects_host() { .contains("cell forwarding task failed") ); } -use std::collections::HashMap; -use std::sync::Arc; -use std::sync::Mutex; -use std::sync::PoisonError; -use std::sync::atomic::AtomicBool; -use std::time::Duration; diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index d2c1f507b4..737787301a 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -14,8 +14,6 @@ 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::EncodedFrame; -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; @@ -34,9 +32,14 @@ use tokio_util::task::TaskTracker; use self::delegate::RemoteDelegate; use self::peer::HostPeer; +use self::transport::ConnectionReader; +use self::transport::ConnectionWriter; + +pub use self::transport::DEFAULT_LISTEN_URL; mod delegate; mod peer; +mod transport; const MAX_IN_FLIGHT_REQUESTS: usize = 256; const MAX_ACTIVE_CELLS: usize = 128; @@ -44,6 +47,25 @@ const MAX_RECENT_REQUEST_IDS: usize = 4096; const MAX_RECENT_SESSION_IDS: usize = 4096; const SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5); +struct HostLimits { + request_permits: Arc, + active_cell_permits: Arc, +} + +impl HostLimits { + fn new() -> Self { + Self { + request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), + active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + } + } +} + +/// Runs the code-mode host on its configured stdio or WebSocket transport. +pub async fn run_main(listen_url: &str) -> Result<()> { + transport::run_transport(listen_url).await +} + /// 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 @@ -55,8 +77,19 @@ where R: AsyncRead + Send + Unpin + 'static, W: AsyncWrite + Send + Unpin + 'static, { - let mut reader = FramedReader::new(reader); - let mut writer = FramedWriter::new(writer); + run_connection( + ConnectionReader::from_reader(reader), + ConnectionWriter::from_writer(writer), + Arc::new(HostLimits::new()), + ) + .await +} + +async fn run_connection( + mut reader: ConnectionReader, + mut writer: ConnectionWriter, + limits: Arc, +) -> Result<()> { if !negotiate(&mut reader, &mut writer).await? { return Ok(()); } @@ -68,8 +101,8 @@ where seen_session_ids: Mutex::new(SeenSessionIds::default()), requests: Mutex::new(RequestRegistry::default()), request_tasks: TaskTracker::new(), - request_permits: Arc::new(Semaphore::new(MAX_IN_FLIGHT_REQUESTS)), - active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)), + request_permits: Arc::clone(&limits.request_permits), + active_cell_permits: Arc::clone(&limits.active_cell_permits), closing: AtomicBool::new(false), peer: Arc::clone(&peer), }); @@ -82,7 +115,11 @@ where let Some(frame) = frame else { return Ok(()); }; - if let Err(err) = writer.write_frame(&frame).await { + let result = tokio::select! { + _ = writer_disconnected.cancelled() => return Ok(()), + result = writer.write_frame(frame) => result, + }; + if let Err(err) = result { return Err( anyhow::Error::new(err) .context("failed to write code-mode host message") @@ -112,7 +149,7 @@ where loop { let message = tokio::select! { _ = peer.disconnected() => break, - message = reader.read::() => message + message = reader.read() => message .context("failed to read code-mode client message")?, }; let Some(message) = message else { @@ -158,13 +195,9 @@ where Ok(()) } -async fn negotiate(reader: &mut FramedReader, writer: &mut FramedWriter) -> Result -where - R: AsyncRead + Unpin, - W: AsyncWrite + Unpin, -{ +async fn negotiate(reader: &mut ConnectionReader, writer: &mut ConnectionWriter) -> Result { let Some(first_message) = reader - .read::() + .read() .await .context("failed to read code-mode client hello")? else { diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index b215869c8a..50ec048288 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1,4 +1,23 @@ +use clap::Parser; + +#[derive(Debug, Parser)] +struct Cli { + /// Transport endpoint: `stdio`, `stdio://`, or `ws://IP:PORT`. + #[arg( + long, + value_name = "URL", + default_value = codex_code_mode_host::DEFAULT_LISTEN_URL + )] + listen: String, +} + #[tokio::main(flavor = "current_thread")] async fn main() -> anyhow::Result<()> { - codex_code_mode_host::run_stdio().await + tracing_subscriber::fmt() + .with_max_level(tracing::Level::INFO) + .with_writer(std::io::stderr) + .with_ansi(false) + .init(); + + codex_code_mode_host::run_main(&Cli::parse().listen).await } diff --git a/codex-rs/code-mode-host/src/transport.rs b/codex-rs/code-mode-host/src/transport.rs new file mode 100644 index 0000000000..3342fc217c --- /dev/null +++ b/codex-rs/code-mode-host/src/transport.rs @@ -0,0 +1,233 @@ +use std::io; +use std::io::Write as _; +use std::net::SocketAddr; +use std::sync::Arc; + +use anyhow::Context; +use anyhow::Result; +use axum::Router; +use axum::body::Body; +use axum::extract::ConnectInfo; +use axum::extract::State; +use axum::extract::ws::Message; +use axum::extract::ws::WebSocket; +use axum::extract::ws::WebSocketUpgrade; +use axum::http::Request; +use axum::http::StatusCode; +use axum::http::header::ORIGIN; +use axum::middleware; +use axum::middleware::Next; +use axum::response::IntoResponse; +use axum::response::Response; +use axum::routing::any; +use axum::routing::get; +use codex_code_mode_protocol::host::ClientToHost; +use codex_code_mode_protocol::host::EncodedFrame; +use codex_code_mode_protocol::host::FramedReader; +use codex_code_mode_protocol::host::FramedWriter; +use codex_code_mode_protocol::host::HostToClient; +use codex_code_mode_protocol::host::MAX_FRAME_BYTES; +use futures::SinkExt; +use futures::StreamExt; +use futures::stream::SplitSink; +use futures::stream::SplitStream; +use tokio::io::AsyncRead; +use tokio::io::AsyncWrite; +use tokio::net::TcpListener; +use tracing::info; +use tracing::warn; + +use crate::HostLimits; + +/// The default transport retains the standalone host's original stdio behavior. +pub const DEFAULT_LISTEN_URL: &str = "stdio"; + +const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::(); + +type BoxedReader = Box; +type BoxedWriter = Box; + +#[derive(Debug, Clone, Eq, PartialEq)] +enum ListenTransport { + Stdio, + WebSocket(SocketAddr), +} + +pub(crate) enum ConnectionReader { + Framed(FramedReader), + WebSocket(SplitStream), +} + +pub(crate) enum ConnectionWriter { + Framed(FramedWriter), + WebSocket(SplitSink), +} + +#[derive(Clone)] +struct WebSocketListenerState { + limits: Arc, +} + +impl ConnectionReader { + pub(crate) fn from_reader(reader: R) -> Self + where + R: AsyncRead + Send + Unpin + 'static, + { + Self::Framed(FramedReader::new(Box::new(reader))) + } + + pub(crate) async fn read(&mut self) -> io::Result> { + match self { + Self::Framed(reader) => reader.read().await, + Self::WebSocket(reader) => loop { + match reader.next().await { + Some(Ok(Message::Binary(bytes))) => { + return EncodedFrame::decode_framed(&bytes).map(Some); + } + Some(Ok(Message::Text(_))) => { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "code-mode websocket messages must be binary framed messages", + )); + } + Some(Ok(Message::Ping(_) | Message::Pong(_))) => {} + Some(Ok(Message::Close(_))) | None => return Ok(None), + Some(Err(err)) => { + return Err(io::Error::other(format!( + "failed to read code-mode websocket message: {err}" + ))); + } + } + }, + } + } +} + +impl ConnectionWriter { + pub(crate) fn from_writer(writer: W) -> Self + where + W: AsyncWrite + Send + Unpin + 'static, + { + Self::Framed(FramedWriter::new(Box::new(writer))) + } + + pub(crate) async fn write(&mut self, message: &HostToClient) -> io::Result<()> { + self.write_frame(EncodedFrame::encode(message)?).await + } + + pub(crate) async fn write_frame(&mut self, frame: EncodedFrame) -> io::Result<()> { + match self { + Self::Framed(writer) => writer.write_frame(&frame).await, + Self::WebSocket(writer) => writer + .send(Message::Binary(frame.into_framed_bytes().into())) + .await + .map_err(|err| { + io::Error::other(format!( + "failed to write code-mode websocket message: {err}" + )) + }), + } + } +} + +pub(crate) async fn run_transport(listen_url: &str) -> Result<()> { + match parse_listen_url(listen_url)? { + ListenTransport::Stdio => crate::run_stdio().await, + ListenTransport::WebSocket(bind_address) => run_websocket_listener(bind_address).await, + } +} + +fn parse_listen_url(listen_url: &str) -> Result { + if matches!(listen_url, "stdio" | "stdio://") { + return Ok(ListenTransport::Stdio); + } + + if let Some(socket_addr) = listen_url.strip_prefix("ws://") { + return socket_addr + .parse::() + .map(ListenTransport::WebSocket) + .with_context(|| { + format!("invalid websocket --listen URL `{listen_url}`; expected `ws://IP:PORT`") + }); + } + + anyhow::bail!( + "unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`, `stdio`, or `stdio://`" + ); +} + +async fn run_websocket_listener(bind_address: SocketAddr) -> Result<()> { + let listener = TcpListener::bind(bind_address) + .await + .with_context(|| format!("failed to bind code-mode host websocket to {bind_address}"))?; + let local_addr = listener + .local_addr() + .context("failed to read code-mode host websocket listen address")?; + let state = WebSocketListenerState { + limits: Arc::new(HostLimits::new()), + }; + info!("codex-code-mode-host listening on ws://{local_addr}"); + println!("ws://{local_addr}"); + io::stdout() + .flush() + .context("failed to publish code-mode host websocket listen address")?; + + let router = Router::new() + .route("/", any(websocket_upgrade_handler)) + .route("/readyz", get(readiness_handler)) + .layer(middleware::from_fn(reject_requests_with_origin_header)) + .with_state(state); + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await + .context("code-mode host websocket listener failed") +} + +async fn readiness_handler() -> StatusCode { + StatusCode::OK +} + +async fn reject_requests_with_origin_header( + request: Request, + next: Next, +) -> Result { + if request.headers().contains_key(ORIGIN) { + warn!( + method = %request.method(), + uri = %request.uri(), + "rejecting code-mode host websocket request with Origin header" + ); + Err(StatusCode::FORBIDDEN) + } else { + Ok(next.run(request).await) + } +} + +async fn websocket_upgrade_handler( + websocket: WebSocketUpgrade, + ConnectInfo(peer_addr): ConnectInfo, + State(state): State, +) -> impl IntoResponse { + websocket + .max_frame_size(MAX_WEBSOCKET_FRAME_BYTES) + .max_message_size(MAX_WEBSOCKET_FRAME_BYTES) + .on_upgrade(move |stream| async move { + info!(%peer_addr, "code-mode host websocket client connected"); + let (writer, reader) = stream.split(); + if let Err(err) = crate::run_connection( + ConnectionReader::WebSocket(reader), + ConnectionWriter::WebSocket(writer), + state.limits, + ) + .await + { + warn!(%peer_addr, "code-mode host websocket connection failed: {err:#}"); + } + }) +} + +#[cfg(test)] +#[path = "transport_tests.rs"] +mod tests; diff --git a/codex-rs/code-mode-host/src/transport_tests.rs b/codex-rs/code-mode-host/src/transport_tests.rs new file mode 100644 index 0000000000..b62b63b380 --- /dev/null +++ b/codex-rs/code-mode-host/src/transport_tests.rs @@ -0,0 +1,53 @@ +use std::net::SocketAddr; + +use pretty_assertions::assert_eq; + +use super::ListenTransport; +use super::parse_listen_url; + +#[test] +fn parse_listen_url_accepts_stdio_transports() { + assert_eq!( + parse_listen_url("stdio").expect("stdio listen URL should parse"), + ListenTransport::Stdio + ); + assert_eq!( + parse_listen_url("stdio://").expect("stdio URL should parse"), + ListenTransport::Stdio + ); +} + +#[test] +fn parse_listen_url_accepts_websocket_addresses() { + assert_eq!( + parse_listen_url("ws://127.0.0.1:0").expect("websocket listen URL should parse"), + ListenTransport::WebSocket( + "127.0.0.1:0" + .parse::() + .expect("valid socket address") + ) + ); + assert_eq!( + parse_listen_url("ws://[::1]:9000").expect("IPv6 websocket listen URL should parse"), + ListenTransport::WebSocket( + "[::1]:9000" + .parse::() + .expect("valid IPv6 socket address") + ) + ); +} + +#[test] +fn parse_listen_url_rejects_invalid_transports() { + let invalid_address = parse_listen_url("ws://localhost:9000") + .expect_err("websocket listener requires an IP address"); + assert!( + invalid_address + .to_string() + .contains("expected `ws://IP:PORT`") + ); + + let unsupported = + parse_listen_url("http://127.0.0.1:9000").expect_err("HTTP is not a listen transport"); + assert!(unsupported.to_string().contains("unsupported --listen URL")); +} diff --git a/codex-rs/code-mode-host/tests/websocket.rs b/codex-rs/code-mode-host/tests/websocket.rs new file mode 100644 index 0000000000..a0f471e737 --- /dev/null +++ b/codex-rs/code-mode-host/tests/websocket.rs @@ -0,0 +1,405 @@ +use std::process::Stdio; +use std::time::Duration; + +use anyhow::Context; +use anyhow::Result; +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::DelegateRequest; +use codex_code_mode_protocol::host::DelegateResponse; +use codex_code_mode_protocol::host::EncodedFrame; +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::MAX_FRAME_BYTES; +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::WireContentItem; +use codex_code_mode_protocol::host::WireExecuteRequest; +use codex_code_mode_protocol::host::WireResult; +use codex_code_mode_protocol::host::WireRuntimeResponse; +use codex_code_mode_protocol::host::WireToolDefinition; +use codex_code_mode_protocol::host::WireToolKind; +use codex_code_mode_protocol::host::WireToolName; +use futures::SinkExt; +use futures::StreamExt; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use tokio::net::TcpStream; +use tokio::process::Child; +use tokio::process::Command; +use tokio::time::timeout; +use tokio_tungstenite::MaybeTlsStream; +use tokio_tungstenite::WebSocketStream; +use tokio_tungstenite::connect_async; +use tokio_tungstenite::connect_async_with_config; +use tokio_tungstenite::tungstenite::Error as WebSocketError; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::HeaderValue; +use tokio_tungstenite::tungstenite::http::StatusCode; +use tokio_tungstenite::tungstenite::http::header::ORIGIN; +use tokio_tungstenite::tungstenite::protocol::WebSocketConfig; + +const TEST_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::(); + +struct HostHarness { + child: Child, + websocket_url: String, +} + +struct HostClient { + websocket: WebSocketStream>, +} + +impl HostHarness { + async fn start() -> Result { + let host_program = codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?; + let mut command = Command::new(host_program); + command + .args(["--listen", "ws://127.0.0.1:0"]) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + let mut child = command.spawn().context("failed to start code-mode host")?; + let stdout = child + .stdout + .take() + .context("code-mode host stdout was not captured")?; + let mut lines = BufReader::new(stdout).lines(); + let websocket_url = timeout(TEST_TIMEOUT, lines.next_line()) + .await + .context("timed out waiting for code-mode host websocket URL")?? + .context("code-mode host exited before publishing its websocket URL")?; + if !websocket_url.starts_with("ws://127.0.0.1:") { + anyhow::bail!("unexpected code-mode host websocket URL `{websocket_url}`"); + } + + Ok(Self { + child, + websocket_url, + }) + } + + async fn connect(&self) -> Result { + let config = WebSocketConfig::default() + .max_frame_size(Some(MAX_WEBSOCKET_FRAME_BYTES)) + .max_message_size(Some(MAX_WEBSOCKET_FRAME_BYTES)); + let (websocket, _) = timeout( + TEST_TIMEOUT, + connect_async_with_config( + self.websocket_url.as_str(), + Some(config), + /*disable_nagle*/ false, + ), + ) + .await + .context("timed out connecting to code-mode host websocket")??; + Ok(HostClient { websocket }) + } +} + +impl HostClient { + async fn send(&mut self, message: &ClientToHost) -> Result<()> { + let frame = EncodedFrame::encode(message)?; + self.send_binary(frame.into_framed_bytes()).await + } + + async fn send_binary(&mut self, bytes: Vec) -> Result<()> { + timeout( + TEST_TIMEOUT, + self.websocket.send(Message::Binary(bytes.into())), + ) + .await + .context("timed out writing code-mode websocket message")? + .context("failed to write code-mode websocket message") + } + + async fn read(&mut self) -> Result { + loop { + let message = timeout(TEST_TIMEOUT, self.websocket.next()) + .await + .context("timed out waiting for code-mode websocket message")? + .context("code-mode websocket closed before returning a message")? + .context("failed to read code-mode websocket message")?; + match message { + Message::Binary(bytes) => { + return EncodedFrame::decode_framed(&bytes) + .context("failed to decode code-mode websocket frame"); + } + Message::Ping(_) | Message::Pong(_) => {} + Message::Close(frame) => { + anyhow::bail!("code-mode websocket closed unexpectedly: {frame:?}"); + } + Message::Text(text) => { + anyhow::bail!("code-mode host returned a text websocket message: {text}"); + } + Message::Frame(_) => { + anyhow::bail!("code-mode host returned an unexpected raw websocket frame"); + } + } + } + } + + async fn negotiate(&mut self, optional_capabilities: CapabilitySet) -> Result<()> { + let hello = ClientHello::new( + SupportedProtocolVersions::try_new([ProtocolVersion::V1])?, + CapabilitySet::empty(), + optional_capabilities, + )?; + self.send(&ClientToHost::ClientHello(hello)).await?; + assert_eq!( + self.read().await?, + HostToClient::HostHello(HostHello::new(ProtocolVersion::V1, CapabilitySet::empty())) + ); + Ok(()) + } + + async fn open_session(&mut self, session_id: SessionId) -> Result<()> { + let id = RequestId::new(/*value*/ 1); + self.send(&ClientToHost::Request { + id, + request: HostRequest::OpenSession { + session_id: session_id.clone(), + }, + }) + .await?; + assert_eq!( + self.read().await?, + HostToClient::Response { + id, + result: WireResult::Ok { + value: HostResponse::SessionReady { session_id }, + }, + } + ); + Ok(()) + } +} + +#[tokio::test] +async fn websocket_listener_serves_readiness_endpoint() -> Result<()> { + let host = HostHarness::start().await?; + let address = host + .websocket_url + .strip_prefix("ws://") + .context("code-mode host websocket URL should use ws://")?; + + let response = timeout(TEST_TIMEOUT, async { + let mut stream = TcpStream::connect(address) + .await + .context("failed to connect to code-mode host readiness endpoint")?; + let request = + format!("GET /readyz HTTP/1.1\r\nHost: {address}\r\nConnection: close\r\n\r\n"); + stream + .write_all(request.as_bytes()) + .await + .context("failed to request code-mode host readiness")?; + + let mut response = String::new(); + stream + .read_to_string(&mut response) + .await + .context("failed to read code-mode host readiness response")?; + Ok::<_, anyhow::Error>(response) + }) + .await + .context("timed out requesting code-mode host readiness")??; + + let status_line = response + .lines() + .next() + .context("code-mode host readiness response is missing a status line")?; + assert_eq!(status_line, "HTTP/1.1 200 OK"); + Ok(()) +} + +#[tokio::test] +async fn websocket_listener_executes_cells_and_forwards_tool_callbacks() -> Result<()> { + let host = HostHarness::start().await?; + let mut client = host.connect().await?; + client.negotiate(CapabilitySet::empty()).await?; + + let session_id = SessionId::new("websocket-session")?; + client.open_session(session_id.clone()).await?; + + let execute_id = RequestId::new(/*value*/ 2); + client + .send(&ClientToHost::Request { + id: execute_id, + request: HostRequest::Execute { + session_id: session_id.clone(), + request: WireExecuteRequest { + tool_call_id: "websocket-call".to_string(), + enabled_tools: vec![WireToolDefinition { + name: "echo".to_string(), + tool_name: WireToolName { + name: "echo".to_string(), + namespace: None, + }, + description: String::new(), + kind: WireToolKind::Function, + input_schema: None, + output_schema: None, + }], + source: + r#"const result = await tools.echo({ value: "ping" }); text(result.value);"# + .to_string(), + yield_time_ms: Some(5_000), + max_output_tokens: Some(1_000), + }, + }, + }) + .await?; + + let started = client.read().await?; + let HostToClient::Response { + id, + result: + WireResult::Ok { + value: HostResponse::ExecutionStarted { cell_id }, + }, + } = started + else { + anyhow::bail!("expected execution-started response, got {started:?}"); + }; + assert_eq!(id, execute_id); + + let callback = client.read().await?; + let HostToClient::DelegateRequest { + id: delegate_id, + session_id: callback_session_id, + request: DelegateRequest::InvokeTool { invocation }, + } = callback + else { + anyhow::bail!("expected tool callback, got {callback:?}"); + }; + assert_eq!(callback_session_id, session_id); + assert_eq!(invocation.input, Some(json!({ "value": "ping" }))); + + client + .send(&ClientToHost::DelegateResponse { + id: delegate_id, + result: WireResult::Ok { + value: DelegateResponse::ToolResult { + result: json!({ "value": "pong" }), + }, + }, + }) + .await?; + + assert_eq!( + client.read().await?, + HostToClient::InitialResponse { + id: execute_id, + result: WireResult::Ok { + value: WireRuntimeResponse::Result { + cell_id, + content_items: vec![WireContentItem::InputText { + text: "pong".to_string(), + }], + error_text: None, + }, + }, + } + ); + Ok(()) +} + +#[tokio::test] +async fn websocket_listener_accepts_frames_larger_than_default_websocket_limit() -> Result<()> { + let host = HostHarness::start().await?; + let mut client = host.connect().await?; + let capability = Capability::new("x".repeat((16 * 1024 * 1024) + 1))?; + + client + .negotiate(CapabilitySet::try_new([capability])?) + .await +} + +#[tokio::test] +async fn websocket_listener_keeps_connections_and_session_ids_isolated() -> Result<()> { + let host = HostHarness::start().await?; + let mut first = host.connect().await?; + let mut second = host.connect().await?; + first.negotiate(CapabilitySet::empty()).await?; + second.negotiate(CapabilitySet::empty()).await?; + + let session_id = SessionId::new("shared-session-name")?; + first.open_session(session_id.clone()).await?; + second.open_session(session_id).await?; + Ok(()) +} + +#[tokio::test] +async fn malformed_websocket_frame_does_not_stop_the_listener() -> Result<()> { + let mut host = HostHarness::start().await?; + let stderr = host + .child + .stderr + .take() + .context("code-mode host stderr was not captured")?; + let mut stderr_lines = BufReader::new(stderr).lines(); + let mut malformed = host.connect().await?; + malformed.send_binary(vec![1, 0, 0, 0, b'{']).await?; + + let close = timeout(TEST_TIMEOUT, malformed.websocket.next()) + .await + .context("timed out waiting for malformed websocket connection to close")?; + if let Some(Ok(message)) = close + && !matches!(message, Message::Close(_)) + { + anyhow::bail!("malformed websocket returned an unexpected message: {message:?}"); + } + + let diagnostic = timeout(TEST_TIMEOUT, async { + loop { + let line = stderr_lines + .next_line() + .await? + .context("code-mode host exited before reporting the malformed frame")?; + if line.contains("code-mode host websocket connection failed") { + return Ok::<_, anyhow::Error>(line); + } + } + }) + .await + .context("timed out waiting for the malformed websocket diagnostic")??; + assert!( + diagnostic.contains("failed to read code-mode client hello"), + "unexpected malformed websocket diagnostic: {diagnostic}", + ); + + let mut recovered = host.connect().await?; + recovered.negotiate(CapabilitySet::empty()).await +} + +#[tokio::test] +async fn websocket_listener_rejects_browser_origin_handshakes() -> Result<()> { + let host = HostHarness::start().await?; + let mut request = host.websocket_url.as_str().into_client_request()?; + request + .headers_mut() + .insert(ORIGIN, HeaderValue::from_static("https://evil.example")); + + let error = match connect_async(request).await { + Ok(_) => anyhow::bail!("browser-origin websocket handshake should be rejected"), + Err(error) => error, + }; + let WebSocketError::Http(response) = error else { + anyhow::bail!("browser-origin websocket handshake failed unexpectedly: {error}"); + }; + assert_eq!(response.status(), StatusCode::FORBIDDEN); + Ok(()) +} diff --git a/codex-rs/code-mode-protocol/src/host/codec.rs b/codex-rs/code-mode-protocol/src/host/codec.rs index ec07b92d41..10d13debe8 100644 --- a/codex-rs/code-mode-protocol/src/host/codec.rs +++ b/codex-rs/code-mode-protocol/src/host/codec.rs @@ -8,7 +8,7 @@ use tokio::io::AsyncReadExt; use tokio::io::AsyncWrite; use tokio::io::AsyncWriteExt; -/// Maximum JSON payload size accepted for one IPC frame. +/// Maximum JSON payload size accepted for one code-mode host frame. pub const MAX_FRAME_BYTES: usize = 64 * 1024 * 1024; /// A serialized IPC frame that has already passed the payload size limit. @@ -39,6 +39,55 @@ impl EncodedFrame { } Ok(Self { payload }) } + + /// Returns the complete length-prefixed representation of this frame. + pub fn into_framed_bytes(self) -> Vec { + let mut bytes = Vec::with_capacity(size_of::() + self.payload.len()); + bytes.extend_from_slice(&(self.payload.len() as u32).to_le_bytes()); + bytes.extend_from_slice(&self.payload); + bytes + } + + /// Decodes exactly one complete length-prefixed frame. + pub fn decode_framed(bytes: &[u8]) -> io::Result + where + T: DeserializeOwned, + { + let length_bytes: [u8; size_of::()] = bytes + .get(..size_of::()) + .and_then(|length_bytes| length_bytes.try_into().ok()) + .ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidData, + "code-mode IPC frame is missing its length prefix", + ) + })?; + 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 payload = &bytes[size_of::()..]; + if payload.len() != length { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + format!( + "code-mode IPC frame declares {length} payload bytes but contains {}", + payload.len() + ), + )); + } + + serde_json::from_slice(payload).map_err(|err| { + io::Error::new( + io::ErrorKind::InvalidData, + format!("failed to decode code-mode IPC frame: {err}"), + ) + }) + } } /// Decodes JSON messages prefixed by a four-byte little-endian payload length. diff --git a/codex-rs/code-mode-protocol/src/host/codec_tests.rs b/codex-rs/code-mode-protocol/src/host/codec_tests.rs index 996bd0679a..332a93766e 100644 --- a/codex-rs/code-mode-protocol/src/host/codec_tests.rs +++ b/codex-rs/code-mode-protocol/src/host/codec_tests.rs @@ -3,10 +3,43 @@ use serde_json::json; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; +use super::EncodedFrame; use super::FramedReader; use super::FramedWriter; use super::MAX_FRAME_BYTES; +#[test] +fn complete_frame_round_trips_without_a_byte_stream() { + let value = json!({"type": "session/open", "sessionId": "session-1"}); + let bytes = EncodedFrame::encode(&value) + .expect("encode frame") + .into_framed_bytes(); + + assert_eq!( + EncodedFrame::decode_framed::(&bytes).expect("decode frame"), + value + ); +} + +#[test] +fn complete_frame_rejects_truncated_and_trailing_payloads() { + let value = json!({"value": 1}); + let bytes = EncodedFrame::encode(&value) + .expect("encode frame") + .into_framed_bytes(); + + let truncated = &bytes[..bytes.len() - 1]; + let truncated_error = EncodedFrame::decode_framed::(truncated) + .expect_err("truncated frame should fail"); + assert_eq!(truncated_error.kind(), std::io::ErrorKind::InvalidData); + + let mut trailing = bytes; + trailing.push(0); + let trailing_error = EncodedFrame::decode_framed::(&trailing) + .expect_err("frame with trailing bytes should fail"); + assert_eq!(trailing_error.kind(), std::io::ErrorKind::InvalidData); +} + #[tokio::test] async fn frame_wire_format_is_little_endian_length_prefixed_json() { let (writer, mut reader) = tokio::io::duplex(/*max_buf_size*/ 128); diff --git a/codex-rs/code-mode-protocol/src/host/mod.rs b/codex-rs/code-mode-protocol/src/host/mod.rs index 3583ccb3e9..e448b670ea 100644 --- a/codex-rs/code-mode-protocol/src/host/mod.rs +++ b/codex-rs/code-mode-protocol/src/host/mod.rs @@ -1,4 +1,4 @@ -//! Messages and local IPC framing for the code-mode host boundary. +//! Messages and framing for the code-mode host boundary. //! //! Protocol version 1 multiplexes session operations and delegate callbacks by //! request ID over one ordered connection. It defines no optional capabilities