gemini_adk_rs/live/
replay.rs

1//! Offline session replay — feed a recorded wire log through the **real**
2//! control plane.
3//!
4//! Recording happens at L0 via
5//! [`SessionConfig::record_wire`](gemini_genai_rs::prelude::SessionConfig::record_wire)
6//! (every wire byte, both directions, as [`WireEntry`] JSONL). This module
7//! closes the loop: [`replay_session`] opens a
8//! [`ReplayTransport`] over the log's inbound frames and attaches the same three-lane processor a
9//! live connection would get — phase machine, extractors, watchers, tool
10//! dispatch, flow governance all run for real. Nothing is mocked above the
11//! transport seam.
12//!
13//! What replay does and does not do:
14//!
15//! - **Does**: re-decode every recorded inbound frame, re-drive the L1
16//!   processor (events, state writes, tool dispatch through whatever
17//!   dispatcher you attach), and collect the outbound frames the processor
18//!   regenerates (setup, tool responses) for comparison against the log.
19//! - **Does not**: re-execute the model. The model's outputs are *in* the
20//!   recorded inbound frames. User-originated sends (text/audio) are in the
21//!   log's outbound entries but are not re-sent — they only ever existed to
22//!   provoke the recorded inbound frames.
23//!
24//! ```rust,no_run
25//! # async fn demo() -> Result<(), Box<dyn std::error::Error>> {
26//! use gemini_adk_rs::live::replay::replay_session;
27//! use gemini_adk_rs::live::LiveSessionBuilder;
28//! use gemini_genai_rs::prelude::SessionConfig;
29//! use gemini_genai_rs::transport::read_wire_log;
30//!
31//! let entries = read_wire_log("session.wire.jsonl")?;
32//! let config = SessionConfig::new("offline");
33//! let builder = LiveSessionBuilder::new(config.clone());
34//! let replay = replay_session(config, builder, &entries).await?;
35//!
36//! let mut events = replay.handle().events();
37//! replay.release(); // start streaming recorded frames
38//! replay.drained().await; // all frames handed to the session loop
39//! # let _ = events.try_recv();
40//! # Ok(())
41//! # }
42//! ```
43
44use std::sync::Arc;
45use std::time::{Duration, UNIX_EPOCH};
46
47use gemini_genai_rs::prelude::{SessionConfig, SessionPhase};
48use gemini_genai_rs::session::{SessionEvent, SessionHandle};
49use gemini_genai_rs::transport::replay::{ReplayControl, ReplayTransport};
50use gemini_genai_rs::transport::{ConnectBuilder, TransportConfig, WireDirection, WireEntry};
51use tokio::sync::broadcast;
52
53use crate::clock::ManualClock;
54
55use crate::error::AgentError;
56
57use super::builder::{LiveSessionBuilder, build_runtime, spawn_lanes};
58use super::events::LiveEvent;
59use super::handle::LiveHandle;
60
61/// Attach the full L1 control plane (three-lane processor, phase machine,
62/// extractors, watchers, tool dispatch, …) to an **already connected** L0
63/// session.
64///
65/// This is the seam that makes replay possible without touching the network:
66/// connect the L0 session over any [`Transport`](gemini_genai_rs::transport::Transport)
67/// (e.g. [`ReplayTransport`] or [`MockTransport`](gemini_genai_rs::transport::MockTransport)), then hand
68/// it here together with a configured [`LiveSessionBuilder`].
69///
70/// Note: the builder's own `SessionConfig` is *not* re-sent — the setup
71/// message was already encoded from the config given to the L0 connect call.
72/// Subscribe to events **after** this returns and only then let the transport
73/// stream (for `ReplayTransport`, call
74/// [`ReplayControl::release`](gemini_genai_rs::transport::replay::ReplayControl::release)),
75/// otherwise early frames race the subscription.
76pub async fn attach_session(
77    builder: LiveSessionBuilder,
78    session: SessionHandle,
79) -> Result<LiveHandle, AgentError> {
80    let plan = builder.into_plan()?;
81    session.wait_for_phase(SessionPhase::Active).await;
82    let runtime = build_runtime(plan, session);
83    spawn_lanes(runtime).await
84}
85
86/// A replayed session: the live handle plus the replay controls.
87pub struct ReplaySession {
88    handle: LiveHandle,
89    control: ReplayControl,
90    clock: Arc<ManualClock>,
91}
92
93impl ReplaySession {
94    /// The live handle — same type a real connection returns. `state()`,
95    /// `events()`, `telemetry()`, `extracted()` all work.
96    pub fn handle(&self) -> &LiveHandle {
97        &self.handle
98    }
99
100    /// Start streaming the recorded inbound frames. Call after subscribing
101    /// to [`LiveHandle::events`].
102    pub fn release(&self) {
103        self.control.release();
104    }
105
106    /// Wait until every recorded inbound frame has been handed to the session
107    /// loop. The last frame's effects may still be propagating through the
108    /// processor — use [`collect_events_until_idle`] (or assert on state) to
109    /// settle.
110    pub async fn drained(&self) {
111        self.control.drained().await;
112    }
113
114    /// Outbound frames the replayed session has sent so far (setup, tool
115    /// responses, …), in send order, for comparison against the recorded log.
116    pub fn outbound_frames(&self) -> Vec<Vec<u8>> {
117        self.control.outbound_frames()
118    }
119
120    /// The replay's clock. It stands at the recorded capture time of the most
121    /// recently delivered frame, measured from the first inbound frame.
122    pub fn clock(&self) -> &Arc<ManualClock> {
123        &self.clock
124    }
125
126    /// Disconnect the replayed session.
127    pub async fn disconnect(&self) -> Result<(), gemini_genai_rs::session::SessionError> {
128        self.handle.disconnect().await
129    }
130}
131
132/// Replay a recorded wire log through the real L1 processor, offline.
133///
134/// - `config` is used to open the replay transport (its re-encoded setup
135///   message becomes the first outbound frame, mirroring the original run).
136///   Use the same configuration as the recorded session for a faithful setup
137///   comparison. No network is touched and no credential is used.
138/// - `builder` supplies the control plane: dispatcher, phases, extractors,
139///   watchers, state, callbacks. Attach the original tool implementations to
140///   re-execute tools deterministically; without a dispatcher, recorded tool
141///   calls surface as events but produce no responses.
142/// - `entries` is the recorded log; only its inbound frames are replayed
143///   (outbound entries are kept in the log purely for comparison/audit).
144///
145/// Frames are delivered as fast as the session consumes them, but the
146/// session reads time from a [`ManualClock`] that moves to each frame's
147/// recorded capture time as it is delivered (this replaces any clock the
148/// builder set). Temporal patterns, phase durations, resolver cache expiry
149/// and the `session:` timing signals therefore see the original gaps between
150/// frames, whatever the replay's own speed.
151///
152/// The session processes frames asynchronously, so before the clock jumps
153/// by [`REPLAY_SYNC_STEP`] or more, the next frame waits until both event
154/// lanes have handled every event of the frames before it (the router runs
155/// in lockstep during a replay). An earlier frame is then evaluated at its
156/// own time, not a later one. Jumps smaller than the step are not waited
157/// for, which keeps dense audio fast and bounds the timing error at the
158/// step. The replay is gated: nothing past
159/// the setup handshake flows until [`ReplaySession::release`] is called, so
160/// subscribe to events first.
161pub async fn replay_session(
162    config: SessionConfig,
163    builder: LiveSessionBuilder,
164    entries: &[WireEntry],
165) -> Result<ReplaySession, AgentError> {
166    let first_ts_ms = entries
167        .iter()
168        .find(|e| e.dir == WireDirection::Inbound)
169        .map_or(0, |e| e.ts_ms);
170    let clock = Arc::new(ManualClock::starting_at(
171        UNIX_EPOCH + Duration::from_millis(first_ts_ms),
172    ));
173    // Once the session exists, the gate counts the events it has emitted
174    // and waits for the lanes to have handled them all (lockstep).
175    let lockstep = Arc::new(super::processor::Lockstep::default());
176    let emitted: Arc<tokio::sync::Mutex<Option<EmittedCount>>> = Arc::default();
177    let synced_ms = Arc::new(std::sync::atomic::AtomicU64::new(first_ts_ms));
178    let gate = {
179        let clock = clock.clone();
180        let emitted = emitted.clone();
181        let lockstep = lockstep.clone();
182        Arc::new(move |ts_ms: u64| {
183            let clock = clock.clone();
184            let emitted = emitted.clone();
185            let lockstep = lockstep.clone();
186            let synced_ms = synced_ms.clone();
187            Box::pin(async move {
188                let step = REPLAY_SYNC_STEP.as_millis() as u64;
189                let synced = synced_ms.load(std::sync::atomic::Ordering::Acquire);
190                if ts_ms.saturating_sub(synced) >= step
191                    && let Some((events, count)) = emitted.lock().await.as_mut()
192                {
193                    // The session loop takes the next frame only after it has
194                    // emitted every event of the previous one, so these are
195                    // all the events so far.
196                    loop {
197                        match events.try_recv() {
198                            Ok(_) => *count += 1,
199                            Err(broadcast::error::TryRecvError::Lagged(n)) => *count += n,
200                            Err(_) => break,
201                        }
202                    }
203                    lockstep.wait_for(*count, Duration::from_secs(5)).await;
204                    synced_ms.store(ts_ms, std::sync::atomic::Ordering::Release);
205                }
206                clock.set_elapsed(Duration::from_millis(ts_ms.saturating_sub(first_ts_ms)));
207            }) as std::pin::Pin<Box<dyn std::future::Future<Output = ()> + Send>>
208        })
209    };
210    let (transport, control) = ReplayTransport::from_wire_log(entries);
211    let transport = transport.with_frame_gate(gate);
212    let builder = builder.clock(clock.clone()).lockstep(lockstep);
213    let transport_config = TransportConfig {
214        max_reconnect_attempts: 0,
215        connect_timeout_secs: 5,
216        setup_timeout_secs: 5,
217        ..TransportConfig::default()
218    };
219    let session = ConnectBuilder::new(config)
220        .transport_config(transport_config)
221        .transport(transport)
222        .connect()
223        .await
224        .map_err(AgentError::Session)?;
225    let handle = attach_session(builder, session).await?;
226    // Count from here, after the router subscribed: it settles everything it
227    // receives, so a count started later can only trail it. Nothing past the
228    // handshake flows before `release`, so no frame's events are missed.
229    *emitted.lock().await = Some((handle.session().subscribe(), 0));
230    Ok(ReplaySession {
231        handle,
232        control,
233        clock,
234    })
235}
236
237/// A subscription to the session's events and how many it has seen.
238type EmittedCount = (broadcast::Receiver<SessionEvent>, u64);
239
240/// The largest clock jump a replay makes without first waiting for the
241/// session to process the frames already delivered. See [`replay_session`].
242pub const REPLAY_SYNC_STEP: Duration = Duration::from_millis(50);
243
244/// Collect [`LiveEvent`]s until the stream stays idle for `idle` (or `max`
245/// elapses). Useful for settling an as-fast-as-possible replay where "done"
246/// means "no more effects are propagating".
247pub async fn collect_events_until_idle(
248    rx: &mut tokio::sync::broadcast::Receiver<LiveEvent>,
249    idle: Duration,
250    max: Duration,
251) -> Vec<LiveEvent> {
252    let mut events = Vec::new();
253    let deadline = tokio::time::Instant::now() + max;
254    loop {
255        let timeout = idle.min(deadline.saturating_duration_since(tokio::time::Instant::now()));
256        if timeout.is_zero() {
257            break;
258        }
259        match tokio::time::timeout(timeout, rx.recv()).await {
260            Ok(Ok(event)) => events.push(event),
261            Ok(Err(tokio::sync::broadcast::error::RecvError::Lagged(_))) => continue,
262            Ok(Err(tokio::sync::broadcast::error::RecvError::Closed)) => break,
263            Err(_) => break, // idle
264        }
265    }
266    events
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272    use gemini_genai_rs::prelude::ModelId;
273
274    #[tokio::test]
275    async fn replay_session_reaches_active_and_emits_events() {
276        let entries = vec![
277            WireEntry {
278                seq: 1,
279                dir: WireDirection::Inbound,
280                ts_ms: 1,
281                payload: br#"{"setupComplete":{}}"#.to_vec(),
282            },
283            WireEntry {
284                seq: 2,
285                dir: WireDirection::Inbound,
286                ts_ms: 2,
287                payload:
288                    br#"{"serverContent":{"modelTurn":{"parts":[{"text":"Hi"}]},"turnComplete":true}}"#
289                        .to_vec(),
290            },
291        ];
292        let config = SessionConfig::new("offline").model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
293        let builder = LiveSessionBuilder::new(config.clone());
294
295        let replay = replay_session(config, builder, &entries).await.unwrap();
296        let mut events = replay.handle().events();
297        replay.release();
298        replay.drained().await;
299
300        let collected = collect_events_until_idle(
301            &mut events,
302            Duration::from_millis(200),
303            Duration::from_secs(5),
304        )
305        .await;
306
307        assert!(
308            collected
309                .iter()
310                .any(|e| matches!(e, LiveEvent::TextDelta(t) if t == "Hi")),
311            "expected replayed TextDelta, got {collected:?}"
312        );
313        assert!(
314            collected
315                .iter()
316                .any(|e| matches!(e, LiveEvent::TurnComplete))
317        );
318
319        // The clock stands at the last frame's recorded time.
320        assert_eq!(replay.clock().elapsed(), Duration::from_millis(1));
321
322        // The replayed session re-encoded and "sent" the setup message.
323        let outbound = replay.outbound_frames();
324        assert!(!outbound.is_empty());
325        assert!(
326            String::from_utf8(outbound[0].clone())
327                .unwrap()
328                .contains("\"setup\"")
329        );
330
331        replay.disconnect().await.unwrap();
332    }
333
334    /// Timing decisions follow the recording, not the replay's speed: a turn
335    /// recorded ten seconds after the handshake is stamped ten seconds later,
336    /// at the recording's wall time, even though the replay takes
337    /// milliseconds.
338    #[tokio::test]
339    async fn replay_reads_time_from_the_recording() {
340        const T0: u64 = 1_000_000_000_000; // 2001-09-09, well before "now"
341        let entries = vec![
342            WireEntry {
343                seq: 1,
344                dir: WireDirection::Inbound,
345                ts_ms: T0,
346                payload: br#"{"setupComplete":{}}"#.to_vec(),
347            },
348            WireEntry {
349                seq: 2,
350                dir: WireDirection::Inbound,
351                ts_ms: T0 + 10_000,
352                payload:
353                    br#"{"serverContent":{"modelTurn":{"parts":[{"text":"Hi"}]},"turnComplete":true}}"#
354                        .to_vec(),
355            },
356        ];
357        let config = SessionConfig::new("offline").model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
358        let replay = replay_session(config.clone(), LiveSessionBuilder::new(config), &entries)
359            .await
360            .unwrap();
361        let mut events = replay.handle().events();
362        replay.release();
363        replay.drained().await;
364        collect_events_until_idle(
365            &mut events,
366            Duration::from_millis(200),
367            Duration::from_secs(5),
368        )
369        .await;
370
371        assert_eq!(replay.clock().elapsed(), Duration::from_secs(10));
372        let mutations = replay.handle().state().recent_mutations();
373        assert!(!mutations.is_empty(), "the turn should write state");
374        let recorded = |ms: u64| UNIX_EPOCH + Duration::from_millis(ms);
375        for m in &mutations {
376            assert!(
377                m.timestamp >= recorded(T0) && m.timestamp <= recorded(T0 + 10_000),
378                "{} stamped outside the recording: {:?}",
379                m.key,
380                m.timestamp
381            );
382        }
383        assert!(
384            mutations
385                .iter()
386                .any(|m| m.timestamp == recorded(T0 + 10_000)),
387            "the turn's writes carry the second frame's recorded time"
388        );
389
390        replay.disconnect().await.unwrap();
391    }
392
393    #[tokio::test]
394    async fn each_frame_is_processed_at_its_own_recorded_time() {
395        const T0: u64 = 1_000_000_000_000;
396        let turn =
397            br#"{"serverContent":{"modelTurn":{"parts":[{"text":"Hi"}]},"turnComplete":true}}"#;
398        let entry = |seq, ts_ms, payload: &[u8]| WireEntry {
399            seq,
400            dir: WireDirection::Inbound,
401            ts_ms,
402            payload: payload.to_vec(),
403        };
404        let entries = vec![
405            entry(1, T0, br#"{"setupComplete":{}}"#),
406            entry(2, T0 + 1_000, turn),
407            entry(3, T0 + 60_000, turn),
408        ];
409        // A slow turn-complete handler: the lane is still on the first turn
410        // when the replay would otherwise deliver the second.
411        let state = crate::state::State::new();
412        let seen = state.clone();
413        let callbacks = crate::live::callbacks::EventCallbacks {
414            on_turn_complete: Some(Arc::new(move || {
415                let seen = seen.clone();
416                Box::pin(async move {
417                    tokio::time::sleep(Duration::from_millis(100)).await;
418                    let n = seen.get::<u32>("turns_seen").unwrap_or(0);
419                    let _ = seen.set("turns_seen", n + 1);
420                })
421            })),
422            ..Default::default()
423        };
424        let config = SessionConfig::new("offline").model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
425        let builder = LiveSessionBuilder::new(config.clone())
426            .state(state)
427            .callbacks(callbacks);
428        let replay = replay_session(config, builder, &entries).await.unwrap();
429        let mut events = replay.handle().events();
430        replay.release();
431        replay.drained().await;
432        collect_events_until_idle(
433            &mut events,
434            Duration::from_millis(400),
435            Duration::from_secs(5),
436        )
437        .await;
438
439        let recorded = |ms: u64| UNIX_EPOCH + Duration::from_millis(ms);
440        let stamps: Vec<_> = replay
441            .handle()
442            .state()
443            .recent_mutations()
444            .iter()
445            .filter(|m| m.key == "turns_seen")
446            .map(|m| m.timestamp)
447            .collect();
448        assert_eq!(
449            stamps,
450            [recorded(T0 + 1_000), recorded(T0 + 60_000)],
451            "each turn's handler runs at its own recorded time"
452        );
453        replay.disconnect().await.unwrap();
454    }
455}
gemini_adk_rs/live/replay.rs

gemini_adk_rs/live/
replay.rs
