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

Unify quota snapshots and oauth refresh handling

Browse Source
This commit is contained in:
fawney19
2026-04-17 18:22:41 +08:00
parent b8702ae124
commit 7eae1f90f6
38 changed files with 3435 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

126
apps/aether-gateway/src/execution_runtime/fallback.rs
View File

@@ -15,8 +15,15 @@ fn local_candidate_index(report_context: Option<&serde_json::Value>) -> Option<u
.and_then(serde_json::Value::as_u64)
}
fn is_retryable_local_upstream_status(status_code: u16) -> bool {
status_code == 429 || status_code >= 500
fn should_failover_local_upstream_status(status_code: u16) -> bool {
status_code >= 400
}
fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
matches!(
plan_kind,
"openai_video_delete_sync" | "openai_video_cancel_sync" | "gemini_video_cancel_sync"
)
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
@@ -54,11 +61,14 @@ impl LocalFailoverDecision {
pub(crate) async fn should_retry_next_local_candidate_sync(
state: &AppState,
plan: &ExecutionPlan,
_plan_kind: &str,
plan_kind: &str,
report_context: Option<&serde_json::Value>,
result: &ExecutionResult,
response_text: Option<&str>,
) -> bool {
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
return false;
}
matches!(
resolve_local_failover_decision(
state,
@@ -75,11 +85,14 @@ pub(crate) async fn should_retry_next_local_candidate_sync(
pub(crate) async fn should_stop_local_candidate_failover_sync(
state: &AppState,
plan: &ExecutionPlan,
_plan_kind: &str,
plan_kind: &str,
report_context: Option<&serde_json::Value>,
result: &ExecutionResult,
response_text: Option<&str>,
) -> bool {
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
return false;
}
matches!(
resolve_local_failover_decision(
state,
@@ -294,7 +307,7 @@ async fn resolve_local_failover_decision(
return LocalFailoverDecision::RetryNextCandidate;
}
if is_retryable_local_upstream_status(status_code) {
if should_failover_local_upstream_status(status_code) {
return LocalFailoverDecision::RetryNextCandidate;
}
@@ -850,6 +863,87 @@ mod tests {
);
}
#[tokio::test]
async fn sync_retry_next_candidate_treats_client_error_as_failover_by_default() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 401,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(None);
let plan = sample_plan();
assert!(
should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_chat_sync",
Some(&local_report_context),
&result,
Some("{\"error\":{\"message\":\"invalid auth token\"}}"),
)
.await
);
}
#[tokio::test]
async fn sync_retry_next_candidate_skips_video_follow_up_plan_kinds() {
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 404,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(None);
let plan = sample_plan();
for plan_kind in [
"openai_video_delete_sync",
"openai_video_cancel_sync",
"gemini_video_cancel_sync",
] {
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
plan_kind,
Some(&local_report_context),
&result,
None,
)
.await,
"{plan_kind} should not retry local failover candidates"
);
assert!(
!should_stop_local_candidate_failover_sync(
&state,
&plan,
plan_kind,
Some(&local_report_context),
&result,
None,
)
.await,
"{plan_kind} should not use local failover stop decisions"
);
}
}
#[tokio::test]
async fn stream_retry_next_candidate_requires_local_candidate_context() {
let local_report_context = serde_json::json!({
@@ -927,6 +1021,28 @@ mod tests {
);
}
#[tokio::test]
async fn stream_retry_next_candidate_treats_client_error_as_failover_by_default() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(None);
let plan = sample_plan();
assert!(
should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_chat_stream",
Some(&local_report_context),
403,
Some("{\"error\":{\"message\":\"invalid auth token\"}}"),
)
.await
);
}
#[test]
fn resolve_local_failover_policy_reads_provider_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({

11
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/quota.rs
View File

@@ -195,10 +195,15 @@ pub(super) async fn maybe_handle(
}
let Some(payload) = (match normalized_provider_type.as_str() {
"codex" => refresh_codex_provider_quota_locally(state, &provider, &endpoint, keys).await?,
"kiro" => refresh_kiro_provider_quota_locally(state, &provider, &endpoint, keys).await?,
"codex" => {
refresh_codex_provider_quota_locally(state, &provider, &endpoint, keys, None).await?
}
"kiro" => {
refresh_kiro_provider_quota_locally(state, &provider, &endpoint, keys, None).await?
}
"antigravity" => {
refresh_antigravity_provider_quota_locally(state, &provider, &endpoint, keys).await?
refresh_antigravity_provider_quota_locally(state, &provider, &endpoint, keys, None)
.await?
}
_ => None,
}) else {

12
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
View File

@@ -11,7 +11,8 @@ use crate::handlers::admin::provider::oauth::provisioning::{
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
@@ -264,9 +265,12 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
}
};
let _ =
refresh_provider_oauth_account_state_after_update(state, &provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
success += 1;
results.push(json!({

12
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs
View File

@@ -5,7 +5,8 @@ use crate::handlers::admin::provider::oauth::provisioning::{
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use crate::handlers::admin::provider::oauth::state::decode_jwt_claims;
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
@@ -548,9 +549,12 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
.get("auth_method")
.cloned()
.unwrap_or(serde_json::Value::Null);
let _ =
refresh_provider_oauth_account_state_after_update(state, &provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
success += 1;
results.push(json!({

3
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
View File

@@ -148,7 +148,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
request_proxy,
request_proxy.clone(),
)
.await
{
@@ -243,6 +243,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
&provider,
&endpoint,
vec![refreshed_key],
request_proxy.clone(),
)
.await?
{

14
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/provider.rs
View File

@@ -5,7 +5,8 @@ use super::super::super::provisioning::{
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use super::super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
@@ -137,7 +138,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
request_proxy,
request_proxy.clone(),
)
.await
{
@@ -231,9 +232,12 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
}
};
let _ = state
.refresh_provider_oauth_account_state_after_update(&provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
Ok(Json(json!({
"key_id": persisted_key.id,

14
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/poll.rs
View File

@@ -8,7 +8,8 @@ use crate::handlers::admin::provider::oauth::provisioning::{
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response, build_kiro_device_key_name,
@@ -149,7 +150,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&session.client_id,
&session.client_secret,
&session.device_code,
request_proxy,
request_proxy.clone(),
)
.await
{
@@ -315,9 +316,12 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
}
};
let _ = state
.refresh_provider_oauth_account_state_after_update(&provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
session.status = "authorized".to_string();
session.key_id = Some(persisted_key.id.clone());

18
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
View File

@@ -5,7 +5,8 @@ use super::super::provisioning::{
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
@@ -120,7 +121,11 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
.await;
let token_payload = match state
.exchange_admin_provider_oauth_refresh_token(template, refresh_token_input, request_proxy)
.exchange_admin_provider_oauth_refresh_token(
template,
refresh_token_input,
request_proxy.clone(),
)
.await
{
Ok(payload) => payload,
@@ -218,9 +223,12 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
}
};
let _ = state
.refresh_provider_oauth_account_state_after_update(&provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
Ok(Json(json!({
"key_id": persisted_key.id,

2
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/execution.rs
View File

@@ -94,7 +94,7 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
refreshed_key.encrypted_auth_config.as_deref(),
);
let (account_state_recheck_attempted, account_state_recheck_error) = state
.refresh_provider_oauth_account_state_after_update(&provider, &key_id)
.refresh_provider_oauth_account_state_after_update(&provider, &key_id, None)
.await?;
Ok(RefreshDispatch::Continue(RefreshSuccessContext {

31
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
View File

@@ -1,14 +1,14 @@
use super::shared::{
coerce_json_f64, coerce_json_string, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
build_quota_snapshot_payload, coerce_json_f64, coerce_json_string,
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_antigravity_usage_response;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -22,6 +22,7 @@ async fn execute_antigravity_quota_plan(
authorization: (String, String),
project_id: &str,
mut identity_headers: BTreeMap<String, String>,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let mut headers = std::mem::take(&mut identity_headers);
headers.insert("authorization".to_string(), authorization.1);
@@ -32,9 +33,14 @@ async fn execute_antigravity_quota_plan(
.or_insert_with(|| "antigravity".to_string());
let body = json!({ "project": project_id });
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
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(
@@ -78,6 +84,7 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
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;
@@ -134,6 +141,7 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
authorization,
&project_id,
identity_headers,
proxy_override.as_ref(),
)
.await?
{
@@ -250,6 +258,13 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
{
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"antigravity",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
results.push(serde_json::Value::Object(payload));
}

47
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
View File

@@ -13,13 +13,15 @@ use self::parse::{
};
use self::plan::{build_codex_refresh_headers, execute_codex_quota_plan};
use super::shared::{
extract_execution_error_message, persist_provider_quota_refresh_state,
provider_auto_remove_banned_keys, quota_refresh_success_invalid_state,
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
quota_refresh_success_invalid_state, should_auto_remove_structured_reason,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -31,6 +33,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let auto_remove_abnormal_keys = provider_auto_remove_banned_keys(provider.config.as_ref());
let mut results = Vec::new();
@@ -77,20 +80,23 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
let result = match execute_codex_quota_plan(state, &transport, headers).await? {
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("wham/usage 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let result =
match execute_codex_quota_plan(state, &transport, headers, proxy_override.as_ref())
.await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("wham/usage 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
@@ -278,6 +284,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
{
payload.insert("metadata".to_string(), metadata_update);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"codex",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
if auto_removed {
payload.insert("auto_removed".to_string(), json!(true));
}

14
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs
View File

@@ -6,7 +6,7 @@ use super::parse::normalize_codex_plan_type;
use crate::handlers::admin::provider::shared::payloads::CODEX_WHAM_USAGE_URL;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use std::collections::BTreeMap;
pub(super) fn build_codex_refresh_headers(
@@ -60,10 +60,16 @@ pub(super) async fn execute_codex_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
headers: BTreeMap<String, String>,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
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(

64
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/mod.rs
View File

@@ -4,11 +4,13 @@ mod plan;
use self::parse::parse_kiro_usage_response;
use self::plan::execute_kiro_quota_plan;
use super::shared::{
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -20,6 +22,7 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
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;
@@ -57,20 +60,22 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
continue;
};
let result = match execute_kiro_quota_plan(state, &transport, &auth).await? {
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("getUsageLimits 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let result =
match execute_kiro_quota_plan(state, &transport, &auth, proxy_override.as_ref()).await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("getUsageLimits 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
@@ -93,7 +98,25 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
metadata_update = parse_kiro_usage_response(body_json, now_unix_secs)
.map(|metadata| json!({ "kiro": metadata }));
if metadata_update.is_some() {
let auth_config_json = auth.auth_config.to_json_value().to_string();
let mut auth_config_object = transport
.key
.decrypted_auth_config
.as_deref()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok())
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
if let Some(refreshed_auth_config) =
auth.auth_config.to_json_value().as_object()
{
for (key, value) in refreshed_auth_config {
auth_config_object.insert(key.clone(), value.clone());
}
}
auth_config_object
.entry("provider_type".to_string())
.or_insert_with(|| json!("kiro"));
let auth_config_json =
serde_json::Value::Object(auth_config_object).to_string();
if let Some(auth_config_json) =
state.encrypt_catalog_secret_with_fallbacks(auth_config_json.as_str())
{
@@ -185,6 +208,13 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
{
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"kiro",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
results.push(serde_json::Value::Object(payload));
}

14
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/plan.rs
View File

@@ -7,7 +7,7 @@ use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use std::collections::BTreeMap;
use url::form_urlencoded;
use uuid::Uuid;
@@ -66,10 +66,16 @@ pub(super) async fn execute_kiro_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
auth: &AdminKiroRequestAuth,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
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(

29
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
View File

@@ -2,6 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::GatewayError;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot};
@@ -89,6 +90,20 @@ pub(super) fn coerce_json_string(value: Option<&serde_json::Value>) -> Option<St
admin_provider_quota_pure::coerce_json_string(value)
}
pub(super) fn build_quota_snapshot_payload(
provider_type: &str,
current_status_snapshot: Option<&serde_json::Value>,
metadata_update: Option<&serde_json::Value>,
) -> Option<serde_json::Value> {
let updated_snapshot = sync_provider_key_quota_status_snapshot(
current_status_snapshot,
provider_type,
metadata_update,
"refresh_api",
)?;
updated_snapshot.get("quota").cloned()
}
pub(crate) async fn persist_provider_quota_refresh_state(
state: &AdminAppState<'_>,
key_id: &str,
@@ -106,17 +121,31 @@ pub(crate) async fn persist_provider_quota_refresh_state(
return Ok(false);
};
let mut quota_snapshot_provider_type = None::<&str>;
if let Some(metadata_update) = metadata_update {
latest_key.upstream_metadata = Some(merge_upstream_metadata(
latest_key.upstream_metadata.as_ref(),
metadata_update,
));
quota_snapshot_provider_type = metadata_update.as_object().and_then(|object| {
["codex", "kiro", "antigravity", "gemini_cli"]
.into_iter()
.find(|provider_type| object.contains_key(*provider_type))
});
}
if let Some(encrypted_auth_config) = encrypted_auth_config {
latest_key.encrypted_auth_config = Some(encrypted_auth_config);
}
latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
latest_key.oauth_invalid_reason = oauth_invalid_reason;
if let Some(provider_type) = quota_snapshot_provider_type {
latest_key.status_snapshot = sync_provider_key_quota_status_snapshot(
latest_key.status_snapshot.as_ref(),
provider_type,
latest_key.upstream_metadata.as_ref(),
"refresh_api",
);
}
latest_key.updated_at_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()

50
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
View File

@@ -3,7 +3,8 @@ use super::quota::codex::refresh_codex_provider_quota_locally;
use super::quota::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use crate::{AppState, GatewayError};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
@@ -65,6 +66,7 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
key_id: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<(bool, Option<String>), GatewayError> {
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if !matches!(provider_type.as_str(), "codex" | "kiro" | "antigravity") {
@@ -90,16 +92,37 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
return Ok((false, None));
}
let proxy_override = proxy_override.cloned();
let payload = match provider_type.as_str() {
"codex" => {
refresh_codex_provider_quota_locally(state, provider, &endpoint, vec![key]).await?
refresh_codex_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override.clone(),
)
.await?
}
"kiro" => {
refresh_kiro_provider_quota_locally(state, provider, &endpoint, vec![key]).await?
refresh_kiro_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override.clone(),
)
.await?
}
"antigravity" => {
refresh_antigravity_provider_quota_locally(state, provider, &endpoint, vec![key])
.await?
refresh_antigravity_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override,
)
.await?
}
_ => None,
};
@@ -123,3 +146,20 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
};
Ok((true, error))
}
pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update(
app: AppState,
provider: StoredProviderCatalogProvider,
key_id: String,
proxy_override: Option<ProxySnapshot>,
) {
tokio::spawn(async move {
let _ = refresh_provider_oauth_account_state_after_update(
&AdminAppState::new(&app),
&provider,
&key_id,
proxy_override.as_ref(),
)
.await;
});
}

467
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
View File

@@ -5,6 +5,7 @@ use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
use crate::provider_key_auth::provider_key_auth_semantics;
use aether_admin::provider::pool as admin_provider_pool_pure;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
@@ -238,88 +239,135 @@ fn admin_pool_format_reset_after(seconds: f64) -> Option<String> {
Some("即将重置".to_string())
}
fn admin_pool_build_codex_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
fn admin_pool_quota_snapshot_matches_provider(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
provider_type: &str,
) -> bool {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
match quota_snapshot
.get("provider_type")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
Some(provider_type) => provider_type.eq_ignore_ascii_case(&normalized_provider_type),
None => {
quota_snapshot
.get("code")
.and_then(serde_json::Value::as_str)
.is_some_and(|code| !code.trim().eq_ignore_ascii_case("unknown"))
|| quota_snapshot
.get("updated_at")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("observed_at")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("usage_ratio")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("reset_seconds")
.is_some_and(|value| !value.is_null())
|| quota_snapshot
.get("windows")
.and_then(serde_json::Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| quota_snapshot
.get("credits")
.and_then(serde_json::Value::as_object)
.is_some_and(|credits| !credits.is_empty())
}
}
}
fn admin_pool_quota_window<'a>(
quota_snapshot: &'a serde_json::Map<String, serde_json::Value>,
code: &str,
) -> Option<&'a serde_json::Map<String, serde_json::Value>> {
quota_snapshot
.get("windows")
.and_then(serde_json::Value::as_array)?
.iter()
.filter_map(serde_json::Value::as_object)
.find(|window| {
window
.get("code")
.and_then(serde_json::Value::as_str)
.is_some_and(|value| value.eq_ignore_ascii_case(code))
})
}
fn admin_pool_quota_window_used_percent(
window: &serde_json::Map<String, serde_json::Value>,
) -> Option<f64> {
admin_pool_json_to_f64(window.get("used_ratio"))
.map(|value| (value * 100.0).clamp(0.0, 100.0))
.or_else(|| {
admin_pool_json_to_f64(window.get("remaining_ratio"))
.map(|value| ((1.0 - value) * 100.0).clamp(0.0, 100.0))
})
}
fn admin_pool_quota_window_reset_seconds(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
window: &serde_json::Map<String, serde_json::Value>,
now_unix_secs: u64,
) -> Option<f64> {
if let Some(reset_at) = admin_pool_json_to_u64(window.get("reset_at")) {
return Some(reset_at.saturating_sub(now_unix_secs) as f64);
}
let remaining = admin_pool_json_to_f64(window.get("reset_seconds"))?;
let observed_at_unix_secs = admin_pool_json_to_u64(quota_snapshot.get("observed_at"))
.or_else(|| admin_pool_json_to_u64(quota_snapshot.get("updated_at")));
let elapsed = observed_at_unix_secs
.map(|observed_at| now_unix_secs.saturating_sub(observed_at) as f64)
.unwrap_or(0.0);
Some((remaining - elapsed).max(0.0))
}
fn admin_pool_codex_quota_part_from_window(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
window_code: &str,
label: &str,
now_unix_secs: u64,
) -> Option<String> {
fn codex_reset_seconds(
data: &serde_json::Map<String, serde_json::Value>,
reset_seconds_key: &str,
reset_after_seconds_key: &str,
reset_at_key: &str,
now_unix_secs: u64,
updated_at_unix_secs: Option<u64>,
) -> Option<f64> {
if let Some(reset_at) = admin_pool_json_to_u64(data.get(reset_at_key)) {
return Some(reset_at.saturating_sub(now_unix_secs) as f64);
}
let window = admin_pool_quota_window(quota_snapshot, window_code)?;
let used_percent = admin_pool_quota_window_used_percent(window)?;
let reset_seconds =
admin_pool_quota_window_reset_seconds(quota_snapshot, window, now_unix_secs);
let effective_used_percent = if reset_seconds.is_some_and(|value| value <= 0.0) {
0.0
} else {
used_percent
};
let remaining = admin_pool_json_to_f64(data.get(reset_seconds_key))
.or_else(|| admin_pool_json_to_f64(data.get(reset_after_seconds_key)))?;
let elapsed = updated_at_unix_secs
.map(|updated_at| now_unix_secs.saturating_sub(updated_at) as f64)
.unwrap_or(0.0);
Some((remaining - elapsed).max(0.0))
let mut part = format!(
"{label}剩余 {}",
admin_pool_format_percent(100.0 - effective_used_percent)
);
if admin_pool_has_quota_consumption(Some(effective_used_percent)) {
if let Some(reset_text) = reset_seconds.and_then(admin_pool_format_reset_after) {
part.push_str(&format!(" ({reset_text})"));
}
}
Some(part)
}
fn codex_effective_used_percent(used_percent: f64, reset_seconds: Option<f64>) -> f64 {
let normalized = used_percent.clamp(0.0, 100.0);
if normalized <= 1e-6 {
return 0.0;
}
if reset_seconds.is_some_and(|value| value <= 0.0) {
return 0.0;
}
normalized
}
let mut parts = Vec::new();
fn admin_pool_build_codex_account_quota_from_snapshot(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
let now_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
let updated_at_unix_secs = admin_pool_json_to_u64(data.get("updated_at"));
let mut parts = Vec::new();
let primary_used_raw = admin_pool_json_to_f64(data.get("primary_used_percent"));
if let Some(primary_used_raw) = primary_used_raw {
let primary_reset_seconds = codex_reset_seconds(
data,
"primary_reset_seconds",
"primary_reset_after_seconds",
"primary_reset_at",
now_unix_secs,
updated_at_unix_secs,
);
let primary_used = codex_effective_used_percent(primary_used_raw, primary_reset_seconds);
let mut part = format!("周剩余 {}", admin_pool_format_percent(100.0 - primary_used));
if admin_pool_has_quota_consumption(Some(primary_used)) {
if let Some(reset_text) = primary_reset_seconds.and_then(admin_pool_format_reset_after)
{
part.push_str(&format!(" ({reset_text})"));
}
}
if let Some(part) =
admin_pool_codex_quota_part_from_window(quota_snapshot, "weekly", "周", now_unix_secs)
{
parts.push(part);
}
let secondary_used_raw = admin_pool_json_to_f64(data.get("secondary_used_percent"));
if let Some(secondary_used_raw) = secondary_used_raw {
let secondary_reset_seconds = codex_reset_seconds(
data,
"secondary_reset_seconds",
"secondary_reset_after_seconds",
"secondary_reset_at",
now_unix_secs,
updated_at_unix_secs,
);
let secondary_used =
codex_effective_used_percent(secondary_used_raw, secondary_reset_seconds);
let mut part = format!(
"5H剩余 {}",
admin_pool_format_percent(100.0 - secondary_used)
);
if admin_pool_has_quota_consumption(Some(secondary_used)) {
if let Some(reset_text) =
secondary_reset_seconds.and_then(admin_pool_format_reset_after)
{
part.push_str(&format!(" ({reset_text})"));
}
}
if let Some(part) =
admin_pool_codex_quota_part_from_window(quota_snapshot, "5h", "5H", now_unix_secs)
{
parts.push(part);
}
@@ -327,11 +375,16 @@ fn admin_pool_build_codex_account_quota(
return Some(parts.join(" | "));
}
let has_credits = data
.get("has_credits")
.and_then(serde_json::Value::as_bool)
let credits = quota_snapshot
.get("credits")
.and_then(serde_json::Value::as_object);
let has_credits = credits
.and_then(|credits| credits.get("has_credits"))
.and_then(admin_provider_quota_pure::coerce_json_bool)
.unwrap_or(false);
let credits_balance = admin_pool_json_to_f64(data.get("credits_balance"));
let credits_balance = credits
.and_then(|credits| credits.get("balance"))
.and_then(admin_provider_quota_pure::coerce_json_f64);
if has_credits && credits_balance.is_some() {
return credits_balance.map(|value| format!("积分 {value:.2}"));
}
@@ -342,73 +395,115 @@ fn admin_pool_build_codex_account_quota(
None
}
fn admin_pool_build_kiro_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
fn admin_pool_quota_windows<'a>(
quota_snapshot: &'a serde_json::Map<String, serde_json::Value>,
) -> Vec<&'a serde_json::Map<String, serde_json::Value>> {
quota_snapshot
.get("windows")
.and_then(serde_json::Value::as_array)
.map(|windows| {
windows
.iter()
.filter_map(serde_json::Value::as_object)
.collect::<Vec<_>>()
})
.unwrap_or_default()
}
fn admin_pool_build_kiro_account_quota_from_snapshot(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
if data
.get("is_banned")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return Some("账号已封禁".to_string());
let code = quota_snapshot
.get("code")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.unwrap_or_default();
if code.eq_ignore_ascii_case("banned") {
return quota_snapshot
.get("label")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| Some("账号已封禁".to_string()));
}
let usage_percentage = admin_pool_json_to_f64(data.get("usage_percentage"));
if let Some(usage_percentage) = usage_percentage {
let remaining = 100.0 - usage_percentage;
let current_usage = admin_pool_json_to_f64(data.get("current_usage"));
let usage_limit = admin_pool_json_to_f64(data.get("usage_limit"));
if let (Some(current_usage), Some(usage_limit)) = (current_usage, usage_limit) {
if usage_limit > 0.0 {
let window = admin_pool_quota_windows(quota_snapshot)
.into_iter()
.next()?;
let used_ratio = admin_pool_json_to_f64(window.get("used_ratio"));
let remaining_ratio = admin_pool_json_to_f64(window.get("remaining_ratio"))
.or_else(|| used_ratio.map(|value| (1.0 - value).max(0.0)));
let used_value = admin_pool_json_to_f64(window.get("used_value"));
let remaining_value = admin_pool_json_to_f64(window.get("remaining_value"));
let limit_value = admin_pool_json_to_f64(window.get("limit_value"));
if let (Some(remaining_value), Some(limit_value)) = (remaining_value, limit_value) {
if limit_value > 0.0 && remaining_value <= 0.0 {
return Some(format!(
"剩余 {}/{}",
admin_pool_format_quota_value(remaining_value),
admin_pool_format_quota_value(limit_value),
));
}
}
if let Some(remaining_ratio) = remaining_ratio {
let remaining_percent = (remaining_ratio * 100.0).clamp(0.0, 100.0);
if let (Some(used_value), Some(limit_value)) = (used_value, limit_value) {
if limit_value > 0.0 {
return Some(format!(
"剩余 {} ({}/{})",
admin_pool_format_percent(remaining),
admin_pool_format_quota_value(current_usage),
admin_pool_format_quota_value(usage_limit),
admin_pool_format_percent(remaining_percent),
admin_pool_format_quota_value(used_value),
admin_pool_format_quota_value(limit_value),
));
}
}
return Some(format!("剩余 {}", admin_pool_format_percent(remaining)));
return Some(format!(
"剩余 {}",
admin_pool_format_percent(remaining_percent)
));
}
let remaining = admin_pool_json_to_f64(data.get("remaining"));
let usage_limit = admin_pool_json_to_f64(data.get("usage_limit"));
match (remaining, usage_limit) {
(Some(remaining), Some(usage_limit)) if usage_limit > 0.0 => Some(format!(
match (remaining_value, limit_value) {
(Some(remaining_value), Some(limit_value)) if limit_value > 0.0 => Some(format!(
"剩余 {}/{}",
admin_pool_format_quota_value(remaining),
admin_pool_format_quota_value(usage_limit),
admin_pool_format_quota_value(remaining_value),
admin_pool_format_quota_value(limit_value),
)),
_ => None,
}
}
fn admin_pool_quota_by_model(
data: &serde_json::Map<String, serde_json::Value>,
) -> Option<&serde_json::Map<String, serde_json::Value>> {
data.get("quota_by_model")?.as_object()
}
fn admin_pool_build_antigravity_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
fn admin_pool_build_antigravity_account_quota_from_snapshot(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
if data
.get("is_forbidden")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
if quota_snapshot
.get("code")
.and_then(serde_json::Value::as_str)
.is_some_and(|code| code.eq_ignore_ascii_case("forbidden"))
{
return Some("访问受限".to_string());
return quota_snapshot
.get("label")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned)
.or_else(|| Some("访问受限".to_string()));
}
let remaining_list = admin_pool_quota_by_model(data)?
.values()
.filter_map(serde_json::Value::as_object)
.filter_map(|item| {
let used_percent = admin_pool_json_to_f64(item.get("used_percent")).or_else(|| {
admin_pool_json_to_f64(item.get("remaining_fraction"))
.map(|value| (1.0 - value) * 100.0)
})?;
Some((100.0 - used_percent).clamp(0.0, 100.0))
let remaining_list = admin_pool_quota_windows(quota_snapshot)
.into_iter()
.filter(|window| {
window
.get("scope")
.and_then(serde_json::Value::as_str)
.is_some_and(|scope| scope.eq_ignore_ascii_case("model"))
})
.filter_map(|window| {
admin_pool_json_to_f64(window.get("remaining_ratio"))
.map(|value| (value * 100.0).clamp(0.0, 100.0))
.or_else(|| {
admin_pool_json_to_f64(window.get("used_ratio"))
.map(|value| ((1.0 - value) * 100.0).clamp(0.0, 100.0))
})
})
.collect::<Vec<_>>();
@@ -427,41 +522,41 @@ fn admin_pool_build_antigravity_account_quota(
))
}
fn admin_pool_gemini_reset_at(item: &serde_json::Map<String, serde_json::Value>) -> Option<i64> {
let reset_at = admin_pool_json_to_u64(item.get("reset_at"))?;
Some(reset_at as i64)
}
fn admin_pool_gemini_model_exhausted(item: &serde_json::Map<String, serde_json::Value>) -> bool {
if item
.get("is_exhausted")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
{
return true;
}
if admin_pool_json_to_f64(item.get("remaining_fraction")).is_some_and(|value| value <= 0.0) {
return true;
}
admin_pool_json_to_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0 - 1e-6)
}
fn admin_pool_build_gemini_cli_account_quota(
data: &serde_json::Map<String, serde_json::Value>,
fn admin_pool_build_gemini_cli_account_quota_from_snapshot(
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
) -> Option<String> {
let now = chrono::Utc::now().timestamp();
let mut active = admin_pool_quota_by_model(data)?
.iter()
.filter_map(|(model_name, item)| {
let item = item.as_object()?;
if !admin_pool_gemini_model_exhausted(item) {
return None;
}
let reset_at = admin_pool_gemini_reset_at(item);
let mut active = admin_pool_quota_windows(quota_snapshot)
.into_iter()
.filter(|window| {
window
.get("scope")
.and_then(serde_json::Value::as_str)
.is_some_and(|scope| scope.eq_ignore_ascii_case("model"))
})
.filter(|window| {
window
.get("is_exhausted")
.and_then(admin_provider_quota_pure::coerce_json_bool)
.or_else(|| {
admin_pool_json_to_f64(window.get("used_ratio"))
.map(|value| value >= 1.0 - 1e-6)
})
.unwrap_or(false)
})
.filter_map(|window| {
let reset_at = admin_pool_json_to_u64(window.get("reset_at")).map(|value| value as i64);
if reset_at.is_some_and(|value| value <= now) {
return None;
}
Some((model_name.as_str(), reset_at))
let label = window
.get("label")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.or_else(|| window.get("model").and_then(serde_json::Value::as_str))
.unwrap_or("模型");
Some((label.to_string(), reset_at))
})
.collect::<Vec<_>>();
@@ -470,10 +565,10 @@ fn admin_pool_build_gemini_cli_account_quota(
}
active.sort_by_key(|(_, reset_at)| reset_at.unwrap_or(i64::MAX));
let (first_model, first_reset_at) = active[0];
let (first_model, first_reset_at) = &active[0];
if active.len() == 1 {
if let Some(reset_at) = first_reset_at {
if let Some(reset_text) = admin_pool_format_reset_after((reset_at - now) as f64) {
if let Some(reset_text) = admin_pool_format_reset_after((*reset_at - now) as f64) {
return Some(format!("{first_model} 冷却中 ({reset_text})"));
}
}
@@ -481,7 +576,7 @@ fn admin_pool_build_gemini_cli_account_quota(
}
if let Some(reset_at) = first_reset_at {
if let Some(reset_text) = admin_pool_format_reset_after((reset_at - now) as f64) {
if let Some(reset_text) = admin_pool_format_reset_after((*reset_at - now) as f64) {
return Some(format!(
"{} 个模型冷却中（最早 {reset_text}）",
active.len()
@@ -493,21 +588,46 @@ fn admin_pool_build_gemini_cli_account_quota(
fn admin_pool_build_account_quota(
provider_type: &str,
upstream_metadata: Option<&serde_json::Value>,
quota_snapshot: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Option<String> {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
let upstream_metadata = upstream_metadata?.as_object()?;
let data = upstream_metadata
.get(&normalized_provider_type)?
.as_object()?;
let quota_snapshot = quota_snapshot.filter(|quota_snapshot| {
admin_pool_quota_snapshot_matches_provider(quota_snapshot, &normalized_provider_type)
})?;
match normalized_provider_type.as_str() {
"codex" => admin_pool_build_codex_account_quota(data),
"kiro" => admin_pool_build_kiro_account_quota(data),
"antigravity" => admin_pool_build_antigravity_account_quota(data),
"gemini_cli" => admin_pool_build_gemini_cli_account_quota(data),
_ => None,
"codex" => {
if let Some(account_quota) =
admin_pool_build_codex_account_quota_from_snapshot(quota_snapshot)
{
return Some(account_quota);
}
}
"kiro" => {
if let Some(account_quota) =
admin_pool_build_kiro_account_quota_from_snapshot(quota_snapshot)
{
return Some(account_quota);
}
}
"antigravity" => {
if let Some(account_quota) =
admin_pool_build_antigravity_account_quota_from_snapshot(quota_snapshot)
{
return Some(account_quota);
}
}
"gemini_cli" => {
if let Some(account_quota) =
admin_pool_build_gemini_cli_account_quota_from_snapshot(quota_snapshot)
{
return Some(account_quota);
}
}
_ => {}
}
None
}
fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
@@ -668,7 +788,7 @@ pub(super) fn build_admin_pool_key_payload(
admin_pool_derive_oauth_expires_at(provider_type, key, auth_config.as_ref());
let oauth_plan_type =
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
let status_snapshot = provider_key_status_snapshot_payload(key);
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
let account_snapshot = status_snapshot
.get("account")
.and_then(serde_json::Value::as_object);
@@ -680,6 +800,7 @@ pub(super) fn build_admin_pool_key_payload(
.and_then(serde_json::Value::as_object);
let quota_updated_at =
admin_pool_json_to_u64(quota_snapshot.and_then(|item| item.get("updated_at")));
let account_quota = admin_pool_build_account_quota(provider_type, quota_snapshot);
let oauth_invalid_at = if auth_semantics.can_show_oauth_metadata() {
admin_pool_json_to_u64(oauth_snapshot.and_then(|item| item.get("invalid_at")))
.or(key.oauth_invalid_at_unix_secs)
@@ -851,13 +972,7 @@ pub(super) fn build_admin_pool_key_payload(
);
payload.insert("proxy".to_string(), json!(key.proxy.clone()));
payload.insert("fingerprint".to_string(), json!(key.fingerprint.clone()));
payload.insert(
"account_quota".to_string(),
json!(admin_pool_build_account_quota(
provider_type,
key.upstream_metadata.as_ref(),
)),
);
payload.insert("account_quota".to_string(), json!(account_quota));
payload.insert("cooldown_reason".to_string(), json!(cooldown_reason));
payload.insert(
"cooldown_ttl_seconds".to_string(),

2
apps/aether-gateway/src/handlers/admin/request/provider/oauth.rs
View File

@@ -514,11 +514,13 @@ impl<'a> AdminAppState<'a> {
&self,
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
key_id: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<(bool, Option<String>), GatewayError> {
crate::handlers::admin::provider::oauth::runtime::refresh_provider_oauth_account_state_after_update(
self,
provider,
key_id,
proxy_override,
)
.await
}

889
apps/aether-gateway/src/handlers/shared/catalog.rs
View File

@@ -1,11 +1,12 @@
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
use crate::provider_key_auth::provider_key_auth_semantics;
use crate::AppState;
use aether_admin::provider::quota as admin_provider_quota_pure;
#[cfg(test)]
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
use aether_crypto::{decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
use serde_json::{json, Map, Value};
pub(crate) fn provider_catalog_key_supports_format(
key: &StoredProviderCatalogKey,
@@ -153,13 +154,683 @@ pub(crate) fn default_provider_key_status_snapshot() -> serde_json::Value {
})
}
fn provider_key_status_snapshot_object(
status_snapshot: Option<&Value>,
) -> Option<Map<String, Value>> {
status_snapshot.and_then(|value| match value {
Value::Object(object) => Some(object.clone()),
_ => None,
})
}
fn provider_quota_metadata_bucket<'a>(
upstream_metadata: Option<&'a Value>,
provider_type: &str,
) -> Option<&'a Map<String, Value>> {
upstream_metadata
.and_then(Value::as_object)
.and_then(|metadata| metadata.get(&provider_type.trim().to_ascii_lowercase()))
.and_then(Value::as_object)
}
fn provider_quota_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
let mut parsed = match value {
Some(Value::Number(number)) => number.as_f64(),
Some(Value::String(text)) => text.trim().parse::<f64>().ok(),
_ => None,
}?;
if !parsed.is_finite() || parsed <= 0.0 {
return None;
}
if parsed > 1_000_000_000_000.0 {
parsed /= 1000.0;
}
Some(parsed.floor() as u64)
}
fn provider_quota_model_bucket(metadata: &Map<String, Value>) -> Option<&Map<String, Value>> {
metadata
.get("quota_by_model")
.or_else(|| metadata.get("models"))
.and_then(Value::as_object)
}
fn quota_window_reset_seconds(
observed_at_unix_secs: Option<u64>,
reset_at_unix_secs: Option<u64>,
) -> Option<u64> {
observed_at_unix_secs
.zip(reset_at_unix_secs)
.map(|(observed_at, reset_at)| reset_at.saturating_sub(observed_at))
}
fn model_quota_window_snapshot(
model_name: &str,
item: &Map<String, Value>,
observed_at_unix_secs: Option<u64>,
) -> Option<Value> {
let used_ratio = item
.get("used_percent")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.map(|value| (value / 100.0).clamp(0.0, 1.0))
.or_else(|| {
item.get("remaining_fraction")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.map(|value| (1.0 - value.clamp(0.0, 1.0)).clamp(0.0, 1.0))
});
let remaining_ratio = item
.get("remaining_fraction")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.map(|value| value.clamp(0.0, 1.0))
.or_else(|| used_ratio.map(|value| (1.0 - value).max(0.0)));
let reset_at = provider_quota_timestamp_unix_secs(
item.get("reset_at").or_else(|| item.get("next_reset_at")),
);
let reset_seconds = quota_window_reset_seconds(observed_at_unix_secs, reset_at);
let is_exhausted = item
.get("is_exhausted")
.and_then(admin_provider_quota_pure::coerce_json_bool)
.or_else(|| used_ratio.map(|value| value >= 1.0 - 1e-6));
if used_ratio.is_none()
&& remaining_ratio.is_none()
&& reset_at.is_none()
&& reset_seconds.is_none()
&& is_exhausted.is_none()
{
return None;
}
let mut window = Map::new();
let label = item
.get("display_name")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(model_name);
window.insert("code".to_string(), json!(format!("model:{model_name}")));
window.insert("label".to_string(), json!(label));
window.insert("scope".to_string(), json!("model"));
window.insert("unit".to_string(), json!("percent"));
window.insert("model".to_string(), json!(model_name));
window.insert("used_ratio".to_string(), json!(used_ratio));
window.insert("remaining_ratio".to_string(), json!(remaining_ratio));
window.insert("reset_at".to_string(), json!(reset_at));
window.insert("reset_seconds".to_string(), json!(reset_seconds));
window.insert("is_exhausted".to_string(), json!(is_exhausted));
Some(Value::Object(window))
}
fn quota_windows_usage_ratio(windows: &[Value]) -> Option<f64> {
windows
.iter()
.filter_map(Value::as_object)
.filter_map(|window| window.get("used_ratio"))
.filter_map(Value::as_f64)
.max_by(f64::total_cmp)
}
fn quota_windows_min_reset_seconds(windows: &[Value]) -> Option<u64> {
windows
.iter()
.filter_map(Value::as_object)
.filter_map(|window| window.get("reset_seconds"))
.filter_map(admin_provider_quota_pure::coerce_json_u64)
.min()
}
fn quota_windows_all_exhausted(windows: &[Value]) -> bool {
let mut total = 0usize;
let mut exhausted = 0usize;
for window in windows.iter().filter_map(Value::as_object) {
total += 1;
let is_exhausted = window
.get("is_exhausted")
.and_then(admin_provider_quota_pure::coerce_json_bool)
.or_else(|| {
window
.get("used_ratio")
.and_then(Value::as_f64)
.map(|value| value >= 1.0 - 1e-6)
})
.unwrap_or(false);
if is_exhausted {
exhausted += 1;
}
}
total > 0 && exhausted == total
}
fn codex_quota_window_snapshot(
metadata: &Map<String, Value>,
prefix: &str,
code: &str,
label: &str,
observed_at_unix_secs: Option<u64>,
) -> Option<Value> {
let used_percent_key = format!("{prefix}_used_percent");
let reset_seconds_key = format!("{prefix}_reset_seconds");
let reset_after_seconds_key = format!("{prefix}_reset_after_seconds");
let reset_at_key = format!("{prefix}_reset_at");
let window_minutes_key = format!("{prefix}_window_minutes");
let used_percent = metadata
.get(&used_percent_key)
.and_then(admin_provider_quota_pure::coerce_json_f64);
let reset_at = metadata
.get(&reset_at_key)
.and_then(admin_provider_quota_pure::coerce_json_u64);
let reset_seconds = metadata
.get(&reset_seconds_key)
.and_then(admin_provider_quota_pure::coerce_json_u64)
.or_else(|| {
metadata
.get(&reset_after_seconds_key)
.and_then(admin_provider_quota_pure::coerce_json_u64)
})
.or_else(|| {
observed_at_unix_secs
.zip(reset_at)
.map(|(observed_at, reset_at)| reset_at.saturating_sub(observed_at))
});
let window_minutes = metadata
.get(&window_minutes_key)
.and_then(admin_provider_quota_pure::coerce_json_u64);
if used_percent.is_none()
&& reset_at.is_none()
&& reset_seconds.is_none()
&& window_minutes.is_none()
{
return None;
}
let used_ratio = used_percent.map(|value| (value / 100.0).clamp(0.0, 1.0));
let remaining_ratio = used_ratio.map(|value| (1.0 - value).max(0.0));
let mut window = Map::new();
window.insert("code".to_string(), json!(code));
window.insert("label".to_string(), json!(label));
window.insert("scope".to_string(), json!("account"));
window.insert("unit".to_string(), json!("percent"));
window.insert("used_ratio".to_string(), json!(used_ratio));
window.insert("remaining_ratio".to_string(), json!(remaining_ratio));
window.insert("reset_at".to_string(), json!(reset_at));
window.insert("reset_seconds".to_string(), json!(reset_seconds));
window.insert("window_minutes".to_string(), json!(window_minutes));
Some(Value::Object(window))
}
fn build_codex_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "codex")?;
let observed_at_unix_secs = metadata
.get("updated_at")
.and_then(admin_provider_quota_pure::coerce_json_u64);
let plan_type = metadata
.get("plan_type")
.and_then(Value::as_str)
.and_then(|value| admin_provider_quota_pure::normalize_codex_plan_type(Some(value)));
let credits_has_credits = metadata
.get("has_credits")
.and_then(admin_provider_quota_pure::coerce_json_bool);
let credits_balance = metadata
.get("credits_balance")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let credits_unlimited = metadata
.get("credits_unlimited")
.and_then(admin_provider_quota_pure::coerce_json_bool);
let windows = [
codex_quota_window_snapshot(metadata, "primary", "weekly", "周", observed_at_unix_secs),
codex_quota_window_snapshot(metadata, "secondary", "5h", "5H", observed_at_unix_secs),
]
.into_iter()
.flatten()
.collect::<Vec<_>>();
if windows.is_empty()
&& plan_type.is_none()
&& credits_has_credits.is_none()
&& credits_balance.is_none()
&& credits_unlimited.is_none()
&& observed_at_unix_secs.is_none()
{
return None;
}
let usage_ratio = windows
.iter()
.filter_map(Value::as_object)
.filter_map(|window| window.get("used_ratio"))
.filter_map(Value::as_f64)
.max_by(f64::total_cmp);
let reset_seconds = windows
.iter()
.filter_map(Value::as_object)
.filter_map(|window| window.get("reset_seconds"))
.filter_map(admin_provider_quota_pure::coerce_json_u64)
.min();
let exhausted_by_credits =
credits_unlimited != Some(true) && credits_has_credits == Some(false);
let exhausted_by_window = usage_ratio.is_some_and(|value| value >= 1.0 - 1e-6);
let exhausted = exhausted_by_credits || exhausted_by_window;
let mut credits = Map::new();
if let Some(value) = credits_has_credits {
credits.insert("has_credits".to_string(), json!(value));
}
if let Some(value) = credits_balance {
credits.insert("balance".to_string(), json!(value));
}
if let Some(value) = credits_unlimited {
credits.insert("unlimited".to_string(), json!(value));
}
let reason = if exhausted_by_credits {
Some("无可用积分")
} else if exhausted_by_window {
Some("额度窗口已耗尽")
} else {
None
};
Some(json!({
"version": 2,
"provider_type": "codex",
"code": if exhausted { "exhausted" } else { "ok" },
"label": if exhausted { Some("额度耗尽") } else { None::<&str> },
"reason": reason,
"freshness": "fresh",
"source": source,
"observed_at": observed_at_unix_secs,
"exhausted": exhausted,
"usage_ratio": usage_ratio,
"updated_at": observed_at_unix_secs,
"reset_seconds": reset_seconds,
"plan_type": plan_type,
"credits": if credits.is_empty() {
Value::Null
} else {
Value::Object(credits)
},
"windows": windows,
}))
}
fn build_kiro_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "kiro")?;
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
let usage_limit = metadata
.get("usage_limit")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let current_usage = metadata
.get("current_usage")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let remaining = metadata
.get("remaining")
.and_then(admin_provider_quota_pure::coerce_json_f64);
let usage_ratio = metadata
.get("usage_percentage")
.and_then(admin_provider_quota_pure::coerce_json_f64)
.map(|value| (value / 100.0).clamp(0.0, 1.0))
.or_else(|| {
current_usage
.zip(usage_limit)
.and_then(|(current_usage, usage_limit)| {
(usage_limit > 0.0).then_some((current_usage / usage_limit).clamp(0.0, 1.0))
})
});
let remaining_ratio = usage_ratio.map(|value| (1.0 - value).max(0.0));
let next_reset_at = provider_quota_timestamp_unix_secs(metadata.get("next_reset_at"));
let reset_seconds = quota_window_reset_seconds(observed_at_unix_secs, next_reset_at);
let plan_type = metadata
.get("subscription_title")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let is_banned = metadata
.get("is_banned")
.and_then(admin_provider_quota_pure::coerce_json_bool)
== Some(true);
let ban_reason = metadata
.get("ban_reason")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let mut windows = Vec::new();
if usage_ratio.is_some()
|| remaining.is_some()
|| usage_limit.is_some()
|| current_usage.is_some()
|| next_reset_at.is_some()
{
windows.push(json!({
"code": "usage",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": usage_ratio,
"remaining_ratio": remaining_ratio,
"used_value": current_usage,
"remaining_value": remaining,
"limit_value": usage_limit,
"reset_at": next_reset_at,
"reset_seconds": reset_seconds,
}));
}
if windows.is_empty() && plan_type.is_none() && observed_at_unix_secs.is_none() && !is_banned {
return None;
}
let exhausted = !is_banned
&& (remaining.is_some_and(|value| value <= 0.0)
|| usage_ratio.is_some_and(|value| value >= 1.0 - 1e-6));
let reason = if is_banned {
ban_reason
} else if exhausted {
Some("额度已耗尽".to_string())
} else {
None
};
let label = if is_banned {
Some("账号已封禁")
} else if exhausted {
Some("额度耗尽")
} else {
None
};
let code = if is_banned {
"banned"
} else if exhausted {
"exhausted"
} else {
"ok"
};
Some(json!({
"version": 2,
"provider_type": "kiro",
"code": code,
"label": label,
"reason": reason,
"freshness": "fresh",
"source": source,
"observed_at": observed_at_unix_secs,
"exhausted": exhausted,
"usage_ratio": usage_ratio,
"updated_at": observed_at_unix_secs,
"reset_seconds": reset_seconds,
"plan_type": plan_type,
"windows": windows,
}))
}
fn build_antigravity_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "antigravity")?;
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
let is_forbidden = metadata
.get("is_forbidden")
.and_then(admin_provider_quota_pure::coerce_json_bool)
== Some(true);
let forbidden_reason = metadata
.get("forbidden_reason")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let windows = provider_quota_model_bucket(metadata)
.map(|models| {
models
.iter()
.filter_map(|(model_name, item)| {
model_quota_window_snapshot(
model_name,
item.as_object()?,
observed_at_unix_secs,
)
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
if windows.is_empty() && observed_at_unix_secs.is_none() && !is_forbidden {
return None;
}
let usage_ratio = quota_windows_usage_ratio(&windows);
let reset_seconds = quota_windows_min_reset_seconds(&windows);
let exhausted = !is_forbidden && quota_windows_all_exhausted(&windows);
let reason = if is_forbidden {
forbidden_reason
} else if exhausted {
Some("所有模型额度已耗尽".to_string())
} else {
None
};
let label = if is_forbidden {
Some("访问受限")
} else if exhausted {
Some("额度耗尽")
} else {
None
};
let code = if is_forbidden {
"forbidden"
} else if exhausted {
"exhausted"
} else {
"ok"
};
Some(json!({
"version": 2,
"provider_type": "antigravity",
"code": code,
"label": label,
"reason": reason,
"freshness": "fresh",
"source": source,
"observed_at": observed_at_unix_secs,
"exhausted": exhausted,
"usage_ratio": usage_ratio,
"updated_at": observed_at_unix_secs,
"reset_seconds": reset_seconds,
"plan_type": serde_json::Value::Null,
"windows": windows,
}))
}
fn build_gemini_cli_quota_status_snapshot(
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let metadata = provider_quota_metadata_bucket(upstream_metadata, "gemini_cli")?;
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
let windows = provider_quota_model_bucket(metadata)
.map(|models| {
models
.iter()
.filter_map(|(model_name, item)| {
model_quota_window_snapshot(
model_name,
item.as_object()?,
observed_at_unix_secs,
)
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
if windows.is_empty() && observed_at_unix_secs.is_none() {
return None;
}
let usage_ratio = quota_windows_usage_ratio(&windows);
let active_exhausted_windows = windows
.iter()
.filter_map(Value::as_object)
.filter(|window| {
window
.get("is_exhausted")
.and_then(admin_provider_quota_pure::coerce_json_bool)
.or_else(|| {
window
.get("used_ratio")
.and_then(Value::as_f64)
.map(|value| value >= 1.0 - 1e-6)
})
.unwrap_or(false)
})
.filter(|window| {
provider_quota_timestamp_unix_secs(window.get("reset_at"))
.zip(observed_at_unix_secs)
.map(|(reset_at, observed_at)| reset_at > observed_at)
.unwrap_or(true)
})
.count();
let exhausted = !windows.is_empty() && active_exhausted_windows == windows.len();
let cooling = active_exhausted_windows > 0;
let reset_seconds = if cooling {
quota_windows_min_reset_seconds(&windows)
} else {
None
};
Some(json!({
"version": 2,
"provider_type": "gemini_cli",
"code": if exhausted {
"exhausted"
} else if cooling {
"cooldown"
} else {
"ok"
},
"label": if cooling { Some("冷却中") } else { None::<&str> },
"reason": if exhausted {
Some("所有模型均处于冷却中")
} else {
None::<&str>
},
"freshness": "fresh",
"source": source,
"observed_at": observed_at_unix_secs,
"exhausted": exhausted,
"usage_ratio": usage_ratio,
"updated_at": observed_at_unix_secs,
"reset_seconds": reset_seconds,
"plan_type": serde_json::Value::Null,
"windows": windows,
}))
}
pub(crate) fn sync_provider_key_quota_status_snapshot(
status_snapshot: Option<&Value>,
provider_type: &str,
upstream_metadata: Option<&Value>,
source: &str,
) -> Option<Value> {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
let quota = match normalized_provider_type.as_str() {
"codex" => build_codex_quota_status_snapshot(upstream_metadata, source),
"kiro" => build_kiro_quota_status_snapshot(upstream_metadata, source),
"antigravity" => build_antigravity_quota_status_snapshot(upstream_metadata, source),
"gemini_cli" => build_gemini_cli_quota_status_snapshot(upstream_metadata, source),
_ => None,
}?;
let default_snapshot = default_provider_key_status_snapshot();
let mut snapshot = provider_key_status_snapshot_object(status_snapshot)
.or_else(|| default_snapshot.as_object().cloned())
.unwrap_or_default();
snapshot.insert("quota".to_string(), quota);
Some(Value::Object(snapshot))
}
fn quota_snapshot_has_materialized_data(
quota_snapshot: Option<&Map<String, Value>>,
provider_type: &str,
) -> bool {
let Some(quota_snapshot) = quota_snapshot else {
return false;
};
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
let snapshot_provider_type = quota_snapshot
.get("provider_type")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.to_ascii_lowercase();
if !snapshot_provider_type.is_empty() && snapshot_provider_type != normalized_provider_type {
return false;
}
if quota_snapshot
.get("windows")
.and_then(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
{
return true;
}
if quota_snapshot
.get("credits")
.is_some_and(|credits| !credits.is_null())
{
return true;
}
quota_snapshot
.get("code")
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|code| {
!code.is_empty()
&& !code.eq_ignore_ascii_case("unknown")
&& !code.eq_ignore_ascii_case("ok")
})
}
pub(crate) fn provider_key_status_snapshot_payload(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> serde_json::Value {
key.status_snapshot
.clone()
.filter(|value| value.is_object())
.unwrap_or_else(default_provider_key_status_snapshot)
let status_snapshot = key
.status_snapshot
.as_ref()
.filter(|value| value.is_object());
let quota_snapshot = status_snapshot
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object);
if quota_snapshot_has_materialized_data(quota_snapshot, provider_type) {
return status_snapshot
.cloned()
.unwrap_or_else(default_provider_key_status_snapshot);
}
sync_provider_key_quota_status_snapshot(
status_snapshot,
provider_type,
key.upstream_metadata.as_ref(),
"catalog_fallback",
)
.or_else(|| status_snapshot.cloned())
.unwrap_or_else(default_provider_key_status_snapshot)
}
pub(crate) fn provider_key_health_summary(
@@ -514,7 +1185,7 @@ pub(crate) fn build_admin_provider_key_response(
);
payload.insert(
"status_snapshot".to_string(),
provider_key_status_snapshot_payload(key),
provider_key_status_snapshot_payload(key, provider_type),
);
payload.insert(
"cache_ttl_minutes".to_string(),
@@ -664,3 +1335,209 @@ pub(crate) fn build_admin_provider_key_response(
);
serde_json::Value::Object(payload)
}
#[cfg(test)]
mod tests {
use super::*;
fn sample_catalog_key() -> StoredProviderCatalogKey {
let encrypted_api_key =
encrypt_python_fernet_plaintext(DEVELOPMENT_ENCRYPTION_KEY, "sk-test-123")
.expect("api key ciphertext should build");
StoredProviderCatalogKey::new(
"key-test".to_string(),
"provider-test".to_string(),
"default".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:chat"])),
encrypted_api_key,
None,
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build")
}
#[test]
fn provider_key_status_snapshot_payload_backfills_missing_quota_from_upstream_metadata() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"codex": {
"updated_at": 1_775_553_285u64,
"plan_type": "plus",
"primary_used_percent": 55.0,
"primary_reset_at": 1_900_000_000u64,
"secondary_used_percent": 12.5,
"secondary_reset_at": 1_900_500_000u64,
"has_credits": true,
"credits_balance": 42.0
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
assert_eq!(quota.get("plan_type"), Some(&json!("plus")));
assert_eq!(quota.get("updated_at"), Some(&json!(1_775_553_285u64)));
assert_eq!(
quota
.get("credits")
.and_then(Value::as_object)
.and_then(|credits| credits.get("balance")),
Some(&json!(42.0))
);
assert_eq!(
quota.get("windows").and_then(Value::as_array).map(Vec::len),
Some(2usize)
);
}
#[test]
fn provider_key_status_snapshot_payload_preserves_existing_materialized_quota_snapshot() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"codex": {
"updated_at": 100u64,
"primary_used_percent": 100.0
}
}));
key.status_snapshot = Some(json!({
"oauth": {
"code": "none",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"expires_at": serde_json::Value::Null,
"invalid_at": serde_json::Value::Null,
"source": serde_json::Value::Null,
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"blocked": false,
"source": serde_json::Value::Null,
"recoverable": false
},
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "refresh_api",
"observed_at": 200u64,
"exhausted": false,
"usage_ratio": 0.25,
"updated_at": 200u64,
"reset_seconds": 3600u64,
"plan_type": "team",
"windows": [{
"code": "weekly",
"label": "周",
"scope": "account",
"unit": "percent",
"used_ratio": 0.25,
"remaining_ratio": 0.75,
"reset_at": 1_900_000_000u64,
"reset_seconds": 3600u64
}]
}
}));
let payload = provider_key_status_snapshot_payload(&key, "codex");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
assert_eq!(quota.get("updated_at"), Some(&json!(200u64)));
assert_eq!(quota.get("plan_type"), Some(&json!("team")));
assert_eq!(
quota
.get("windows")
.and_then(Value::as_array)
.and_then(|windows| windows.first())
.and_then(Value::as_object)
.and_then(|window| window.get("used_ratio")),
Some(&json!(0.25))
);
}
#[test]
fn provider_key_status_snapshot_payload_backfills_thin_ok_snapshot_from_upstream_metadata() {
let mut key = sample_catalog_key();
key.upstream_metadata = Some(json!({
"antigravity": {
"updated_at": 1_775_553_285u64,
"quota_by_model": {
"gemini-2.5-pro": { "used_percent": 0.0 },
"gemini-2.5-flash": { "used_percent": 25.0 }
}
}
}));
key.status_snapshot = Some(json!({
"oauth": {
"code": "none",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"expires_at": serde_json::Value::Null,
"invalid_at": serde_json::Value::Null,
"source": serde_json::Value::Null,
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"blocked": false,
"source": serde_json::Value::Null,
"recoverable": false
},
"quota": {
"version": 2,
"provider_type": "antigravity",
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "refresh_api",
"observed_at": 100u64,
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": 100u64,
"reset_seconds": serde_json::Value::Null,
"plan_type": serde_json::Value::Null
}
}));
let payload = provider_key_status_snapshot_payload(&key, "antigravity");
let quota = payload
.get("quota")
.and_then(Value::as_object)
.expect("quota snapshot should be object");
assert_eq!(quota.get("provider_type"), Some(&json!("antigravity")));
assert_eq!(quota.get("updated_at"), Some(&json!(1_775_553_285u64)));
assert_eq!(
quota.get("windows").and_then(Value::as_array).map(Vec::len),
Some(2usize)
);
}
}

2
apps/aether-gateway/src/handlers/shared/mod.rs
View File

@@ -22,7 +22,7 @@ pub(crate) use self::catalog::{
default_provider_key_status_snapshot, effective_catalog_encryption_key,
encrypt_catalog_secret_with_fallbacks, masked_catalog_api_key, parse_catalog_auth_config_json,
provider_catalog_key_supports_format, provider_key_health_summary,
provider_key_status_snapshot_payload,
provider_key_status_snapshot_payload, sync_provider_key_quota_status_snapshot,
};
pub(crate) use self::email_templates::{
admin_email_template_definition, admin_email_template_html_key,

101
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
View File

@@ -165,6 +165,21 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 1);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(
payload["results"][0]["quota_snapshot"]["provider_type"],
"codex"
);
assert_eq!(payload["results"][0]["quota_snapshot"]["plan_type"], "plus");
assert_eq!(
payload["results"][0]["quota_snapshot"]["credits"]["balance"],
json!(42.0)
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["windows"]
.as_array()
.map(Vec::len),
Some(2usize)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
@@ -517,6 +532,18 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_ad
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 1);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(
payload["results"][0]["quota_snapshot"]["provider_type"],
"kiro"
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["plan_type"],
"KIRO PRO+"
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["windows"][0]["remaining_value"],
json!(15.0)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
@@ -566,6 +593,40 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_ad
.and_then(|value| value.get("email")),
Some(&json!("dev@example.com"))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("provider_type")),
Some(&json!("kiro"))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("usage_ratio")),
Some(&json!(0.25))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("plan_type")),
Some(&json!("KIRO PRO+"))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("windows"))
.and_then(|value| value.get(0))
.and_then(|value| value.get("remaining_value")),
Some(&json!(15.0))
);
gateway_handle.abort();
execution_runtime_handle.abort();
@@ -830,6 +891,20 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
assert_eq!(payload["failed"], 0);
assert_eq!(payload["total"], 1);
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(
payload["results"][0]["quota_snapshot"]["provider_type"],
"antigravity"
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["usage_ratio"],
json!(0.75)
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["windows"]
.as_array()
.map(Vec::len),
Some(1usize)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
@@ -880,6 +955,32 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
.and_then(|value| value.get("used_percent")),
Some(&json!(75.0))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("provider_type")),
Some(&json!("antigravity"))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("usage_ratio")),
Some(&json!(0.75))
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.get("quota"))
.and_then(|value| value.get("windows"))
.and_then(|value| value.as_array())
.map(Vec::len),
Some(1usize)
);
gateway_handle.abort();
execution_runtime_handle.abort();

7
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -2564,6 +2564,13 @@ async fn gateway_batch_imports_admin_provider_oauth_kiro_via_execution_runtime_p
assert_eq!(payload["results"][0]["key_id"], "key-kiro-batch-runtime");
assert_eq!(payload["results"][0]["replaced"], true);
for _ in 0..40 {
let plan_count = execution_plans.lock().expect("mutex should lock").len();
if plan_count == 2 {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
}
{
let plans = execution_plans.lock().expect("mutex should lock");
assert_eq!(plans.len(), 2);

438
apps/aether-gateway/src/tests/control/admin/pool.rs
View File

@@ -1037,6 +1037,224 @@ async fn gateway_includes_pool_quota_and_compat_fields_in_list_keys_response() {
assert_eq!(keys[0]["account_status_blocked"], json!(false));
}
#[tokio::test]
async fn gateway_prefers_status_snapshot_antigravity_quota_over_stale_metadata() {
let mut provider = sample_provider("provider-antigravity", "antigravity", 10)
.with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "antigravity".to_string();
let mut key = sample_key(
"key-antigravity-snapshot-fresh",
"provider-antigravity",
"gemini:chat",
"oauth-placeholder",
);
key.name = "antigravity snapshot fresh".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"antigravity": {
"quota_by_model": {
"gemini-2.5-pro": { "used_percent": 100.0 },
"gemini-2.5-flash": { "used_percent": 100.0 }
}
}
}));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "antigravity",
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "refresh_api",
"observed_at": 1_775_553_285u64,
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": 1_775_553_285u64,
"reset_seconds": serde_json::Value::Null,
"plan_type": serde_json::Value::Null,
"windows": [
{
"code": "model:gemini-2.5-pro",
"label": "Gemini 2.5 Pro",
"scope": "model",
"unit": "percent",
"model": "gemini-2.5-pro",
"used_ratio": 0.0,
"remaining_ratio": 1.0,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
"is_exhausted": false
},
{
"code": "model:gemini-2.5-flash",
"label": "Gemini 2.5 Flash",
"scope": "model",
"unit": "percent",
"model": "gemini-2.5-flash",
"used_ratio": 0.0,
"remaining_ratio": 1.0,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
"is_exhausted": false
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-antigravity/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("available"));
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
assert_eq!(keys[0]["quota_updated_at"], json!(1_775_553_285u64));
assert_eq!(keys[0]["account_quota"], json!("最低剩余 100.0% (2 模型)"));
}
#[tokio::test]
async fn gateway_renders_gemini_cli_account_quota_from_status_snapshot() {
let mut provider = sample_provider("provider-gemini-cli", "gemini_cli", 10)
.with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "gemini_cli".to_string();
let mut key = sample_key(
"key-gemini-cli-snapshot",
"provider-gemini-cli",
"gemini:chat",
"oauth-placeholder",
);
key.name = "gemini cli snapshot".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"gemini_cli": {
"quota_by_model": {
"gemini-2.5-pro": {
"is_exhausted": false
}
}
}
}));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "gemini_cli",
"code": "cooldown",
"label": "冷却中",
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "background_refresh",
"observed_at": 1_775_553_285u64,
"exhausted": false,
"usage_ratio": 1.0,
"updated_at": 1_775_553_285u64,
"reset_seconds": serde_json::Value::Null,
"plan_type": serde_json::Value::Null,
"windows": [
{
"code": "model:gemini-2.5-pro",
"label": "Gemini 2.5 Pro",
"scope": "model",
"unit": "percent",
"model": "gemini-2.5-pro",
"used_ratio": 1.0,
"remaining_ratio": 0.0,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
"is_exhausted": true
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-gemini-cli/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("available"));
assert_eq!(keys[0]["quota_updated_at"], json!(1_775_553_285u64));
assert_eq!(keys[0]["account_quota"], json!("Gemini 2.5 Pro 冷却中"));
}
#[tokio::test]
async fn gateway_formats_codex_quota_countdown_from_reset_after_seconds() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
@@ -1192,6 +1410,124 @@ async fn gateway_marks_exhausted_codex_pool_key_as_blocked_when_flag_enabled() {
assert_eq!(keys[0]["account_quota"], json!("5H剩余 0.0%"));
}
#[tokio::test]
async fn gateway_prefers_status_snapshot_codex_quota_over_stale_metadata() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "codex".to_string();
let mut key = sample_key(
"key-codex-snapshot-fresh",
"provider-codex",
"openai:cli",
"oauth-placeholder",
);
key.name = "codex snapshot fresh".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"codex": {
"plan_type": "plus",
"secondary_used_percent": 100.0
}
}));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "response_headers",
"observed_at": 1_775_553_285u64,
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": 1_775_553_285u64,
"reset_seconds": serde_json::Value::Null,
"plan_type": "plus",
"credits": {
"has_credits": true,
"balance": 12.5,
"unlimited": false
},
"windows": [
{
"code": "weekly",
"label": "周",
"scope": "account",
"unit": "percent",
"used_ratio": 0.0,
"remaining_ratio": 1.0,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
"window_minutes": 10_080
},
{
"code": "5h",
"label": "5H",
"scope": "account",
"unit": "percent",
"used_ratio": 0.0,
"remaining_ratio": 1.0,
"reset_at": serde_json::Value::Null,
"reset_seconds": serde_json::Value::Null,
"window_minutes": 300
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-codex/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("available"));
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
assert_eq!(keys[0]["quota_updated_at"], json!(1_775_553_285u64));
assert_eq!(
keys[0]["account_quota"],
json!("周剩余 100.0% | 5H剩余 100.0%")
);
}
#[tokio::test]
async fn gateway_marks_exhausted_kiro_pool_key_as_blocked_when_flag_enabled() {
let mut provider = sample_provider("provider-kiro", "kiro", 10).with_transport_fields(
@@ -1265,6 +1601,108 @@ async fn gateway_marks_exhausted_kiro_pool_key_as_blocked_when_flag_enabled() {
assert_eq!(keys[0]["account_quota"], json!("剩余 0/100"));
}
#[tokio::test]
async fn gateway_prefers_status_snapshot_kiro_quota_over_stale_metadata() {
let mut provider = sample_provider("provider-kiro", "kiro", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true,
"skip_exhausted_accounts": true
}
})),
);
provider.provider_type = "kiro".to_string();
let mut key = sample_key(
"key-kiro-snapshot-fresh",
"provider-kiro",
"claude:cli",
"oauth-placeholder",
);
key.name = "kiro snapshot fresh".to_string();
key.auth_type = "oauth".to_string();
key.upstream_metadata = Some(json!({
"kiro": {
"remaining": 0.0,
"usage_limit": 100.0,
"current_usage": 100.0
}
}));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "kiro",
"code": "ok",
"label": serde_json::Value::Null,
"reason": serde_json::Value::Null,
"freshness": "fresh",
"source": "refresh_api",
"observed_at": 1_775_553_285u64,
"exhausted": false,
"usage_ratio": 0.25,
"updated_at": 1_775_553_285u64,
"reset_seconds": 86_400u64,
"plan_type": "KIRO PRO+",
"windows": [
{
"code": "usage",
"label": "额度",
"scope": "account",
"unit": "count",
"used_ratio": 0.25,
"remaining_ratio": 0.75,
"used_value": 5.0,
"remaining_value": 15.0,
"limit_value": 20.0,
"reset_at": 1_775_639_685u64,
"reset_seconds": 86_400u64
}
]
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key],
));
let state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(GatewayDataState::with_provider_catalog_reader_for_tests(
provider_catalog_repository,
));
let response = local_admin_pool_response(
&state,
http::Method::GET,
"/api/admin/pool/provider-kiro/keys?page=1&page_size=50&status=all",
None,
)
.await;
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = serde_json::from_slice(
&to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("json body should parse");
let keys = payload["keys"].as_array().expect("keys should be array");
assert_eq!(keys[0]["scheduling_status"], json!("available"));
assert_eq!(keys[0]["scheduling_reason"], json!("available"));
assert_eq!(keys[0]["quota_updated_at"], json!(1_775_553_285u64));
assert_eq!(keys[0]["account_quota"], json!("剩余 75.0% (5/20)"));
}
#[tokio::test]
async fn gateway_codex_quota_resets_to_full_after_countdown_elapsed() {
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(

27
apps/aether-gateway/src/tests/video/openai_sync_task.rs
View File

@@ -288,15 +288,15 @@ async fn gateway_executes_openai_video_delete_via_reconstructed_data_backed_loca
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_video_task_truth_source_mode(VideoTaskTruthSourceMode::RustAuthoritative)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_video_task_provider_transport_and_request_candidate_repository_for_tests(
repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
.with_video_task_truth_source_mode(VideoTaskTruthSourceMode::RustAuthoritative)
.with_data_state_for_tests(
crate::data::GatewayDataState::with_video_task_provider_transport_and_request_candidate_repository_for_tests(
repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -308,8 +308,15 @@ async fn gateway_executes_openai_video_delete_via_reconstructed_data_backed_loca
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let response_json: serde_json::Value = response.json().await.expect("body should parse");
let response_status = response.status();
let response_text = response.text().await.expect("body should read");
assert_eq!(
response_status,
StatusCode::OK,
"unexpected response body: {response_text}"
);
let response_json: serde_json::Value =
serde_json::from_str(&response_text).expect("body should parse");
assert_eq!(
response_json,
json!({

111
apps/aether-gateway/src/usage/reporting/codex_realtime_quota.rs
View File

@@ -3,6 +3,7 @@ use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
use crate::clock::current_unix_secs;
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::{AppState, GatewayError};
use aether_admin::provider::quota as admin_provider_quota_pure;
use serde_json::Value;
@@ -185,14 +186,21 @@ pub(super) async fn sync_codex_quota_from_response_headers(
let updated_upstream_metadata =
merge_metadata_object(key.upstream_metadata.as_ref(), "codex", parsed);
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 mut updated_key = key;
updated_key.upstream_metadata = updated_upstream_metadata;
updated_key.status_snapshot = updated_status_snapshot;
updated_key.updated_at_unix_secs = Some(now_unix_secs);
let updated = state
.update_provider_catalog_key_upstream_metadata(
&key_id,
updated_upstream_metadata.as_ref(),
Some(now_unix_secs),
)
.await?;
.update_provider_catalog_key(&updated_key)
.await?
.is_some();
if updated {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
}
@@ -270,6 +278,15 @@ mod tests {
key
}
fn quota_snapshot<'a>(key: &'a StoredProviderCatalogKey) -> &'a serde_json::Map<String, Value> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.expect("quota snapshot should exist")
}
fn paid_headers(
primary_used_percent: &str,
secondary_used_percent: &str,
@@ -323,25 +340,43 @@ mod tests {
async fn sync_codex_quota_replaces_existing_codex_fields_and_preserves_other_sections() {
clear_codex_quota_fingerprint_cache();
let mut key = sample_key(
"key-codex-1",
"provider-codex",
Some(json!({
"codex": {
"legacy_marker": "drop-me",
"secondary_used_percent": 2.0,
"credits_balance": 42.0,
"account_disabled": true,
"reason": "deactivated_workspace"
},
"other": {
"value": true
}
})),
);
key.status_snapshot = Some(json!({
"oauth": {
"code": "valid",
"label": "有效",
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"blocked": false,
"recoverable": false
},
"quota": {
"code": "unknown",
"exhausted": false
}
}));
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key(
"key-codex-1",
"provider-codex",
Some(json!({
"codex": {
"legacy_marker": "drop-me",
"secondary_used_percent": 2.0,
"credits_balance": 42.0,
"account_disabled": true,
"reason": "deactivated_workspace"
},
"other": {
"value": true
}
})),
)],
vec![key],
));
let state = build_state(Arc::clone(&repository));
@@ -386,6 +421,38 @@ mod tests {
.and_then(|metadata| metadata.get("other")),
Some(&json!({"value": true}))
);
let quota = quota_snapshot(&reloaded[0]);
assert_eq!(quota.get("version"), Some(&json!(2)));
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
assert_eq!(quota.get("source"), Some(&json!("response_headers")));
assert_eq!(quota.get("code"), Some(&json!("exhausted")));
assert_eq!(quota.get("exhausted"), Some(&json!(true)));
assert_eq!(quota.get("plan_type"), Some(&json!("team")));
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
let windows = quota
.get("windows")
.and_then(Value::as_array)
.expect("windows should be array");
assert_eq!(windows.len(), 2);
assert_eq!(windows[0].get("code"), Some(&json!("weekly")));
assert_eq!(windows[1].get("code"), Some(&json!("5h")));
let oauth = reloaded[0]
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("oauth"))
.and_then(Value::as_object)
.expect("oauth snapshot should exist");
assert_eq!(oauth.get("code"), Some(&json!("valid")));
let account = reloaded[0]
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object)
.expect("account snapshot should exist");
assert_eq!(account.get("code"), Some(&json!("ok")));
}
#[tokio::test]

22
apps/aether-gateway/src/usage/reporting/mod.rs
View File

@@ -873,6 +873,17 @@ mod tests {
.expect("codex metadata should exist");
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
let quota = reloaded[0]
.status_snapshot
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(serde_json::Value::as_object)
.expect("quota snapshot should exist");
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
assert_eq!(quota.get("source"), Some(&json!("response_headers")));
assert_eq!(quota.get("code"), Some(&json!("exhausted")));
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
}
#[tokio::test]
@@ -925,6 +936,17 @@ mod tests {
.expect("codex metadata should exist");
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
let quota = reloaded[0]
.status_snapshot
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(serde_json::Value::as_object)
.expect("quota snapshot should exist");
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
assert_eq!(quota.get("source"), Some(&json!("response_headers")));
assert_eq!(quota.get("code"), Some(&json!("exhausted")));
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
}
#[tokio::test]