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, 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    /// Signal activity start (user started speaking).
181    pub async fn signal_activity_start(&self) -> Result<(), SessionError> {
182        self.send_command(SessionCommand::ActivityStart).await
183    }
184
185    /// Signal activity end (user stopped speaking).
186    pub async fn signal_activity_end(&self) -> Result<(), SessionError> {
187        self.send_command(SessionCommand::ActivityEnd).await
188    }
189
190    /// Send client content (turns + turn_complete flag).
191    /// Used for injecting conversation history, context, or multi-turn text.
192    pub async fn send_client_content(
193        &self,
194        turns: Vec<Content>,
195        turn_complete: bool,
196    ) -> Result<(), SessionError> {
197        self.send_command(SessionCommand::SendClientContent {
198            turns,
199            turn_complete,
200        })
201        .await
202    }
203
204    /// Gracefully disconnect the session.
205    pub async fn disconnect(&self) -> Result<(), SessionError> {
206        self.send_command(SessionCommand::Disconnect).await
207    }
208
209    /// Send a command to the transport.
210    async fn send_command(&self, cmd: SessionCommand) -> Result<(), SessionError> {
211        self.command_tx
212            .send(cmd)
213            .await
214            .map_err(|_| SessionError::ChannelClosed)
215    }
216}
217
218impl std::fmt::Debug for SessionHandle {
219    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
220        f.debug_struct("SessionHandle")
221            .field("session_id", &self.state.session_id)
222            .field("phase", &self.state.phase())
223            .finish()
224    }
225}
226
227// ---------------------------------------------------------------------------
228// Trait implementations for SessionHandle
229// ---------------------------------------------------------------------------
230
231#[async_trait]
232impl SessionWriter for SessionHandle {
233    async fn send_audio(&self, data: Bytes) -> Result<(), SessionError> {
234        SessionHandle::send_audio(self, data).await
235    }
236
237    async fn send_text(&self, text: String) -> Result<(), SessionError> {
238        self.send_command(SessionCommand::SendText(text)).await
239    }
240
241    async fn send_tool_response(
242        &self,
243        responses: Vec<FunctionResponse>,
244    ) -> Result<(), SessionError> {
245        self.send_command(SessionCommand::SendToolResponse(responses))
246            .await
247    }
248
249    async fn send_client_content(
250        &self,
251        turns: Vec<Content>,
252        turn_complete: bool,
253    ) -> Result<(), SessionError> {
254        self.send_command(SessionCommand::SendClientContent {
255            turns,
256            turn_complete,
257        })
258        .await
259    }
260
261    async fn send_video(&self, jpeg_data: Bytes) -> Result<(), SessionError> {
262        self.send_command(SessionCommand::SendVideo(jpeg_data))
263            .await
264    }
265
266    async fn update_instruction(&self, instruction: String) -> Result<(), SessionError> {
267        self.send_command(SessionCommand::UpdateInstruction(instruction))
268            .await
269    }
270
271    async fn signal_activity_start(&self) -> Result<(), SessionError> {
272        self.send_command(SessionCommand::ActivityStart).await
273    }
274
275    async fn signal_activity_end(&self) -> Result<(), SessionError> {
276        self.send_command(SessionCommand::ActivityEnd).await
277    }
278
279    async fn disconnect(&self) -> Result<(), SessionError> {
280        self.send_command(SessionCommand::Disconnect).await
281    }
282}
283
284impl SessionReader for SessionHandle {
285    fn subscribe(&self) -> broadcast::Receiver<SessionEvent> {
286        self.event_tx.subscribe()
287    }
288
289    fn phase(&self) -> SessionPhase {
290        self.state.phase()
291    }
292
293    fn session_id(&self) -> &str {
294        &self.state.session_id
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[tokio::test(start_paused = true)]
303    async fn audio_pacing_throttles_producer_to_sustained_rate() {
304        let (command_tx, mut command_rx) = mpsc::channel(64);
305        let (event_tx, _) = broadcast::channel(16);
306        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Active);
307        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
308
309        // 1000-byte burst allowance, 1000 B/s sustained.
310        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx).with_audio_pacing(
311            crate::transport::BackpressureConfig {
312                bucket_capacity: 1000,
313                refill_rate_bps: 1000,
314            },
315        );
316
317        let start = tokio::time::Instant::now();
318        // First 1000 bytes ride the burst; the next 1000 must wait ~1s.
319        handle.send_audio(vec![0u8; 1000]).await.unwrap();
320        let after_burst = start.elapsed();
321        handle.send_audio(vec![0u8; 1000]).await.unwrap();
322        let after_paced = start.elapsed();
323
324        assert!(after_burst < std::time::Duration::from_millis(50));
325        assert!(
326            after_paced >= std::time::Duration::from_millis(900),
327            "second send should be paced ~1s, was {after_paced:?}"
328        );
329        // Both frames were enqueued.
330        assert!(command_rx.recv().await.is_some());
331        assert!(command_rx.recv().await.is_some());
332    }
333
334    #[tokio::test]
335    async fn session_handle_join_returns_ok_after_task_completes() {
336        let (command_tx, _command_rx) = mpsc::channel(8);
337        let (event_tx, _) = broadcast::channel(16);
338        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
339        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
340
341        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
342
343        // Spawn a trivial task that completes immediately
344        let task = tokio::spawn(async {});
345        handle.set_task(task);
346
347        // join() should return Ok(())
348        let result = handle.join().await;
349        assert!(
350            result.is_ok(),
351            "join() should return Ok after task completes"
352        );
353    }
354
355    #[tokio::test]
356    async fn session_handle_join_without_task_returns_ok() {
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        // join() without set_task should return Ok immediately
365        let result = handle.join().await;
366        assert!(result.is_ok(), "join() without task should return Ok");
367    }
368
369    #[tokio::test]
370    async fn session_handle_join_idempotent() {
371        let (command_tx, _command_rx) = mpsc::channel(8);
372        let (event_tx, _) = broadcast::channel(16);
373        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
374        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
375
376        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
377
378        let task = tokio::spawn(async {});
379        handle.set_task(task);
380
381        // First join takes the handle
382        assert!(handle.join().await.is_ok());
383        // Second join returns Ok immediately (handle already taken)
384        assert!(handle.join().await.is_ok());
385    }
386
387    #[tokio::test]
388    async fn session_handle_join_works_on_clone() {
389        let (command_tx, _command_rx) = mpsc::channel(8);
390        let (event_tx, _) = broadcast::channel(16);
391        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
392        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
393
394        let handle = SessionHandle::new(command_tx, event_tx, state, phase_rx);
395        let handle_clone = handle.clone();
396
397        let task = tokio::spawn(async {});
398        handle.set_task(task);
399
400        // join() on clone should work (shares the Arc)
401        let result = handle_clone.join().await;
402        assert!(result.is_ok(), "join() on clone should work");
403
404        // Original handle's join should now return Ok (handle already taken)
405        assert!(handle.join().await.is_ok());
406    }
407
408    // PhaseChanged event emission tests
409
410    #[tokio::test]
411    async fn phase_changed_event_emitted_on_transition() {
412        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Disconnected);
413        let (event_tx, mut event_rx) = broadcast::channel(16);
414        let state = SessionState::with_events(phase_tx, event_tx);
415
416        state.transition_to(SessionPhase::Connecting).unwrap();
417
418        match event_rx.try_recv() {
419            Ok(SessionEvent::PhaseChanged(SessionPhase::Connecting)) => {}
420            other => panic!("expected PhaseChanged(Connecting), got {other:?}"),
421        }
422    }
423
424    #[test]
425    fn phase_changed_not_emitted_without_event_tx() {
426        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Disconnected);
427        let state = SessionState::new(phase_tx);
428        // Should not panic even though no event_tx
429        state.transition_to(SessionPhase::Connecting).unwrap();
430    }
431}
gemini_genai_rs/session/handle.rs

gemini_genai_rs/session/
handle.rs
