From fa2bec67d1402844046fd1a716e419a660461328 Mon Sep 17 00:00:00 2001 From: fix2it Date: Sun, 18 Oct 2026 22:37:08 +0300 Subject: [PATCH] fix(antigravity): read the Gemini pool from the quota summary reliably The summary picked the first group with windows when the Gemini group did not match, so the widget showed the "Claude and GPT models" quota under the Antigravity name. - The first endpoint that returns retrieveUserQuotaSummary is authoritative and returned at once; endpoints are no longer merged and fetchAvailableModels is skipped when the summary answered. - Groups are classified as Gemini or external (claude/gpt/3p-* buckets); an unrecognised group is used only when it is the sole group. - A summary without recognised windows (or with an empty group list) is "no limit data", not a failure. - The model-quota fallback uses Gemini models only, because per-model quotas do not say which window they belong to. - Send an explicit Antigravity User-Agent; try the sandbox endpoint last. - Absent windows are cached with `percentage: null` so external readers see "unknown" rather than 0% used. Co-Authored-By: Claude Opus 5.5 --- src/models.rs | 33 +++++- src/poller/antigravity.rs | 203 +++++++++++++++++++++++++++------- src/poller/tests.rs | 224 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 416 insertions(+), 44 deletions(-) diff --git a/src/models.rs b/src/models.rs index 3e11d9a7..1cddc13c 100644 --- a/src/models.rs +++ b/src/models.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Deserializer, Serialize, Serializer}; use crate::providers::ProviderId; -#[derive(Clone, Debug, Default, PartialEq, Serialize)] +#[derive(Clone, Debug, Default, PartialEq)] pub struct UsageSection { /// The provider reported this window, even if unused and without a reset time. pub available: bool, @@ -14,6 +14,21 @@ pub struct UsageSection { pub resets_at: Option, } +impl Serialize for UsageSection { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + // An absent window carries no figure: external cache readers that only + // look at `percentage` must see "unknown", not 0% used. + let mut map = serializer.serialize_map(Some(3))?; + map.serialize_entry("available", &self.available)?; + map.serialize_entry("percentage", &self.available.then_some(self.percentage))?; + map.serialize_entry("resets_at", &self.resets_at)?; + map.end() + } +} + impl<'de> Deserialize<'de> for UsageSection { fn deserialize(deserializer: D) -> Result where @@ -22,17 +37,18 @@ impl<'de> Deserialize<'de> for UsageSection { #[derive(Deserialize)] struct StoredSection { available: Option, - percentage: f64, + percentage: Option, resets_at: Option, } let stored = StoredSection::deserialize(deserializer)?; + let percentage = stored.percentage.unwrap_or(0.0); Ok(Self { // Older caches lost the distinction between an idle window and an // absent one. Preserve evidence of presence until a fresh poll. available: stored .available - .unwrap_or(stored.resets_at.is_some() || stored.percentage != 0.0), - percentage: stored.percentage, + .unwrap_or(stored.resets_at.is_some() || percentage != 0.0), + percentage, resets_at: stored.resets_at, }) } @@ -464,4 +480,13 @@ mod tests { assert!(decoded.get(ProviderId::Antigravity).is_none()); assert!(decoded.get(ProviderId::OpenCode).is_none()); } + + #[test] + fn absent_window_is_cached_without_a_percentage() { + let section = UsageSection::default(); + let json = serde_json::to_value(§ion).unwrap(); + assert!(json["percentage"].is_null(), "readers must not see 0% for an absent window"); + let decoded: UsageSection = serde_json::from_value(json).unwrap(); + assert_eq!(decoded, section); + } } diff --git a/src/poller/antigravity.rs b/src/poller/antigravity.rs index 7b503ab9..9c427139 100644 --- a/src/poller/antigravity.rs +++ b/src/poller/antigravity.rs @@ -14,10 +14,11 @@ use crate::models::{UsageData, 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"; const ANTIGRAVITY_ENDPOINTS: &[&str] = &[ "https://daily-cloudcode-pa.googleapis.com", - "https://daily-cloudcode-pa.sandbox.googleapis.com", "https://cloudcode-pa.googleapis.com", + "https://daily-cloudcode-pa.sandbox.googleapis.com", ]; #[derive(Deserialize)] @@ -44,17 +45,17 @@ struct AntigravityLoadResponse { } #[derive(Deserialize)] -struct AntigravityModelsResponse { - models: HashMap, +pub(super) struct AntigravityModelsResponse { + pub(super) models: HashMap, } #[derive(Deserialize)] -struct AntigravityModelInfo { +pub(super) struct AntigravityModelInfo { #[serde(rename = "quotaInfo")] - quota_info: Option, + pub(super) quota_info: Option, } -#[derive(Deserialize)] +#[derive(Clone, Deserialize)] pub(super) struct AntigravityQuotaInfo { #[serde(rename = "remainingFraction")] remaining_fraction: Option, @@ -314,17 +315,37 @@ pub(super) fn antigravity_credential_watch_signature() -> String { } pub(super) fn fetch_antigravity_usage(token: &str) -> Result { + select_antigravity_usage(ANTIGRAVITY_ENDPOINTS, |base_url| { + fetch_antigravity_usage_from_endpoint(base_url, token) + }) +} + +pub(super) fn select_antigravity_usage( + endpoints: &[&str], + mut fetch: impl FnMut(&str) -> Result<(UsageData, AntigravitySource), PollError>, +) -> Result { + let mut model_fallback: Option = None; let mut auth_error = false; let mut last_error = PollError::RequestFailed; - for base_url in ANTIGRAVITY_ENDPOINTS { - match fetch_antigravity_usage_from_endpoint(base_url, token) { - Ok(data) => return Ok(data), + for base_url in endpoints { + match fetch(base_url) { + // A quota summary names every window it has; it is authoritative, + // so later endpoints are neither asked nor merged. + Ok((data, AntigravitySource::Summary)) => return Ok(data), + // Model quotas do not name their window; keep them only as a last + // resort in case no endpoint answers with a summary. + Ok((data, AntigravitySource::Models)) => { + model_fallback.get_or_insert(data); + } Err(PollError::AuthRequired) => auth_error = true, Err(error) => last_error = error, } } + if let Some(data) = model_fallback { + return Ok(data); + } if auth_error { Err(PollError::AuthRequired) } else { @@ -332,35 +353,70 @@ pub(super) fn fetch_antigravity_usage(token: &str) -> Result Result { +) -> Result<(UsageData, AntigravitySource), PollError> { let project = fetch_antigravity_project(base_url, token)?; - if let Some(project) = project.as_deref() { + let summary_result = if let Some(project) = project.as_deref() { match fetch_antigravity_quota_summary(base_url, token, project) { - Ok(data) => return Ok(data), + Ok(data) => Some(data), Err(PollError::AuthRequired) => return Err(PollError::AuthRequired), - Err(error) => diagnose::log(format!( - "Antigravity retrieveUserQuotaSummary failed, falling back to model quota: {error:?}" - )), + Err(error) => { + diagnose::log(format!( + "Antigravity retrieveUserQuotaSummary failed: {error:?}" + )); + None + } } + } else { + None + }; + + // Per-model quotas do not say which window (5h or weekly) they belong to, + // so when the summary answered it is used alone and the models request is + // skipped. + if let Some(data) = summary_result { + return Ok((data, AntigravitySource::Summary)); + } + let models_result = fetch_antigravity_models_response(base_url, token, project.as_deref()); + + match models_result { + Ok(models) => antigravity_usage_from_models(&models) + .map(|data| (data, AntigravitySource::Models)) + .ok_or(PollError::RequestFailed), + Err(err) => Err(err), } +} - let session = fetch_antigravity_model_quota(base_url, token, project.as_deref())?; - let weekly = UsageSection::default(); - - Ok(UsageData { - limits: Vec::new(), - session, - weekly, - weekly_label: None, - monthly: None, - credits: None, - stale: false, +/// Fallback when the quota summary is unavailable: Gemini models only. +pub(super) fn antigravity_usage_from_models(models: &AntigravityModelsResponse) -> Option { + let gemini = best_gemini_model_quota(models); + // Claude/GPT model quotas do not say whether they are the 5h or the weekly + // window, so the external pool is left as "no limit data" here rather + // than guessed. + Some(UsageData { + session: gemini?, + ..Default::default() }) } +pub(super) fn best_gemini_model_quota(models: &AntigravityModelsResponse) -> Option { + best_antigravity_section(models.models.iter().filter_map(|(model, info)| { + if !model.starts_with("gemini") { + return None; + } + let quota = info.quota_info.as_ref().cloned()?; + antigravity_section_from_quota(quota) + })) +} + pub(super) fn fetch_antigravity_project( base_url: &str, token: &str, @@ -376,7 +432,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) { @@ -404,11 +460,11 @@ pub(super) fn fetch_antigravity_project( Ok(response.project.filter(|project| !project.is_empty())) } -pub(super) fn fetch_antigravity_model_quota( +pub(super) fn fetch_antigravity_models_response( base_url: &str, token: &str, project: Option<&str>, -) -> Result { +) -> Result { let agent = build_agent()?; let body = match project { Some(project) => serde_json::json!({ "project": project }), @@ -419,7 +475,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) { @@ -447,6 +503,17 @@ pub(super) fn fetch_antigravity_model_quota( } }; + Ok(response) +} + +#[allow(dead_code)] +pub(super) fn fetch_antigravity_model_quota( + base_url: &str, + token: &str, + project: Option<&str>, +) -> Result { + let response = fetch_antigravity_models_response(base_url, token, project)?; + best_antigravity_section(response.models.into_iter().filter_map(|(model, info)| { let quota = info.quota_info?; if !is_antigravity_display_model(&model) { @@ -469,7 +536,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 +587,77 @@ pub(super) fn antigravity_section_from_summary_bucket( pub(super) fn antigravity_usage_from_summary( response: AntigravityQuotaSummaryResponse, ) -> Option { - let mut fallback = None; - - for group in response.groups.unwrap_or_default() { + let mut gemini_usage = None; + let mut external_usage = None; + let mut other_groups = Vec::new(); + + // A summary that returns a group list (even an empty one) answered + // authoritatively: no recognised window means "no limit data", not a + // failure to be papered over with model quotas or another endpoint. + let answered = response.groups.is_some(); + let groups = response.groups.unwrap_or_default(); + for group in groups { let is_gemini = is_antigravity_gemini_summary_group(&group); + let is_external = is_antigravity_external_summary_group(&group); let usage = antigravity_usage_from_summary_group(group); - if is_gemini && usage.is_some() { - return usage; + if is_gemini && gemini_usage.is_none() { + gemini_usage = usage; + } else if is_external && external_usage.is_none() { + external_usage = usage; + } else if let Some(u) = usage { + other_groups.push(u); } + } - if fallback.is_none() { - fallback = usage; - } + // An unrecognised group is assigned only when it is the sole group with + // quota and no named pool exists; otherwise a position guess could show one + // pool's usage under the other's label. + if gemini_usage.is_none() && external_usage.is_none() && other_groups.len() == 1 { + gemini_usage = other_groups.pop(); } - fallback + // The external "Claude and GPT models" group is recognised only so it is + // never shown as the Gemini pool. + match gemini_usage { + Some(gemini) => Some(gemini), + None => answered.then(UsageData::default), + } +} + +pub(super) fn is_antigravity_external_summary_group(group: &AntigravityQuotaSummaryGroup) -> bool { + group + .display_name + .as_deref() + .is_some_and(|name| { + let n = name.to_ascii_lowercase(); + n.contains("claude") || n.contains("gpt") || n.contains("external") + }) + || group + .description + .as_deref() + .is_some_and(|description| { + let d = description.to_ascii_lowercase(); + d.contains("claude") || d.contains("gpt") || d.contains("external") + }) + || group.buckets.as_ref().is_some_and(|buckets| { + buckets.iter().any(|bucket| { + bucket + .bucket_id + .as_deref() + .is_some_and(|id| { + let id = id.to_ascii_lowercase(); + id.starts_with("claude-") || id.starts_with("gpt-") || id.starts_with("3p-") + }) + || bucket + .display_name + .as_deref() + .is_some_and(|name| { + let n = name.to_ascii_lowercase(); + n.contains("claude") || n.contains("gpt") + }) + }) + }) } pub(super) fn antigravity_usage_from_summary_group( @@ -826,3 +948,4 @@ mod auth_tests { assert_eq!(current.token.refresh_token.as_deref(), Some("refresh")); } } + diff --git a/src/poller/tests.rs b/src/poller/tests.rs index ff8d515d..866ff7e0 100644 --- a/src/poller/tests.rs +++ b/src/poller/tests.rs @@ -468,3 +468,227 @@ fn all_failed_providers_can_carry_their_previous_readings() { assert_eq!(claude.session.percentage, 21.0); assert!(claude.stale, "the carried reading must be marked stale"); } + +#[test] +fn antigravity_gemini_weekly_reset_projection_tests() { + use std::time::{Duration, SystemTime}; + + fn format_iso8601_for_test(time: SystemTime) -> String { + let secs = time + .duration_since(SystemTime::UNIX_EPOCH) + .unwrap() + .as_secs(); + let days = secs / 86400; + let rem = secs % 86400; + let hours = rem / 3600; + let minutes = (rem % 3600) / 60; + let seconds = rem % 60; + let mut year = 1970; + let mut d = days; + loop { + let leap = (year % 4 == 0 && year % 100 != 0) || (year % 400 == 0); + let y_days = if leap { 366 } else { 365 }; + if d < y_days { + let mut month = 1; + let month_days = [ + 31, + if leap { 29 } else { 28 }, + 31, + 30, + 31, + 30, + 31, + 31, + 30, + 31, + 30, + 31, + ]; + for &m_len in &month_days { + if d < m_len { + return format!( + "{year:04}-{month:02}-{:02}T{hours:02}:{minutes:02}:{seconds:02}Z", + d + 1 + ); + } + d -= m_len; + month += 1; + } + } + d -= y_days; + year += 1; + } + } + + let now = SystemTime::now(); + + // 1. A weekly reset that is close is reported as-is: the API's resetTime is + // authoritative and no "+6 days" projection is applied. + let reset_16h = now + Duration::from_secs(16 * 3600); + let reset_iso_16h = format_iso8601_for_test(reset_16h); + let bucket_json = serde_json::json!({ + "bucketId": "gemini-weekly", + "displayName": "Weekly Limit", + "window": "weekly", + "resetTime": reset_iso_16h, + "remainingFraction": 0.80 + }); + let bucket: antigravity::AntigravityQuotaSummaryBucket = + serde_json::from_value(bucket_json).unwrap(); + let section = antigravity::antigravity_section_from_summary_bucket(&bucket) + .expect("section exists"); + assert!((section.percentage - 20.0).abs() < 0.001); + let reported_reset = section.resets_at.expect("reset exists"); + let diff = reported_reset + .duration_since(reset_16h) + .unwrap_or_else(|e| e.duration()); + assert!(diff < Duration::from_secs(2), "close weekly reset must not be projected"); + + // 2. gemini-weekly with reset in the past (-1 hour) + let reset_past = now - Duration::from_secs(3600); + let reset_iso_past = format_iso8601_for_test(reset_past); + let bucket_past_json = serde_json::json!({ + "bucketId": "gemini-weekly", + "window": "weekly", + "resetTime": reset_iso_past, + "remainingFraction": 0.80 + }); + let bucket_past: antigravity::AntigravityQuotaSummaryBucket = + serde_json::from_value(bucket_past_json).unwrap(); + let section_past = antigravity::antigravity_section_from_summary_bucket(&bucket_past) + .expect("section exists"); + let past_reset = section_past.resets_at.expect("reset exists"); + let diff_past = past_reset + .duration_since(reset_past) + .unwrap_or_else(|e| e.duration()); + assert!(diff_past < Duration::from_secs(2), "Past reset must not have +6d added"); + + // 3. 3p-weekly with reset in 16 hours (MUST NOT have +6d added) + let bucket_3p_json = serde_json::json!({ + "bucketId": "3p-weekly", + "window": "weekly", + "resetTime": reset_iso_16h, + "remainingFraction": 0.50 + }); + let bucket_3p: antigravity::AntigravityQuotaSummaryBucket = + serde_json::from_value(bucket_3p_json).unwrap(); + let section_3p = antigravity::antigravity_section_from_summary_bucket(&bucket_3p) + .expect("section exists"); + let reset_3p = section_3p.resets_at.expect("reset exists"); + let diff_3p = reset_3p + .duration_since(reset_16h) + .unwrap_or_else(|e| e.duration()); + assert!(diff_3p < Duration::from_secs(2), "3p-weekly must not have +6d added"); +} + +#[test] +fn antigravity_models_quota_extracts_gemini_and_external_correctly() { + let models_json = r#"{ + "models": { + "gemini-2.5-flash": { + "quotaInfo": { + "remainingFraction": 0.66, + "resetTime": "2026-10-17T16:00:00Z" + } + }, + "gemini-3-flash": { + "quotaInfo": { + "remainingFraction": 0.80, + "resetTime": "2026-10-17T16:00:00Z" + } + }, + "claude-sonnet-4-6": { + "quotaInfo": { + "remainingFraction": 0.02, + "resetTime": "2099-10-22T18:00:00Z" + } + } + } + }"#; + let resp: antigravity::AntigravityModelsResponse = serde_json::from_str(models_json).unwrap(); + let gemini = antigravity::best_gemini_model_quota(&resp).expect("gemini quota exists"); + assert!((gemini.percentage - 34.0).abs() < 0.001); + +} + +#[test] +fn antigravity_first_summary_wins_and_model_quotas_are_only_a_last_resort() { + use antigravity::AntigravitySource::{Models, Summary}; + use std::cell::RefCell; + + let with_session = |percentage| UsageData { + session: UsageSection { + available: true, + percentage, + resets_at: None, + }, + ..Default::default() + }; + let endpoints = ["a", "b", "c"]; + + // A weekly-only summary is final: the next endpoint's model quota must not + // fill its 5h window, and later endpoints are not asked at all. + let asked = RefCell::new(Vec::new()); + let data = antigravity::select_antigravity_usage(&endpoints, |url| { + asked.borrow_mut().push(url.to_string()); + match url { + "a" => Err(PollError::RequestFailed), + "b" => Ok((UsageData::default(), Summary)), + _ => Ok((with_session(80.0), Models)), + } + }) + .unwrap(); + assert!(!data.session.available); + assert_eq!(*asked.borrow(), ["a", "b"]); + + // Without any summary the model quota is used. + let data = antigravity::select_antigravity_usage(&endpoints, |url| match url { + "a" => Ok((with_session(30.0), Models)), + _ => Err(PollError::RequestFailed), + }) + .unwrap(); + assert!((data.session.percentage - 30.0).abs() < 0.001); +} + +#[test] +fn antigravity_model_fallback_reports_only_gemini_models() { + let models: antigravity::AntigravityModelsResponse = serde_json::from_str( + r#"{"models":{ + "gemini-3-pro":{"quotaInfo":{"remainingFraction":0.9}}, + "claude-opus-4-6-thinking":{"quotaInfo":{"remainingFraction":0.1}}, + "gpt-oss-120b-medium":{"quotaInfo":{"remainingFraction":0.5}}, + "imagen-4":{"quotaInfo":{"remainingFraction":0.0}} + }}"#, + ) + .unwrap(); + let usage = antigravity::antigravity_usage_from_models(&models).expect("usage"); + assert!((usage.session.percentage - 10.0).abs() < 0.001, "Gemini pool must ignore Claude/GPT/image models"); +} + +#[test] +fn antigravity_summary_does_not_guess_pools_by_position() { + // An unnamed group with 3p-* buckets is the external pool, never Gemini. + let response: antigravity::AntigravityQuotaSummaryResponse = serde_json::from_str( + r#"{"groups":[{"buckets":[ + {"bucketId":"3p-weekly","window":"weekly","remainingFraction":0.4}, + {"bucketId":"3p-5h","window":"5h","remainingFraction":1} + ]}]}"#, + ) + .unwrap(); + let usage = antigravity::antigravity_usage_from_summary(response).expect("usage"); + assert!(!usage.session.available && !usage.weekly.available, "Gemini pool must stay empty"); +} + +#[test] +fn antigravity_summary_without_windows_is_an_answer_not_a_failure() { + let response: antigravity::AntigravityQuotaSummaryResponse = serde_json::from_str( + r#"{"groups":[{"displayName":"Gemini models","buckets":[]}]}"#, + ) + .unwrap(); + let usage = antigravity::antigravity_usage_from_summary(response).expect("authoritative"); + assert!(!usage.session.available && !usage.weekly.available); + + let empty: antigravity::AntigravityQuotaSummaryResponse = + serde_json::from_str(r#"{"groups":[]}"#).unwrap(); + assert!(antigravity::antigravity_usage_from_summary(empty).is_some()); +}