diff --git a/src/poller/antigravity.rs b/src/poller/antigravity.rs index f4ac41a8..3ca898cf 100644 --- a/src/poller/antigravity.rs +++ b/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 { + 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 { diff --git a/src/poller/tests.rs b/src/poller/tests.rs index e29c58b9..64014fa3 100644 --- a/src/poller/tests.rs +++ b/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); +}