diff --git a/src/poller/antigravity.rs b/src/poller/antigravity.rs index 7b503ab9..5ac3606a 100644 --- a/src/poller/antigravity.rs +++ b/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 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,14 +541,18 @@ pub(super) fn antigravity_section_from_summary_bucket( pub(super) fn antigravity_usage_from_summary( response: AntigravityQuotaSummaryResponse, ) -> Option { + 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() { @@ -535,7 +560,57 @@ pub(super) fn antigravity_usage_from_summary( } } - 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 { + 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, extra: Vec) { + let known: Vec> = 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); diff --git a/src/poller/tests.rs b/src/poller/tests.rs index ff8d515d..6fed4199 100644 --- a/src/poller/tests.rs +++ b/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 { + 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 = [