feat(claude-code): 支持通过 /api/oauth/usage 查询账号 5H/周额度并在号池展示

2026-10-07 18:07:47 +08:00 · 2026-09-29 22:40:44 +08:00
parent 00a315e3bf
commit 8093899b5a
15 changed files with 653 additions and 33 deletions
@@ -0,0 +1,221 @@
+use super::shared::{
+    build_provider_quota_execution_plan, build_quota_snapshot_payload,
+    default_provider_quota_execution_timeouts, execute_provider_quota_plan,
+    extract_execution_error_message, oauth_refresh_auto_removed_result,
+    persist_provider_quota_refresh_state, quota_key_auto_removed,
+    quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
+};
+use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
+use crate::GatewayError;
+use aether_admin::provider::quota::parse_claude_code_oauth_usage_response;
+use aether_admin::provider::redaction::admin_provider_metadata_bucket_safe_json;
+use aether_contracts::ProxySnapshot;
+use aether_data_contracts::repository::provider_catalog::{
+    StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
+};
+use aether_provider_pool::build_claude_code_pool_quota_request;
+use serde_json::json;
+use std::time::{SystemTime, UNIX_EPOCH};
+
+async fn execute_claude_code_quota_plan(
+    state: &AdminAppState<'_>,
+    transport: &AdminGatewayProviderTransportSnapshot,
+    authorization: (String, String),
+    proxy_override: Option<&ProxySnapshot>,
+) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
+    let proxy = match proxy_override {
+        Some(proxy) => Some(proxy.clone()),
+        None => {
+            state
+                .resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
+                .await
+        }
+    };
+    let timeouts = state
+        .resolve_transport_execution_timeouts(transport)
+        .or(Some(default_provider_quota_execution_timeouts(
+            proxy.as_ref(),
+        )));
+    let spec = build_claude_code_pool_quota_request(&transport.key.id, authorization);
+    let plan = build_provider_quota_execution_plan(
+        transport,
+        spec,
+        proxy,
+        state.resolve_transport_profile(transport),
+        timeouts,
+    );
+
+    execute_provider_quota_plan(state, transport, plan, "claude_code").await
+}
+
+pub(crate) async fn refresh_claude_code_provider_quota_locally(
+    state: &AdminAppState<'_>,
+    provider: &StoredProviderCatalogProvider,
+    endpoint: &StoredProviderCatalogEndpoint,
+    keys: Vec<StoredProviderCatalogKey>,
+    proxy_override: Option<ProxySnapshot>,
+) -> Result<Option<serde_json::Value>, GatewayError> {
+    let mut results = Vec::new();
+    let mut success_count = 0usize;
+    let mut failed_count = 0usize;
+    let mut auto_removed_count = 0usize;
+
+    for key in keys {
+        let transport = match state
+            .read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
+            .await?
+        {
+            Some(transport) => transport,
+            None => {
+                failed_count += 1;
+                results.push(json!({
+                    "key_id": key.id,
+                    "key_name": key.name,
+                    "status": "error",
+                    "message": "Provider transport snapshot unavailable",
+                }));
+                continue;
+            }
+        };
+
+        let authorization = match state.resolve_local_oauth_header_auth(&transport).await? {
+            Some(auth) => auth,
+            _ => {
+                if quota_key_auto_removed(state, &key.id).await? {
+                    auto_removed_count += 1;
+                    results.push(oauth_refresh_auto_removed_result(&key));
+                    continue;
+                }
+                failed_count += 1;
+                results.push(json!({
+                    "key_id": key.id,
+                    "key_name": key.name,
+                    "status": "error",
+                    "message": "缺少 OAuth 认证信息，请先授权/刷新 Token",
+                }));
+                continue;
+            }
+        };
+
+        let result = match execute_claude_code_quota_plan(
+            state,
+            &transport,
+            authorization,
+            proxy_override.as_ref(),
+        )
+        .await?
+        {
+            ProviderQuotaExecutionOutcome::Response(result) => result,
+            ProviderQuotaExecutionOutcome::Failure(_) => {
+                failed_count += 1;
+                results.push(json!({
+                    "key_id": key.id,
+                    "key_name": key.name,
+                    "status": "error",
+                    "message": "oauth/usage 请求执行失败",
+                    "status_code": 502,
+                }));
+                continue;
+            }
+        };
+
+        let now_unix_secs = SystemTime::now()
+            .duration_since(UNIX_EPOCH)
+            .ok()
+            .map(|duration| duration.as_secs())
+            .unwrap_or(0);
+        let mut metadata_update = None::<serde_json::Value>;
+        let (oauth_invalid_at_unix_secs, oauth_invalid_reason) =
+            quota_refresh_success_invalid_state(&key);
+        let mut status = "error".to_string();
+        let mut message = None::<String>;
+
+        if result.status_code == 200 {
+            if let Some(body_json) = result
+                .body
+                .as_ref()
+                .and_then(|body| body.json_body.as_ref())
+            {
+                metadata_update = parse_claude_code_oauth_usage_response(body_json, now_unix_secs)
+                    .map(|metadata| json!({ "claude_code": metadata }));
+                if metadata_update.is_some() {
+                    status = "success".to_string();
+                } else {
+                    status = "no_metadata".to_string();
+                    message = Some("响应中未包含额度窗口".to_string());
+                }
+            } else {
+                status = "no_metadata".to_string();
+                message = Some("响应中未包含配额信息".to_string());
+            }
+        } else {
+            message = Some(match result.status_code {
+                401 => "oauth/usage 返回 401，Token 可能已失效，请刷新 Token".to_string(),
+                403 => "oauth/usage 返回 403，该账号缺少 user:profile 权限（如 Setup Token），无法查询额度"
+                    .to_string(),
+                429 => "oauth/usage 被限流，请稍后重试".to_string(),
+                code => format!("oauth/usage 返回状态码 {code}"),
+            });
+        }
+
+        if !persist_provider_quota_refresh_state(
+            state,
+            &key.id,
+            metadata_update.as_ref(),
+            oauth_invalid_at_unix_secs,
+            oauth_invalid_reason,
+            None,
+        )
+        .await?
+        {
+            failed_count += 1;
+            results.push(json!({
+                "key_id": key.id,
+                "key_name": key.name,
+                "status": "error",
+                "message": "Key 状态写入失败",
+            }));
+            continue;
+        }
+
+        if status == "success" {
+            success_count += 1;
+        } else {
+            failed_count += 1;
+        }
+
+        let mut payload = serde_json::Map::new();
+        payload.insert("key_id".to_string(), json!(key.id));
+        payload.insert("key_name".to_string(), json!(key.name));
+        payload.insert("status".to_string(), json!(status));
+        if let Some(message) = message {
+            payload.insert("message".to_string(), json!(message));
+        }
+        if let Some(metadata) = metadata_update
+            .as_ref()
+            .and_then(|value| value.get("claude_code"))
+        {
+            payload.insert(
+                "metadata".to_string(),
+                admin_provider_metadata_bucket_safe_json("claude_code", Some(metadata)),
+            );
+        }
+        if let Some(quota_snapshot) = build_quota_snapshot_payload(
+            "claude_code",
+            key.status_snapshot.as_ref(),
+            metadata_update.as_ref(),
+        ) {
+            payload.insert("quota_snapshot".to_string(), quota_snapshot);
+        }
+        results.push(serde_json::Value::Object(payload));
+    }
+
+    Ok(Some(json!({
+        "success": success_count,
+        "failed": failed_count,
+        "total": results.len(),
+        "results": results,
+        "message": format!("已处理 {} 个 Key", results.len()),
+        "auto_removed": auto_removed_count,
+    })))
+}
@@ -3,6 +3,7 @@ use std::pin::Pin;

 use super::antigravity::refresh_antigravity_provider_quota_locally;
 use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
+use super::claude_code::refresh_claude_code_provider_quota_locally;
 use super::codex::refresh_codex_provider_quota_locally;
 use super::gemini_cli::refresh_gemini_cli_provider_quota_locally;
 use super::grok::refresh_grok_provider_quota_locally;
@@ -36,6 +37,10 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
        "chatgpt_web",
        refresh_chatgpt_web_provider_quota_locally_boxed,
    ),
+    (
+        "claude_code",
+        refresh_claude_code_provider_quota_locally_boxed,
+    ),
    ("codex", refresh_codex_provider_quota_locally_boxed),
    (
        "gemini_cli",
@@ -113,6 +118,22 @@ fn refresh_codex_provider_quota_locally_boxed<'a>(
    ))
 }

+fn refresh_claude_code_provider_quota_locally_boxed<'a>(
+    state: &'a AdminAppState<'a>,
+    provider: &'a StoredProviderCatalogProvider,
+    endpoint: &'a StoredProviderCatalogEndpoint,
+    keys: Vec<StoredProviderCatalogKey>,
+    proxy_override: Option<ProxySnapshot>,
+) -> ProviderQuotaRefreshFuture<'a> {
+    Box::pin(refresh_claude_code_provider_quota_locally(
+        state,
+        provider,
+        endpoint,
+        keys,
+        proxy_override,
+    ))
+}
+
 fn refresh_gemini_cli_provider_quota_locally_boxed<'a>(
    state: &'a AdminAppState<'a>,
    provider: &'a StoredProviderCatalogProvider,
@@ -1,5 +1,6 @@
 pub(crate) mod antigravity;
 pub(crate) mod chatgpt_web;
+pub(crate) mod claude_code;
 pub(crate) mod codex;
 pub(crate) mod dispatch;
 pub(crate) mod gemini_cli;
@@ -1713,6 +1713,7 @@ fn provider_quota_url_has_allowed_origin(provider_name: &str, value: &str) -> bo
                | "daily-cloudcode-pa.sandbox.googleapis.com"
        ),
        "gemini_cli" => host == "cloudcode-pa.googleapis.com",
+        "claude_code" => host == "api.anthropic.com",
        "chatgpt_web" | "codex" => host == "chatgpt.com",
        "grok" => host == "grok.com",
        "xai" => host == "cli-chat-proxy.grok.com",
@@ -2294,6 +2294,107 @@ fn build_gemini_cli_quota_status_snapshot(
    }))
 }

+fn build_claude_code_quota_status_snapshot(
+    upstream_metadata: Option<&Value>,
+    source: &str,
+) -> Option<Value> {
+    let metadata = provider_quota_metadata_bucket(upstream_metadata, "claude_code")?;
+    let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
+    // (metadata prefix, window code, window minutes, account-wide?). Display labels are
+    // resolved by the frontend from `code` so they follow the UI locale.
+    let definitions: [(&str, &str, u64, bool); 4] = [
+        ("five_hour", "5h", 300, true),
+        ("seven_day", "weekly", 10_080, true),
+        ("seven_day_sonnet", "weekly_sonnet", 10_080, false),
+        ("seven_day_fable", "weekly_fable", 10_080, false),
+    ];
+    let mut windows = Vec::new();
+    for (prefix, code, window_minutes, account_wide) in definitions {
+        let used_percent = metadata
+            .get(&format!("{prefix}_used_percent"))
+            .and_then(Value::as_f64);
+        let reset_at =
+            provider_quota_timestamp_unix_secs(metadata.get(&format!("{prefix}_reset_at")));
+        // A window whose reset time already passed no longer describes current usage.
+        let expired = reset_at
+            .zip(observed_at_unix_secs)
+            .is_some_and(|(reset_at, observed_at)| reset_at <= observed_at);
+        let Some(used_percent) = used_percent else {
+            continue;
+        };
+        let used_ratio = if expired {
+            0.0
+        } else {
+            (used_percent / 100.0).clamp(0.0, 1.0)
+        };
+        let reset_seconds = reset_at
+            .zip(observed_at_unix_secs)
+            .map(|(reset_at, observed_at)| reset_at.saturating_sub(observed_at));
+        let mut window = json!({
+            "code": code,
+            "scope": if account_wide { "account" } else { "model" },
+            "unit": "percent",
+            "used_ratio": used_ratio,
+            "remaining_ratio": 1.0 - used_ratio,
+            "reset_at": reset_at,
+            "reset_seconds": reset_seconds,
+            "window_minutes": window_minutes,
+            "is_exhausted": used_ratio >= 1.0 - 1e-6,
+        });
+        if !account_wide {
+            window["quota_group"] = json!(code);
+        }
+        windows.push(window);
+    }
+    if windows.is_empty() {
+        return None;
+    }
+
+    let account_windows = windows
+        .iter()
+        .filter(|window| window.get("scope").and_then(Value::as_str) == Some("account"))
+        .cloned()
+        .collect::<Vec<_>>();
+    let blocking_windows = account_windows
+        .iter()
+        .filter(|window| window.get("is_exhausted").and_then(Value::as_bool) == Some(true))
+        .cloned()
+        .collect::<Vec<_>>();
+    let exhausted = !blocking_windows.is_empty();
+    // The account is usable again only once every exhausted window resets.
+    let reset_at = if exhausted {
+        blocking_windows
+            .iter()
+            .filter_map(|window| provider_quota_timestamp_unix_secs(window.get("reset_at")))
+            .max()
+    } else {
+        None
+    };
+    let reset_seconds = if exhausted {
+        blocking_windows
+            .iter()
+            .filter_map(|window| window.get("reset_seconds").and_then(Value::as_u64))
+            .max()
+    } else {
+        None
+    };
+
+    Some(json!({
+        "version": 2,
+        "provider_type": "claude_code",
+        "code": if exhausted { "exhausted" } else { "ok" },
+        "freshness": "fresh",
+        "source": source,
+        "observed_at": observed_at_unix_secs,
+        "exhausted": exhausted,
+        "usage_ratio": quota_windows_usage_ratio(&account_windows),
+        "updated_at": observed_at_unix_secs,
+        "reset_at": reset_at,
+        "reset_seconds": reset_seconds,
+        "windows": windows,
+    }))
+}
+
 fn build_codex_reset_credits_status_snapshot(
    metadata: &Map<String, Value>,
    observed_at_unix_secs: Option<u64>,
@@ -2423,6 +2524,7 @@ pub(crate) fn sync_provider_key_quota_status_snapshot(
        "antigravity" => build_antigravity_quota_status_snapshot(upstream_metadata, source),
        "grok" => build_grok_quota_status_snapshot(upstream_metadata, source),
        "gemini_cli" => build_gemini_cli_quota_status_snapshot(upstream_metadata, source),
+        "claude_code" => build_claude_code_quota_status_snapshot(upstream_metadata, source),
        _ => None,
    }?;
    if normalized_provider_type == "codex" {
@@ -3737,6 +3839,44 @@ mod tests {
        assert_eq!(window.get("used_value"), Some(&json!(0.0)));
    }

+    #[test]
+    fn provider_key_status_snapshot_payload_backfills_claude_code_usage_windows() {
+        let mut key = sample_catalog_key();
+        key.upstream_metadata = Some(json!({
+            "claude_code": {
+                "updated_at": 1_800_000_000u64,
+                "five_hour_used_percent": 100.0,
+                "five_hour_reset_at": 1_800_003_600u64,
+                "seven_day_used_percent": 40.0,
+                "seven_day_reset_at": 1_800_400_000u64,
+                "seven_day_sonnet_used_percent": 10.0,
+                "seven_day_sonnet_reset_at": 1_800_400_000u64
+            }
+        }));
+
+        let payload = provider_key_status_snapshot_payload(&key, "claude_code");
+        let quota = payload
+            .get("quota")
+            .and_then(Value::as_object)
+            .expect("quota snapshot should be object");
+        assert_eq!(quota.get("provider_type"), Some(&json!("claude_code")));
+        // An exhausted 5h window blocks the whole account until it resets.
+        assert_eq!(quota.get("exhausted"), Some(&json!(true)));
+        assert_eq!(quota.get("reset_at"), Some(&json!(1_800_003_600u64)));
+        let windows = quota
+            .get("windows")
+            .and_then(Value::as_array)
+            .expect("windows should exist");
+        assert_eq!(windows.len(), 3);
+        assert_eq!(windows[0]["code"], json!("5h"));
+        assert_eq!(windows[0]["scope"], json!("account"));
+        assert_eq!(windows[0]["window_minutes"], json!(300));
+        assert_eq!(windows[1]["code"], json!("weekly"));
+        assert_eq!(windows[1]["used_ratio"], json!(0.4));
+        assert_eq!(windows[2]["code"], json!("weekly_sonnet"));
+        assert_eq!(windows[2]["scope"], json!("model"));
+    }
+
    #[test]
    fn provider_key_status_snapshot_payload_backfills_grok_model_quota() {
        let mut key = sample_catalog_key();