gemini_genai_rs/transport/auth/
mod.rs

1//! Authentication providers for Gemini API connections.
2//!
3//! This module defines the [`AuthProvider`] trait and built-in implementations
4//! for Google AI (API key and OAuth2 token) and Vertex AI (Bearer token).
5//!
6//! The [`ServiceEndpoint`] enum allows constructing URLs for both WebSocket (Live)
7//! and REST API endpoints from the same auth provider.
8
9pub mod google_ai;
10#[cfg(feature = "http")]
11pub mod google_token;
12pub(crate) mod url_builders;
13pub mod vertex;
14
15pub use google_ai::*;
16#[cfg(feature = "http")]
17pub use google_token::GoogleAccessToken;
18pub use vertex::*;
19
20use async_trait::async_trait;
21
22use crate::protocol::types::ModelId;
23use crate::session::AuthError;
24
25/// Identifies which Gemini API service to connect to.
26///
27/// Used by `AuthProvider::rest_url` to construct the correct REST endpoint URL.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
29pub enum ServiceEndpoint {
30    /// WebSocket Live/Bidi streaming endpoint.
31    LiveWs,
32    /// POST /models/{model}:generateContent
33    GenerateContent,
34    /// POST /models/{model}:streamGenerateContent
35    StreamGenerateContent,
36    /// POST /models/{model}:embedContent
37    EmbedContent,
38    /// POST /models/{model}:countTokens
39    CountTokens,
40    /// POST /models/{model}:computeTokens
41    ComputeTokens,
42    /// GET /models
43    ListModels,
44    /// GET /models/{model}
45    GetModel,
46    /// Files CRUD (upload, get, list, delete)
47    Files,
48    /// Cached content CRUD
49    CachedContents,
50    /// Tuning jobs CRUD
51    TuningJobs,
52    /// Batch jobs CRUD
53    BatchJobs,
54}
55
56impl ServiceEndpoint {
57    /// REST method suffix appended to the model path (e.g., `:generateContent`).
58    /// Returns `None` for endpoints that don't use a model suffix.
59    pub fn model_method(&self) -> Option<&'static str> {
60        match self {
61            Self::GenerateContent => Some("generateContent"),
62            Self::StreamGenerateContent => Some("streamGenerateContent"),
63            Self::EmbedContent => Some("embedContent"),
64            Self::CountTokens => Some("countTokens"),
65            Self::ComputeTokens => Some("computeTokens"),
66            _ => None,
67        }
68    }
69
70    /// Whether this endpoint requires a model ID in the path.
71    pub fn requires_model(&self) -> bool {
72        matches!(
73            self,
74            Self::GenerateContent
75                | Self::StreamGenerateContent
76                | Self::EmbedContent
77                | Self::CountTokens
78                | Self::ComputeTokens
79                | Self::GetModel
80        )
81    }
82}
83
84/// Provides authentication credentials and URL construction for Gemini API connections.
85#[async_trait]
86pub trait AuthProvider: Send + Sync + 'static {
87    /// Build the WebSocket URL for the given model.
88    fn ws_url(&self, model: &ModelId) -> String;
89
90    /// HTTP headers for the WebSocket upgrade request (e.g., Bearer token).
91    async fn auth_headers(&self) -> Result<Vec<(String, String)>, AuthError>;
92
93    /// Query parameters to append to the URL (e.g., API key).
94    fn query_params(&self) -> Vec<(String, String)> {
95        vec![]
96    }
97
98    /// Called on auth failure to allow token refresh. Default: no-op.
99    async fn refresh(&self) -> Result<(), AuthError> {
100        Ok(())
101    }
102}
103
104/// Auth providers that additionally support REST endpoint URL construction.
105///
106/// Split from [`AuthProvider`] so that Live-only providers are not forced to
107/// implement REST URL building, and so the REST [`Client`](crate::client::Client)
108/// can require it at the type level — replacing the previous runtime
109/// `unimplemented!()` default with a compile-time guarantee.
110pub trait RestAuth: AuthProvider {
111    /// Build a REST API URL for the given service endpoint and model.
112    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&ModelId>) -> String;
113}
114
115// ---------------------------------------------------------------------------
116// Tests
117// ---------------------------------------------------------------------------
118
119#[cfg(test)]
120mod tests {
121    use super::*;
122    use crate::protocol::types::ModelId;
123
124    #[test]
125    fn google_ai_auth_url() {
126        let auth = GoogleAIAuth::new("test-key-123");
127        let url = auth.ws_url(&ModelId::FLASH_LATEST);
128        assert!(url.contains("generativelanguage.googleapis.com"));
129        assert!(url.contains("v1beta"));
130        assert!(url.contains("key=test-key-123"));
131    }
132
133    #[test]
134    fn google_ai_auth_query_params() {
135        let auth = GoogleAIAuth::new("my-api-key");
136        let params = auth.query_params();
137        assert_eq!(params.len(), 1);
138        assert_eq!(params[0].0, "key");
139        assert_eq!(params[0].1, "my-api-key");
140    }
141
142    #[tokio::test]
143    async fn google_ai_rest_key_travels_in_a_header_not_the_url() {
144        let auth = GoogleAIAuth::new("test-key");
145        let headers = auth.auth_headers().await.unwrap();
146        assert_eq!(
147            headers,
148            vec![("x-goog-api-key".to_string(), "test-key".to_string())]
149        );
150        let url = auth.rest_url(
151            ServiceEndpoint::GenerateContent,
152            Some(&ModelId::FLASH_LATEST),
153        );
154        assert!(!url.contains("test-key"), "{url}");
155    }
156
157    #[test]
158    fn google_ai_token_auth_url() {
159        let auth = GoogleAITokenAuth::new("oauth2-token-abc");
160        let url = auth.ws_url(&ModelId::FLASH_LATEST);
161        assert!(url.contains("generativelanguage.googleapis.com"));
162        assert!(url.contains("access_token=oauth2-token-abc"));
163        assert!(url.contains("v1alpha"));
164    }
165
166    #[test]
167    fn vertex_ai_auth_url_regional() {
168        let auth = VertexAIAuth::new("my-project", "us-central1", "token");
169        let url = auth.ws_url(&ModelId::FLASH_LATEST);
170        assert!(url.contains("us-central1-aiplatform.googleapis.com"));
171        assert!(url.contains("v1beta1"));
172        assert!(url.contains("x-goog-project-id=my-project"));
173    }
174
175    #[test]
176    fn vertex_ai_auth_url_global() {
177        let auth = VertexAIAuth::new("my-project", "global", "token");
178        let url = auth.ws_url(&ModelId::FLASH_LATEST);
179        // Global uses aiplatform.googleapis.com without location prefix.
180        assert!(url.starts_with("wss://aiplatform.googleapis.com/"));
181        assert!(!url.contains("global-aiplatform"));
182    }
183
184    #[tokio::test]
185    async fn vertex_ai_auth_headers() {
186        let auth = VertexAIAuth::new("proj", "us-central1", "my-bearer-token");
187        let headers = auth.auth_headers().await.unwrap();
188        assert_eq!(headers.len(), 1);
189        assert_eq!(headers[0].0, "Authorization");
190        assert_eq!(headers[0].1, "Bearer my-bearer-token");
191    }
192
193    #[test]
194    fn vertex_ai_auth_url_contains_model() {
195        let auth = VertexAIAuth::new("proj", "us-central1", "tok");
196        let url = auth.ws_url(&ModelId::from_static("models/gemini-2.0-flash-live-001"));
197        assert!(url.contains("model=gemini-2.0-flash-live-001"));
198    }
199
200    #[test]
201    fn auth_provider_is_object_safe() {
202        fn _assert(_: &dyn AuthProvider) {}
203    }
204
205    #[tokio::test]
206    async fn default_refresh_is_noop() {
207        let auth = GoogleAIAuth::new("key");
208        // Should succeed without error.
209        auth.refresh().await.unwrap();
210    }
211
212    #[tokio::test]
213    async fn default_query_params_empty_for_vertex() {
214        let auth = VertexAIAuth::new("proj", "loc", "tok");
215        let params = auth.query_params();
216        assert!(params.is_empty());
217    }
218
219    // -----------------------------------------------------------------------
220    // REST URL tests
221    // -----------------------------------------------------------------------
222
223    #[test]
224    fn google_ai_rest_url_generate_content() {
225        let auth = GoogleAIAuth::new("test-key");
226        let model = ModelId::from_static("models/gemini-2.0-flash-live-001");
227        let url = auth.rest_url(ServiceEndpoint::GenerateContent, Some(&model));
228        assert!(url.starts_with("https://generativelanguage.googleapis.com/v1beta/"));
229        assert!(url.contains(":generateContent"));
230        assert!(
231            !url.contains("test-key"),
232            "the key rides in a header: {url}"
233        );
234    }
235
236    #[test]
237    fn google_ai_rest_url_list_models() {
238        let auth = GoogleAIAuth::new("key123");
239        let url = auth.rest_url(ServiceEndpoint::ListModels, None);
240        assert!(url.ends_with("/models"), "{url}");
241    }
242
243    #[test]
244    fn google_ai_rest_url_files() {
245        let auth = GoogleAIAuth::new("key");
246        let url = auth.rest_url(ServiceEndpoint::Files, None);
247        assert!(url.ends_with("/files"), "{url}");
248    }
249
250    #[test]
251    fn google_ai_token_rest_url_no_key_in_url() {
252        let auth = GoogleAITokenAuth::new("oauth-token");
253        let url = auth.rest_url(ServiceEndpoint::CountTokens, Some(&ModelId::FLASH_LATEST));
254        assert!(url.contains(":countTokens"));
255        assert!(!url.contains("key="));
256        assert!(!url.contains("access_token="));
257    }
258
259    #[test]
260    fn vertex_rest_url_generate_content() {
261        let auth = VertexAIAuth::new("my-project", "us-central1", "token");
262        let model = ModelId::from_static("models/gemini-2.0-flash-live-001");
263        let url = auth.rest_url(ServiceEndpoint::GenerateContent, Some(&model));
264        assert!(url.starts_with("https://us-central1-aiplatform.googleapis.com/v1beta1/"));
265        assert!(url.contains("projects/my-project/locations/us-central1"));
266        assert!(url.contains(":generateContent"));
267    }
268
269    #[test]
270    fn vertex_rest_url_list_models() {
271        let auth = VertexAIAuth::new("proj", "us-east1", "tok");
272        let url = auth.rest_url(ServiceEndpoint::ListModels, None);
273        assert!(url.contains("publishers/google/models"));
274    }
275
276    #[test]
277    fn vertex_rest_url_global() {
278        let auth = VertexAIAuth::new("proj", "global", "tok");
279        let model = ModelId::FLASH_LATEST;
280        let url = auth.rest_url(ServiceEndpoint::EmbedContent, Some(&model));
281        assert!(url.starts_with("https://aiplatform.googleapis.com/"));
282        assert!(!url.contains("global-aiplatform"));
283        assert!(url.contains(":embedContent"));
284    }
285
286    #[test]
287    fn service_endpoint_model_method() {
288        assert_eq!(
289            ServiceEndpoint::GenerateContent.model_method(),
290            Some("generateContent")
291        );
292        assert_eq!(
293            ServiceEndpoint::StreamGenerateContent.model_method(),
294            Some("streamGenerateContent")
295        );
296        assert_eq!(ServiceEndpoint::ListModels.model_method(), None);
297        assert_eq!(ServiceEndpoint::Files.model_method(), None);
298    }
299
300    #[tokio::test]
301    async fn vertex_ai_refreshable_token() {
302        use std::sync::atomic::{AtomicU32, Ordering};
303        let counter = std::sync::Arc::new(AtomicU32::new(0));
304        let c = counter.clone();
305        let auth = VertexAIAuth::with_token_refresher("proj", "us-central1", move || {
306            c.fetch_add(1, Ordering::SeqCst);
307            format!("token-{}", c.load(Ordering::SeqCst))
308        });
309        let h1 = auth.auth_headers().await.unwrap();
310        assert!(h1[0].1.starts_with("Bearer token-"));
311        let h2 = auth.auth_headers().await.unwrap();
312        assert!(h2[0].1.starts_with("Bearer token-"));
313        // Refresher called twice (once per auth_headers)
314        assert_eq!(counter.load(Ordering::SeqCst), 2);
315    }
316
317    #[test]
318    fn service_endpoint_requires_model() {
319        assert!(ServiceEndpoint::GenerateContent.requires_model());
320        assert!(ServiceEndpoint::CountTokens.requires_model());
321        assert!(!ServiceEndpoint::ListModels.requires_model());
322        assert!(!ServiceEndpoint::Files.requires_model());
323    }
324}
gemini_genai_rs/transport/auth/mod.rs

gemini_genai_rs/transport/auth/
mod.rs
