i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-09 04:30:20 +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/ai_serving/planner/standard/openai/chat/decision/request.rs
+2 -2
View File
@@ -2741,7 +2741,7 @@ mod tests {
.provider_request_headers
.get("x-client-version")
.map(String::as_str),
Some("1.2.3")
Some("4.3.0")
);
assert_eq!(
payload
@@ -2761,7 +2761,7 @@ mod tests {
assert_eq!(payload.provider_request_body["model"], "gemini-2.5-pro");
assert_eq!(
payload.provider_request_body["userAgent"],
"antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)"
"vscode/1.X.X (Antigravity/4.3.0)"
);
assert_eq!(payload.provider_request_body["requestType"], "agent");
assert!(payload.provider_request_body.get("contents").is_none());
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();
apps/aether-gateway/src/tests/ai_execute/finalize_local.rs
+1 -1
View File
@@ -1949,7 +1949,7 @@ async fn gateway_executes_openai_chat_antigravity_cross_format_sync_via_local_fi
"Bearer imported-antigravity-chat-token"
);
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0");
assert_eq!(
seen_execution_runtime_request.x_vscode_sessionid,
"sess-antigravity-chat-local-123"
apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs
+1 -1
View File
@@ -1537,7 +1537,7 @@ async fn gateway_executes_openai_responses_antigravity_cross_format_upstream_str
);
assert_eq!(
seen_remote_execution_runtime_request.x_client_version,
"1.2.3"
"4.3.0"
);
assert_eq!(
seen_remote_execution_runtime_request.x_vscode_sessionid,
apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs
+1 -1
View File
@@ -2123,7 +2123,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_upstream_stream_via_local_
);
assert_eq!(
seen_remote_execution_runtime_request.x_client_version,
"1.2.3"
"4.3.0"
);
assert_eq!(
seen_remote_execution_runtime_request.x_vscode_sessionid,
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs
+1 -1
View File
@@ -2017,7 +2017,7 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
"Bearer refreshed-antigravity-cli-stream-access-token"
);
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0");
assert_eq!(
seen_execution_runtime_request.x_vscode_sessionid,
"sess-antigravity-stream-local-123"
apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs
+1 -1
View File
@@ -2321,7 +2321,7 @@ async fn gateway_executes_antigravity_gemini_cli_sync_via_local_decision_gate_af
"Bearer refreshed-antigravity-cli-access-token"
);
assert_eq!(seen_execution_runtime_request.x_client_name, "antigravity");
assert_eq!(seen_execution_runtime_request.x_client_version, "1.2.3");
assert_eq!(seen_execution_runtime_request.x_client_version, "4.3.0");
assert_eq!(
seen_execution_runtime_request.x_vscode_sessionid,
"sess-antigravity-local-123"
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+133 -31
View File
@@ -2264,6 +2264,8 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
user_agent: String,
x_client_version: String,
provider_api_format: String,
request_body: Option<serde_json::Value>,
}
@@ -2281,7 +2283,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
}),
);
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime = Arc::new(Mutex::new(Vec::<SeenExecutionRuntimeRequest>::new()));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
@@ -2294,26 +2296,30 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
let request_body = plan.body.json_body.clone();
seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
request_body: plan.body.json_body.clone(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
.expect("mutex should lock")
.push(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
user_agent: plan.headers.get("user-agent").cloned().unwrap_or_default(),
x_client_version: plan
.headers
.get("x-client-version")
.cloned()
.unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
request_body: request_body.clone(),
});
let (status_code, json_body) = match plan.provider_api_format.as_str() {
"antigravity:fetch_available_models" => (
200,
json!({
"models": {
"claude-sonnet-4": {
"displayName": "Claude Sonnet 4",
@@ -2326,7 +2332,47 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
"displayName": "Gemini 2.5 Pro"
}
}
})),
}),
),
"antigravity:retrieve_user_quota_summary"
if request_body
.as_ref()
.and_then(|body| body.get("project"))
.is_some() =>
{
(403, json!({"error": {"message": "project not accepted"}}))
}
"antigravity:retrieve_user_quota_summary" => (
200,
json!({
"groups": [{
"displayName": "Claude and GPT models",
"description": "Shared quota",
"buckets": [{
"bucketId": "3p-5h",
"window": "5h",
"remainingFraction": 0.25,
"resetTime": "2026-05-05T05:00:00Z",
"displayName": "5 hour"
}, {
"bucketId": "3p-weekly",
"window": "weekly",
"remainingFraction": 0.8,
"resetTime": "2026-05-11T00:00:00Z"
}]
}]
}),
),
unexpected => panic!("unexpected quota request format: {unexpected}"),
};
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json_body),
body_bytes_b64: None,
}),
telemetry: None,
@@ -2429,31 +2475,55 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
payload["results"][0]["quota_snapshot"]["windows"]
.as_array()
.map(Vec::len),
Some(1usize)
Some(3usize)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let seen_execution_runtime_request = seen_execution_runtime
let seen_execution_runtime_requests = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
.clone();
assert_eq!(seen_execution_runtime_requests.len(), 3);
let fetch_models_request = &seen_execution_runtime_requests[0];
assert_eq!(
seen_execution_runtime_request.url,
fetch_models_request.url,
"https://daily-cloudcode-pa.googleapis.com/v1internal:fetchAvailableModels"
);
assert_eq!(fetch_models_request.authorization, "Bearer ya29.ant-token");
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer ya29.ant-token"
fetch_models_request.user_agent,
"vscode/1.X.X (Antigravity/4.3.0)"
);
assert_eq!(fetch_models_request.x_client_version, "4.3.0");
assert_eq!(
seen_execution_runtime_request.provider_api_format,
fetch_models_request.provider_api_format,
"antigravity:fetch_available_models"
);
assert_eq!(
seen_execution_runtime_request.request_body,
fetch_models_request.request_body,
Some(json!({ "project": "project-ant-123" }))
);
let grouped_with_project = &seen_execution_runtime_requests[1];
let grouped_without_project = &seen_execution_runtime_requests[2];
assert_eq!(
grouped_with_project.url,
"https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary"
);
assert_eq!(
grouped_with_project.provider_api_format,
"antigravity:retrieve_user_quota_summary"
);
assert_eq!(
grouped_with_project.request_body,
Some(json!({"project": "project-ant-123"}))
);
assert_eq!(grouped_without_project.url, grouped_with_project.url);
assert_eq!(grouped_without_project.request_body, Some(json!({})));
assert!(seen_execution_runtime_requests.iter().all(|request| {
request.authorization == "Bearer ya29.ant-token"
&& request.user_agent == "vscode/1.X.X (Antigravity/4.3.0)"
&& request.x_client_version == "4.3.0"
}));
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-antigravity-a".to_string()])
@@ -2481,6 +2551,26 @@ 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]
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/antigravity/quota_groups/0/buckets/0/bucket_id")),
Some(&json!("3p-5h"))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/antigravity/project_id")),
Some(&json!("project-ant-123"))
);
assert!(reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.pointer("/antigravity/quota_groups_updated_at"))
.and_then(serde_json::Value::as_u64)
.is_some());
assert_eq!(
reloaded[0]
.status_snapshot
@@ -2505,7 +2595,19 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
.and_then(|value| value.get("windows"))
.and_then(|value| value.as_array())
.map(Vec::len),
Some(1usize)
Some(3usize)
);
assert_eq!(
reloaded[0]
.status_snapshot
.as_ref()
.and_then(|value| value.pointer("/quota/windows"))
.and_then(serde_json::Value::as_array)
.and_then(|windows| windows
.iter()
.find(|window| window["code"] == "group:0:3p-5h"))
.and_then(|window| window.get("remaining_ratio")),
Some(&json!(0.25))
);
gateway_handle.abort();
crates/aether-admin/src/provider/quota.rs
+165 -10
View File
@@ -286,6 +286,124 @@ pub fn parse_antigravity_usage_response(
}))
}
pub fn parse_antigravity_quota_summary_response(
value: &serde_json::Value,
) -> Option<serde_json::Value> {
let groups = value.get("groups")?.as_array()?;
let mut parsed_groups = Vec::new();
for (group_index, group) in groups.iter().enumerate() {
let Some(group) = group.as_object() else {
continue;
};
let group_id = coerce_json_string(
group
.get("groupId")
.or_else(|| group.get("group_id"))
.or_else(|| group.get("id")),
);
let display_name = coerce_json_string(
group
.get("displayName")
.or_else(|| group.get("display_name")),
)
.unwrap_or_else(|| format!("Quota group {}", group_index + 1));
let description = coerce_json_string(group.get("description"));
let mut parsed_buckets = Vec::new();
for (bucket_index, bucket) in group
.get("buckets")
.and_then(serde_json::Value::as_array)
.into_iter()
.flatten()
.enumerate()
{
let Some(bucket) = bucket.as_object() else {
continue;
};
let bucket_id = coerce_json_string(
bucket
.get("bucketId")
.or_else(|| bucket.get("bucket_id"))
.or_else(|| bucket.get("id")),
)
.unwrap_or_else(|| format!("bucket-{}", bucket_index + 1));
let window = coerce_json_string(bucket.get("window"));
let remaining_fraction = bucket
.get("remainingFraction")
.or_else(|| bucket.get("remaining_fraction"))
.and_then(coerce_json_f64)
.map(|value| value.clamp(0.0, 1.0));
let reset_time = bucket
.get("resetTime")
.or_else(|| bucket.get("reset_time"))
.cloned()
.filter(|value| !value.is_null());
let bucket_display_name = coerce_json_string(
bucket
.get("displayName")
.or_else(|| bucket.get("display_name")),
);
let bucket_description = coerce_json_string(bucket.get("description"));
if window.is_none()
&& remaining_fraction.is_none()
&& reset_time.is_none()
&& bucket_display_name.is_none()
&& bucket_description.is_none()
{
continue;
}
let mut parsed_bucket = serde_json::Map::new();
parsed_bucket.insert("bucket_id".to_string(), json!(bucket_id));
if let Some(window) = window {
parsed_bucket.insert("window".to_string(), json!(window));
}
if let Some(remaining_fraction) = remaining_fraction {
parsed_bucket.insert("remaining_fraction".to_string(), json!(remaining_fraction));
parsed_bucket.insert(
"used_percent".to_string(),
json!((1.0 - remaining_fraction) * 100.0),
);
parsed_bucket.insert(
"is_exhausted".to_string(),
json!(remaining_fraction <= 1e-6),
);
}
if let Some(reset_time) = reset_time {
parsed_bucket.insert("reset_time".to_string(), reset_time);
}
if let Some(display_name) = bucket_display_name {
parsed_bucket.insert("display_name".to_string(), json!(display_name));
}
if let Some(description) = bucket_description {
parsed_bucket.insert("description".to_string(), json!(description));
}
parsed_buckets.push(serde_json::Value::Object(parsed_bucket));
}
if parsed_buckets.is_empty() {
continue;
}
let mut parsed_group = serde_json::Map::new();
if let Some(group_id) = group_id {
parsed_group.insert("group_id".to_string(), json!(group_id));
}
parsed_group.insert("display_name".to_string(), json!(display_name));
if let Some(description) = description {
parsed_group.insert("description".to_string(), json!(description));
}
parsed_group.insert(
"buckets".to_string(),
serde_json::Value::Array(parsed_buckets),
);
parsed_groups.push(serde_json::Value::Object(parsed_group));
}
(!parsed_groups.is_empty()).then_some(serde_json::Value::Array(parsed_groups))
}
pub fn parse_gemini_cli_retrieve_user_quota_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
@@ -2042,6 +2160,13 @@ fn codex_additional_quota_windows(
let Some(item_object) = item.as_object() else {
continue;
};
if item_object
.get("limit_name")
.and_then(serde_json::Value::as_str)
.is_some_and(|name| name.trim() == CODEX_SPARK_LIMIT_NAME)
{
continue;
}
let Some(limit_name) = ["limit_name", "metered_feature", "name", "id"]
.iter()
.find_map(|key| item_object.get(*key).and_then(serde_json::Value::as_str))
@@ -3944,11 +4069,11 @@ mod tests {
codex_rate_limit_metadata_exhausted, codex_runtime_invalid_reason,
codex_websocket_response_has_usage_limit_error, codex_websocket_usage_limit_reset_at,
extract_execution_error_detail, merge_codex_quota_metadata_snapshot,
normalize_codex_reset_credit_consume_outcome, parse_antigravity_usage_response,
parse_chatgpt_web_conversation_init_response, parse_codex_backend_me_response,
parse_codex_usage_headers, parse_codex_websocket_rate_limits_response,
parse_codex_wham_reset_credits_detail_response, parse_codex_wham_usage_response,
parse_gemini_cli_retrieve_user_quota_response,
normalize_codex_reset_credit_consume_outcome, parse_antigravity_quota_summary_response,
parse_antigravity_usage_response, parse_chatgpt_web_conversation_init_response,
parse_codex_backend_me_response, parse_codex_usage_headers,
parse_codex_websocket_rate_limits_response, parse_codex_wham_reset_credits_detail_response,
parse_codex_wham_usage_response, parse_gemini_cli_retrieve_user_quota_response,
parse_gemini_cli_v1internal_credits_response, parse_windsurf_model_configs_response,
parse_windsurf_rate_limit_response, parse_windsurf_user_status_response,
provider_auto_remove_quota_exhausted_keys, quota_refresh_success_invalid_state,
@@ -6368,11 +6493,10 @@ mod tests {
.get("additional_quota_windows")
.and_then(serde_json::Value::as_array)
.expect("all additional quota windows should be preserved");
assert_eq!(additional.len(), 2);
assert!(additional.iter().all(|window| {
window.get("model") == Some(&json!("GPT-5.3-Codex-Spark"))
&& window.get("scope") == Some(&json!("model"))
}));
assert!(
additional.is_empty(),
"the Spark windows are already normalized into the dedicated spark fields"
);
}
#[test]
@@ -6916,6 +7040,37 @@ mod tests {
);
}
#[test]
fn parses_antigravity_grouped_weekly_and_five_hour_quota() {
let groups = parse_antigravity_quota_summary_response(&json!({
"groups": [{
"displayName": "Claude and GPT models",
"description": "Shared quota",
"buckets": [{
"bucketId": "3p-5h",
"window": "5h",
"remainingFraction": 0.25,
"resetTime": "2026-05-05T05:00:00Z",
"displayName": "5 hour"
}, {
"bucketId": "3p-weekly",
"window": "weekly",
"remainingFraction": 0.8,
"resetTime": "2026-05-11T00:00:00Z"
}]
}]
}))
.expect("grouped Antigravity quota should parse");
assert_eq!(groups[0]["display_name"], json!("Claude and GPT models"));
assert_eq!(groups[0]["description"], json!("Shared quota"));
assert_eq!(groups[0]["buckets"][0]["bucket_id"], json!("3p-5h"));
assert_eq!(groups[0]["buckets"][0]["window"], json!("5h"));
assert_eq!(groups[0]["buckets"][0]["remaining_fraction"], json!(0.25));
assert_eq!(groups[0]["buckets"][0]["used_percent"], json!(75.0));
assert_eq!(groups[0]["buckets"][1]["bucket_id"], json!("3p-weekly"));
}
#[test]
fn parses_gemini_cli_retrieve_user_quota_buckets() {
let parsed = parse_gemini_cli_retrieve_user_quota_response(
crates/aether-model-fetch/src/logic.rs
+22
View File
@@ -624,6 +624,15 @@ pub fn merge_upstream_metadata(current: Option<&Value>, incoming: &Value) -> Val
next_value.as_object_mut(),
merged.get(namespace).and_then(Value::as_object),
) {
if namespace.eq_ignore_ascii_case("antigravity") {
for field in ["quota_groups", "quota_groups_updated_at"] {
if !next_namespace.contains_key(field) {
if let Some(value) = old_namespace.get(field) {
next_namespace.insert(field.to_string(), value.clone());
}
}
}
}
if let (Some(new_quota), Some(old_quota)) = (
next_namespace
.get_mut("quota_by_model")
@@ -1738,6 +1747,11 @@ mod tests {
let merged = merge_upstream_metadata(
Some(&json!({
"antigravity": {
"quota_groups": [{
"display_name": "Claude and GPT models",
"buckets": [{"bucket_id": "3p-5h", "window": "5h"}]
}],
"quota_groups_updated_at": 1_777_000_000u64,
"quota_by_model": {
"gemini-2.5-pro": {
"remaining_fraction": 0.3,
@@ -1767,6 +1781,14 @@ mod tests {
assert!(merged["antigravity"]["quota_by_model"]
.get("stale-model")
.is_none());
assert_eq!(
merged["antigravity"]["quota_groups"][0]["buckets"][0]["bucket_id"],
"3p-5h"
);
assert_eq!(
merged["antigravity"]["quota_groups_updated_at"],
json!(1_777_000_000u64)
);
}
#[test]
crates/aether-model-fetch/src/strategy.rs
+72 -20
View File
@@ -1037,7 +1037,9 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec<Value>, Option
}));
let quota_payload = build_antigravity_quota_payload(model_object.get("quotaInfo"));
quota_by_model.insert(model_id.to_string(), Value::Object(quota_payload));
if !quota_payload.is_empty() {
quota_by_model.insert(model_id.to_string(), Value::Object(quota_payload));
}
}
let upstream_metadata = (!quota_by_model.is_empty()).then(|| {
@@ -1141,31 +1143,37 @@ fn infer_kiro_model_owner(model_id: &str) -> &'static str {
}
fn build_antigravity_quota_payload(quota_info: Option<&Value>) -> serde_json::Map<String, Value> {
let quota_info = quota_info.and_then(Value::as_object);
let Some(quota_info) = quota_info.and_then(Value::as_object) else {
return serde_json::Map::new();
};
let reset_time = quota_info
.and_then(|value| value.get("resetTime"))
.get("resetTime")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let remaining_fraction = quota_info
.and_then(|value| value.get("remainingFraction"))
.and_then(Value::as_f64);
.get("remainingFraction")
.and_then(|value| {
value.as_f64().or_else(|| {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| value.parse::<f64>().ok())
})
})
.filter(|value| value.is_finite())
.map(|value| value.clamp(0.0, 1.0));
let mut payload = serde_json::Map::new();
match remaining_fraction {
Some(remaining_fraction) => {
let used_percent = ((1.0 - remaining_fraction) * 100.0).clamp(0.0, 100.0);
payload.insert(
"remaining_fraction".to_string(),
Value::from(remaining_fraction),
);
payload.insert("used_percent".to_string(), Value::from(used_percent));
}
None => {
payload.insert("remaining_fraction".to_string(), Value::from(0.0));
payload.insert("used_percent".to_string(), Value::from(100.0));
}
if let Some(remaining_fraction) = remaining_fraction {
let used_percent = (1.0 - remaining_fraction) * 100.0;
payload.insert(
"remaining_fraction".to_string(),
Value::from(remaining_fraction),
);
payload.insert("used_percent".to_string(), Value::from(used_percent));
}
if let Some(reset_time) = reset_time {
payload.insert("reset_time".to_string(), Value::String(reset_time));
@@ -1618,8 +1626,8 @@ mod tests {
use super::{
build_vertex_google_list_url, build_vertex_service_account_list_url,
parse_codex_models_response_for_request, select_model_fetch_strategy, ModelFetchStrategy,
ModelFetchStrategyKind,
parse_antigravity_models_response, parse_codex_models_response_for_request,
select_model_fetch_strategy, ModelFetchStrategy, ModelFetchStrategyKind,
};
use crate::transport::ModelFetchTransportRuntime;
use crate::{fetch_models_from_transports, fetch_models_from_transports_for_client_version};
@@ -2700,6 +2708,50 @@ mod tests {
);
}
#[test]
fn antigravity_models_without_explicit_quota_are_not_marked_exhausted() {
let (models, metadata) = parse_antigravity_models_response(&json!({
"models": {
"gemini-3.7-flash-tiered": {
"displayName": "Gemini 3.7 Flash"
},
"gemini-3.7-flash-high": {
"displayName": "Gemini 3.7 Flash High",
"quotaInfo": {
"remainingFraction": "0.75",
"resetTime": "2030-01-01T00:00:00Z"
}
},
"gemini-3.7-flash-low": {
"displayName": "Gemini 3.7 Flash Low",
"quotaInfo": {
"remainingFraction": 0.0
}
}
}
}))
.expect("Antigravity models should parse");
assert_eq!(models.len(), 3);
let metadata = metadata.expect("explicit quota should produce metadata");
let antigravity = &metadata["antigravity"];
assert!(antigravity["quota_by_model"]
.get("gemini-3.7-flash-tiered")
.is_none());
assert_eq!(
antigravity["quota_by_model"]["gemini-3.7-flash-high"]["remaining_fraction"],
json!(0.75)
);
assert_eq!(
antigravity["quota_by_model"]["gemini-3.7-flash-high"]["used_percent"],
json!(25.0)
);
assert_eq!(
antigravity["quota_by_model"]["gemini-3.7-flash-low"]["used_percent"],
json!(100.0)
);
}
#[tokio::test]
async fn kiro_transport_fetches_list_available_models() {
let executed_urls = Arc::new(Mutex::new(Vec::new()));
crates/aether-model-fetch/src/transport.rs
+1 -1
View File
@@ -1185,7 +1185,7 @@ mod tests {
);
assert_eq!(
plan.headers.get("x-client-version").map(String::as_str),
Some("1.2.3")
Some("4.3.0")
);
assert_eq!(
plan.headers.get("x-vscode-sessionid").map(String::as_str),
crates/aether-provider/pool/src/lib.rs
+71 -9
View File
@@ -15,10 +15,11 @@ pub use presets::{
};
pub use provider::{ProviderPoolAdapter, ProviderPoolMemberInput};
pub use providers::{
build_antigravity_pool_quota_request, build_chatgpt_web_pool_quota_request,
build_codex_pool_quota_request, build_codex_pool_reset_credit_consume_request,
build_codex_pool_reset_credits_request, build_gemini_cli_pool_quota_request,
build_kiro_pool_quota_request, build_windsurf_pool_model_configs_request,
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
build_chatgpt_web_pool_quota_request, build_codex_pool_quota_request,
build_codex_pool_reset_credit_consume_request, build_codex_pool_reset_credits_request,
build_gemini_cli_pool_quota_request, build_kiro_pool_quota_request,
build_windsurf_pool_model_configs_request,
build_windsurf_pool_model_configs_request_with_base_url, build_windsurf_pool_quota_request,
build_windsurf_pool_quota_request_with_base_url, build_windsurf_pool_rate_limit_request,
build_windsurf_pool_rate_limit_request_with_base_url, enrich_chatgpt_web_quota_metadata,
@@ -27,11 +28,12 @@ pub use providers::{
AntigravityProviderPoolAdapter, ChatGptWebProviderPoolAdapter, CodexProviderPoolAdapter,
DefaultProviderPoolAdapter, GeminiCliProviderPoolAdapter, GrokProviderPoolAdapter,
KiroPoolQuotaAuthInput, KiroProviderPoolAdapter, UnsupportedQuotaProviderPoolAdapter,
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, CHATGPT_WEB_CONVERSATION_INIT_PATH,
CHATGPT_WEB_DEFAULT_BASE_URL, CODEX_WHAM_RESET_CREDITS_CONSUME_URL,
CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL, GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH,
GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH, WINDSURF_USER_STATUS_PATH,
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH,
CHATGPT_WEB_CONVERSATION_INIT_PATH, CHATGPT_WEB_DEFAULT_BASE_URL,
CODEX_WHAM_RESET_CREDITS_CONSUME_URL, CODEX_WHAM_RESET_CREDITS_URL, CODEX_WHAM_USAGE_URL,
GEMINI_CLI_RETRIEVE_USER_QUOTA_PATH, GEMINI_CLI_USER_AGENT, KIRO_USAGE_LIMITS_PATH,
KIRO_USAGE_SDK_VERSION, WINDSURF_MODEL_CONFIGS_PATH, WINDSURF_RATE_LIMIT_PATH,
WINDSURF_USER_STATUS_PATH,
};
pub use quota::{
provider_pool_key_account_quota_exhausted, provider_pool_key_model_quota_exhausted,
@@ -883,6 +885,66 @@ mod tests {
assert!(!available.quota_exhausted);
}
#[test]
fn antigravity_tiered_model_quota_wins_over_display_only_group_windows() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "antigravity",
"exhausted": true,
"windows": [{
"code": "model:gemini-3.7-flash-tiered",
"scope": "model",
"model": "gemini-3.7-flash-tiered",
"remaining_ratio": 0.906,
"used_ratio": 0.094,
"is_exhausted": false
}, {
"code": "group:0:3p-5h",
"scope": "quota_group",
"quota_group": "group:0",
"bucket_id": "3p-5h",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals =
service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered"));
assert!(!signals.quota_exhausted);
assert!(!signals.quota_hard_blocked);
}
#[test]
fn antigravity_tiered_variant_does_not_match_another_model_family() {
let service = ProviderPoolService::with_builtin_adapters();
let mut key = sample_key(None);
key.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "antigravity",
"exhausted": false,
"windows": [{
"code": "model:gemini-3.7-pro-tiered",
"scope": "model",
"model": "gemini-3.7-pro-tiered",
"used_ratio": 1.0,
"is_exhausted": true
}]
}
}));
let signals =
service.member_signals("antigravity", &key, None, Some("gemini-3.7-flash-tiered"));
assert!(!signals.quota_exhausted);
assert!(!signals.quota_hard_blocked);
}
#[test]
fn codex_standard_and_spark_quota_families_are_independent() {
let service = ProviderPoolService::with_builtin_adapters();
crates/aether-provider/pool/src/providers/antigravity.rs
+84
View File
@@ -12,6 +12,8 @@ use crate::quota::provider_pool_model_quota_exhausted;
use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
pub const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";
pub const ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH: &str =
"/v1internal:retrieveUserQuotaSummary";
#[derive(Debug, Clone, Default)]
pub struct AntigravityProviderPoolAdapter;
@@ -92,3 +94,85 @@ pub fn build_antigravity_pool_quota_request(
accept_invalid_certs: false,
}
}
pub fn build_antigravity_pool_quota_summary_request(
key_id: &str,
endpoint_base_url: &str,
authorization: (String, String),
project_id: Option<&str>,
mut identity_headers: BTreeMap<String, String>,
) -> ProviderPoolQuotaRequestSpec {
let mut headers = std::mem::take(&mut identity_headers);
headers.insert("authorization".to_string(), authorization.1);
headers.insert("content-type".to_string(), "application/json".to_string());
headers.insert("accept".to_string(), "application/json".to_string());
headers
.entry("user-agent".to_string())
.or_insert_with(|| "antigravity".to_string());
let json_body = project_id
.map(str::trim)
.filter(|project_id| !project_id.is_empty())
.map_or_else(|| json!({}), |project_id| json!({ "project": project_id }));
ProviderPoolQuotaRequestSpec {
request_id: format!("antigravity-quota-summary:{key_id}"),
provider_name: "antigravity".to_string(),
quota_kind: "antigravity".to_string(),
method: "POST".to_string(),
url: format!(
"{}{}",
endpoint_base_url.trim_end_matches('/'),
ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH
),
headers,
content_type: Some("application/json".to_string()),
json_body: Some(json_body),
client_api_format: "gemini:generate_content".to_string(),
provider_api_format: "antigravity:retrieve_user_quota_summary".to_string(),
model_name: Some("retrieveUserQuotaSummary".to_string()),
accept_invalid_certs: false,
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use super::{
build_antigravity_pool_quota_summary_request, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH,
};
#[test]
fn grouped_quota_request_can_retry_without_project_on_the_same_endpoint() {
let with_project = build_antigravity_pool_quota_summary_request(
"key-1",
"https://daily-cloudcode-pa.googleapis.com/",
("authorization".to_string(), "Bearer token".to_string()),
Some("project-1"),
BTreeMap::new(),
);
let without_project = build_antigravity_pool_quota_summary_request(
"key-1",
"https://daily-cloudcode-pa.googleapis.com/",
("authorization".to_string(), "Bearer token".to_string()),
None,
BTreeMap::new(),
);
assert_eq!(
with_project.url,
format!(
"https://daily-cloudcode-pa.googleapis.com{ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH}"
)
);
assert_eq!(
with_project.json_body,
Some(json!({"project": "project-1"}))
);
assert_eq!(without_project.url, with_project.url);
assert_eq!(without_project.json_body, Some(json!({})));
}
}
crates/aether-provider/pool/src/providers/mod.rs
+2 -1
View File
@@ -10,7 +10,8 @@ pub mod windsurf;
pub use antigravity::AntigravityProviderPoolAdapter;
pub use antigravity::{
build_antigravity_pool_quota_request, ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH,
build_antigravity_pool_quota_request, build_antigravity_pool_quota_summary_request,
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH, ANTIGRAVITY_RETRIEVE_USER_QUOTA_SUMMARY_PATH,
};
pub use chatgpt_web::ChatGptWebProviderPoolAdapter;
pub use chatgpt_web::{
crates/aether-provider/pool/src/quota.rs
+1
View File
@@ -398,6 +398,7 @@ fn provider_pool_is_specific_model_token(token: &str) -> bool {
"team",
"think",
"thinking",
"tiered",
"vendor",
]
.contains(&token)
crates/aether-provider/transport/src/antigravity/auth.rs
+11 -12
View File
@@ -5,8 +5,8 @@ use serde_json::Value;
use super::super::snapshot::GatewayProviderTransportSnapshot;
pub const ANTIGRAVITY_PROVIDER_TYPE: &str = "antigravity";
pub const ANTIGRAVITY_REQUEST_USER_AGENT: &str =
"antigravity/cli/1.0.16 (aidev_client; os_type=linux; arch=arm64; auth_method=consumer)";
pub const ANTIGRAVITY_CLIENT_VERSION: &str = "4.3.0";
pub const ANTIGRAVITY_REQUEST_USER_AGENT: &str = "vscode/1.X.X (Antigravity/4.3.0)";
const ANTIGRAVITY_CLIENT_NAME: &str = "antigravity";
const ANTIGRAVITY_GOOG_API_CLIENT: &str = "gl-node/18.18.2 fire/0.8.6 grpc/1.10.x";
@@ -109,7 +109,7 @@ pub fn build_antigravity_static_identity_headers(
}
pub fn build_antigravity_static_client_headers(
client_version: Option<&str>,
_client_version: Option<&str>,
session_id: Option<&str>,
) -> BTreeMap<String, String> {
let mut headers = BTreeMap::from([
@@ -125,14 +125,12 @@ pub fn build_antigravity_static_client_headers(
String::from("user-agent"),
String::from(ANTIGRAVITY_REQUEST_USER_AGENT),
),
(
String::from("x-client-version"),
String::from(ANTIGRAVITY_CLIENT_VERSION),
),
]);
if let Some(client_version) = client_version
.map(str::trim)
.filter(|value| !value.is_empty())
{
headers.insert(String::from("x-client-version"), client_version.to_string());
}
if let Some(session_id) = session_id.map(str::trim).filter(|value| !value.is_empty()) {
headers.insert(String::from("x-vscode-sessionid"), session_id.to_string());
}
@@ -273,7 +271,8 @@ mod tests {
use super::{
build_antigravity_static_client_headers, resolve_local_antigravity_request_auth,
AntigravityRequestAuth, AntigravityRequestAuthSupport, ANTIGRAVITY_REQUEST_USER_AGENT,
AntigravityRequestAuth, AntigravityRequestAuthSupport, ANTIGRAVITY_CLIENT_VERSION,
ANTIGRAVITY_REQUEST_USER_AGENT,
};
use crate::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
@@ -383,7 +382,7 @@ mod tests {
}
#[test]
fn static_client_headers_use_native_antigravity_cli_user_agent() {
fn static_client_headers_pin_the_known_good_antigravity_identity() {
let headers = build_antigravity_static_client_headers(Some("1.0.16"), Some("session-abc"));
assert_eq!(
@@ -396,7 +395,7 @@ mod tests {
);
assert_eq!(
headers.get("x-client-version").map(String::as_str),
Some("1.0.16")
Some(ANTIGRAVITY_CLIENT_VERSION)
);
assert_eq!(
headers.get("x-vscode-sessionid").map(String::as_str),
frontend/src/features/providers/components/AntigravityQuotaDialog.vue
+46 -2
View File
@@ -45,11 +45,11 @@
<div class="py-2">
<div
v-if="items.length > 0"
v-if="displayItems.length > 0"
class="grid grid-cols-2 gap-3"
>
<div
v-for="item in items"
v-for="item in displayItems"
:key="item.model"
>
<div class="flex items-center justify-between text-[10px] mb-0.5">
@@ -238,6 +238,48 @@ function buildItemsFromQuotaSnapshot(quota: QuotaStatusSnapshot | null | undefin
return dedupeAntigravityQuotaItemsByLabel(items)
}
function buildGroupedItemsFromQuotaSnapshot(
quota: QuotaStatusSnapshot | null | undefined,
): QuotaItem[] {
if (!quota) return []
const providerType = String(quota.provider_type || '').trim().toLowerCase()
if (providerType && providerType !== 'antigravity') return []
const windows = Array.isArray(quota.windows)
? quota.windows.filter(window => String(window?.scope || '').trim().toLowerCase() === 'quota_group')
: []
return windows
.map((window) => {
const code = String(window.code || '').trim()
const label = String(window.label || code).trim()
if (!code || !label) return null
const usedPercent =
typeof window.used_ratio === 'number'
? Math.max(Math.min(window.used_ratio * 100, 100), 0)
: typeof window.remaining_ratio === 'number'
? Math.max(Math.min((1 - window.remaining_ratio) * 100, 100), 0)
: null
if (usedPercent == null) return null
const remainingPercent =
typeof window.remaining_ratio === 'number'
? Math.max(Math.min(window.remaining_ratio * 100, 100), 0)
: Math.max(100 - usedPercent, 0)
return {
model: code,
label,
usedPercent,
remainingPercent,
resetSeconds: getQuotaWindowLiveResetSeconds(quota, window),
} satisfies QuotaItem
})
.filter((item): item is QuotaItem => item !== null)
}
const rawItems = computed<QuotaItem[]>(() => {
const snapshotItems = buildItemsFromQuotaSnapshot(props.quotaSnapshot)
if (snapshotItems.length > 0) return snapshotItems
@@ -291,6 +333,8 @@ const rawItems = computed<QuotaItem[]>(() => {
})
const items = computed<QuotaItem[]>(() => summarizeAntigravityQuotaItems(rawItems.value))
const groupedItems = computed<QuotaItem[]>(() => buildGroupedItemsFromQuotaSnapshot(props.quotaSnapshot))
const displayItems = computed<QuotaItem[]>(() => [...groupedItems.value, ...items.value])
async function handleTestModel(modelName: string) {
if (!props.providerId || testingModel.value) return
frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts
+32 -2
View File
@@ -2,7 +2,7 @@ import { describe, expect, it, vi } from 'vitest'
import { createApp, defineComponent, h } from 'vue'
import AntigravityQuotaDialog from '@/features/providers/components/AntigravityQuotaDialog.vue'
import type { UpstreamMetadata } from '@/api/endpoints/types'
import type { QuotaStatusSnapshot, UpstreamMetadata } from '@/api/endpoints/types'
vi.mock('@/components/ui', async () => {
const { defineComponent, h } = await import('vue')
@@ -77,7 +77,7 @@ vi.mock('@/utils/errorParser', () => ({
parseApiError: (value: unknown) => String(value),
}))
function mount(metadata: UpstreamMetadata) {
function mount(metadata: UpstreamMetadata, quotaSnapshot?: QuotaStatusSnapshot) {
const root = document.createElement('div')
document.body.appendChild(root)
@@ -86,6 +86,7 @@ function mount(metadata: UpstreamMetadata) {
return () => h(AntigravityQuotaDialog, {
open: true,
metadata,
quotaSnapshot,
keyName: 'Key-1',
})
},
@@ -102,6 +103,35 @@ function mount(metadata: UpstreamMetadata) {
}
describe('AntigravityQuotaDialog', () => {
it('renders grouped five-hour and weekly quota windows', () => {
const { root, unmount } = mount({}, {
code: 'ok',
provider_type: 'antigravity',
exhausted: false,
windows: [{
code: 'group:0:3p-5h',
label: 'Claude and GPT models · 5 hour',
scope: 'quota_group',
used_ratio: 0.75,
remaining_ratio: 0.25,
}, {
code: 'group:0:3p-weekly',
label: 'Claude and GPT models · weekly',
scope: 'quota_group',
used_ratio: 0.2,
remaining_ratio: 0.8,
}],
})
const text = root.textContent || ''
expect(text).toContain('Claude and GPT models · 5 hour')
expect(text).toContain('25.0%')
expect(text).toContain('Claude and GPT models · weekly')
expect(text).toContain('80.0%')
unmount()
})
it('hides quota buckets outside the shared pool summary families', () => {
const rawIdentifier = 'RateLimitResetCredit_05cbb6eeeb9c81918e011d8300f9ebfb'
const { root, unmount } = mount({