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()); +}