From 64ead6a83a6ed348229bc98a9b5d8b0c550d8305 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Sun, 17 May 2026 20:07:32 -0700 Subject: [PATCH] Add exec-server websocket keepalive (#23226) ## Summary - send periodic websocket Ping frames from outbound exec-server websocket clients - cover direct exec-server websocket clients plus rendezvous harness/executor websocket connections - keep inbound axum-accepted exec-server websocket connections passive - add focused keepalive coverage for direct and relay websocket paths ## Validation - /Users/starr/code/openai/project/dotslash-gen/bin/bazel test //codex-rs/exec-server:exec-server-unit-tests --test_filter='websocket_connection_sends_keepalive_ping|harness_connection_sends_keepalive_ping|multiplexed_executor_sends_keepalive_ping' - /Users/starr/code/openai/project/dotslash-gen/bin/bazel test //codex-rs/exec-server:exec-server-relay-test --test_filter=multiplexed_remote_executor_routes_independent_virtual_streams --- codex-rs/exec-server/src/connection.rs | 191 ++++++++++++++++++++++--- codex-rs/exec-server/src/relay.rs | 159 +++++++++++++++++--- 2 files changed, 304 insertions(+), 46 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index d5ff1f966..cf504bbea 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -30,6 +30,10 @@ use tokio::io::BufWriter; pub(crate) const CHANNEL_CAPACITY: usize = 128; const STDIO_TERMINATION_GRACE_PERIOD: Duration = Duration::from_secs(2); +#[cfg(test)] +pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_millis(25); +#[cfg(not(test))] +pub(crate) const WEBSOCKET_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(30); #[derive(Debug)] pub(crate) enum JsonRpcConnectionEvent { @@ -320,18 +324,32 @@ impl JsonRpcConnection { S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { let (websocket_writer, websocket_reader) = stream.split(); - Self::from_websocket_parts(websocket_writer, websocket_reader, connection_label) + Self::from_websocket_parts( + websocket_writer, + websocket_reader, + connection_label, + Some(WEBSOCKET_KEEPALIVE_INTERVAL), + ) } 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) + Self::from_websocket_parts( + websocket_writer, + websocket_reader, + connection_label, + // Axum only wraps inbound exec-server websocket accepts. Outbound websocket clients + // own keepalive pings so one side does not accidentally create redundant traffic. + /*keepalive_interval*/ + None, + ) } fn from_websocket_parts( mut websocket_writer: W, mut websocket_reader: R, connection_label: String, + keepalive_interval: Option, ) -> Self where W: Sink + Unpin + Send + 'static, @@ -404,30 +422,54 @@ impl JsonRpcConnection { }); let writer_task = tokio::spawn(async move { - while let Some(message) = outgoing_rx.recv().await { - match serialize_jsonrpc_message(&message) { - Ok(encoded) => { - if let Err(err) = websocket_writer.send(M::from_text(encoded)).await { - send_disconnected( - &incoming_tx, - &disconnected_tx, - Some(format!( - "failed to write websocket JSON-RPC message to {connection_label}: {err}" - )), + if let Some(keepalive_interval) = keepalive_interval { + let mut keepalive = tokio::time::interval_at( + tokio::time::Instant::now() + keepalive_interval, + keepalive_interval, + ); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + maybe_message = outgoing_rx.recv() => { + let Some(message) = maybe_message else { + break; + }; + if let Err(reason) = send_websocket_jsonrpc_message( + &mut websocket_writer, + &connection_label, + &message, ) - .await; - break; + .await + { + send_disconnected(&incoming_tx, &disconnected_tx, Some(reason)).await; + break; + } + } + _ = keepalive.tick() => { + if let Err(err) = websocket_writer.send(M::ping()).await { + send_disconnected( + &incoming_tx, + &disconnected_tx, + Some(format!( + "failed to write websocket ping to {connection_label}: {err}" + )), + ) + .await; + break; + } } } - Err(err) => { - send_disconnected( - &incoming_tx, - &disconnected_tx, - Some(format!( - "failed to serialize JSON-RPC message for {connection_label}: {err}" - )), - ) - .await; + } + } else { + while let Some(message) = outgoing_rx.recv().await { + if let Err(reason) = send_websocket_jsonrpc_message( + &mut websocket_writer, + &connection_label, + &message, + ) + .await + { + send_disconnected(&incoming_tx, &disconnected_tx, Some(reason)).await; break; } } @@ -458,6 +500,7 @@ enum JsonRpcWebSocketFrame { trait JsonRpcWebSocketMessage: Send + 'static { fn parse_jsonrpc_frame(self) -> Result; fn from_text(text: String) -> Self; + fn ping() -> Self; } impl JsonRpcWebSocketMessage for Message { @@ -479,6 +522,10 @@ impl JsonRpcWebSocketMessage for Message { fn from_text(text: String) -> Self { Self::Text(text.into()) } + + fn ping() -> Self { + Self::Ping(Vec::new().into()) + } } impl JsonRpcWebSocketMessage for AxumWebSocketMessage { @@ -500,6 +547,10 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { fn from_text(text: String) -> Self { Self::Text(text.into()) } + + fn ping() -> Self { + Self::Ping(Vec::new().into()) + } } async fn send_disconnected( @@ -538,6 +589,100 @@ where writer.flush().await } +async fn send_websocket_jsonrpc_message( + websocket_writer: &mut W, + connection_label: &str, + message: &JSONRPCMessage, +) -> Result<(), String> +where + W: Sink + Unpin, + M: JsonRpcWebSocketMessage, + E: std::fmt::Display, +{ + match serialize_jsonrpc_message(message) { + Ok(encoded) => websocket_writer + .send(M::from_text(encoded)) + .await + .map_err(|err| { + format!("failed to write websocket JSON-RPC message to {connection_label}: {err}") + }), + Err(err) => Err(format!( + "failed to serialize JSON-RPC message for {connection_label}: {err}" + )), + } +} + fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result { serde_json::to_string(message) } + +#[cfg(test)] +mod tests { + use std::pin::Pin; + + use futures::channel::mpsc as futures_mpsc; + use futures::stream; + use futures::task::Context; + use futures::task::Poll; + use tokio::time::timeout; + + use super::*; + + struct TestWebSocketSink { + message_tx: futures_mpsc::UnboundedSender, + } + + impl Sink for TestWebSocketSink { + type Error = std::convert::Infallible; + + fn poll_ready( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + + fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + self.get_mut() + .message_tx + .unbounded_send(item) + .expect("test websocket receiver should stay open"); + Ok(()) + } + + fn poll_flush( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_close( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + } + + #[tokio::test] + async fn websocket_connection_sends_keepalive_ping() { + let (message_tx, mut message_rx) = futures_mpsc::unbounded::(); + let websocket_writer = TestWebSocketSink { message_tx }; + let websocket_reader = stream::pending::>(); + let connection = JsonRpcConnection::from_websocket_parts( + websocket_writer, + websocket_reader, + "test".into(), + Some(WEBSOCKET_KEEPALIVE_INTERVAL), + ); + + let message = timeout(Duration::from_secs(1), message_rx.next()) + .await + .expect("keepalive ping should arrive before timeout") + .expect("keepalive ping should be sent"); + assert!(matches!(message, Message::Ping(_))); + + drop(connection); + } +} diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index bce787cfc..7470a6290 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -19,6 +19,7 @@ use crate::connection::CHANNEL_CAPACITY; use crate::connection::JsonRpcConnection; use crate::connection::JsonRpcConnectionEvent; use crate::connection::JsonRpcTransport; +use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; use crate::relay_proto::RelayResume; @@ -262,24 +263,42 @@ where return; } + let mut keepalive = tokio::time::interval_at( + tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, + WEBSOCKET_KEEPALIVE_INTERVAL, + ); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); let mut next_seq = 0u32; - while let Some(message) = outgoing_rx.recv().await { - let payload = match jsonrpc_payload(&message) { - Ok(payload) => payload, - Err(err) => { - warn!("failed to serialize JSON-RPC payload for relay transport: {err}"); - break; + loop { + tokio::select! { + maybe_message = outgoing_rx.recv() => { + let Some(message) = maybe_message else { + break; + }; + let payload = match jsonrpc_payload(&message) { + Ok(payload) => payload, + Err(err) => { + warn!("failed to serialize JSON-RPC payload for relay transport: {err}"); + break; + } + }; + let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload); + next_seq = next_seq.wrapping_add(1); + if websocket_writer + .send(Message::Binary(encode_relay_message_frame(&frame).into())) + .await + .is_err() + { + let _ = disconnected_tx.send(true); + break; + } + } + _ = keepalive.tick() => { + if websocket_writer.send(Message::Ping(Vec::new().into())).await.is_err() { + let _ = disconnected_tx.send(true); + break; + } } - }; - let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload); - next_seq = next_seq.wrapping_add(1); - if websocket_writer - .send(Message::Binary(encode_relay_message_frame(&frame).into())) - .await - .is_err() - { - let _ = disconnected_tx.send(true); - break; } } }); @@ -303,13 +322,30 @@ pub(crate) async fn run_multiplexed_executor( let (physical_outgoing_tx, mut physical_outgoing_rx) = mpsc::channel::>(CHANNEL_CAPACITY); let writer_task = tokio::spawn(async move { - while let Some(encoded) = physical_outgoing_rx.recv().await { - if websocket_writer - .send(Message::Binary(encoded.into())) - .await - .is_err() - { - break; + let mut keepalive = tokio::time::interval_at( + tokio::time::Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL, + WEBSOCKET_KEEPALIVE_INTERVAL, + ); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + maybe_encoded = physical_outgoing_rx.recv() => { + let Some(encoded) = maybe_encoded else { + break; + }; + if websocket_writer + .send(Message::Binary(encoded.into())) + .await + .is_err() + { + break; + } + } + _ = keepalive.tick() => { + if websocket_writer.send(Message::Ping(Vec::new().into())).await.is_err() { + break; + } + } } } }); @@ -453,3 +489,80 @@ fn spawn_virtual_stream( disconnected_tx, } } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use tokio::net::TcpListener; + use tokio::time::timeout; + use tokio_tungstenite::accept_async; + use tokio_tungstenite::connect_async; + use tokio_tungstenite::tungstenite::Message; + + use super::*; + + fn test_runtime_paths() -> anyhow::Result { + crate::ExecServerRuntimePaths::new( + std::env::current_exe()?, + /*codex_linux_sandbox_exe*/ None, + ) + .map_err(anyhow::Error::from) + } + + #[tokio::test] + async fn multiplexed_executor_sends_keepalive_ping() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let executor_task = tokio::spawn(run_multiplexed_executor( + client_websocket, + ConnectionProcessor::new(test_runtime_paths()?), + )); + + read_keepalive_ping(&mut server_websocket).await?; + + executor_task.abort(); + let _ = executor_task.await; + Ok(()) + } + + #[tokio::test] + async fn harness_connection_sends_keepalive_ping() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let connection = harness_connection_from_websocket(client_websocket, "test".to_string()); + + read_keepalive_ping(&mut server_websocket).await?; + + drop(connection); + Ok(()) + } + + async fn websocket_pair() -> anyhow::Result<( + WebSocketStream>, + WebSocketStream, + )> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let server_task = tokio::spawn(async move { + let (stream, _) = listener.accept().await?; + accept_async(stream).await.map_err(anyhow::Error::from) + }); + let (client_websocket, _) = connect_async(websocket_url).await?; + let server_websocket = server_task.await??; + Ok((client_websocket, server_websocket)) + } + + async fn read_keepalive_ping( + websocket: &mut WebSocketStream, + ) -> anyhow::Result<()> { + loop { + let Some(message) = timeout(Duration::from_secs(1), websocket.next()).await? else { + anyhow::bail!("websocket closed before keepalive ping"); + }; + match message? { + Message::Ping(_) => return Ok(()), + Message::Binary(_) | Message::Text(_) | Message::Pong(_) | Message::Frame(_) => {} + Message::Close(_) => anyhow::bail!("websocket closed before keepalive ping"), + } + } + } +}