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, 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    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&GeminiModel>) -> String {
41        let base = "https://generativelanguage.googleapis.com/v1beta";
42        build_google_ai_rest_url(base, endpoint, model, &self.api_key)
43    }
44
45    async fn auth_headers(&self) -> Result<Vec<(String, String)>, AuthError> {
46        Ok(vec![]) // API key is in the URL
47    }
48
49    fn query_params(&self) -> Vec<(String, String)> {
50        vec![("key".to_string(), self.api_key.clone())]
51    }
52}
53
54// ---------------------------------------------------------------------------
55// Google AI — OAuth2 access token authentication
56// ---------------------------------------------------------------------------
57
58/// Google AI OAuth2 access token authentication.
59///
60/// The access token is included directly in the WebSocket URL.
61pub struct GoogleAITokenAuth {
62    access_token: String,
63}
64
65impl GoogleAITokenAuth {
66    /// Create a new Google AI token auth provider with the given access token.
67    pub fn new(access_token: impl Into<String>) -> Self {
68        Self {
69            access_token: access_token.into(),
70        }
71    }
72}
73
74#[async_trait]
75impl AuthProvider for GoogleAITokenAuth {
76    fn ws_url(&self, _model: &GeminiModel) -> String {
77        format!(
78            "wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1alpha.GenerativeService.BidiGenerateContentConstrained?access_token={}",
79            self.access_token
80        )
81    }
82
83    fn rest_url(&self, endpoint: ServiceEndpoint, model: Option<&GeminiModel>) -> String {
84        let base = "https://generativelanguage.googleapis.com/v1beta";
85        // Token auth uses Bearer header, not query param — build URL without key
86        build_google_ai_rest_url_no_key(base, endpoint, model)
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}