Time out stalled code-mode host requests (#36830)

## Why

Code-mode `wait` and `terminate` requests can remain pending when the host
transport stalls.

## What changed

- Add a 60-second transport allowance to the runtime timeout for `wait`, and
  apply the same transport deadline to `terminate`.
- Return a model-visible timeout error and invalidate the connection when the
  deadline expires, so the next execution reconnects to the host.

## Testing

- Cover queued `wait` and `terminate` requests that exceed their deadlines.
- Verify that a timed-out `wait` reports the error and reconnects on the next
  code-mode execution.

GitOrigin-RevId: 5d772e5a6f3793aa8865a1160639b851fd4824fb

codex-rs/code-mode/src/remote_session/connection.rs

@@ -1,4 +1,5 @@
 use std::fmt;
+use std::future::Future;
 use std::io;
 use std::path::Path;
 use std::path::PathBuf;
@@ -62,6 +63,8 @@ mod transport;
 
 const IPC_CHANNEL_CAPACITY: usize = 128;
 const HOST_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(10);
+// TODO(anp) make this timeout configurable if 60 seconds is insufficient.
+const DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT: Duration = Duration::from_secs(60);
 const MAX_WEBSOCKET_FRAME_BYTES: usize = MAX_FRAME_BYTES + std::mem::size_of::<u32>();
 // Host spawn errors become model-visible tool output. Bound configured paths
 // while preserving the executable-bearing suffix needed to diagnose failures.
@@ -526,16 +529,23 @@ impl Connection {
         session: RemoteSession,
         request: WaitRequest,
     ) -> Result<WaitOutcome, String> {
+        // Account for the runtime's one-second yield grace separately from transport.
+        let runtime_timeout =
+            Duration::from_millis(request.yield_time_ms).saturating_add(Duration::from_secs(1));
         let cancellation = CallerCancellation::new();
         let (response_tx, response_rx) = oneshot::channel();
-        self.send(DriverCommand::Wait {
-            session,
-            request,
-            caller_cancellation: cancellation.token(),
-            response_tx,
-        })
-        .await?;
-        let result = self.receive(response_rx).await;
+        let result = self
+            .with_transport_deadline(runtime_timeout, "wait", async {
+                self.send(DriverCommand::Wait {
+                    session,
+                    request,
+                    caller_cancellation: cancellation.token(),
+                    response_tx,
+                })
+                .await?;
+                self.receive(response_rx).await
+            })
+            .await;
         cancellation.disarm();
         result
     }
@@ -546,13 +556,16 @@ impl Connection {
         cell_id: CellId,
     ) -> Result<WaitOutcome, String> {
         let (response_tx, response_rx) = oneshot::channel();
-        self.send(DriverCommand::Terminate {
-            session,
-            cell_id,
-            response_tx,
+        self.with_transport_deadline(Duration::ZERO, "terminate", async {
+            self.send(DriverCommand::Terminate {
+                session,
+                cell_id,
+                response_tx,
+            })
+            .await?;
+            self.receive(response_rx).await
         })
-        .await?;
-        self.receive(response_rx).await
+        .await
     }
 
     pub(super) async fn shutdown_session(&self, session: RemoteSession) -> Result<(), String> {
@@ -565,6 +578,26 @@ impl Connection {
         self.receive(response_rx).await
     }
 
+    async fn with_transport_deadline<T>(
+        &self,
+        runtime_timeout: Duration,
+        request_type: &str,
+        request: impl Future<Output = Result<T, String>>,
+    ) -> Result<T, String> {
+        let deadline = runtime_timeout.saturating_add(DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT);
+        match tokio::time::timeout(deadline, request).await {
+            Ok(result) => result,
+            Err(_) => {
+                warn!(request_type, "code-mode host request exceeded its deadline");
+                let reason =
+                    format!("code-mode host timed out waiting for {request_type} response");
+                mark_connection_dead(&self.alive, &self.failure, reason.clone());
+                self.cancellation.cancel();
+                Err(reason)
+            }
+        }
+    }
+
     async fn send(&self, command: DriverCommand) -> Result<(), String> {
         if !self.is_alive() {
             return Err(self.failure_message());

codex-rs/code-mode/src/remote_session/connection/driver_tests.rs

@@ -32,6 +32,8 @@ use tokio::sync::mpsc;
 use tokio::sync::oneshot;
 use tokio_util::sync::CancellationToken;
 
+use super::super::Connection;
+use super::super::DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT;
 use super::ConnectionDriver;
 use super::DriverCommand;
 use super::DriverEvent;
@@ -46,6 +48,7 @@ struct DriverHarness {
     outgoing_rx: mpsc::Receiver<codex_code_mode_protocol::host::EncodedFrame>,
     cancellation: CancellationToken,
     alive: Arc<AtomicBool>,
+    failure: Arc<StdMutex<Option<String>>>,
     driver_task: tokio::task::JoinHandle<()>,
 }
 
@@ -56,6 +59,7 @@ impl DriverHarness {
         let (outgoing_tx, outgoing_rx) = mpsc::channel(/*max_capacity*/ 16);
         let cancellation = CancellationToken::new();
         let alive = Arc::new(AtomicBool::new(true));
+        let failure = Arc::new(StdMutex::new(None));
         let (driver, execute_claim_tx) = ConnectionDriver::new(
             command_rx,
             event_rx,
@@ -63,7 +67,7 @@ impl DriverHarness {
             outgoing_tx,
             DriverLifecycle {
                 alive: Arc::clone(&alive),
-                failure: Arc::new(StdMutex::new(None)),
+                failure: Arc::clone(&failure),
                 cancellation: cancellation.clone(),
             },
         );
@@ -75,6 +79,7 @@ impl DriverHarness {
             outgoing_rx,
             cancellation,
             alive,
+            failure,
             driver_task,
         }
     }
@@ -647,6 +652,7 @@ async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bul
     let (bulk_tx, mut bulk_rx) = mpsc::channel(MAX_PENDING_DELEGATE_CALLS);
     let cancellation = CancellationToken::new();
     let alive = Arc::new(AtomicBool::new(true));
+    let failure = Arc::new(StdMutex::new(None));
     let (driver, execute_claim_tx) = ConnectionDriver::new(
         command_rx,
         event_rx,
@@ -654,7 +660,7 @@ async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bul
         outgoing_tx,
         DriverLifecycle {
             alive: Arc::clone(&alive),
-            failure: Arc::new(StdMutex::new(None)),
+            failure: Arc::clone(&failure),
             cancellation: cancellation.clone(),
         },
     );
@@ -666,6 +672,7 @@ async fn concurrent_large_delegate_results_do_not_disconnect_a_backpressured_bul
         outgoing_rx,
         cancellation,
         alive,
+        failure,
         driver_task,
     };
     let session = remote_session();
@@ -1328,6 +1335,90 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() {
     );
 }
 
+/// A stalled remote wait must expire even when its request never reaches the host.
+#[tokio::test(start_paused = true)]
+async fn queued_remote_wait_times_out_and_invalidates_the_connection() {
+    let mut harness = DriverHarness::start();
+    let session = remote_session();
+    harness
+        .open(session.clone(), Arc::new(RecordingDelegate::default()))
+        .await;
+    let _started = harness
+        .start_cell(session.clone(), /*request_id*/ 2, "1")
+        .await;
+    let connection = Connection {
+        command_tx: harness.command_tx.clone(),
+        execute_claim_tx: harness.execute_claim_tx.clone(),
+        alive: Arc::clone(&harness.alive),
+        failure: Arc::clone(&harness.failure),
+        cancellation: harness.cancellation.clone(),
+    };
+    let response = tokio::spawn(async move {
+        let result = connection
+            .wait(
+                session,
+                WaitRequest {
+                    cell_id: CellId::new("1".to_string()),
+                    yield_time_ms: 1,
+                },
+            )
+            .await;
+        (connection, result)
+    });
+    while harness.outgoing_rx.is_empty() {
+        tokio::task::yield_now().await;
+    }
+    tokio::time::advance(DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT + Duration::from_secs(2)).await;
+
+    let (connection, result) = response.await.expect("wait task");
+    assert_eq!(
+        result,
+        Err("code-mode host timed out waiting for wait response".to_string())
+    );
+    assert!(!harness.alive.load(Ordering::Acquire));
+    assert!(harness.cancellation.is_cancelled());
+    drop(connection);
+}
+
+/// A stalled termination must expire and invalidate the same connection as a stalled wait.
+#[tokio::test(start_paused = true)]
+async fn queued_remote_termination_times_out_and_invalidates_the_connection() {
+    let mut harness = DriverHarness::start();
+    let session = remote_session();
+    harness
+        .open(session.clone(), Arc::new(RecordingDelegate::default()))
+        .await;
+    let _started = harness
+        .start_cell(session.clone(), /*request_id*/ 2, "1")
+        .await;
+    let connection = Connection {
+        command_tx: harness.command_tx.clone(),
+        execute_claim_tx: harness.execute_claim_tx.clone(),
+        alive: Arc::clone(&harness.alive),
+        failure: Arc::clone(&harness.failure),
+        cancellation: harness.cancellation.clone(),
+    };
+    let response = tokio::spawn(async move {
+        let result = connection
+            .terminate(session, CellId::new("1".to_string()))
+            .await;
+        (connection, result)
+    });
+    while harness.outgoing_rx.is_empty() {
+        tokio::task::yield_now().await;
+    }
+    tokio::time::advance(DEFAULT_HOST_WAIT_TRANSPORT_TIMEOUT + Duration::from_secs(1)).await;
+
+    let (connection, result) = response.await.expect("termination task");
+    assert_eq!(
+        result,
+        Err("code-mode host timed out waiting for terminate response".to_string())
+    );
+    assert!(!harness.alive.load(Ordering::Acquire));
+    assert!(harness.cancellation.is_cancelled());
+    drop(connection);
+}
+
 #[tokio::test]
 async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
     let mut harness = DriverHarness::start();