gemini_genai_rs/transport/
ws.rs

1//! Transport abstraction — bidirectional message transport.
2//!
3//! The [`Transport`] trait defines a pluggable transport layer for sending and
4//! receiving raw bytes. The default implementation [`TungsteniteTransport`] wraps
5//! `tokio-tungstenite` for WebSocket connectivity. [`MockTransport`] enables
6//! deterministic unit testing without a network.
7
8use async_trait::async_trait;
9
10/// A bidirectional message transport.
11///
12/// The default is WebSocket ([`TungsteniteTransport`]); [`MockTransport`] enables
13/// unit testing without a real server.
14///
15/// # Implementors
16///
17/// - [`TungsteniteTransport`] -- Production WebSocket transport using `tokio-tungstenite`.
18///   Handles both Text and Binary frames (Vertex AI sends Binary).
19/// - [`MockTransport`] -- Deterministic test transport. Records sent data and replays
20///   scripted responses. When the queue is empty, `recv()` pends indefinitely.
21#[async_trait]
22pub trait Transport: Send + 'static {
23    /// The error type produced by this transport.
24    type Error: std::error::Error + Send + Sync + 'static;
25
26    /// Connect to the given URL with optional headers.
27    async fn connect(
28        &mut self,
29        url: &str,
30        headers: Vec<(String, String)>,
31    ) -> Result<(), Self::Error>;
32
33    /// Send raw bytes.
34    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error>;
35
36    /// Receive raw bytes. Returns `None` when the connection is closed.
37    async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error>;
38
39    /// Close the transport.
40    async fn close(&mut self) -> Result<(), Self::Error>;
41
42    /// Why the peer closed the connection, if it said.
43    ///
44    /// Read after [`recv`](Self::recv) returns `Ok(None)`. A server-initiated
45    /// close on the Live API carries a status code and a reason — a context
46    /// window exhausted, an invalid argument, a quota — and without this the
47    /// only thing that survives is the fact of the close. That is the difference
48    /// between a session that "just dropped" and one that dropped for a stated
49    /// reason: a governed voice evaluation lost several turns to a close whose
50    /// reason had been logged at `warn` and then discarded, so it was never
51    /// diagnosed.
52    ///
53    /// Defaults to `None` so an existing transport implementation keeps
54    /// compiling; it merely reports nothing.
55    fn close_reason(&self) -> Option<String> {
56        None
57    }
58}
59
60// ---------------------------------------------------------------------------
61// TungsteniteTransport — WebSocket transport using tokio-tungstenite
62// ---------------------------------------------------------------------------
63
64use futures_util::{SinkExt, StreamExt};
65use tokio_tungstenite::tungstenite::client::IntoClientRequest;
66use tokio_tungstenite::tungstenite::Message;
67
68type WsStream =
69    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
70
71/// WebSocket transport using `tokio-tungstenite`.
72pub struct TungsteniteTransport {
73    ws_write: Option<futures_util::stream::SplitSink<WsStream, Message>>,
74    ws_read: Option<futures_util::stream::SplitStream<WsStream>>,
75    /// Code and reason from the peer's close frame, kept so the session loop can
76    /// report *why* the connection ended rather than only that it did.
77    close_reason: Option<String>,
78}
79
80impl TungsteniteTransport {
81    /// Create a new, disconnected transport.
82    pub fn new() -> Self {
83        Self {
84            ws_write: None,
85            ws_read: None,
86            close_reason: None,
87        }
88    }
89}
90
91impl Default for TungsteniteTransport {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97/// Errors from the [`TungsteniteTransport`].
98#[derive(Debug, thiserror::Error)]
99pub enum TungsteniteError {
100    /// The transport is not connected.
101    #[error("Not connected")]
102    NotConnected,
103
104    /// WebSocket protocol error from tungstenite.
105    #[error("WebSocket error: {0}")]
106    WebSocket(#[from] tokio_tungstenite::tungstenite::Error),
107
108    /// Failed to construct the HTTP request (e.g. bad URL or header).
109    #[error("Request error: {0}")]
110    Request(String),
111}
112
113#[async_trait]
114impl Transport for TungsteniteTransport {
115    type Error = TungsteniteError;
116
117    async fn connect(
118        &mut self,
119        url: &str,
120        headers: Vec<(String, String)>,
121    ) -> Result<(), Self::Error> {
122        let mut request = url
123            .into_client_request()
124            .map_err(|e| TungsteniteError::Request(e.to_string()))?;
125
126        for (name, value) in headers {
127            let header_name: tokio_tungstenite::tungstenite::http::HeaderName =
128                name.parse().map_err(
129                    |e: tokio_tungstenite::tungstenite::http::header::InvalidHeaderName| {
130                        TungsteniteError::Request(format!("invalid header name: {e}"))
131                    },
132                )?;
133            let header_value: tokio_tungstenite::tungstenite::http::HeaderValue =
134                value.parse().map_err(
135                    |e: tokio_tungstenite::tungstenite::http::header::InvalidHeaderValue| {
136                        TungsteniteError::Request(format!("invalid header value: {e}"))
137                    },
138                )?;
139            request.headers_mut().insert(header_name, header_value);
140        }
141
142        let (ws_stream, _response) = tokio_tungstenite::connect_async(request).await?;
143        let (ws_write, ws_read) = ws_stream.split();
144        self.ws_write = Some(ws_write);
145        self.ws_read = Some(ws_read);
146        Ok(())
147    }
148
149    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
150        let ws_write = self
151            .ws_write
152            .as_mut()
153            .ok_or(TungsteniteError::NotConnected)?;
154        // Convert bytes to a UTF-8 text frame. The wire protocol sends JSON as text.
155        let text = String::from_utf8(data)
156            .map_err(|e| TungsteniteError::Request(format!("invalid UTF-8: {e}")))?;
157        ws_write.send(Message::Text(text)).await?;
158        Ok(())
159    }
160
161    async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
162        let ws_read = self
163            .ws_read
164            .as_mut()
165            .ok_or(TungsteniteError::NotConnected)?;
166        loop {
167            match ws_read.next().await {
168                Some(Ok(Message::Text(t))) => return Ok(Some(t.into_bytes())),
169                // IMPORTANT: Vertex AI sends JSON in Binary frames.
170                Some(Ok(Message::Binary(b))) => return Ok(Some(b)),
171                Some(Ok(Message::Close(frame))) => {
172                    if let Some(ref cf) = frame {
173                        tracing::warn!(code = %cf.code, reason = %cf.reason, "WebSocket close frame received");
174                        // Kept, not merely logged: a `warn` is invisible unless
175                        // the application configured a subscriber, and the
176                        // reason is the only account of why the session ended.
177                        self.close_reason = Some(if cf.reason.is_empty() {
178                            format!("server closed the connection ({})", cf.code)
179                        } else {
180                            format!("server closed the connection ({}): {}", cf.code, cf.reason)
181                        });
182                    }
183                    return Ok(None);
184                }
185                // Ping/Pong are handled internally by tungstenite; skip them.
186                Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue,
187                // Frame is a low-level variant; skip.
188                Some(Ok(Message::Frame(_))) => continue,
189                Some(Err(e)) => return Err(TungsteniteError::WebSocket(e)),
190                None => return Ok(None),
191            }
192        }
193    }
194
195    async fn close(&mut self) -> Result<(), Self::Error> {
196        if let Some(ref mut ws_write) = self.ws_write {
197            ws_write.send(Message::Close(None)).await?;
198        }
199        self.ws_write = None;
200        self.ws_read = None;
201        Ok(())
202    }
203
204    fn close_reason(&self) -> Option<String> {
205        self.close_reason.clone()
206    }
207}
208
209// ---------------------------------------------------------------------------
210// MockTransport — for unit testing
211// ---------------------------------------------------------------------------
212
213/// Mock transport for unit testing.
214///
215/// Records sent data and replays scripted responses from a queue. When the
216/// queue is empty and the transport is connected, [`recv`](Transport::recv)
217/// will pend indefinitely — simulating a connected-but-idle transport.
218/// Call [`close`](Transport::close) to signal connection closure (returns `None`).
219pub struct MockTransport {
220    sent: Vec<Vec<u8>>,
221    recv_queue: std::collections::VecDeque<Vec<u8>>,
222    /// Whether connect() has been called (and close() has not).
223    connected: bool,
224    /// Scripted peer close reason, surfaced by [`Transport::close_reason`].
225    close_reason: Option<String>,
226}
227
228impl MockTransport {
229    /// Create a new, disconnected mock transport.
230    pub fn new() -> Self {
231        Self {
232            sent: Vec::new(),
233            recv_queue: std::collections::VecDeque::new(),
234            connected: false,
235            close_reason: None,
236        }
237    }
238
239    /// Queue a message to be returned by [`Transport::recv`].
240    pub fn script_recv(&mut self, data: Vec<u8>) {
241        self.recv_queue.push_back(data);
242    }
243
244    /// Script the reason the peer gives for closing, so a test can assert that
245    /// it reaches the application rather than being swallowed.
246    pub fn script_close_reason(&mut self, reason: impl Into<String>) {
247        self.close_reason = Some(reason.into());
248    }
249
250    /// Take all sent data (for assertions). Drains the internal buffer.
251    pub fn take_sent(&mut self) -> Vec<Vec<u8>> {
252        std::mem::take(&mut self.sent)
253    }
254}
255
256impl Default for MockTransport {
257    fn default() -> Self {
258        Self::new()
259    }
260}
261
262/// Errors from the [`MockTransport`].
263#[derive(Debug, thiserror::Error)]
264pub enum MockTransportError {
265    /// Operation attempted while not connected.
266    #[error("Not connected")]
267    NotConnected,
268
269    /// A custom error injected for testing.
270    #[error("Mock error: {0}")]
271    Custom(String),
272}
273
274#[async_trait]
275impl Transport for MockTransport {
276    type Error = MockTransportError;
277
278    async fn connect(
279        &mut self,
280        _url: &str,
281        _headers: Vec<(String, String)>,
282    ) -> Result<(), Self::Error> {
283        self.connected = true;
284        Ok(())
285    }
286
287    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
288        if !self.connected {
289            return Err(MockTransportError::NotConnected);
290        }
291        self.sent.push(data);
292        Ok(())
293    }
294
295    async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
296        if !self.connected {
297            return Err(MockTransportError::NotConnected);
298        }
299        // Yield to the scheduler so tests can observe intermediate states
300        // (phase transitions, events) before the next message is processed.
301        tokio::task::yield_now().await;
302
303        if let Some(data) = self.recv_queue.pop_front() {
304            return Ok(Some(data));
305        }
306
307        // Queue is empty: pend indefinitely, simulating a connected-but-idle
308        // transport waiting for the next message from the server.
309        // The connection loop uses `tokio::select!` so this future is dropped
310        // when a command (e.g., Disconnect) arrives on the command channel.
311        std::future::pending().await
312    }
313
314    async fn close(&mut self) -> Result<(), Self::Error> {
315        self.connected = false;
316        Ok(())
317    }
318
319    fn close_reason(&self) -> Option<String> {
320        self.close_reason.clone()
321    }
322}
323
324// ---------------------------------------------------------------------------
325// Tests
326// ---------------------------------------------------------------------------
327
328#[cfg(test)]
329mod tests {
330    use super::*;
331
332    /// A transport that says nothing about why it closed still compiles and
333    /// still reports nothing — the default keeps third-party implementations
334    /// working rather than forcing them to answer a question they cannot.
335    #[test]
336    fn a_transport_that_says_nothing_reports_nothing() {
337        assert_eq!(MockTransport::new().close_reason(), None);
338    }
339
340    /// When the peer gives a reason, it survives to the caller.
341    ///
342    /// It used to be logged at `warn` and dropped: `recv` returned a bare
343    /// `Ok(None)` and the session loop turned that into the literal string
344    /// "Transport closed". A `warn` is invisible without a configured
345    /// subscriber, so in practice the only account of why a session ended was
346    /// discarded at the point it was produced — which is why a governed voice
347    /// evaluation lost turns to a close it could never diagnose.
348    #[test]
349    fn a_stated_close_reason_survives_to_the_caller() {
350        let mut transport = MockTransport::new();
351        transport.script_close_reason("server closed the connection (1011): quota exceeded");
352        assert_eq!(
353            transport.close_reason().as_deref(),
354            Some("server closed the connection (1011): quota exceeded")
355        );
356    }
357
358    #[tokio::test]
359    async fn mock_transport_round_trip() {
360        let mut transport = MockTransport::new();
361        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
362
363        transport
364            .connect("wss://example.com", vec![])
365            .await
366            .unwrap();
367        transport.send(b"hello".to_vec()).await.unwrap();
368        let data = transport.recv().await.unwrap();
369        assert!(data.is_some());
370        let text = String::from_utf8(data.unwrap()).unwrap();
371        assert!(text.contains("setupComplete"));
372    }
373
374    #[tokio::test]
375    async fn mock_transport_records_sent() {
376        let mut transport = MockTransport::new();
377        transport
378            .connect("wss://example.com", vec![])
379            .await
380            .unwrap();
381        transport.send(b"msg1".to_vec()).await.unwrap();
382        transport.send(b"msg2".to_vec()).await.unwrap();
383        let sent = transport.take_sent();
384        assert_eq!(sent.len(), 2);
385        assert_eq!(sent[0], b"msg1");
386    }
387
388    #[tokio::test]
389    async fn mock_transport_recv_pends_when_queue_empty() {
390        let mut transport = MockTransport::new();
391        transport
392            .connect("wss://example.com", vec![])
393            .await
394            .unwrap();
395        // recv() should pend when queue is empty (simulating idle transport)
396        let result =
397            tokio::time::timeout(std::time::Duration::from_millis(50), transport.recv()).await;
398        assert!(result.is_err(), "recv should pend when queue is empty");
399    }
400
401    #[tokio::test]
402    async fn mock_transport_recv_errors_when_not_connected() {
403        let mut transport = MockTransport::new();
404        // Not connected yet — recv should error
405        let result = transport.recv().await;
406        assert!(result.is_err());
407    }
408
409    #[tokio::test]
410    async fn mock_transport_not_connected_error() {
411        let mut transport = MockTransport::new();
412        let result = transport.send(b"data".to_vec()).await;
413        assert!(result.is_err());
414    }
415
416    #[test]
417    fn transport_trait_is_object_safe_check() {
418        // Transport has an associated type, so it's not directly object-safe
419        // but can be used as generic bounds. This test just verifies compilation.
420        fn _assert_transport<T: Transport>() {}
421        _assert_transport::<MockTransport>();
422    }
423}
gemini_genai_rs/transport/ws.rs

gemini_genai_rs/transport/
ws.rs
