mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-31 19:22:15 +08:00
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
This commit is contained in:
Jason
@@ -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<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": ... }
|
||||
/// ```
|
||||
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(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<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") {
|
||||
"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<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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 转换为 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<QuotaTier> = 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<SubscriptionQuota, Str
|
||||
}
|
||||
}
|
||||
}
|
||||
// Gemini: 暂不支持
|
||||
"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)),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user