gemini_adk_fluent_rs/live/
scripted.rs

1//! A scripted Gemini Live server for offline session tests; see
2//! [`ScriptedServer`].
3
4use std::time::Duration;
5
6use gemini_adk_rs::State;
7use gemini_adk_rs::error::AgentError;
8use gemini_adk_rs::live::replay::collect_events_until_idle;
9use gemini_adk_rs::live::{LiveEvent, LiveHandle};
10use gemini_genai_rs::transport::{ReplayControl, ReplayTransport};
11use serde_json::{Value, json};
12
13use super::Live;
14
15/// A scripted Gemini Live server, for testing a [`Live`] session offline.
16///
17/// [`ScriptedServer`] is a script of what the server says: model text, what
18/// it heard, tool calls, interruptions. [`play`](ScriptedServer::play)
19/// connects your fully configured [`Live`] builder to it over an in-memory
20/// transport and runs the script. Your tools, phases, extractors, watchers,
21/// governance and callbacks all run for real. What comes back is a
22/// [`ScriptedRun`]: the events the session emitted, what it sent to the
23/// "server" (setup, tool responses, context), and its state.
24///
25/// No network or credential is used, and the model is not involved. The
26/// script is its stand-in.
27///
28/// ```
29/// # tokio_test::block_on(async {
30/// use gemini_adk_fluent_rs::prelude::*;
31/// use gemini_adk_fluent_rs::testing::ScriptedServer;
32/// use serde_json::json;
33///
34/// let run = ScriptedServer::new()
35///     .hears("What's the weather in Paris?")
36///     .calls("get_weather", json!({ "city": "Paris" }))
37///     .says("It is sunny in Paris.")
38///     .play(
39///         Live::builder()
40///             .instruction("You are a weather assistant.")
41///             .tools(T::simple("get_weather", "Weather for a city", |args| async move {
42///                 Ok(json!({ "city": args["city"], "sky": "sunny" }))
43///             })),
44///     )
45///     .await
46///     .unwrap();
47///
48/// let responses = run.tool_responses();
49/// assert_eq!(responses[0]["name"], "get_weather");
50/// assert_eq!(responses[0]["response"]["sky"], "sunny");
51/// assert!(run.transcript_text().contains("It is sunny"));
52/// run.disconnect().await;
53/// # });
54/// ```
55#[derive(Debug, Clone)]
56pub struct ScriptedServer {
57    frames: Vec<Value>,
58    next_call: usize,
59    idle: Duration,
60}
61
62impl Default for ScriptedServer {
63    fn default() -> Self {
64        Self::new()
65    }
66}
67
68impl ScriptedServer {
69    /// A script that starts with the setup handshake.
70    pub fn new() -> Self {
71        Self {
72            frames: vec![json!({ "setupComplete": {} })],
73            next_call: 0,
74            idle: Duration::from_millis(200),
75        }
76    }
77
78    /// The model says `text` and ends its turn.
79    pub fn says(self, text: impl Into<String>) -> Self {
80        self.text(text).turn_complete()
81    }
82
83    /// The model says `text` without ending its turn.
84    pub fn text(self, text: impl Into<String>) -> Self {
85        self.frame(json!({
86            "serverContent": { "modelTurn": { "parts": [{ "text": text.into() }] } }
87        }))
88    }
89
90    /// The model ends its turn.
91    pub fn turn_complete(self) -> Self {
92        self.frame(json!({ "serverContent": { "turnComplete": true } }))
93    }
94
95    /// The server transcribes the user saying `text`.
96    pub fn hears(self, text: impl Into<String>) -> Self {
97        self.frame(json!({
98            "serverContent": { "inputTranscription": { "text": text.into() } }
99        }))
100    }
101
102    /// The server transcribes the model's speech as `text`.
103    pub fn speaks(self, text: impl Into<String>) -> Self {
104        self.frame(json!({
105            "serverContent": { "outputTranscription": { "text": text.into() } }
106        }))
107    }
108
109    /// The model calls tool `name` with `args`. Call ids are `call-1`,
110    /// `call-2`, … in script order.
111    pub fn calls(mut self, name: impl Into<String>, args: Value) -> Self {
112        self.next_call += 1;
113        let id = format!("call-{}", self.next_call);
114        self.frame(json!({
115            "toolCall": { "functionCalls": [{ "name": name.into(), "args": args, "id": id }] }
116        }))
117    }
118
119    /// The user barges in: the server reports the model's turn interrupted.
120    pub fn interrupts(self) -> Self {
121        self.frame(json!({ "serverContent": { "interrupted": true } }))
122    }
123
124    /// Any raw server message, for what the helpers above do not cover.
125    pub fn frame(mut self, message: Value) -> Self {
126        self.frames.push(message);
127        self
128    }
129
130    /// How long the session must stay quiet after the last frame before the
131    /// run counts as settled (default 200 ms).
132    pub fn settle_after(mut self, idle: Duration) -> Self {
133        self.idle = idle;
134        self
135    }
136
137    /// The ids of the scripted tool calls the session must answer: every
138    /// call not listed in a later cancellation.
139    fn awaited_call_ids(&self) -> Vec<String> {
140        let ids = |frame: &Value, pointer: &str| -> Vec<String> {
141            frame
142                .pointer(pointer)
143                .and_then(Value::as_array)
144                .into_iter()
145                .flatten()
146                .filter_map(|v| v.get("id").or(Some(v)).and_then(Value::as_str))
147                .map(str::to_owned)
148                .collect()
149        };
150        let cancelled: Vec<String> = self
151            .frames
152            .iter()
153            .flat_map(|f| ids(f, "/toolCallCancellation/ids"))
154            .collect();
155        self.frames
156            .iter()
157            .flat_map(|f| ids(f, "/toolCall/functionCalls"))
158            .filter(|id| !cancelled.contains(id))
159            .collect()
160    }
161
162    /// The script as an in-memory transport plus its control handle, for a
163    /// test that drives the connection itself. Nothing past the handshake
164    /// flows until [`ReplayControl::release`].
165    pub fn into_transport(self) -> (ReplayTransport, ReplayControl) {
166        ReplayTransport::from_frames(
167            self.frames
168                .iter()
169                .map(|frame| frame.to_string().into_bytes())
170                .collect(),
171        )
172    }
173
174    /// Connect `live` to this script, play every frame, and return once the
175    /// session has settled: every scripted tool call that was not cancelled
176    /// has been answered, and no event arrived for the
177    /// [`settle_after`](Self::settle_after) window.
178    pub async fn play(self, live: Live) -> Result<ScriptedRun, AgentError> {
179        let idle = self.idle;
180        let mut awaiting = self.awaited_call_ids();
181        let (transport, control) = self.into_transport();
182        let handle = live.connect_with_transport(transport).await?;
183        let mut rx = handle.events();
184        control.release();
185        control.drained().await;
186        // A slow tool emits nothing while it runs, so a quiet window alone
187        // can end the run before the tool answers.
188        let deadline = tokio::time::Instant::now() + Duration::from_secs(30);
189        let mut events = Vec::new();
190        loop {
191            let left = deadline.saturating_duration_since(tokio::time::Instant::now());
192            events.extend(collect_events_until_idle(&mut rx, idle, left).await);
193            awaiting.retain(|id| !answered(&control, id));
194            if awaiting.is_empty() || left.is_zero() {
195                break;
196            }
197        }
198        Ok(ScriptedRun {
199            handle,
200            events,
201            control,
202        })
203    }
204}
205
206/// Whether the session has sent a function response for call `id`.
207fn answered(control: &ReplayControl, id: &str) -> bool {
208    control.outbound_frames().iter().any(|frame| {
209        serde_json::from_slice::<Value>(frame)
210            .ok()
211            .and_then(|m| m.pointer("/toolResponse/functionResponses").cloned())
212            .and_then(|r| r.as_array().cloned())
213            .is_some_and(|rs| rs.iter().any(|r| r["id"] == id))
214    })
215}
216
217/// The outcome of [`ScriptedServer::play`].
218pub struct ScriptedRun {
219    handle: LiveHandle,
220    events: Vec<LiveEvent>,
221    control: ReplayControl,
222}
223
224impl ScriptedRun {
225    /// The live session, still connected.
226    pub fn handle(&self) -> &LiveHandle {
227        &self.handle
228    }
229
230    /// The session's state.
231    pub fn state(&self) -> &State {
232        self.handle.state()
233    }
234
235    /// Every event the session emitted while the script played.
236    pub fn events(&self) -> &[LiveEvent] {
237        &self.events
238    }
239
240    /// The model text the session delivered, concatenated.
241    pub fn transcript_text(&self) -> String {
242        self.events
243            .iter()
244            .filter_map(|event| match event {
245                LiveEvent::TextDelta(text) => Some(text.as_str()),
246                _ => None,
247            })
248            .collect()
249    }
250
251    /// Every message the session sent to the server, parsed, in order: the
252    /// setup first, then tool responses, context and anything else.
253    pub fn sent(&self) -> Vec<Value> {
254        self.control
255            .outbound_frames()
256            .iter()
257            .filter_map(|frame| serde_json::from_slice(frame).ok())
258            .collect()
259    }
260
261    /// The setup message the session opened with.
262    pub fn setup(&self) -> Option<Value> {
263        self.sent()
264            .into_iter()
265            .find_map(|m| m.get("setup").cloned())
266    }
267
268    /// Every function response the session sent, in order.
269    pub fn tool_responses(&self) -> Vec<Value> {
270        self.sent()
271            .iter()
272            .filter_map(|m| m.pointer("/toolResponse/functionResponses"))
273            .filter_map(Value::as_array)
274            .flatten()
275            .cloned()
276            .collect()
277    }
278
279    /// Disconnect the session.
280    pub async fn disconnect(self) {
281        let _ = self.handle.disconnect().await;
282    }
283}
gemini_adk_fluent_rs/live/scripted.rs

gemini_adk_fluent_rs/live/
scripted.rs
