gemini_genai_rs/transport/auth/
google_ai.rs

1//! Google AI authentication providers (API key and OAuth2 token).
2
3use async_trait::async_trait;
4
5use crate::protocol::types::GeminiModel;
6use crate::session::AuthError;
7
8use super::url_builders::{build_google_ai_rest_url, build_google_ai_rest_url_no_key};
9use super::{AuthProvider, RestAuth, ServiceEndpoint};
10
11// ---------------------------------------------------------------------------
12// Google AI — API key authentication
13// ---------------------------------------------------------------------------
14
15/// Google AI API key authentication.
16///
17/// The API key is included as a query parameter in the WebSocket URL.
18pub struct GoogleAIAuth {
19    api_key: String,
20}
21
22impl GoogleAIAuth {
23    /// Create a new Google AI auth provider with the given API key.
24    pub fn new(api_key: impl Into<String>) -> Self {
25        Self {
26            api_key: api_key.into(),
27        }
28    }
29}
30
31#[async_trait]
32impl AuthProvider for GoogleAIAuth {
33    fn ws_url(&self, _model: &GeminiModel) -> String {
34        format!(
35            "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={}",
36            self.api_key
37        )
38    }
39
40    async fn auth_headers(&self) -> Result<Vec<(String, String)>, AuthError> {
41        Ok(vec![]) // API key is in the URL
42    }
43
44    fn query_params(&self) -> Vec<(String, String)> {
45        vec![("key".to_string(), self.api_key.clone())]
46    }
47}
48
49impl RestAuth for GoogleAIAuth {
50    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&GeminiModel>) -> String {
51        let base = "https://generativelanguage.googleapis.com/v1beta";
52        build_google_ai_rest_url(base, endpoint, model, &self.api_key)
53    }
54}
55
56// ---------------------------------------------------------------------------
57// Google AI — OAuth2 access token authentication
58// ---------------------------------------------------------------------------
59
60/// Google AI OAuth2 access token authentication.
61///
62/// The access token is included directly in the WebSocket URL.
63pub struct GoogleAITokenAuth {
64    access_token: String,
65}
66
67impl GoogleAITokenAuth {
68    /// Create a new Google AI token auth provider with the given access token.
69    pub fn new(access_token: impl Into<String>) -> Self {
70        Self {
71            access_token: access_token.into(),
72        }
73    }
74}
75
76#[async_trait]
77impl AuthProvider for GoogleAITokenAuth {
78    fn ws_url(&self, _model: &GeminiModel) -> String {
79        format!(
80            "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContentConstrained?access_token={}",
81            self.access_token
82        )
83    }
84
85    async fn auth_headers(&self) -> Result<Vec<(String, String)>, AuthError> {
86        Ok(vec![(
87            "Authorization".to_string(),
88            format!("Bearer {}", self.access_token),
89        )])
90    }
91}
92
93impl RestAuth for GoogleAITokenAuth {
94    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&GeminiModel>) -> String {
95        let base = "https://generativelanguage.googleapis.com/v1beta";
96        // Token auth uses Bearer header, not query param — build URL without key
97        build_google_ai_rest_url_no_key(base, endpoint, model)
98    }
99}