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

gemini_genai_rs/transport/
ws.rs
