gemini_genai_rs/transport/
builder.rs

1//! ConnectBuilder — ergonomic builder for advanced transport/codec configuration.
2
3use crate::protocol::types::SessionConfig;
4use crate::session::{SessionError, SessionHandle};
5use crate::transport::codec::{Codec, JsonCodec};
6use crate::transport::connection::connect_with;
7use crate::transport::ws::{Transport, TungsteniteTransport};
8use crate::transport::TransportConfig;
9
10/// Builder for advanced connection configuration.
11///
12/// Allows customizing the transport and codec used for the connection.
13/// Defaults to TungsteniteTransport + JsonCodec.
14///
15/// # Example
16/// ```rust,no_run
17/// use gemini_genai_rs::prelude::*;
18///
19/// # async fn example() {
20/// let config = SessionConfig::new("key");
21/// let handle = ConnectBuilder::new(config)
22///     .transport_config(TransportConfig { connect_timeout_secs: 30, ..Default::default() })
23///     .build()
24///     .await
25///     .unwrap();
26/// # }
27/// ```
28pub struct ConnectBuilder<T = TungsteniteTransport, C = JsonCodec> {
29    config: SessionConfig,
30    transport_config: TransportConfig,
31    transport: T,
32    codec: C,
33}
34
35impl ConnectBuilder {
36    /// Create a new builder with default transport and codec.
37    pub fn new(config: SessionConfig) -> Self {
38        Self {
39            config,
40            transport_config: TransportConfig::default(),
41            transport: TungsteniteTransport::new(),
42            codec: JsonCodec,
43        }
44    }
45}
46
47impl<T: Transport, C: Codec> ConnectBuilder<T, C> {
48    /// Set the transport configuration.
49    pub fn transport_config(mut self, tc: TransportConfig) -> Self {
50        self.transport_config = tc;
51        self
52    }
53
54    /// Use a custom transport implementation.
55    pub fn transport<T2: Transport>(self, transport: T2) -> ConnectBuilder<T2, C> {
56        ConnectBuilder {
57            config: self.config,
58            transport_config: self.transport_config,
59            transport,
60            codec: self.codec,
61        }
62    }
63
64    /// Record every wire byte (both directions) to the given recorder.
65    ///
66    /// Equivalent to [`SessionConfig::record_wire`]: the connection wraps the
67    /// codec in a [`RecordingCodec`](crate::transport::recording::RecordingCodec).
68    pub fn record_wire(
69        mut self,
70        recorder: std::sync::Arc<dyn crate::transport::recording::WireRecorder>,
71    ) -> Self {
72        self.config = self.config.record_wire(recorder);
73        self
74    }
75
76    /// Use a custom codec implementation.
77    pub fn codec<C2: Codec>(self, codec: C2) -> ConnectBuilder<T, C2> {
78        ConnectBuilder {
79            config: self.config,
80            transport_config: self.transport_config,
81            transport: self.transport,
82            codec,
83        }
84    }
85
86    /// Build the connection and return a SessionHandle.
87    pub async fn build(self) -> Result<SessionHandle, SessionError> {
88        connect_with(
89            self.config,
90            self.transport_config,
91            self.transport,
92            self.codec,
93        )
94        .await
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::protocol::types::*;
102    use crate::transport::ws::MockTransport;
103
104    #[test]
105    fn builder_compiles_with_defaults() {
106        let config = SessionConfig::new("key").model(GeminiModel::Gemini2_0FlashLive);
107        let _builder = ConnectBuilder::new(config);
108    }
109
110    #[test]
111    fn builder_with_custom_transport_config() {
112        let config = SessionConfig::new("key");
113        let _builder = ConnectBuilder::new(config).transport_config(TransportConfig {
114            connect_timeout_secs: 30,
115            ..Default::default()
116        });
117    }
118
119    #[test]
120    fn builder_with_mock_transport() {
121        let config = SessionConfig::new("key");
122        let mock = MockTransport::new();
123        let _builder = ConnectBuilder::new(config).transport(mock);
124    }
125
126    #[test]
127    fn builder_with_custom_codec() {
128        let config = SessionConfig::new("key");
129        let _builder = ConnectBuilder::new(config).codec(JsonCodec);
130    }
131
132    #[tokio::test]
133    async fn builder_with_mock_builds() {
134        let mut mock = MockTransport::new();
135        mock.script_recv(br#"{"setupComplete":{}}"#.to_vec());
136
137        let config = SessionConfig::new("key").model(GeminiModel::Gemini2_0FlashLive);
138        let handle = ConnectBuilder::new(config)
139            .transport(mock)
140            .build()
141            .await
142            .unwrap();
143
144        handle
145            .wait_for_phase(crate::session::SessionPhase::Active)
146            .await;
147        assert_eq!(handle.phase(), crate::session::SessionPhase::Active);
148    }
149}
gemini_genai_rs/transport/builder.rs

gemini_genai_rs/transport/
builder.rs
