//! 官方订阅额度查询服务 //! //! 读取 CLI 工具的已有 OAuth 凭据,查询官方订阅额度。 //! 第一层:仅读取凭据,不实现登录/刷新。 use serde::{Deserialize, Serialize}; use std::time::{SystemTime, UNIX_EPOCH}; use std::collections::HashMap; use crate::config; // ── 数据类型 ────────────────────────────────────────────── /// 凭据状态 #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum CredentialStatus { Valid, Expired, NotFound, ParseError, } /// 单个限速窗口(如 5小时会话、7天周期) #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct QuotaTier { /// 窗口标识:five_hour, seven_day, seven_day_opus, seven_day_sonnet 等 pub name: String, /// 使用百分比 0–100 pub utilization: f64, /// ISO 8601 重置时间 pub resets_at: Option, /// ZenMux: 已用额度(USD) #[serde(skip_serializing_if = "Option::is_none")] pub used_value_usd: Option, /// ZenMux: 窗口上限(USD) #[serde(skip_serializing_if = "Option::is_none")] pub max_value_usd: Option, } /// 超额使用信息 #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExtraUsage { pub is_enabled: bool, pub monthly_limit: Option, pub used_credits: Option, pub utilization: Option, pub currency: Option, } /// 订阅额度查询结果 #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct SubscriptionQuota { pub tool: String, pub credential_status: CredentialStatus, pub credential_message: Option, pub success: bool, pub tiers: Vec, pub extra_usage: Option, pub error: Option, pub queried_at: Option, } impl SubscriptionQuota { pub(crate) fn not_found(tool: &str) -> Self { Self { tool: tool.to_string(), credential_status: CredentialStatus::NotFound, credential_message: None, success: false, tiers: vec![], extra_usage: None, error: None, queried_at: None, } } pub(crate) fn error(tool: &str, status: CredentialStatus, message: String) -> Self { Self { tool: tool.to_string(), credential_status: status, credential_message: Some(message.clone()), success: false, tiers: vec![], extra_usage: None, error: Some(message), queried_at: Some(now_millis()), } } } // ── Claude 凭据读取 ────────────────────────────────────── /// Claude OAuth 凭据文件中的嵌套结构 #[derive(Deserialize)] struct ClaudeOAuthEntry { #[serde(rename = "accessToken")] access_token: Option, #[serde(rename = "expiresAt")] expires_at: Option, } /// 读取 Claude OAuth 凭据 /// /// 按优先级尝试以下来源: /// 1. macOS Keychain (service: "Claude Code-credentials") /// 2. 凭据文件 ~/.claude/.credentials.json /// /// JSON 格式(两种 key 都兼容): /// {"claudeAiOauth": {"accessToken": "...", "expiresAt": ...}} /// {"claude.ai_oauth": {"accessToken": "...", "expiresAt": ...}} fn read_claude_credentials() -> (Option, CredentialStatus, Option) { // 来源 1: macOS Keychain #[cfg(target_os = "macos")] { if let Some(result) = read_claude_credentials_from_keychain() { return result; } } // 来源 2: 凭据文件 read_claude_credentials_from_file() } /// 从 macOS Keychain 读取 Claude 凭据 #[cfg(target_os = "macos")] fn read_claude_credentials_from_keychain( ) -> Option<(Option, CredentialStatus, Option)> { let output = std::process::Command::new("security") .args([ "find-generic-password", "-s", "Claude Code-credentials", "-w", ]) .output() .ok()?; if !output.status.success() { return None; // Keychain 中无此条目,回退到文件 } let json_str = String::from_utf8(output.stdout).ok()?; let json_str = json_str.trim(); if json_str.is_empty() { return None; } Some(parse_claude_credentials_json(json_str)) } /// 从文件读取 Claude 凭据 fn read_claude_credentials_from_file() -> (Option, CredentialStatus, Option) { let cred_path = config::get_claude_config_dir().join(".credentials.json"); if !cred_path.exists() { return (None, CredentialStatus::NotFound, None); } let content = match std::fs::read_to_string(&cred_path) { Ok(c) => c, Err(e) => { return ( None, CredentialStatus::ParseError, Some(format!("Failed to read credentials file: {e}")), ); } }; parse_claude_credentials_json(&content) } /// 解析 Claude 凭据 JSON(Keychain 和文件共用) fn parse_claude_credentials_json( content: &str, ) -> (Option, CredentialStatus, Option) { let parsed: serde_json::Value = match serde_json::from_str(content) { Ok(v) => v, Err(e) => { return ( None, CredentialStatus::ParseError, Some(format!("Failed to parse credentials JSON: {e}")), ); } }; // 兼容两种 key 名 let entry_value = parsed .get("claudeAiOauth") .or_else(|| parsed.get("claude.ai_oauth")); let entry_value = match entry_value { Some(v) => v, None => { return ( None, CredentialStatus::ParseError, Some("No OAuth entry found in credentials".to_string()), ); } }; let entry: ClaudeOAuthEntry = match serde_json::from_value(entry_value.clone()) { Ok(e) => e, Err(e) => { return ( None, CredentialStatus::ParseError, Some(format!("Failed to parse OAuth entry: {e}")), ); } }; let access_token = match entry.access_token { Some(t) if !t.is_empty() => t, _ => { return ( None, CredentialStatus::ParseError, Some("accessToken is empty or missing".to_string()), ); } }; // 检查 token 是否过期 if let Some(expires_at) = entry.expires_at { if is_token_expired(&expires_at) { return ( Some(access_token), CredentialStatus::Expired, Some("OAuth token has expired".to_string()), ); } } (Some(access_token), CredentialStatus::Valid, None) } /// 判断 token 是否过期,兼容 Unix 时间戳(秒/毫秒)和 ISO 字符串 fn is_token_expired(expires_at: &serde_json::Value) -> bool { let now_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs(); match expires_at { serde_json::Value::Number(n) => { if let Some(ts) = n.as_u64() { // 区分秒和毫秒(毫秒级时间戳大于 1e12) let ts_secs = if ts > 1_000_000_000_000 { ts / 1000 } else { ts }; ts_secs < now_secs } else { false } } serde_json::Value::String(s) => { // 尝试解析 ISO 8601 格式 if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(s) { (dt.timestamp() as u64) < now_secs } else if let Ok(dt) = chrono::NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f") { (dt.and_utc().timestamp() as u64) < now_secs } else { false // 无法解析时不视为过期 } } _ => false, } } // ── Claude API 查询 ────────────────────────────────────── /// Claude OAuth 用量 API 响应中的单个窗口 #[derive(Deserialize)] struct ApiUsageWindow { utilization: Option, resets_at: Option, } /// Claude OAuth 用量 API 响应中的超额用量 #[derive(Deserialize)] struct ApiExtraUsage { is_enabled: Option, monthly_limit: Option, used_credits: Option, utilization: Option, currency: Option, } /// 已知的 Claude 用量窗口名称。`QuotaTier::name` 会是其中之一。 pub const TIER_FIVE_HOUR: &str = "five_hour"; pub const TIER_SEVEN_DAY: &str = "seven_day"; pub const TIER_SEVEN_DAY_OPUS: &str = "seven_day_opus"; pub const TIER_SEVEN_DAY_SONNET: &str = "seven_day_sonnet"; /// Coding Plan(Kimi / MiniMax)的周窗口 tier 名。与 `coding_plan::query_*` /// 写入、tray 渲染、commands::provider 扁平化三处共用同一标识。 pub const TIER_WEEKLY_LIMIT: &str = "weekly_limit"; /// Gemini 用量分组名称(按模型而非时间窗口)。`classify_gemini_model` 输出。 pub const TIER_GEMINI_PRO: &str = "gemini_pro"; pub const TIER_GEMINI_FLASH: &str = "gemini_flash"; pub const TIER_GEMINI_FLASH_LITE: &str = "gemini_flash_lite"; const KNOWN_TIERS: &[&str] = &[ TIER_FIVE_HOUR, TIER_SEVEN_DAY, TIER_SEVEN_DAY_OPUS, TIER_SEVEN_DAY_SONNET, ]; /// 查询 Claude 官方订阅额度 async fn query_claude_quota(access_token: &str) -> SubscriptionQuota { let client = crate::proxy::http_client::get(); let resp = client .get("https://api.anthropic.com/api/oauth/usage") .header("Authorization", format!("Bearer {access_token}")) .header("anthropic-beta", "oauth-2025-04-20") .header("Accept", "application/json") .timeout(std::time::Duration::from_secs(15)) .send() .await; let resp = match resp { Ok(r) => r, Err(e) => { return SubscriptionQuota::error( "claude", CredentialStatus::Valid, format!("Network error: {e}"), ); } }; let status = resp.status(); if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { return SubscriptionQuota::error( "claude", CredentialStatus::Expired, format!("Authentication failed (HTTP {status}). Please re-login with Claude CLI."), ); } if !status.is_success() { let body = resp.text().await.unwrap_or_default(); return SubscriptionQuota::error( "claude", CredentialStatus::Valid, format!("API error (HTTP {status}): {body}"), ); } let body: serde_json::Value = match resp.json().await { Ok(v) => v, Err(e) => { return SubscriptionQuota::error( "claude", CredentialStatus::Valid, format!("Failed to parse API response: {e}"), ); } }; // 解析已知的 tier 窗口 let mut tiers = Vec::new(); for &tier_name in KNOWN_TIERS { if let Some(window) = body.get(tier_name) { if let Ok(w) = serde_json::from_value::(window.clone()) { if let Some(util) = w.utilization { tiers.push(QuotaTier { name: tier_name.to_string(), utilization: util, resets_at: w.resets_at, used_value_usd: None, max_value_usd: None, }); } } } } // 也解析未知窗口(API 可能返回新的窗口类型) if let Some(obj) = body.as_object() { for (key, value) in obj { if key == "extra_usage" || KNOWN_TIERS.contains(&key.as_str()) { continue; } if let Ok(w) = serde_json::from_value::(value.clone()) { if let Some(util) = w.utilization { tiers.push(QuotaTier { name: key.clone(), utilization: util, resets_at: w.resets_at, used_value_usd: None, max_value_usd: None, }); } } } } // 解析超额使用 let extra_usage = body.get("extra_usage").and_then(|v| { serde_json::from_value::(v.clone()) .ok() .map(|e| ExtraUsage { is_enabled: e.is_enabled.unwrap_or(false), monthly_limit: e.monthly_limit, used_credits: e.used_credits, utilization: e.utilization, currency: e.currency, }) }); SubscriptionQuota { tool: "claude".to_string(), credential_status: CredentialStatus::Valid, credential_message: None, success: true, tiers, extra_usage, error: None, queried_at: Some(now_millis()), } } // ── Codex 凭据读取 ────────────────────────────────────── #[derive(Deserialize)] struct CodexAuthJson { auth_mode: Option, tokens: Option, last_refresh: Option, } #[derive(Deserialize)] struct CodexTokens { access_token: Option, account_id: Option, } /// (access_token, account_id, status, message) type CodexCredentials = ( Option, Option, CredentialStatus, Option, ); /// 读取 Codex OAuth 凭据 /// /// 按优先级尝试以下来源: /// 1. macOS Keychain (service: "Codex Auth") /// 2. 凭据文件 ~/.codex/auth.json /// /// 仅 auth_mode == "chatgpt" (OAuth) 时有效,API key 模式不支持用量查询。 fn read_codex_credentials() -> CodexCredentials { #[cfg(target_os = "macos")] { if let Some(result) = read_codex_credentials_from_keychain() { return result; } } read_codex_credentials_from_file() } /// 从 macOS Keychain 读取 Codex 凭据 #[cfg(target_os = "macos")] fn read_codex_credentials_from_keychain() -> Option { let output = std::process::Command::new("security") .args(["find-generic-password", "-s", "Codex Auth", "-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_codex_credentials_json(json_str)) } /// 从文件读取 Codex 凭据 fn read_codex_credentials_from_file() -> CodexCredentials { let auth_path = crate::codex_config::get_codex_auth_path(); if !auth_path.exists() { return (None, None, CredentialStatus::NotFound, None); } let content = match std::fs::read_to_string(&auth_path) { Ok(c) => c, Err(e) => { return ( None, None, CredentialStatus::ParseError, Some(format!("Failed to read Codex auth file: {e}")), ); } }; parse_codex_credentials_json(&content) } /// 解析 Codex 凭据 JSON(Keychain 和文件共用) fn parse_codex_credentials_json(content: &str) -> CodexCredentials { let auth: CodexAuthJson = match serde_json::from_str(content) { Ok(a) => a, Err(e) => { return ( None, None, CredentialStatus::ParseError, Some(format!("Failed to parse Codex auth JSON: {e}")), ); } }; // 仅 OAuth 模式有用量数据 if auth.auth_mode.as_deref() != Some("chatgpt") { return ( None, None, CredentialStatus::NotFound, Some("Codex not using OAuth mode".to_string()), ); } let tokens = match auth.tokens { Some(t) => t, None => { return ( None, None, CredentialStatus::ParseError, Some("No tokens in Codex auth".to_string()), ); } }; let access_token = match tokens.access_token { Some(t) if !t.is_empty() => t, _ => { return ( None, None, CredentialStatus::ParseError, Some("access_token is empty or missing".to_string()), ); } }; // 检查 token 是否可能过期(距上次刷新 > 8 天) if let Some(ref last_refresh) = auth.last_refresh { if is_codex_token_stale(last_refresh) { return ( Some(access_token), tokens.account_id, CredentialStatus::Expired, Some("Codex token may be stale (>8 days since last refresh)".to_string()), ); } } ( Some(access_token), tokens.account_id, CredentialStatus::Valid, None, ) } /// 判断 Codex token 是否可能过期(Codex CLI 在 >8 天时自动刷新) fn is_codex_token_stale(last_refresh: &str) -> bool { let now_secs = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs(); if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(last_refresh) { let age_secs = now_secs.saturating_sub(dt.timestamp() as u64); age_secs > 8 * 24 * 3600 } else { false } } // ── Codex API 查询 ────────────────────────────────────── #[derive(Deserialize)] struct CodexRateLimitWindow { used_percent: Option, limit_window_seconds: Option, reset_at: Option, } #[derive(Deserialize)] struct CodexRateLimit { primary_window: Option, secondary_window: Option, } #[derive(Deserialize)] struct CodexUsageResponse { rate_limit: Option, } /// 根据窗口秒数映射到 tier 名称(与 Claude 的命名兼容以复用前端 i18n) fn window_seconds_to_tier_name(secs: i64) -> String { match secs { 18000 => "five_hour".to_string(), 604800 => "seven_day".to_string(), s => { let hours = s / 3600; if hours >= 24 { format!("{}_day", hours / 24) } else { format!("{}_hour", hours) } } } } /// Unix 时间戳(秒)转 ISO 8601 字符串 fn unix_ts_to_iso(ts: i64) -> Option { chrono::DateTime::from_timestamp(ts, 0).map(|dt| dt.to_rfc3339()) } /// 查询 Codex / ChatGPT 反代订阅额度 /// /// 参数化 `tool_label` 和 `expired_message` 让该函数可被两个调用点共用: /// - `"codex"` + "Please re-login with Codex CLI."(CLI 凭据路径) /// - `"codex_oauth"` + "Please re-login via cc-switch."(cc-switch 自管 OAuth 路径) pub(crate) async fn query_codex_quota( access_token: &str, account_id: Option<&str>, tool_label: &str, expired_message: &str, ) -> SubscriptionQuota { let client = crate::proxy::http_client::get(); let mut req = client .get("https://chatgpt.com/backend-api/wham/usage") .header("Authorization", format!("Bearer {access_token}")) .header("User-Agent", "codex-cli") .header("Accept", "application/json"); if let Some(id) = account_id { req = req.header("ChatGPT-Account-Id", id); } let resp = match req.timeout(std::time::Duration::from_secs(15)).send().await { Ok(r) => r, Err(e) => { return SubscriptionQuota::error( tool_label, CredentialStatus::Valid, format!("Network error: {e}"), ); } }; let status = resp.status(); if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { return SubscriptionQuota::error( tool_label, CredentialStatus::Expired, format!("{expired_message} (HTTP {status})"), ); } if !status.is_success() { let body = resp.text().await.unwrap_or_default(); return SubscriptionQuota::error( tool_label, CredentialStatus::Valid, format!("API error (HTTP {status}): {body}"), ); } let body: CodexUsageResponse = match resp.json().await { Ok(v) => v, Err(e) => { return SubscriptionQuota::error( tool_label, CredentialStatus::Valid, format!("Failed to parse API response: {e}"), ); } }; let mut tiers = Vec::new(); if let Some(rate_limit) = body.rate_limit { for window in [rate_limit.primary_window, rate_limit.secondary_window] .into_iter() .flatten() { if let Some(used) = window.used_percent { tiers.push(QuotaTier { name: window .limit_window_seconds .map(window_seconds_to_tier_name) .unwrap_or_else(|| "unknown".to_string()), utilization: used, resets_at: window.reset_at.and_then(unix_ts_to_iso), used_value_usd: None, max_value_usd: None, }); } } } SubscriptionQuota { tool: tool_label.to_string(), credential_status: CredentialStatus::Valid, credential_message: None, success: true, tiers, extra_usage: None, error: None, queried_at: Some(now_millis()), } } // ── 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": ... } /// ``` #[cfg(target_os = "macos")] 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(15)) .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") { TIER_GEMINI_FLASH_LITE } else if model_id.contains("flash") { TIER_GEMINI_FLASH } else if model_id.contains("pro") { TIER_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(15)) .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(15)) .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 { TIER_GEMINI_PRO => 0, TIER_GEMINI_FLASH => 1, TIER_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, used_value_usd: None, max_value_usd: None, }) .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 工具的官方订阅额度 pub async fn get_subscription_quota(tool: &str) -> Result { match tool { "claude" => { let (token, status, message) = read_claude_credentials(); match status { CredentialStatus::NotFound => Ok(SubscriptionQuota::not_found("claude")), CredentialStatus::ParseError => Ok(SubscriptionQuota::error( "claude", CredentialStatus::ParseError, message.unwrap_or_else(|| "Failed to parse credentials".to_string()), )), CredentialStatus::Expired => { // 即使过期也尝试调用 API(token 可能实际上仍有效) if let Some(token) = token { let result = query_claude_quota(&token).await; if result.success { return Ok(result); } } Ok(SubscriptionQuota::error( "claude", CredentialStatus::Expired, message.unwrap_or_else(|| "OAuth token has expired".to_string()), )) } CredentialStatus::Valid => { let token = token.expect("token must be Some when status is Valid"); Ok(query_claude_quota(&token).await) } } } "codex" => { let (token, account_id, status, message) = read_codex_credentials(); match status { CredentialStatus::NotFound => Ok(SubscriptionQuota::not_found("codex")), CredentialStatus::ParseError => Ok(SubscriptionQuota::error( "codex", CredentialStatus::ParseError, message.unwrap_or_else(|| "Failed to parse credentials".to_string()), )), CredentialStatus::Expired => { // 即使可能过期也尝试调用 API if let Some(token) = token { let result = query_codex_quota( &token, account_id.as_deref(), "codex", "Authentication failed. Please re-login with Codex CLI.", ) .await; if result.success { return Ok(result); } } Ok(SubscriptionQuota::error( "codex", CredentialStatus::Expired, message.unwrap_or_else(|| "Codex OAuth token may be stale".to_string()), )) } CredentialStatus::Valid => { let token = token.expect("token must be Some when status is Valid"); Ok(query_codex_quota( &token, account_id.as_deref(), "codex", "Authentication failed. Please re-login with Codex CLI.", ) .await) } } } "gemini" => { 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)), } } // ── 辅助函数 ────────────────────────────────────────────── fn now_millis() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_millis() as i64 }