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

@@ -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]]

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 }

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<oneshot::Sender<()>>,
+    handshake_flushed: bool,
+}
+
+impl AsyncWrite for BlockingWriter {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        bytes: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        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<std::io::Result<()>> {
+        self.handshake_flushed = true;
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        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;

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<Semaphore>,
+    active_cell_permits: Arc<Semaphore>,
+}
+
+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);
+    run_connection(
-    let mut writer = FramedWriter::new(writer);
+        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<HostLimits>,
+) -> 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)),
+        request_permits: Arc::clone(&limits.request_permits),
-        active_cell_permits: Arc::new(Semaphore::new(MAX_ACTIVE_CELLS)),
+        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::<ClientToHost>() => 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<R, W>(reader: &mut FramedReader<R>, writer: &mut FramedWriter<W>) -> Result<bool>
+async fn negotiate(reader: &mut ConnectionReader, writer: &mut ConnectionWriter) -> Result<bool> {
-where
-    R: AsyncRead + Unpin,
-    W: AsyncWrite + Unpin,
-{
     let Some(first_message) = reader
-        .read::<ClientToHost>()
+        .read()
         .await
         .context("failed to read code-mode client hello")?
     else {

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
 }

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::<u32>();
+
+type BoxedReader = Box<dyn AsyncRead + Send + Unpin>;
+type BoxedWriter = Box<dyn AsyncWrite + Send + Unpin>;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+enum ListenTransport {
+    Stdio,
+    WebSocket(SocketAddr),
+}
+
+pub(crate) enum ConnectionReader {
+    Framed(FramedReader<BoxedReader>),
+    WebSocket(SplitStream<WebSocket>),
+}
+
+pub(crate) enum ConnectionWriter {
+    Framed(FramedWriter<BoxedWriter>),
+    WebSocket(SplitSink<WebSocket, Message>),
+}
+
+#[derive(Clone)]
+struct WebSocketListenerState {
+    limits: Arc<HostLimits>,
+}
+
+impl ConnectionReader {
+    pub(crate) fn from_reader<R>(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<Option<ClientToHost>> {
+        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<W>(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<ListenTransport> {
+    if matches!(listen_url, "stdio" | "stdio://") {
+        return Ok(ListenTransport::Stdio);
+    }
+
+    if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
+        return socket_addr
+            .parse::<SocketAddr>()
+            .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::<SocketAddr>(),
+    )
+    .await
+    .context("code-mode host websocket listener failed")
+}
+
+async fn readiness_handler() -> StatusCode {
+    StatusCode::OK
+}
+
+async fn reject_requests_with_origin_header(
+    request: Request<Body>,
+    next: Next,
+) -> Result<Response, StatusCode> {
+    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<SocketAddr>,
+    State(state): State<WebSocketListenerState>,
+) -> 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;

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::<SocketAddr>()
+                .expect("valid socket address")
+        )
+    );
+    assert_eq!(
+        parse_listen_url("ws://[::1]:9000").expect("IPv6 websocket listen URL should parse"),
+        ListenTransport::WebSocket(
+            "[::1]:9000"
+                .parse::<SocketAddr>()
+                .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"));
+}

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::<u32>();
+
+struct HostHarness {
+    child: Child,
+    websocket_url: String,
+}
+
+struct HostClient {
+    websocket: WebSocketStream<MaybeTlsStream<TcpStream>>,
+}
+
+impl HostHarness {
+    async fn start() -> Result<Self> {
+        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<HostClient> {
+        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<u8>) -> 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<HostToClient> {
+        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(())
+}

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<u8> {
+        let mut bytes = Vec::with_capacity(size_of::<u32>() + 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<T>(bytes: &[u8]) -> io::Result<T>
+    where
+        T: DeserializeOwned,
+    {
+        let length_bytes: [u8; size_of::<u32>()] = bytes
+            .get(..size_of::<u32>())
+            .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::<u32>()..];
+        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.

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::<serde_json::Value>(&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::<serde_json::Value>(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::<serde_json::Value>(&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);

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