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::ModelId;
6use crate::session::AuthError;
7
8use super::url_builders::build_google_ai_rest_url;
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: &ModelId) -> String {
34        format!(
35            "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={}",
36            self.api_key
37        )
38    }
39
40    /// REST requests carry the key in the `x-goog-api-key` header rather
41    /// than the query string, so it never lands in access logs, proxies, or
42    /// error messages that echo the URL. (The Live WebSocket upgrade still
43    /// passes it as `?key=`, which is what that endpoint accepts.)
44    async fn auth_headers(&self) -> Result<Vec<(String, String)>, AuthError> {
45        Ok(vec![("x-goog-api-key".to_string(), self.api_key.clone())])
46    }
47
48    fn query_params(&self) -> Vec<(String, String)> {
49        vec![("key".to_string(), self.api_key.clone())]
50    }
51}
52
53impl RestAuth for GoogleAIAuth {
54    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&ModelId>) -> String {
55        let base = "https://generativelanguage.googleapis.com/v1beta";
56        build_google_ai_rest_url(base, endpoint, model)
57    }
58}
59
60// ---------------------------------------------------------------------------
61// Google AI — OAuth2 access token authentication
62// ---------------------------------------------------------------------------
63
64/// Google AI OAuth2 access token authentication.
65///
66/// The access token is included directly in the WebSocket URL.
67pub struct GoogleAITokenAuth {
68    access_token: String,
69}
70
71impl GoogleAITokenAuth {
72    /// Create a new Google AI token auth provider with the given access token.
73    pub fn new(access_token: impl Into<String>) -> Self {
74        Self {
75            access_token: access_token.into(),
76        }
77    }
78}
79
80#[async_trait]
81impl AuthProvider for GoogleAITokenAuth {
82    fn ws_url(&self, _model: &ModelId) -> String {
83        format!(
84            "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContentConstrained?access_token={}",
85            self.access_token
86        )
87    }
88
89    async fn auth_headers(&self) -> Result<Vec<(String, String)>, AuthError> {
90        Ok(vec![(
91            "Authorization".to_string(),
92            format!("Bearer {}", self.access_token),
93        )])
94    }
95}
96
97impl RestAuth for GoogleAITokenAuth {
98    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&ModelId>) -> String {
99        let base = "https://generativelanguage.googleapis.com/v1beta";
100        // Token auth uses Bearer header, not query param — build URL without key
101        build_google_ai_rest_url(base, endpoint, model)
102    }
103}