gemini_genai_rs/transport/
replay.rs

1//! Replay transport — feed a recorded wire log back through the session loop.
2//!
3//! [`ReplayTransport`] implements [`Transport`] over a recorded wire log (see
4//! [`crate::transport::recording`]): `recv()` yields the recorded **inbound**
5//! frames in order (as fast as the session loop consumes them), and `send()`
6//! collects outbound frames for later comparison instead of touching a network.
7//!
8//! Because the session loop broadcasts events as soon as frames arrive, a
9//! replay that starts streaming before the application has subscribed would
10//! lose events nondeterministically. The transport is therefore *gated*: the
11//! first `ungated_prefix` frames (default 1 — the `setupComplete` handshake)
12//! are delivered immediately so the connection can reach `Active`, and the
13//! rest are held until [`ReplayControl::release`] is called. Once the inbound
14//! queue is exhausted the [`ReplayControl::drained`] signal fires and `recv()`
15//! pends (like [`MockTransport`](super::ws::MockTransport)), keeping the
16//! session alive until it is disconnected.
17
18use std::collections::VecDeque;
19use std::sync::Arc;
20
21use async_trait::async_trait;
22use tokio::sync::watch;
23
24use super::recording::{WireDirection, WireEntry};
25use super::ws::Transport;
26
27/// Shared collection of frames "sent" during a replay.
28pub type OutboundFrames = Arc<parking_lot::Mutex<Vec<Vec<u8>>>>;
29
30/// Called with each frame's recorded capture time (milliseconds since the
31/// epoch) before the frame is delivered. The frame waits for the returned
32/// future, so a replay can hold a frame back until the session has finished
33/// processing the frames before it.
34pub type FrameGate = Arc<
35    dyn Fn(u64) -> std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>> + Send + Sync,
36>;
37
38/// Errors from the [`ReplayTransport`].
39#[derive(Debug, thiserror::Error)]
40pub enum ReplayTransportError {
41    /// Operation attempted while not connected.
42    #[error("Not connected")]
43    NotConnected,
44}
45
46/// Control handle for a [`ReplayTransport`] that has been moved into a
47/// session loop.
48#[derive(Clone)]
49pub struct ReplayControl {
50    gate_tx: Arc<watch::Sender<bool>>,
51    drained_rx: watch::Receiver<bool>,
52    outbound: OutboundFrames,
53}
54
55impl ReplayControl {
56    /// Release the gated frames: inbound replay starts flowing.
57    ///
58    /// Call this after subscribing to the session's events so none are lost.
59    pub fn release(&self) {
60        let _ = self.gate_tx.send(true);
61    }
62
63    /// Wait until every recorded inbound frame has been handed to the session
64    /// loop. Note: the *last* frame may still be in flight through the
65    /// processor when this returns — wait for its observable effects (events,
66    /// state) before asserting.
67    pub async fn drained(&self) {
68        let mut rx = self.drained_rx.clone();
69        while !*rx.borrow() {
70            if rx.changed().await.is_err() {
71                break;
72            }
73        }
74    }
75
76    /// Snapshot the outbound frames collected so far (in send order).
77    pub fn outbound_frames(&self) -> Vec<Vec<u8>> {
78        self.outbound.lock().clone()
79    }
80}
81
82/// A [`Transport`] that replays recorded inbound frames and collects outbound
83/// frames. See the [module docs](self) for gating and drain semantics.
84pub struct ReplayTransport {
85    inbound: VecDeque<Vec<u8>>,
86    timestamps: VecDeque<u64>,
87    on_frame: Option<FrameGate>,
88    ungated_prefix: usize,
89    delivered: usize,
90    gate_rx: watch::Receiver<bool>,
91    drained_tx: watch::Sender<bool>,
92    outbound: OutboundFrames,
93    connected: bool,
94}
95
96impl ReplayTransport {
97    /// Build a replay transport from raw inbound frames.
98    ///
99    /// The first frame should be the `setupComplete` handshake; it is
100    /// delivered ungated so the connection can reach `Active`.
101    pub fn from_frames(frames: Vec<Vec<u8>>) -> (Self, ReplayControl) {
102        let (gate_tx, gate_rx) = watch::channel(false);
103        let (drained_tx, drained_rx) = watch::channel(false);
104        let outbound: OutboundFrames = Arc::new(parking_lot::Mutex::new(Vec::new()));
105        let control = ReplayControl {
106            gate_tx: Arc::new(gate_tx),
107            drained_rx,
108            outbound: outbound.clone(),
109        };
110        (
111            Self {
112                inbound: frames.into(),
113                timestamps: VecDeque::new(),
114                on_frame: None,
115                ungated_prefix: 1,
116                delivered: 0,
117                gate_rx,
118                drained_tx,
119                outbound,
120                connected: false,
121            },
122            control,
123        )
124    }
125
126    /// Build a replay transport from a recorded wire log, keeping only the
127    /// [`WireDirection::Inbound`] entries (in log order).
128    ///
129    /// Each frame keeps its recorded capture time, reported to a
130    /// [`with_frame_gate`](Self::with_frame_gate) callback.
131    pub fn from_wire_log(entries: &[WireEntry]) -> (Self, ReplayControl) {
132        let inbound: Vec<&WireEntry> = entries
133            .iter()
134            .filter(|e| e.dir == WireDirection::Inbound)
135            .collect();
136        let (mut transport, control) =
137            Self::from_frames(inbound.iter().map(|e| e.payload.clone()).collect());
138        transport.timestamps = inbound.iter().map(|e| e.ts_ms).collect();
139        (transport, control)
140    }
141
142    /// Pass each frame's recorded capture time to `gate` and wait for it
143    /// before delivering the frame. Frames built with
144    /// [`from_frames`](Self::from_frames) carry no timestamps, so the gate is
145    /// not called for them.
146    pub fn with_frame_gate(mut self, gate: FrameGate) -> Self {
147        self.on_frame = Some(gate);
148        self
149    }
150
151    /// Override how many leading frames are delivered before
152    /// [`ReplayControl::release`] (default 1: the setup handshake).
153    pub fn with_ungated_prefix(mut self, n: usize) -> Self {
154        self.ungated_prefix = n;
155        self
156    }
157}
158
159#[async_trait]
160impl Transport for ReplayTransport {
161    type Error = ReplayTransportError;
162
163    async fn connect(
164        &mut self,
165        _url: &str,
166        _headers: Vec<(String, String)>,
167    ) -> Result<(), Self::Error> {
168        self.connected = true;
169        Ok(())
170    }
171
172    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
173        if !self.connected {
174            return Err(ReplayTransportError::NotConnected);
175        }
176        self.outbound.lock().push(data);
177        Ok(())
178    }
179
180    async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
181        if !self.connected {
182            return Err(ReplayTransportError::NotConnected);
183        }
184        // Yield so observers can see intermediate states between frames
185        // (mirrors MockTransport).
186        tokio::task::yield_now().await;
187
188        if self.inbound.is_empty() {
189            let _ = self.drained_tx.send(true);
190            // Stay connected-but-idle; the session loop's `select!` drops this
191            // future when a command (e.g. Disconnect) arrives.
192            std::future::pending::<()>().await;
193            unreachable!("pending() never resolves");
194        }
195
196        if self.delivered >= self.ungated_prefix {
197            let mut gate = self.gate_rx.clone();
198            while !*gate.borrow() {
199                // If the control handle is dropped, proceed ungated rather
200                // than deadlocking the replay.
201                if gate.changed().await.is_err() {
202                    break;
203                }
204            }
205        }
206
207        let frame = self
208            .inbound
209            .pop_front()
210            .expect("checked non-empty inbound queue");
211        if let Some(ts_ms) = self.timestamps.pop_front()
212            && let Some(gate) = &self.on_frame
213        {
214            gate(ts_ms).await;
215        }
216        self.delivered += 1;
217        if self.inbound.is_empty() {
218            let _ = self.drained_tx.send(true);
219        }
220        Ok(Some(frame))
221    }
222
223    async fn close(&mut self) -> Result<(), Self::Error> {
224        self.connected = false;
225        Ok(())
226    }
227}
228
229#[cfg(test)]
230mod tests {
231    use super::*;
232    use std::time::Duration;
233
234    #[tokio::test]
235    async fn replay_delivers_first_frame_ungated_then_waits_for_release() {
236        let (mut transport, control) = ReplayTransport::from_frames(vec![
237            br#"{"setupComplete":{}}"#.to_vec(),
238            br#"{"serverContent":{"turnComplete":true}}"#.to_vec(),
239        ]);
240        transport.connect("replay://", vec![]).await.unwrap();
241
242        // First frame (handshake) flows without release.
243        let first = transport.recv().await.unwrap().unwrap();
244        assert!(String::from_utf8(first).unwrap().contains("setupComplete"));
245
246        // Second frame is gated.
247        let gated = tokio::time::timeout(Duration::from_millis(50), transport.recv()).await;
248        assert!(gated.is_err(), "second frame should be gated");
249
250        control.release();
251        let second = transport.recv().await.unwrap().unwrap();
252        assert!(String::from_utf8(second).unwrap().contains("turnComplete"));
253
254        // Drained fires once the queue is exhausted.
255        tokio::time::timeout(Duration::from_millis(100), control.drained())
256            .await
257            .expect("drained should be signalled");
258
259        // And recv() pends from then on.
260        let idle = tokio::time::timeout(Duration::from_millis(50), transport.recv()).await;
261        assert!(idle.is_err(), "recv should pend after drain");
262    }
263
264    #[tokio::test]
265    async fn replay_collects_outbound_frames() {
266        let (mut transport, control) =
267            ReplayTransport::from_frames(vec![br#"{"setupComplete":{}}"#.to_vec()]);
268        transport.connect("replay://", vec![]).await.unwrap();
269        transport.send(b"{\"setup\":{}}".to_vec()).await.unwrap();
270        transport
271            .send(b"{\"toolResponse\":{}}".to_vec())
272            .await
273            .unwrap();
274
275        let sent = control.outbound_frames();
276        assert_eq!(sent.len(), 2);
277        assert_eq!(sent[0], b"{\"setup\":{}}".to_vec());
278    }
279
280    #[tokio::test]
281    async fn replay_from_wire_log_keeps_inbound_only() {
282        let entries = vec![
283            WireEntry {
284                seq: 1,
285                dir: WireDirection::Outbound,
286                ts_ms: 1,
287                payload: b"{\"setup\":{}}".to_vec(),
288            },
289            WireEntry {
290                seq: 2,
291                dir: WireDirection::Inbound,
292                ts_ms: 2,
293                payload: br#"{"setupComplete":{}}"#.to_vec(),
294            },
295        ];
296        let (mut transport, _control) = ReplayTransport::from_wire_log(&entries);
297        transport.connect("replay://", vec![]).await.unwrap();
298        let first = transport.recv().await.unwrap().unwrap();
299        assert!(String::from_utf8(first).unwrap().contains("setupComplete"));
300    }
301
302    #[tokio::test]
303    async fn replay_errors_when_not_connected() {
304        let (mut transport, _control) = ReplayTransport::from_frames(vec![]);
305        assert!(transport.recv().await.is_err());
306        assert!(transport.send(vec![1]).await.is_err());
307    }
308}
gemini_genai_rs/transport/replay.rs

gemini_genai_rs/transport/
replay.rs
