Bound relay connection attempts and handle pongs during blocked writes (#51502)

## Why

A stalled rendezvous WebSocket upgrade can prevent reconnects, while short-lived connections reset reconnect backoff too early. Blocked relay writes also prevent received pongs from clearing their heartbeat deadline.

## What changed

- Bound rendezvous WebSocket connection attempts with `DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT`.
- Retain exponential reconnect backoff across short-lived connections, resetting it only after a connection lasts at least 120 seconds.
- Process pongs during blocked data writes while preserving the in-flight send and its independent write timeout. Prefer queued pongs over an expired heartbeat deadline.
- Abort the relay writer when its owning connection task is dropped.

## Testing

Add coverage for pongs arriving during blocked writes, missing and queued pongs, independent write timeouts, and writer cancellation. Extend reconnect tests to cover stalled WebSocket upgrades without re-registration and backoff reset after a stable session.

GitOrigin-RevId: 0efa350c2b4ada9b89e23df8909063fbb69f33dc

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

@@ -15,6 +15,7 @@ use tokio::task::JoinSet;
 use tokio::time::Instant;
 use tokio::time::timeout;
 use tokio_tungstenite::tungstenite::Message;
+use tokio_util::task::AbortOnDropHandle;
 use tracing::debug;
 use tracing::info;
 use tracing::warn;
@@ -46,9 +47,10 @@ use crate::relay_proto::relay_message_frame;
 #[cfg(test)]
 use crate::server::ConnectionProcessor;
 use crate::telemetry::ExecutorRegistration;
-use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT;
 use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT_REASON;
-use crate::websocket_pong_watchdog::WebSocketPongWatchdog;
+
+#[path = "relay_writer.rs"]
+mod writer;
 
 const RELAY_MESSAGE_FRAME_VERSION: u32 = 1;
 const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128;
@@ -526,77 +528,17 @@ where
     let executor_registration =
         ExecutorRegistration::new(environment_id.clone(), executor_registration_id.clone())
             .map(Arc::new);
-    let (mut websocket_sink, mut websocket_stream) = stream.split();
-    let (physical_outgoing_tx, mut physical_outgoing_rx) =
-        mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);
+    let (websocket_sink, mut websocket_stream) = stream.split();
+    let (physical_outgoing_tx, physical_outgoing_rx) = mpsc::channel::<Vec<u8>>(CHANNEL_CAPACITY);
     let (closed_stream_tx, mut closed_stream_rx) =
         mpsc::channel::<ClosedNoiseVirtualStream>(MAX_ACTIVE_NOISE_RELAY_STREAMS);
-    let (pong_tx, mut pong_rx) = mpsc::channel(1);
+    let (pong_tx, pong_rx) = mpsc::channel(1);
     // Use a separate writer so this loop never waits on the channel it drains.
-    let mut physical_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);
-        let mut pong_watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT);
-        let pong_deadline = tokio::time::sleep(WEBSOCKET_PONG_TIMEOUT);
-        tokio::pin!(pong_deadline);
-        loop {
-            let message = tokio::select! {
-                pong = pong_rx.recv() => {
-                    let Some(()) = pong else {
-                        break RendezvousDisconnectReason::LocalShutdown;
-                    };
-                    pong_watchdog.received_pong();
-                    continue;
-                }
-                _ = &mut pong_deadline, if pong_watchdog.deadline().is_some() => {
-                    match pong_rx.try_recv() {
-                        Ok(()) => {
-                            pong_watchdog.received_pong();
-                            continue;
-                        }
-                        Err(tokio::sync::mpsc::error::TryRecvError::Empty) => {
-                            break RendezvousDisconnectReason::PongTimeout;
-                        }
-                        Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {
-                            break RendezvousDisconnectReason::LocalShutdown;
-                        }
-                    }
-                }
-                _ = keepalive.tick(), if pong_watchdog.deadline().is_none() => {
-                    Message::Ping(Vec::new().into())
-                }
-                encoded = physical_outgoing_rx.recv() => {
-                    let Some(encoded) = encoded else {
-                        break RendezvousDisconnectReason::LocalShutdown;
-                    };
-                    Message::Binary(encoded.into())
-                }
-            };
-            let is_keepalive_ping = matches!(message, Message::Ping(_));
-            let write_deadline = pong_watchdog.write_deadline(tokio::time::Instant::now());
-            match tokio::time::timeout_at(write_deadline, websocket_sink.send(message)).await {
-                Ok(Ok(())) => {
-                    if is_keepalive_ping {
-                        pong_watchdog.ping_sent(tokio::time::Instant::now());
-                        if let Some(deadline) = pong_watchdog.deadline() {
-                            pong_deadline.as_mut().reset(deadline);
-                        }
-                    }
-                }
-                Ok(Err(error)) => {
-                    warn!("Noise multiplexed environment websocket write failed: {error}");
-                    break RendezvousDisconnectReason::WriteError;
-                }
-                Err(_) => {
-                    warn!("Noise multiplexed environment websocket write timed out");
-                    break RendezvousDisconnectReason::WriteError;
-                }
-            }
-        }
-    });
+    let mut physical_writer_task = AbortOnDropHandle::new(tokio::spawn(writer::run(
+        websocket_sink,
+        physical_outgoing_rx,
+        pong_rx,
+    )));
     let mut streams: HashMap<String, NoiseVirtualStream<H>> = HashMap::new();
     let mut pending_handshakes: HashMap<String, PendingHandshake> = HashMap::new();
     let mut validation_tasks: JoinSet<HarnessKeyValidationResult> = JoinSet::new();
@@ -1113,9 +1055,11 @@ mod tests {
         Ok(())
     }
 
-    #[tokio::test]
-    async fn multiplexed_environment_sends_keepalive() -> anyhow::Result<()> {
-        let (client_websocket, mut server_websocket) = websocket_pair().await?;
+    #[tokio::test(start_paused = true)]
+    async fn multiplexed_environment_sends_keepalive_and_cancels_blocked_writer()
+    -> anyhow::Result<()> {
+        let (client_websocket, control, mut outbound_rx) =
+            ControlledWebSocket::new(/*write_ready*/ true);
         let runtime_paths = crate::ExecServerRuntimeOptions::new(
             std::env::current_exe()?,
             /*codex_linux_sandbox_exe*/ None,
@@ -1130,10 +1074,17 @@ mod tests {
             AllowHarnessKeyValidator,
         ));
 
-        read_keepalive_ping(&mut server_websocket).await?;
+        assert!(matches!(outbound_rx.next().await, Some(Message::Ping(_))));
+        control.set_write_blocked();
+        control.send_inbound(Message::Pong(Vec::new().into()))?;
+        control.wait_for_blocked_write().await?;
 
         environment_task.abort();
         let _ = environment_task.await;
+        assert_eq!(
+            timeout(Duration::from_millis(1), outbound_rx.next()).await?,
+            None
+        );
         Ok(())
     }
 

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

@@ -0,0 +1,107 @@
+//! Keep executor writes bounded while processing pongs independently of write backpressure.
+
+use futures::Sink;
+use futures::SinkExt;
+use tokio::sync::mpsc;
+use tokio::time::Instant;
+use tokio_tungstenite::tungstenite::Message;
+use tracing::warn;
+
+use crate::connection::WEBSOCKET_KEEPALIVE_INTERVAL;
+use crate::relay::RendezvousDisconnectReason;
+use crate::websocket_pong_watchdog::WEBSOCKET_PONG_TIMEOUT;
+use crate::websocket_pong_watchdog::WebSocketPongWatchdog;
+
+pub(super) async fn run<S, E>(
+    mut sink: S,
+    mut outgoing_rx: mpsc::Receiver<Vec<u8>>,
+    mut pong_rx: mpsc::Receiver<()>,
+) -> RendezvousDisconnectReason
+where
+    S: Sink<Message, Error = E> + Unpin,
+    E: std::fmt::Display,
+{
+    let mut keepalive = tokio::time::interval_at(
+        Instant::now() + WEBSOCKET_KEEPALIVE_INTERVAL,
+        WEBSOCKET_KEEPALIVE_INTERVAL,
+    );
+    keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+    let mut watchdog = WebSocketPongWatchdog::new(WEBSOCKET_PONG_TIMEOUT);
+    loop {
+        let message = tokio::select! {
+            result = receive_pong(&mut watchdog, &mut pong_rx) => {
+                if let Err(reason) = result {
+                    return reason;
+                }
+                continue;
+            }
+            _ = keepalive.tick(), if watchdog.deadline().is_none() => {
+                Message::Ping(Vec::new().into())
+            }
+            encoded = outgoing_rx.recv() => {
+                let Some(encoded) = encoded else {
+                    return RendezvousDisconnectReason::LocalShutdown;
+                };
+                Message::Binary(encoded.into())
+            }
+        };
+        let is_keepalive_ping = matches!(message, Message::Ping(_));
+        // A pong clears only its own deadline, never the fixed write budget.
+        let write = tokio::time::timeout(WEBSOCKET_PONG_TIMEOUT, sink.send(message));
+        tokio::pin!(write);
+        loop {
+            tokio::select! {
+                // Keep the same send future alive when a pong arrives: cancelling
+                // and restarting a partially flushed send can duplicate a frame.
+                result = receive_pong(&mut watchdog, &mut pong_rx), if !is_keepalive_ping => {
+                    if let Err(reason) = result {
+                        return reason;
+                    }
+                }
+                result = &mut write => {
+                    match result {
+                        Ok(Ok(())) => {
+                            if is_keepalive_ping {
+                                // Pongs received during the ping flush stay queued
+                                // until the watchdog is armed here.
+                                watchdog.ping_sent(Instant::now());
+                            }
+                            break;
+                        }
+                        Ok(Err(error)) => {
+                            warn!("Noise multiplexed environment websocket write failed: {error}");
+                            return RendezvousDisconnectReason::WriteError;
+                        }
+                        Err(_) => {
+                            warn!("Noise multiplexed environment websocket write timed out");
+                            return RendezvousDisconnectReason::WriteError;
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
+async fn receive_pong(
+    watchdog: &mut WebSocketPongWatchdog,
+    pong_rx: &mut mpsc::Receiver<()>,
+) -> Result<(), RendezvousDisconnectReason> {
+    let deadline = watchdog.deadline();
+    tokio::select! {
+        biased;
+        // Honor a pong already read by the reader even if the timer is also ready.
+        pong = pong_rx.recv() => {
+            pong.ok_or(RendezvousDisconnectReason::LocalShutdown)?;
+            watchdog.received_pong();
+            Ok(())
+        }
+        _ = tokio::time::sleep_until(deadline.unwrap_or_else(Instant::now)), if deadline.is_some() => {
+            Err(RendezvousDisconnectReason::PongTimeout)
+        }
+    }
+}
+
+#[cfg(test)]
+#[path = "relay_writer_tests.rs"]
+mod tests;

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

@@ -0,0 +1,98 @@
+//! Exercise heartbeat deadlines while a single physical write is stalled.
+
+use std::convert::Infallible;
+use std::sync::Arc;
+use std::time::Duration;
+
+use pretty_assertions::assert_eq;
+use tokio::sync::Notify;
+use tokio::sync::mpsc;
+use tokio_util::task::AbortOnDropHandle;
+
+use super::*;
+
+#[tokio::test(start_paused = true)]
+async fn pong_during_write_clears_only_the_pong_deadline() {
+    for release_write in [true, false] {
+        let (outgoing_tx, outgoing_rx) = mpsc::channel(1);
+        let (pong_tx, pong_rx) = mpsc::channel(1);
+        let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
+        let release = Arc::new(Notify::new());
+        let completed = Arc::new(Notify::new());
+        let sink = futures::sink::unfold(
+            (observed_tx, Arc::clone(&release), Arc::clone(&completed)),
+            |(observed_tx, release, completed), message| async move {
+                let is_data = matches!(message, Message::Binary(_));
+                observed_tx.send(message).unwrap();
+                if is_data {
+                    release.notified().await;
+                    completed.notify_one();
+                }
+                Ok::<_, Infallible>((observed_tx, release, completed))
+            },
+        );
+        let task = AbortOnDropHandle::new(tokio::spawn(run(Box::pin(sink), outgoing_rx, pong_rx)));
+        assert!(matches!(observed_rx.recv().await, Some(Message::Ping(_))));
+        tokio::time::advance(WEBSOCKET_PONG_TIMEOUT * 3 / 4).await;
+        outgoing_tx.send(vec![42]).await.unwrap();
+        assert_eq!(
+            observed_rx.recv().await,
+            Some(Message::Binary(vec![42].into()))
+        );
+        pong_tx.send(()).await.unwrap();
+
+        // Let the old pong deadline expire while the write remains blocked.
+        tokio::time::advance(WEBSOCKET_PONG_TIMEOUT / 2).await;
+        tokio::task::yield_now().await;
+        assert!(!task.is_finished());
+        if release_write {
+            release.notify_one();
+            completed.notified().await;
+            drop(outgoing_tx);
+            drop(pong_tx);
+            assert_eq!(
+                task.await.unwrap(),
+                RendezvousDisconnectReason::LocalShutdown
+            );
+        } else {
+            // Receiving a pong must not extend the independent write budget.
+            tokio::time::advance(WEBSOCKET_PONG_TIMEOUT).await;
+            assert_eq!(task.await.unwrap(), RendezvousDisconnectReason::WriteError);
+        }
+    }
+}
+
+#[tokio::test(start_paused = true)]
+async fn missing_pong_times_out_during_a_blocked_write() {
+    let (outgoing_tx, outgoing_rx) = mpsc::channel(1);
+    let (_pong_tx, pong_rx) = mpsc::channel(1);
+    let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
+    let sink = futures::sink::unfold(observed_tx, |observed_tx, message| async move {
+        let is_data = matches!(message, Message::Binary(_));
+        observed_tx.send(message).unwrap();
+        if is_data {
+            std::future::pending::<()>().await;
+        }
+        Ok::<_, Infallible>(observed_tx)
+    });
+    let task = AbortOnDropHandle::new(tokio::spawn(run(Box::pin(sink), outgoing_rx, pong_rx)));
+    assert!(matches!(observed_rx.recv().await, Some(Message::Ping(_))));
+    tokio::time::advance(WEBSOCKET_PONG_TIMEOUT * 3 / 4).await;
+    outgoing_tx.send(vec![42]).await.unwrap();
+    assert_eq!(
+        observed_rx.recv().await,
+        Some(Message::Binary(vec![42].into()))
+    );
+    assert_eq!(task.await.unwrap(), RendezvousDisconnectReason::PongTimeout);
+}
+
+#[tokio::test(start_paused = true)]
+async fn queued_pong_wins_over_an_expired_deadline() {
+    let (pong_tx, mut pong_rx) = mpsc::channel(1);
+    let mut watchdog = WebSocketPongWatchdog::new(Duration::from_secs(1));
+    watchdog.ping_sent(Instant::now());
+    tokio::time::advance(Duration::from_secs(1)).await;
+    pong_tx.send(()).await.unwrap();
+    assert_eq!(receive_pong(&mut watchdog, &mut pong_rx).await, Ok(()));
+    assert_eq!(watchdog.deadline(), None);
+}

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

@@ -709,7 +709,7 @@ async fn run_remote_environment_connections<H: NoiseStreamHandler>(
     let identity = NoiseChannelIdentity::generate().map_err(|error| {
         ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}"))
     })?;
-    let mut backoff = Duration::from_secs(1);
+    let mut backoff = reconnect_backoff::ReconnectBackoff::new();
     let mut response = client
         .register_environment_with_retry(&config.environment_id, &identity.public_key())
         .await?;
@@ -727,7 +727,7 @@ async fn run_remote_environment_connections<H: NoiseStreamHandler>(
         .await
         {
             Ok(websocket) => {
-                backoff = Duration::from_secs(1);
+                let connected_at = tokio::time::Instant::now();
                 let executor_registration_id = response.executor_registration_id.clone();
                 info!(
                     noise_event = "rendezvous_connection",
@@ -749,6 +749,7 @@ async fn run_remote_environment_connections<H: NoiseStreamHandler>(
                     },
                 )
                 .await;
+                backoff.connection_closed(connected_at.elapsed());
                 info!(
                     noise_event = "rendezvous_connection",
                     noise_outcome = "disconnected",
@@ -794,12 +795,7 @@ async fn run_remote_environment_connections<H: NoiseStreamHandler>(
             }
         }
 
-        sleep(reconnect_backoff::reconnect_delay(
-            backoff,
-            uuid::Uuid::new_v4().as_u64_pair().1,
-        ))
-        .await;
-        backoff = (backoff * 2).min(Duration::from_secs(30));
+        sleep(backoff.next_delay(uuid::Uuid::new_v4().as_u64_pair().1)).await;
     }
 }
 
@@ -820,7 +816,7 @@ async fn connect_rendezvous(
     http_client_factory: &HttpClientFactory,
 ) -> Result<WebSocketConnection, tokio_tungstenite::tungstenite::Error> {
     let started_at = Instant::now();
-    let result = async {
+    let result = tokio::time::timeout(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, async {
         let mut request = url.into_client_request()?;
         request.headers_mut().extend(current_rendezvous_headers());
         request
@@ -836,8 +832,16 @@ async fn connect_rendezvous(
             .connect(request, noise_relay_websocket_config())
             .await
             .map(|(websocket, _)| websocket)
-    }
-    .await;
+    })
+    .await
+    .unwrap_or_else(|_| {
+        Err(tokio_tungstenite::tungstenite::Error::Io(
+            std::io::Error::new(
+                std::io::ErrorKind::TimedOut,
+                "rendezvous websocket connection timed out",
+            ),
+        ))
+    });
     let result_name = if result.is_ok() { "success" } else { "error" };
     tracing::Span::current().record("result", result_name);
     telemetry.remote_rendezvous_completed(result_name, started_at.elapsed());

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

@@ -186,22 +186,32 @@ async fn reconnect_reuses_registration_until_url_is_rejected() -> Result<()> {
         static_registry_auth_provider(),
         HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
     )?;
-    let environment_task = tokio::spawn(run_remote_environment(
-        config,
-        ExecServerRuntimeOptions::new(
-            std::env::current_exe()?,
-            /*codex_linux_sandbox_exe*/ None,
-        )?,
-    ));
+    let environment_task =
+        tokio_util::task::AbortOnDropHandle::new(tokio::spawn(run_remote_environment(
+            config,
+            ExecServerRuntimeOptions::new(
+                std::env::current_exe()?,
+                /*codex_linux_sandbox_exe*/ None,
+            )?,
+        )));
 
     let (first_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??;
     let mut first_websocket = accept_async(first_socket).await?;
     first_websocket.close(None).await?;
 
-    // An ordinary disconnect retries the same URL without registering again.
-    let (mut rejected_socket, _peer_addr) =
+    // An ordinary disconnect and a stalled upgrade both retry without registering again.
+    let (mut stalled_socket, _peer_addr) =
         timeout(Duration::from_secs(5), listener.accept()).await??;
     let mut request = [0u8; 4096];
+    assert!(stalled_socket.read(&mut request).await? > 0);
+    tokio::time::pause();
+    tokio::time::advance(DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT).await;
+    tokio::time::resume();
+
+    // Keep the stalled socket open until the executor starts its next attempt.
+    let (mut rejected_socket, _peer_addr) =
+        timeout(Duration::from_secs(5), listener.accept()).await??;
+    drop(stalled_socket);
     let _ = rejected_socket.read(&mut request).await?;
     rejected_socket
         .write_all(b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\n\r\n")
@@ -209,8 +219,8 @@ async fn reconnect_reuses_registration_until_url_is_rejected() -> Result<()> {
     rejected_socket.shutdown().await?;
 
     // The 4xx response discards the old registration before this attempt.
-    let (third_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??;
-    let _third_websocket = accept_async(third_socket).await?;
+    let (final_socket, _peer_addr) = timeout(Duration::from_secs(5), listener.accept()).await??;
+    let _final_websocket = accept_async(final_socket).await?;
     registry.verify().await;
 
     environment_task.abort();

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

@@ -1,8 +1,37 @@
-//! Spread executor reconnects after a shared outage without exceeding the existing backoff cap.
+//! Spread executor reconnects and retain backoff until a connection has been stable.
 
 use std::time::Duration;
 
-pub(super) fn reconnect_delay(backoff: Duration, random_sample: u64) -> Duration {
+const INITIAL_BACKOFF: Duration = Duration::from_secs(1);
+const MAX_BACKOFF: Duration = Duration::from_secs(30);
+// Longer than the first keepalive interval plus its pong timeout (30s + 60s).
+const STABLE_CONNECTION_DURATION: Duration = Duration::from_secs(120);
+
+pub(super) struct ReconnectBackoff {
+    delay: Duration,
+}
+
+impl ReconnectBackoff {
+    pub(super) fn new() -> Self {
+        Self {
+            delay: INITIAL_BACKOFF,
+        }
+    }
+
+    pub(super) fn connection_closed(&mut self, connected_for: Duration) {
+        if connected_for >= STABLE_CONNECTION_DURATION {
+            self.delay = INITIAL_BACKOFF;
+        }
+    }
+
+    pub(super) fn next_delay(&mut self, random_sample: u64) -> Duration {
+        let delay = reconnect_delay(self.delay, random_sample);
+        self.delay = (self.delay * 2).min(MAX_BACKOFF);
+        delay
+    }
+}
+
+fn reconnect_delay(backoff: Duration, random_sample: u64) -> Duration {
     let upper_ms = backoff.as_millis() as u64;
     let lower_ms = upper_ms / 2;
     Duration::from_millis(lower_ms + random_sample % (upper_ms - lower_ms + 1))

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

@@ -4,6 +4,8 @@ use std::time::Duration;
 
 use pretty_assertions::assert_eq;
 
+use super::ReconnectBackoff;
+use super::STABLE_CONNECTION_DURATION;
 use super::reconnect_delay;
 
 #[test]
@@ -23,3 +25,24 @@ fn reconnect_jitter_stays_within_each_backoff_window() {
         );
     }
 }
+
+#[test]
+fn flapping_connections_retain_backoff_until_a_stable_session() {
+    let mut backoff = ReconnectBackoff::new();
+    for seconds in [1, 2, 4, 8, 16, 30, 30] {
+        backoff.connection_closed(STABLE_CONNECTION_DURATION - Duration::from_millis(1));
+        assert_eq!(
+            backoff.next_delay(/*random_sample*/ 0),
+            Duration::from_millis(seconds * 500)
+        );
+    }
+    backoff.connection_closed(STABLE_CONNECTION_DURATION);
+    assert_eq!(
+        backoff.next_delay(/*random_sample*/ 0),
+        Duration::from_millis(500)
+    );
+    assert_eq!(
+        backoff.next_delay(/*random_sample*/ 0),
+        Duration::from_secs(1)
+    );
+}