i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-08 10:27:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #871 from dalamudx/feat/claude-code-usage-quota

Browse Source 
feat(claude-code): 支持查询 Claude Code 账号 5H/周额度并在号池和提供商详情展示
This commit is contained in:
ZheFox
2026-09-30 09:00:53 +08:00
committed by GitHub
parent d068b62966 017d08d45e
commit 54fbcc25a1
22 changed files with 1334 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/claude_code.rs
+221
View File
@@ -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,
})))
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/dispatch.rs
+21
View File
@@ -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,
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/mod.rs
+1
View File
@@ -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;
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+1
View File
@@ -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",
apps/aether-gateway/src/handlers/shared/catalog.rs
+159
View File
@@ -2294,6 +2294,111 @@ 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,
"reset_credits": build_codex_reset_credits_status_snapshot(
metadata,
observed_at_unix_secs,
),
"windows": windows,
}))
}
fn build_codex_reset_credits_status_snapshot(
metadata: &Map<String, Value>,
observed_at_unix_secs: Option<u64>,
@@ -2423,6 +2528,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 +3843,59 @@ 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,
"reset_credits": {
"available_count": 2,
"updated_at": 1_800_000_000u64,
"detail_source": "claude_oauth_usage",
"credits": [{
"display_key": "Key-1",
"status": "available",
"expires_at": 1_800_144_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"));
assert_eq!(quota["reset_credits"]["available_count"], json!(2));
assert_eq!(
quota["reset_credits"]["credits"][0]["remaining_seconds"],
json!(144_000u64)
);
}
#[test]
fn provider_key_status_snapshot_payload_backfills_grok_model_quota() {
let mut key = sample_catalog_key();
apps/aether-gateway/src/orchestration/report_effects.rs
+164
View File
@@ -551,6 +551,24 @@ async fn sync_grok_quota_from_report_context(
async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncReportRequest) {
apply_local_gemini_file_mapping_report_effect(state, payload).await;
if claude_code_quota_headers_reportable(payload.status_code) {
if let Err(err) = sync_claude_code_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "claude_code_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist claude_code realtime quota from sync response headers"
);
}
}
if (200..300).contains(&payload.status_code) {
if let Err(err) = sync_codex_quota_from_response_headers(
state,
@@ -641,6 +659,24 @@ async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStr
);
}
}
if claude_code_quota_headers_reportable(payload.status_code) {
if let Err(err) = sync_claude_code_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "claude_code_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist claude_code realtime quota from stream response headers"
);
}
}
if let Err(err) = sync_grok_quota_from_report_context(
state,
payload.report_context.as_ref(),
@@ -896,6 +932,134 @@ async fn sync_codex_quota_from_response_headers(
.await
}
fn claude_code_quota_headers_reportable(status_code: u16) -> bool {
// Real limit 429s carry the unified headers (fingerprint-rejection 429s do not, and then
// the parser finds nothing to record).
(200..300).contains(&status_code) || status_code == 429
}
/// Passive sampling of Anthropic's `anthropic-ratelimit-unified-*` response headers into the
/// `claude_code` quota metadata, so the 5H / weekly windows stay fresh between active refreshes.
async fn sync_claude_code_quota_from_response_headers(
state: &AppState,
report_context: Option<&Value>,
headers: &BTreeMap<String, String>,
) -> Result<bool, GatewayError> {
let observed_at_unix_secs = report_context_u64(
report_context,
"provider_response_headers_observed_at_unix_ms",
)
.map(|value| value / 1_000)
.filter(|value| *value > 0)
.unwrap_or_else(current_unix_secs);
let parsed = report_context_provider_response_headers(report_context)
.and_then(|headers| {
admin_provider_quota_pure::parse_claude_code_usage_headers(
&headers,
observed_at_unix_secs,
)
})
.or_else(|| {
admin_provider_quota_pure::parse_claude_code_usage_headers(
headers,
observed_at_unix_secs,
)
});
let Some(parsed) = parsed else {
return Ok(false);
};
let Some(key_id) = report_context_key_id(report_context) else {
return Ok(false);
};
for attempt in 0..RUNTIME_METADATA_CAS_MAX_ATTEMPTS {
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if !provider
.provider_type
.trim()
.eq_ignore_ascii_case("claude_code")
{
return Ok(false);
}
let expected_namespace_value =
upstream_metadata_namespace_value(key.upstream_metadata.as_ref(), "claude_code");
let mut bucket = expected_namespace_value
.as_ref()
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
// Never let an older observation overwrite a newer active refresh.
if bucket
.get("updated_at")
.and_then(admin_provider_quota_pure::coerce_json_u64)
.is_some_and(|stored| stored > observed_at_unix_secs)
{
return Ok(false);
}
let Some(patch) = parsed.as_object() else {
return Ok(false);
};
// Headers can describe only some windows; absent windows keep their stored value.
for (field, value) in patch {
bucket.insert(field.clone(), value.clone());
}
let next_bucket = Value::Object(bucket);
if expected_namespace_value.as_ref() == Some(&next_bucket) {
return Ok(false);
}
let updated_upstream_metadata = merge_metadata_object(
key.upstream_metadata.as_ref(),
"claude_code",
next_bucket.clone(),
);
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"response_headers",
);
let updated = state
.update_provider_catalog_key_runtime_metadata(
&ProviderCatalogKeyRuntimeMetadataUpdate {
key_id: key_id.clone(),
namespace: "claude_code".to_string(),
expected_upstream_metadata_value: expected_namespace_value,
upstream_metadata_value: next_bucket,
status_snapshot_patch: quota_status_snapshot_patch(
updated_status_snapshot.as_ref(),
),
updated_at_unix_secs: Some(observed_at_unix_secs),
},
)
.await?;
if updated {
return Ok(true);
}
if attempt + 1 < RUNTIME_METADATA_CAS_MAX_ATTEMPTS {
let backoff_us = 50_u64.saturating_mul((attempt + 1) as u64).min(1_000);
tokio::time::sleep(Duration::from_micros(backoff_us)).await;
}
}
Ok(false)
}
async fn sync_codex_websocket_quota_from_stream_payload(
state: &AppState,
payload: &GatewayStreamReportRequest,
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+10 -18
View File
@@ -2041,24 +2041,16 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_gemini_cli_with_trus
#[tokio::test]
async fn gateway_refresh_quota_reconciles_unsupported_fixed_provider_endpoints_before_clear_message(
) {
let cases = [
(
"provider-claude-code-reconcile",
"claude_code",
1usize,
"claude:messages",
"https://api.anthropic.com/v1",
"Claude Code 暂不支持自动刷新额度",
),
(
"provider-vertex-ai-reconcile",
"vertex_ai",
2usize,
"gemini:generate_content",
"https://aiplatform.googleapis.com",
"Vertex AI 暂不支持自动刷新额度",
),
];
// Claude Code supports quota refresh now, so Vertex AI is the remaining fixed provider
// whose refresh is unsupported.
let cases = [(
"provider-vertex-ai-reconcile",
"vertex_ai",
2usize,
"gemini:generate_content",
"https://aiplatform.googleapis.com",
"Vertex AI 暂不支持自动刷新额度",
)];
let providers = cases
.iter()
apps/aether-gateway/src/tests/control/admin/oauth.rs
+13 -1
View File
@@ -410,6 +410,10 @@ async fn gateway_authorizes_claude_cookie_without_persisting_cookie_impl() {
"email_address": "[email protected]"
}
}),
// Newly authorized accounts get their 5H/weekly quota fetched right away.
quota if quota.starts_with("claude-code-quota:") => {
json!({"five_hour": {"utilization": 10.0}})
}
unexpected => panic!("unexpected execution plan: {unexpected}"),
};
Json(json!({
@@ -811,7 +815,15 @@ async fn gateway_batch_authorizes_claude_cookies_as_redacted_task_impl() {
}
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 6);
// Post-authorization quota refreshes are fire-and-forget, so their count is not
// deterministic here; only the OAuth flow plans are asserted.
assert_eq!(
plans
.iter()
.filter(|plan| !plan.request_id.starts_with("claude-code-quota:"))
.count(),
6
);
assert_eq!(
plans
.iter()