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::time::Duration;
45
46use gemini_genai_rs::prelude::{SessionConfig, SessionPhase};
47use gemini_genai_rs::session::SessionHandle;
48use gemini_genai_rs::transport::replay::{ReplayControl, ReplayTransport};
49use gemini_genai_rs::transport::{ConnectBuilder, TransportConfig, WireEntry};
50
51use crate::error::AgentError;
52
53use super::builder::{LiveSessionBuilder, build_runtime, spawn_lanes};
54use super::events::LiveEvent;
55use super::handle::LiveHandle;
56
57/// Attach the full L1 control plane (three-lane processor, phase machine,
58/// extractors, watchers, tool dispatch, …) to an **already connected** L0
59/// session.
60///
61/// This is the seam that makes replay possible without touching the network:
62/// connect the L0 session over any [`Transport`](gemini_genai_rs::transport::Transport)
63/// (e.g. [`ReplayTransport`] or [`MockTransport`](gemini_genai_rs::transport::MockTransport)), then hand
64/// it here together with a configured [`LiveSessionBuilder`].
65///
66/// Note: the builder's own `SessionConfig` is *not* re-sent — the setup
67/// message was already encoded from the config given to the L0 connect call.
68/// Subscribe to events **after** this returns and only then let the transport
69/// stream (for `ReplayTransport`, call
70/// [`ReplayControl::release`](gemini_genai_rs::transport::replay::ReplayControl::release)),
71/// otherwise early frames race the subscription.
72pub async fn attach_session(
73    builder: LiveSessionBuilder,
74    session: SessionHandle,
75) -> Result<LiveHandle, AgentError> {
76    let plan = builder.into_plan()?;
77    session.wait_for_phase(SessionPhase::Active).await;
78    let runtime = build_runtime(plan, session);
79    spawn_lanes(runtime).await
80}
81
82/// A replayed session: the live handle plus the replay controls.
83pub struct ReplaySession {
84    handle: LiveHandle,
85    control: ReplayControl,
86}
87
88impl ReplaySession {
89    /// The live handle — same type a real connection returns. `state()`,
90    /// `events()`, `telemetry()`, `extracted()` all work.
91    pub fn handle(&self) -> &LiveHandle {
92        &self.handle
93    }
94
95    /// Start streaming the recorded inbound frames. Call after subscribing
96    /// to [`LiveHandle::events`].
97    pub fn release(&self) {
98        self.control.release();
99    }
100
101    /// Wait until every recorded inbound frame has been handed to the session
102    /// loop. The last frame's effects may still be propagating through the
103    /// processor — use [`collect_events_until_idle`] (or assert on state) to
104    /// settle.
105    pub async fn drained(&self) {
106        self.control.drained().await;
107    }
108
109    /// Outbound frames the replayed session has sent so far (setup, tool
110    /// responses, …), in send order, for comparison against the recorded log.
111    pub fn outbound_frames(&self) -> Vec<Vec<u8>> {
112        self.control.outbound_frames()
113    }
114
115    /// Disconnect the replayed session.
116    pub async fn disconnect(&self) -> Result<(), gemini_genai_rs::session::SessionError> {
117        self.handle.disconnect().await
118    }
119}
120
121/// Replay a recorded wire log through the real L1 processor, offline.
122///
123/// - `config` is used to open the replay transport (its re-encoded setup
124///   message becomes the first outbound frame, mirroring the original run).
125///   Use the same configuration as the recorded session for a faithful setup
126///   comparison. No network is touched and no credential is used.
127/// - `builder` supplies the control plane: dispatcher, phases, extractors,
128///   watchers, state, callbacks. Attach the original tool implementations to
129///   re-execute tools deterministically; without a dispatcher, recorded tool
130///   calls surface as events but produce no responses.
131/// - `entries` is the recorded log; only its inbound frames are replayed
132///   (outbound entries are kept in the log purely for comparison/audit).
133///
134/// Frames are delivered as fast as the session loop consumes them (no
135/// original-timing pacing). The replay is gated: nothing past the setup
136/// handshake flows until [`ReplaySession::release`] is called, so subscribe
137/// to events first.
138pub async fn replay_session(
139    config: SessionConfig,
140    builder: LiveSessionBuilder,
141    entries: &[WireEntry],
142) -> Result<ReplaySession, AgentError> {
143    let (transport, control) = ReplayTransport::from_wire_log(entries);
144    let transport_config = TransportConfig {
145        max_reconnect_attempts: 0,
146        connect_timeout_secs: 5,
147        setup_timeout_secs: 5,
148        ..TransportConfig::default()
149    };
150    let session = ConnectBuilder::new(config)
151        .transport_config(transport_config)
152        .transport(transport)
153        .connect()
154        .await
155        .map_err(AgentError::Session)?;
156    let handle = attach_session(builder, session).await?;
157    Ok(ReplaySession { handle, control })
158}
159
160/// Collect [`LiveEvent`]s until the stream stays idle for `idle` (or `max`
161/// elapses). Useful for settling an as-fast-as-possible replay where "done"
162/// means "no more effects are propagating".
163pub async fn collect_events_until_idle(
164    rx: &mut tokio::sync::broadcast::Receiver<LiveEvent>,
165    idle: Duration,
166    max: Duration,
167) -> Vec<LiveEvent> {
168    let mut events = Vec::new();
169    let deadline = tokio::time::Instant::now() + max;
170    loop {
171        let timeout = idle.min(deadline.saturating_duration_since(tokio::time::Instant::now()));
172        if timeout.is_zero() {
173            break;
174        }
175        match tokio::time::timeout(timeout, rx.recv()).await {
176            Ok(Ok(event)) => events.push(event),
177            Ok(Err(tokio::sync::broadcast::error::RecvError::Lagged(_))) => continue,
178            Ok(Err(tokio::sync::broadcast::error::RecvError::Closed)) => break,
179            Err(_) => break, // idle
180        }
181    }
182    events
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188    use gemini_genai_rs::prelude::ModelId;
189    use gemini_genai_rs::transport::WireDirection;
190
191    #[tokio::test]
192    async fn replay_session_reaches_active_and_emits_events() {
193        let entries = vec![
194            WireEntry {
195                seq: 1,
196                dir: WireDirection::Inbound,
197                ts_ms: 1,
198                payload: br#"{"setupComplete":{}}"#.to_vec(),
199            },
200            WireEntry {
201                seq: 2,
202                dir: WireDirection::Inbound,
203                ts_ms: 2,
204                payload:
205                    br#"{"serverContent":{"modelTurn":{"parts":[{"text":"Hi"}]},"turnComplete":true}}"#
206                        .to_vec(),
207            },
208        ];
209        let config = SessionConfig::new("offline").model(ModelId::LIVE_2_5_FLASH_NATIVE_AUDIO);
210        let builder = LiveSessionBuilder::new(config.clone());
211
212        let replay = replay_session(config, builder, &entries).await.unwrap();
213        let mut events = replay.handle().events();
214        replay.release();
215        replay.drained().await;
216
217        let collected = collect_events_until_idle(
218            &mut events,
219            Duration::from_millis(200),
220            Duration::from_secs(5),
221        )
222        .await;
223
224        assert!(
225            collected
226                .iter()
227                .any(|e| matches!(e, LiveEvent::TextDelta(t) if t == "Hi")),
228            "expected replayed TextDelta, got {collected:?}"
229        );
230        assert!(
231            collected
232                .iter()
233                .any(|e| matches!(e, LiveEvent::TurnComplete))
234        );
235
236        // The replayed session re-encoded and "sent" the setup message.
237        let outbound = replay.outbound_frames();
238        assert!(!outbound.is_empty());
239        assert!(
240            String::from_utf8(outbound[0].clone())
241                .unwrap()
242                .contains("\"setup\"")
243        );
244
245        replay.disconnect().await.unwrap();
246    }
247}
gemini_adk_rs/live/replay.rs

gemini_adk_rs/live/
replay.rs
