i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +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

95
crates/aether-admin/src/provider/pool.rs
View File

@@ -121,10 +121,72 @@ fn admin_pool_json_f64(value: Option<&Value>) -> Option<f64> {
.filter(|value| value.is_finite())
}
fn admin_pool_quota_snapshot_matches_provider(
quota_snapshot: &serde_json::Map<String, Value>,
provider_type: &str,
) -> bool {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
match quota_snapshot
.get("provider_type")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
Some(provider_type) => provider_type.eq_ignore_ascii_case(&normalized_provider_type),
None => {
admin_pool_json_bool(quota_snapshot.get("exhausted")) == Some(true)
|| quota_snapshot
.get("code")
.and_then(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(Value::as_array)
.is_some_and(|windows| !windows.is_empty())
|| quota_snapshot
.get("credits")
.and_then(Value::as_object)
.is_some_and(|credits| !credits.is_empty())
}
}
}
fn admin_pool_key_quota_snapshot<'a>(
key: &'a StoredProviderCatalogKey,
provider_type: &str,
) -> Option<&'a serde_json::Map<String, Value>> {
let quota_snapshot = key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)?;
admin_pool_quota_snapshot_matches_provider(quota_snapshot, provider_type)
.then_some(quota_snapshot)
}
pub fn admin_pool_key_account_quota_exhausted(
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> bool {
if let Some(exhausted) = admin_pool_key_quota_snapshot(key, provider_type)
.and_then(|quota_snapshot| admin_pool_json_bool(quota_snapshot.get("exhausted")))
{
return exhausted;
}
let provider_type = provider_type.trim().to_ascii_lowercase();
let Some(bucket) = admin_pool_metadata_bucket(key.upstream_metadata.as_ref(), &provider_type)
else {
@@ -663,6 +725,39 @@ mod tests {
));
}
#[test]
fn prefers_quota_snapshot_over_metadata_for_codex_exhaustion() {
let mut key = sample_key(Some(json!({
"codex": {
"secondary_used_percent": 100.0
}
})));
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "ok",
"exhausted": false,
"usage_ratio": 0.0,
"updated_at": 1_776_395_200u64,
"windows": [
{
"code": "weekly",
"used_ratio": 0.0,
"remaining_ratio": 1.0
},
{
"code": "5h",
"used_ratio": 0.0,
"remaining_ratio": 1.0
}
]
}
}));
assert!(!admin_pool_key_account_quota_exhausted(&key, "codex"));
}
#[test]
fn detects_kiro_exhaustion_from_metadata() {
assert!(admin_pool_key_account_quota_exhausted(

13
crates/aether-admin/src/provider/quota.rs
View File

@@ -145,14 +145,11 @@ pub fn parse_antigravity_usage_response(
let remaining_fraction = quota_info
.and_then(|object| object.get("remainingFraction"))
.and_then(coerce_json_f64);
let used_percent = remaining_fraction
.map(|value| ((1.0 - value).max(0.0) * 100.0).min(100.0))
.unwrap_or(100.0);
payload.insert(
"remaining_fraction".to_string(),
json!(remaining_fraction.unwrap_or(0.0)),
);
payload.insert("used_percent".to_string(), json!(used_percent));
if let Some(remaining_fraction) = remaining_fraction {
let used_percent = ((1.0 - remaining_fraction).max(0.0) * 100.0).min(100.0);
payload.insert("remaining_fraction".to_string(), json!(remaining_fraction));
payload.insert("used_percent".to_string(), json!(used_percent));
}
if let Some(reset_time) = quota_info
.and_then(|object| object.get("resetTime"))
.cloned()