diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 8f98962..b2291ac 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -4,6 +4,14 @@ ### Added +- Add stable protocol v1 `session/load` and `session/resume` restore builders + (`ConnectionTo::load_session`, `ConnectionTo::resume_session`, + `ConnectionTo::restore_session_from`) that install session routing *before* + publishing the request, so replay or early notifications for the restored + session are captured, and return the complete restore response alongside the + `ActiveSession` in a `RestoredSession`. On failure or cancellation the + routing is dropped, leaving no stale handler behind. Mirrors the v2 + `OpenedV2Session` shape. - *(unstable-v2)* Add `Proxy::protocol_router` and `ProxyProtocolRouter` to compose strict v1 and v2 proxy implementations behind one connection. Routing requires the exact version selected by the conductor and preserves diff --git a/src/agent-client-protocol/src/session.rs b/src/agent-client-protocol/src/session.rs index 47200c0..be35942 100644 --- a/src/agent-client-protocol/src/session.rs +++ b/src/agent-client-protocol/src/session.rs @@ -3,15 +3,17 @@ use std::{future::Future, marker::PhantomData, path::Path}; use futures::channel::{mpsc, oneshot}; use crate::{ - Agent, Client, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Responder, Role, + Agent, Client, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, JsonRpcRequest, Responder, + Role, jsonrpc::{ DynamicHandlerGuard, run::{NullRun, RunWithConnectionTo}, }, role::{HasPeer, acp::ProxySessionMessages}, schema::v1::{ - ContentBlock, ContentChunk, NewSessionRequest, NewSessionResponse, PromptRequest, - PromptResponse, SessionId, SessionModeState, SessionNotification, SessionUpdate, + ContentBlock, ContentChunk, LoadSessionRequest, LoadSessionResponse, Meta, + NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse, ResumeSessionRequest, + ResumeSessionResponse, SessionId, SessionModeState, SessionNotification, SessionUpdate, StopReason, }, util::{MatchDispatch, MatchDispatchFrom, run_until}, @@ -39,6 +41,64 @@ impl SessionBlockState for NonBlocking {} /// See [`SessionBuilder::block_task`]. pub trait SessionBlockState: Send + 'static + Sync + std::fmt::Debug {} +/// A stable protocol v1 "restore an existing session" request — either +/// `session/load` or `session/resume` — reduced to what the SDK needs to attach +/// an [`ActiveSession`] to the restored session. +/// +/// Implemented for [`LoadSessionRequest`] and [`ResumeSessionRequest`]. The +/// full request stays on the builder, so interception helpers such as MCP +/// injection keep working unchanged. +pub trait RestoreSessionRequest: JsonRpcRequest { + /// The session this request loads or resumes. + fn session_id(&self) -> &SessionId; +} + +/// The response to a stable protocol v1 restore request, reduced to what an +/// [`ActiveSession`] carries from the operation. +/// +/// Implemented for [`LoadSessionResponse`] and [`ResumeSessionResponse`]. The +/// complete response remains available on [`RestoredSession::response`], so no +/// operation-specific data is lost by this reduction. +pub trait RestoreSessionResponse { + /// Initial mode state reported by the agent, if any. + fn modes(&self) -> Option; + + /// `_meta` data reported by the agent, if any. + fn meta(&self) -> Option; +} + +impl RestoreSessionRequest for LoadSessionRequest { + fn session_id(&self) -> &SessionId { + &self.session_id + } +} + +impl RestoreSessionResponse for LoadSessionResponse { + fn modes(&self) -> Option { + self.modes.clone() + } + + fn meta(&self) -> Option { + self.meta.clone() + } +} + +impl RestoreSessionRequest for ResumeSessionRequest { + fn session_id(&self) -> &SessionId { + &self.session_id + } +} + +impl RestoreSessionResponse for ResumeSessionResponse { + fn modes(&self) -> Option { + self.modes.clone() + } + + fn meta(&self) -> Option { + self.meta.clone() + } +} + impl ConnectionTo where Counterpart: HasPeer, @@ -75,6 +135,44 @@ where SessionBuilder::new(self, request) } + /// Stable protocol v1 session builder that loads an existing session. + /// + /// [`ActiveSession`] routing is installed as with + /// [`build_session`](Self::build_session); `session/load` carries the + /// session id itself, so the builder hands it to the request rather than + /// reading it back from the response. + pub fn load_session( + &self, + session_id: impl Into, + cwd: impl AsRef, + ) -> RestoreSessionBuilder { + RestoreSessionBuilder::new(self, LoadSessionRequest::new(session_id, cwd.as_ref())) + } + + /// Stable protocol v1 session builder that resumes an existing session. + /// + /// The `session/resume` counterpart of [`load_session`](Self::load_session). + pub fn resume_session( + &self, + session_id: impl Into, + cwd: impl AsRef, + ) -> RestoreSessionBuilder { + RestoreSessionBuilder::new(self, ResumeSessionRequest::new(session_id, cwd.as_ref())) + } + + /// Stable protocol v1 session builder starting from an existing restore + /// request — either `session/load` or `session/resume`. + /// + /// Use this when you've intercepted a restore request and want to modify + /// it (e.g., inject MCP servers) before forwarding. + pub fn restore_session_from(&self, request: Req) -> RestoreSessionBuilder + where + Req: RestoreSessionRequest, + Req::Response: RestoreSessionResponse, + { + RestoreSessionBuilder::new(self, request) + } + /// Given a session response received from the agent, /// attach a handler to process messages related to this session /// and let you access them. @@ -97,22 +195,283 @@ where .. } = response; - let (update_tx, update_rx) = mpsc::unbounded(); - let handler = ActiveSessionHandler::new(session_id.clone(), update_tx.clone()); - let session_handler_registration = self.add_dynamic_handler(handler)?; + let prepared = self.prepare_session_routing(&session_id)?; Ok(ActiveSession { session_id, modes, meta, - update_rx, - update_tx, + update_rx: prepared.update_rx, + update_tx: prepared.update_tx, connection: self.clone(), - session_handler_registration, + session_handler_registration: prepared.session_handler_registration, mcp_handler_registrations, _runner: PhantomData, }) } + + /// Install the update channel and session handler for `session_id`, + /// returning the pieces an [`ActiveSession`] claims on success. + /// + /// Call this *before* publishing a restore request so routing is in place + /// for replay or early notifications; dropping the returned registration + /// removes routing on failure or cancellation. + fn prepare_session_routing( + &self, + session_id: &SessionId, + ) -> Result, crate::Error> { + let (update_tx, update_rx) = mpsc::unbounded(); + let handler = ActiveSessionHandler::new(session_id.clone(), update_tx.clone()); + let session_handler_registration = self.add_dynamic_handler(handler)?; + + Ok(PreparedSession { + update_rx, + update_tx, + session_handler_registration, + }) + } +} + +/// Routing pieces installed for a restore request before it is published. +/// +/// `prepare_session_routing` returns this so the caller can publish the +/// request with routing already in place; an [`ActiveSession`] claims the +/// pieces on success, and the registration drops routing on failure or +/// cancellation. +struct PreparedSession +where + Counterpart: HasPeer, +{ + update_rx: mpsc::UnboundedReceiver, + update_tx: mpsc::UnboundedSender, + session_handler_registration: DynamicHandlerGuard, +} + +/// Stable protocol v1 session builder that restores an existing session via +/// either `session/load` or `session/resume`. +/// +/// The `BlockState` type parameter tracks whether blocking methods are +/// available, mirroring [`SessionBuilder`]: +/// - `NonBlocking` (default): Only [`on_session_start`](Self::on_session_start) is available +/// - `Blocking` (after calling [`block_task`](Self::block_task)): +/// [`start_session`](Self::start_session) becomes available +/// +/// The routing that delivers session messages to the returned +/// [`ActiveSession`] is installed before the request is published, so replay +/// or early notifications for the restored session are captured; on failure +/// or cancellation the routing is dropped with the request. +#[must_use = "use `start_session` or `on_session_start` to start the session"] +#[derive(Debug)] +pub struct RestoreSessionBuilder< + Counterpart, + Req: RestoreSessionRequest, + BlockState: SessionBlockState = NonBlocking, +> where + Counterpart: HasPeer, + Req::Response: RestoreSessionResponse, +{ + connection: ConnectionTo, + request: Req, + block_state: PhantomData, +} + +impl RestoreSessionBuilder +where + Counterpart: HasPeer, + Req: RestoreSessionRequest, + Req::Response: RestoreSessionResponse, +{ + fn new(connection: &ConnectionTo, request: Req) -> Self { + RestoreSessionBuilder { + connection: connection.clone(), + request, + block_state: PhantomData, + } + } + + /// Restore the session in a spawned background task, running `op` with + /// the restored session once the agent confirms. + /// + /// Mirrors [`SessionBuilder::on_session_start`]: returns immediately and + /// the handshake runs in a background task. Because routing is installed + /// before the request is published, session messages that arrive before + /// the restore response still reach `op`'s [`ActiveSession`]. + /// + /// On an error response the session-handler routing is dropped, so a + /// failed restore leaves no stale handler behind. + pub fn on_session_start(self, op: F) -> Result<(), crate::Error> + where + F: FnOnce(RestoredSession<'static, Counterpart, Req::Response>) -> Fut + Send + 'static, + Fut: Future> + Send, + { + ensure_v1_session_protocol(&self.connection)?; + + let RestoreSessionBuilder { + connection, + request, + block_state: _, + } = self; + + // Install routing before publishing so replay or early notifications + // for this session are captured from the first inbound dispatch. + let session_id = request.session_id().clone(); + let prepared = connection.prepare_session_routing(&session_id)?; + + connection + .send_ordered_request_to(Agent, request) + .on_receiving_result({ + let connection = connection.clone(); + async move |result| { + // Failure or cancellation: `result?` returns and `prepared` + // is dropped, removing the session-handler routing. + let response = result?; + + let session = ActiveSession { + session_id, + modes: response.modes(), + meta: response.meta(), + update_rx: prepared.update_rx, + update_tx: prepared.update_tx, + connection: connection.clone(), + session_handler_registration: prepared.session_handler_registration, + mcp_handler_registrations: Vec::new(), + _runner: PhantomData, + }; + + connection.spawn(async move { op(RestoredSession { session, response }).await }) + } + }) + } +} + +impl RestoreSessionBuilder +where + Counterpart: HasPeer, + Req: RestoreSessionRequest, + Req::Response: RestoreSessionResponse, +{ + /// Mark this restore builder as able to block the current task. + /// + /// After calling this, you can use [`start_session`](Self::start_session) + /// which blocks the current task. + /// + /// This should not be used from inside a message handler like + /// [`Builder::on_receive_request`](`crate::Builder::on_receive_request`). + pub fn block_task(self) -> RestoreSessionBuilder { + RestoreSessionBuilder { + connection: self.connection, + request: self.request, + block_state: PhantomData, + } + } +} + +impl RestoreSessionBuilder +where + Counterpart: HasPeer, + Req: RestoreSessionRequest, + Req::Response: RestoreSessionResponse, +{ + /// Send the restore request, block until the agent confirms, and return + /// the restored session handle. + /// + /// Mirrors [`SessionBuilder::start_session`]. Routing is installed before + /// the request is published; on an error response it is dropped. + /// + /// Requires calling [`block_task`](Self::block_task) first. + pub async fn start_session( + self, + ) -> Result, crate::Error> { + ensure_v1_session_protocol(&self.connection)?; + + let RestoreSessionBuilder { + connection, + request, + block_state: _, + } = self; + + let session_id = request.session_id().clone(); + let prepared = connection.prepare_session_routing(&session_id)?; + + // On error `prepared` is dropped here, removing the session routing. + let response = match connection + .send_request_to(Agent, request) + .block_task() + .await + { + Ok(response) => response, + Err(err) => return Err(err), + }; + + let session = ActiveSession { + session_id, + modes: response.modes(), + meta: response.meta(), + update_rx: prepared.update_rx, + update_tx: prepared.update_tx, + connection, + session_handler_registration: prepared.session_handler_registration, + mcp_handler_registrations: Vec::new(), + _runner: PhantomData, + }; + + Ok(RestoredSession { session, response }) + } +} + +/// A restored session and the operation-specific response that opened it. +/// +/// The response is kept separate from the [`ActiveSession`] because v1 +/// `session/load` and `session/resume` responses carry no session id — it +/// lives on the request — so the complete response stays available without +/// being folded into session state. Mirrors the protocol v2 `OpenedV2Session`. +pub struct RestoredSession<'runner, Link, Response> +where + Link: HasPeer, +{ + session: ActiveSession<'runner, Link>, + response: Response, +} + +impl<'runner, Link, Response> RestoredSession<'runner, Link, Response> +where + Link: HasPeer, +{ + /// Access the command handle for the restored session. + pub fn session(&self) -> &ActiveSession<'runner, Link> { + &self.session + } + + /// Access the command handle mutably, e.g. to read updates that arrive + /// after the restore. + pub fn session_mut(&mut self) -> &mut ActiveSession<'runner, Link> { + &mut self.session + } + + /// Access the complete response from the restore operation. + pub fn response(&self) -> &Response { + &self.response + } + + /// Split this result into the command handle and complete setup response. + pub fn into_parts(self) -> (ActiveSession<'runner, Link>, Response) { + (self.session, self.response) + } +} + +impl std::fmt::Debug for RestoredSession<'_, Link, Response> +where + Link: HasPeer, + Response: std::fmt::Debug, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // [`ActiveSession`] holds live channels and is intentionally not + // `Debug`; its session id is the identity worth printing here. + f.debug_struct("RestoredSession") + .field("session_id", self.session.session_id()) + .field("response", &self.response) + .finish() + } } /// Stable protocol v1 session builder for a new session request. diff --git a/src/agent-client-protocol/tests/session_restore.rs b/src/agent-client-protocol/tests/session_restore.rs new file mode 100644 index 0000000..13179ba --- /dev/null +++ b/src/agent-client-protocol/tests/session_restore.rs @@ -0,0 +1,342 @@ +//! Stable protocol v1 session restore: `session/load` and `session/resume` +//! (issue #323). +//! +//! v1 restore responses carry no session id — the id lives on the request — +//! so the SDK's restore path takes the id from the request, installs the +//! session routing *before* the request is published, and hands the complete +//! response back alongside the [`ActiveSession`]. These tests pin the two +//! properties that fall out of that design: +//! +//! - an early notification sharing a batch with the restore response is routed +//! to the session even though it precedes the response in dispatch order; +//! - a failed restore drops the routing, so a later update for the failed +//! session id cannot reach a stale handler. +//! +//! The raw-frame peers drive the exact transport wiring a real agent sees. + +use std::time::Duration; + +use agent_client_protocol::{ + Channel, Client, Error, RawJsonRpcMessage, SessionMessage, TransportBatch, TransportFrame, + schema::v1::{ + ContentBlock, ContentChunk, LoadSessionRequest, LoadSessionResponse, ResumeSessionRequest, + ResumeSessionResponse, SessionId, SessionMode, SessionModeId, SessionModeState, + SessionNotification, SessionUpdate, TextContent, + }, +}; +use futures::{StreamExt as _, channel::oneshot}; + +const TIMEOUT: Duration = Duration::from_secs(10); + +fn modes_state() -> SessionModeState { + SessionModeState::new( + SessionModeId::new("default"), + vec![SessionMode::new(SessionModeId::new("default"), "Default")], + ) +} + +/// A `session/update` notification for `session_id`, serialized the way the +/// agent would put it on the wire. +fn update_notification(session_id: SessionId) -> RawJsonRpcMessage { + RawJsonRpcMessage::notification( + "session/update".to_string(), + serde_json::to_value(SessionNotification::new( + session_id, + SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text( + TextContent::new("restored update"), + ))), + )) + .expect("session notification should serialize"), + ) + .expect("session notification should form valid JSON-RPC parameters") +} + +/// The headline restore path: routing is installed before the request is +/// published, so an update that arrives *before* the restore response in the +/// same batch is still captured, and the complete load response round-trips. +#[tokio::test(flavor = "current_thread")] +async fn load_session_routes_early_notification_and_returns_exact_response() { + let session_id = SessionId::new("restore-load"); + let response_session_id = session_id.clone(); + let peer_response_id = response_session_id.clone(); + let notification_session_id = session_id.clone(); + let expected = LoadSessionResponse::new().modes(modes_state()).meta( + serde_json::json!({ "trace": "abc" }) + .as_object() + .unwrap() + .clone(), + ); + let wire_response = expected.clone(); + let (transport, mut peer) = Channel::duplex(); + let (result_tx, result_rx) = oneshot::channel(); + + let client = Client + .builder() + .connect_with(transport, async move |connection| { + connection + .load_session(session_id.clone(), "/restore/cwd") + .on_session_start(async move |mut restored| { + // The notification was dispatched before the response; it + // must still reach the session's update stream. + let update = restored.session_mut().read_update().await?; + assert!(matches!(update, SessionMessage::SessionMessage(_))); + assert_eq!(restored.response(), &expected); + assert_eq!(restored.session().session_id(), &response_session_id); + result_tx.send(()).map_err(|()| Error::internal_error()) + })?; + + result_rx.await.map_err(|_| Error::internal_error()) + }); + + let peer = async move { + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = + peer.rx.next().await + else { + panic!("expected a session/load request"); + }; + assert_eq!(request.method.as_ref(), "session/load"); + let req: LoadSessionRequest = serde_json::from_value( + request + .params + .expect("session/load request carries params") + .into_value(), + ) + .expect("session/load params should parse"); + assert_eq!(req.session_id, peer_response_id); + assert_eq!(req.cwd, std::path::PathBuf::from("/restore/cwd")); + + // Update BEFORE response in the same batch: with routing installed + // before publish, dispatch can route it even though it precedes the + // restore response. + let notification = update_notification(notification_session_id); + let response = RawJsonRpcMessage::response( + request.id, + Ok(serde_json::to_value(wire_response).expect("load response should serialize")), + ); + let batch = TransportBatch::from_messages([notification, response]) + .expect("test batch should be non-empty"); + peer.tx + .unbounded_send(TransportFrame::Batch(batch)) + .expect("client should accept the restore batch"); + + while peer.rx.next().await.is_some() {} + Ok::<(), Error>(()) + }; + + tokio::time::timeout(TIMEOUT, async { futures::try_join!(client, peer) }) + .await + .expect("same-batch restore update was not routed") + .expect("load session connection failed"); +} + +/// The blocking path: `resume_session().block_task().start_session()` returns +/// the restored session and the exact resume response. +#[tokio::test(flavor = "current_thread")] +async fn resume_session_round_trips_exact_response() { + let session_id = SessionId::new("restore-resume"); + let client_session_id = session_id.clone(); + let expected = ResumeSessionResponse::new().modes(modes_state()); + let wire_response = expected.clone(); + let (transport, mut peer) = Channel::duplex(); + + let client = Client + .builder() + .connect_with(transport, async move |connection| { + let restored = connection + .resume_session(client_session_id.clone(), "/resume/cwd") + .block_task() + .start_session() + .await?; + assert_eq!(restored.session().session_id(), &client_session_id); + assert_eq!(restored.response(), &expected); + + let (active, response) = restored.into_parts(); + assert_eq!(active.session_id(), &client_session_id); + assert_eq!(response, expected); + + Ok(()) + }); + + let peer = async move { + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = + peer.rx.next().await + else { + panic!("expected a session/resume request"); + }; + assert_eq!(request.method.as_ref(), "session/resume"); + let req: ResumeSessionRequest = serde_json::from_value( + request + .params + .expect("session/resume request carries params") + .into_value(), + ) + .expect("session/resume params should parse"); + assert_eq!(req.session_id, session_id); + + let response = RawJsonRpcMessage::response( + request.id, + Ok(serde_json::to_value(wire_response).expect("resume response should serialize")), + ); + peer.tx + .unbounded_send(TransportFrame::Single(response)) + .expect("client should accept the resume response"); + + while peer.rx.next().await.is_some() {} + Ok::<(), Error>(()) + }; + + tokio::time::timeout(TIMEOUT, async { futures::try_join!(client, peer) }) + .await + .expect("resume session connection failed") + .expect("resume session errored"); +} + +/// The interception entry point: a request built elsewhere (e.g. decoded from +/// a client's `session/load`) is handed to the SDK unchanged and forwarded +/// verbatim. +#[tokio::test(flavor = "current_thread")] +async fn restore_session_from_forwards_an_intercepted_request() { + let session_id = SessionId::new("restore-from"); + let client_session_id = session_id.clone(); + let expected = LoadSessionResponse::new().modes(modes_state()); + let wire_response = expected.clone(); + let (transport, mut peer) = Channel::duplex(); + let (result_tx, result_rx) = oneshot::channel(); + + let client = Client + .builder() + .connect_with(transport, async move |connection| { + let request = LoadSessionRequest::new(client_session_id, "/intercepted/cwd"); + connection + .restore_session_from(request) + .on_session_start(async move |restored| { + assert_eq!(restored.response(), &expected); + result_tx.send(()).map_err(|()| Error::internal_error()) + })?; + + result_rx.await.map_err(|_| Error::internal_error()) + }); + + let peer = async move { + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = + peer.rx.next().await + else { + panic!("expected a session/load request"); + }; + assert_eq!(request.method.as_ref(), "session/load"); + let req: LoadSessionRequest = serde_json::from_value( + request + .params + .expect("session/load request carries params") + .into_value(), + ) + .expect("session/load params should parse"); + assert_eq!(req.session_id, session_id); + assert_eq!(req.cwd, std::path::PathBuf::from("/intercepted/cwd")); + + let response = RawJsonRpcMessage::response( + request.id, + Ok(serde_json::to_value(wire_response).expect("load response should serialize")), + ); + peer.tx + .unbounded_send(TransportFrame::Single(response)) + .expect("client should accept the restore response"); + + while peer.rx.next().await.is_some() {} + Ok::<(), Error>(()) + }; + + tokio::time::timeout(TIMEOUT, async { futures::try_join!(client, peer) }) + .await + .expect("restore_session_from connection failed") + .expect("restore_session_from errored"); +} + +/// A failed restore returns `Err` and drops the session routing: a subsequent +/// update for the failed session id must not reach a stale handler (an +/// `ActiveSessionHandler` whose receiver is gone would error dispatch), and +/// the connection must remain usable for a fresh restore. +#[tokio::test(flavor = "current_thread")] +async fn failed_restore_returns_err_and_drops_routing() { + let failed_id = SessionId::new("restore-failed"); + let client_failed_id = failed_id.clone(); + let (transport, mut peer) = Channel::duplex(); + let (failed_tx, failed_rx) = oneshot::channel(); + let (proceed_tx, proceed_rx) = oneshot::channel(); + + let client = Client + .builder() + .connect_with(transport, async move |connection| { + let err = connection + .load_session(client_failed_id, "/restore/fail") + .block_task() + .start_session() + .await; + assert!(err.is_err(), "an error response must fail the restore"); + + // The failure consumed the guard; let the peer exercise the stale + // handler question now. + failed_tx.send(()).map_err(|()| Error::internal_error())?; + proceed_rx.await.map_err(|_| Error::internal_error())?; + + // The connection is still healthy: a fresh restore succeeds, and + // had the failed session's handler survived, the update the peer + // sent would have errored dispatch before we got here. + let recovered = SessionId::new("restore-recovery"); + let restored = connection + .load_session(recovered.clone(), "/restore/ok") + .block_task() + .start_session() + .await + .expect("a second restore must succeed"); + assert_eq!(restored.session().session_id(), &recovered); + Ok(()) + }); + + let peer = async move { + // 1. Fail the first restore. + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = + peer.rx.next().await + else { + panic!("expected the first session/load request"); + }; + assert_eq!(request.method.as_ref(), "session/load"); + peer.tx + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::response( + request.id, + Err(Error::internal_error().data("restore refused")), + ))) + .expect("client should accept the error response"); + + // 2. Once the failure has been consumed, send an update for the failed + // session id. With routing dropped it passes through unhandled. + failed_rx.await.expect("client reported the failure"); + peer.tx + .unbounded_send(TransportFrame::Single(update_notification(failed_id))) + .expect("client should accept the late update"); + proceed_tx.send(()).map_err(|()| Error::internal_error())?; + + // 3. Serve the recovery restore. + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(recovery))) = + peer.rx.next().await + else { + panic!("expected the second session/load request"); + }; + assert_eq!(recovery.method.as_ref(), "session/load"); + peer.tx + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::response( + recovery.id, + Ok(serde_json::to_value(LoadSessionResponse::new()) + .expect("recovery response should serialize")), + ))) + .expect("client should accept the recovery response"); + + while peer.rx.next().await.is_some() {} + Ok::<(), Error>(()) + }; + + tokio::time::timeout(TIMEOUT, async { futures::try_join!(client, peer) }) + .await + .expect("failed restore left the connection unhealthy") + .expect("failed restore connection errored"); +}