gemini_genai_rs/session/
handle.rs

1//! [`SessionHandle`] — the public API surface for a Gemini Live session.
2//!
3//! Cheaply cloneable (wraps `Arc`). Provides methods to send commands,
4//! subscribe to events, and observe session state.
5
6use super::errors::SessionError;
7use super::events::{SessionCommand, SessionEvent};
8use super::state::{SessionPhase, SessionState};
9use super::traits::{SessionReader, SessionWriter};
10use crate::protocol::{Content, ContextUpdate, FunctionResponse};
11use async_trait::async_trait;
12use bytes::Bytes;
13use std::sync::Arc;
14use tokio::sync::{broadcast, mpsc, watch};
15use tokio::task::JoinHandle;
16
17/// The public API surface for a Gemini Live session.
18///
19/// Cheaply cloneable (wraps `Arc`). Provides methods to send commands,
20/// subscribe to events, and observe session state.
21#[derive(Clone)]
22pub struct SessionHandle {
23    /// Channel for sending commands to the transport layer.
24    command_tx: mpsc::Sender<SessionCommand>,
25    /// Broadcast channel for session events.
26    event_tx: broadcast::Sender<SessionEvent>,
27    /// Shared session state.
28    state: Arc<SessionState>,
29    /// Phase watch receiver for async observation.
30    phase_rx: watch::Receiver<SessionPhase>,
31    /// Handle to the spawned connection loop task.
32    ///
33    /// Wrapped in `Arc<Mutex<Option<...>>>` so that `SessionHandle` remains
34    /// `Clone` (since `JoinHandle` is not `Clone`). The first call to
35    /// [`join()`](Self::join) takes the handle; subsequent calls return `Ok(())`.
36    task: Arc<tokio::sync::Mutex<Option<JoinHandle<()>>>>,
37    /// Optional producer-side audio pacing (see [`SessionConfig::audio_pacing`](crate::protocol::SessionConfig::audio_pacing)).
38    ///
39    /// Shared across handle clones so all producers draw from one bucket. The
40    /// tokio mutex is held across the pacing wait deliberately: concurrent
41    /// audio producers serialize, which is the desired backpressure semantic.
42    audio_pacer: Option<Arc<tokio::sync::Mutex<crate::transport::TokenBucket>>>,
43}
44
45impl SessionHandle {
46    /// Create a new session handle from its components.
47    pub fn new(
48        command_tx: mpsc::Sender<SessionCommand>,
49        event_tx: broadcast::Sender<SessionEvent>,
50        state: Arc<SessionState>,
51        phase_rx: watch::Receiver<SessionPhase>,
52    ) -> Self {
53        Self {
54            command_tx,
55            event_tx,
56            state,
57            phase_rx,
58            task: Arc::new(tokio::sync::Mutex::new(None)),
59            audio_pacer: None,
60        }
61    }
62
63    /// Enable producer-side audio send pacing (token bucket).
64    ///
65    /// Installed by the connection layer when
66    /// [`SessionConfig::audio_pacing`](crate::protocol::SessionConfig::audio_pacing) is set.
67    pub fn with_audio_pacing(mut self, config: crate::transport::BackpressureConfig) -> Self {
68        self.audio_pacer = Some(Arc::new(tokio::sync::Mutex::new(
69            crate::transport::TokenBucket::new(config),
70        )));
71        self
72    }
73
74    /// Store the connection loop task handle.
75    ///
76    /// Called by the transport layer after spawning the connection loop.
77    pub fn set_task(&self, handle: JoinHandle<()>) {
78        // Use try_lock to avoid blocking — this is only called once at startup.
79        if let Ok(mut guard) = self.task.try_lock() {
80            *guard = Some(handle);
81        }
82    }
83
84    /// Wait for the session connection loop to complete.
85    ///
86    /// Returns `Ok(())` when the session disconnects normally.
87    /// Returns `Err` if the connection task panicked.
88    ///
89    /// Only the first call across all clones actually awaits the task;
90    /// subsequent calls return `Ok(())` immediately.
91    pub async fn join(&self) -> Result<(), tokio::task::JoinError> {
92        let task = self.task.lock().await.take();
93        if let Some(handle) = task {
94            handle.await
95        } else {
96            Ok(())
97        }
98    }
99
100    /// Subscribe to session events.
101    pub fn subscribe(&self) -> broadcast::Receiver<SessionEvent> {
102        self.event_tx.subscribe()
103    }
104
105    /// The event broadcaster, for a runtime that injects its own events
106    /// into the session stream (the L1 agent session does). Application code
107    /// wants [`subscribe`](Self::subscribe).
108    pub fn event_sender(&self) -> &broadcast::Sender<SessionEvent> {
109        &self.event_tx
110    }
111
112    /// The most recent session-resumption handle the server issued, if any.
113    pub fn resume_handle(&self) -> Option<String> {
114        self.state.resume_handle.lock().clone()
115    }
116
117    /// Current session phase.
118    pub fn phase(&self) -> SessionPhase {
119        self.state.phase()
120    }
121
122    /// Session ID.
123    pub fn session_id(&self) -> &str {
124        &self.state.session_id
125    }
126
127    /// Wait for the session to reach a specific phase.
128    pub async fn wait_for_phase(&self, target: SessionPhase) {
129        let mut rx = self.phase_rx.clone();
130        while *rx.borrow_and_update() != target {
131            if rx.changed().await.is_err() {
132                break;
133            }
134        }
135    }
136
137    /// Send audio data (raw PCM16 bytes).
138    ///
139    /// When [`SessionConfig::audio_pacing`](crate::protocol::SessionConfig::audio_pacing) is configured, this paces the
140    /// caller: pushing audio faster than the configured sustained rate waits
141    /// here instead of overflowing the send queue.
142    pub async fn send_audio(&self, data: impl Into<Bytes>) -> Result<(), SessionError> {
143        let data: Bytes = data.into();
144        if let Some(pacer) = &self.audio_pacer {
145            pacer.lock().await.consume(data.len()).await;
146        }
147        self.send_command(SessionCommand::SendAudio(data)).await
148    }
149
150    /// Send a text message.
151    pub async fn send_text(&self, text: impl Into<String>) -> Result<(), SessionError> {
152        self.send_command(SessionCommand::SendText(text.into()))
153            .await
154    }
155
156    /// Send tool responses.
157    pub async fn send_tool_response(
158        &self,
159        responses: Vec<FunctionResponse>,
160    ) -> Result<(), SessionError> {
161        self.send_command(SessionCommand::SendToolResponse(responses))
162            .await
163    }
164
165    /// Send a video/image frame (raw JPEG bytes).
166    pub async fn send_video(&self, jpeg_data: impl Into<Bytes>) -> Result<(), SessionError> {
167        self.send_command(SessionCommand::SendVideo(jpeg_data.into()))
168            .await
169    }
170
171    /// Update the system instruction mid-session.
172    pub async fn update_instruction(
173        &self,
174        instruction: impl Into<String>,
175    ) -> Result<(), SessionError> {
176        self.send_command(SessionCommand::UpdateInstruction(instruction.into()))
177            .await
178    }
179
180    /// Replace the declared tools and/or the system instruction mid-session
181    /// with a `contextUpdate` message (Gemini 3.8 Live).
182    ///
183    /// Only the fields set on `update` change, and an empty update sends
184    /// nothing. Each update invalidates the server's prefix cache, so send
185    /// one only when the tool set or instruction actually changes. Once sent,
186    /// the update also shapes the setup message of any reconnect.
187    ///
188    /// On a model without `contextUpdate` (Gemini 2.5 Live; see
189    /// [`SessionConfig::supports_context_update`](crate::protocol::SessionConfig::supports_context_update))
190    /// nothing is sent and the session reports a
191    /// [`SessionError::Codec`] error event; the session itself stays up.
192    pub async fn update_context(&self, update: ContextUpdate) -> Result<(), SessionError> {
193        self.send_command(SessionCommand::UpdateContext(update))
194            .await
195    }
196
197    /// Signal activity start (user started speaking).
198    pub async fn signal_activity_start(&self) -> Result<(), SessionError> {
199        self.send_command(SessionCommand::ActivityStart).await
200    }
201
202    /// Signal activity end (user stopped speaking).
203    pub async fn signal_activity_end(&self) -> Result<(), SessionError> {
204        self.send_command(SessionCommand::ActivityEnd).await
205    }
206
207    /// Send client content (turns + turn_complete flag).
208    /// Used for injecting conversation history, context, or multi-turn text.
209    pub async fn send_client_content(
210        &self,
211        turns: Vec<Content>,
212        turn_complete: bool,
213    ) -> Result<(), SessionError> {
214        self.send_command(SessionCommand::SendClientContent {
215            turns,
216            turn_complete,
217        })
218        .await
219    }
220
221    /// Gracefully disconnect the session.
222    pub async fn disconnect(&self) -> Result<(), SessionError> {
223        self.send_command(SessionCommand::Disconnect).await
224    }
225
226    /// Send a command to the transport.
227    async fn send_command(&self, cmd: SessionCommand) -> Result<(), SessionError> {
228        self.command_tx
229            .send(cmd)
230            .await
231            .map_err(|_| SessionError::ChannelClosed)
232    }
233}
234
235impl std::fmt::Debug for SessionHandle {
236    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
237        f.debug_struct("SessionHandle")
238            .field("session_id", &self.state.session_id)
239            .field("phase", &self.state.phase())
240            .finish()
241    }
242}
243
244// ---------------------------------------------------------------------------
245// Trait implementations for SessionHandle
246// ---------------------------------------------------------------------------
247
248#[async_trait]
249impl SessionWriter for SessionHandle {
250    async fn send_audio(&self, data: Bytes) -> Result<(), SessionError> {
251        SessionHandle::send_audio(self, data).await
252    }
253
254    async fn send_text(&self, text: String) -> Result<(), SessionError> {
255        self.send_command(SessionCommand::SendText(text)).await
256    }
257
258    async fn send_tool_response(
259        &self,
260        responses: Vec<FunctionResponse>,
261    ) -> Result<(), SessionError> {
262        self.send_command(SessionCommand::SendToolResponse(responses))
263            .await
264    }
265
266    async fn send_client_content(
267        &self,
268        turns: Vec<Content>,
269        turn_complete: bool,
270    ) -> Result<(), SessionError> {
271        self.send_command(SessionCommand::SendClientContent {
272            turns,
273            turn_complete,
274        })
275        .await
276    }
277
278    async fn send_video(&self, jpeg_data: Bytes) -> Result<(), SessionError> {
279        self.send_command(SessionCommand::SendVideo(jpeg_data))
280            .await
281    }
282
283    async fn update_instruction(&self, instruction: String) -> Result<(), SessionError> {
284        self.send_command(SessionCommand::UpdateInstruction(instruction))
285            .await
286    }
287
288    async fn update_context(&self, update: ContextUpdate) -> Result<(), SessionError> {
289        SessionHandle::update_context(self, update).await
290    }
291
292    async fn signal_activity_start(&self) -> Result<(), SessionError> {
293        self.send_command(SessionCommand::ActivityStart).await
294    }
295
296    async fn signal_activity_end(&self) -> Result<(), SessionError> {
297        self.send_command(SessionCommand::ActivityEnd).await
298    }
299
300    async fn disconnect(&self) -> Result<(), SessionError> {
301        self.send_command(SessionCommand::Disconnect).await
302    }
303}
304
305impl SessionReader for SessionHandle {
306    fn subscribe(&self) -> broadcast::Receiver<SessionEvent> {
307        self.event_tx.subscribe()
308    }
309
310    fn phase(&self) -> SessionPhase {
311        self.state.phase()
312    }
313
314    fn session_id(&self) -> &str {
315        &self.state.session_id
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    #[tokio::test(start_paused = true)]
324    async fn audio_pacing_throttles_producer_to_sustained_rate() {
325        let (command_tx, mut command_rx) = mpsc::channel(64);
326        let (event_tx, _) = broadcast::channel(16);
327        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Active);
328        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
329
330        // 1000-byte burst allowance, 1000 B/s sustained.
331        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx).with_audio_pacing(
332            crate::transport::BackpressureConfig {
333                bucket_capacity: 1000,
334                refill_rate_bps: 1000,
335            },
336        );
337
338        let start = tokio::time::Instant::now();
339        // First 1000 bytes ride the burst; the next 1000 must wait ~1s.
340        handle.send_audio(vec![0u8; 1000]).await.unwrap();
341        let after_burst = start.elapsed();
342        handle.send_audio(vec![0u8; 1000]).await.unwrap();
343        let after_paced = start.elapsed();
344
345        assert!(after_burst < std::time::Duration::from_millis(50));
346        assert!(
347            after_paced >= std::time::Duration::from_millis(900),
348            "second send should be paced ~1s, was {after_paced:?}"
349        );
350        // Both frames were enqueued.
351        assert!(command_rx.recv().await.is_some());
352        assert!(command_rx.recv().await.is_some());
353    }
354
355    #[tokio::test]
356    async fn session_handle_join_returns_ok_after_task_completes() {
357        let (command_tx, _command_rx) = mpsc::channel(8);
358        let (event_tx, _) = broadcast::channel(16);
359        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
360        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
361
362        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
363
364        // Spawn a trivial task that completes immediately
365        let task = tokio::spawn(async {});
366        handle.set_task(task);
367
368        // join() should return Ok(())
369        let result = handle.join().await;
370        assert!(
371            result.is_ok(),
372            "join() should return Ok after task completes"
373        );
374    }
375
376    #[tokio::test]
377    async fn session_handle_join_without_task_returns_ok() {
378        let (command_tx, _command_rx) = mpsc::channel(8);
379        let (event_tx, _) = broadcast::channel(16);
380        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
381        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
382
383        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
384
385        // join() without set_task should return Ok immediately
386        let result = handle.join().await;
387        assert!(result.is_ok(), "join() without task should return Ok");
388    }
389
390    #[tokio::test]
391    async fn session_handle_join_idempotent() {
392        let (command_tx, _command_rx) = mpsc::channel(8);
393        let (event_tx, _) = broadcast::channel(16);
394        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
395        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
396
397        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
398
399        let task = tokio::spawn(async {});
400        handle.set_task(task);
401
402        // First join takes the handle
403        assert!(handle.join().await.is_ok());
404        // Second join returns Ok immediately (handle already taken)
405        assert!(handle.join().await.is_ok());
406    }
407
408    #[tokio::test]
409    async fn session_handle_join_works_on_clone() {
410        let (command_tx, _command_rx) = mpsc::channel(8);
411        let (event_tx, _) = broadcast::channel(16);
412        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
413        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
414
415        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
416        let handle_clone = handle.clone();
417
418        let task = tokio::spawn(async {});
419        handle.set_task(task);
420
421        // join() on clone should work (shares the Arc)
422        let result = handle_clone.join().await;
423        assert!(result.is_ok(), "join() on clone should work");
424
425        // Original handle's join should now return Ok (handle already taken)
426        assert!(handle.join().await.is_ok());
427    }
428
429    // PhaseChanged event emission tests
430
431    #[tokio::test]
432    async fn phase_changed_event_emitted_on_transition() {
433        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Disconnected);
434        let (event_tx, mut event_rx) = broadcast::channel(16);
435        let state = SessionState::with_events(phase_tx, event_tx);
436
437        state.transition_to(SessionPhase::Connecting).unwrap();
438
439        match event_rx.try_recv() {
440            Ok(SessionEvent::PhaseChanged(SessionPhase::Connecting)) => {}
441            other => panic!("expected PhaseChanged(Connecting), got {other:?}"),
442        }
443    }
444
445    #[test]
446    fn phase_changed_not_emitted_without_event_tx() {
447        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Disconnected);
448        let state = SessionState::new(phase_tx);
449        // Should not panic even though no event_tx
450        state.transition_to(SessionPhase::Connecting).unwrap();
451    }
452}
gemini_genai_rs/session/handle.rs

gemini_genai_rs/session/
handle.rs
