From 95d8669ab2e55d7c68cf469c152ebce37b395017 Mon Sep 17 00:00:00 2001 From: Ruslan Nigmatullin Date: Mon, 11 May 2026 17:04:21 -0700 Subject: [PATCH] [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 | 5 + codex-rs/exec-server/Cargo.toml | 2 + codex-rs/exec-server/src/client_transport.rs | 3 + codex-rs/exec-server/src/connection.rs | 160 ++++++++++++------ codex-rs/exec-server/src/remote.rs | 3 + codex-rs/exec-server/src/server/transport.rs | 70 +++++--- .../exec-server/tests/common/exec_server.rs | 5 + codex-rs/exec-server/tests/health.rs | 21 +++ codex-rs/exec-server/tests/websocket.rs | 36 ++++ 9 files changed, 226 insertions(+), 79 deletions(-) create mode 100644 codex-rs/exec-server/tests/health.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d15a1062d..bbb61f833 100644 --- a/codex-rs/Cargo.lock +++ b/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", diff --git a/codex-rs/exec-server/Cargo.toml b/codex-rs/exec-server/Cargo.toml index 9fbdd9111..936fa412f 100644 --- a/codex-rs/exec-server/Cargo.toml +++ b/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"] } diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 23dc0bc7b..8ca1eb028 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/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 { + 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())) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index bb2599ea4..d5ff1f966 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/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(stream: WebSocketStream, 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( + mut websocket_writer: W, + mut websocket_reader: R, + connection_label: String, + ) -> Self + where + W: Sink + Unpin + Send + 'static, + R: Stream> + 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))) => { - match serde_json::from_str::(text.as_ref()) { - Ok(message) => { - if incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Message(message)) - .await - .is_err() - { - break; - } - } - Err(err) => { - send_malformed_message( - &incoming_tx_for_reader, - Some(format!( - "failed to parse websocket JSON-RPC message from {reader_label}: {err}" - )), - ) - .await; + Some(Ok(message)) => match message.parse_jsonrpc_frame() { + Ok(JsonRpcWebSocketFrame::Message(message)) => { + if incoming_tx_for_reader + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; } } - } - Some(Ok(Message::Binary(bytes))) => { - match serde_json::from_slice::(bytes.as_ref()) { - Ok(message) => { - if incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Message(message)) - .await - .is_err() - { - 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) => { + send_malformed_message( + &incoming_tx_for_reader, + Some(format!( + "failed to parse websocket JSON-RPC message from {reader_label}: {err}" + )), + ) + .await; } - } - Some(Ok(Message::Close(_))) => { - send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - /*reason*/ None, - ) - .await; - break; - } - Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {} - Some(Ok(_)) => {} + Ok(JsonRpcWebSocketFrame::Close) => { + send_disconnected( + &incoming_tx_for_reader, + &disconnected_tx_for_reader, + /*reason*/ None, + ) + .await; + break; + } + Ok(JsonRpcWebSocketFrame::Ignore) => {} + }, 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; + fn from_text(text: String) -> Self; +} + +impl JsonRpcWebSocketMessage for Message { + fn parse_jsonrpc_frame(self) -> Result { + 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 { + 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, disconnected_tx: &watch::Sender, diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 43c424a14..bb22105c1 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/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); diff --git a/codex-rs/exec-server/src/server/transport.rs b/codex-rs/exec-server/src/server/transport.rs index d284bf64b..92b2c77aa 100644 --- a/codex-rs/exec-server/src/server/transport.rs +++ b/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::warn; +use tracing::info; 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 (stream, peer_addr) = listener.accept().await?; - let processor = processor.clone(); - tokio::spawn(async move { - match accept_async(stream).await { - Ok(websocket) => { - processor - .run_connection(JsonRpcConnection::from_websocket( - websocket, - format!("exec-server websocket {peer_addr}"), - )) - .await; - } - Err(err) => { - warn!( - "failed to accept exec-server websocket connection from {peer_addr}: {err}" - ); - } - } - }); - } + let router = Router::new() + .route("/", any(websocket_upgrade_handler)) + .route("/readyz", get(readiness_handler)) + .with_state(ExecServerWebSocketState { processor }); + axum::serve( + listener, + router.into_make_service_with_connect_info::(), + ) + .await?; + Ok(()) +} + +#[derive(Clone)] +struct ExecServerWebSocketState { + processor: ConnectionProcessor, +} + +async fn readiness_handler() -> StatusCode { + StatusCode::OK +} + +async fn websocket_upgrade_handler( + websocket: WebSocketUpgrade, + ConnectInfo(peer_addr): ConnectInfo, + State(state): State, +) -> 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)] diff --git a/codex-rs/exec-server/tests/common/exec_server.rs b/codex-rs/exec-server/tests/common/exec_server.rs index 4ff740871..f1bf03d25 100644 --- a/codex-rs/exec-server/tests/common/exec_server.rs +++ b/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) -> anyhow::Result<()> { + self.websocket.send(Message::Binary(bytes.into())).await?; + Ok(()) + } + pub(crate) async fn next_event(&mut self) -> anyhow::Result { self.next_event_with_timeout(EVENT_TIMEOUT).await } diff --git a/codex-rs/exec-server/tests/health.rs b/codex-rs/exec-server/tests/health.rs new file mode 100644 index 000000000..91b3806a2 --- /dev/null +++ b/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(()) +} diff --git a/codex-rs/exec-server/tests/websocket.rs b/codex-rs/exec-server/tests/websocket.rs index 64c9438b8..03b2cabfc 100644 --- a/codex-rs/exec-server/tests/websocket.rs +++ b/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(()) +}