[exec-server] serve websocket listener via HTTP upgrade (#21963)

## Why

`codex exec-server` should keep the existing public `ws://IP:PORT` URL
shape while serving that websocket connection through an HTTP upgrade
path internally. That keeps the client-facing configuration simple and
allows the listener to work through intermediate HTTP-aware
infrastructure.

## What changed

- keep the emitted and configured exec-server URL as `ws://IP:PORT`
- serve that websocket endpoint through Axum HTTP upgrade handling on
`/`
- expose `GET /readyz` from the same listener for readiness checks
- route upgraded Axum websocket streams through the shared JSON-RPC
connection machinery
- initialize the rustls crypto provider before websocket client
connections
- preserve inbound binary websocket JSON-RPC parsing for compatibility
with the prior transport behavior

## Verification

- `cargo test -p codex-exec-server --test health --test process --test
websocket --test initialize --test exec_process`

codex-rs/Cargo.lock

@@ -1145,6 +1145,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8"
 dependencies = [
  "axum-core",
+ "base64 0.22.1",
  "bytes",
  "form_urlencoded",
  "futures-util",
@@ -1163,8 +1164,10 @@ dependencies = [
  "serde_json",
  "serde_path_to_error",
  "serde_urlencoded",
+ "sha1",
  "sync_wrapper",
  "tokio",
+ "tokio-tungstenite",
  "tower",
  "tower-layer",
  "tower-service",
@@ -2722,6 +2725,7 @@ dependencies = [
  "anyhow",
  "arc-swap",
  "async-trait",
+ "axum",
  "base64 0.22.1",
  "bytes",
  "codex-app-server-protocol",
@@ -2732,6 +2736,7 @@ dependencies = [
  "codex-test-binary-support",
  "codex-utils-absolute-path",
  "codex-utils-pty",
+ "codex-utils-rustls-provider",
  "ctor 0.6.3",
  "futures",
  "pretty_assertions",

codex-rs/exec-server/Cargo.toml

@@ -13,6 +13,7 @@ workspace = true
 [dependencies]
 arc-swap = { workspace = true }
 async-trait = { workspace = true }
+axum = { workspace = true, features = ["http1", "tokio", "ws"] }
 base64 = { workspace = true }
 bytes = { workspace = true }
 codex-app-server-protocol = { workspace = true }
@@ -22,6 +23,7 @@ codex-protocol = { workspace = true }
 codex-sandboxing = { workspace = true }
 codex-utils-absolute-path = { workspace = true }
 codex-utils-pty = { workspace = true }
+codex-utils-rustls-provider = { workspace = true }
 futures = { workspace = true }
 reqwest = { workspace = true, features = ["json", "rustls-tls", "stream"] }
 serde = { workspace = true, features = ["derive"] }

codex-rs/exec-server/src/client_transport.rs

@@ -7,6 +7,8 @@ use tokio_tungstenite::connect_async;
 use tracing::debug;
 use tracing::warn;
 
+use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
+
 use crate::ExecServerClient;
 use crate::ExecServerError;
 use crate::client_api::RemoteExecServerConnectArgs;
@@ -53,6 +55,7 @@ impl ExecServerClient {
     pub async fn connect_websocket(
         args: RemoteExecServerConnectArgs,
     ) -> Result<Self, ExecServerError> {
+        ensure_rustls_crypto_provider();
         let websocket_url = args.websocket_url.clone();
         let connect_timeout = args.connect_timeout;
         let (stream, _) = timeout(connect_timeout, connect_async(websocket_url.as_str()))

codex-rs/exec-server/src/connection.rs

@@ -5,8 +5,12 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::time::Duration;
 
+use axum::extract::ws::Message as AxumWebSocketMessage;
+use axum::extract::ws::WebSocket as AxumWebSocket;
 use codex_app_server_protocol::JSONRPCMessage;
+use futures::Sink;
 use futures::SinkExt;
+use futures::Stream;
 use futures::StreamExt;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
@@ -314,11 +318,30 @@ impl JsonRpcConnection {
     pub(crate) fn from_websocket<S>(stream: WebSocketStream<S>, connection_label: String) -> Self
     where
         S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
+    {
+        let (websocket_writer, websocket_reader) = stream.split();
+        Self::from_websocket_parts(websocket_writer, websocket_reader, connection_label)
+    }
+
+    pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self {
+        let (websocket_writer, websocket_reader) = stream.split();
+        Self::from_websocket_parts(websocket_writer, websocket_reader, connection_label)
+    }
+
+    fn from_websocket_parts<W, R, M, E>(
+        mut websocket_writer: W,
+        mut websocket_reader: R,
+        connection_label: String,
+    ) -> Self
+    where
+        W: Sink<M, Error = E> + Unpin + Send + 'static,
+        R: Stream<Item = Result<M, E>> + Unpin + Send + 'static,
+        M: JsonRpcWebSocketMessage,
+        E: std::fmt::Display + Send + 'static,
     {
         let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY);
         let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY);
         let (disconnected_tx, disconnected_rx) = watch::channel(false);
-        let (mut websocket_writer, mut websocket_reader) = stream.split();
 
         let reader_label = connection_label.clone();
         let incoming_tx_for_reader = incoming_tx.clone();
@@ -326,61 +349,36 @@ impl JsonRpcConnection {
         let reader_task = tokio::spawn(async move {
             loop {
                 match websocket_reader.next().await {
-                    Some(Ok(Message::Text(text))) => {
+                    Some(Ok(message)) => match message.parse_jsonrpc_frame() {
-                        match serde_json::from_str::<JSONRPCMessage>(text.as_ref()) {
+                        Ok(JsonRpcWebSocketFrame::Message(message)) => {
-                            Ok(message) => {
+                            if incoming_tx_for_reader
-                                if incoming_tx_for_reader
+                                .send(JsonRpcConnectionEvent::Message(message))
-                                    .send(JsonRpcConnectionEvent::Message(message))
+                                .await
-                                    .await
+                                .is_err()
-                                    .is_err()
+                            {
-                                {
+                                break;
-                                    break;
-                                }
-                            }
-                            Err(err) => {
-                                send_malformed_message(
-                                    &incoming_tx_for_reader,
-                                    Some(format!(
-                                        "failed to parse websocket JSON-RPC message from {reader_label}: {err}"
-                                    )),
-                                )
-                                .await;
                             }
                         }
-                    }
+                        Err(err) => {
-                    Some(Ok(Message::Binary(bytes))) => {
+                            send_malformed_message(
-                        match serde_json::from_slice::<JSONRPCMessage>(bytes.as_ref()) {
+                                &incoming_tx_for_reader,
-                            Ok(message) => {
+                                Some(format!(
-                                if incoming_tx_for_reader
+                                    "failed to parse websocket JSON-RPC message from {reader_label}: {err}"
-                                    .send(JsonRpcConnectionEvent::Message(message))
+                                )),
-                                    .await
+                            )
-                                    .is_err()
+                            .await;
-                                {
-                                    break;
-                                }
-                            }
-                            Err(err) => {
-                                send_malformed_message(
-                                    &incoming_tx_for_reader,
-                                    Some(format!(
-                                        "failed to parse websocket JSON-RPC message from {reader_label}: {err}"
-                                    )),
-                                )
-                                .await;
-                            }
                         }
-                    }
+                        Ok(JsonRpcWebSocketFrame::Close) => {
-                    Some(Ok(Message::Close(_))) => {
+                            send_disconnected(
-                        send_disconnected(
+                                &incoming_tx_for_reader,
-                            &incoming_tx_for_reader,
+                                &disconnected_tx_for_reader,
-                            &disconnected_tx_for_reader,
+                                /*reason*/ None,
-                            /*reason*/ None,
+                            )
-                        )
+                            .await;
-                        .await;
+                            break;
-                        break;
+                        }
-                    }
+                        Ok(JsonRpcWebSocketFrame::Ignore) => {}
-                    Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {}
+                    },
-                    Some(Ok(_)) => {}
                     Some(Err(err)) => {
                         send_disconnected(
                             &incoming_tx_for_reader,
@@ -409,8 +407,7 @@ impl JsonRpcConnection {
             while let Some(message) = outgoing_rx.recv().await {
                 match serialize_jsonrpc_message(&message) {
                     Ok(encoded) => {
-                        if let Err(err) = websocket_writer.send(Message::Text(encoded.into())).await
+                        if let Err(err) = websocket_writer.send(M::from_text(encoded)).await {
-                        {
                             send_disconnected(
                                 &incoming_tx,
                                 &disconnected_tx,
@@ -452,6 +449,59 @@ impl JsonRpcConnection {
     }
 }
 
+enum JsonRpcWebSocketFrame {
+    Message(JSONRPCMessage),
+    Close,
+    Ignore,
+}
+
+trait JsonRpcWebSocketMessage: Send + 'static {
+    fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error>;
+    fn from_text(text: String) -> Self;
+}
+
+impl JsonRpcWebSocketMessage for Message {
+    fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
+        match self {
+            Message::Text(text) => {
+                serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            Message::Binary(bytes) => {
+                serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            Message::Close(_) => Ok(JsonRpcWebSocketFrame::Close),
+            Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => {
+                Ok(JsonRpcWebSocketFrame::Ignore)
+            }
+        }
+    }
+
+    fn from_text(text: String) -> Self {
+        Self::Text(text.into())
+    }
+}
+
+impl JsonRpcWebSocketMessage for AxumWebSocketMessage {
+    fn parse_jsonrpc_frame(self) -> Result<JsonRpcWebSocketFrame, serde_json::Error> {
+        match self {
+            AxumWebSocketMessage::Text(text) => {
+                serde_json::from_str(text.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            AxumWebSocketMessage::Binary(bytes) => {
+                serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message)
+            }
+            AxumWebSocketMessage::Close(_) => Ok(JsonRpcWebSocketFrame::Close),
+            AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => {
+                Ok(JsonRpcWebSocketFrame::Ignore)
+            }
+        }
+    }
+
+    fn from_text(text: String) -> Self {
+        Self::Text(text.into())
+    }
+}
+
 async fn send_disconnected(
     incoming_tx: &mpsc::Sender<JsonRpcConnectionEvent>,
     disconnected_tx: &watch::Sender<bool>,

codex-rs/exec-server/src/remote.rs

@@ -7,6 +7,8 @@ use tokio::time::sleep;
 use tokio_tungstenite::connect_async;
 use tracing::warn;
 
+use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
+
 use crate::ExecServerError;
 use crate::ExecServerRuntimePaths;
 use crate::connection::JsonRpcConnection;
@@ -133,6 +135,7 @@ pub async fn run_remote_executor(
     config: RemoteExecutorConfig,
     runtime_paths: ExecServerRuntimePaths,
 ) -> Result<(), ExecServerError> {
+    ensure_rustls_crypto_provider();
     let client = ExecutorRegistryClient::new(config.base_url.clone(), config.bearer_token.clone())?;
     let processor = ConnectionProcessor::new(runtime_paths);
     let mut backoff = Duration::from_secs(1);

codex-rs/exec-server/src/server/transport.rs

@@ -1,11 +1,18 @@
+use axum::Router;
+use axum::extract::ConnectInfo;
+use axum::extract::State;
+use axum::extract::ws::WebSocketUpgrade;
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use axum::routing::any;
+use axum::routing::get;
 use std::io::Write as _;
 use std::net::SocketAddr;
 use tokio::io;
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
 use tokio::net::TcpListener;
-use tokio_tungstenite::accept_async;
+use tracing::info;
-use tracing::warn;
 
 use crate::ExecServerRuntimePaths;
 use crate::connection::JsonRpcConnection;
@@ -109,31 +116,46 @@ async fn run_websocket_listener(
     let listener = TcpListener::bind(bind_address).await?;
     let local_addr = listener.local_addr()?;
     let processor = ConnectionProcessor::new(runtime_paths);
-    tracing::info!("codex-exec-server listening on ws://{local_addr}");
+    info!("codex-exec-server listening on ws://{local_addr}");
     println!("ws://{local_addr}");
     std::io::stdout().flush()?;
 
-    loop {
+    let router = Router::new()
-        let (stream, peer_addr) = listener.accept().await?;
+        .route("/", any(websocket_upgrade_handler))
-        let processor = processor.clone();
+        .route("/readyz", get(readiness_handler))
-        tokio::spawn(async move {
+        .with_state(ExecServerWebSocketState { processor });
-            match accept_async(stream).await {
+    axum::serve(
-                Ok(websocket) => {
+        listener,
-                    processor
+        router.into_make_service_with_connect_info::<SocketAddr>(),
-                        .run_connection(JsonRpcConnection::from_websocket(
+    )
-                            websocket,
+    .await?;
-                            format!("exec-server websocket {peer_addr}"),
+    Ok(())
-                        ))
+}
-                        .await;
+
-                }
+#[derive(Clone)]
-                Err(err) => {
+struct ExecServerWebSocketState {
-                    warn!(
+    processor: ConnectionProcessor,
-                        "failed to accept exec-server websocket connection from {peer_addr}: {err}"
+}
-                    );
+
-                }
+async fn readiness_handler() -> StatusCode {
-            }
+    StatusCode::OK
-        });
+}
-    }
+
+async fn websocket_upgrade_handler(
+    websocket: WebSocketUpgrade,
+    ConnectInfo(peer_addr): ConnectInfo<SocketAddr>,
+    State(state): State<ExecServerWebSocketState>,
+) -> impl IntoResponse {
+    info!(%peer_addr, "exec-server websocket client connected");
+    websocket.on_upgrade(move |stream| async move {
+        state
+            .processor
+            .run_connection(JsonRpcConnection::from_axum_websocket(
+                stream,
+                format!("exec-server websocket {peer_addr}"),
+            ))
+            .await;
+    })
 }
 
 #[cfg(test)]

codex-rs/exec-server/tests/common/exec_server.rs

@@ -142,6 +142,11 @@ impl ExecServerHarness {
         Ok(())
     }
 
+    pub(crate) async fn send_raw_binary(&mut self, bytes: Vec<u8>) -> anyhow::Result<()> {
+        self.websocket.send(Message::Binary(bytes.into())).await?;
+        Ok(())
+    }
+
     pub(crate) async fn next_event(&mut self) -> anyhow::Result<JSONRPCMessage> {
         self.next_event_with_timeout(EVENT_TIMEOUT).await
     }

codex-rs/exec-server/tests/health.rs

@@ -0,0 +1,21 @@
+#![cfg(unix)]
+
+mod common;
+
+use common::exec_server::exec_server;
+use pretty_assertions::assert_eq;
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn exec_server_serves_readyz_alongside_websocket_endpoint() -> anyhow::Result<()> {
+    let mut server = exec_server().await?;
+    let http_base_url = server
+        .websocket_url()
+        .strip_prefix("ws://")
+        .expect("websocket URL should use ws://");
+
+    let response = reqwest::get(format!("http://{http_base_url}/readyz")).await?;
+    assert_eq!(response.status(), reqwest::StatusCode::OK);
+
+    server.shutdown().await?;
+    Ok(())
+}

codex-rs/exec-server/tests/websocket.rs

@@ -60,3 +60,39 @@ async fn exec_server_reports_malformed_websocket_json_and_keeps_running() -> any
     server.shutdown().await?;
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn exec_server_accepts_binary_websocket_json() -> anyhow::Result<()> {
+    let mut server = exec_server().await?;
+    let initialize_id = codex_app_server_protocol::RequestId::Integer(1);
+    let initialize = JSONRPCMessage::Request(codex_app_server_protocol::JSONRPCRequest {
+        id: initialize_id.clone(),
+        method: "initialize".to_string(),
+        params: Some(serde_json::to_value(InitializeParams {
+            client_name: "exec-server-binary-test".to_string(),
+            resume_session_id: None,
+        })?),
+        trace: None,
+    });
+    server
+        .send_raw_binary(serde_json::to_vec(&initialize)?)
+        .await?;
+
+    let response = server
+        .wait_for_event(|event| {
+            matches!(
+                event,
+                JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
+            )
+        })
+        .await?;
+    let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
+        panic!("expected initialize response for binary input");
+    };
+    assert_eq!(id, initialize_id);
+    let initialize_response: InitializeResponse = serde_json::from_value(result)?;
+    Uuid::parse_str(&initialize_response.session_id)?;
+
+    server.shutdown().await?;
+    Ok(())
+}