diff --git a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs index 4300a1733..36e3d8b4c 100644 --- a/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs +++ b/apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs @@ -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()); diff --git a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs index 3e60995f1..d3e26fc79 100644 --- a/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs +++ b/apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs @@ -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, + proxy_override: Option<&ProxySnapshot>, +) -> Option { + 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(); diff --git a/apps/aether-gateway/src/handlers/shared/catalog.rs b/apps/aether-gateway/src/handlers/shared/catalog.rs index fe8791567..bd842358c 100644 --- a/apps/aether-gateway/src/handlers/shared/catalog.rs +++ b/apps/aether-gateway/src/handlers/shared/catalog.rs @@ -654,6 +654,112 @@ fn antigravity_model_quota_window_snapshot( Some(window) } +fn antigravity_grouped_quota_window_snapshots( + metadata: &Map, + observed_at_unix_secs: Option, +) -> Vec { + 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, 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::>() }) .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::>(); + + 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(); diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs index c2103c842..2a699a237 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local.rs @@ -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" diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs index 03471fe8b..0f1a1338c 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_cli/cross_format.rs @@ -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, diff --git a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs index 8fa7f5745..509186eab 100644 --- a/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs +++ b/apps/aether-gateway/src/tests/ai_execute/finalize_local_provider/gemini.rs @@ -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, diff --git a/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs b/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs index 7b7d98ad1..5de65b509 100644 --- a/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs @@ -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" diff --git a/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs b/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs index 770bd5846..6beb48a01 100644 --- a/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs +++ b/apps/aether-gateway/src/tests/ai_execute/sync/gemini/cli.rs @@ -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" diff --git a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs index 0c6bc508b..c4d52ce31 100644 --- a/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs +++ b/apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs @@ -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, } @@ -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::)); + let seen_execution_runtime = Arc::new(Mutex::new(Vec::::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(); diff --git a/crates/aether-admin/src/provider/quota.rs b/crates/aether-admin/src/provider/quota.rs index e70af0420..7fde73523 100644 --- a/crates/aether-admin/src/provider/quota.rs +++ b/crates/aether-admin/src/provider/quota.rs @@ -286,6 +286,124 @@ pub fn parse_antigravity_usage_response( })) } +pub fn parse_antigravity_quota_summary_response( + value: &serde_json::Value, +) -> Option { + 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( diff --git a/crates/aether-model-fetch/src/logic.rs b/crates/aether-model-fetch/src/logic.rs index b4c330257..e09f69057 100644 --- a/crates/aether-model-fetch/src/logic.rs +++ b/crates/aether-model-fetch/src/logic.rs @@ -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] diff --git a/crates/aether-model-fetch/src/strategy.rs b/crates/aether-model-fetch/src/strategy.rs index 1412f7331..64951adb8 100644 --- a/crates/aether-model-fetch/src/strategy.rs +++ b/crates/aether-model-fetch/src/strategy.rs @@ -1037,7 +1037,9 @@ fn parse_antigravity_models_response(body: &Value) -> Result<(Vec, 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 { - 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::().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())); diff --git a/crates/aether-model-fetch/src/transport.rs b/crates/aether-model-fetch/src/transport.rs index 15fa42366..1fffba4af 100644 --- a/crates/aether-model-fetch/src/transport.rs +++ b/crates/aether-model-fetch/src/transport.rs @@ -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), diff --git a/crates/aether-provider/pool/src/lib.rs b/crates/aether-provider/pool/src/lib.rs index f8c813614..7b29abc51 100644 --- a/crates/aether-provider/pool/src/lib.rs +++ b/crates/aether-provider/pool/src/lib.rs @@ -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(); diff --git a/crates/aether-provider/pool/src/providers/antigravity.rs b/crates/aether-provider/pool/src/providers/antigravity.rs index 24d9bb797..9b87ae7c6 100644 --- a/crates/aether-provider/pool/src/providers/antigravity.rs +++ b/crates/aether-provider/pool/src/providers/antigravity.rs @@ -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, +) -> 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!({}))); + } +} diff --git a/crates/aether-provider/pool/src/providers/mod.rs b/crates/aether-provider/pool/src/providers/mod.rs index 0239dde79..72734c81e 100644 --- a/crates/aether-provider/pool/src/providers/mod.rs +++ b/crates/aether-provider/pool/src/providers/mod.rs @@ -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::{ diff --git a/crates/aether-provider/pool/src/quota.rs b/crates/aether-provider/pool/src/quota.rs index f19200fc8..9d18d506c 100644 --- a/crates/aether-provider/pool/src/quota.rs +++ b/crates/aether-provider/pool/src/quota.rs @@ -398,6 +398,7 @@ fn provider_pool_is_specific_model_token(token: &str) -> bool { "team", "think", "thinking", + "tiered", "vendor", ] .contains(&token) diff --git a/crates/aether-provider/transport/src/antigravity/auth.rs b/crates/aether-provider/transport/src/antigravity/auth.rs index 52925427d..c8043a7b5 100644 --- a/crates/aether-provider/transport/src/antigravity/auth.rs +++ b/crates/aether-provider/transport/src/antigravity/auth.rs @@ -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 { 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), diff --git a/frontend/src/features/providers/components/AntigravityQuotaDialog.vue b/frontend/src/features/providers/components/AntigravityQuotaDialog.vue index 2fec683f7..0a60ba705 100644 --- a/frontend/src/features/providers/components/AntigravityQuotaDialog.vue +++ b/frontend/src/features/providers/components/AntigravityQuotaDialog.vue @@ -45,11 +45,11 @@
@@ -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(() => { const snapshotItems = buildItemsFromQuotaSnapshot(props.quotaSnapshot) if (snapshotItems.length > 0) return snapshotItems @@ -291,6 +333,8 @@ const rawItems = computed(() => { }) const items = computed(() => summarizeAntigravityQuotaItems(rawItems.value)) +const groupedItems = computed(() => buildGroupedItemsFromQuotaSnapshot(props.quotaSnapshot)) +const displayItems = computed(() => [...groupedItems.value, ...items.value]) async function handleTestModel(modelName: string) { if (!props.providerId || testingModel.value) return diff --git a/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts b/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts index 8a7bda80b..ded9c2119 100644 --- a/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts +++ b/frontend/src/features/providers/components/__tests__/AntigravityQuotaDialog.spec.ts @@ -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({