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::Message;
66use tokio_tungstenite::tungstenite::client::IntoClientRequest;
67
68type WsStream =
69    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
70
71/// Bytes tungstenite reads from the socket per frame read.
72///
73/// Its default, 128 KiB, is allocated and zero-filled for every connection:
74/// over a thousand sessions that is 128 MB for buffers that hold one
75/// message at a time. 16 KiB still takes a 40 ms audio message (about
76/// 2.7 KB) or an 8 KB-PCM one (about 11 KB) in a single read; anything
77/// larger is read in more than one.
78const READ_BUFFER_SIZE: usize = 16 * 1024;
79
80/// WebSocket transport using `tokio-tungstenite`.
81pub struct TungsteniteTransport {
82    ws_write: Option<futures_util::stream::SplitSink<WsStream, Message>>,
83    ws_read: Option<futures_util::stream::SplitStream<WsStream>>,
84    /// Code and reason from the peer's close frame, kept so the session loop can
85    /// report *why* the connection ended rather than only that it did.
86    close_reason: Option<String>,
87}
88
89impl TungsteniteTransport {
90    /// Create a new, disconnected transport.
91    pub fn new() -> Self {
92        Self {
93            ws_write: None,
94            ws_read: None,
95            close_reason: None,
96        }
97    }
98}
99
100impl Default for TungsteniteTransport {
101    fn default() -> Self {
102        Self::new()
103    }
104}
105
106/// Errors from the [`TungsteniteTransport`].
107#[derive(Debug, thiserror::Error)]
108pub enum TungsteniteError {
109    /// The transport is not connected.
110    #[error("Not connected")]
111    NotConnected,
112
113    /// WebSocket protocol error from tungstenite.
114    #[error("WebSocket error: {0}")]
115    WebSocket(#[from] tokio_tungstenite::tungstenite::Error),
116
117    /// Failed to construct the HTTP request (e.g. bad URL or header).
118    #[error("Request error: {0}")]
119    Request(String),
120
121    /// A `wss://` URL was dialed but the crate was compiled with neither TLS
122    /// backend.
123    ///
124    /// Without `tls-native` or `tls-rustls`, `tokio-tungstenite` can open
125    /// plain `ws://` sockets only; the Live API is TLS-only, so this build
126    /// cannot reach it. The check runs before any socket is opened, so the
127    /// failure names the fix instead of surfacing as a generic handshake error
128    /// from deep inside the WebSocket stack.
129    #[error(
130        "cannot dial {url}: this build of gemini-genai-rs has no TLS backend; \
131         enable the `tls-native` (default) or `tls-rustls` feature"
132    )]
133    NoTlsBackend {
134        /// The URL that was refused.
135        url: String,
136    },
137}
138
139/// Whether a TLS backend was compiled in. `wss://` and `https://` need one.
140pub const HAS_TLS_BACKEND: bool = cfg!(any(feature = "tls-native", feature = "tls-rustls"));
141
142#[async_trait]
143impl Transport for TungsteniteTransport {
144    type Error = TungsteniteError;
145
146    async fn connect(
147        &mut self,
148        url: &str,
149        headers: Vec<(String, String)>,
150    ) -> Result<(), Self::Error> {
151        if !HAS_TLS_BACKEND && url.trim_start().to_ascii_lowercase().starts_with("wss://") {
152            return Err(TungsteniteError::NoTlsBackend {
153                url: url.to_string(),
154            });
155        }
156
157        let mut request = url
158            .into_client_request()
159            .map_err(|e| TungsteniteError::Request(e.to_string()))?;
160
161        for (name, value) in headers {
162            let header_name: tokio_tungstenite::tungstenite::http::HeaderName =
163                name.parse().map_err(
164                    |e: tokio_tungstenite::tungstenite::http::header::InvalidHeaderName| {
165                        TungsteniteError::Request(format!("invalid header name: {e}"))
166                    },
167                )?;
168            let header_value: tokio_tungstenite::tungstenite::http::HeaderValue =
169                value.parse().map_err(
170                    |e: tokio_tungstenite::tungstenite::http::header::InvalidHeaderValue| {
171                        TungsteniteError::Request(format!("invalid header value: {e}"))
172                    },
173                )?;
174            request.headers_mut().insert(header_name, header_value);
175        }
176
177        let config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default()
178            .read_buffer_size(READ_BUFFER_SIZE);
179        // Nagle off: audio goes out as a stream of small messages, and a
180        // small write must not wait for the previous one to be acknowledged.
181        let (ws_stream, _response) =
182            tokio_tungstenite::connect_async_with_config(request, Some(config), true).await?;
183        let (ws_write, ws_read) = ws_stream.split();
184        self.ws_write = Some(ws_write);
185        self.ws_read = Some(ws_read);
186        Ok(())
187    }
188
189    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
190        let ws_write = self
191            .ws_write
192            .as_mut()
193            .ok_or(TungsteniteError::NotConnected)?;
194        // Convert bytes to a UTF-8 text frame. The wire protocol sends JSON as text.
195        let text = String::from_utf8(data)
196            .map_err(|e| TungsteniteError::Request(format!("invalid UTF-8: {e}")))?;
197        ws_write.send(Message::Text(text.into())).await?;
198        Ok(())
199    }
200
201    async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
202        let ws_read = self
203            .ws_read
204            .as_mut()
205            .ok_or(TungsteniteError::NotConnected)?;
206        loop {
207            match ws_read.next().await {
208                Some(Ok(Message::Text(t))) => return Ok(Some(bytes::Bytes::from(t).into())),
209                // IMPORTANT: Vertex AI sends JSON in Binary frames.
210                Some(Ok(Message::Binary(b))) => return Ok(Some(b.into())),
211                Some(Ok(Message::Close(frame))) => {
212                    if let Some(ref cf) = frame {
213                        tracing::warn!(code = %cf.code, reason = %cf.reason, "WebSocket close frame received");
214                        // Kept, not merely logged: a `warn` is invisible unless
215                        // the application configured a subscriber, and the
216                        // reason is the only account of why the session ended.
217                        self.close_reason = Some(if cf.reason.is_empty() {
218                            format!("server closed the connection ({})", cf.code)
219                        } else {
220                            format!("server closed the connection ({}): {}", cf.code, cf.reason)
221                        });
222                    }
223                    return Ok(None);
224                }
225                // Ping/Pong are handled internally by tungstenite; skip them.
226                Some(Ok(Message::Ping(_) | Message::Pong(_))) => continue,
227                // Frame is a low-level variant; skip.
228                Some(Ok(Message::Frame(_))) => continue,
229                Some(Err(e)) => return Err(e.into()),
230                None => return Ok(None),
231            }
232        }
233    }
234
235    async fn close(&mut self) -> Result<(), Self::Error> {
236        if let Some(ref mut ws_write) = self.ws_write {
237            ws_write.send(Message::Close(None)).await?;
238        }
239        self.ws_write = None;
240        self.ws_read = None;
241        Ok(())
242    }
243
244    fn close_reason(&self) -> Option<String> {
245        self.close_reason.clone()
246    }
247}
248
249// ---------------------------------------------------------------------------
250// MockTransport — for unit testing
251// ---------------------------------------------------------------------------
252
253/// Mock transport for unit testing.
254///
255/// Records sent data and replays scripted responses from a queue. When the
256/// queue is empty and the transport is connected, [`recv`](Transport::recv)
257/// will pend indefinitely — simulating a connected-but-idle transport.
258/// Call [`close`](Transport::close) to signal connection closure (returns `None`).
259pub struct MockTransport {
260    sent: Vec<Vec<u8>>,
261    recv_queue: std::collections::VecDeque<Vec<u8>>,
262    /// Whether connect() has been called (and close() has not).
263    connected: bool,
264    /// Scripted peer close reason, surfaced by [`Transport::close_reason`].
265    close_reason: Option<String>,
266}
267
268impl MockTransport {
269    /// Create a new, disconnected mock transport.
270    pub fn new() -> Self {
271        Self {
272            sent: Vec::new(),
273            recv_queue: std::collections::VecDeque::new(),
274            connected: false,
275            close_reason: None,
276        }
277    }
278
279    /// Queue a message to be returned by [`Transport::recv`].
280    pub fn script_recv(&mut self, data: Vec<u8>) {
281        self.recv_queue.push_back(data);
282    }
283
284    /// Script the reason the peer gives for closing, so a test can assert that
285    /// it reaches the application rather than being swallowed.
286    pub fn script_close_reason(&mut self, reason: impl Into<String>) {
287        self.close_reason = Some(reason.into());
288    }
289
290    /// Take all sent data (for assertions). Drains the internal buffer.
291    pub fn take_sent(&mut self) -> Vec<Vec<u8>> {
292        std::mem::take(&mut self.sent)
293    }
294}
295
296impl Default for MockTransport {
297    fn default() -> Self {
298        Self::new()
299    }
300}
301
302/// Errors from the [`MockTransport`].
303#[derive(Debug, thiserror::Error)]
304pub enum MockTransportError {
305    /// Operation attempted while not connected.
306    #[error("Not connected")]
307    NotConnected,
308
309    /// A custom error injected for testing.
310    #[error("Mock error: {0}")]
311    Custom(String),
312}
313
314#[async_trait]
315impl Transport for MockTransport {
316    type Error = MockTransportError;
317
318    async fn connect(
319        &mut self,
320        _url: &str,
321        _headers: Vec<(String, String)>,
322    ) -> Result<(), Self::Error> {
323        self.connected = true;
324        Ok(())
325    }
326
327    async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
328        if !self.connected {
329            return Err(MockTransportError::NotConnected);
330        }
331        self.sent.push(data);
332        Ok(())
333    }
334
335    async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
336        if !self.connected {
337            return Err(MockTransportError::NotConnected);
338        }
339        // Yield to the scheduler so tests can observe intermediate states
340        // (phase transitions, events) before the next message is processed.
341        tokio::task::yield_now().await;
342
343        if let Some(data) = self.recv_queue.pop_front() {
344            return Ok(Some(data));
345        }
346
347        // Queue is empty: pend indefinitely, simulating a connected-but-idle
348        // transport waiting for the next message from the server.
349        // The connection loop uses `tokio::select!` so this future is dropped
350        // when a command (e.g., Disconnect) arrives on the command channel.
351        std::future::pending().await
352    }
353
354    async fn close(&mut self) -> Result<(), Self::Error> {
355        self.connected = false;
356        Ok(())
357    }
358
359    fn close_reason(&self) -> Option<String> {
360        self.close_reason.clone()
361    }
362}
363
364// ---------------------------------------------------------------------------
365// Tests
366// ---------------------------------------------------------------------------
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    /// A transport that says nothing about why it closed still compiles and
373    /// still reports nothing — the default keeps third-party implementations
374    /// working rather than forcing them to answer a question they cannot.
375    #[test]
376    fn a_transport_that_says_nothing_reports_nothing() {
377        assert_eq!(MockTransport::new().close_reason(), None);
378    }
379
380    /// When the peer gives a reason, it survives to the caller.
381    ///
382    /// It used to be logged at `warn` and dropped: `recv` returned a bare
383    /// `Ok(None)` and the session loop turned that into the literal string
384    /// "Transport closed". A `warn` is invisible without a configured
385    /// subscriber, so in practice the only account of why a session ended was
386    /// discarded at the point it was produced — which is why a governed voice
387    /// evaluation lost turns to a close it could never diagnose.
388    #[test]
389    fn a_stated_close_reason_survives_to_the_caller() {
390        let mut transport = MockTransport::new();
391        transport.script_close_reason("server closed the connection (1011): quota exceeded");
392        assert_eq!(
393            transport.close_reason().as_deref(),
394            Some("server closed the connection (1011): quota exceeded")
395        );
396    }
397
398    /// Only meaningful in a build with no TLS backend (`--no-default-features`):
399    /// the `wss://` dial must fail before any socket is opened, with an error
400    /// that names the feature to enable.
401    #[cfg(not(any(feature = "tls-native", feature = "tls-rustls")))]
402    #[tokio::test]
403    async fn wss_without_tls_backend_fails_before_dialing() {
404        let mut transport = TungsteniteTransport::new();
405        let err = transport
406            .connect("wss://example.invalid/ws", vec![])
407            .await
408            .expect_err("no TLS backend compiled in");
409        assert!(matches!(err, TungsteniteError::NoTlsBackend { .. }));
410        let text = err.to_string();
411        assert!(text.contains("tls-native"), "{text}");
412        assert!(text.contains("tls-rustls"), "{text}");
413    }
414
415    #[test]
416    fn has_tls_backend_tracks_features() {
417        assert_eq!(
418            HAS_TLS_BACKEND,
419            cfg!(any(feature = "tls-native", feature = "tls-rustls"))
420        );
421    }
422
423    #[tokio::test]
424    async fn mock_transport_round_trip() {
425        let mut transport = MockTransport::new();
426        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
427
428        transport
429            .connect("wss://example.com", vec![])
430            .await
431            .unwrap();
432        transport.send(b"hello".to_vec()).await.unwrap();
433        let data = transport.recv().await.unwrap();
434        assert!(data.is_some());
435        let text = String::from_utf8(data.unwrap()).unwrap();
436        assert!(text.contains("setupComplete"));
437    }
438
439    #[tokio::test]
440    async fn mock_transport_records_sent() {
441        let mut transport = MockTransport::new();
442        transport
443            .connect("wss://example.com", vec![])
444            .await
445            .unwrap();
446        transport.send(b"msg1".to_vec()).await.unwrap();
447        transport.send(b"msg2".to_vec()).await.unwrap();
448        let sent = transport.take_sent();
449        assert_eq!(sent.len(), 2);
450        assert_eq!(sent[0], b"msg1");
451    }
452
453    #[tokio::test]
454    async fn mock_transport_recv_pends_when_queue_empty() {
455        let mut transport = MockTransport::new();
456        transport
457            .connect("wss://example.com", vec![])
458            .await
459            .unwrap();
460        // recv() should pend when queue is empty (simulating idle transport)
461        let result =
462            tokio::time::timeout(std::time::Duration::from_millis(50), transport.recv()).await;
463        assert!(result.is_err(), "recv should pend when queue is empty");
464    }
465
466    #[tokio::test]
467    async fn mock_transport_recv_errors_when_not_connected() {
468        let mut transport = MockTransport::new();
469        // Not connected yet — recv should error
470        let result = transport.recv().await;
471        assert!(result.is_err());
472    }
473
474    #[tokio::test]
475    async fn mock_transport_not_connected_error() {
476        let mut transport = MockTransport::new();
477        let result = transport.send(b"data".to_vec()).await;
478        assert!(result.is_err());
479    }
480
481    #[test]
482    fn transport_trait_is_object_safe_check() {
483        // Transport has an associated type, so it's not directly object-safe
484        // but can be used as generic bounds. This test just verifies compilation.
485        fn _assert_transport<T: Transport>() {}
486        _assert_transport::<MockTransport>();
487    }
488}
gemini_genai_rs/transport/ws.rs

gemini_genai_rs/transport/
ws.rs
