Show every Antigravity quota group, including Gemini 4 Argon

src/poller/antigravity.rs

@@ -9,12 +9,14 @@ 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);
-const ANTIGRAVITY_USER_AGENT: &str = "antigravity/2.24.0";
+// The daily summary returns the "Gemini 4 Argon" group only for the agy CLI
+// form of the UA; "antigravity", "antigravity/cli/0.0.1" and other forms omit it.
+const ANTIGRAVITY_USER_AGENT: &str = "antigravity/cli/1.4.0";
 const ANTIGRAVITY_ENDPOINTS: &[&str] = &[
     "https://daily-cloudcode-pa.googleapis.com",
     "https://cloudcode-pa.googleapis.com",
@@ -599,8 +601,11 @@ pub(super) fn antigravity_usage_from_summary(
     // failure to be papered over with model quotas or another endpoint.
     let answered = response.groups.is_some();
     let groups = response.groups.unwrap_or_default();
+    let limits = antigravity_limits_from_summary_groups(&groups);
     for group in groups {
-        let is_gemini = is_antigravity_gemini_summary_group(&group);
+        // Argon has its own pool; it must not stand in for the main Gemini pool.
+        let is_gemini = is_antigravity_gemini_summary_group(&group)
+            && !is_antigravity_separate_gemini_group(&group);
         let is_external = is_antigravity_external_summary_group(&group);
         has_named_pool |= is_gemini || is_external;
         let usage = antigravity_usage_from_summary_group(group);
@@ -625,20 +630,69 @@ pub(super) fn antigravity_usage_from_summary(
         gemini_usage = other_groups.pop();
     }
 
-    if let Some(mut gemini) = gemini_usage {
+    let mut usage = if let Some(mut gemini) = gemini_usage {
         if let Some(external) = external_usage {
             gemini.external_pool = Some(Box::new(external));
         }
-        Some(gemini)
+        gemini
     } else if let Some(external) = external_usage {
         let mut gemini = UsageData::default();
         gemini.session.available = false;
         gemini.weekly.available = false;
         gemini.external_pool = Some(Box::new(external));
-        Some(gemini)
+        gemini
+    } else if answered {
+        UsageData::default()
     } else {
-        answered.then(UsageData::default)
+        return None;
+    };
+    usage.limits = limits;
+    Some(usage)
+}
+
+pub(super) fn is_antigravity_separate_gemini_group(group: &AntigravityQuotaSummaryGroup) -> bool {
+    group
+        .display_name
+        .as_deref()
+        .is_some_and(|name| name.to_ascii_lowercase().contains("argon"))
+}
+
+// The main and external pools show two groups only; themes need every group
+// (e.g. "Gemini 4 Argon") and every window 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
 }
 
 pub(super) fn is_antigravity_external_summary_group(group: &AntigravityQuotaSummaryGroup) -> bool {

src/poller/tests.rs

@@ -838,3 +838,41 @@ fn failed_antigravity_poll_preserves_and_marks_both_pools_stale() {
         "stale external resets must not cause a poll loop"
     );
 }
+
+#[test]
+fn antigravity_summary_exposes_every_group_and_window_as_limits() {
+    let summary = serde_json::from_value(serde_json::json!({
+        "groups": [
+            {"displayName": "Gemini 4 Argon",
+             "buckets": [{"bucketId": "gemini-4-argon", "window": "weekly", "remainingFraction": 0.9}]},
+            {"displayName": "Claude and GPT models",
+             "buckets": [{"window": "5h", "remainingFraction": 0.4},
+                         {"window": "weekly", "remainingFraction": 0.2}]},
+            {"displayName": "Gemini Models",
+             "buckets": [{"window": "5h", "remainingFraction": 0.85},
+                         {"window": "weekly", "remainingFraction": 0.70}]}
+        ]
+    }))
+    .unwrap();
+    let data = super::antigravity::antigravity_usage_from_summary(summary).unwrap();
+
+    // Argon comes first but has its own pool; the main pool stays "Gemini Models".
+    assert!((data.session.percentage - 15.0).abs() < 1e-4);
+    assert!((data.weekly.percentage - 30.0).abs() < 1e-4);
+
+    let keys: Vec<&str> = data.limits.iter().map(|limit| limit.key.as_str()).collect();
+    assert_eq!(
+        keys,
+        [
+            "gemini_4_argon_weekly",
+            "claude_and_gpt_models_session",
+            "claude_and_gpt_models_weekly",
+            "gemini_models_session",
+            "gemini_models_weekly",
+        ]
+    );
+    let argon = &data.limits[0];
+    assert_eq!(argon.kind, "weekly");
+    assert_eq!(argon.model_id.as_deref(), Some("gemini-4-argon"));
+    assert!((argon.usage.percentage - 10.0).abs() < 1e-4);
+}