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

fix(providers): refresh Antigravity grouped quotas

Browse Source
This commit is contained in:
ZheFox
2026-09-03 10:46:44 +08:00
parent cc6f5e89b6
commit 45c840b8d3
20 changed files with 976 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+107 -7
View File
@@ -7,15 +7,20 @@ use super::shared::{
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_antigravity_usage_response;
use aether_admin::provider::quota::{
parse_antigravity_quota_summary_response, parse_antigravity_usage_response,
};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_pool::build_antigravity_pool_quota_request;
use aether_provider_pool::{
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
};
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
async fn execute_antigravity_quota_plan(
state: &AdminAppState<'_>,
@@ -55,6 +60,79 @@ async fn execute_antigravity_quota_plan(
execute_provider_quota_plan(state, transport, plan, "antigravity").await
}
async fn fetch_antigravity_quota_summary_best_effort(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
authorization: (String, String),
project_id: &str,
identity_headers: BTreeMap<String, String>,
proxy_override: Option<&ProxySnapshot>,
) -> Option<serde_json::Value> {
let mut request_project_id = Some(project_id);
loop {
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
};
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let spec = build_antigravity_pool_quota_summary_request(
&transport.key.id,
&transport.endpoint.base_url,
authorization.clone(),
request_project_id,
identity_headers.clone(),
);
let plan = build_provider_quota_execution_plan(
transport,
spec,
proxy,
state.resolve_transport_profile(transport),
timeouts,
);
let outcome = match execute_provider_quota_plan(state, transport, plan, "antigravity").await
{
Ok(outcome) => outcome,
Err(error) => {
warn!(error = ?error, "Antigravity grouped quota request failed");
return None;
}
};
let result = match outcome {
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
warn!(detail = %detail, "Antigravity grouped quota execution failed");
return None;
}
};
if result.status_code == 200 {
return result
.body
.as_ref()
.and_then(|body| body.json_body.as_ref())
.and_then(parse_antigravity_quota_summary_response);
}
if result.status_code == 403 && request_project_id.is_some() {
request_project_id = None;
continue;
}
warn!(
status_code = result.status_code,
"Antigravity grouped quota request returned a non-success status"
);
return None;
}
}
pub(crate) async fn refresh_antigravity_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -130,9 +208,9 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
let result = match execute_antigravity_quota_plan(
state,
&transport,
authorization,
authorization.clone(),
&project_id,
identity_headers,
identity_headers.clone(),
proxy_override.as_ref(),
)
.await?
@@ -168,9 +246,31 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
.as_ref()
.and_then(|body| body.json_body.as_ref())
{
metadata_update = parse_antigravity_usage_response(body_json, now_unix_secs)
.map(|metadata| json!({ "antigravity": metadata }));
if metadata_update.is_some() {
if let Some(mut metadata) =
parse_antigravity_usage_response(body_json, now_unix_secs)
{
if let Some(metadata) = metadata.as_object_mut() {
metadata.insert("project_id".to_string(), json!(project_id));
}
if let Some(quota_groups) = fetch_antigravity_quota_summary_best_effort(
state,
&transport,
authorization,
&project_id,
identity_headers,
proxy_override.as_ref(),
)
.await
{
if let Some(metadata) = metadata.as_object_mut() {
metadata.insert("quota_groups".to_string(), quota_groups);
metadata.insert(
"quota_groups_updated_at".to_string(),
json!(now_unix_secs),
);
}
}
metadata_update = Some(json!({ "antigravity": metadata }));
status = "success".to_string();
} else {
status = "no_metadata".to_string();
apps/aether-gateway/src/handlers/shared/catalog.rs
+222 -1
View File
@@ -654,6 +654,112 @@ fn antigravity_model_quota_window_snapshot(
Some(window)
}
fn antigravity_grouped_quota_window_snapshots(
metadata: &Map<String, Value>,
observed_at_unix_secs: Option<u64>,
) -> Vec<Value> {
let mut windows = Vec::new();
let Some(groups) = metadata.get("quota_groups").and_then(Value::as_array) else {
return windows;
};
for (group_index, group) in groups.iter().filter_map(Value::as_object).enumerate() {
let group_code = group
.get("group_id")
.or_else(|| group.get("groupId"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("group:{group_index}"));
let group_label = group
.get("display_name")
.or_else(|| group.get("displayName"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("Quota group {}", group_index + 1));
let group_description = group
.get("description")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
for (bucket_index, bucket) in group
.get("buckets")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.enumerate()
{
let bucket_id = bucket
.get("bucket_id")
.or_else(|| bucket.get("bucketId"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.unwrap_or_else(|| format!("bucket-{}", bucket_index + 1));
let Some(mut window) =
model_quota_window_snapshot(&bucket_id, bucket, observed_at_unix_secs)
else {
continue;
};
let Some(window) = window.as_object_mut() else {
continue;
};
let period = bucket
.get("window")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let bucket_detail = bucket
.get("display_name")
.or_else(|| bucket.get("displayName"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| period.clone());
let bucket_label = bucket_detail
.filter(|detail| !detail.eq_ignore_ascii_case(&group_label))
.map(|detail| format!("{group_label} · {detail}"))
.unwrap_or_else(|| group_label.clone());
let description = bucket
.get("description")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.or_else(|| group_description.clone());
window.insert(
"code".to_string(),
json!(format!("group:{group_index}:{bucket_id}")),
);
window.insert("label".to_string(), json!(bucket_label));
window.insert("scope".to_string(), json!("quota_group"));
window.remove("model");
window.insert("quota_group".to_string(), json!(group_code));
window.insert("quota_group_label".to_string(), json!(group_label));
window.insert("bucket_id".to_string(), json!(bucket_id));
if let Some(period) = period {
window.insert("window".to_string(), json!(period));
}
if let Some(description) = description {
window.insert("description".to_string(), json!(description));
}
windows.push(Value::Object(window.clone()));
}
}
windows
}
fn provider_quota_metadata_string(
metadata: &Map<String, Value>,
fields: &[&str],
@@ -1622,7 +1728,7 @@ fn build_antigravity_quota_status_snapshot(
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let windows = provider_quota_model_bucket(metadata)
let mut windows = provider_quota_model_bucket(metadata)
.map(|models| {
models
.iter()
@@ -1636,6 +1742,13 @@ fn build_antigravity_quota_status_snapshot(
.collect::<Vec<_>>()
})
.unwrap_or_default();
let grouped_observed_at_unix_secs =
provider_quota_timestamp_unix_secs(metadata.get("quota_groups_updated_at"))
.or(observed_at_unix_secs);
windows.extend(antigravity_grouped_quota_window_snapshots(
metadata,
grouped_observed_at_unix_secs,
));
if windows.is_empty() && observed_at_unix_secs.is_none() && !is_forbidden {
return None;
@@ -3132,6 +3245,52 @@ mod tests {
assert_eq!(spark_weekly.get("remaining_ratio"), Some(&json!(0.95)));
}
#[test]
fn codex_wham_snapshot_does_not_duplicate_normalized_spark_windows() {
let codex = admin_provider_quota_pure::parse_codex_wham_usage_response(
&json!({
"plan_type": "plus",
"rate_limit": {
"primary_window": {"used_percent": 25.0},
"secondary_window": {"used_percent": 10.0}
},
"additional_rate_limits": [{
"limit_name": "GPT-5.3-Codex-Spark",
"rate_limit": {
"primary_window": {
"used_percent": 40.0,
"limit_window_seconds": 18_000
},
"secondary_window": {
"used_percent": 5.0,
"limit_window_seconds": 604_800
}
}
}]
}),
1_777_000_000,
)
.expect("Codex WHAM quota should parse");
let upstream_metadata = json!({"codex": codex});
let payload = sync_provider_key_quota_status_snapshot(
None,
"codex",
Some(&upstream_metadata),
"refresh_api",
)
.expect("Codex quota snapshot should sync");
let windows = payload["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
let spark_codes = windows
.iter()
.filter_map(|window| window.get("code").and_then(Value::as_str))
.filter(|code| code.starts_with("spark_"))
.collect::<Vec<_>>();
assert_eq!(spark_codes, vec!["spark_5h", "spark_weekly"]);
}
#[test]
fn provider_key_status_snapshot_payload_keeps_codex_free_window_quota_available() {
let mut key = sample_catalog_key();
@@ -4166,6 +4325,68 @@ mod tests {
assert_eq!(claude_window.get("reset_seconds"), Some(&json!(12_011u64)));
}
#[test]
fn sync_provider_key_quota_status_snapshot_materializes_antigravity_group_windows() {
let upstream_metadata = json!({
"antigravity": {
"updated_at": 1_777_000_000u64,
"models": {
"gemini-3.7-flash-tiered": {
"remaining_fraction": 0.9
}
},
"quota_groups": [{
"display_name": "Claude and GPT models",
"description": "Shared quota",
"buckets": [{
"bucket_id": "3p-5h",
"window": "5h",
"remaining_fraction": 0.25,
"reset_time": "2026-05-05T05:00:00Z",
"display_name": "5 hour"
}, {
"bucket_id": "3p-weekly",
"window": "weekly",
"remaining_fraction": 0.8,
"reset_time": "2026-05-11T00:00:00Z"
}]
}]
}
});
let payload = sync_provider_key_quota_status_snapshot(
None,
"antigravity",
Some(&upstream_metadata),
"refresh_api",
)
.expect("Antigravity quota snapshot should sync");
let windows = payload["quota"]["windows"]
.as_array()
.expect("quota windows should exist");
let five_hour = windows
.iter()
.find(|window| window["code"] == "group:0:3p-5h")
.expect("5h grouped quota window should exist");
let weekly = windows
.iter()
.find(|window| window["code"] == "group:0:3p-weekly")
.expect("weekly grouped quota window should exist");
assert_eq!(windows.len(), 3);
assert_eq!(five_hour["scope"], json!("quota_group"));
assert_eq!(five_hour["bucket_id"], json!("3p-5h"));
assert_eq!(five_hour["label"], json!("Claude and GPT models · 5 hour"));
assert_eq!(
five_hour["quota_group_label"],
json!("Claude and GPT models")
);
assert_eq!(five_hour["remaining_ratio"], json!(0.25));
assert_eq!(five_hour["used_ratio"], json!(0.75));
assert!(five_hour.get("model").is_none());
assert_eq!(weekly["window"], json!("weekly"));
}
#[test]
fn provider_key_status_snapshot_payload_backfills_account_block_from_oauth_invalid_reason() {
let mut key = sample_catalog_key();