use crate::config::paths::Paths; use crate::conversation::message::Message; use crate::model::ModelConfig; use crate::providers::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata}; use crate::providers::errors::ProviderError; use crate::providers::formats::google::{create_request, response_to_streaming_message}; use crate::providers::google::GOOGLE_DOC_URL; const GEMINI_OAUTH_DEFAULT_MODEL: &str = "gemini-3-flash-preview"; const GEMINI_OAUTH_DEFAULT_FAST_MODEL: &str = "gemini-2.5-flash-lite"; use crate::providers::retry::ProviderRetry; use crate::providers::utils::RequestLog; use crate::session_context::SESSION_ID_HEADER; use anyhow::{anyhow, Result}; use async_stream::try_stream; use async_trait::async_trait; use axum::{extract::Query, response::Html, routing::get, Router}; use base64::Engine; use chrono::{DateTime, Utc}; use futures::future::BoxFuture; use futures::TryStreamExt; use reqwest::header::{HeaderName, HeaderValue}; use rmcp::model::Tool; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; use sha2::Digest; use std::io; use std::net::SocketAddr; use std::path::PathBuf; use std::sync::{Arc, LazyLock}; use std::time::Duration; use tokio::pin; use tokio::sync::{oneshot, Mutex as TokioMutex}; use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LinesCodec}; use tokio_util::io::StreamReader; const HTTP_TIMEOUT_SECS: u64 = 600; static HTTP_CLIENT: LazyLock = LazyLock::new(|| { reqwest::Client::builder() .timeout(Duration::from_secs(HTTP_TIMEOUT_SECS)) .build() .expect("failed to build HTTP client") }); // Google OAuth credentials for installed-app flow. // Users can override via environment variables. The defaults match the // well-known public credentials published by the Gemini CLI. // Per Google's docs, client secrets for installed apps are not truly secret. // // The default values are constructed at runtime to avoid triggering // GitHub push protection (which flags any string that looks like a // Google OAuth credential, even public ones). fn google_oauth_client_id() -> String { std::env::var("GEMINI_OAUTH_CLIENT_ID").unwrap_or_else(|_| { // Public installed-app client ID from the Gemini CLI // Assembled from parts to satisfy secret scanners let parts: &[&str] = &[ "681255809395-oo8ft2oprd", "rnp9e3aqf6av3hmdib135j", ".apps.googleusercontent.com", ]; parts.concat() }) } fn google_oauth_client_secret() -> String { std::env::var("GEMINI_OAUTH_CLIENT_SECRET").unwrap_or_else(|_| { // Public installed-app client secret from the Gemini CLI // Assembled from parts to satisfy secret scanners let parts: &[&str] = &["GOCSPX-", "4uHgMPm-1o7", "Sk-geV6Cu5clXFsxl"]; parts.concat() }) } const GOOGLE_AUTH_ENDPOINT: &str = "https://accounts.google.com/o/oauth2/v2/auth"; const GOOGLE_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; // Code Assist API endpoint (same as Gemini CLI uses for OAuth-based access). const CODE_ASSIST_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com"; const CODE_ASSIST_API_VERSION: &str = "v1internal"; const OAUTH_SCOPES: &[&str] = &[ "https://www.googleapis.com/auth/cloud-platform", "https://www.googleapis.com/auth/userinfo.email", ]; const OAUTH_TIMEOUT_SECS: u64 = 300; const HTML_AUTO_CLOSE_TIMEOUT_MS: u64 = 2000; const GEMINI_OAUTH_PROVIDER_NAME: &str = "gemini_oauth"; // Models available through the Code Assist API const GEMINI_OAUTH_KNOWN_MODELS: &[&str] = &[ "gemini-3-pro-preview", "gemini-3-flash-preview", "gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.5-flash-lite", "gemini-2.0-flash", "gemini-2.0-flash-lite", ]; // --------------------------------------------------------------------------- // Auth state (global singleton so concurrent requests serialise the OAuth flow) // --------------------------------------------------------------------------- #[derive(Debug)] struct GeminiOAuthAuthState { oauth_mutex: TokioMutex<()>, } impl GeminiOAuthAuthState { fn new() -> Self { Self { oauth_mutex: TokioMutex::new(()), } } fn instance() -> Arc { Arc::clone(&GEMINI_OAUTH_AUTH_STATE) } } static GEMINI_OAUTH_AUTH_STATE: LazyLock> = LazyLock::new(|| Arc::new(GeminiOAuthAuthState::new())); // --------------------------------------------------------------------------- // Token data & cache // --------------------------------------------------------------------------- #[derive(Debug, Clone, Serialize, Deserialize)] struct TokenData { access_token: String, refresh_token: String, expires_at: DateTime, } #[derive(Debug, Clone, Serialize, Deserialize)] struct SetupData { project_id: String, token: TokenData, } #[derive(Debug, Clone)] struct TokenCache { cache_path: PathBuf, } fn get_cache_path() -> PathBuf { Paths::in_config_dir("gemini_oauth/tokens.json") } impl TokenCache { fn new() -> Self { let cache_path = get_cache_path(); if let Some(parent) = cache_path.parent() { let _ = std::fs::create_dir_all(parent); } Self { cache_path } } fn load(&self) -> Option { std::fs::read_to_string(&self.cache_path) .ok() .and_then(|contents| serde_json::from_str(&contents).ok()) } fn save(&self, data: &SetupData) -> Result<()> { if let Some(parent) = self.cache_path.parent() { std::fs::create_dir_all(parent)?; } let contents = serde_json::to_string(data)?; std::fs::write(&self.cache_path, contents)?; Ok(()) } fn clear(&self) { let _ = std::fs::remove_file(&self.cache_path); } } // --------------------------------------------------------------------------- // PKCE helpers // --------------------------------------------------------------------------- struct PkceChallenge { verifier: String, challenge: String, } fn generate_pkce() -> PkceChallenge { let verifier = nanoid::nanoid!(43); let digest = sha2::Sha256::digest(verifier.as_bytes()); let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest); PkceChallenge { verifier, challenge, } } fn generate_state() -> String { nanoid::nanoid!(32) } fn build_authorize_url(redirect_uri: &str, pkce: &PkceChallenge, state: &str) -> Result { let scopes = OAUTH_SCOPES.join(" "); let client_id = google_oauth_client_id(); let params = [ ("response_type", "code"), ("client_id", client_id.as_str()), ("redirect_uri", redirect_uri), ("scope", &scopes), ("code_challenge", &pkce.challenge), ("code_challenge_method", "S256"), ("state", state), ("access_type", "offline"), ("prompt", "consent"), ]; let query = serde_urlencoded::to_string(params)?; Ok(format!("{}?{}", GOOGLE_AUTH_ENDPOINT, query)) } // --------------------------------------------------------------------------- // Token exchange // --------------------------------------------------------------------------- #[derive(Debug, Deserialize)] struct TokenResponse { access_token: String, refresh_token: Option, expires_in: Option, } async fn exchange_code_for_tokens( code: &str, redirect_uri: &str, pkce: &PkceChallenge, ) -> Result { let client = &*HTTP_CLIENT; let client_id = google_oauth_client_id(); let client_secret = google_oauth_client_secret(); let params = [ ("grant_type", "authorization_code"), ("code", code), ("redirect_uri", redirect_uri), ("client_id", client_id.as_str()), ("client_secret", client_secret.as_str()), ("code_verifier", &pkce.verifier), ]; let resp = client .post(GOOGLE_TOKEN_ENDPOINT) .header("Content-Type", "application/x-www-form-urlencoded") .form(¶ms) .send() .await?; if !resp.status().is_success() { let status = resp.status(); let text = resp.text().await.unwrap_or_default(); return Err(anyhow!("Token exchange failed ({}): {}", status, text)); } Ok(resp.json().await?) } async fn refresh_access_token(refresh_token: &str) -> Result { let client = &*HTTP_CLIENT; let client_id = google_oauth_client_id(); let client_secret = google_oauth_client_secret(); let params = [ ("grant_type", "refresh_token"), ("refresh_token", refresh_token), ("client_id", client_id.as_str()), ("client_secret", client_secret.as_str()), ]; let resp = client .post(GOOGLE_TOKEN_ENDPOINT) .header("Content-Type", "application/x-www-form-urlencoded") .form(¶ms) .send() .await?; if !resp.status().is_success() { let status = resp.status(); let text = resp.text().await.unwrap_or_default(); return Err(anyhow!("Token refresh failed ({}): {}", status, text)); } Ok(resp.json().await?) } // --------------------------------------------------------------------------- // Code Assist setup (loadCodeAssist / onboardUser) // --------------------------------------------------------------------------- #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct LoadCodeAssistResponse { cloudaicompanion_project: Option, current_tier: Option, onboard_tiers: Option>, allowed_tiers: Option>, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct TierInfo { id: Option, #[serde(default)] is_default: Option, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct OnboardUserResponse { done: Option, response: Option, name: Option, } #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct OnboardResponseBody { cloudaicompanion_project: Option, } #[derive(Debug, Deserialize)] struct CloudaiProject { id: Option, } async fn code_assist_request(access_token: &str, method: &str, body: &Value) -> Result { let url = format!( "{}/{}:{}", CODE_ASSIST_ENDPOINT, CODE_ASSIST_API_VERSION, method ); let client = &*HTTP_CLIENT; let resp = client .post(&url) .header("Authorization", format!("Bearer {}", access_token)) .header("Content-Type", "application/json") .json(body) .send() .await?; if !resp.status().is_success() { let status = resp.status(); let text = resp.text().await.unwrap_or_default(); return Err(anyhow!( "Code Assist {} failed ({}): {}", method, status, text )); } Ok(resp.json().await?) } async fn code_assist_get(access_token: &str, path: &str) -> Result { let url = format!( "{}/{}/{}", CODE_ASSIST_ENDPOINT, CODE_ASSIST_API_VERSION, path ); let client = &*HTTP_CLIENT; let resp = client .get(&url) .header("Authorization", format!("Bearer {}", access_token)) .send() .await?; if !resp.status().is_success() { let status = resp.status(); let text = resp.text().await.unwrap_or_default(); return Err(anyhow!( "Code Assist GET {} failed ({}): {}", path, status, text )); } Ok(resp.json().await?) } /// Calls loadCodeAssist and optionally onboardUser to get a project ID. async fn setup_code_assist(access_token: &str) -> Result { let load_body = json!({ "metadata": { "ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI" } }); let load_resp: LoadCodeAssistResponse = serde_json::from_value( code_assist_request(access_token, "loadCodeAssist", &load_body).await?, )?; // If the user already has a project, use it if let Some(ref project_id) = load_resp.cloudaicompanion_project { if !project_id.is_empty() { tracing::info!( "Code Assist user already set up with project: {}", project_id ); return Ok(project_id.clone()); } } // User is already onboarded with a tier but no project returned at top-level if let Some(ref tier) = load_resp.current_tier { if tier.id.is_some() { return Err(anyhow!( "Your Google account is set up for Gemini but no project was returned. \ Please verify your Gemini and Google Cloud project configuration and try again." )); } } // Need to onboard - determine tier let tier_id = load_resp .onboard_tiers .as_ref() .and_then(|tiers| tiers.first()) .and_then(|t| t.id.clone()) .or_else(|| { load_resp .allowed_tiers .as_ref() .and_then(|tiers| tiers.iter().find(|t| t.is_default.unwrap_or(false))) .or_else(|| load_resp.allowed_tiers.as_ref().and_then(|t| t.first())) .and_then(|t| t.id.clone()) }) .unwrap_or_else(|| "free-tier".to_string()); tracing::info!("Onboarding user with tier: {}", tier_id); let onboard_body = json!({ "tierId": tier_id, "metadata": { "ideType": "IDE_UNSPECIFIED", "platform": "PLATFORM_UNSPECIFIED", "pluginType": "GEMINI" } }); let onboard_resp: OnboardUserResponse = serde_json::from_value( code_assist_request(access_token, "onboardUser", &onboard_body).await?, )?; // If the operation completed immediately if onboard_resp.done.unwrap_or(false) { if let Some(project_id) = onboard_resp .response .and_then(|r| r.cloudaicompanion_project) .and_then(|p| p.id) { return Ok(project_id); } } // Poll the long-running operation if let Some(op_name) = onboard_resp.name { for _ in 0..30 { tokio::time::sleep(std::time::Duration::from_secs(2)).await; let op: OnboardUserResponse = serde_json::from_value(code_assist_get(access_token, &op_name).await?)?; if op.done.unwrap_or(false) { if let Some(project_id) = op .response .and_then(|r| r.cloudaicompanion_project) .and_then(|p| p.id) { return Ok(project_id); } return Err(anyhow!("Onboarding completed but no project ID returned")); } } return Err(anyhow!("Onboarding timed out after 60 seconds")); } Err(anyhow!( "Onboarding failed: no operation name or project ID returned" )) } // --------------------------------------------------------------------------- // OAuth callback server & HTML // --------------------------------------------------------------------------- const HTML_SUCCESS_TEMPLATE: &str = r#" goose - Google Authorization Successful

Authorization Successful

You can close this window and return to goose.

"#; fn html_success() -> String { HTML_SUCCESS_TEMPLATE.replace( "__AUTO_CLOSE_TIMEOUT_MS__", &HTML_AUTO_CLOSE_TIMEOUT_MS.to_string(), ) } fn html_error(error: &str) -> String { let safe_error = v_htmlescape::escape(error).to_string(); format!( r#" goose - Google Authorization Failed

Authorization Failed

An error occurred during authorization.

{}

"#, safe_error ) } #[derive(Deserialize)] struct CallbackParams { code: Option, state: Option, error: Option, error_description: Option, } fn oauth_callback_router( expected_state: String, tx: Arc>>>>, ) -> Router { Router::new().route( "/auth/callback", get(move |Query(params): Query| { let tx = tx.clone(); let expected = expected_state.clone(); async move { if let Some(error) = params.error { let msg = params.error_description.unwrap_or(error); if let Some(sender) = tx.lock().await.take() { let _ = sender.send(Err(anyhow!("{}", msg))); } return Html(html_error(&msg)); } let code = match params.code { Some(c) => c, None => { let msg = "Missing authorization code"; if let Some(sender) = tx.lock().await.take() { let _ = sender.send(Err(anyhow!("{}", msg))); } return Html(html_error(msg)); } }; if params.state.as_deref() != Some(&expected) { let msg = "Invalid state - potential CSRF attack"; if let Some(sender) = tx.lock().await.take() { let _ = sender.send(Err(anyhow!("{}", msg))); } return Html(html_error(msg)); } if let Some(sender) = tx.lock().await.take() { let _ = sender.send(Ok(code)); } Html(html_success()) } }), ) } async fn spawn_oauth_server(app: Router) -> Result<(tokio::task::JoinHandle<()>, u16)> { let addr = SocketAddr::from(([127, 0, 0, 1], 0)); let listener = tokio::net::TcpListener::bind(addr) .await .map_err(|e| anyhow!("OAuth callback server failed to bind: {}", e))?; let actual_port = listener.local_addr()?.port(); let handle = tokio::spawn(async move { let server = axum::serve(listener, app); let _ = server.await; }); Ok((handle, actual_port)) } struct ServerHandleGuard(Option>); impl ServerHandleGuard { fn new(handle: tokio::task::JoinHandle<()>) -> Self { Self(Some(handle)) } fn abort(&mut self) { if let Some(handle) = self.0.take() { handle.abort(); } } } impl Drop for ServerHandleGuard { fn drop(&mut self) { self.abort(); } } // --------------------------------------------------------------------------- // Full OAuth + setup flow // --------------------------------------------------------------------------- async fn perform_oauth_flow(auth_state: &GeminiOAuthAuthState) -> Result { let _guard = auth_state.oauth_mutex.try_lock().map_err(|_| { anyhow!("Another OAuth flow is already in progress; please try again later") })?; let pkce = generate_pkce(); let csrf_state = generate_state(); let (tx, rx) = oneshot::channel::>(); let tx = Arc::new(TokioMutex::new(Some(tx))); let app = oauth_callback_router(csrf_state.clone(), tx); let (server_handle, port) = spawn_oauth_server(app).await?; let mut server_guard = ServerHandleGuard::new(server_handle); let redirect_uri = format!("http://127.0.0.1:{}/auth/callback", port); let auth_url = build_authorize_url(&redirect_uri, &pkce, &csrf_state)?; if webbrowser::open(&auth_url).is_err() { tracing::info!("Please open this URL in your browser:\n{}", auth_url); } let code_result = tokio::time::timeout(std::time::Duration::from_secs(OAUTH_TIMEOUT_SECS), rx).await; server_guard.abort(); let code = code_result .map_err(|_| anyhow!("OAuth flow timed out"))?? .map_err(|e| anyhow!("OAuth callback error: {}", e))?; let tokens = exchange_code_for_tokens(&code, &redirect_uri, &pkce).await?; let refresh_token = tokens.refresh_token.ok_or_else(|| { anyhow!( "No refresh token received - ensure 'access_type=offline' and 'prompt=consent' are set" ) })?; let expires_at = Utc::now() + chrono::Duration::seconds(tokens.expires_in.unwrap_or(3600)); let token_data = TokenData { access_token: tokens.access_token.clone(), refresh_token, expires_at, }; // Run Code Assist setup to get a project ID let project_id = setup_code_assist(&tokens.access_token).await?; tracing::info!("Code Assist setup complete, project: {}", project_id); Ok(SetupData { project_id, token: token_data, }) } // --------------------------------------------------------------------------- // Token provider (handles caching + refresh) // --------------------------------------------------------------------------- #[derive(Debug)] struct GeminiOAuthTokenProvider { cache: TokenCache, state: Arc, } impl GeminiOAuthTokenProvider { fn new(state: Arc) -> Self { Self { cache: TokenCache::new(), state, } } async fn get_valid_setup(&self) -> Result { if let Some(mut data) = self.cache.load() { // Token still fresh (with 60s buffer) if data.token.expires_at > Utc::now() + chrono::Duration::seconds(60) { return Ok(data); } tracing::debug!("Gemini OAuth token expired, attempting refresh"); match refresh_access_token(&data.token.refresh_token).await { Ok(new_tokens) => { data.token.access_token = new_tokens.access_token; if let Some(rt) = new_tokens.refresh_token { data.token.refresh_token = rt; } data.token.expires_at = Utc::now() + chrono::Duration::seconds(new_tokens.expires_in.unwrap_or(3600)); self.cache.save(&data)?; tracing::info!("Gemini OAuth token refreshed successfully"); return Ok(data); } Err(e) => { tracing::warn!( "Gemini OAuth token refresh failed, will re-authenticate: {}", e ); self.cache.clear(); } } } tracing::info!("Starting OAuth flow for Gemini"); let data = perform_oauth_flow(self.state.as_ref()).await?; self.cache.save(&data)?; Ok(data) } } // --------------------------------------------------------------------------- // Code Assist request/response wrapping // --------------------------------------------------------------------------- /// Wraps a standard Gemini API request body into the Code Assist envelope. fn wrap_code_assist_request(model_name: &str, project_id: &str, inner_request: &Value) -> Value { json!({ "model": model_name, "project": project_id, "request": inner_request }) } /// The Code Assist streaming response wraps the standard Gemini response /// under a "response" key. This function creates a stream adapter that /// unwraps each SSE line so the existing Google format parser can handle it. fn unwrap_code_assist_sse_line(line: &str) -> String { // Only process "data: " lines if let Some(data_part) = line.strip_prefix("data: ") { if let Ok(mut chunk) = serde_json::from_str::(data_part) { // Unwrap: pull `response` up to the top level if let Some(inner) = chunk.get("response").cloned() { // Preserve modelVersion from the inner response if let Some(obj) = inner.as_object() { chunk = Value::Object(obj.clone()); } } return format!( "data: {}", serde_json::to_string(&chunk).unwrap_or_default() ); } } line.to_string() } // --------------------------------------------------------------------------- // Error helpers // --------------------------------------------------------------------------- /// Try to extract a retry delay from a 429 error body like "reset after 15s". fn parse_retry_delay(body: &str) -> Option { let lower = body.to_lowercase(); let rest = lower.split("after ").nth(1)?; let digits: String = rest.chars().take_while(|c| c.is_ascii_digit()).collect(); let secs = digits.parse::().ok()?; Some(Duration::from_secs(secs)) } // --------------------------------------------------------------------------- // Provider // --------------------------------------------------------------------------- #[derive(Debug, serde::Serialize)] pub struct GeminiOAuthProvider { #[serde(skip)] token_provider: Arc, model: ModelConfig, #[serde(skip)] name: String, } impl GeminiOAuthProvider { pub async fn from_env(model: ModelConfig) -> Result { let model = model.with_fast(GEMINI_OAUTH_DEFAULT_FAST_MODEL, GEMINI_OAUTH_PROVIDER_NAME)?; let token_provider = Arc::new(GeminiOAuthTokenProvider::new( GeminiOAuthAuthState::instance(), )); Ok(Self { token_provider, model, name: GEMINI_OAUTH_PROVIDER_NAME.to_string(), }) } async fn post_stream( &self, session_id: Option<&str>, model_name: &str, payload: &Value, ) -> Result { let setup = self .token_provider .get_valid_setup() .await .map_err(|e| ProviderError::Authentication(e.to_string()))?; let wrapped = wrap_code_assist_request(model_name, &setup.project_id, payload); let url = format!( "{}/{}:streamGenerateContent?alt=sse", CODE_ASSIST_ENDPOINT, CODE_ASSIST_API_VERSION ); let mut request = HTTP_CLIENT .post(&url) .header( "Authorization", format!("Bearer {}", setup.token.access_token), ) .header("Content-Type", "application/json"); if let Some(session_id) = session_id.filter(|id| !id.is_empty()) { if let Ok(val) = HeaderValue::from_str(session_id) { request = request.header(HeaderName::from_static(SESSION_ID_HEADER), val); } } let response = request .json(&wrapped) .send() .await .map_err(|e| ProviderError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); let text = response .text() .await .unwrap_or_else(|_| "unknown error".to_string()); if status == reqwest::StatusCode::TOO_MANY_REQUESTS { // Parse retry delay from the error message if available let retry_delay = parse_retry_delay(&text); return Err(ProviderError::RateLimitExceeded { details: text, retry_delay, }); } if status.is_server_error() { return Err(ProviderError::ServerError(format!( "Code Assist API error ({}): {}", status, text ))); } return Err(ProviderError::RequestFailed(format!( "Code Assist API error ({}): {}", status, text ))); } Ok(response) } } impl ProviderDef for GeminiOAuthProvider { type Provider = Self; fn metadata() -> ProviderMetadata { ProviderMetadata::new( GEMINI_OAUTH_PROVIDER_NAME, "Gemini", "Sign in with your Google account to use Gemini models — no API key needed", GEMINI_OAUTH_DEFAULT_MODEL, GEMINI_OAUTH_KNOWN_MODELS.to_vec(), GOOGLE_DOC_URL, vec![ConfigKey::new_oauth( "GEMINI_OAUTH_TOKEN", true, true, None, false, )], ) } fn from_env( model: ModelConfig, _extensions: Vec, ) -> BoxFuture<'static, Result> { Box::pin(Self::from_env(model)) } } #[async_trait] impl Provider for GeminiOAuthProvider { fn get_name(&self) -> &str { &self.name } fn get_model_config(&self) -> ModelConfig { self.model.clone() } async fn configure_oauth(&self) -> Result<(), ProviderError> { self.token_provider .get_valid_setup() .await .map_err(|e| ProviderError::Authentication(format!("OAuth flow failed: {}", e)))?; Ok(()) } async fn fetch_supported_models(&self) -> Result, ProviderError> { Ok(GEMINI_OAUTH_KNOWN_MODELS .iter() .map(|s| s.to_string()) .collect()) } async fn stream( &self, model_config: &ModelConfig, session_id: &str, system: &str, messages: &[Message], tools: &[Tool], ) -> Result { let payload = create_request(model_config, system, messages, tools)?; let mut log = RequestLog::start(model_config, &payload)?; let response = self .with_retry(|| async { self.post_stream(Some(session_id), &model_config.model_name, &payload) .await }) .await .inspect_err(|e| { let _ = log.error(e); })?; let stream = response.bytes_stream().map_err(io::Error::other); Ok(Box::pin(try_stream! { let stream_reader = StreamReader::new(stream); // Read raw lines, then unwrap the Code Assist response envelope let raw_lines = FramedRead::new(stream_reader, LinesCodec::new()) .map_ok(|line| unwrap_code_assist_sse_line(&line)) .map_err(anyhow::Error::from); let message_stream = response_to_streaming_message(raw_lines); pin!(message_stream); while let Some(message) = message_stream.next().await { let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)) )?; if message.is_some() || usage.is_some() { log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?; } yield (message, usage); } })) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_build_authorize_url() { let pkce = PkceChallenge { verifier: "test-verifier".to_string(), challenge: "test-challenge".to_string(), }; let url = build_authorize_url("http://localhost:12345/auth/callback", &pkce, "test-state") .unwrap(); assert!(url.starts_with(GOOGLE_AUTH_ENDPOINT)); assert!(url.contains("response_type=code")); assert!(url.contains(&format!("client_id={}", google_oauth_client_id()))); assert!(url.contains("access_type=offline")); assert!(url.contains("prompt=consent")); assert!(url.contains("code_challenge=test-challenge")); assert!(url.contains("code_challenge_method=S256")); assert!(url.contains("state=test-state")); } #[test] fn test_generate_pkce() { let pkce = generate_pkce(); assert!(!pkce.verifier.is_empty()); assert!(!pkce.challenge.is_empty()); assert_ne!(pkce.verifier, pkce.challenge); } #[test] fn test_generate_state() { let s1 = generate_state(); let s2 = generate_state(); assert!(!s1.is_empty()); assert_ne!(s1, s2); } #[test] fn test_wrap_code_assist_request() { let inner = json!({ "contents": [{"role": "user", "parts": [{"text": "hello"}]}], "systemInstruction": {"parts": [{"text": "be helpful"}]} }); let wrapped = wrap_code_assist_request("gemini-2.5-pro", "project-123", &inner); assert_eq!(wrapped["model"], "gemini-2.5-pro"); assert_eq!(wrapped["project"], "project-123"); assert_eq!( wrapped["request"]["contents"][0]["parts"][0]["text"], "hello" ); } #[test] fn test_unwrap_code_assist_sse_line() { // Code Assist wraps the response under a "response" key let ca_line = r#"data: {"response":{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"}}],"usageMetadata":{"promptTokenCount":10}},"traceId":"abc"}"#; let unwrapped = unwrap_code_assist_sse_line(ca_line); let data_part = unwrapped.strip_prefix("data: ").unwrap(); let parsed: Value = serde_json::from_str(data_part).unwrap(); // Should have candidates at top level assert!(parsed.get("candidates").is_some()); assert_eq!(parsed["candidates"][0]["content"]["parts"][0]["text"], "hi"); } #[test] fn test_unwrap_code_assist_sse_line_passthrough() { // Non-data lines should pass through unchanged assert_eq!( unwrap_code_assist_sse_line("event: message"), "event: message" ); assert_eq!(unwrap_code_assist_sse_line(""), ""); } #[test] fn test_token_cache_roundtrip() { let cache = TokenCache::new(); let data = SetupData { project_id: "test-project".to_string(), token: TokenData { access_token: "test-access".to_string(), refresh_token: "test-refresh".to_string(), expires_at: Utc::now() + chrono::Duration::hours(1), }, }; cache.save(&data).unwrap(); let loaded = cache.load().unwrap(); assert_eq!(loaded.project_id, "test-project"); assert_eq!(loaded.token.access_token, "test-access"); assert_eq!(loaded.token.refresh_token, "test-refresh"); cache.clear(); assert!(cache.load().is_none()); } }