From 0630502479deb228065eb09bf736beb48e996043 Mon Sep 17 00:00:00 2001 From: NessZerra Date: Sun, 12 Jul 2026 06:26:04 +0700 Subject: [PATCH 1/2] Fix Gemini paid-tier plan display Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/cli/usage.rs | 13 ++ rust/src/providers/gemini/api.rs | 203 +++++++++++++++++++++++++++++-- rust/src/providers/gemini/mod.rs | 4 +- 3 files changed, 211 insertions(+), 9 deletions(-) diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 6c1302db5c..f3e716578a 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -601,4 +601,17 @@ mod tests { "Claude: Session (5h) <1%, Weekly 100%, resets n/a, Pro" ); } + + #[test] + fn gemini_plan_preserves_acronym_casing() { + let result = fetch_result( + UsageSnapshot::new(RateWindow::new(0.0)) + .with_login_method("Gemini Code Assist in Google One AI Pro"), + ); + + let output = render_text(ProviderId::Gemini, &result, false); + + assert!(output.contains("Plan: Gemini Code Assist in Google One AI Pro")); + assert!(!output.contains("Google One Ai Pro")); + } } diff --git a/rust/src/providers/gemini/api.rs b/rust/src/providers/gemini/api.rs index bfa4a2dd7c..15acc54f73 100755 --- a/rust/src/providers/gemini/api.rs +++ b/rust/src/providers/gemini/api.rs @@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; const QUOTA_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota"; +const CODE_ASSIST_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist"; const TOKEN_REFRESH_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; /// Gemini API client @@ -25,12 +26,20 @@ impl GeminiApi { } /// Fetch quota information from the Gemini API - /// Returns (primary RateWindow, optional model-specific RateWindow, optional email) + /// Returns (primary RateWindow, optional model-specific RateWindow, optional email, optional plan) /// Note: Gemini quota API requires OAuth tokens, not API keys pub async fn fetch_quota( &self, _ctx: &FetchContext, - ) -> Result<(RateWindow, Option, Option), ProviderError> { + ) -> Result< + ( + RateWindow, + Option, + Option, + Option, + ), + ProviderError, + > { // Gemini quota endpoint requires OAuth credentials (not API keys) // Always load OAuth credentials from ~/.gemini/oauth_creds.json let mut creds = self.load_credentials()?; @@ -46,6 +55,8 @@ impl GeminiApi { .clone() .ok_or_else(|| ProviderError::AuthRequired)?; + let code_assist = self.load_code_assist_status(&access_token).await; + // Fetch quota let response = self .client @@ -74,7 +85,47 @@ impl GeminiApi { .map_err(|e| ProviderError::Parse(e.to_string()))?; // Since we use OAuth, we can use the credentials we already loaded for email extraction - self.parse_quota_response(quota_response, Some(&creds)) + let (primary, model_specific, email) = + self.parse_quota_response(quota_response, Some(&creds))?; + let hosted_domain = creds + .id_token + .as_deref() + .and_then(extract_hosted_domain_from_jwt); + let plan = resolve_account_plan(&code_assist, hosted_domain.as_deref()); + + Ok((primary, model_specific, email, plan)) + } + + async fn load_code_assist_status(&self, access_token: &str) -> CodeAssistStatus { + let response = self + .client + .post(CODE_ASSIST_ENDPOINT) + .header("Authorization", format!("Bearer {}", access_token)) + .header("Content-Type", "application/json") + .body(r#"{"metadata":{"ideType":"GEMINI_CLI","pluginType":"GEMINI"}}"#) + .timeout(std::time::Duration::from_secs(10)) + .send() + .await; + + let response = match response { + Ok(response) if response.status().is_success() => response, + Ok(response) => { + tracing::warn!(status = %response.status(), "Gemini loadCodeAssist request failed"); + return CodeAssistStatus::default(); + } + Err(error) => { + tracing::warn!(%error, "Gemini loadCodeAssist request failed"); + return CodeAssistStatus::default(); + } + }; + + match response.json::().await { + Ok(response) => code_assist_status(response), + Err(error) => { + tracing::warn!(%error, "Gemini loadCodeAssist response was invalid"); + CodeAssistStatus::default() + } + } } fn load_credentials(&self) -> Result { @@ -481,6 +532,78 @@ struct QuotaBucket { token_type: Option, } +#[derive(Default)] +struct CodeAssistStatus { + tier: Option, + paid_tier_name: Option, +} + +#[derive(Clone, Copy)] +enum GeminiUserTier { + Free, + Legacy, + Standard, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct CodeAssistResponse { + current_tier: Option, + paid_tier: Option, +} + +#[derive(Deserialize)] +struct CodeAssistTier { + id: Option, +} + +#[derive(Deserialize)] +struct PaidTier { + name: Option, +} + +fn code_assist_status(response: CodeAssistResponse) -> CodeAssistStatus { + let tier = response + .current_tier + .and_then(|tier| match tier.id.as_deref() { + Some("free-tier") => Some(GeminiUserTier::Free), + Some("legacy-tier") => Some(GeminiUserTier::Legacy), + Some("standard-tier") => Some(GeminiUserTier::Standard), + _ => None, + }); + let paid_tier_name = response + .paid_tier + .and_then(|tier| tier.name) + .map(|name| name.trim().to_string()) + .filter(|name| !name.is_empty()); + + CodeAssistStatus { + tier, + paid_tier_name, + } +} + +#[cfg(test)] +fn parse_code_assist_status(body: &str) -> CodeAssistStatus { + serde_json::from_str(body) + .map(code_assist_status) + .unwrap_or_default() +} + +fn resolve_account_plan(status: &CodeAssistStatus, hosted_domain: Option<&str>) -> Option { + if let Some(plan) = &status.paid_tier_name { + return Some(plan.clone()); + } + + match status.tier { + Some(GeminiUserTier::Standard) => Some("Paid".to_string()), + Some(GeminiUserTier::Free) if hosted_domain.is_some() => Some("Workspace".to_string()), + Some(GeminiUserTier::Free) => Some("Free".to_string()), + Some(GeminiUserTier::Legacy) => Some("Legacy".to_string()), + None => None, + } +} + // --- Helper functions --- fn parse_iso_date(s: &str) -> Option> { @@ -498,6 +621,20 @@ fn parse_iso_date(s: &str) -> Option> { } fn extract_email_from_jwt(token: &str) -> Option { + jwt_payload(token)? + .get("email") + .and_then(|v| v.as_str()) + .map(str::to_owned) +} + +fn extract_hosted_domain_from_jwt(token: &str) -> Option { + jwt_payload(token)? + .get("hd") + .and_then(|v| v.as_str()) + .map(str::to_owned) +} + +fn jwt_payload(token: &str) -> Option { let parts: Vec<&str> = token.split('.').collect(); if parts.len() < 2 { return None; @@ -515,8 +652,60 @@ fn extract_email_from_jwt(token: &str) -> Option { let decoded = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, &payload).ok()?; - let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; - json.get("email") - .and_then(|v| v.as_str()) - .map(|s| s.to_string()) + serde_json::from_slice(&decoded).ok() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn paid_tier_name_overrides_generic_tier_fallbacks() { + let status = parse_code_assist_status( + r#"{ + "currentTier": { "id": "free-tier" }, + "paidTier": { "name": "Gemini Code Assist in Google One AI Pro" } + }"#, + ); + + assert_eq!( + resolve_account_plan(&status, Some("example.com")), + Some("Gemini Code Assist in Google One AI Pro".to_string()) + ); + + let standard = parse_code_assist_status( + r#"{ + "currentTier": { "id": "standard-tier" }, + "paidTier": { "name": "Plus" } + }"#, + ); + + assert_eq!( + resolve_account_plan(&standard, None), + Some("Plus".to_string()) + ); + } + + #[test] + fn generic_tier_fallbacks_remain_when_paid_tier_is_absent() { + let free_tier = parse_code_assist_status(r#"{"currentTier":{"id":"free-tier"}}"#); + let paid = parse_code_assist_status(r#"{"currentTier":{"id":"standard-tier"}}"#); + + assert_eq!( + resolve_account_plan(&free_tier, Some("example.com")), + Some("Workspace".to_string()) + ); + assert_eq!( + resolve_account_plan(&free_tier, None), + Some("Free".to_string()) + ); + assert_eq!(resolve_account_plan(&paid, None), Some("Paid".to_string())); + } + + #[test] + fn invalid_code_assist_response_does_not_create_a_generic_plan() { + let status = parse_code_assist_status("not json"); + + assert_eq!(resolve_account_plan(&status, Some("example.com")), None); + } } diff --git a/rust/src/providers/gemini/mod.rs b/rust/src/providers/gemini/mod.rs index ec938e94a7..30364c234e 100755 --- a/rust/src/providers/gemini/mod.rs +++ b/rust/src/providers/gemini/mod.rs @@ -60,7 +60,7 @@ impl Provider for GeminiProvider { tracing::debug!("Fetching Gemini usage via API"); match self.api.fetch_quota(ctx).await { - Ok((primary, model_specific, email)) => { + Ok((primary, model_specific, email, plan)) => { let mut usage = UsageSnapshot::new(primary); if let Some(ms) = model_specific { usage = usage.with_model_specific(ms); @@ -68,7 +68,7 @@ impl Provider for GeminiProvider { if let Some(e) = email { usage = usage.with_email(e); } - usage = usage.with_login_method("Gemini CLI"); + usage = usage.with_login_method(plan.unwrap_or_else(|| "Gemini CLI".to_string())); Ok(ProviderFetchResult::new(usage, "cli")) } From 5226d076233d7b81c794a3938d62064fef64d885 Mon Sep 17 00:00:00 2001 From: NessZerra Date: Sun, 12 Jul 2026 07:29:00 +0700 Subject: [PATCH 2/2] Preserve Gemini tier fallbacks Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/providers/gemini/api.rs | 69 +++++++++++++++----------------- 1 file changed, 32 insertions(+), 37 deletions(-) diff --git a/rust/src/providers/gemini/api.rs b/rust/src/providers/gemini/api.rs index 15acc54f73..41747abb5c 100755 --- a/rust/src/providers/gemini/api.rs +++ b/rust/src/providers/gemini/api.rs @@ -119,8 +119,8 @@ impl GeminiApi { } }; - match response.json::().await { - Ok(response) => code_assist_status(response), + match response.text().await { + Ok(body) => parse_code_assist_status(&body), Err(error) => { tracing::warn!(%error, "Gemini loadCodeAssist response was invalid"); CodeAssistStatus::default() @@ -545,36 +545,27 @@ enum GeminiUserTier { Standard, } -#[derive(Deserialize)] -#[serde(rename_all = "camelCase")] -struct CodeAssistResponse { - current_tier: Option, - paid_tier: Option, -} - -#[derive(Deserialize)] -struct CodeAssistTier { - id: Option, -} - -#[derive(Deserialize)] -struct PaidTier { - name: Option, -} - -fn code_assist_status(response: CodeAssistResponse) -> CodeAssistStatus { - let tier = response - .current_tier - .and_then(|tier| match tier.id.as_deref() { - Some("free-tier") => Some(GeminiUserTier::Free), - Some("legacy-tier") => Some(GeminiUserTier::Legacy), - Some("standard-tier") => Some(GeminiUserTier::Standard), +fn parse_code_assist_status(body: &str) -> CodeAssistStatus { + let Ok(json) = serde_json::from_str::(body) else { + return CodeAssistStatus::default(); + }; + + let tier = json + .get("currentTier") + .and_then(|tier| tier.get("id")) + .and_then(serde_json::Value::as_str) + .and_then(|tier| match tier { + "free-tier" => Some(GeminiUserTier::Free), + "legacy-tier" => Some(GeminiUserTier::Legacy), + "standard-tier" => Some(GeminiUserTier::Standard), _ => None, }); - let paid_tier_name = response - .paid_tier - .and_then(|tier| tier.name) - .map(|name| name.trim().to_string()) + let paid_tier_name = json + .get("paidTier") + .and_then(|tier| tier.get("name")) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .map(str::to_owned) .filter(|name| !name.is_empty()); CodeAssistStatus { @@ -583,13 +574,6 @@ fn code_assist_status(response: CodeAssistResponse) -> CodeAssistStatus { } } -#[cfg(test)] -fn parse_code_assist_status(body: &str) -> CodeAssistStatus { - serde_json::from_str(body) - .map(code_assist_status) - .unwrap_or_default() -} - fn resolve_account_plan(status: &CodeAssistStatus, hosted_domain: Option<&str>) -> Option { if let Some(plan) = &status.paid_tier_name { return Some(plan.clone()); @@ -708,4 +692,15 @@ mod tests { assert_eq!(resolve_account_plan(&status, Some("example.com")), None); } + + #[test] + fn malformed_paid_tier_preserves_current_tier_fallback() { + let status = + parse_code_assist_status(r#"{"currentTier":{"id":"free-tier"},"paidTier":[]}"#); + + assert_eq!( + resolve_account_plan(&status, Some("example.com")), + Some("Workspace".to_string()) + ); + } }