From 98230c3970f4769a2a631d5119d371db6e17e534 Mon Sep 17 00:00:00 2001 From: Jason Date: Sat, 4 Apr 2026 18:59:50 +0800 Subject: [PATCH] feat: add official subscription quota display for Gemini - Read Gemini OAuth credentials from macOS Keychain (gemini-cli-oauth) or legacy file (~/.gemini/oauth_creds.json) - Auto-refresh expired access tokens using refresh_token (Google OAuth tokens expire in ~1h, unlike Claude/Codex) - Two-step API: loadCodeAssist for project ID, then retrieveUserQuota for per-model quota buckets - Classify models into Pro/Flash/Flash Lite categories, show min remaining fraction as utilization percentage - Extend isOfficialProvider() for Gemini (no API key + no base URL) - Parameterize expiredHint i18n key with tool name for all three apps --- src-tauri/src/services/subscription.rs | 503 ++++++++++++++++++++- src/components/SubscriptionQuotaFooter.tsx | 6 +- src/components/providers/ProviderCard.tsx | 9 + src/i18n/locales/en.json | 5 +- src/i18n/locales/ja.json | 5 +- src/i18n/locales/zh.json | 5 +- 6 files changed, 528 insertions(+), 5 deletions(-) diff --git a/src-tauri/src/services/subscription.rs b/src-tauri/src/services/subscription.rs index 727413d10..364ee1fa7 100644 --- a/src-tauri/src/services/subscription.rs +++ b/src-tauri/src/services/subscription.rs @@ -6,6 +6,8 @@ use serde::{Deserialize, Serialize}; use std::time::{SystemTime, UNIX_EPOCH}; +use std::collections::HashMap; + use crate::config; // ── 数据类型 ────────────────────────────────────────────── @@ -706,6 +708,470 @@ async fn query_codex_quota(access_token: &str, account_id: Option<&str>) -> Subs } } +// ── Gemini 凭据读取 ────────────────────────────────────── + +/// Gemini OAuth 凭据文件格式（~/.gemini/oauth_creds.json） +#[derive(Deserialize)] +struct GeminiOAuthCredsFile { + access_token: Option, + refresh_token: Option, + expiry_date: Option, // 毫秒时间戳 +} + +/// (access_token, refresh_token, status, message) +type GeminiCredentials = ( + Option, + Option, + CredentialStatus, + Option, +); + +/// 读取 Gemini OAuth 凭据 +/// +/// 按优先级尝试以下来源： +/// 1. macOS Keychain (service: "gemini-cli-oauth", account: "main-account") +/// 2. 凭据文件 ~/.gemini/oauth_creds.json（遗留格式） +/// +/// 仅 OAuth 认证模式（`oauth-personal`）有效；API key 模式无法查询官方用量。 +fn read_gemini_credentials() -> GeminiCredentials { + #[cfg(target_os = "macos")] + { + if let Some(result) = read_gemini_credentials_from_keychain() { + return result; + } + } + + read_gemini_credentials_from_file() +} + +/// 从 macOS Keychain 读取 Gemini 凭据 +#[cfg(target_os = "macos")] +fn read_gemini_credentials_from_keychain() -> Option { + let output = std::process::Command::new("security") + .args([ + "find-generic-password", + "-s", + "gemini-cli-oauth", + "-a", + "main-account", + "-w", + ]) + .output() + .ok()?; + + if !output.status.success() { + return None; + } + + let json_str = String::from_utf8(output.stdout).ok()?; + let json_str = json_str.trim(); + if json_str.is_empty() { + return None; + } + + Some(parse_gemini_keychain_json(json_str)) +} + +/// 解析 Keychain 格式的 Gemini 凭据 +/// +/// Keychain 格式（keytar）： +/// ```json +/// { "token": { "accessToken": "...", "refreshToken": "...", "expiresAt": 1234 }, "updatedAt": ... } +/// ``` +fn parse_gemini_keychain_json(content: &str) -> GeminiCredentials { + let parsed: serde_json::Value = match serde_json::from_str(content) { + Ok(v) => v, + Err(e) => { + return ( + None, + None, + CredentialStatus::ParseError, + Some(format!("Failed to parse Gemini keychain JSON: {e}")), + ) + } + }; + + let token = match parsed.get("token") { + Some(t) => t, + None => { + // Keychain 中可能是扁平格式，尝试文件格式解析 + return parse_gemini_file_json(content); + } + }; + + let access_token = token + .get("accessToken") + .and_then(|v| v.as_str()) + .map(String::from); + let refresh_token = token + .get("refreshToken") + .and_then(|v| v.as_str()) + .map(String::from); + let expires_at = token.get("expiresAt").and_then(|v| v.as_i64()); + + match access_token { + Some(at) if !at.is_empty() => { + // expiresAt 是毫秒时间戳 + if let Some(exp_ms) = expires_at { + if exp_ms < now_millis() { + return ( + Some(at), + refresh_token, + CredentialStatus::Expired, + Some("Gemini access token has expired".to_string()), + ); + } + } + (Some(at), refresh_token, CredentialStatus::Valid, None) + } + _ => ( + None, + refresh_token, + CredentialStatus::ParseError, + Some("accessToken is empty or missing".to_string()), + ), + } +} + +/// 从文件读取 Gemini 凭据 +fn read_gemini_credentials_from_file() -> GeminiCredentials { + let cred_path = crate::gemini_config::get_gemini_dir().join("oauth_creds.json"); + if !cred_path.exists() { + return (None, None, CredentialStatus::NotFound, None); + } + + let content = match std::fs::read_to_string(&cred_path) { + Ok(c) => c, + Err(e) => { + return ( + None, + None, + CredentialStatus::ParseError, + Some(format!("Failed to read Gemini credentials: {e}")), + ) + } + }; + + parse_gemini_file_json(&content) +} + +/// 解析文件格式的 Gemini 凭据 +/// +/// 文件格式（oauth_creds.json）： +/// ```json +/// { "access_token": "...", "refresh_token": "...", "expiry_date": 1234 } +/// ``` +fn parse_gemini_file_json(content: &str) -> GeminiCredentials { + let creds: GeminiOAuthCredsFile = match serde_json::from_str(content) { + Ok(c) => c, + Err(e) => { + return ( + None, + None, + CredentialStatus::ParseError, + Some(format!("Failed to parse Gemini credentials: {e}")), + ) + } + }; + + let access_token = match creds.access_token { + Some(t) if !t.is_empty() => t, + _ => { + return ( + None, + creds.refresh_token, + CredentialStatus::ParseError, + Some("access_token is empty or missing".to_string()), + ) + } + }; + + // expiry_date 是毫秒时间戳 + if let Some(exp_ms) = creds.expiry_date { + if exp_ms < now_millis() { + return ( + Some(access_token), + creds.refresh_token, + CredentialStatus::Expired, + Some("Gemini access token has expired".to_string()), + ); + } + } + + ( + Some(access_token), + creds.refresh_token, + CredentialStatus::Valid, + None, + ) +} + +// ── Gemini Token 刷新 ────────────────────────────────────── + +/// Gemini OAuth Client 凭据（公开值，来自 Gemini CLI 源码 google-gemini/gemini-cli） +const GEMINI_OAUTH_CLIENT_ID: &str = + "681255809395-oo8ft2oprdrnp9e3aqf6av3hmdib135j.apps.googleusercontent.com"; +const GEMINI_OAUTH_CLIENT_SECRET: &str = "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"; + +/// 使用 refresh_token 刷新 Gemini access token +/// +/// Google OAuth access_token 仅有 ~1h 有效期，需要定期用 refresh_token 刷新。 +/// refresh_token 本身不过期（除非用户撤销授权）。 +async fn refresh_gemini_token(refresh_token: &str) -> Option { + let client = crate::proxy::http_client::get(); + + let resp = client + .post("https://oauth2.googleapis.com/token") + .form(&[ + ("client_id", GEMINI_OAUTH_CLIENT_ID), + ("client_secret", GEMINI_OAUTH_CLIENT_SECRET), + ("refresh_token", refresh_token), + ("grant_type", "refresh_token"), + ]) + .timeout(std::time::Duration::from_secs(10)) + .send() + .await + .ok()?; + + if !resp.status().is_success() { + return None; + } + + let body: serde_json::Value = resp.json().await.ok()?; + body.get("access_token")?.as_str().map(String::from) +} + +// ── Gemini API 查询 ────────────────────────────────────── + +/// loadCodeAssist 响应 +#[derive(Deserialize)] +struct GeminiLoadCodeAssistResponse { + #[serde(rename = "cloudaicompanionProject")] + cloudaicompanion_project: Option, +} + +/// 配额 bucket +#[derive(Deserialize)] +struct GeminiBucketInfo { + #[serde(rename = "remainingFraction")] + remaining_fraction: Option, + #[serde(rename = "resetTime")] + reset_time: Option, + #[serde(rename = "modelId")] + model_id: Option, +} + +/// retrieveUserQuota 响应 +#[derive(Deserialize)] +struct GeminiQuotaResponse { + buckets: Option>, +} + +/// 从 loadCodeAssist 响应中提取项目 ID +fn extract_project_id(value: &serde_json::Value) -> Option { + match value { + serde_json::Value::String(s) => Some(s.clone()), + serde_json::Value::Object(obj) => obj + .get("id") + .or_else(|| obj.get("projectId")) + .and_then(|v| v.as_str()) + .map(String::from), + _ => None, + } +} + +/// 将 Gemini 模型 ID 分类为 Pro / Flash / Flash Lite +fn classify_gemini_model(model_id: &str) -> &str { + if model_id.contains("flash-lite") { + "gemini_flash_lite" + } else if model_id.contains("flash") { + "gemini_flash" + } else if model_id.contains("pro") { + "gemini_pro" + } else { + model_id + } +} + +/// 查询 Gemini 官方订阅额度 +/// +/// 两步 API 调用： +/// 1. loadCodeAssist → 获取 cloudaicompanionProject +/// 2. retrieveUserQuota → 获取按模型分桶的配额数据 +async fn query_gemini_quota(access_token: &str) -> SubscriptionQuota { + let client = crate::proxy::http_client::get(); + + // ── Step 1: loadCodeAssist 获取项目 ID ── + let load_resp = client + .post("https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist") + .header("Authorization", format!("Bearer {access_token}")) + .header("Content-Type", "application/json") + .json(&serde_json::json!({ + "metadata": { + "ideType": "GEMINI_CLI", + "pluginType": "GEMINI" + } + })) + .timeout(std::time::Duration::from_secs(10)) + .send() + .await; + + let load_resp = match load_resp { + Ok(r) => r, + Err(e) => { + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Valid, + format!("Network error (loadCodeAssist): {e}"), + ); + } + }; + + let load_status = load_resp.status(); + if load_status == reqwest::StatusCode::UNAUTHORIZED + || load_status == reqwest::StatusCode::FORBIDDEN + { + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Expired, + format!("Authentication failed (HTTP {load_status}). Please re-login with Gemini CLI."), + ); + } + if !load_status.is_success() { + let body = load_resp.text().await.unwrap_or_default(); + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Valid, + format!("loadCodeAssist failed (HTTP {load_status}): {body}"), + ); + } + + let load_body: GeminiLoadCodeAssistResponse = match load_resp.json().await { + Ok(v) => v, + Err(e) => { + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Valid, + format!("Failed to parse loadCodeAssist response: {e}"), + ); + } + }; + + let project_id = load_body + .cloudaicompanion_project + .as_ref() + .and_then(extract_project_id); + + // ── Step 2: retrieveUserQuota 获取配额 ── + let mut quota_body = serde_json::json!({}); + if let Some(ref pid) = project_id { + quota_body["project"] = serde_json::Value::String(pid.clone()); + } + + let quota_resp = client + .post("https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota") + .header("Authorization", format!("Bearer {access_token}")) + .header("Content-Type", "application/json") + .json("a_body) + .timeout(std::time::Duration::from_secs(10)) + .send() + .await; + + let quota_resp = match quota_resp { + Ok(r) => r, + Err(e) => { + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Valid, + format!("Network error (retrieveUserQuota): {e}"), + ); + } + }; + + let quota_status = quota_resp.status(); + if quota_status == reqwest::StatusCode::UNAUTHORIZED + || quota_status == reqwest::StatusCode::FORBIDDEN + { + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Expired, + format!("Authentication failed (HTTP {quota_status})."), + ); + } + if !quota_status.is_success() { + let body = quota_resp.text().await.unwrap_or_default(); + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Valid, + format!("retrieveUserQuota failed (HTTP {quota_status}): {body}"), + ); + } + + let quota_data: GeminiQuotaResponse = match quota_resp.json().await { + Ok(v) => v, + Err(e) => { + return SubscriptionQuota::error( + "gemini", + CredentialStatus::Valid, + format!("Failed to parse quota response: {e}"), + ); + } + }; + + // ── 按模型分类汇总，每类取最低 remainingFraction ── + let mut category_map: HashMap)> = HashMap::new(); + + if let Some(buckets) = quota_data.buckets { + for bucket in buckets { + let model_id = bucket.model_id.as_deref().unwrap_or("unknown"); + let category = classify_gemini_model(model_id).to_string(); + let remaining = bucket.remaining_fraction.unwrap_or(1.0).clamp(0.0, 1.0); + + let entry = category_map + .entry(category) + .or_insert((remaining, bucket.reset_time.clone())); + if remaining < entry.0 { + entry.0 = remaining; + if bucket.reset_time.is_some() { + entry.1.clone_from(&bucket.reset_time); + } + } + } + } + + // 转换为 tiers（remainingFraction → utilization: 已用百分比） + let sort_order = |name: &str| -> usize { + match name { + "gemini_pro" => 0, + "gemini_flash" => 1, + "gemini_flash_lite" => 2, + _ => 3, + } + }; + + let mut tiers: Vec = category_map + .into_iter() + .map(|(name, (remaining, reset_time))| QuotaTier { + name, + utilization: (1.0 - remaining) * 100.0, + resets_at: reset_time, + }) + .collect(); + + tiers.sort_by_key(|t| sort_order(&t.name)); + + SubscriptionQuota { + tool: "gemini".to_string(), + credential_status: CredentialStatus::Valid, + credential_message: None, + success: true, + tiers, + extra_usage: None, + error: None, + queried_at: Some(now_millis()), + } +} + // ── 入口函数 ────────────────────────────────────────────── /// 查询指定 CLI 工具的官方订阅额度 @@ -771,7 +1237,42 @@ pub async fn get_subscription_quota(tool: &str) -> Result { + let (token, refresh_token, status, message) = read_gemini_credentials(); + + match status { + CredentialStatus::NotFound => Ok(SubscriptionQuota::not_found("gemini")), + CredentialStatus::ParseError => Ok(SubscriptionQuota::error( + "gemini", + CredentialStatus::ParseError, + message.unwrap_or_else(|| "Failed to parse credentials".to_string()), + )), + CredentialStatus::Expired => { + // Gemini access_token 仅 ~1h 有效，尝试用 refresh_token 刷新 + if let Some(ref rt) = refresh_token { + if let Some(new_token) = refresh_gemini_token(rt).await { + return Ok(query_gemini_quota(&new_token).await); + } + } + // 刷新失败，尝试用旧 token + if let Some(ref token) = token { + let result = query_gemini_quota(token).await; + if result.success { + return Ok(result); + } + } + Ok(SubscriptionQuota::error( + "gemini", + CredentialStatus::Expired, + message.unwrap_or_else(|| "Gemini OAuth token has expired".to_string()), + )) + } + CredentialStatus::Valid => { + let token = token.expect("token must be Some when status is Valid"); + Ok(query_gemini_quota(&token).await) + } + } + } _ => Ok(SubscriptionQuota::not_found(tool)), } } diff --git a/src/components/SubscriptionQuotaFooter.tsx b/src/components/SubscriptionQuotaFooter.tsx index 6666f74ba..10d1906d2 100644 --- a/src/components/SubscriptionQuotaFooter.tsx +++ b/src/components/SubscriptionQuotaFooter.tsx @@ -16,6 +16,10 @@ const TIER_I18N_KEYS: Record = { seven_day: "subscription.sevenDay", seven_day_opus: "subscription.sevenDayOpus", seven_day_sonnet: "subscription.sevenDaySonnet", + // Gemini 模型分类 + gemini_pro: "subscription.geminiPro", + gemini_flash: "subscription.geminiFlash", + gemini_flash_lite: "subscription.geminiFlashLite", }; /** 根据使用百分比返回颜色 class */ @@ -123,7 +127,7 @@ const SubscriptionQuotaFooter: React.FC = ({
{t("subscription.expired")} - {t("subscription.expiredHint")} + {t("subscription.expiredHint", { tool: appId })}
diff --git a/src/components/providers/ProviderCard.tsx b/src/components/providers/ProviderCard.tsx index f327d8d84..f3569c363 100644 --- a/src/components/providers/ProviderCard.tsx +++ b/src/components/providers/ProviderCard.tsx @@ -68,6 +68,15 @@ function isOfficialProvider(provider: Provider, appId: AppId): boolean { const apiKey = config?.auth?.OPENAI_API_KEY; return !apiKey || (typeof apiKey === "string" && apiKey.trim() === ""); } + if (appId === "gemini") { + // 无 GEMINI_API_KEY 且无 GOOGLE_GEMINI_BASE_URL → Google OAuth 官方模式 + const apiKey = config?.env?.GEMINI_API_KEY; + const baseUrl = config?.env?.GOOGLE_GEMINI_BASE_URL; + return ( + (!apiKey || (typeof apiKey === "string" && apiKey.trim() === "")) && + (!baseUrl || (typeof baseUrl === "string" && baseUrl.trim() === "")) + ); + } return false; } diff --git a/src/i18n/locales/en.json b/src/i18n/locales/en.json index e3c0ffe5d..d5a15b69e 100644 --- a/src/i18n/locales/en.json +++ b/src/i18n/locales/en.json @@ -2283,11 +2283,14 @@ "sevenDay": "7-Day", "sevenDayOpus": "7-Day (Opus)", "sevenDaySonnet": "7-Day (Sonnet)", + "geminiPro": "Pro", + "geminiFlash": "Flash", + "geminiFlashLite": "Flash Lite", "utilization": "{{value}}%", "resetsIn": "Resets in {{time}}", "extraUsage": "Extra Usage", "expired": "Session expired", - "expiredHint": "Run the claude command to refresh your login", + "expiredHint": "Run the {{tool}} command to refresh your login", "queryFailed": "Query failed", "refresh": "Refresh" } diff --git a/src/i18n/locales/ja.json b/src/i18n/locales/ja.json index 46519e9f5..2babea0ff 100644 --- a/src/i18n/locales/ja.json +++ b/src/i18n/locales/ja.json @@ -2283,11 +2283,14 @@ "sevenDay": "7日間", "sevenDayOpus": "7日間(Opus)", "sevenDaySonnet": "7日間(Sonnet)", + "geminiPro": "Pro", + "geminiFlash": "Flash", + "geminiFlashLite": "Flash Lite", "utilization": "{{value}}%", "resetsIn": "{{time}}後にリセット", "extraUsage": "超過使用量", "expired": "セッション期限切れ", - "expiredHint": "claudeコマンドを実行してログインを更新してください", + "expiredHint": "{{tool}}コマンドを実行してログインを更新してください", "queryFailed": "クエリに失敗しました", "refresh": "更新" } diff --git a/src/i18n/locales/zh.json b/src/i18n/locales/zh.json index 5d27d3e29..ea08109a8 100644 --- a/src/i18n/locales/zh.json +++ b/src/i18n/locales/zh.json @@ -2283,11 +2283,14 @@ "sevenDay": "7天", "sevenDayOpus": "7天(Opus)", "sevenDaySonnet": "7天(Sonnet)", + "geminiPro": "Pro", + "geminiFlash": "Flash", + "geminiFlashLite": "Flash Lite", "utilization": "{{value}}%", "resetsIn": "{{time}}后重置", "extraUsage": "超额用量", "expired": "会话已过期", - "expiredHint": "请运行 claude 命令刷新登录", + "expiredHint": "请运行 {{tool}} 命令刷新登录", "queryFailed": "查询失败", "refresh": "刷新" }