gemini_adk_rs/live/
session_signals.rs

1//! Auto-tracked session-level state signals.
2//!
3//! [`SessionSignals`] is called on every [`SessionEvent`] and transparently
4//! updates keys under the `session:` prefix in the shared [`State`]. The
5//! router touches the activity timestamp for audio and text inline; the
6//! telemetry lane handles every event that writes state.
7//!
8//! Hot-path timestamps use [`AtomicU64`] (nanos since start) instead of
9//! `Mutex<Instant>`, eliminating per-event mutex contention. Derived
10//! timing signals (`silence_ms`, `elapsed_ms`, `remaining_budget_ms`)
11//! are flushed periodically via `flush_timing()` rather than on every event.
12
13use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
14use std::time::Instant;
15
16use crate::clock::SharedClock;
17
18use gemini_genai_rs::prelude::{SessionEvent, SessionPhase};
19use parking_lot::Mutex;
20
21use crate::state::State;
22
23// ---------------------------------------------------------------------------
24// Public types
25// ---------------------------------------------------------------------------
26
27/// Session type determines the server-side duration limit.
28#[derive(Debug, Clone, Copy, PartialEq, Eq)]
29pub enum SessionType {
30    /// Audio-only session (~15 min limit).
31    AudioOnly,
32    /// Audio + video session (~2 min limit).
33    AudioVideo,
34}
35
36// ---------------------------------------------------------------------------
37// SessionSignals
38// ---------------------------------------------------------------------------
39
40/// Tracks session-level signals automatically from events.
41///
42/// Every call to [`on_event`](SessionSignals::on_event) updates the
43/// corresponding keys under `session:` in the shared [`State`], making
44/// them available to instruction templates, watchers, and computed vars.
45///
46/// **Performance**: Timestamps use `AtomicU64` (nanos since session start)
47/// instead of `Mutex<Instant>`. Derived timing signals are flushed
48/// periodically via `flush_timing()` (100ms interval) rather than per-event.
49pub struct SessionSignals {
50    state: State,
51    /// Session start time — used as epoch for all atomic timestamps.
52    start: Instant,
53    /// The state's clock, captured at construction.
54    clock: SharedClock,
55    /// Nanos since start when connected.
56    connected_at_ns: AtomicU64,
57    /// Whether currently connected.
58    is_connected: AtomicBool,
59    /// Nanos since start of last activity.
60    last_activity_ns: AtomicU64,
61    /// Whether the session includes video input.
62    has_video: AtomicBool,
63    /// Server-sent GoAway timestamp, if received.
64    go_away_at: Mutex<Option<Instant>>,
65    /// Latest resumption handle from server (persisted for reconnection).
66    latest_resume_handle: Mutex<Option<String>>,
67}
68
69impl SessionSignals {
70    /// Create a new `SessionSignals` backed by the given [`State`].
71    pub fn new(state: State) -> Self {
72        let clock = state.clock();
73        Self {
74            state,
75            start: clock.now(),
76            clock,
77            connected_at_ns: AtomicU64::new(0),
78            is_connected: AtomicBool::new(false),
79            last_activity_ns: AtomicU64::new(0),
80            has_video: AtomicBool::new(false),
81            go_away_at: Mutex::new(None),
82            latest_resume_handle: Mutex::new(None),
83        }
84    }
85
86    /// Process an event — updates state keys and atomic timestamps.
87    ///
88    /// This is the per-event handler. It updates boolean flags, counters,
89    /// and atomic timestamps. **Derived timing** (silence_ms, elapsed_ms,
90    /// remaining_budget_ms) is NOT computed here — call `flush_timing()`
91    /// periodically instead.
92    pub fn on_event(&self, event: &SessionEvent) {
93        match event {
94            SessionEvent::Connected => {
95                let now_ns = self.elapsed_ns();
96                self.connected_at_ns.store(now_ns, Ordering::Relaxed);
97                self.is_connected.store(true, Ordering::Relaxed);
98                self.last_activity_ns.store(now_ns, Ordering::Relaxed);
99                let _ = self.state.session().set("connected_at_ms", 0u64);
100                let _ = self.state.session().set("interrupt_count", 0u64);
101                let _ = self.state.session().set("error_count", 0u64);
102                let _ = self.state.session().set("is_user_speaking", false);
103                let _ = self.state.session().set("is_model_speaking", false);
104                let _ = self.state.session().set("go_away_received", false);
105                let _ = self.state.session().set("resumable", false);
106                let _ = self.state.session().set("session_type", "audio_only");
107            }
108
109            SessionEvent::VoiceActivityStart => {
110                let _ = self.state.session().set("is_user_speaking", true);
111                self.touch_activity();
112            }
113
114            SessionEvent::VoiceActivityEnd => {
115                let _ = self.state.session().set("is_user_speaking", false);
116                self.touch_activity();
117            }
118
119            SessionEvent::Interrupted => {
120                let count: u64 = self.state.session().get("interrupt_count").unwrap_or(0);
121                let _ = self.state.session().set("interrupt_count", count + 1);
122                self.touch_activity();
123            }
124
125            SessionEvent::Error(msg) => {
126                let count: u64 = self.state.session().get("error_count").unwrap_or(0);
127                let _ = self.state.session().set("error_count", count + 1);
128                let _ = self.state.session().set("last_error", msg.to_string());
129            }
130
131            SessionEvent::PhaseChanged(phase) => {
132                let _ = self
133                    .state
134                    .session()
135                    .set("is_model_speaking", *phase == SessionPhase::ModelSpeaking);
136                let _ = self.state.session().set("phase", phase.to_string());
137                self.touch_activity();
138            }
139
140            SessionEvent::GoAway(time_left) => {
141                let _ = self.state.session().set("go_away_received", true);
142                if let Some(tl) = time_left {
143                    *self.go_away_at.lock() = Some(self.clock.now() + *tl);
144                    let _ = self
145                        .state
146                        .session()
147                        .set("go_away_time_left_ms", tl.as_millis() as u64);
148                }
149            }
150
151            SessionEvent::SessionResumeUpdate(info) => {
152                *self.latest_resume_handle.lock() = Some(info.handle.clone());
153                let _ = self.state.session().set("resumable", info.resumable);
154                if let Some(ref idx) = info.last_consumed_index {
155                    let _ = self
156                        .state
157                        .session()
158                        .set("last_consumed_client_index", idx.clone());
159                }
160            }
161
162            SessionEvent::Usage(usage) => {
163                if let Some(total) = usage.total_token_count {
164                    let _ = self.state.session().set("total_token_count", total);
165                }
166                if let Some(prompt) = usage.prompt_token_count {
167                    let _ = self.state.session().set("prompt_token_count", prompt);
168                }
169                if let Some(response) = usage.response_token_count {
170                    let _ = self.state.session().set("response_token_count", response);
171                }
172                if let Some(cached) = usage.cached_content_token_count {
173                    let _ = self
174                        .state
175                        .session()
176                        .set("cached_content_token_count", cached);
177                }
178                if let Some(thoughts) = usage.thoughts_token_count {
179                    let _ = self.state.session().set("thoughts_token_count", thoughts);
180                }
181            }
182
183            SessionEvent::GenerationComplete => {
184                // No-op for signals — generation complete is handled by control lane
185            }
186
187            SessionEvent::InputTranscription(text) => {
188                let _ = self
189                    .state
190                    .session()
191                    .set("last_input_transcription", text.clone());
192                self.touch_activity();
193            }
194
195            SessionEvent::OutputTranscription(text) => {
196                let _ = self
197                    .state
198                    .session()
199                    .set("last_output_transcription", text.clone());
200                self.touch_activity();
201            }
202
203            SessionEvent::AudioData(_)
204            | SessionEvent::TextDelta(_)
205            | SessionEvent::TextComplete(_) => {
206                // High-frequency events: only touch the atomic timestamp.
207                // No DashMap writes, no mutex locks.
208                self.touch_activity();
209            }
210
211            SessionEvent::TurnComplete => {
212                self.touch_activity();
213            }
214
215            SessionEvent::Disconnected(_reason) => {
216                self.is_connected.store(false, Ordering::Relaxed);
217                let _ = self.state.session().set("disconnected", true);
218            }
219
220            _ => {}
221        }
222    }
223
224    /// Flush derived timing signals to state.
225    ///
226    /// Call this periodically (e.g., every 100ms) from the telemetry lane.
227    /// Computes `silence_ms`, `elapsed_ms`, and `remaining_budget_ms` from
228    /// atomic timestamps without any mutex locks.
229    pub fn flush_timing(&self) {
230        let last_activity = self.last_activity_ns.load(Ordering::Relaxed);
231        if last_activity > 0 {
232            let now_ns = self.elapsed_ns();
233            let silence_ms = now_ns.saturating_sub(last_activity) / 1_000_000;
234            let _ = self.state.session().set("silence_ms", silence_ms);
235        }
236
237        if self.is_connected.load(Ordering::Relaxed) {
238            let connected_ns = self.connected_at_ns.load(Ordering::Relaxed);
239            let now_ns = self.elapsed_ns();
240            let elapsed_ms = now_ns.saturating_sub(connected_ns) / 1_000_000;
241            let _ = self.state.session().set("elapsed_ms", elapsed_ms);
242
243            let limit_ms: u64 = match self.session_type() {
244                SessionType::AudioOnly => 15 * 60 * 1000,
245                SessionType::AudioVideo => 2 * 60 * 1000,
246            };
247            let remaining = limit_ms.saturating_sub(elapsed_ms);
248            let _ = self.state.session().set("remaining_budget_ms", remaining);
249        }
250    }
251
252    #[inline]
253    pub(crate) fn touch_activity(&self) {
254        self.last_activity_ns
255            .store(self.elapsed_ns(), Ordering::Relaxed);
256    }
257
258    #[inline]
259    fn elapsed_ns(&self) -> u64 {
260        self.clock.since(self.start).as_nanos() as u64
261    }
262
263    /// Record the turn's response latency (end of user speech, or text send,
264    /// to the model's first output) as `session:last_response_latency_ms`, so
265    /// instruction templates, watchers and computed vars can react to a slow
266    /// turn. Called by the telemetry lane once per measured turn.
267    ///
268    /// The write happens when the turn's *first* output event is observed, so
269    /// for a spoken turn it lands a whole model response before the boundary.
270    /// The telemetry and control lanes are independent readers of the same
271    /// event stream, though, with no barrier between them: if the telemetry
272    /// lane is running behind, a turn-boundary reader can still see the
273    /// previous turn's value. Treat the key as the latest measurement, not as
274    /// one synchronised to the current turn.
275    pub fn record_response_latency(&self, latency: std::time::Duration) {
276        let _ = self
277            .state
278            .session()
279            .set("last_response_latency_ms", latency.as_millis() as u64);
280        self.touch_activity();
281    }
282
283    /// Returns the current session type based on whether video has been sent.
284    pub fn session_type(&self) -> SessionType {
285        if self.has_video.load(Ordering::Relaxed) {
286            SessionType::AudioVideo
287        } else {
288            SessionType::AudioOnly
289        }
290    }
291
292    /// Returns the latest resumption handle for reconnection.
293    pub fn latest_resume_handle(&self) -> Option<String> {
294        self.latest_resume_handle.lock().clone()
295    }
296
297    /// Mark that video has been sent (changes session type to `AudioVideo`).
298    pub fn mark_video_sent(&self) {
299        if !self.has_video.swap(true, Ordering::Relaxed) {
300            let _ = self.state.session().set("session_type", "audio_video");
301        }
302    }
303}
304
305#[cfg(test)]
306mod tests {
307    use super::*;
308    use bytes::Bytes;
309    use gemini_genai_rs::prelude::SessionEvent;
310
311    fn signals() -> SessionSignals {
312        SessionSignals::new(State::new())
313    }
314
315    #[test]
316    fn connected_initializes_state() {
317        let s = signals();
318        s.on_event(&SessionEvent::Connected);
319
320        assert_eq!(s.state.session().get::<u64>("connected_at_ms"), Some(0));
321        assert_eq!(s.state.session().get::<u64>("interrupt_count"), Some(0));
322        assert_eq!(s.state.session().get::<u64>("error_count"), Some(0));
323        assert_eq!(
324            s.state.session().get::<bool>("is_user_speaking"),
325            Some(false)
326        );
327        assert_eq!(
328            s.state.session().get::<bool>("is_model_speaking"),
329            Some(false)
330        );
331        assert_eq!(
332            s.state.session().get::<bool>("go_away_received"),
333            Some(false)
334        );
335        assert_eq!(s.state.session().get::<bool>("resumable"), Some(false));
336        assert_eq!(
337            s.state.session().get::<String>("session_type"),
338            Some("audio_only".to_string())
339        );
340        assert!(s.is_connected.load(Ordering::Relaxed));
341    }
342
343    #[test]
344    fn voice_activity_toggles_user_speaking() {
345        let s = signals();
346        s.on_event(&SessionEvent::Connected);
347        s.on_event(&SessionEvent::VoiceActivityStart);
348        assert_eq!(
349            s.state.session().get::<bool>("is_user_speaking"),
350            Some(true)
351        );
352        s.on_event(&SessionEvent::VoiceActivityEnd);
353        assert_eq!(
354            s.state.session().get::<bool>("is_user_speaking"),
355            Some(false)
356        );
357    }
358
359    #[test]
360    fn interrupted_increments_count() {
361        let s = signals();
362        s.on_event(&SessionEvent::Connected);
363        s.on_event(&SessionEvent::Interrupted);
364        assert_eq!(s.state.session().get::<u64>("interrupt_count"), Some(1));
365        s.on_event(&SessionEvent::Interrupted);
366        assert_eq!(s.state.session().get::<u64>("interrupt_count"), Some(2));
367        s.on_event(&SessionEvent::Interrupted);
368        assert_eq!(s.state.session().get::<u64>("interrupt_count"), Some(3));
369    }
370
371    #[test]
372    fn error_increments_count() {
373        let s = signals();
374        s.on_event(&SessionEvent::Connected);
375        let err = |m: &str| {
376            SessionEvent::Error(gemini_genai_rs::session::SessionError::WebSocket(
377                gemini_genai_rs::session::WebSocketError::ProtocolError(m.into()),
378            ))
379        };
380        s.on_event(&err("oops"));
381        assert_eq!(s.state.session().get::<u64>("error_count"), Some(1));
382        let last: String = s.state.session().get("last_error").unwrap();
383        assert!(last.contains("oops"), "{last}");
384        s.on_event(&err("oops2"));
385        assert_eq!(s.state.session().get::<u64>("error_count"), Some(2));
386        let last: String = s.state.session().get("last_error").unwrap();
387        assert!(last.contains("oops2"), "{last}");
388    }
389
390    #[test]
391    fn phase_changed_sets_model_speaking() {
392        let s = signals();
393        s.on_event(&SessionEvent::Connected);
394        s.on_event(&SessionEvent::PhaseChanged(SessionPhase::ModelSpeaking));
395        assert_eq!(
396            s.state.session().get::<bool>("is_model_speaking"),
397            Some(true)
398        );
399        assert_eq!(
400            s.state.session().get::<String>("phase"),
401            Some("ModelSpeaking".into())
402        );
403        s.on_event(&SessionEvent::PhaseChanged(SessionPhase::Active));
404        assert_eq!(
405            s.state.session().get::<bool>("is_model_speaking"),
406            Some(false)
407        );
408        assert_eq!(
409            s.state.session().get::<String>("phase"),
410            Some("Active".into())
411        );
412    }
413
414    #[test]
415    fn go_away_sets_state() {
416        let s = signals();
417        s.on_event(&SessionEvent::Connected);
418        s.on_event(&SessionEvent::GoAway(Some(std::time::Duration::from_secs(
419            60,
420        ))));
421        assert_eq!(
422            s.state.session().get::<bool>("go_away_received"),
423            Some(true)
424        );
425        assert_eq!(
426            s.state.session().get::<u64>("go_away_time_left_ms"),
427            Some(60_000)
428        );
429        assert!(s.go_away_at.lock().is_some());
430    }
431
432    #[test]
433    fn go_away_without_time_left() {
434        let s = signals();
435        s.on_event(&SessionEvent::Connected);
436        s.on_event(&SessionEvent::GoAway(None));
437        assert_eq!(
438            s.state.session().get::<bool>("go_away_received"),
439            Some(true)
440        );
441        assert_eq!(s.state.session().get::<u64>("go_away_time_left_ms"), None);
442        assert!(s.go_away_at.lock().is_none());
443    }
444
445    #[test]
446    fn session_resume_handle_stored() {
447        let s = signals();
448        s.on_event(&SessionEvent::Connected);
449        s.on_event(&SessionEvent::SessionResumeUpdate(
450            gemini_genai_rs::session::ResumeInfo {
451                handle: "handle-abc".into(),
452                resumable: true,
453                last_consumed_index: None,
454            },
455        ));
456        assert_eq!(s.state.session().get::<bool>("resumable"), Some(true));
457        assert_eq!(s.latest_resume_handle(), Some("handle-abc".to_string()));
458    }
459
460    #[test]
461    fn transcription_stores_last() {
462        let s = signals();
463        s.on_event(&SessionEvent::Connected);
464        s.on_event(&SessionEvent::InputTranscription("hello".into()));
465        assert_eq!(
466            s.state.session().get::<String>("last_input_transcription"),
467            Some("hello".into())
468        );
469        s.on_event(&SessionEvent::OutputTranscription("hi there".into()));
470        assert_eq!(
471            s.state.session().get::<String>("last_output_transcription"),
472            Some("hi there".into())
473        );
474        s.on_event(&SessionEvent::InputTranscription("bye".into()));
475        assert_eq!(
476            s.state.session().get::<String>("last_input_transcription"),
477            Some("bye".into())
478        );
479    }
480
481    #[test]
482    fn session_type_defaults_to_audio_only() {
483        let s = signals();
484        assert_eq!(s.session_type(), SessionType::AudioOnly);
485    }
486
487    #[test]
488    fn mark_video_sent_changes_session_type() {
489        let s = signals();
490        s.on_event(&SessionEvent::Connected);
491        assert_eq!(s.session_type(), SessionType::AudioOnly);
492        s.mark_video_sent();
493        assert_eq!(s.session_type(), SessionType::AudioVideo);
494        assert_eq!(
495            s.state.session().get::<String>("session_type"),
496            Some("audio_video".into())
497        );
498    }
499
500    #[test]
501    fn mark_video_sent_idempotent() {
502        let s = signals();
503        s.on_event(&SessionEvent::Connected);
504        s.mark_video_sent();
505        s.mark_video_sent();
506        assert_eq!(s.session_type(), SessionType::AudioVideo);
507    }
508
509    #[test]
510    fn flush_timing_after_connected() {
511        let s = signals();
512        s.on_event(&SessionEvent::Connected);
513        s.flush_timing();
514        let elapsed: u64 = s.state.session().get("elapsed_ms").unwrap_or(0);
515        assert!(elapsed < 100, "elapsed should be near zero, got {elapsed}");
516        let remaining: u64 = s.state.session().get("remaining_budget_ms").unwrap();
517        let limit = 15 * 60 * 1000u64;
518        assert!(
519            remaining > limit - 1000,
520            "remaining should be near limit, got {remaining}"
521        );
522    }
523
524    #[test]
525    fn flush_timing_respects_video_budget() {
526        let s = signals();
527        s.on_event(&SessionEvent::Connected);
528        s.flush_timing();
529        let remaining_audio: u64 = s.state.session().get("remaining_budget_ms").unwrap();
530        assert!(remaining_audio > 14 * 60 * 1000);
531        s.mark_video_sent();
532        s.flush_timing();
533        let remaining_video: u64 = s.state.session().get("remaining_budget_ms").unwrap();
534        assert!(
535            remaining_video <= 2 * 60 * 1000,
536            "video remaining should be <= 120_000, got {remaining_video}"
537        );
538    }
539
540    #[test]
541    fn latest_resume_handle_initially_none() {
542        let s = signals();
543        assert_eq!(s.latest_resume_handle(), None);
544    }
545
546    #[test]
547    fn latest_resume_handle_updates() {
548        let s = signals();
549        s.on_event(&SessionEvent::SessionResumeUpdate(
550            gemini_genai_rs::session::ResumeInfo {
551                handle: "h1".into(),
552                resumable: true,
553                last_consumed_index: None,
554            },
555        ));
556        assert_eq!(s.latest_resume_handle(), Some("h1".to_string()));
557        s.on_event(&SessionEvent::SessionResumeUpdate(
558            gemini_genai_rs::session::ResumeInfo {
559                handle: "h2".into(),
560                resumable: true,
561                last_consumed_index: Some("5".into()),
562            },
563        ));
564        assert_eq!(s.latest_resume_handle(), Some("h2".to_string()));
565    }
566
567    #[test]
568    fn silence_ms_tracked() {
569        let s = signals();
570        s.on_event(&SessionEvent::Connected);
571        s.flush_timing();
572        let silence: u64 = s.state.session().get("silence_ms").unwrap_or(u64::MAX);
573        assert!(silence < 100, "silence should be near zero, got {silence}");
574    }
575
576    #[test]
577    fn audio_data_updates_activity() {
578        let s = signals();
579        s.on_event(&SessionEvent::Connected);
580        s.on_event(&SessionEvent::AudioData(Bytes::from_static(b"pcm")));
581        s.flush_timing();
582        let silence: u64 = s.state.session().get("silence_ms").unwrap_or(u64::MAX);
583        assert!(silence < 100);
584    }
585
586    #[test]
587    fn turn_complete_updates_activity() {
588        let s = signals();
589        s.on_event(&SessionEvent::Connected);
590        s.on_event(&SessionEvent::TurnComplete);
591        s.flush_timing();
592        let silence: u64 = s.state.session().get("silence_ms").unwrap_or(u64::MAX);
593        assert!(silence < 100);
594    }
595
596    #[test]
597    fn text_complete_updates_activity() {
598        let s = signals();
599        s.on_event(&SessionEvent::Connected);
600        s.on_event(&SessionEvent::TextComplete("done".into()));
601        s.flush_timing();
602        let silence: u64 = s.state.session().get("silence_ms").unwrap_or(u64::MAX);
603        assert!(silence < 100);
604    }
605
606    #[test]
607    fn disconnected_clears_connected_and_sets_flag() {
608        let s = signals();
609        s.on_event(&SessionEvent::Connected);
610        assert!(s.is_connected.load(Ordering::Relaxed));
611        s.on_event(&SessionEvent::Disconnected(Some("server closed".into())));
612        assert!(!s.is_connected.load(Ordering::Relaxed));
613        assert_eq!(s.state.session().get::<bool>("disconnected"), Some(true));
614    }
615
616    #[test]
617    fn disconnected_without_reason() {
618        let s = signals();
619        s.on_event(&SessionEvent::Connected);
620        s.on_event(&SessionEvent::Disconnected(None));
621        assert!(!s.is_connected.load(Ordering::Relaxed));
622        assert_eq!(s.state.session().get::<bool>("disconnected"), Some(true));
623    }
624}
gemini_adk_rs/live/session_signals.rs

gemini_adk_rs/live/
session_signals.rs
