Show every Antigravity quota group, including Gemini 4 Argon

src/poller/antigravity.rs

@@ -9,15 +9,23 @@ use serde::Deserialize;
 
 use super::{build_agent, parse_iso8601, windows_credentials, PollError};
 use crate::diagnose;
-use crate::models::{UsageData, UsageSection};
+use crate::models::{limit_slug, UsageData, UsageLimit, UsageSection};
 
 const ANTIGRAVITY_CREDENTIAL_TARGET: &str = "gemini:antigravity";
 const GOOGLE_TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
 const EXPIRY_SKEW: Duration = Duration::from_secs(60);
+// The server returns the "Gemini 4 Argon" group only for this UA form, and
+// only from the daily endpoint; old CLI versions are rejected.
+const ANTIGRAVITY_USER_AGENT: &str = "antigravity/cli/1.4.0";
+// The production endpoint goes first: daily-cloudcode-pa answers
+// retrieveUserQuotaSummary with the Claude and GPT group only, so stopping
+// there hides the Gemini group.
+const ANTIGRAVITY_PRODUCTION_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com";
+const ANTIGRAVITY_DAILY_ENDPOINT: &str = "https://daily-cloudcode-pa.googleapis.com";
 const ANTIGRAVITY_ENDPOINTS: &[&str] = &[
-    "https://daily-cloudcode-pa.googleapis.com",
+    ANTIGRAVITY_PRODUCTION_ENDPOINT,
+    ANTIGRAVITY_DAILY_ENDPOINT,
     "https://daily-cloudcode-pa.sandbox.googleapis.com",
-    "https://cloudcode-pa.googleapis.com",
 ];
 
 #[derive(Deserialize)]
@@ -319,7 +327,20 @@ pub(super) fn fetch_antigravity_usage(token: &str) -> Result<UsageData, PollErro
 
     for base_url in ANTIGRAVITY_ENDPOINTS {
         match fetch_antigravity_usage_from_endpoint(base_url, token) {
-            Ok(data) => return Ok(data),
+            Ok(mut data) => {
+                if *base_url == ANTIGRAVITY_PRODUCTION_ENDPOINT {
+                    // Production never returns the Argon group, and its other
+                    // groups are the ones that match the official display.
+                    match fetch_antigravity_usage_from_endpoint(ANTIGRAVITY_DAILY_ENDPOINT, token)
+                    {
+                        Ok(daily) => merge_antigravity_limits(&mut data.limits, daily.limits),
+                        Err(error) => diagnose::log(format!(
+                            "Antigravity daily endpoint failed, skipping its extra groups: {error:?}"
+                        )),
+                    }
+                }
+                return Ok(data);
+            }
             Err(PollError::AuthRequired) => auth_error = true,
             Err(error) => last_error = error,
         }
@@ -376,7 +397,7 @@ pub(super) fn fetch_antigravity_project(
         .post(&format!("{base_url}/v1internal:loadCodeAssist"))
         .header("Authorization", &format!("Bearer {token}"))
         .header("Content-Type", "application/json")
-        .header("User-Agent", "antigravity")
+        .header("User-Agent", ANTIGRAVITY_USER_AGENT)
         .send_json(&body)
         .and_then(super::check_http_status)
     {
@@ -419,7 +440,7 @@ pub(super) fn fetch_antigravity_model_quota(
         .post(&format!("{base_url}/v1internal:fetchAvailableModels"))
         .header("Authorization", &format!("Bearer {token}"))
         .header("Content-Type", "application/json")
-        .header("User-Agent", "antigravity")
+        .header("User-Agent", ANTIGRAVITY_USER_AGENT)
         .send_json(&body)
         .and_then(super::check_http_status)
     {
@@ -469,7 +490,7 @@ pub(super) fn fetch_antigravity_quota_summary(
         .post(&format!("{base_url}/v1internal:retrieveUserQuotaSummary"))
         .header("Authorization", &format!("Bearer {token}"))
         .header("Content-Type", "application/json")
-        .header("User-Agent", "antigravity")
+        .header("User-Agent", ANTIGRAVITY_USER_AGENT)
         .send_json(&body)
         .and_then(super::check_http_status)
     {
@@ -520,22 +541,76 @@ pub(super) fn antigravity_section_from_summary_bucket(
 pub(super) fn antigravity_usage_from_summary(
     response: AntigravityQuotaSummaryResponse,
 ) -> Option<UsageData> {
+    let groups = response.groups.unwrap_or_default();
+    let limits = antigravity_limits_from_summary_groups(&groups);
+    let mut primary = None;
     let mut fallback = None;
 
-    for group in response.groups.unwrap_or_default() {
+    for group in groups {
         let is_gemini = is_antigravity_gemini_summary_group(&group);
         let usage = antigravity_usage_from_summary_group(group);
 
         if is_gemini && usage.is_some() {
-            return usage;
+            primary = usage;
+            break;
         }
 
         if fallback.is_none() {
             fallback = usage;
         }
     }
 
-    fallback
+    let mut usage = primary.or(fallback)?;
+    usage.limits = limits;
+    Some(usage)
+}
+
+// The primary 5h/weekly pair shows only one group; themes need every group
+// (e.g. "Claude and GPT", "Gemini 4 Argon") to show per-group quota.
+pub(super) fn antigravity_limits_from_summary_groups(
+    groups: &[AntigravityQuotaSummaryGroup],
+) -> Vec<UsageLimit> {
+    let mut limits = Vec::new();
+    for group in groups {
+        let Some(name) = group
+            .display_name
+            .as_deref()
+            .filter(|name| !name.trim().is_empty())
+        else {
+            continue;
+        };
+        for bucket in group.buckets.iter().flatten() {
+            let kind = match bucket.window.as_deref() {
+                Some(window) if window.eq_ignore_ascii_case("5h") => "session",
+                Some(window) if window.eq_ignore_ascii_case("weekly") => "weekly",
+                _ => continue,
+            };
+            let Some(usage) = antigravity_section_from_summary_bucket(bucket) else {
+                continue;
+            };
+            limits.push(UsageLimit {
+                key: format!("{}_{kind}", limit_slug(name)),
+                kind: kind.to_string(),
+                label: name.to_string(),
+                model: Some(name.to_string()),
+                model_id: bucket.bucket_id.clone(),
+                scope: None,
+                is_active: false,
+                usage,
+            });
+        }
+    }
+    limits
+}
+
+// Adds the groups only `extra` has; a group both sides know keeps `base`'s buckets.
+pub(super) fn merge_antigravity_limits(base: &mut Vec<UsageLimit>, extra: Vec<UsageLimit>) {
+    let known: Vec<Option<String>> = base.iter().map(|limit| limit.model.clone()).collect();
+    base.extend(
+        extra
+            .into_iter()
+            .filter(|limit| !known.contains(&limit.model)),
+    );
 }
 
 pub(super) fn antigravity_usage_from_summary_group(
@@ -633,6 +708,23 @@ mod auth_tests {
         }
     }
 
+    #[test]
+    fn production_endpoint_is_tried_before_daily() {
+        assert_eq!(
+            ANTIGRAVITY_ENDPOINTS.first(),
+            Some(&"https://cloudcode-pa.googleapis.com")
+        );
+    }
+
+    #[test]
+    fn user_agent_uses_the_cli_form_that_unlocks_argon() {
+        assert!(ANTIGRAVITY_USER_AGENT.starts_with("antigravity/cli/"));
+        assert_eq!(
+            ANTIGRAVITY_ENDPOINTS.get(1),
+            Some(&ANTIGRAVITY_DAILY_ENDPOINT)
+        );
+    }
+
     #[test]
     fn valid_token_uses_existing_request() {
         let calls = Cell::new(0);

src/poller/tests.rs

@@ -405,6 +405,122 @@ fn antigravity_summary_prefers_gemini_group() {
     assert!(usage.session.resets_at.is_some());
 }
 
+#[test]
+fn antigravity_summary_exposes_every_group_as_limits() {
+    let response: antigravity::AntigravityQuotaSummaryResponse = serde_json::from_str(
+        r#"{
+                "groups": [
+                    {
+                        "displayName": "Gemini Models",
+                        "buckets": [
+                            {"bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.75},
+                            {"bucketId": "gemini-weekly", "window": "weekly", "remainingFraction": 0.5}
+                        ]
+                    },
+                    {
+                        "displayName": "Claude and GPT",
+                        "buckets": [
+                            {"bucketId": "3p-5h", "window": "5h", "remainingFraction": 1},
+                            {"bucketId": "3p-weekly", "window": "weekly", "remainingFraction": 0}
+                        ]
+                    },
+                    {
+                        "displayName": "Gemini 4 Argon",
+                        "buckets": [
+                            {"bucketId": "argon-weekly", "window": "weekly", "remainingFraction": 0.9},
+                            {"bucketId": "argon-daily", "window": "daily", "remainingFraction": 0.1}
+                        ]
+                    }
+                ]
+            }"#,
+    )
+    .expect("summary response should deserialize");
+
+    let usage = antigravity::antigravity_usage_from_summary(response)
+        .expect("Gemini quota should be selected");
+
+    assert_eq!(usage.session.percentage, 25.0);
+    assert_eq!(usage.weekly.percentage, 50.0);
+    let expected = [
+        ("gemini_models_session", 25.0),
+        ("gemini_models_weekly", 50.0),
+        ("claude_and_gpt_session", 0.0),
+        ("claude_and_gpt_weekly", 100.0),
+        ("gemini_4_argon_weekly", 10.0),
+    ];
+    assert_eq!(usage.limits.len(), expected.len());
+    for (limit, (key, percentage)) in usage.limits.iter().zip(expected) {
+        assert_eq!(limit.key, key);
+        assert!((limit.usage.percentage - percentage).abs() < 0.000001);
+    }
+}
+
+fn antigravity_limits_from_json(json: &str) -> Vec<crate::models::UsageLimit> {
+    let response: antigravity::AntigravityQuotaSummaryResponse =
+        serde_json::from_str(json).expect("summary response should deserialize");
+    antigravity::antigravity_usage_from_summary(response)
+        .expect("summary should have quota")
+        .limits
+}
+
+const ANTIGRAVITY_PRODUCTION_SUMMARY: &str = r#"{"groups": [
+    {"displayName": "Gemini Models", "buckets": [
+        {"bucketId": "gemini-5h", "window": "5h", "remainingFraction": 1},
+        {"bucketId": "gemini-weekly", "window": "weekly", "remainingFraction": 0.975}]},
+    {"displayName": "Claude and GPT models", "buckets": [
+        {"bucketId": "3p-weekly", "window": "weekly", "remainingFraction": 0.5}]}
+]}"#;
+
+const ANTIGRAVITY_DAILY_SUMMARY: &str = r#"{"groups": [
+    {"displayName": "Gemini Models", "buckets": [
+        {"bucketId": "gemini-5h", "window": "5h", "remainingFraction": 1},
+        {"bucketId": "gemini-weekly", "window": "weekly", "remainingFraction": 1}]},
+    {"displayName": "Claude and GPT models", "buckets": [
+        {"bucketId": "3p-weekly", "window": "weekly", "remainingFraction": 0.1}]},
+    {"displayName": "Gemini 4 Argon", "buckets": [
+        {"bucketId": "gemini-argon-wtu-weekly", "window": "weekly", "remainingFraction": 0.816}]}
+]}"#;
+
+#[test]
+fn antigravity_merge_prefers_production_and_adds_daily_only_groups() {
+    let mut limits = antigravity_limits_from_json(ANTIGRAVITY_PRODUCTION_SUMMARY);
+    antigravity::merge_antigravity_limits(
+        &mut limits,
+        antigravity_limits_from_json(ANTIGRAVITY_DAILY_SUMMARY),
+    );
+
+    let percentage = |key: &str| {
+        limits
+            .iter()
+            .find(|limit| limit.key == key)
+            .unwrap_or_else(|| panic!("missing {key}"))
+            .usage
+            .percentage
+    };
+    assert_eq!(limits.len(), 4);
+    assert!((percentage("gemini_models_weekly") - 2.5).abs() < 0.000001);
+    assert!((percentage("claude_and_gpt_models_weekly") - 50.0).abs() < 0.000001);
+    assert!((percentage("gemini_4_argon_weekly") - 18.4).abs() < 0.000001);
+}
+
+#[test]
+fn antigravity_merge_with_failed_endpoint_keeps_the_other_side() {
+    let production = antigravity_limits_from_json(ANTIGRAVITY_PRODUCTION_SUMMARY);
+
+    let mut only_production = production.clone();
+    antigravity::merge_antigravity_limits(&mut only_production, Vec::new());
+    assert_eq!(only_production.len(), production.len());
+
+    let mut only_daily = Vec::new();
+    antigravity::merge_antigravity_limits(
+        &mut only_daily,
+        antigravity_limits_from_json(ANTIGRAVITY_DAILY_SUMMARY),
+    );
+    assert!(only_daily
+        .iter()
+        .any(|limit| limit.key == "gemini_4_argon_weekly"));
+}
+
 #[test]
 fn one_provider_failing_does_not_blank_its_row() {
     let previous: AppUsageData = [