Defer environment disconnect notifications until recovery fails (#55401)

## Why

A temporary exec-server transport loss could emit a thread environment disconnect notification while the session was still recovering.

## What changed

- Add `EnvironmentConnectionState::Recovering` for interrupted transports and pending session recovery.
- Suppress thread lifecycle notifications during recovery, emit a connected notification after successful recovery, and report disconnection when recovery fails.
- Ignore recovery callbacks from retired clients so they cannot overwrite the connection state during refresh.

## Testing

Add coverage for fast and near-deadline recovery, recovery timeout, rejected session resume, and retired-client callbacks. Extend app-server coverage to verify one lifecycle notification per selected thread and usable `environment/info` after recovery.

GitOrigin-RevId: d6dce5cc5264af29b6f993daa209b93102122e40

codex-rs/app-server/tests/suite/v2/environment_add.rs

@@ -5,20 +5,25 @@ use app_test_support::TestAppServer;
 use app_test_support::to_response;
 use codex_app_server_protocol::EnvironmentAddResponse;
 use codex_app_server_protocol::EnvironmentConnectionNotification;
+use codex_app_server_protocol::EnvironmentInfoResponse;
+use codex_app_server_protocol::EnvironmentShellInfo;
 use codex_app_server_protocol::JSONRPCResponse;
 use codex_app_server_protocol::RequestId;
 use codex_app_server_protocol::ThreadStartParams;
 use codex_app_server_protocol::ThreadStartResponse;
 use codex_app_server_protocol::TurnEnvironmentParams;
+use futures::SinkExt;
 use pretty_assertions::assert_eq;
 use serde_json::json;
 use tempfile::TempDir;
 use tokio::io::AsyncReadExt;
 use tokio::net::TcpListener;
-use tokio::sync::oneshot;
+use tokio::sync::mpsc;
 use tokio::time::timeout;
+use tokio_tungstenite::accept_async;
+use tokio_tungstenite::tungstenite::Message;
 
-use super::exec_server_test_support::accept_exec_server_environment;
+use super::exec_server_test_support::read_exec_server_json;
 
 const RPC_TIMEOUT: Duration = Duration::from_secs(10);
 const CONNECTION_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
@@ -115,15 +120,60 @@ async fn selected_environment_emits_connection_lifecycle_notifications() -> Resu
     .await??;
     let ThreadStartResponse { thread, .. } = to_response(response)?;
 
-    let (disconnect_tx, disconnect_rx) = oneshot::channel();
+    let (disconnect_tx, mut disconnect_rx) = mpsc::channel(/*buffer*/ 1);
     let exec_server = tokio::spawn(async move {
-        let mut websocket = accept_exec_server_environment(
-            listener,
-            json!({"shell": {"name": "zsh", "path": "/bin/zsh"}}),
-        )
-        .await?;
-        disconnect_rx.await?;
-        websocket.close(None).await?;
+        // Recover the first disconnect, then reject the next resume. This
+        // process boundary cannot share Tokio's paused clock; deadline
+        // exhaustion is covered by core's connection-event tests instead.
+        for attempt in 0..3 {
+            let mut websocket = accept_async(listener.accept().await?.0).await?;
+            let initialize = read_exec_server_json(&mut websocket).await?;
+            assert_eq!(initialize["method"], "initialize");
+            assert_eq!(
+                initialize["params"]["resumeSessionId"],
+                json!((attempt > 0).then_some("test-session"))
+            );
+            let response = if attempt == 2 {
+                json!({"id": initialize["id"], "error": {
+                    "code": -32000, "message": "session no longer exists"
+                }})
+            } else {
+                json!({"id": initialize["id"], "result": {"sessionId": "test-session"}})
+            };
+            websocket
+                .send(Message::Text(response.to_string().into()))
+                .await?;
+            if attempt == 2 {
+                break;
+            }
+            assert_eq!(
+                read_exec_server_json(&mut websocket).await?["method"],
+                "initialized"
+            );
+            loop {
+                let request = tokio::select! {
+                    command = disconnect_rx.recv() => {
+                        command.expect("disconnect command");
+                        break;
+                    }
+                    request = read_exec_server_json(&mut websocket) => request?,
+                };
+                let response = if request["method"] == "environment/info" {
+                    json!({"id": request["id"], "result": {
+                        "shell": {"name": "zsh", "path": "/bin/zsh"}
+                    }})
+                } else {
+                    // Startup discovery can probe files on the resumed transport.
+                    json!({"id": request["id"], "error": {
+                        "code": -32601, "message": "method not supported by test executor"
+                    }})
+                };
+                websocket
+                    .send(Message::Text(response.to_string().into()))
+                    .await?;
+            }
+            websocket.close(None).await?;
+        }
         Ok::<_, anyhow::Error>(())
     });
 
@@ -158,23 +208,6 @@ async fn selected_environment_emits_connection_lifecycle_notifications() -> Resu
         ..
     } = to_response(response)?;
 
-    disconnect_tx
-        .send(())
-        .map_err(|_| anyhow::anyhow!("exec-server disconnect receiver closed"))?;
-    let mut disconnected = Vec::new();
-    for _ in 0..2 {
-        let notification = timeout(
-            RPC_TIMEOUT,
-            app_server.read_stream_until_notification_message("thread/environment/disconnected"),
-        )
-        .await??;
-        disconnected.push(serde_json::from_value::<EnvironmentConnectionNotification>(
-            notification
-                .params
-                .expect("disconnected notification params"),
-        )?);
-    }
-    disconnected.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
     let mut expected = vec![
         EnvironmentConnectionNotification {
             thread_id: thread.id,
@@ -186,7 +219,6 @@ async fn selected_environment_emits_connection_lifecycle_notifications() -> Resu
         },
     ];
     expected.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
-    assert_eq!(disconnected, expected);
     assert!(
         !app_server
             .pending_notification_methods()
@@ -195,6 +227,63 @@ async fn selected_environment_emits_connection_lifecycle_notifications() -> Resu
         "connection state should not be replayed when a thread starts"
     );
 
+    for method in [
+        "thread/environment/connected",
+        "thread/environment/disconnected",
+    ] {
+        disconnect_tx.send(()).await?;
+        let mut notifications = Vec::new();
+        for _ in 0..2 {
+            notifications.push(
+                timeout(
+                    RPC_TIMEOUT,
+                    app_server.read_notification::<EnvironmentConnectionNotification>(method),
+                )
+                .await??,
+            );
+        }
+        notifications.sort_by(|left, right| left.thread_id.cmp(&right.thread_id));
+        assert_eq!(notifications, expected);
+        if method == "thread/environment/connected" {
+            let request_id = app_server
+                .send_raw_request(
+                    "environment/info",
+                    Some(json!({"environmentId": "remote-a"})),
+                )
+                .await?;
+            let response = timeout(
+                RPC_TIMEOUT,
+                app_server.read_stream_until_response_message(RequestId::Integer(request_id)),
+            )
+            .await??;
+            assert_eq!(
+                to_response::<EnvironmentInfoResponse>(response)?,
+                EnvironmentInfoResponse {
+                    shell: EnvironmentShellInfo {
+                        name: "zsh".to_string(),
+                        path: "/bin/zsh".to_string()
+                    },
+                    cwd: None,
+                }
+            );
+        }
+        assert!(
+            timeout(
+                Duration::from_millis(100),
+                app_server.read_stream_until_matching_notification(
+                    "unexpected lifecycle notification",
+                    |notification| matches!(
+                        notification.method.as_str(),
+                        "thread/environment/connected" | "thread/environment/disconnected"
+                    ),
+                ),
+            )
+            .await
+            .is_err(),
+            "expected exactly one {method} event per selected thread"
+        );
+    }
+
     timeout(RPC_TIMEOUT, exec_server).await???;
     Ok(())
 }

codex-rs/core/src/environment_connection_tests.rs

@@ -0,0 +1,180 @@
+//! Exercise lifecycle forwarding against real recovering exec-server transports.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use codex_exec_server::Environment;
+use codex_exec_server::EnvironmentConnectionState;
+use codex_protocol::protocol::EventMsg;
+use futures::SinkExt;
+use futures::StreamExt;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use tokio::net::TcpListener;
+use tokio::net::TcpStream;
+use tokio::sync::mpsc;
+use tokio::sync::oneshot;
+use tokio_tungstenite::WebSocketStream;
+use tokio_tungstenite::accept_async;
+use tokio_tungstenite::tungstenite::Message;
+use tokio_util::task::AbortOnDropHandle;
+
+use super::ThreadEnvironments;
+use super::tests::read_websocket_json;
+
+// Keep paused time from advancing during socket I/O; only the test advances it.
+fn freeze_clock() -> std::sync::mpsc::Sender<()> {
+    tokio::time::pause();
+    let (guard, dropped) = std::sync::mpsc::channel();
+    tokio::task::spawn_blocking(move || {
+        let _ = dropped.recv();
+    });
+    guard
+}
+
+async fn initialize(
+    socket: &mut WebSocketStream<TcpStream>,
+    resume_session_id: Option<&str>,
+    result: Result<(), i64>,
+) {
+    let request = read_websocket_json(socket).await;
+    assert_eq!(request["method"], "initialize");
+    assert_eq!(
+        request["params"]["resumeSessionId"],
+        json!(resume_session_id)
+    );
+    let response = match result {
+        Ok(()) => json!({"id": request["id"], "result": {"sessionId": "session"}}),
+        Err(code) => json!({"id": request["id"], "error": {
+            "code": code, "message": "cannot resume"
+        }}),
+    };
+    socket
+        .send(Message::Text(response.to_string().into()))
+        .await
+        .unwrap();
+    if result.is_ok() {
+        assert_eq!(read_websocket_json(socket).await["method"], "initialized");
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq)]
+enum Recovery {
+    Succeeds,
+    Rejected,
+    TimesOut,
+}
+
+#[test_case::test_case(0, Recovery::Succeeds; "fast")]
+#[test_case::test_case(3, Recovery::Succeeds; "near_deadline")]
+#[test_case::test_case(0, Recovery::TimesOut; "retry_exhaustion")]
+#[test_case::test_case(0, Recovery::Rejected; "nonretryable_failure")]
+#[tokio::test]
+async fn lifecycle_waits_for_recovery(retries: usize, recovery: Recovery) {
+    let _clock = freeze_clock();
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let environment = Arc::new(
+        Environment::create_for_tests(Some(format!("ws://{}", listener.local_addr().unwrap())))
+            .unwrap(),
+    );
+    let (events, received) = async_channel::unbounded();
+    let _forwarding = ThreadEnvironments::spawn_connection_event_listener(
+        &environment,
+        "remote".to_owned(),
+        events,
+    )
+    .unwrap();
+    let mut raw = environment.subscribe_connection_state().unwrap();
+    let (disconnect, disconnected) = oneshot::channel();
+    let (resume, resumed) = oneshot::channel();
+    let (ready, ready_rx) = oneshot::channel();
+    let (retrying, mut retried) = mpsc::unbounded_channel();
+    let (response_sent, response_rx) = oneshot::channel();
+    let (finish, finished) = oneshot::channel();
+    let server = AbortOnDropHandle::new(tokio::spawn(async move {
+        let mut socket = accept_async(listener.accept().await.unwrap().0)
+            .await
+            .unwrap();
+        initialize(&mut socket, /*resume_session_id*/ None, Ok(())).await;
+        disconnected.await.unwrap();
+        socket.close(None).await.unwrap();
+        for _ in 0..retries {
+            let mut socket = accept_async(listener.accept().await.unwrap().0)
+                .await
+                .unwrap();
+            initialize(&mut socket, Some("session"), Err(-32010)).await;
+            // Wait for the client to process the rejection before advancing time.
+            while let Some(Ok(message)) = socket.next().await {
+                if matches!(message, Message::Close(_)) {
+                    break;
+                }
+            }
+            retrying.send(()).unwrap();
+        }
+        let mut socket = accept_async(listener.accept().await.unwrap().0)
+            .await
+            .unwrap();
+        ready.send(()).unwrap();
+        resumed.await.unwrap();
+        let result = match recovery {
+            Recovery::Succeeds => Ok(()),
+            Recovery::Rejected => Err(-32000),
+            Recovery::TimesOut => Err(-32010),
+        };
+        initialize(&mut socket, Some("session"), result).await;
+        let _ = response_sent.send(());
+        finished.await.unwrap();
+    }));
+    environment.wait_until_ready().await.unwrap();
+    assert!(matches!(
+        received.recv().await.unwrap().msg,
+        EventMsg::EnvironmentConnected(_)
+    ));
+    raw.borrow_and_update();
+    let disconnected_at = tokio::time::Instant::now();
+    disconnect.send(()).unwrap();
+    for _ in 0..retries {
+        retried.recv().await.unwrap();
+        assert_eq!(
+            *raw.borrow_and_update(),
+            EnvironmentConnectionState::Recovering
+        );
+        assert!(received.try_recv().is_err());
+        tokio::time::advance(Duration::from_secs(8)).await;
+    }
+    ready_rx.await.unwrap();
+    assert_eq!(
+        disconnected_at.elapsed(),
+        Duration::from_secs(8 * retries as u64)
+    );
+    raw.wait_for(|state| *state == EnvironmentConnectionState::Recovering)
+        .await
+        .unwrap();
+    assert_eq!(
+        *raw.borrow_and_update(),
+        EnvironmentConnectionState::Recovering
+    );
+    assert!(received.try_recv().is_err());
+    resume.send(()).unwrap();
+    if recovery == Recovery::TimesOut {
+        response_rx.await.unwrap();
+        assert!(received.try_recv().is_err());
+        tokio::time::advance(Duration::from_secs(26)).await;
+    }
+    let event = received.recv().await.unwrap().msg;
+    match recovery {
+        Recovery::Succeeds => {
+            assert!(matches!(event, EventMsg::EnvironmentConnected(_)));
+        }
+        Recovery::Rejected | Recovery::TimesOut => {
+            assert!(matches!(event, EventMsg::EnvironmentDisconnected(_)));
+            assert_eq!(
+                *raw.borrow_and_update(),
+                EnvironmentConnectionState::Disconnected
+            );
+        }
+    }
+    assert!(received.try_recv().is_err());
+    finish.send(()).unwrap();
+    server.await.unwrap();
+}

codex-rs/core/src/environment_selection.rs

@@ -867,6 +867,7 @@ impl ThreadEnvironments {
                     }
                 };
                 let msg = match state {
+                    EnvironmentConnectionState::Recovering => continue,
                     EnvironmentConnectionState::Connected => {
                         EventMsg::EnvironmentConnected(EnvironmentConnectionEvent {
                             environment_id: environment_id.clone(),
@@ -1426,6 +1427,10 @@ impl TurnEnvironmentSnapshot {
 #[path = "environment_selection_snapshot_tests.rs"]
 mod snapshot_tests;
 
+#[cfg(test)]
+#[path = "environment_connection_tests.rs"]
+mod connection_tests;
+
 #[cfg(test)]
 mod tests {
     use std::time::Duration;
@@ -1504,7 +1509,7 @@ mod tests {
         .expect("runtime paths")
     }
 
-    async fn read_websocket_json(websocket: &mut WebSocketStream<TcpStream>) -> Value {
+    pub(super) async fn read_websocket_json(websocket: &mut WebSocketStream<TcpStream>) -> Value {
         loop {
             match timeout(std::time::Duration::from_secs(5), websocket.next())
                 .await

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

@@ -327,9 +327,10 @@ impl ConnectionState {
             ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => {
                 EnvironmentConnectionState::Connected
             }
-            ConnectionStatus::Connected(_)
-            | ConnectionStatus::Recovering
-            | ConnectionStatus::Failed(_) => EnvironmentConnectionState::Disconnected,
+            ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => {
+                EnvironmentConnectionState::Recovering
+            }
+            ConnectionStatus::Failed(_) => EnvironmentConnectionState::Disconnected,
         };
         let _ = self
             .environment_connection_state_tx

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

@@ -336,7 +336,7 @@ impl Inner {
                 .unwrap_or_else(std::sync::PoisonError::into_inner);
             match &connection.status {
                 ConnectionStatus::Connected(current)
-                    if Arc::ptr_eq(current, &failed_rpc_client) =>
+                    if Arc::ptr_eq(current, &failed_rpc_client) && !self.retired.is_cancelled() =>
                 {
                     connection.set_status(ConnectionStatus::Recovering);
                     true

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

@@ -17,6 +17,7 @@ use tokio::task::JoinSet;
 use tokio_util::task::AbortOnDropHandle;
 
 use super::*;
+use crate::EnvironmentConnectionState;
 use crate::ExecServerRuntimeOptions;
 use crate::NoiseChannelIdentity;
 use crate::ProcessId;
@@ -263,12 +264,25 @@ async fn refresh_cancels_old_recovery_and_connects_without_resuming_old_session(
     assert!(futures::poll!(refresh.as_mut()).is_pending());
     assert!(original.inner.retired.is_cancelled());
     assert!(original.is_disconnected());
+    let mut state = client.subscribe_connection_state();
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Disconnected
+    );
+    {
+        let mut connection = original.inner.connection.lock().unwrap();
+        connection.set_status(ConnectionStatus::Recovering);
+        connection.set_status(ConnectionStatus::Failed(
+            super::super::ConnectionFailure::Disconnected("late retired failure".to_owned()),
+        ));
+    }
+    assert!(!state.has_changed()?);
     assert_eq!(started.elapsed(), Duration::ZERO);
     refresh.await?;
     let replacement = client.get().await?;
     assert_ne!(original.session_id(), replacement.session_id());
     assert_eq!(
-        *client.environment_connection_state_tx.borrow(),
+        *client.subscribe_connection_state().borrow(),
         EnvironmentConnectionState::Connected
     );
     assert!(matches!(
@@ -279,6 +293,21 @@ async fn refresh_cancels_old_recovery_and_connects_without_resuming_old_session(
     let _ = blocked_recovery.send(());
     tokio::task::yield_now().await;
     assert!(Arc::ptr_eq(&client.get().await?.inner, &replacement.inner));
+    let mut state = client.subscribe_connection_state();
+    original
+        .inner
+        .connection
+        .lock()
+        .unwrap()
+        .set_status(ConnectionStatus::Failed(
+            super::super::ConnectionFailure::Disconnected("late old failure".to_owned()),
+        ));
+    assert!(!state.has_changed()?);
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Connected
+    );
+    assert!(matches!(client.readiness_result(), Some(Ok(()))));
     replacement.environment_status().await?;
     Ok(())
 }
@@ -588,7 +617,14 @@ async fn refresh_cancels_a_stalled_initial_lookup() -> Result<()> {
     assert!(futures::poll!(initial.as_mut()).is_pending());
     *registry.target.lock().unwrap() = new.target.clone();
     client.refresh_connection().await?;
+    let mut state = client.subscribe_connection_state();
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Connected
+    );
+    // Finish the cancelled attempt only after its replacement has connected.
     assert!(initial.await.is_err());
+    assert!(!state.has_changed()?);
     let _ = release.send(());
     client.get().await?.environment_status().await?;
     Ok(())
@@ -613,7 +649,7 @@ async fn refresh_cancels_a_stalled_noise_handshake() -> Result<()> {
     validator.handshake.unwrap().notify_one();
     client.get().await?.environment_status().await?;
     assert_eq!(
-        *client.environment_connection_state_tx.borrow(),
+        *client.subscribe_connection_state().borrow(),
         EnvironmentConnectionState::Connected
     );
     Ok(())
@@ -667,7 +703,7 @@ async fn failed_replacement_connection_keeps_old_handles_retired_and_get_retries
         Err(ExecServerError::Disconnected(_))
     ));
     assert_eq!(
-        *client.environment_connection_state_tx.borrow(),
+        *client.subscribe_connection_state().borrow(),
         EnvironmentConnectionState::Disconnected
     );
 
@@ -677,7 +713,7 @@ async fn failed_replacement_connection_keeps_old_handles_retired_and_get_retries
     assert_ne!(original.session_id(), replacement.session_id());
     replacement.environment_status().await?;
     assert_eq!(
-        *client.environment_connection_state_tx.borrow(),
+        *client.subscribe_connection_state().borrow(),
         EnvironmentConnectionState::Connected
     );
     assert!(matches!(
@@ -1067,3 +1103,80 @@ async fn application_policy_registry_denial_does_not_poison_later_noise_startup(
     client.get().await?;
     Ok(())
 }
+
+/// A closed transport is recovery, even before the recovery task changes session status.
+#[tokio::test]
+async fn connection_watch_reports_recovery_before_recovery_task_runs() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, registry) = executor.client()?;
+    let original = client.get().await?;
+    let mut state = client.subscribe_connection_state();
+    let release = registry.block_next_lookup();
+    // Closing does not yield, so publishing here precedes the RPC reader's recovery task.
+    disconnect(&original).await;
+    assert!(matches!(
+        original.inner.connection.lock().unwrap().status,
+        ConnectionStatus::Connected(_)
+    ));
+    original.attach_environment_connection_state(client.environment_connection_state_tx.clone());
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Recovering
+    );
+    assert!(client.fail_fast().get().await.is_err());
+
+    registry.lookup_started.notified().await;
+    executor.detached.notified().await;
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Recovering
+    );
+    let waiting = original.inner.rpc_client();
+    tokio::pin!(waiting);
+    assert!(futures::poll!(waiting.as_mut()).is_pending());
+    release.send(()).unwrap();
+    waiting.await?;
+    assert!(state.has_changed()?);
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Connected
+    );
+    assert!(!state.has_changed()?);
+    original.environment_status().await?;
+    Ok(())
+}
+
+/// Refresh announces retirement before asynchronous transport cleanup. A delayed
+/// callback from that transport must not leave the shared watch in Recovering.
+#[tokio::test]
+async fn late_recovery_callback_cannot_overwrite_retirement() -> Result<()> {
+    let _clock = freeze_clock();
+    let executor = Executor::start(Validator::default()).await?;
+    let (client, _registry) = executor.client()?;
+    let original = client.get().await?;
+    let mut state = client.subscribe_connection_state();
+    let rpc = {
+        let current = client.current_client.lock().unwrap();
+        let connection = current.as_ref().unwrap().inner.connection.lock().unwrap();
+        let ConnectionStatus::Connected(rpc) = &connection.status else {
+            panic!("expected connected session");
+        };
+        // Stop at refresh's retirement/publication boundary.
+        original.inner.retired.cancel();
+        client
+            .environment_connection_state_tx
+            .send_replace(EnvironmentConnectionState::Disconnected);
+        Arc::clone(rpc)
+    };
+    original
+        .inner
+        .request_recovery(rpc, "late old transport callback".to_owned());
+    original.inner.retire().await;
+    // A failed replacement adds no further wakeup in the base implementation.
+    assert_eq!(
+        *state.borrow_and_update(),
+        EnvironmentConnectionState::Disconnected
+    );
+    Ok(())
+}

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

@@ -71,6 +71,8 @@ pub const CODEX_EXEC_SERVER_NOISE_CHATGPT_ACCOUNT_ID_ENV_VAR: &str =
 pub enum EnvironmentConnectionState {
     /// An initialized exec-server connection is available.
     Connected,
+    /// An established session is unavailable while its transport is being recovered.
+    Recovering,
     /// No initialized exec-server connection is currently available.
     Disconnected,
 }

codex-rs/exec-server/tests/accepted_websocket.rs

@@ -319,7 +319,7 @@ async fn accepted_websocket_interoperates_and_recovers_with_real_direct_executor
         }
         timeout(
             TEST_TIMEOUT,
-            connection_state.wait_for(|state| *state == EnvironmentConnectionState::Disconnected),
+            connection_state.wait_for(|state| *state == EnvironmentConnectionState::Recovering),
         )
         .await??;
         tokio::fs::write(&gate_path, b"").await?;

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

@@ -304,7 +304,7 @@ async fn failed_noise_environment_recovers_and_reconnects_after_ready_report() -
     let _ = first_relay.await;
     assert_eq!(
         next_connection_state(&mut connection_state).await?,
-        EnvironmentConnectionState::Disconnected
+        EnvironmentConnectionState::Recovering
     );
     let first_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;
     let second_reconnected_websocket = accept_websocket(&listener, "reconnected peer").await?;

codex-rs/ext/multi-env-copy-file/tests/suite/connectivity.rs

@@ -217,7 +217,7 @@ async fn copy_file_reports_disconnection_after_step_capture() -> Result<()> {
     streaming.wait_for_request_count(/*count*/ 1).await;
     drop(executor);
     state
-        .wait_for(|state| *state == EnvironmentConnectionState::Disconnected)
+        .wait_for(|state| *state == EnvironmentConnectionState::Recovering)
         .await?;
     release_response.send(()).unwrap();
     // Allow the client's recovery window to finish before checking the completed tool results.