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

This commit is contained in:
dalamudx
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",