From 8a9300e92ad191d0268a4c946bf97e2a5f2cea5f Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 10:04:34 -0700 Subject: [PATCH 1/8] Refactor exec-server websocket pump --- codex-rs/exec-server/src/connection.rs | 329 ++++++++++++------------- codex-rs/exec-server/src/relay.rs | 325 ++++++++++++------------ 2 files changed, 311 insertions(+), 343 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index cf504bbea7b7..0974a21c3a24 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -323,37 +323,20 @@ impl JsonRpcConnection { where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - let (websocket_writer, websocket_reader) = stream.split(); - Self::from_websocket_parts( - websocket_writer, - websocket_reader, - connection_label, - Some(WEBSOCKET_KEEPALIVE_INTERVAL), - ) + Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) } 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, - // 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, + Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL)) + } + + fn from_websocket_stream( + mut websocket: T, connection_label: String, - keepalive_interval: Option, + ping_interval: Option, ) -> Self where - W: Sink + Unpin + Send + 'static, - R: Stream> + Unpin + Send + 'static, + T: Sink + Stream> + Unpin + Send + 'static, M: JsonRpcWebSocketMessage, E: std::fmt::Display + Send + 'static, { @@ -361,118 +344,119 @@ impl JsonRpcConnection { let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); let (disconnected_tx, disconnected_rx) = watch::channel(false); - let reader_label = connection_label.clone(); - let incoming_tx_for_reader = incoming_tx.clone(); - let disconnected_tx_for_reader = disconnected_tx.clone(); - let reader_task = tokio::spawn(async move { + let websocket_task = tokio::spawn(async move { + let mut ping_interval = ping_interval.map(|ping_interval| { + let mut interval = tokio::time::interval_at( + tokio::time::Instant::now() + ping_interval, + ping_interval, + ); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + interval + }); + loop { - match websocket_reader.next().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; - } + tokio::select! { + maybe_message = outgoing_rx.recv() => { + let Some(message) = maybe_message else { + break; + }; + if let Err(reason) = send_websocket_jsonrpc_message( + &mut websocket, + &connection_label, + &message, + ) + .await + { + send_disconnected(&incoming_tx, &disconnected_tx, Some(reason)).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; + } + _ = async { + match ping_interval.as_mut() { + Some(interval) => interval.tick().await, + None => std::future::pending().await, } - Ok(JsonRpcWebSocketFrame::Close) => { + } => { + if let Err(err) = websocket.send(M::ping()).await { send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - /*reason*/ None, + &incoming_tx, + &disconnected_tx, + Some(format!( + "failed to write websocket ping to {connection_label}: {err}" + )), ) .await; break; } - Ok(JsonRpcWebSocketFrame::Ignore) => {} - }, - Some(Err(err)) => { - send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - Some(format!( - "failed to read websocket JSON-RPC message from {reader_label}: {err}" - )), - ) - .await; - break; } - None => { - send_disconnected( - &incoming_tx_for_reader, - &disconnected_tx_for_reader, - /*reason*/ None, - ) - .await; - break; - } - } - } - }); - - let writer_task = tokio::spawn(async move { - 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 - { - send_disconnected(&incoming_tx, &disconnected_tx, Some(reason)).await; - break; - } - } - _ = keepalive.tick() => { - if let Err(err) = websocket_writer.send(M::ping()).await { + incoming_message = websocket.next() => { + match incoming_message { + Some(Ok(message)) => match message.parse_jsonrpc_frame() { + Ok(JsonRpcWebSocketFrame::Message(message)) => { + if incoming_tx + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; + } + } + Ok(JsonRpcWebSocketFrame::Ping(payload)) => { + if let Err(err) = websocket.send(M::pong(payload)).await { + send_disconnected( + &incoming_tx, + &disconnected_tx, + Some(format!( + "failed to write websocket pong to {connection_label}: {err}" + )), + ) + .await; + break; + } + } + Ok(JsonRpcWebSocketFrame::Close) => { + send_disconnected( + &incoming_tx, + &disconnected_tx, + /*reason*/ None, + ) + .await; + break; + } + Ok(JsonRpcWebSocketFrame::Ignore) => {} + Err(err) => { + send_malformed_message( + &incoming_tx, + Some(format!( + "failed to parse websocket JSON-RPC message from {connection_label}: {err}" + )), + ) + .await; + } + }, + Some(Err(err)) => { send_disconnected( &incoming_tx, &disconnected_tx, Some(format!( - "failed to write websocket ping to {connection_label}: {err}" + "failed to read websocket JSON-RPC message from {connection_label}: {err}" )), ) .await; break; } + None => { + send_disconnected( + &incoming_tx, + &disconnected_tx, + /*reason*/ None, + ) + .await; + break; + } } } } - } 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; - } - } } }); @@ -480,7 +464,7 @@ impl JsonRpcConnection { outgoing_tx, incoming_rx, disconnected_rx, - task_handles: vec![reader_task, writer_task], + task_handles: vec![websocket_task], transport: JsonRpcTransport::Plain, } } @@ -493,6 +477,7 @@ impl JsonRpcConnection { enum JsonRpcWebSocketFrame { Message(JSONRPCMessage), + Ping(bytes::Bytes), Close, Ignore, } @@ -501,6 +486,7 @@ trait JsonRpcWebSocketMessage: Send + 'static { fn parse_jsonrpc_frame(self) -> Result; fn from_text(text: String) -> Self; fn ping() -> Self; + fn pong(payload: bytes::Bytes) -> Self; } impl JsonRpcWebSocketMessage for Message { @@ -513,9 +499,8 @@ impl JsonRpcWebSocketMessage for Message { serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message) } Message::Close(_) => Ok(JsonRpcWebSocketFrame::Close), - Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => { - Ok(JsonRpcWebSocketFrame::Ignore) - } + Message::Ping(payload) => Ok(JsonRpcWebSocketFrame::Ping(payload)), + Message::Pong(_) | Message::Frame(_) => Ok(JsonRpcWebSocketFrame::Ignore), } } @@ -526,6 +511,10 @@ impl JsonRpcWebSocketMessage for Message { fn ping() -> Self { Self::Ping(Vec::new().into()) } + + fn pong(payload: bytes::Bytes) -> Self { + Self::Pong(payload) + } } impl JsonRpcWebSocketMessage for AxumWebSocketMessage { @@ -538,9 +527,8 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message) } AxumWebSocketMessage::Close(_) => Ok(JsonRpcWebSocketFrame::Close), - AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => { - Ok(JsonRpcWebSocketFrame::Ignore) - } + AxumWebSocketMessage::Ping(payload) => Ok(JsonRpcWebSocketFrame::Ping(payload)), + AxumWebSocketMessage::Pong(_) => Ok(JsonRpcWebSocketFrame::Ignore), } } @@ -551,6 +539,10 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { fn ping() -> Self { Self::Ping(Vec::new().into()) } + + fn pong(payload: bytes::Bytes) -> Self { + Self::Pong(payload) + } } async fn send_disconnected( @@ -618,71 +610,60 @@ fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result, - } - - 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(())) - } + #[tokio::test] + async fn websocket_connection_pongs_server_ping() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let connection = JsonRpcConnection::from_websocket(client_websocket, "test".into()); + + server_websocket + .send(Message::Ping(b"check".to_vec().into())) + .await?; + let message = timeout(Duration::from_secs(1), server_websocket.next()) + .await? + .expect("websocket should stay open")?; + assert_eq!(message, Message::Pong(b"check".to_vec().into())); - fn poll_close( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } + drop(connection); + 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, + async fn websocket_connection_sends_configured_ping() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let connection = JsonRpcConnection::from_websocket_stream( + client_websocket, "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"); + let message = timeout(Duration::from_secs(1), server_websocket.next()) + .await? + .expect("websocket should stay open")?; assert!(matches!(message, Message::Ping(_))); 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)) } } diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 7470a6290c95..471a2e658ec8 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -19,7 +19,6 @@ 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; @@ -148,113 +147,16 @@ where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { let stream_id = Uuid::new_v4().to_string(); - let (mut websocket_writer, mut websocket_reader) = stream.split(); 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 reader_label = connection_label; - let reader_stream_id = stream_id.clone(); - let incoming_tx_for_reader = incoming_tx; - let disconnected_tx_for_reader = disconnected_tx.clone(); - let reader_task = tokio::spawn(async move { - loop { - match websocket_reader.next().await { - Some(Ok(Message::Binary(payload))) => { - let frame = match decode_relay_message_frame(payload.as_ref()) { - Ok(frame) => frame, - Err(err) => { - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::MalformedMessage { - reason: format!( - "failed to parse relay message frame from {reader_label}: {err}" - ), - }) - .await; - continue; - } - }; - if frame.stream_id != reader_stream_id { - continue; - } - let kind = match frame.validate() { - Ok(kind) => kind, - Err(err) => { - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::MalformedMessage { - reason: err.to_string(), - }) - .await; - continue; - } - }; - match kind { - RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() { - Ok(message) => { - if incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Message(message)) - .await - .is_err() - { - break; - } - } - Err(err) => { - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::MalformedMessage { - reason: err.to_string(), - }) - .await; - } - }, - RelayFrameBodyKind::Reset => { - let _ = disconnected_tx_for_reader.send(true); - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Disconnected { - reason: frame.into_reset_reason(), - }) - .await; - break; - } - RelayFrameBodyKind::Ack - | RelayFrameBodyKind::Resume - | RelayFrameBodyKind::Heartbeat => {} - } - } - Some(Ok(Message::Close(_))) | None => { - let _ = disconnected_tx_for_reader.send(true); - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Disconnected { reason: None }) - .await; - break; - } - Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {} - Some(Ok(Message::Text(_))) => { - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::MalformedMessage { - reason: "relay exec-server transport expects binary protobuf frames" - .to_string(), - }) - .await; - } - Some(Err(err)) => { - let _ = disconnected_tx_for_reader.send(true); - let _ = incoming_tx_for_reader - .send(JsonRpcConnectionEvent::Disconnected { - reason: Some(format!( - "failed to read relay websocket frame from {reader_label}: {err}" - )), - }) - .await; - break; - } - } - } - }); - - let writer_task = tokio::spawn(async move { + let websocket_task = tokio::spawn(async move { + let mut websocket = stream; + let reader_label = connection_label; + let reader_stream_id = stream_id.clone(); let resume = RelayMessageFrame::resume(stream_id.clone()); - if websocket_writer + if websocket .send(Message::Binary(encode_relay_message_frame(&resume).into())) .await .is_err() @@ -263,11 +165,6 @@ 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; loop { tokio::select! { @@ -284,7 +181,7 @@ where }; let frame = RelayMessageFrame::data(stream_id.clone(), next_seq, payload); next_seq = next_seq.wrapping_add(1); - if websocket_writer + if websocket .send(Message::Binary(encode_relay_message_frame(&frame).into())) .await .is_err() @@ -293,10 +190,102 @@ where break; } } - _ = keepalive.tick() => { - if websocket_writer.send(Message::Ping(Vec::new().into())).await.is_err() { - let _ = disconnected_tx.send(true); - break; + incoming_message = websocket.next() => { + match incoming_message { + Some(Ok(Message::Binary(payload))) => { + let frame = match decode_relay_message_frame(payload.as_ref()) { + Ok(frame) => frame, + Err(err) => { + let _ = incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: format!( + "failed to parse relay message frame from {reader_label}: {err}" + ), + }) + .await; + continue; + } + }; + if frame.stream_id != reader_stream_id { + continue; + } + let kind = match frame.validate() { + Ok(kind) => kind, + Err(err) => { + let _ = incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: err.to_string(), + }) + .await; + continue; + } + }; + match kind { + RelayFrameBodyKind::Data => match frame.into_jsonrpc_message() { + Ok(message) => { + if incoming_tx + .send(JsonRpcConnectionEvent::Message(message)) + .await + .is_err() + { + break; + } + } + Err(err) => { + let _ = incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: err.to_string(), + }) + .await; + } + }, + RelayFrameBodyKind::Reset => { + let _ = disconnected_tx.send(true); + let _ = incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { + reason: frame.into_reset_reason(), + }) + .await; + break; + } + RelayFrameBodyKind::Ack + | RelayFrameBodyKind::Resume + | RelayFrameBodyKind::Heartbeat => {} + } + } + Some(Ok(Message::Ping(payload))) => { + if websocket.send(Message::Pong(payload)).await.is_err() { + let _ = disconnected_tx.send(true); + break; + } + } + Some(Ok(Message::Close(_))) | None => { + let _ = disconnected_tx.send(true); + let _ = incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { reason: None }) + .await; + break; + } + Some(Ok(Message::Pong(_) | Message::Frame(_))) => {} + Some(Ok(Message::Text(_))) => { + let _ = incoming_tx + .send(JsonRpcConnectionEvent::MalformedMessage { + reason: "relay exec-server transport expects binary protobuf frames" + .to_string(), + }) + .await; + } + Some(Err(err)) => { + let _ = disconnected_tx.send(true); + let _ = incoming_tx + .send(JsonRpcConnectionEvent::Disconnected { + reason: Some(format!( + "failed to read relay websocket frame from {reader_label}: {err}" + )), + }) + .await; + break; + } } } } @@ -307,7 +296,7 @@ where outgoing_tx, incoming_rx, disconnected_rx, - task_handles: vec![reader_task, writer_task], + task_handles: vec![websocket_task], transport: JsonRpcTransport::Plain, } } @@ -318,59 +307,48 @@ pub(crate) async fn run_multiplexed_executor( ) where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - let (mut websocket_writer, mut websocket_reader) = stream.split(); + let mut websocket = stream; let (physical_outgoing_tx, mut physical_outgoing_rx) = mpsc::channel::>(CHANNEL_CAPACITY); - let writer_task = tokio::spawn(async move { - 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; - } - } - } - } - }); let mut streams: HashMap = HashMap::new(); loop { - let frame = match websocket_reader.next().await { - Some(Ok(Message::Binary(payload))) => { - match decode_relay_message_frame(payload.as_ref()) { - Ok(frame) => frame, - Err(err) => { - warn!("dropping malformed relay message frame from harness: {err}"); - continue; - } + let frame = tokio::select! { + maybe_encoded = physical_outgoing_rx.recv() => { + let Some(encoded) = maybe_encoded else { + break; + }; + if websocket.send(Message::Binary(encoded.into())).await.is_err() { + break; } - } - Some(Ok(Message::Close(_))) | None => break, - Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue, - Some(Ok(Message::Text(_))) => { - warn!("dropping non-binary relay message frame from harness"); continue; } - Some(Err(err)) => { - debug!("multiplexed executor websocket read failed: {err}"); - break; + incoming_message = websocket.next() => match incoming_message { + Some(Ok(Message::Binary(payload))) => { + match decode_relay_message_frame(payload.as_ref()) { + Ok(frame) => frame, + Err(err) => { + warn!("dropping malformed relay message frame from harness: {err}"); + continue; + } + } + } + Some(Ok(Message::Ping(payload))) => { + if websocket.send(Message::Pong(payload)).await.is_err() { + break; + } + continue; + } + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(Message::Pong(_) | Message::Frame(_))) => continue, + Some(Ok(Message::Text(_))) => { + warn!("dropping non-binary relay message frame from harness"); + continue; + } + Some(Err(err)) => { + debug!("multiplexed executor websocket read failed: {err}"); + break; + } } }; @@ -423,7 +401,6 @@ pub(crate) async fn run_multiplexed_executor( stream.disconnect(/*reason*/ None).await; } drop(physical_outgoing_tx); - let _ = writer_task.await; } struct VirtualStream { @@ -511,14 +488,17 @@ mod tests { } #[tokio::test] - async fn multiplexed_executor_sends_keepalive_ping() -> anyhow::Result<()> { + async fn multiplexed_executor_pongs_server_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?; + server_websocket + .send(Message::Ping(b"check".to_vec().into())) + .await?; + read_pong(&mut server_websocket).await?; executor_task.abort(); let _ = executor_task.await; @@ -526,11 +506,14 @@ mod tests { } #[tokio::test] - async fn harness_connection_sends_keepalive_ping() -> anyhow::Result<()> { + async fn harness_connection_pongs_server_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?; + server_websocket + .send(Message::Ping(b"check".to_vec().into())) + .await?; + read_pong(&mut server_websocket).await?; drop(connection); Ok(()) @@ -551,17 +534,21 @@ mod tests { Ok((client_websocket, server_websocket)) } - async fn read_keepalive_ping( + async fn read_pong( 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"); + anyhow::bail!("websocket closed before pong"); }; match message? { - Message::Ping(_) => return Ok(()), - Message::Binary(_) | Message::Text(_) | Message::Pong(_) | Message::Frame(_) => {} - Message::Close(_) => anyhow::bail!("websocket closed before keepalive ping"), + Message::Pong(payload) if payload.as_ref() == b"check" => return Ok(()), + Message::Binary(_) + | Message::Text(_) + | Message::Ping(_) + | Message::Pong(_) + | Message::Frame(_) => {} + Message::Close(_) => anyhow::bail!("websocket closed before pong"), } } } From 3673b69a2ab5f10edaf10ba8e36db29463198840 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 10:31:48 -0700 Subject: [PATCH 2/8] Add exec-server websocket pump tests --- codex-rs/exec-server/src/connection.rs | 58 +++++++++++++++++ codex-rs/exec-server/src/relay.rs | 87 ++++++++++++++++++++++++++ 2 files changed, 145 insertions(+) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 0974a21c3a24..ca1b70b85dc9 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -610,6 +610,8 @@ fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into()); + + server_websocket + .send(Message::Pong(b"check".to_vec().into())) + .await?; + assert!( + timeout(Duration::from_millis(50), connection.incoming_rx.recv()) + .await + .is_err() + ); + + drop(connection); + Ok(()) + } + + #[tokio::test] + async fn websocket_connection_reports_server_close() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into()); + + server_websocket.close(None).await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Disconnected { reason: None }) + )); + + drop(connection); + Ok(()) + } + + #[tokio::test] + async fn websocket_connection_accepts_binary_jsonrpc_message() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = JsonRpcConnection::from_websocket(client_websocket, "test".into()); + let message = JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "test".to_string(), + params: None, + trace: None, + }); + + server_websocket + .send(Message::Binary(serde_json::to_vec(&message)?.into())) + .await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Message(actual)) if actual == message + )); + + drop(connection); + Ok(()) + } + async fn websocket_pair() -> anyhow::Result<( WebSocketStream>, WebSocketStream, diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 471a2e658ec8..6a6fef9d6238 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -471,6 +471,8 @@ fn spawn_virtual_stream( mod tests { use std::time::Duration; + use codex_app_server_protocol::JSONRPCRequest; + use codex_app_server_protocol::RequestId; use tokio::net::TcpListener; use tokio::time::timeout; use tokio_tungstenite::accept_async; @@ -519,6 +521,68 @@ mod tests { Ok(()) } + #[tokio::test] + async fn harness_connection_receives_relay_data() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = + harness_connection_from_websocket(client_websocket, "test".to_string()); + let stream_id = read_resume_stream_id(&mut server_websocket).await?; + let message = test_jsonrpc_message(); + + server_websocket + .send(Message::Binary( + encode_relay_message_frame(&RelayMessageFrame::data( + stream_id, + /*seq*/ 0, + jsonrpc_payload(&message)?, + )) + .into(), + )) + .await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Message(actual)) if actual == message + )); + + drop(connection); + Ok(()) + } + + #[tokio::test] + async fn harness_connection_reports_text_frames_as_malformed() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = + harness_connection_from_websocket(client_websocket, "test".to_string()); + + read_resume_stream_id(&mut server_websocket).await?; + server_websocket.send(Message::Text("nope".into())).await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::MalformedMessage { reason }) + if reason == "relay exec-server transport expects binary protobuf frames" + )); + + drop(connection); + Ok(()) + } + + #[tokio::test] + async fn harness_connection_reports_server_close() -> anyhow::Result<()> { + let (client_websocket, mut server_websocket) = websocket_pair().await?; + let mut connection = + harness_connection_from_websocket(client_websocket, "test".to_string()); + + read_resume_stream_id(&mut server_websocket).await?; + server_websocket.close(None).await?; + assert!(matches!( + timeout(Duration::from_secs(1), connection.incoming_rx.recv()).await?, + Some(JsonRpcConnectionEvent::Disconnected { reason: None }) + )); + + drop(connection); + Ok(()) + } + async fn websocket_pair() -> anyhow::Result<( WebSocketStream>, WebSocketStream, @@ -534,6 +598,29 @@ mod tests { Ok((client_websocket, server_websocket)) } + async fn read_resume_stream_id( + websocket: &mut WebSocketStream, + ) -> anyhow::Result { + let message = timeout(Duration::from_secs(1), websocket.next()) + .await? + .expect("websocket should stay open")?; + let Message::Binary(payload) = message else { + anyhow::bail!("expected relay resume frame, got {message:?}"); + }; + let frame = decode_relay_message_frame(payload.as_ref())?; + assert_eq!(frame.validate()?, RelayFrameBodyKind::Resume); + Ok(frame.stream_id) + } + + fn test_jsonrpc_message() -> JSONRPCMessage { + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "test".to_string(), + params: None, + trace: None, + }) + } + async fn read_pong( websocket: &mut WebSocketStream, ) -> anyhow::Result<()> { From d94300d782c45ea66aa95afacf789f0b9779c03a Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 10:42:43 -0700 Subject: [PATCH 3/8] Preserve exec-server websocket keepalive ownership --- codex-rs/exec-server/src/connection.rs | 4 ++-- codex-rs/exec-server/src/relay.rs | 23 +++++++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index ca1b70b85dc9..d1e2ad7720c4 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -323,11 +323,11 @@ impl JsonRpcConnection { where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) + Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL)) } pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self { - Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL)) + Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) } fn from_websocket_stream( diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 6a6fef9d6238..7a676e346988 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; @@ -166,6 +167,11 @@ where } let mut next_seq = 0u32; + 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_message = outgoing_rx.recv() => { @@ -190,6 +196,12 @@ where break; } } + _ = keepalive.tick() => { + if websocket.send(Message::Ping(Vec::new().into())).await.is_err() { + let _ = disconnected_tx.send(true); + break; + } + } incoming_message = websocket.next() => { match incoming_message { Some(Ok(Message::Binary(payload))) => { @@ -312,6 +324,11 @@ pub(crate) async fn run_multiplexed_executor( mpsc::channel::>(CHANNEL_CAPACITY); let mut streams: HashMap = HashMap::new(); + 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 { let frame = tokio::select! { maybe_encoded = physical_outgoing_rx.recv() => { @@ -323,6 +340,12 @@ pub(crate) async fn run_multiplexed_executor( } continue; } + _ = keepalive.tick() => { + if websocket.send(Message::Ping(Vec::new().into())).await.is_err() { + break; + } + continue; + } incoming_message = websocket.next() => match incoming_message { Some(Ok(Message::Binary(payload))) => { match decode_relay_message_frame(payload.as_ref()) { From 9215e15ee3d76c565be678e255072ee31c6a5e11 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 15:24:27 -0700 Subject: [PATCH 4/8] Add websocket backpressure regression tests --- codex-rs/exec-server/src/connection.rs | 145 ++++++++++++++++++++++ codex-rs/exec-server/src/relay.rs | 162 ++++++++++++++++++++++++- 2 files changed, 304 insertions(+), 3 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index d1e2ad7720c4..4c1ea630621a 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -610,8 +610,17 @@ fn serialize_jsonrpc_message(message: &JSONRPCMessage) -> Result anyhow::Result<()> { + let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(false); + let mut connection = + JsonRpcConnection::from_websocket_stream(websocket, "test".into(), None); + let message = test_jsonrpc_message(); + + connection.outgoing_tx.send(message.clone()).await?; + control.send_inbound(Message::Ping(b"check".to_vec().into()))?; + assert!( + timeout(Duration::from_millis(50), connection.incoming_rx.recv()) + .await + .is_err() + ); + + control.set_write_ready(); + assert!(matches!( + timeout(Duration::from_secs(1), outbound_rx.next()).await?, + Some(Message::Text(text)) if serde_json::from_str::(&text)? == message + )); + assert!(matches!( + timeout(Duration::from_secs(1), outbound_rx.next()).await?, + Some(Message::Pong(payload)) if payload.as_ref() == b"check" + )); + + drop(connection); + Ok(()) + } + async fn websocket_pair() -> anyhow::Result<( WebSocketStream>, WebSocketStream, @@ -724,4 +763,110 @@ mod tests { let server_websocket = server_task.await??; Ok((client_websocket, server_websocket)) } + + fn test_jsonrpc_message() -> JSONRPCMessage { + JSONRPCMessage::Request(JSONRPCRequest { + id: RequestId::Integer(1), + method: "test".to_string(), + params: None, + trace: None, + }) + } + + struct ControlledWebSocket { + inbound_rx: futures_mpsc::UnboundedReceiver>, + outbound_tx: futures_mpsc::UnboundedSender, + write_ready: Arc, + write_waker: Arc, + } + + struct ControlledWebSocketHandle { + inbound_tx: futures_mpsc::UnboundedSender>, + write_ready: Arc, + write_waker: Arc, + } + + impl ControlledWebSocket { + fn new( + write_ready: bool, + ) -> ( + Self, + ControlledWebSocketHandle, + futures_mpsc::UnboundedReceiver, + ) { + let (inbound_tx, inbound_rx) = futures_mpsc::unbounded(); + let (outbound_tx, outbound_rx) = futures_mpsc::unbounded(); + let write_ready = Arc::new(AtomicBool::new(write_ready)); + let write_waker = Arc::new(AtomicWaker::new()); + ( + Self { + inbound_rx, + outbound_tx, + write_ready: Arc::clone(&write_ready), + write_waker: Arc::clone(&write_waker), + }, + ControlledWebSocketHandle { + inbound_tx, + write_ready, + write_waker, + }, + outbound_rx, + ) + } + } + + impl ControlledWebSocketHandle { + fn send_inbound(&self, message: Message) -> anyhow::Result<()> { + self.inbound_tx + .unbounded_send(Ok(message)) + .map_err(anyhow::Error::from) + } + + fn set_write_ready(&self) { + self.write_ready.store(true, Ordering::Release); + self.write_waker.wake(); + } + } + + impl Sink for ControlledWebSocket { + type Error = std::convert::Infallible; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.write_ready.load(Ordering::Acquire) { + Poll::Ready(Ok(())) + } else { + self.write_waker.register(cx.waker()); + Poll::Pending + } + } + + fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + self.outbound_tx + .unbounded_send(item) + .expect("test outbound 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(())) + } + } + + impl Stream for ControlledWebSocket { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inbound_rx).poll_next(cx) + } + } } diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 7a676e346988..7d3751c67df1 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,7 +1,9 @@ use std::collections::HashMap; use codex_app_server_protocol::JSONRPCMessage; +use futures::Sink; use futures::SinkExt; +use futures::Stream; use futures::StreamExt; use prost::Message as ProstMessage; use tokio::io::AsyncRead; @@ -140,12 +142,13 @@ fn jsonrpc_payload(message: &JSONRPCMessage) -> Result, ExecServerError> serde_json::to_vec(message).map_err(ExecServerError::Json) } -pub(crate) fn harness_connection_from_websocket( - stream: WebSocketStream, +pub(crate) fn harness_connection_from_websocket( + stream: T, connection_label: String, ) -> JsonRpcConnection where - S: AsyncRead + AsyncWrite + Unpin + Send + 'static, + T: Sink + Stream> + Unpin + Send + 'static, + E: std::fmt::Display + Send + 'static, { let stream_id = Uuid::new_v4().to_string(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); @@ -492,10 +495,20 @@ fn spawn_virtual_stream( #[cfg(test)] mod tests { + use std::pin::Pin; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; + use std::task::Context; + use std::task::Poll; use std::time::Duration; use codex_app_server_protocol::JSONRPCRequest; use codex_app_server_protocol::RequestId; + use futures::Sink; + use futures::Stream; + use futures::channel::mpsc as futures_mpsc; + use futures::task::AtomicWaker; use tokio::net::TcpListener; use tokio::time::timeout; use tokio_tungstenite::accept_async; @@ -606,6 +619,48 @@ mod tests { Ok(()) } + #[tokio::test] + async fn harness_connection_keeps_outbound_frame_while_send_is_backpressured() + -> anyhow::Result<()> { + let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(true); + let mut connection = harness_connection_from_websocket(websocket, "test".to_string()); + let Message::Binary(resume_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .expect("resume frame") + else { + anyhow::bail!("expected relay resume frame"); + }; + let stream_id = decode_relay_message_frame(resume_payload.as_ref())?.stream_id; + let message = test_jsonrpc_message(); + + control.set_write_blocked(); + connection.outgoing_tx.send(message.clone()).await?; + control.send_inbound(Message::Ping(b"check".to_vec().into()))?; + assert!( + timeout(Duration::from_millis(50), connection.incoming_rx.recv()) + .await + .is_err() + ); + + control.set_write_ready(); + let Message::Binary(data_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) + .await? + .expect("data frame") + else { + anyhow::bail!("expected relay data frame"); + }; + let frame = decode_relay_message_frame(data_payload.as_ref())?; + assert_eq!(frame.stream_id, stream_id); + assert_eq!(frame.into_jsonrpc_message()?, message); + assert!(matches!( + timeout(Duration::from_secs(1), outbound_rx.next()).await?, + Some(Message::Pong(payload)) if payload.as_ref() == b"check" + )); + + drop(connection); + Ok(()) + } + async fn websocket_pair() -> anyhow::Result<( WebSocketStream>, WebSocketStream, @@ -662,4 +717,105 @@ mod tests { } } } + + struct ControlledWebSocket { + inbound_rx: futures_mpsc::UnboundedReceiver>, + outbound_tx: futures_mpsc::UnboundedSender, + write_ready: Arc, + write_waker: Arc, + } + + struct ControlledWebSocketHandle { + inbound_tx: futures_mpsc::UnboundedSender>, + write_ready: Arc, + write_waker: Arc, + } + + impl ControlledWebSocket { + fn new( + write_ready: bool, + ) -> ( + Self, + ControlledWebSocketHandle, + futures_mpsc::UnboundedReceiver, + ) { + let (inbound_tx, inbound_rx) = futures_mpsc::unbounded(); + let (outbound_tx, outbound_rx) = futures_mpsc::unbounded(); + let write_ready = Arc::new(AtomicBool::new(write_ready)); + let write_waker = Arc::new(AtomicWaker::new()); + ( + Self { + inbound_rx, + outbound_tx, + write_ready: Arc::clone(&write_ready), + write_waker: Arc::clone(&write_waker), + }, + ControlledWebSocketHandle { + inbound_tx, + write_ready, + write_waker, + }, + outbound_rx, + ) + } + } + + impl ControlledWebSocketHandle { + fn send_inbound(&self, message: Message) -> anyhow::Result<()> { + self.inbound_tx + .unbounded_send(Ok(message)) + .map_err(anyhow::Error::from) + } + + fn set_write_blocked(&self) { + self.write_ready.store(false, Ordering::Release); + } + + fn set_write_ready(&self) { + self.write_ready.store(true, Ordering::Release); + self.write_waker.wake(); + } + } + + impl Sink for ControlledWebSocket { + type Error = std::convert::Infallible; + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.write_ready.load(Ordering::Acquire) { + Poll::Ready(Ok(())) + } else { + self.write_waker.register(cx.waker()); + Poll::Pending + } + } + + fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> { + self.outbound_tx + .unbounded_send(item) + .expect("test outbound 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(())) + } + } + + impl Stream for ControlledWebSocket { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inbound_rx).poll_next(cx) + } + } } From 6d6cdeb128960a4a3b0673e4e31d07e147493902 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 15:36:41 -0700 Subject: [PATCH 5/8] Fix websocket backpressure regression tests --- codex-rs/exec-server/src/connection.rs | 39 ++++++++++++++++++++++++-- codex-rs/exec-server/src/relay.rs | 32 ++++++++++++++++++++- 2 files changed, 67 insertions(+), 4 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 4c1ea630621a..9d88a41a0cdd 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -722,12 +722,17 @@ mod tests { #[tokio::test] async fn websocket_connection_keeps_outbound_message_while_send_is_backpressured() -> anyhow::Result<()> { - let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(false); - let mut connection = - JsonRpcConnection::from_websocket_stream(websocket, "test".into(), None); + let (websocket, control, mut outbound_rx) = + ControlledWebSocket::new(/*write_ready*/ false); + let mut connection = JsonRpcConnection::from_websocket_stream( + websocket, + "test".into(), + /*ping_interval*/ None, + ); let message = test_jsonrpc_message(); connection.outgoing_tx.send(message.clone()).await?; + control.wait_for_blocked_write().await?; control.send_inbound(Message::Ping(b"check".to_vec().into()))?; assert!( timeout(Duration::from_millis(50), connection.incoming_rx.recv()) @@ -777,12 +782,16 @@ mod tests { inbound_rx: futures_mpsc::UnboundedReceiver>, outbound_tx: futures_mpsc::UnboundedSender, write_ready: Arc, + write_blocked: Arc, + write_blocked_waker: Arc, write_waker: Arc, } struct ControlledWebSocketHandle { inbound_tx: futures_mpsc::UnboundedSender>, write_ready: Arc, + write_blocked: Arc, + write_blocked_waker: Arc, write_waker: Arc, } @@ -797,17 +806,23 @@ mod tests { let (inbound_tx, inbound_rx) = futures_mpsc::unbounded(); let (outbound_tx, outbound_rx) = futures_mpsc::unbounded(); let write_ready = Arc::new(AtomicBool::new(write_ready)); + let write_blocked = Arc::new(AtomicBool::new(false)); + let write_blocked_waker = Arc::new(AtomicWaker::new()); let write_waker = Arc::new(AtomicWaker::new()); ( Self { inbound_rx, outbound_tx, write_ready: Arc::clone(&write_ready), + write_blocked: Arc::clone(&write_blocked), + write_blocked_waker: Arc::clone(&write_blocked_waker), write_waker: Arc::clone(&write_waker), }, ControlledWebSocketHandle { inbound_tx, write_ready, + write_blocked, + write_blocked_waker, write_waker, }, outbound_rx, @@ -826,6 +841,22 @@ mod tests { self.write_ready.store(true, Ordering::Release); self.write_waker.wake(); } + + async fn wait_for_blocked_write(&self) -> anyhow::Result<()> { + timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + if self.write_blocked.load(Ordering::Acquire) { + Poll::Ready(()) + } else { + self.write_blocked_waker.register(cx.waker()); + Poll::Pending + } + }), + ) + .await?; + Ok(()) + } } impl Sink for ControlledWebSocket { @@ -835,6 +866,8 @@ mod tests { if self.write_ready.load(Ordering::Acquire) { Poll::Ready(Ok(())) } else { + self.write_blocked.store(true, Ordering::Release); + self.write_blocked_waker.wake(); self.write_waker.register(cx.waker()); Poll::Pending } diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 7d3751c67df1..4dd607a305c0 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -622,7 +622,8 @@ mod tests { #[tokio::test] async fn harness_connection_keeps_outbound_frame_while_send_is_backpressured() -> anyhow::Result<()> { - let (websocket, control, mut outbound_rx) = ControlledWebSocket::new(true); + let (websocket, control, mut outbound_rx) = + ControlledWebSocket::new(/*write_ready*/ true); let mut connection = harness_connection_from_websocket(websocket, "test".to_string()); let Message::Binary(resume_payload) = timeout(Duration::from_secs(1), outbound_rx.next()) .await? @@ -635,6 +636,7 @@ mod tests { control.set_write_blocked(); connection.outgoing_tx.send(message.clone()).await?; + control.wait_for_blocked_write().await?; control.send_inbound(Message::Ping(b"check".to_vec().into()))?; assert!( timeout(Duration::from_millis(50), connection.incoming_rx.recv()) @@ -722,12 +724,16 @@ mod tests { inbound_rx: futures_mpsc::UnboundedReceiver>, outbound_tx: futures_mpsc::UnboundedSender, write_ready: Arc, + write_blocked: Arc, + write_blocked_waker: Arc, write_waker: Arc, } struct ControlledWebSocketHandle { inbound_tx: futures_mpsc::UnboundedSender>, write_ready: Arc, + write_blocked: Arc, + write_blocked_waker: Arc, write_waker: Arc, } @@ -742,17 +748,23 @@ mod tests { let (inbound_tx, inbound_rx) = futures_mpsc::unbounded(); let (outbound_tx, outbound_rx) = futures_mpsc::unbounded(); let write_ready = Arc::new(AtomicBool::new(write_ready)); + let write_blocked = Arc::new(AtomicBool::new(false)); + let write_blocked_waker = Arc::new(AtomicWaker::new()); let write_waker = Arc::new(AtomicWaker::new()); ( Self { inbound_rx, outbound_tx, write_ready: Arc::clone(&write_ready), + write_blocked: Arc::clone(&write_blocked), + write_blocked_waker: Arc::clone(&write_blocked_waker), write_waker: Arc::clone(&write_waker), }, ControlledWebSocketHandle { inbound_tx, write_ready, + write_blocked, + write_blocked_waker, write_waker, }, outbound_rx, @@ -775,6 +787,22 @@ mod tests { self.write_ready.store(true, Ordering::Release); self.write_waker.wake(); } + + async fn wait_for_blocked_write(&self) -> anyhow::Result<()> { + timeout( + Duration::from_secs(1), + futures::future::poll_fn(|cx| { + if self.write_blocked.load(Ordering::Acquire) { + Poll::Ready(()) + } else { + self.write_blocked_waker.register(cx.waker()); + Poll::Pending + } + }), + ) + .await?; + Ok(()) + } } impl Sink for ControlledWebSocket { @@ -784,6 +812,8 @@ mod tests { if self.write_ready.load(Ordering::Acquire) { Poll::Ready(Ok(())) } else { + self.write_blocked.store(true, Ordering::Release); + self.write_blocked_waker.wake(); self.write_waker.register(cx.waker()); Poll::Pending } From 6f9640533a9b3d2b1be9b4f292de1f6c1d92d384 Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 16:00:04 -0700 Subject: [PATCH 6/8] Restore server-owned websocket keepalive --- codex-rs/exec-server/src/connection.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 9d88a41a0cdd..750c9e09955e 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -323,11 +323,11 @@ impl JsonRpcConnection { where S: AsyncRead + AsyncWrite + Unpin + Send + 'static, { - Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL)) + Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) } pub(crate) fn from_axum_websocket(stream: AxumWebSocket, connection_label: String) -> Self { - Self::from_websocket_stream(stream, connection_label, /*ping_interval*/ None) + Self::from_websocket_stream(stream, connection_label, Some(WEBSOCKET_KEEPALIVE_INTERVAL)) } fn from_websocket_stream( From 90804bb2eb03b096470641c55f823142de7613bc Mon Sep 17 00:00:00 2001 From: starr-openai Date: Mon, 18 May 2026 16:14:53 -0700 Subject: [PATCH 7/8] Remove manual websocket pong handling --- codex-rs/exec-server/src/connection.rs | 57 +++--------------- codex-rs/exec-server/src/relay.rs | 82 +------------------------- 2 files changed, 10 insertions(+), 129 deletions(-) diff --git a/codex-rs/exec-server/src/connection.rs b/codex-rs/exec-server/src/connection.rs index 750c9e09955e..b211c504acab 100644 --- a/codex-rs/exec-server/src/connection.rs +++ b/codex-rs/exec-server/src/connection.rs @@ -401,19 +401,6 @@ impl JsonRpcConnection { break; } } - Ok(JsonRpcWebSocketFrame::Ping(payload)) => { - if let Err(err) = websocket.send(M::pong(payload)).await { - send_disconnected( - &incoming_tx, - &disconnected_tx, - Some(format!( - "failed to write websocket pong to {connection_label}: {err}" - )), - ) - .await; - break; - } - } Ok(JsonRpcWebSocketFrame::Close) => { send_disconnected( &incoming_tx, @@ -477,7 +464,6 @@ impl JsonRpcConnection { enum JsonRpcWebSocketFrame { Message(JSONRPCMessage), - Ping(bytes::Bytes), Close, Ignore, } @@ -486,7 +472,6 @@ trait JsonRpcWebSocketMessage: Send + 'static { fn parse_jsonrpc_frame(self) -> Result; fn from_text(text: String) -> Self; fn ping() -> Self; - fn pong(payload: bytes::Bytes) -> Self; } impl JsonRpcWebSocketMessage for Message { @@ -499,8 +484,9 @@ impl JsonRpcWebSocketMessage for Message { serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message) } Message::Close(_) => Ok(JsonRpcWebSocketFrame::Close), - Message::Ping(payload) => Ok(JsonRpcWebSocketFrame::Ping(payload)), - Message::Pong(_) | Message::Frame(_) => Ok(JsonRpcWebSocketFrame::Ignore), + Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => { + Ok(JsonRpcWebSocketFrame::Ignore) + } } } @@ -511,10 +497,6 @@ impl JsonRpcWebSocketMessage for Message { fn ping() -> Self { Self::Ping(Vec::new().into()) } - - fn pong(payload: bytes::Bytes) -> Self { - Self::Pong(payload) - } } impl JsonRpcWebSocketMessage for AxumWebSocketMessage { @@ -527,8 +509,9 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { serde_json::from_slice(bytes.as_ref()).map(JsonRpcWebSocketFrame::Message) } AxumWebSocketMessage::Close(_) => Ok(JsonRpcWebSocketFrame::Close), - AxumWebSocketMessage::Ping(payload) => Ok(JsonRpcWebSocketFrame::Ping(payload)), - AxumWebSocketMessage::Pong(_) => Ok(JsonRpcWebSocketFrame::Ignore), + AxumWebSocketMessage::Ping(_) | AxumWebSocketMessage::Pong(_) => { + Ok(JsonRpcWebSocketFrame::Ignore) + } } } @@ -539,10 +522,6 @@ impl JsonRpcWebSocketMessage for AxumWebSocketMessage { fn ping() -> Self { Self::Ping(Vec::new().into()) } - - fn pong(payload: bytes::Bytes) -> Self { - Self::Pong(payload) - } } async fn send_disconnected( @@ -628,23 +607,6 @@ mod tests { use super::*; - #[tokio::test] - async fn websocket_connection_pongs_server_ping() -> anyhow::Result<()> { - let (client_websocket, mut server_websocket) = websocket_pair().await?; - let connection = JsonRpcConnection::from_websocket(client_websocket, "test".into()); - - server_websocket - .send(Message::Ping(b"check".to_vec().into())) - .await?; - let message = timeout(Duration::from_secs(1), server_websocket.next()) - .await? - .expect("websocket should stay open")?; - assert_eq!(message, Message::Pong(b"check".to_vec().into())); - - drop(connection); - Ok(()) - } - #[tokio::test] async fn websocket_connection_sends_configured_ping() -> anyhow::Result<()> { let (client_websocket, mut server_websocket) = websocket_pair().await?; @@ -733,7 +695,7 @@ mod tests { connection.outgoing_tx.send(message.clone()).await?; control.wait_for_blocked_write().await?; - control.send_inbound(Message::Ping(b"check".to_vec().into()))?; + control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( timeout(Duration::from_millis(50), connection.incoming_rx.recv()) .await @@ -745,11 +707,6 @@ mod tests { timeout(Duration::from_secs(1), outbound_rx.next()).await?, Some(Message::Text(text)) if serde_json::from_str::(&text)? == message )); - assert!(matches!( - timeout(Duration::from_secs(1), outbound_rx.next()).await?, - Some(Message::Pong(payload)) if payload.as_ref() == b"check" - )); - drop(connection); Ok(()) } diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index 4dd607a305c0..da7a7422bb92 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -268,12 +268,6 @@ where | RelayFrameBodyKind::Heartbeat => {} } } - Some(Ok(Message::Ping(payload))) => { - if websocket.send(Message::Pong(payload)).await.is_err() { - let _ = disconnected_tx.send(true); - break; - } - } Some(Ok(Message::Close(_))) | None => { let _ = disconnected_tx.send(true); let _ = incoming_tx @@ -281,7 +275,7 @@ where .await; break; } - Some(Ok(Message::Pong(_) | Message::Frame(_))) => {} + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {} Some(Ok(Message::Text(_))) => { let _ = incoming_tx .send(JsonRpcConnectionEvent::MalformedMessage { @@ -359,14 +353,8 @@ pub(crate) async fn run_multiplexed_executor( } } } - Some(Ok(Message::Ping(payload))) => { - if websocket.send(Message::Pong(payload)).await.is_err() { - break; - } - continue; - } Some(Ok(Message::Close(_))) | None => break, - Some(Ok(Message::Pong(_) | Message::Frame(_))) => continue, + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => continue, Some(Ok(Message::Text(_))) => { warn!("dropping non-binary relay message frame from harness"); continue; @@ -517,46 +505,6 @@ mod tests { 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_pongs_server_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()?), - )); - - server_websocket - .send(Message::Ping(b"check".to_vec().into())) - .await?; - read_pong(&mut server_websocket).await?; - - executor_task.abort(); - let _ = executor_task.await; - Ok(()) - } - - #[tokio::test] - async fn harness_connection_pongs_server_ping() -> anyhow::Result<()> { - let (client_websocket, mut server_websocket) = websocket_pair().await?; - let connection = harness_connection_from_websocket(client_websocket, "test".to_string()); - - server_websocket - .send(Message::Ping(b"check".to_vec().into())) - .await?; - read_pong(&mut server_websocket).await?; - - drop(connection); - Ok(()) - } - #[tokio::test] async fn harness_connection_receives_relay_data() -> anyhow::Result<()> { let (client_websocket, mut server_websocket) = websocket_pair().await?; @@ -637,7 +585,7 @@ mod tests { control.set_write_blocked(); connection.outgoing_tx.send(message.clone()).await?; control.wait_for_blocked_write().await?; - control.send_inbound(Message::Ping(b"check".to_vec().into()))?; + control.send_inbound(Message::Pong(b"check".to_vec().into()))?; assert!( timeout(Duration::from_millis(50), connection.incoming_rx.recv()) .await @@ -654,11 +602,6 @@ mod tests { let frame = decode_relay_message_frame(data_payload.as_ref())?; assert_eq!(frame.stream_id, stream_id); assert_eq!(frame.into_jsonrpc_message()?, message); - assert!(matches!( - timeout(Duration::from_secs(1), outbound_rx.next()).await?, - Some(Message::Pong(payload)) if payload.as_ref() == b"check" - )); - drop(connection); Ok(()) } @@ -701,25 +644,6 @@ mod tests { }) } - async fn read_pong( - websocket: &mut WebSocketStream, - ) -> anyhow::Result<()> { - loop { - let Some(message) = timeout(Duration::from_secs(1), websocket.next()).await? else { - anyhow::bail!("websocket closed before pong"); - }; - match message? { - Message::Pong(payload) if payload.as_ref() == b"check" => return Ok(()), - Message::Binary(_) - | Message::Text(_) - | Message::Ping(_) - | Message::Pong(_) - | Message::Frame(_) => {} - Message::Close(_) => anyhow::bail!("websocket closed before pong"), - } - } - } - struct ControlledWebSocket { inbound_rx: futures_mpsc::UnboundedReceiver>, outbound_tx: futures_mpsc::UnboundedSender, From 7830b18af78b72f71344f95cd4afc7f47cc44e5f Mon Sep 17 00:00:00 2001 From: starr-openai Date: Tue, 19 May 2026 10:54:00 -0700 Subject: [PATCH 8/8] Remove relay websocket keepalive pings --- codex-rs/exec-server/src/relay.rs | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index da7a7422bb92..b3708f1d8371 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -21,7 +21,6 @@ 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; @@ -170,11 +169,6 @@ where } let mut next_seq = 0u32; - 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_message = outgoing_rx.recv() => { @@ -199,12 +193,6 @@ where break; } } - _ = keepalive.tick() => { - if websocket.send(Message::Ping(Vec::new().into())).await.is_err() { - let _ = disconnected_tx.send(true); - break; - } - } incoming_message = websocket.next() => { match incoming_message { Some(Ok(Message::Binary(payload))) => { @@ -321,11 +309,6 @@ pub(crate) async fn run_multiplexed_executor( mpsc::channel::>(CHANNEL_CAPACITY); let mut streams: HashMap = HashMap::new(); - 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 { let frame = tokio::select! { maybe_encoded = physical_outgoing_rx.recv() => { @@ -337,12 +320,6 @@ pub(crate) async fn run_multiplexed_executor( } continue; } - _ = keepalive.tick() => { - if websocket.send(Message::Ping(Vec::new().into())).await.is_err() { - break; - } - continue; - } incoming_message = websocket.next() => match incoming_message { Some(Ok(Message::Binary(payload))) => { match decode_relay_message_frame(payload.as_ref()) {