Files
CC-Switch/src-tauri/src/services/subscription.rs
T
Jason 89ff2d58d1 fix(usage): improve usage-query resilience and error surfacing
- useUsageQuery: retry once + keep-last-good — show the last successful
  result for up to 10min when a query fails transiently (network/timeout/
  HTTP 5xx), so a single blip no longer flips the card to red. Deterministic
  failures (auth, empty key, unknown provider, 4xx) surface immediately and
  clear the snapshot so a stale quota can't resurface after credentials change.
- bump native balance/coding-plan/subscription request timeouts 10s -> 15s
  for slow cross-border endpoints.
- coding_plan: return explicit errors ("API key is empty" / "Unknown coding
  plan provider") instead of a blank failure, mirroring balance.
- add unit tests for keep-last-good and transient/deterministic classification.
2026-06-14 21:21:45 +08:00

1338 lines
42 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! 官方订阅额度查询服务
//!
//! 读取 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,
/// 使用百分比 0100
pub utilization: f64,
/// ISO 8601 重置时间
pub resets_at: Option<String>,
/// ZenMux: 已用额度(USD
#[serde(skip_serializing_if = "Option::is_none")]
pub used_value_usd: Option<f64>,
/// ZenMux: 窗口上限(USD
#[serde(skip_serializing_if = "Option::is_none")]
pub max_value_usd: Option<f64>,
}
/// 超额使用信息
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExtraUsage {
pub is_enabled: bool,
pub monthly_limit: Option<f64>,
pub used_credits: Option<f64>,
pub utilization: Option<f64>,
pub currency: Option<String>,
}
/// 订阅额度查询结果
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SubscriptionQuota {
pub tool: String,
pub credential_status: CredentialStatus,
pub credential_message: Option<String>,
pub success: bool,
pub tiers: Vec<QuotaTier>,
pub extra_usage: Option<ExtraUsage>,
pub error: Option<String>,
pub queried_at: Option<i64>,
}
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<String>,
#[serde(rename = "expiresAt")]
expires_at: Option<serde_json::Value>,
}
/// 读取 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<String>, CredentialStatus, Option<String>) {
// 来源 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<String>, CredentialStatus, Option<String>)> {
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<String>, CredentialStatus, Option<String>) {
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 凭据 JSONKeychain 和文件共用)
fn parse_claude_credentials_json(
content: &str,
) -> (Option<String>, CredentialStatus, Option<String>) {
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<f64>,
resets_at: Option<String>,
}
/// Claude OAuth 用量 API 响应中的超额用量
#[derive(Deserialize)]
struct ApiExtraUsage {
is_enabled: Option<bool>,
monthly_limit: Option<f64>,
used_credits: Option<f64>,
utilization: Option<f64>,
currency: Option<String>,
}
/// 已知的 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 PlanKimi / 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::<ApiUsageWindow>(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::<ApiUsageWindow>(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::<ApiExtraUsage>(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<String>,
tokens: Option<CodexTokens>,
last_refresh: Option<String>,
}
#[derive(Deserialize)]
struct CodexTokens {
access_token: Option<String>,
account_id: Option<String>,
}
/// (access_token, account_id, status, message)
type CodexCredentials = (
Option<String>,
Option<String>,
CredentialStatus,
Option<String>,
);
/// 读取 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<CodexCredentials> {
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 凭据 JSONKeychain 和文件共用)
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<f64>,
limit_window_seconds: Option<i64>,
reset_at: Option<i64>,
}
#[derive(Deserialize)]
struct CodexRateLimit {
primary_window: Option<CodexRateLimitWindow>,
secondary_window: Option<CodexRateLimitWindow>,
}
#[derive(Deserialize)]
struct CodexUsageResponse {
rate_limit: Option<CodexRateLimit>,
}
/// 根据窗口秒数映射到 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<String> {
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<String>,
refresh_token: Option<String>,
expiry_date: Option<i64>, // 毫秒时间戳
}
/// (access_token, refresh_token, status, message)
type GeminiCredentials = (
Option<String>,
Option<String>,
CredentialStatus,
Option<String>,
);
/// 读取 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<GeminiCredentials> {
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<String> {
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<serde_json::Value>,
}
/// 配额 bucket
#[derive(Deserialize)]
struct GeminiBucketInfo {
#[serde(rename = "remainingFraction")]
remaining_fraction: Option<f64>,
#[serde(rename = "resetTime")]
reset_time: Option<String>,
#[serde(rename = "modelId")]
model_id: Option<String>,
}
/// retrieveUserQuota 响应
#[derive(Deserialize)]
struct GeminiQuotaResponse {
buckets: Option<Vec<GeminiBucketInfo>>,
}
/// 从 loadCodeAssist 响应中提取项目 ID
fn extract_project_id(value: &serde_json::Value) -> Option<String> {
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(&quota_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<String, (f64, Option<String>)> = 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);
}
}
}
}
// 转换为 tiersremainingFraction → 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<QuotaTier> = 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<SubscriptionQuota, String> {
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
}