gemini_genai_rs/transport/connection/
mod.rs

1//! WebSocket connection lifecycle — connect, setup, full-duplex split, reconnection.
2
3mod message_handler;
4mod reconnect;
5mod session_loop;
6
7use std::sync::Arc;
8
9use tokio::sync::{broadcast, mpsc, watch};
10
11use crate::protocol::types::*;
12use crate::session::{SessionHandle, SessionPhase, SessionState};
13use crate::transport::TransportConfig;
14use crate::transport::codec::{Codec, JsonCodec};
15use crate::transport::ws::{Transport, TungsteniteTransport};
16
17/// Connect to the Gemini Multimodal Live API with the default transport and
18/// return a session handle.
19///
20/// Timeouts, reconnection, a custom transport or codec, or a wire recorder:
21/// [`ConnectBuilder`](crate::transport::ConnectBuilder), of which this is the
22/// zero-option form.
23pub async fn connect(config: SessionConfig) -> Result<SessionHandle, crate::session::SessionError> {
24    connect_with(
25        config,
26        TransportConfig::default(),
27        TungsteniteTransport::new(),
28        JsonCodec,
29    )
30    .await
31}
32
33/// Connect with an explicit transport config, transport, and codec — the
34/// one path every public entry point ends in.
35pub(crate) async fn connect_with<T, C>(
36    config: SessionConfig,
37    transport_config: TransportConfig,
38    transport: T,
39    codec: C,
40) -> Result<SessionHandle, crate::session::SessionError>
41where
42    T: Transport,
43    C: Codec,
44{
45    let (command_tx, command_rx) = mpsc::channel(transport_config.send_queue_depth);
46    let (event_tx, _) = broadcast::channel(transport_config.event_channel_capacity);
47    let (phase_tx, phase_rx) = watch::channel(SessionPhase::Disconnected);
48
49    let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
50
51    let mut handle = SessionHandle::new(command_tx, event_tx.clone(), state.clone(), phase_rx);
52    if let Some(pacing) = config.audio_pacing.clone() {
53        handle = handle.with_audio_pacing(pacing);
54    }
55
56    // Honor a config-installed wire recorder by wrapping the codec. The
57    // boxed indirection keeps `connect_with` generic while letting the
58    // recorder be a runtime decision.
59    let task = if let Some(recorder) = config.wire_recorder.clone() {
60        let codec: Box<dyn Codec> = Box::new(crate::transport::recording::RecordingCodec::new(
61            codec,
62            recorder.recorder(),
63        ));
64        tokio::spawn(async move {
65            session_loop::generic_connection_loop(
66                config,
67                transport_config,
68                state,
69                command_rx,
70                event_tx,
71                transport,
72                codec,
73            )
74            .await;
75        })
76    } else {
77        tokio::spawn(async move {
78            session_loop::generic_connection_loop(
79                config,
80                transport_config,
81                state,
82                command_rx,
83                event_tx,
84                transport,
85                codec,
86            )
87            .await;
88        })
89    };
90    handle.set_task(task);
91
92    Ok(handle)
93}
94
95#[cfg(test)]
96mod tests {
97    use super::message_handler::{MessageAction, handle_server_msg};
98    use super::reconnect::reconnect_delay;
99    use super::*;
100
101    use std::time::Duration;
102
103    use crate::protocol::messages::ServerMessage;
104    use crate::session::{SessionEvent, SessionPhase, SessionState};
105    use crate::transport::codec::JsonCodec;
106    use crate::transport::ws::MockTransport;
107
108    /// TransportConfig that disables reconnection for mock tests.
109    fn no_reconnect_config() -> TransportConfig {
110        TransportConfig {
111            max_reconnect_attempts: 0,
112            connect_timeout_secs: 5,
113            setup_timeout_secs: 5,
114            ..TransportConfig::default()
115        }
116    }
117
118    /// A transport that plays one script per connection and records every
119    /// setup message it is sent.
120    struct Reconnecting {
121        scripts: std::collections::VecDeque<Vec<Vec<u8>>>,
122        current: std::collections::VecDeque<Vec<u8>>,
123        setups: std::sync::Arc<parking_lot::Mutex<Vec<serde_json::Value>>>,
124    }
125
126    #[async_trait::async_trait]
127    impl crate::transport::ws::Transport for Reconnecting {
128        type Error = std::io::Error;
129
130        async fn connect(
131            &mut self,
132            _url: &str,
133            _headers: Vec<(String, String)>,
134        ) -> Result<(), Self::Error> {
135            self.current = self.scripts.pop_front().unwrap_or_default().into();
136            Ok(())
137        }
138
139        async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
140            if let Ok(message) = serde_json::from_slice::<serde_json::Value>(&data)
141                && message.get("setup").is_some()
142            {
143                self.setups.lock().push(message);
144            }
145            Ok(())
146        }
147
148        async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
149            match self.current.pop_front() {
150                Some(frame) => Ok(Some(frame)),
151                // Out of script: stay open and quiet.
152                None => std::future::pending().await,
153            }
154        }
155
156        async fn close(&mut self) -> Result<(), Self::Error> {
157            Ok(())
158        }
159    }
160
161    /// A server that refuses every setup, closing with a status and reason.
162    struct Refusing {
163        connects: std::sync::Arc<std::sync::atomic::AtomicUsize>,
164        reason: &'static str,
165    }
166
167    #[async_trait::async_trait]
168    impl crate::transport::ws::Transport for Refusing {
169        type Error = std::io::Error;
170
171        async fn connect(
172            &mut self,
173            _url: &str,
174            _headers: Vec<(String, String)>,
175        ) -> Result<(), Self::Error> {
176            self.connects
177                .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
178            Ok(())
179        }
180
181        async fn send(&mut self, _data: Vec<u8>) -> Result<(), Self::Error> {
182            Ok(())
183        }
184
185        async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
186            // Give the test time to subscribe before the refusal.
187            tokio::time::sleep(Duration::from_millis(50)).await;
188            Ok(None)
189        }
190
191        async fn close(&mut self) -> Result<(), Self::Error> {
192            Ok(())
193        }
194
195        fn close_reason(&self) -> Option<String> {
196            Some(self.reason.to_string())
197        }
198    }
199
200    #[tokio::test]
201    async fn an_invalid_setup_is_reported_with_the_servers_reason_and_not_retried() {
202        let connects = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
203        let transport = Refusing {
204            connects: connects.clone(),
205            reason: "server closed the connection (1007): The requested combination of response modalities (TEXT) is not supported by the model.",
206        };
207        let transport_config = TransportConfig {
208            max_reconnect_attempts: 3,
209            reconnect_base_delay_ms: 10,
210            reconnect_max_delay_ms: 10,
211            ..no_reconnect_config()
212        };
213        let handle = connect_with(
214            SessionConfig::new("test-key"),
215            transport_config,
216            transport,
217            JsonCodec,
218        )
219        .await
220        .unwrap();
221        let mut events = handle.subscribe();
222        let mut error = None;
223        let mut disconnected = None;
224        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
225        while disconnected.is_none() && tokio::time::Instant::now() < deadline {
226            match tokio::time::timeout(Duration::from_millis(200), events.recv()).await {
227                Ok(Ok(SessionEvent::Error(e))) => error = Some(e.to_string()),
228                Ok(Ok(SessionEvent::Disconnected(reason))) => disconnected = Some(reason),
229                _ => {}
230            }
231        }
232        let error = error.expect("the refusal is reported");
233        assert!(error.contains("response modalities (TEXT)"), "{error}");
234        assert!(disconnected.flatten().unwrap_or_default().contains("1007"));
235        assert_eq!(
236            connects.load(std::sync::atomic::Ordering::SeqCst),
237            1,
238            "no retry"
239        );
240    }
241
242    #[test]
243    fn close_codes_are_read_from_the_reason() {
244        assert_eq!(
245            super::session_loop::close_code("server closed the connection (1007): x"),
246            Some(1007)
247        );
248        assert_eq!(
249            super::session_loop::close_code("server closed the connection (1011)"),
250            Some(1011)
251        );
252        assert_eq!(super::session_loop::close_code("no code"), None);
253    }
254
255    /// A server that accepts every setup and, on the first connection, sends
256    /// `goAway` as soon as it has received a `contextUpdate`.
257    struct GoAwayAfterUpdate {
258        connects: usize,
259        current: std::collections::VecDeque<Vec<u8>>,
260        sent: std::sync::Arc<parking_lot::Mutex<Vec<serde_json::Value>>>,
261        /// Issue a resumption handle on the first connection, so the
262        /// reconnect resumes.
263        issue_handle: bool,
264    }
265
266    #[async_trait::async_trait]
267    impl crate::transport::ws::Transport for GoAwayAfterUpdate {
268        type Error = std::io::Error;
269
270        async fn connect(
271            &mut self,
272            _url: &str,
273            _headers: Vec<(String, String)>,
274        ) -> Result<(), Self::Error> {
275            self.connects += 1;
276            self.current = vec![br#"{"setupComplete":{}}"#.to_vec()].into();
277            if self.issue_handle && self.connects == 1 {
278                self.current.push_back(
279                    br#"{"sessionResumptionUpdate":{"newHandle":"h-1","resumable":true}}"#.to_vec(),
280                );
281            }
282            Ok(())
283        }
284
285        async fn send(&mut self, data: Vec<u8>) -> Result<(), Self::Error> {
286            let message: serde_json::Value = serde_json::from_slice(&data).unwrap();
287            if self.connects == 1 && message.get("contextUpdate").is_some() {
288                self.current
289                    .push_back(br#"{"goAway":{"timeLeft":"0s"}}"#.to_vec());
290            }
291            self.sent.lock().push(message);
292            Ok(())
293        }
294
295        async fn recv(&mut self) -> Result<Option<Vec<u8>>, Self::Error> {
296            match self.current.pop_front() {
297                Some(frame) => Ok(Some(frame)),
298                None => std::future::pending().await,
299            }
300        }
301
302        async fn close(&mut self) -> Result<(), Self::Error> {
303            Ok(())
304        }
305    }
306
307    #[tokio::test]
308    async fn a_reconnect_after_a_context_update_declares_the_updated_preamble() {
309        let sent = std::sync::Arc::new(parking_lot::Mutex::new(Vec::new()));
310        let transport = GoAwayAfterUpdate {
311            connects: 0,
312            current: Default::default(),
313            sent: sent.clone(),
314            issue_handle: false,
315        };
316        let config = SessionConfig::from_vertex("p", "us-central1", "t")
317            .model(ModelId::LIVE_3_8)
318            .system_instruction("Unverified caller.")
319            .add_tool(Tool::functions(vec![FunctionDeclaration {
320                name: "verify_identity".into(),
321                description: "Verify the caller.".into(),
322                parameters: None,
323                behavior: None,
324            }]));
325        let transport_config = TransportConfig {
326            max_reconnect_attempts: 1,
327            reconnect_base_delay_ms: 10,
328            reconnect_max_delay_ms: 10,
329            ..no_reconnect_config()
330        };
331        let handle = connect_with(config, transport_config, transport, JsonCodec)
332            .await
333            .unwrap();
334        handle.wait_for_phase(SessionPhase::Active).await;
335        handle
336            .update_context(
337                ContextUpdate::new()
338                    .system_instruction("Verified caller.")
339                    .tools(vec![Tool::functions(vec![FunctionDeclaration {
340                        name: "get_balance".into(),
341                        description: "Fetch the balance.".into(),
342                        parameters: None,
343                        behavior: None,
344                    }])]),
345            )
346            .await
347            .unwrap();
348
349        let setups = || -> Vec<serde_json::Value> {
350            sent.lock()
351                .iter()
352                .filter(|m| m.get("setup").is_some())
353                .cloned()
354                .collect()
355        };
356        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
357        while setups().len() < 2 && tokio::time::Instant::now() < deadline {
358            tokio::time::sleep(Duration::from_millis(10)).await;
359        }
360        let setups = setups();
361        assert_eq!(setups.len(), 2, "one setup per connection");
362        let declared = |i: usize| {
363            setups[i]
364                .pointer("/setup/tools/0/functionDeclarations/0/name")
365                .cloned()
366        };
367        let instruction = |i: usize| {
368            setups[i]
369                .pointer("/setup/systemInstruction/parts/0/text")
370                .cloned()
371        };
372        assert_eq!(declared(0), Some("verify_identity".into()));
373        assert_eq!(instruction(0), Some("Unverified caller.".into()));
374        assert_eq!(
375            declared(1),
376            Some("get_balance".into()),
377            "the reconnect declares the tools in effect, not the starting ones"
378        );
379        assert_eq!(instruction(1), Some("Verified caller.".into()));
380        let updates = sent
381            .lock()
382            .iter()
383            .filter(|m| m.get("contextUpdate").is_some())
384            .count();
385        assert_eq!(
386            updates, 1,
387            "a fresh session needs no replay: its setup declares the update"
388        );
389    }
390
391    #[tokio::test]
392    async fn a_resume_replays_the_fields_updates_replaced() {
393        let sent = std::sync::Arc::new(parking_lot::Mutex::new(Vec::new()));
394        let transport = GoAwayAfterUpdate {
395            connects: 0,
396            current: Default::default(),
397            sent: sent.clone(),
398            issue_handle: true,
399        };
400        let mut config = SessionConfig::from_vertex("p", "us-central1", "t")
401            .model(ModelId::LIVE_3_8)
402            .system_instruction("Unverified caller.")
403            .add_tool(Tool::functions(vec![FunctionDeclaration {
404                name: "verify_identity".into(),
405                description: "Verify the caller.".into(),
406                parameters: None,
407                behavior: None,
408            }]));
409        config.session_resumption = Some(SessionResumptionConfig {
410            handle: None,
411            transparent: None,
412        });
413        let transport_config = TransportConfig {
414            max_reconnect_attempts: 1,
415            reconnect_base_delay_ms: 10,
416            reconnect_max_delay_ms: 10,
417            ..no_reconnect_config()
418        };
419        let handle = connect_with(config, transport_config, transport, JsonCodec)
420            .await
421            .unwrap();
422        handle.wait_for_phase(SessionPhase::Active).await;
423        handle
424            .update_context(ContextUpdate::new().tools(vec![Tool::functions(vec![
425                FunctionDeclaration {
426                    name: "get_balance".into(),
427                    description: "Fetch the balance.".into(),
428                    parameters: None,
429                    behavior: None,
430                },
431            ])]))
432            .await
433            .unwrap();
434
435        // A resume ignores the setup's tools, and the client cannot tell
436        // whether the update reached the server before the goAway.
437        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
438        // The frames sent after the second (resumed) setup, once there is one.
439        let after_resume = || -> Option<Vec<serde_json::Value>> {
440            let sent = sent.lock();
441            let resume_at = sent
442                .iter()
443                .enumerate()
444                .filter(|(_, m)| m.get("setup").is_some())
445                .nth(1)?
446                .0;
447            assert_eq!(
448                sent[resume_at].pointer("/setup/sessionResumption/handle"),
449                Some(&serde_json::json!("h-1"))
450            );
451            Some(sent[resume_at + 1..].to_vec())
452        };
453        let replay = loop {
454            let found = after_resume()
455                .and_then(|after| after.into_iter().find(|m| m.get("contextUpdate").is_some()));
456            if let Some(replay) = found {
457                break replay;
458            }
459            assert!(
460                tokio::time::Instant::now() < deadline,
461                "no replay after the resume"
462            );
463            tokio::time::sleep(Duration::from_millis(10)).await;
464        };
465        assert_eq!(
466            replay.pointer("/contextUpdate/tools/tools/0/functionDeclarations/0/name"),
467            Some(&serde_json::json!("get_balance")),
468            "the resumed session is brought back to the updated tools"
469        );
470        assert!(
471            replay.pointer("/contextUpdate/systemInstruction").is_none(),
472            "only the fields an update replaced are sent again"
473        );
474    }
475
476    #[tokio::test]
477    async fn a_reconnect_after_go_away_resumes_with_the_latest_handle() {
478        let setups = std::sync::Arc::new(parking_lot::Mutex::new(Vec::new()));
479        let transport = Reconnecting {
480            scripts: vec![
481                vec![
482                    br#"{"setupComplete":{}}"#.to_vec(),
483                    br#"{"sessionResumptionUpdate":{"newHandle":"h-1","resumable":true}}"#.to_vec(),
484                    br#"{"goAway":{"timeLeft":"0s"}}"#.to_vec(),
485                ],
486                vec![br#"{"setupComplete":{}}"#.to_vec()],
487            ]
488            .into(),
489            current: Default::default(),
490            setups: setups.clone(),
491        };
492        let mut config = SessionConfig::new("test-key")
493            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
494        config.session_resumption = Some(SessionResumptionConfig {
495            handle: None,
496            transparent: None,
497        });
498        let transport_config = TransportConfig {
499            max_reconnect_attempts: 1,
500            reconnect_base_delay_ms: 10,
501            reconnect_max_delay_ms: 10,
502            ..no_reconnect_config()
503        };
504        let _handle = connect_with(config, transport_config, transport, JsonCodec)
505            .await
506            .unwrap();
507
508        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
509        while setups.lock().len() < 2 && tokio::time::Instant::now() < deadline {
510            tokio::time::sleep(Duration::from_millis(10)).await;
511        }
512        let setups = setups.lock();
513        assert_eq!(setups.len(), 2, "one setup per connection");
514        let resumption = |i: usize| setups[i].pointer("/setup/sessionResumption").cloned();
515        assert_eq!(
516            resumption(0),
517            Some(serde_json::json!({})),
518            "first connect: no handle yet"
519        );
520        assert_eq!(
521            resumption(1),
522            Some(serde_json::json!({ "handle": "h-1" })),
523            "the reconnect presents the handle the server issued"
524        );
525    }
526
527    #[tokio::test]
528    async fn connect_with_mock_transport() {
529        let mut transport = MockTransport::new();
530        // Script setupComplete response
531        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
532        // Script a text response then turn complete
533        transport.script_recv(
534            br#"{"serverContent":{"modelTurn":{"parts":[{"text":"Hello!"}]},"turnComplete":true}}"#
535                .to_vec(),
536        );
537
538        let config = SessionConfig::new("test-key")
539            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
540
541        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
542            .await
543            .unwrap();
544
545        // Should reach Active phase after setup completes
546        handle.wait_for_phase(SessionPhase::Active).await;
547        assert_eq!(handle.phase(), SessionPhase::Active);
548    }
549
550    #[tokio::test]
551    async fn connect_with_mock_receives_text_events() {
552        let mut transport = MockTransport::new();
553        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
554        transport.script_recv(
555            br#"{"serverContent":{"modelTurn":{"parts":[{"text":"Hello from mock!"}]},"turnComplete":true}}"#
556                .to_vec(),
557        );
558
559        let config = SessionConfig::new("test-key")
560            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
561        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
562            .await
563            .unwrap();
564
565        let mut events = handle.subscribe();
566
567        // Wait for the session to become active
568        handle.wait_for_phase(SessionPhase::Active).await;
569
570        // Collect events until TurnComplete
571        let mut got_text_delta = false;
572        let mut got_text_complete = false;
573        let mut got_turn_complete = false;
574
575        for _ in 0..20 {
576            match tokio::time::timeout(Duration::from_millis(100), events.recv()).await {
577                Ok(Ok(SessionEvent::TextDelta(t))) => {
578                    assert_eq!(t, "Hello from mock!");
579                    got_text_delta = true;
580                }
581                Ok(Ok(SessionEvent::TextComplete(t))) => {
582                    assert_eq!(t, "Hello from mock!");
583                    got_text_complete = true;
584                }
585                Ok(Ok(SessionEvent::TurnComplete)) => {
586                    got_turn_complete = true;
587                    break;
588                }
589                Ok(Ok(_)) => continue,
590                Ok(Err(_)) => break,
591                Err(_) => break,
592            }
593        }
594
595        assert!(got_text_delta, "should have received TextDelta");
596        assert!(got_text_complete, "should have received TextComplete");
597        assert!(got_turn_complete, "should have received TurnComplete");
598    }
599
600    #[tokio::test]
601    async fn connect_with_mock_decodes_audio_and_standalone_usage() {
602        // 40 ms of 24 kHz PCM16 with every byte value, so a decoder that
603        // mangles any of the 64 base64 symbols shows up.
604        let pcm: Vec<u8> = (0..1920u32).map(|i| (i * 7 % 256) as u8).collect();
605        let b64 = {
606            use base64::Engine as _;
607            base64::engine::general_purpose::STANDARD.encode(&pcm)
608        };
609        let mut transport = MockTransport::new();
610        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
611        transport.script_recv(
612            format!(
613                r#"{{"serverContent":{{"modelTurn":{{"parts":[{{"inlineData":{{"mimeType":"audio/pcm;rate=24000","data":"{b64}"}}}}]}}}}}}"#
614            )
615            .into_bytes(),
616        );
617        // Usage with no content alongside it used to parse as an unknown
618        // message and be dropped.
619        transport.script_recv(
620            br#"{"usageMetadata":{"promptTokenCount":12,"totalTokenCount":30}}"#.to_vec(),
621        );
622        transport.script_recv(br#"{"serverContent":{"turnComplete":true}}"#.to_vec());
623
624        let config = SessionConfig::new("test-key")
625            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
626        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
627            .await
628            .unwrap();
629        let mut events = handle.subscribe();
630        handle.wait_for_phase(SessionPhase::Active).await;
631
632        let mut audio = None;
633        let mut usage = None;
634        for _ in 0..20 {
635            match tokio::time::timeout(Duration::from_millis(100), events.recv()).await {
636                Ok(Ok(SessionEvent::AudioData(data))) => audio = Some(data),
637                Ok(Ok(SessionEvent::Usage(u))) => usage = Some(u),
638                Ok(Ok(SessionEvent::TurnComplete)) => break,
639                Ok(Ok(_)) => continue,
640                Ok(Err(_)) | Err(_) => break,
641            }
642        }
643
644        assert_eq!(audio.as_deref(), Some(pcm.as_slice()));
645        let usage = usage.expect("standalone usage should reach the session");
646        assert_eq!(usage.prompt_token_count, Some(12));
647        assert_eq!(usage.total_token_count, Some(30));
648    }
649
650    #[tokio::test]
651    async fn connect_with_mock_tool_call() {
652        let mut transport = MockTransport::new();
653        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
654        transport.script_recv(
655            br#"{"toolCall":{"functionCalls":[{"name":"get_weather","args":{"city":"London"},"id":"call-1"}]}}"#
656                .to_vec(),
657        );
658
659        let config = SessionConfig::new("test-key")
660            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
661        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
662            .await
663            .unwrap();
664
665        let mut events = handle.subscribe();
666        handle.wait_for_phase(SessionPhase::Active).await;
667
668        // Look for the ToolCall event
669        let mut got_tool_call = false;
670        for _ in 0..20 {
671            match tokio::time::timeout(Duration::from_millis(100), events.recv()).await {
672                Ok(Ok(SessionEvent::ToolCall(calls))) => {
673                    assert_eq!(calls.len(), 1);
674                    assert_eq!(calls[0].name, "get_weather");
675                    got_tool_call = true;
676                    break;
677                }
678                Ok(Ok(_)) => continue,
679                Ok(Err(_)) => break,
680                Err(_) => break,
681            }
682        }
683
684        assert!(got_tool_call, "should have received ToolCall event");
685    }
686
687    #[tokio::test]
688    async fn connect_with_mock_graceful_disconnect() {
689        let mut transport = MockTransport::new();
690        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
691        // Keep the connection alive with a message that arrives before disconnect
692        transport.script_recv(
693            br#"{"serverContent":{"modelTurn":{"parts":[{"text":"hi"}]},"turnComplete":true}}"#
694                .to_vec(),
695        );
696
697        let config = SessionConfig::new("test-key")
698            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
699        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
700            .await
701            .unwrap();
702
703        handle.wait_for_phase(SessionPhase::Active).await;
704        // Small delay to let the background task process
705        tokio::time::sleep(Duration::from_millis(50)).await;
706
707        // Disconnect gracefully
708        handle.disconnect().await.unwrap();
709
710        // Wait for disconnected phase
711        handle.wait_for_phase(SessionPhase::Disconnected).await;
712        assert_eq!(handle.phase(), SessionPhase::Disconnected);
713    }
714
715    #[test]
716    fn handle_server_msg_preserves_interruption() {
717        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Active);
718        let (event_tx, mut event_rx) = broadcast::channel(16);
719        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
720
721        let json = r#"{"serverContent":{"interrupted":true}}"#;
722        let msg = ServerMessage::parse(json).unwrap();
723        let action = handle_server_msg(msg, &state, &event_tx);
724
725        assert!(matches!(action, MessageAction::Continue));
726        // Should have emitted Interrupted event
727        let mut found_interrupted = false;
728        while let Ok(evt) = event_rx.try_recv() {
729            if matches!(evt, SessionEvent::Interrupted) {
730                found_interrupted = true;
731            }
732        }
733        assert!(found_interrupted, "should emit Interrupted event");
734    }
735
736    #[test]
737    fn handle_server_msg_routes_avatar_video_away_from_audio() {
738        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Active);
739        let (event_tx, mut event_rx) = broadcast::channel(16);
740        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
741
742        // "AAEC" = [0, 1, 2]; "AwQF" = [3, 4, 5].
743        let json = r#"{"serverContent":{"modelTurn":{"parts":[
744            {"inlineData":{"mimeType":"audio/pcm;rate=24000","data":"AAEC"}},
745            {"inlineData":{"mimeType":"video/mp4","data":"AwQF"}}
746        ]}}}"#;
747        let msg = ServerMessage::parse(json).unwrap();
748        handle_server_msg(msg, &state, &event_tx);
749
750        let mut audio = Vec::new();
751        let mut media = Vec::new();
752        while let Ok(evt) = event_rx.try_recv() {
753            match evt {
754                SessionEvent::AudioData(bytes) => audio.push(bytes.to_vec()),
755                SessionEvent::Media(m) => media.push(m),
756                _ => {}
757            }
758        }
759        assert_eq!(audio, [vec![0u8, 1, 2]], "only the audio part is audio");
760        assert_eq!(media.len(), 1);
761        assert!(media[0].is_video());
762        assert_eq!(media[0].mime_type, "video/mp4");
763        assert_eq!(media[0].data.as_ref(), [3u8, 4, 5]);
764    }
765
766    #[test]
767    fn a_text_session_on_a_speech_only_model_reads_the_transcript_as_text() {
768        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Active);
769        let (event_tx, mut event_rx) = broadcast::channel(32);
770        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
771        state.set_text_from_transcription(true);
772
773        for json in [
774            r#"{"serverContent":{"modelTurn":{"parts":[{"inlineData":{"mimeType":"audio/pcm;rate=24000","data":"AAEC"}}]}}}"#,
775            r#"{"serverContent":{"outputTranscription":{"text":"A table "}}}"#,
776            r#"{"serverContent":{"outputTranscription":{"text":"for two."}}}"#,
777            r#"{"serverContent":{"turnComplete":true}}"#,
778        ] {
779            handle_server_msg(ServerMessage::parse(json).unwrap(), &state, &event_tx);
780        }
781        let mut deltas = Vec::new();
782        let mut complete = None;
783        while let Ok(evt) = event_rx.try_recv() {
784            match evt {
785                SessionEvent::TextDelta(t) => deltas.push(t),
786                SessionEvent::TextComplete(t) => complete = Some(t),
787                SessionEvent::AudioData(_) => panic!("a text session plays no audio"),
788                SessionEvent::OutputTranscription(_) => panic!("the transcript is the text"),
789                _ => {}
790            }
791        }
792        assert_eq!(deltas, ["A table ", "for two."]);
793        assert_eq!(complete.as_deref(), Some("A table for two."));
794    }
795
796    #[test]
797    fn handle_server_msg_go_away() {
798        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Active);
799        let (event_tx, _event_rx) = broadcast::channel(16);
800        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
801
802        let json = r#"{"goAway":{"timeLeft":"30s"}}"#;
803        let msg = ServerMessage::parse(json).unwrap();
804        let action = handle_server_msg(msg, &state, &event_tx);
805
806        assert!(matches!(action, MessageAction::GoAway(Some(_))));
807    }
808
809    #[test]
810    fn handle_server_msg_unknown_is_continue() {
811        let (phase_tx, _phase_rx) = watch::channel(SessionPhase::Active);
812        let (event_tx, _event_rx) = broadcast::channel(16);
813        let state = Arc::new(SessionState::with_events(phase_tx, event_tx.clone()));
814
815        let json = r#"{"unknownField":{"data":"test"}}"#;
816        let msg = ServerMessage::parse(json).unwrap();
817        let action = handle_server_msg(msg, &state, &event_tx);
818
819        assert!(matches!(action, MessageAction::Continue));
820    }
821
822    #[tokio::test]
823    async fn session_handle_join_after_disconnect() {
824        let mut transport = MockTransport::new();
825        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
826
827        let config = SessionConfig::new("test-key")
828            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
829        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
830            .await
831            .unwrap();
832
833        handle.wait_for_phase(SessionPhase::Active).await;
834
835        // Disconnect to end the connection loop task
836        handle.disconnect().await.unwrap();
837        handle.wait_for_phase(SessionPhase::Disconnected).await;
838
839        // join() should return Ok after the task completes
840        let result = handle.join().await;
841        assert!(result.is_ok(), "join() should succeed after disconnect");
842    }
843
844    #[tokio::test]
845    async fn session_handle_join_after_command_channel_closed() {
846        let mut transport = MockTransport::new();
847        transport.script_recv(br#"{"setupComplete":{}}"#.to_vec());
848
849        let config = SessionConfig::new("test-key")
850            .model(ModelId::from_static("models/gemini-2.0-flash-live-001"));
851        let handle = connect_with(config, no_reconnect_config(), transport, JsonCodec)
852            .await
853            .unwrap();
854
855        handle.wait_for_phase(SessionPhase::Active).await;
856
857        // Drop all senders to close the command channel, which triggers disconnect
858        // We need to get the handle before dropping the original
859        let join_handle = handle.clone();
860
861        // Drop command_tx by dropping the handle — but we cloned it first.
862        // Instead, disconnect and then join.
863        handle.disconnect().await.unwrap();
864
865        let result = join_handle.join().await;
866        assert!(result.is_ok(), "join() should succeed after channel close");
867    }
868
869    #[test]
870    fn reconnect_delay_exponential_backoff() {
871        let config = TransportConfig::default();
872        let d1 = reconnect_delay(1, &config);
873        let d2 = reconnect_delay(2, &config);
874        let d3 = reconnect_delay(3, &config);
875        // Each step should roughly double (plus jitter)
876        assert!(d2 > d1);
877        assert!(d3 > d2);
878        // Should not exceed max
879        let d_large = reconnect_delay(100, &config);
880        let max_with_jitter = Duration::from_millis(
881            config.reconnect_max_delay_ms as u64 + config.reconnect_max_delay_ms as u64 / 4,
882        );
883        assert!(d_large <= max_with_jitter);
884    }
885}
gemini_genai_rs/transport/connection/mod.rs

gemini_genai_rs/transport/connection/
mod.rs
