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

gemini_adk_rs/
agent_session.rs
