diff --git a/codex-rs/app-server/tests/suite/v2/environment_add.rs b/codex-rs/app-server/tests/suite/v2/environment_add.rs index 477c628132..7c6e66211f 100644 --- a/codex-rs/app-server/tests/suite/v2/environment_add.rs +++ b/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::( - 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::(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::(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(()) } diff --git a/codex-rs/core/src/environment_connection_tests.rs b/codex-rs/core/src/environment_connection_tests.rs new file mode 100644 index 0000000000..a4094268fe --- /dev/null +++ b/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, + 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(); +} diff --git a/codex-rs/core/src/environment_selection.rs b/codex-rs/core/src/environment_selection.rs index ef2b50848c..1d81033895 100644 --- a/codex-rs/core/src/environment_selection.rs +++ b/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) -> Value { + pub(super) async fn read_websocket_json(websocket: &mut WebSocketStream) -> Value { loop { match timeout(std::time::Duration::from_secs(5), websocket.next()) .await diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index 6f1f74114b..13f0e82f81 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/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 diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index fa353e4a38..7c6fdc86d4 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/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 diff --git a/codex-rs/exec-server/src/client_refresh_tests.rs b/codex-rs/exec-server/src/client_refresh_tests.rs index ac79cb8c40..9a5cb0a453 100644 --- a/codex-rs/exec-server/src/client_refresh_tests.rs +++ b/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(()) +} diff --git a/codex-rs/exec-server/src/environment.rs b/codex-rs/exec-server/src/environment.rs index 8c3d89b88d..0b046061df 100644 --- a/codex-rs/exec-server/src/environment.rs +++ b/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, } diff --git a/codex-rs/exec-server/tests/accepted_websocket.rs b/codex-rs/exec-server/tests/accepted_websocket.rs index 259ea8dfac..99eb447549 100644 --- a/codex-rs/exec-server/tests/accepted_websocket.rs +++ b/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?; diff --git a/codex-rs/exec-server/tests/relay.rs b/codex-rs/exec-server/tests/relay.rs index e7ff5360dc..8bca066ef1 100644 --- a/codex-rs/exec-server/tests/relay.rs +++ b/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?; diff --git a/codex-rs/ext/multi-env-copy-file/tests/suite/connectivity.rs b/codex-rs/ext/multi-env-copy-file/tests/suite/connectivity.rs index 9e1f129cc7..c179498a85 100644 --- a/codex-rs/ext/multi-env-copy-file/tests/suite/connectivity.rs +++ b/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.