gemini_adk_rs/
agent_session.rs

1//! AgentSession — intercepting wrapper around SessionHandle.
2//!
3//! Replaces ADK Python's LiveRequestQueue. Instead of adding a second queue
4//! on top of SessionHandle's existing mpsc channel, this wraps a SessionWriter
5//! and intercepts sends for: (1) input fan-out to streaming tools,
6//! (2) middleware hooks, (3) state tracking.
7//!
8//! Data flow: App → AgentSession → SessionWriter → WebSocket
9//!                                ↘ broadcast to input-streaming tools
10//!
11//! ONE queue, ONE consumer task, zero-copy on the hot path.
12
13use gemini_genai_rs::prelude::{Content, FunctionResponse};
14use gemini_genai_rs::session::{SessionError, SessionEvent, SessionHandle, SessionWriter};
15use std::sync::Arc;
16use tokio::sync::broadcast;
17
18use crate::error::AgentError;
19use crate::state::State;
20
21/// Input events broadcast to input-streaming tools.
22/// Distinct from SessionCommand — this is observation-only.
23#[derive(Debug, Clone)]
24pub enum InputEvent {
25    /// Raw PCM16 audio bytes.
26    Audio(bytes::Bytes),
27    /// Text content.
28    Text(String),
29    /// User started speaking.
30    ActivityStart,
31    /// User stopped speaking.
32    ActivityEnd,
33}
34
35/// Intercepting wrapper around a SessionWriter.
36///
37/// Adds input fan-out, middleware hooks, and state tracking without
38/// introducing a second queue (avoids double-queuing).
39#[derive(Clone)]
40pub struct AgentSession {
41    /// The underlying wire-level session writer (Layer 0).
42    writer: Arc<dyn SessionWriter>,
43    /// Event subscription source.
44    event_tx: broadcast::Sender<SessionEvent>,
45    /// Fan-out for input-streaming tools.
46    /// Zero-cost when no tools are subscribed (receiver_count == 0).
47    input_broadcast: broadcast::Sender<InputEvent>,
48    /// Conversation state container.
49    state: State,
50}
51
52impl AgentSession {
53    /// Create a new AgentSession wrapping a SessionHandle.
54    pub fn new(session: SessionHandle) -> Self {
55        let (input_broadcast, _) = broadcast::channel(256);
56        let event_tx = session.event_sender().clone();
57        Self {
58            writer: Arc::new(session),
59            event_tx,
60            input_broadcast,
61            state: State::new(),
62        }
63    }
64
65    /// Create from a trait-object writer (enables mock testing and middleware injection).
66    pub fn from_writer(
67        writer: Arc<dyn SessionWriter>,
68        event_tx: broadcast::Sender<SessionEvent>,
69    ) -> Self {
70        let (input_broadcast, _) = broadcast::channel(256);
71        Self {
72            writer,
73            event_tx,
74            input_broadcast,
75            state: State::new(),
76        }
77    }
78
79    /// Send audio data. Fans out to input-streaming tools ONLY if listeners exist.
80    pub async fn send_audio(&self, data: impl Into<bytes::Bytes>) -> Result<(), AgentError> {
81        let data: bytes::Bytes = data.into();
82        // Fan-out ONLY if input-streaming tools are listening
83        if self.input_broadcast.receiver_count() > 0 {
84            let _ = self.input_broadcast.send(InputEvent::Audio(data.clone()));
85        }
86        // Forward directly to Layer 0 (ONE hop to WebSocket)
87        self.writer
88            .send_audio(data)
89            .await
90            .map_err(AgentError::Session)
91    }
92
93    /// Send a text message.
94    pub async fn send_text(&self, text: impl Into<String>) -> Result<(), AgentError> {
95        let t = text.into();
96        if self.input_broadcast.receiver_count() > 0 {
97            let _ = self.input_broadcast.send(InputEvent::Text(t.clone()));
98        }
99        self.writer.send_text(t).await.map_err(AgentError::Session)
100    }
101
102    /// Send tool responses.
103    pub async fn send_tool_response(
104        &self,
105        responses: Vec<FunctionResponse>,
106    ) -> Result<(), AgentError> {
107        self.writer
108            .send_tool_response(responses)
109            .await
110            .map_err(AgentError::Session)
111    }
112
113    /// Send client content (conversation history or context injection).
114    pub async fn send_client_content(
115        &self,
116        turns: Vec<Content>,
117        turn_complete: bool,
118    ) -> Result<(), AgentError> {
119        self.writer
120            .send_client_content(turns, turn_complete)
121            .await
122            .map_err(AgentError::Session)
123    }
124
125    /// Send video/image data (raw JPEG bytes).
126    pub async fn send_video(&self, jpeg_data: impl Into<bytes::Bytes>) -> Result<(), AgentError> {
127        self.writer
128            .send_video(jpeg_data.into())
129            .await
130            .map_err(AgentError::Session)
131    }
132
133    /// Update the system instruction mid-session.
134    pub async fn update_instruction(
135        &self,
136        instruction: impl Into<String>,
137    ) -> Result<(), AgentError> {
138        self.writer
139            .update_instruction(instruction.into())
140            .await
141            .map_err(AgentError::Session)
142    }
143
144    /// Signal activity start (user started speaking).
145    pub async fn signal_activity_start(&self) -> Result<(), AgentError> {
146        if self.input_broadcast.receiver_count() > 0 {
147            let _ = self.input_broadcast.send(InputEvent::ActivityStart);
148        }
149        self.writer
150            .signal_activity_start()
151            .await
152            .map_err(AgentError::Session)
153    }
154
155    /// Signal activity end (user stopped speaking).
156    pub async fn signal_activity_end(&self) -> Result<(), AgentError> {
157        if self.input_broadcast.receiver_count() > 0 {
158            let _ = self.input_broadcast.send(InputEvent::ActivityEnd);
159        }
160        self.writer
161            .signal_activity_end()
162            .await
163            .map_err(AgentError::Session)
164    }
165
166    /// Gracefully disconnect.
167    pub async fn disconnect(&self) -> Result<(), AgentError> {
168        self.writer.disconnect().await.map_err(AgentError::Session)
169    }
170
171    /// Subscribe to input events (for input-streaming tools).
172    pub fn subscribe_input(&self) -> broadcast::Receiver<InputEvent> {
173        self.input_broadcast.subscribe()
174    }
175
176    /// Subscribe to session events.
177    pub fn subscribe_events(&self) -> broadcast::Receiver<SessionEvent> {
178        self.event_tx.subscribe()
179    }
180
181    /// Access the underlying session writer.
182    pub fn writer(&self) -> &dyn SessionWriter {
183        &*self.writer
184    }
185
186    /// Access conversation state.
187    pub fn state(&self) -> &State {
188        &self.state
189    }
190
191    /// Number of input-streaming subscribers (for diagnostics).
192    pub fn input_subscriber_count(&self) -> usize {
193        self.input_broadcast.receiver_count()
194    }
195}
196
197/// A SessionWriter that discards all writes.
198/// Used for isolated agent execution (AgentTool) where no real WebSocket exists.
199pub struct NoOpSessionWriter;
200
201#[async_trait::async_trait]
202impl SessionWriter for NoOpSessionWriter {
203    async fn send_audio(&self, _data: bytes::Bytes) -> Result<(), SessionError> {
204        Ok(())
205    }
206    async fn send_text(&self, _text: String) -> Result<(), SessionError> {
207        Ok(())
208    }
209    async fn send_tool_response(
210        &self,
211        _responses: Vec<FunctionResponse>,
212    ) -> Result<(), SessionError> {
213        Ok(())
214    }
215    async fn send_client_content(
216        &self,
217        _turns: Vec<Content>,
218        _turn_complete: bool,
219    ) -> Result<(), SessionError> {
220        Ok(())
221    }
222    async fn send_video(&self, _jpeg_data: bytes::Bytes) -> Result<(), SessionError> {
223        Ok(())
224    }
225    async fn update_instruction(&self, _instruction: String) -> Result<(), SessionError> {
226        Ok(())
227    }
228    async fn signal_activity_start(&self) -> Result<(), SessionError> {
229        Ok(())
230    }
231    async fn signal_activity_end(&self) -> Result<(), SessionError> {
232        Ok(())
233    }
234    async fn disconnect(&self) -> Result<(), SessionError> {
235        Ok(())
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use gemini_genai_rs::session::{SessionHandle, SessionPhase, SessionState};
243    use std::sync::Arc;
244    use tokio::sync::{broadcast, mpsc, watch};
245
246    fn mock_session_handle() -> SessionHandle {
247        let (cmd_tx, _cmd_rx) = mpsc::channel(16);
248        let (evt_tx, _) = broadcast::channel(16);
249        let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
250        let state = Arc::new(SessionState::new(phase_tx));
251        SessionHandle::new(cmd_tx, evt_tx, state, phase_rx)
252    }
253
254    #[tokio::test]
255    async fn send_audio_without_subscribers_no_broadcast() {
256        let handle = mock_session_handle();
257        let session = AgentSession::new(handle);
258        assert_eq!(session.input_subscriber_count(), 0);
259    }
260
261    #[tokio::test]
262    async fn send_audio_with_subscriber_broadcasts() {
263        let handle = mock_session_handle();
264        let session = AgentSession::new(handle);
265        let mut input_rx = session.subscribe_input();
266        assert_eq!(session.input_subscriber_count(), 1);
267
268        // send_audio will fail at SessionHandle level (no real WS), but
269        // the broadcast should still fire
270        let data = vec![1, 2, 3, 4];
271        let _ = session.send_audio(bytes::Bytes::from(data.clone())).await;
272
273        match input_rx.try_recv() {
274            Ok(InputEvent::Audio(received)) => assert_eq!(received, data),
275            other => panic!("expected Audio, got {other:?}"),
276        }
277    }
278
279    #[test]
280    fn agent_session_is_clone() {
281        let handle = mock_session_handle();
282        let session = AgentSession::new(handle);
283        let _clone = session.clone();
284    }
285
286    #[test]
287    fn state_accessible() {
288        let handle = mock_session_handle();
289        let session = AgentSession::new(handle);
290        let _ = session.state().set("key", "value");
291        assert_eq!(
292            session.state().get::<String>("key"),
293            Some("value".to_string())
294        );
295    }
296
297    #[tokio::test]
298    async fn text_broadcast() {
299        let handle = mock_session_handle();
300        let session = AgentSession::new(handle);
301        let mut input_rx = session.subscribe_input();
302
303        let _ = session.send_text("hello").await;
304
305        match input_rx.try_recv() {
306            Ok(InputEvent::Text(t)) => assert_eq!(t, "hello"),
307            other => panic!("expected Text, got {other:?}"),
308        }
309    }
310
311    #[tokio::test]
312    async fn activity_signals_broadcast() {
313        let handle = mock_session_handle();
314        let session = AgentSession::new(handle);
315        let mut input_rx = session.subscribe_input();
316
317        let _ = session.signal_activity_start().await;
318        let _ = session.signal_activity_end().await;
319
320        assert!(matches!(input_rx.try_recv(), Ok(InputEvent::ActivityStart)));
321        assert!(matches!(input_rx.try_recv(), Ok(InputEvent::ActivityEnd)));
322    }
323
324    #[tokio::test]
325    async fn from_writer_with_mock() {
326        // Create a mock writer using a real SessionHandle (simplest mock available)
327        let handle = mock_session_handle();
328        let event_tx = handle.event_sender().clone();
329        let writer: Arc<dyn SessionWriter> = Arc::new(handle);
330        let session = AgentSession::from_writer(writer, event_tx);
331        assert_eq!(session.input_subscriber_count(), 0);
332    }
333}
gemini_adk_rs/agent_session.rs

gemini_adk_rs/
agent_session.rs
