gemini_genai_rs/session/
state.rs

1//! Session phase finite state machine and shared session state.
2//!
3//! [`SessionPhase`] — lifecycle phase enum with validated transitions.
4//! [`SessionState`] — shared state struct (phase, turns, resume handle).
5//!
6//! Invalid transitions return `Err(SessionError::InvalidTransition)`.
7//! The phase is observable via a `watch::Receiver<SessionPhase>` channel.
8
9use std::fmt;
10use std::time::Instant;
11use tokio::sync::{broadcast, watch};
12
13use super::errors::SessionError;
14use super::events::{SessionEvent, Turn};
15
16/// The lifecycle phase of a Gemini Live session.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
18pub enum SessionPhase {
19    /// Not connected to the server.
20    Disconnected,
21    /// WebSocket connection in progress.
22    Connecting,
23    /// Setup message sent, awaiting setupComplete.
24    SetupSent,
25    /// Session is active and ready for interaction.
26    Active,
27    /// User is currently speaking (client VAD or server signal).
28    UserSpeaking,
29    /// Model is currently generating a response.
30    ModelSpeaking,
31    /// Model was interrupted by user barge-in.
32    Interrupted,
33    /// Model requested tool calls, awaiting dispatch.
34    ToolCallPending,
35    /// Tool calls are executing.
36    ToolCallExecuting,
37    /// Session is shutting down gracefully.
38    Disconnecting,
39}
40
41impl fmt::Display for SessionPhase {
42    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43        match self {
44            Self::Disconnected => write!(f, "Disconnected"),
45            Self::Connecting => write!(f, "Connecting"),
46            Self::SetupSent => write!(f, "SetupSent"),
47            Self::Active => write!(f, "Active"),
48            Self::UserSpeaking => write!(f, "UserSpeaking"),
49            Self::ModelSpeaking => write!(f, "ModelSpeaking"),
50            Self::Interrupted => write!(f, "Interrupted"),
51            Self::ToolCallPending => write!(f, "ToolCallPending"),
52            Self::ToolCallExecuting => write!(f, "ToolCallExecuting"),
53            Self::Disconnecting => write!(f, "Disconnecting"),
54        }
55    }
56}
57
58impl SessionPhase {
59    /// Check whether a transition from this phase to `to` is valid.
60    pub fn can_transition_to(&self, to: &SessionPhase) -> bool {
61        matches!(
62            (self, to),
63            // Connection lifecycle
64            (SessionPhase::Disconnected, SessionPhase::Connecting)
65                | (SessionPhase::Connecting, SessionPhase::SetupSent)
66                | (SessionPhase::SetupSent, SessionPhase::Active)
67                // Conversation flow
68                | (SessionPhase::Active, SessionPhase::UserSpeaking)
69                | (SessionPhase::Active, SessionPhase::ModelSpeaking)
70                | (SessionPhase::Active, SessionPhase::ToolCallPending)
71                // User speaking transitions
72                | (SessionPhase::UserSpeaking, SessionPhase::Active)
73                | (SessionPhase::UserSpeaking, SessionPhase::ModelSpeaking)
74                // Model speaking transitions
75                | (SessionPhase::ModelSpeaking, SessionPhase::Active)
76                | (SessionPhase::ModelSpeaking, SessionPhase::Interrupted)
77                | (SessionPhase::ModelSpeaking, SessionPhase::ToolCallPending)
78                // Barge-in recovery
79                | (SessionPhase::Interrupted, SessionPhase::Active)
80                | (SessionPhase::Interrupted, SessionPhase::UserSpeaking)
81                // Tool flow
82                | (SessionPhase::ToolCallPending, SessionPhase::ToolCallExecuting)
83                | (SessionPhase::ToolCallExecuting, SessionPhase::Active)
84                | (SessionPhase::ToolCallExecuting, SessionPhase::ModelSpeaking)
85                // Graceful shutdown
86                | (SessionPhase::Active, SessionPhase::Disconnecting)
87                | (SessionPhase::UserSpeaking, SessionPhase::Disconnecting)
88                | (SessionPhase::ModelSpeaking, SessionPhase::Disconnecting)
89                | (SessionPhase::Interrupted, SessionPhase::Disconnecting)
90                | (SessionPhase::ToolCallPending, SessionPhase::Disconnecting)
91                | (SessionPhase::ToolCallExecuting, SessionPhase::Disconnecting)
92                | (SessionPhase::Disconnecting, SessionPhase::Disconnected)
93                // Force-disconnect from any state
94                | (_, SessionPhase::Disconnected)
95        )
96    }
97}
98
99// ---------------------------------------------------------------------------
100// Session state (shared, read-mostly)
101// ---------------------------------------------------------------------------
102
103/// Shared session state, accessible from the SessionHandle.
104#[derive(Debug)]
105pub struct SessionState {
106    /// Current phase (updated atomically via watch channel).
107    phase_tx: watch::Sender<SessionPhase>,
108    /// Optional broadcast sender to emit `PhaseChanged` events on transitions.
109    event_tx: Option<broadcast::Sender<SessionEvent>>,
110    /// Session ID.
111    pub session_id: String,
112    /// Session resume handle from server.
113    pub resume_handle: parking_lot::Mutex<Option<String>>,
114    /// Turn history.
115    pub turns: parking_lot::Mutex<Vec<Turn>>,
116    /// Current in-progress turn.
117    pub current_turn: parking_lot::Mutex<Option<Turn>>,
118    /// The session is text-only on a speech-only model: the output
119    /// transcription is its text, and audio is dropped.
120    text_from_transcription: std::sync::atomic::AtomicBool,
121}
122
123impl SessionState {
124    /// Create new session state (no `PhaseChanged` event emission).
125    pub fn new(phase_tx: watch::Sender<SessionPhase>) -> Self {
126        Self {
127            phase_tx,
128            event_tx: None,
129            session_id: uuid::Uuid::new_v4().to_string(),
130            resume_handle: parking_lot::Mutex::new(None),
131            turns: parking_lot::Mutex::new(Vec::new()),
132            current_turn: parking_lot::Mutex::new(None),
133            text_from_transcription: std::sync::atomic::AtomicBool::new(false),
134        }
135    }
136
137    /// Deliver the output transcription as the session's text (and drop
138    /// audio), for a text-only session on a speech-only model.
139    pub(crate) fn set_text_from_transcription(&self, enabled: bool) {
140        self.text_from_transcription
141            .store(enabled, std::sync::atomic::Ordering::Relaxed);
142    }
143
144    pub(crate) fn text_from_transcription(&self) -> bool {
145        self.text_from_transcription
146            .load(std::sync::atomic::Ordering::Relaxed)
147    }
148
149    /// Create new session state that emits `PhaseChanged` events on transitions.
150    pub fn with_events(
151        phase_tx: watch::Sender<SessionPhase>,
152        event_tx: broadcast::Sender<SessionEvent>,
153    ) -> Self {
154        Self {
155            phase_tx,
156            event_tx: Some(event_tx),
157            session_id: uuid::Uuid::new_v4().to_string(),
158            resume_handle: parking_lot::Mutex::new(None),
159            turns: parking_lot::Mutex::new(Vec::new()),
160            current_turn: parking_lot::Mutex::new(None),
161            text_from_transcription: std::sync::atomic::AtomicBool::new(false),
162        }
163    }
164
165    /// Get the current phase.
166    pub fn phase(&self) -> SessionPhase {
167        *self.phase_tx.borrow()
168    }
169
170    /// Attempt a validated phase transition.
171    ///
172    /// If an `event_tx` was provided via [`with_events`](Self::with_events),
173    /// a [`SessionEvent::PhaseChanged`] is broadcast after a successful transition.
174    pub fn transition_to(&self, to: SessionPhase) -> Result<SessionPhase, SessionError> {
175        let from = self.phase();
176        if !from.can_transition_to(&to) {
177            return Err(SessionError::InvalidTransition { from, to });
178        }
179        self.phase_tx.send_replace(to);
180        if let Some(ref tx) = self.event_tx {
181            let _ = tx.send(SessionEvent::PhaseChanged(to));
182        }
183        Ok(to)
184    }
185
186    /// Force transition (bypasses validation — use only for disconnect).
187    pub fn force_phase(&self, phase: SessionPhase) {
188        self.phase_tx.send_replace(phase);
189    }
190
191    /// Start a new turn.
192    pub fn start_turn(&self) {
193        let mut current = self.current_turn.lock();
194        if let Some(prev) = current.take() {
195            self.turns.lock().push(prev);
196        }
197        *current = Some(Turn::new());
198    }
199
200    /// Append text to the current turn.
201    pub fn append_text(&self, text: &str) {
202        if let Some(turn) = self.current_turn.lock().as_mut() {
203            turn.text.push_str(text);
204        }
205    }
206
207    /// Mark audio received in the current turn.
208    pub fn mark_audio(&self) {
209        if let Some(turn) = self.current_turn.lock().as_mut() {
210            turn.has_audio = true;
211        }
212    }
213
214    /// Complete the current turn.
215    pub fn complete_turn(&self) -> Option<Turn> {
216        let mut current = self.current_turn.lock();
217        if let Some(turn) = current.as_mut() {
218            turn.completed_at = Some(Instant::now());
219        }
220        let completed = current.take();
221        if let Some(ref t) = completed {
222            self.turns.lock().push(t.clone());
223        }
224        completed
225    }
226
227    /// Mark the current turn as interrupted.
228    pub fn interrupt_turn(&self) {
229        if let Some(turn) = self.current_turn.lock().as_mut() {
230            turn.interrupted = true;
231            turn.completed_at = Some(Instant::now());
232        }
233    }
234}
235
236#[cfg(test)]
237mod tests {
238    use super::*;
239
240    #[test]
241    fn valid_connection_lifecycle() {
242        assert!(SessionPhase::Disconnected.can_transition_to(&SessionPhase::Connecting));
243        assert!(SessionPhase::Connecting.can_transition_to(&SessionPhase::SetupSent));
244        assert!(SessionPhase::SetupSent.can_transition_to(&SessionPhase::Active));
245    }
246
247    #[test]
248    fn valid_conversation_flow() {
249        assert!(SessionPhase::Active.can_transition_to(&SessionPhase::UserSpeaking));
250        assert!(SessionPhase::Active.can_transition_to(&SessionPhase::ModelSpeaking));
251        assert!(SessionPhase::UserSpeaking.can_transition_to(&SessionPhase::Active));
252        assert!(SessionPhase::ModelSpeaking.can_transition_to(&SessionPhase::Active));
253    }
254
255    #[test]
256    fn valid_barge_in() {
257        assert!(SessionPhase::ModelSpeaking.can_transition_to(&SessionPhase::Interrupted));
258        assert!(SessionPhase::Interrupted.can_transition_to(&SessionPhase::Active));
259        assert!(SessionPhase::Interrupted.can_transition_to(&SessionPhase::UserSpeaking));
260    }
261
262    #[test]
263    fn valid_tool_flow() {
264        assert!(SessionPhase::Active.can_transition_to(&SessionPhase::ToolCallPending));
265        assert!(SessionPhase::ModelSpeaking.can_transition_to(&SessionPhase::ToolCallPending));
266        assert!(SessionPhase::ToolCallPending.can_transition_to(&SessionPhase::ToolCallExecuting));
267        assert!(SessionPhase::ToolCallExecuting.can_transition_to(&SessionPhase::Active));
268        assert!(SessionPhase::ToolCallExecuting.can_transition_to(&SessionPhase::ModelSpeaking));
269    }
270
271    #[test]
272    fn valid_disconnect_from_any() {
273        let phases = [
274            SessionPhase::Disconnected,
275            SessionPhase::Connecting,
276            SessionPhase::SetupSent,
277            SessionPhase::Active,
278            SessionPhase::UserSpeaking,
279            SessionPhase::ModelSpeaking,
280            SessionPhase::Interrupted,
281            SessionPhase::ToolCallPending,
282            SessionPhase::ToolCallExecuting,
283            SessionPhase::Disconnecting,
284        ];
285
286        for phase in &phases {
287            assert!(
288                phase.can_transition_to(&SessionPhase::Disconnected),
289                "{phase} should be able to force-disconnect"
290            );
291        }
292    }
293
294    #[test]
295    fn invalid_transitions() {
296        assert!(!SessionPhase::Disconnected.can_transition_to(&SessionPhase::Active));
297        assert!(!SessionPhase::Connecting.can_transition_to(&SessionPhase::Active));
298        assert!(!SessionPhase::Active.can_transition_to(&SessionPhase::SetupSent));
299        assert!(!SessionPhase::UserSpeaking.can_transition_to(&SessionPhase::ToolCallExecuting));
300        assert!(!SessionPhase::Disconnecting.can_transition_to(&SessionPhase::Active));
301    }
302
303    #[test]
304    fn display_impl() {
305        assert_eq!(format!("{}", SessionPhase::Active), "Active");
306        assert_eq!(format!("{}", SessionPhase::ModelSpeaking), "ModelSpeaking");
307    }
308}