1use std::time::Duration;
48
49use chrono::{DateTime, Utc};
50
51use super::cloud;
52use super::credential::{self, StoredToken};
53use crate::cache::{Cache, acquire_lock_async};
54use crate::error::{AppError, Result};
55use crate::usage::{AntigravitySnapshot, AntigravitySource, ModelPoolWindow, UsageWindow};
56
57const HTTP_TIMEOUT: Duration = Duration::from_secs(5);
58const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
59
60const QUOTA_RPC: &str = "exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary";
61const STATUS_RPC: &str = "exa.language_server_pb.LanguageServerService/GetUserStatus";
62
63const DEFAULT_PLAN: &str = "Antigravity";
64
65const NO_LOCAL_SERVER: &str = "Antigravity: no local server found. Quota is only served while \
66 Antigravity is running — open the Antigravity app, or an interactive \
67 `agy` session, or point ANTIGRAVITY_LS_ADDRESS at a host:port.";
68
69const AGY_CSRF_UNAVAILABLE: &str = "Antigravity: the running `agy` server requires a CSRF token \
70 that it does not publish.";
71
72const NO_SAVED_SESSION: &str = "Or sign in to Antigravity once, so its saved Google session can \
76 be used while it is closed.";
77
78const SESSION_EXPIRED: &str =
79 "Antigravity's saved Google session expired; open Antigravity to sign in again";
80
81pub type FetchOutcome = crate::outcome::Outcome<AntigravitySnapshot>;
84
85impl From<FetchOutcome> for crate::vendor::VendorOutcome {
86 fn from(o: FetchOutcome) -> Self {
87 o.map(crate::usage::VendorSnapshot::Antigravity)
88 }
89}
90
91#[derive(Debug, Clone, Copy, Default)]
95pub enum SavedCredential<'a> {
96 #[default]
98 Keyring,
99 Blob(&'a str),
101 File(&'a std::path::Path),
104 Absent,
106}
107
108pub struct RemoteOverride<'a> {
113 pub credential: SavedCredential<'a>,
115 pub endpoints: Option<&'a cloud::Endpoints>,
117 pub local_bases: Option<Vec<Candidate>>,
122 pub agy: super::agy::AgyCommand,
126 pub ide_language_server: bool,
130}
131
132impl Default for RemoteOverride<'_> {
133 fn default() -> Self {
134 Self {
135 credential: SavedCredential::default(),
136 endpoints: None,
137 local_bases: None,
138 agy: super::agy::AgyCommand::default(),
139 ide_language_server: true,
140 }
141 }
142}
143
144#[derive(Debug, Clone, PartialEq, Eq)]
149pub struct Candidate {
150 pub base: String,
152 pub pid: Option<u32>,
155 pub ide: bool,
162}
163
164impl From<String> for Candidate {
165 fn from(base: String) -> Self {
166 Self {
167 base,
168 pid: None,
169 ide: false,
170 }
171 }
172}
173
174pub async fn fetch_snapshot(
177 client: &reqwest::Client,
178 cache: &Cache,
179 cache_ttl: Duration,
180 oauth: Option<&cloud::OauthClient>,
181 agy_binary: Option<&std::path::Path>,
182 ide_language_server: bool,
183) -> Result<FetchOutcome> {
184 fetch_snapshot_at(
185 client,
186 cache,
187 cache_ttl,
188 oauth,
189 RemoteOverride {
190 agy: agy_binary.map_or(super::agy::AgyCommand::Discover, |path| {
191 super::agy::AgyCommand::At(path.to_path_buf())
192 }),
193 ide_language_server,
194 ..Default::default()
195 },
196 Utc::now(),
197 )
198 .await
199}
200
201pub async fn fetch_snapshot_at(
205 client: &reqwest::Client,
206 cache: &Cache,
207 cache_ttl: Duration,
208 oauth: Option<&cloud::OauthClient>,
209 remote: RemoteOverride<'_>,
210 now: DateTime<Utc>,
211) -> Result<FetchOutcome> {
212 cache.ensure_dir()?;
213 let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
214
215 let origin = match open_session(
225 client,
226 remote.local_bases.as_deref(),
227 remote.ide_language_server,
228 )
229 .await
230 {
231 Ok(session) => Origin::Local(Ok(session)),
232 Err(error) => match remote_fallback_reason(&error) {
233 Some(reason) => Origin::Remote(saved_session(remote.credential, reason)),
234 None => Origin::Local(Err(error)),
235 },
236 };
237 let account = origin.account();
238
239 let cache_ttl = match origin {
242 Origin::Remote(Ok(_)) if remote.agy.is_available() => {
243 cache_ttl.max(super::agy::QUOTA_INTERVAL)
244 }
245 _ => cache_ttl,
246 };
247 let auth_rejected = matches!(&origin, Origin::Local(Err(e)) if is_auth_rejection(e));
248 if !auth_rejected
249 && let Some(bytes) = cache.fresh_payload(cache_ttl)?
250 && let Ok(outcome) = reuse_cache(bytes, cache, false, account.as_deref(), now)
251 {
252 return Ok(outcome);
253 }
254
255 let default_endpoints = cloud::Endpoints::default();
256 let endpoints = remote.endpoints.unwrap_or(&default_endpoints);
257 let live = match origin {
258 Origin::Local(session) => fetch_live(client, session).await,
259 Origin::Remote(token) => {
260 let renewal = Renewal {
261 agy: &remote.agy,
262 credential: remote.credential,
263 };
264 match fetch_remote(client, cache, oauth, endpoints, token, renewal, now).await {
265 Ok(RemoteReading::Live(snap)) => Ok(snap),
266 Ok(RemoteReading::Kept(outcome)) => return Ok(outcome),
267 Ok(RemoteReading::Backoff(remaining)) => {
268 return Err(AppError::Http {
269 status: 429,
270 body: format!(
271 "rate limited; next attempt in {}",
272 crate::cache::human_backoff(remaining)
273 ),
274 });
275 }
276 Err(e) => Err(e),
277 }
278 }
279 };
280
281 match live {
282 Ok(snap) => {
283 let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
284 cache.write_payload(&bytes)?;
285 Ok(crate::outcome::Outcome::fresh(snap))
286 }
287 Err(e) if e.is_transient() => fallback_silent(cache, now, e),
288 Err(AppError::Http { status, body }) if matches!(status, 401 | 403) => {
289 cache.mark_stale();
290 let _ = cache.write_last_error(status, &body);
291 Err(AppError::Http { status, body })
292 }
293 Err(AppError::Http { status, body }) => {
294 cache.mark_stale();
295 let last_error = Some(cache.write_last_error(status, &body));
296 let reason = AppError::Http { status, body };
297 fallback_with_error(cache, last_error, reason, now)
298 }
299 Err(e) => {
300 cache.mark_stale();
301 let last_error = Some(cache.write_last_error(0, &e.to_string()));
302 fallback_with_error(cache, last_error, e, now)
303 }
304 }
305}
306
307struct Session {
310 base: String,
311 csrf: Option<String>,
312 plan: String,
313 account: String,
314}
315
316enum Origin {
319 Local(Result<Session>),
320 Remote(Result<StoredToken>),
321}
322
323impl Origin {
324 fn account(&self) -> Option<String> {
326 match self {
327 Origin::Local(Ok(session)) => Some(session.account.clone()),
328 Origin::Remote(Ok(token)) => Some(remote_account(&token.fingerprint)),
329 Origin::Local(Err(_)) | Origin::Remote(Err(_)) => None,
330 }
331 }
332}
333
334fn no_local_server() -> AppError {
335 AppError::Credentials(NO_LOCAL_SERVER.into())
336}
337
338#[derive(Clone, Copy)]
339enum RemoteFallbackReason {
340 NoLocalServer,
341 AgyMissingCsrf,
342}
343
344impl RemoteFallbackReason {
345 fn message(self) -> &'static str {
346 match self {
347 Self::NoLocalServer => NO_LOCAL_SERVER,
348 Self::AgyMissingCsrf => AGY_CSRF_UNAVAILABLE,
349 }
350 }
351}
352
353fn remote_fallback_reason(error: &AppError) -> Option<RemoteFallbackReason> {
359 if matches!(error, AppError::Credentials(message) if message == NO_LOCAL_SERVER) {
360 Some(RemoteFallbackReason::NoLocalServer)
361 } else if is_missing_csrf(error) {
362 Some(RemoteFallbackReason::AgyMissingCsrf)
363 } else {
364 None
365 }
366}
367
368async fn open_session(
379 client: &reqwest::Client,
380 bases: Option<&[Candidate]>,
381 ide_language_server: bool,
382) -> Result<Session> {
383 let candidates = bases.map_or_else(
384 || candidate_bases(ide_language_server),
385 <[Candidate]>::to_vec,
386 );
387 let candidates: Vec<Candidate> = candidates.into_iter().filter(|c| !c.ide).collect();
388 if candidates.is_empty() {
389 return Err(no_local_server());
390 }
391
392 let mut errors = Vec::new();
393 let mut blocked_pids = std::collections::HashSet::new();
398 for candidate in candidates {
399 if candidate.pid.is_some_and(|pid| blocked_pids.contains(&pid)) {
400 continue;
401 }
402 let csrf = fetch_csrf(client, &candidate.base).await;
403 match post_rpc(client, &candidate.base, csrf.as_deref(), STATUS_RPC).await {
404 Ok(v) => {
405 return Ok(Session {
406 base: candidate.base,
407 csrf,
408 plan: plan_from_status(&v),
409 account: account_key(&v),
410 });
411 }
412 Err(e) => {
413 if is_missing_csrf(&e)
417 && let Some(pid) = candidate.pid
418 {
419 blocked_pids.insert(pid);
420 }
421 errors.push(e);
422 }
423 }
424 }
425 Err(select_probe_error(errors))
426}
427
428fn select_probe_error(errors: Vec<AppError>) -> AppError {
445 let mut actionable = None;
446 let mut missing_csrf = None;
447 let mut last = None;
448 let mut echo = None;
449 for e in errors {
450 if missing_csrf.is_none() && is_missing_csrf(&e) {
451 missing_csrf = Some(e);
452 } else if actionable.is_none() && is_actionable(&e) {
453 actionable = Some(e);
454 } else if is_tls_echo(&e) {
455 echo = Some(e);
456 } else {
457 last = Some(e);
458 }
459 }
460 actionable
461 .or(missing_csrf)
462 .or(last)
463 .or(echo)
464 .unwrap_or_else(|| {
465 AppError::Other("antigravity: no local server answered GetUserStatus".into())
466 })
467}
468
469fn is_missing_csrf(error: &AppError) -> bool {
475 let AppError::Http { status: 401, body } = error else {
476 return false;
477 };
478 let Ok(body) = serde_json::from_str::<serde_json::Value>(body) else {
479 return false;
480 };
481 matches!(
482 (body["code"].as_str(), body["message"].as_str()),
483 (Some(code), Some(message))
484 if code.eq_ignore_ascii_case("unauthenticated")
485 && is_missing_csrf_message(message)
486 )
487}
488
489fn is_missing_csrf_message(message: &str) -> bool {
492 message
493 .rsplit(':')
494 .next()
495 .is_some_and(|part| part.trim().eq_ignore_ascii_case("missing CSRF token"))
496}
497
498fn is_actionable(e: &AppError) -> bool {
503 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
504 && !is_missing_csrf(e)
505}
506
507fn is_tls_echo(e: &AppError) -> bool {
520 matches!(
521 e,
522 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
523 )
524}
525
526async fn fetch_live(
527 client: &reqwest::Client,
528 session: Result<Session>,
529) -> Result<AntigravitySnapshot> {
530 let session = session?;
531 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
532 let mut snap = parse_quota_summary("a, session.plan)?;
533 snap.account = session.account;
534 Ok(snap)
535}
536
537fn saved_session(
547 credential: SavedCredential<'_>,
548 reason: RemoteFallbackReason,
549) -> Result<StoredToken> {
550 let raw = match credential {
551 SavedCredential::Keyring => credential::read()?,
552 SavedCredential::Blob(blob) => Some(blob.to_string()),
553 SavedCredential::File(path) => std::fs::read_to_string(path).ok(),
554 SavedCredential::Absent => None,
555 };
556 let Some(raw) = raw else {
557 return Err(AppError::Credentials(format!(
558 "{} {NO_SAVED_SESSION}",
559 reason.message()
560 )));
561 };
562 credential::parse_keyring_blob(&raw)
563}
564
565fn remote_account(fingerprint: &str) -> String {
569 format!("acct:{fingerprint}")
570}
571
572fn session_expired() -> AppError {
573 AppError::Credentials(SESSION_EXPIRED.into())
574}
575
576fn refresh_unconfigured() -> AppError {
580 AppError::Credentials(
581 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
582 refresh it; open Antigravity (or run `agy`) to sign in again, or set [antigravity] oauth_client_id \
583 and oauth_client_secret in config.toml"
584 .into(),
585 )
586}
587
588fn is_auth_rejection(e: &AppError) -> bool {
589 matches!(
590 e,
591 AppError::Http {
592 status: 401 | 403,
593 ..
594 }
595 )
596}
597
598struct AccessToken {
602 value: String,
603 just_refreshed: bool,
604}
605
606fn current_access_token(
611 oauth_path: &std::path::Path,
612 token: &StoredToken,
613 now: DateTime<Utc>,
614) -> Option<String> {
615 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
616 && token
617 .expires_at
618 .is_none_or(|keyring| persisted.expires_at > keyring)
619 && !cloud::needs_refresh(Some(persisted.expires_at), now)
620 {
621 return Some(persisted.access_token);
622 }
623 (!cloud::needs_refresh(token.expires_at, now)).then(|| token.access_token.clone())
624}
625
626async fn resolve_access_token(
628 client: &reqwest::Client,
629 oauth: Option<&cloud::OauthClient>,
630 endpoints: &cloud::Endpoints,
631 oauth_path: &std::path::Path,
632 token: &StoredToken,
633 now: DateTime<Utc>,
634) -> Result<AccessToken> {
635 if let Some(value) = current_access_token(oauth_path, token, now) {
636 return Ok(AccessToken {
637 value,
638 just_refreshed: false,
639 });
640 }
641 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
642}
643
644async fn refresh_and_persist(
648 client: &reqwest::Client,
649 oauth: Option<&cloud::OauthClient>,
650 endpoints: &cloud::Endpoints,
651 oauth_path: &std::path::Path,
652 token: &StoredToken,
653) -> Result<AccessToken> {
654 let Some(refresh_token) = token.refresh_token.as_deref() else {
655 return Err(session_expired());
656 };
657 let Some(oauth) = oauth else {
658 return Err(refresh_unconfigured());
659 };
660 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
661 cloud::write_persisted(
662 oauth_path,
663 &cloud::PersistedOAuth {
664 fingerprint: token.fingerprint.clone(),
665 access_token: refreshed.access_token.clone(),
666 expires_at: refreshed.expires_at,
667 },
668 )?;
669 Ok(AccessToken {
670 value: refreshed.access_token,
671 just_refreshed: true,
672 })
673}
674
675struct Renewal<'a> {
678 agy: &'a super::agy::AgyCommand,
679 credential: SavedCredential<'a>,
680}
681
682impl Renewal<'_> {
683 async fn renewed(&self, cache: &Cache, now: DateTime<Utc>) -> Option<StoredToken> {
687 if !super::agy::renew_session(self.agy, cache, now).await {
688 return None;
689 }
690 let raw = match self.credential {
691 SavedCredential::Keyring => credential::read().ok().flatten()?,
692 SavedCredential::Blob(blob) => blob.to_string(),
693 SavedCredential::File(path) => std::fs::read_to_string(path).ok()?,
694 SavedCredential::Absent => return None,
695 };
696 credential::parse_keyring_blob(&raw).ok()
697 }
698}
699
700enum RemoteReading {
704 Live(AntigravitySnapshot),
705 Kept(FetchOutcome),
706 Backoff(Duration),
709}
710
711async fn fetch_remote(
737 client: &reqwest::Client,
738 cache: &Cache,
739 oauth: Option<&cloud::OauthClient>,
740 endpoints: &cloud::Endpoints,
741 token: Result<StoredToken>,
742 renewal: Renewal<'_>,
743 now: DateTime<Utc>,
744) -> Result<RemoteReading> {
745 let mut token = token?;
746 let oauth_path = cloud::oauth_cache_path(cache);
747 match super::agy::quota_snapshot(renewal.agy, cache, now).await {
748 Some(mut snap) => {
749 let access = current_access_token(&oauth_path, &token, now)
750 .filter(|_| cache.backoff_remaining().is_none());
751 snap.plan = match access {
752 Some(access) => cloud::fetch_plan(client, endpoints, &access).await,
753 None => None,
754 }
755 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
756 snap.account = remote_account(&token.fingerprint);
757 snap.source = AntigravitySource::Agy;
758 return Ok(RemoteReading::Live(snap));
759 }
760 None if renewal.agy.is_available() => {
761 if let Some(kept) = kept_agy_reading(cache, &remote_account(&token.fingerprint), now) {
762 return Ok(RemoteReading::Kept(kept));
763 }
764 }
765 None => {}
766 }
767 if let Some(remaining) = cache.backoff_remaining() {
771 return Ok(RemoteReading::Backoff(remaining));
772 }
773 let mut access =
774 match resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await {
775 Ok(access) => access,
776 Err(e) if oauth.is_none() => {
779 let Some(renewed) = renewal.renewed(cache, now).await else {
780 return Err(e);
781 };
782 token = renewed;
783 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
784 }
785 Err(e) => return Err(e),
786 };
787
788 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
789 Err(e) if cloud::is_no_plan_quota(&e) => return Err(e),
792 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
793 access = if oauth.is_none() {
794 let Some(renewed) = renewal.renewed(cache, now).await else {
795 return Err(refresh_unconfigured());
796 };
797 token = renewed;
798 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
799 } else {
800 refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?
801 };
802 cloud::fetch_quota(client, endpoints, &access.value)
803 .await
804 .map_err(|e| {
805 if cloud::is_no_plan_quota(&e) {
806 e
807 } else if is_auth_rejection(&e) {
808 session_expired()
809 } else {
810 e
811 }
812 })?
813 }
814 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
815 other => other?,
816 };
817
818 let plan = cloud::fetch_plan(client, endpoints, &access.value)
819 .await
820 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
821 let mut snap = parse_quota_summary("a, plan)?;
822 snap.account = remote_account(&token.fingerprint);
823 snap.source = AntigravitySource::Remote;
824 Ok(RemoteReading::Live(snap))
825}
826
827fn kept_agy_reading(cache: &Cache, account: &str, now: DateTime<Utc>) -> Option<FetchOutcome> {
836 let not_agys = || AppError::Other("Antigravity: no reading of `agy`'s own to keep".into());
837 crate::outcome::fallback(cache, None, not_agys(), |bytes| {
838 let snap = parse_cache_at(bytes, Some(account), now)?;
839 (snap.source == AntigravitySource::Agy)
840 .then_some(snap)
841 .ok_or_else(not_agys)
842 })
843 .ok()
844}
845
846fn account_key(user_status: &serde_json::Value) -> String {
851 let email = user_status["userStatus"]["email"]
852 .as_str()
853 .filter(|s| !s.is_empty());
854 match email {
855 Some(e) => {
856 use std::hash::{Hash, Hasher};
857 let mut h = std::collections::hash_map::DefaultHasher::new();
858 e.hash(&mut h);
859 format!("acct:{:016x}", h.finish())
860 }
861 None => "acct:unknown".to_string(),
862 }
863}
864
865async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
869 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
870 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
873 .await
874 .ok()?;
875 let html = String::from_utf8_lossy(&bytes);
876 html.split("csrfToken\":\"")
877 .nth(1)
878 .and_then(|s| s.split('"').next())
879 .filter(|t| !t.is_empty())
880 .map(|t| t.to_string())
881}
882
883async fn post_rpc(
884 client: &reqwest::Client,
885 base: &str,
886 csrf: Option<&str>,
887 rpc: &str,
888) -> Result<serde_json::Value> {
889 let mut req = client
890 .post(format!("{base}/{rpc}"))
891 .header("Content-Type", "application/json")
892 .body("{}")
893 .timeout(HTTP_TIMEOUT);
894 if let Some(token) = csrf {
895 req = req.header("x-codeium-csrf-token", token);
896 }
897 let resp = req.send().await?;
898
899 let status = resp.status();
900 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
904 if !status.is_success() {
905 let body = String::from_utf8_lossy(&bytes).into_owned();
906 return Err(AppError::Http {
907 status: status.as_u16(),
908 body,
909 });
910 }
911
912 Ok(serde_json::from_slice(&bytes)?)
913}
914
915pub fn plan_from_status(v: &serde_json::Value) -> String {
916 v["userStatus"]["userTier"]["name"]
917 .as_str()
918 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
919 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
920 .filter(|s| !s.is_empty())
921 .unwrap_or(DEFAULT_PLAN)
922 .to_string()
923}
924
925fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
937 let id = if id.is_empty() { "<unnamed>" } else { id };
938 let group = if group_name.is_empty() {
939 String::new()
940 } else {
941 format!(" in {group_name:?}")
942 };
943 match window {
944 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
945 _ => format!("{id}{group}"),
946 }
947}
948
949fn no_usable_bucket(seen: &[String]) -> AppError {
955 if seen.is_empty() {
956 return AppError::Other(
957 "antigravity: quota summary has no buckets at all — the running product may \
958 not have a quota for this account yet"
959 .into(),
960 );
961 }
962 AppError::Other(format!(
963 "antigravity: quota summary has no bucket in a window we recognise (5h or \
964 weekly, Gemini or Claude/GPT); it offered: {}",
965 seen.join(", ")
966 ))
967}
968
969pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
970 let groups = v["response"]["groups"]
971 .as_array()
972 .or_else(|| v["groups"].as_array())
973 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
974
975 const SLOT_NAMES: [&str; 4] = [
977 "Gemini 5h",
978 "Gemini weekly",
979 "Claude/GPT 5h",
980 "Claude/GPT weekly",
981 ];
982 const CANONICAL_IDS: [&str; 4] = ["gemini-5h", "gemini-weekly", "3p-5h", "3p-weekly"];
983 let is_canonical = |bucket: &serde_json::Value| {
984 CANONICAL_IDS.contains(&bucket["bucketId"].as_str().unwrap_or_default())
985 };
986 let shares_a_slot = |group: &serde_json::Value| {
987 group["buckets"]
988 .as_array()
989 .is_some_and(|buckets| buckets.iter().any(is_canonical))
990 };
991 let renamed = !groups.iter().any(shares_a_slot);
998
999 let mut exact: [Vec<&serde_json::Value>; 4] = Default::default();
1002 let mut inferred: [Vec<&serde_json::Value>; 4] = Default::default();
1003 let mut pools: Vec<(&str, &serde_json::Value, bool)> = Vec::new();
1004 let mut seen: Vec<String> = Vec::new();
1008
1009 for group in groups {
1010 let group_name = group["displayName"].as_str().unwrap_or_default();
1011 let Some(buckets) = group["buckets"].as_array() else {
1012 continue;
1013 };
1014 let own_pool = !renamed && !shares_a_slot(group);
1015 for bucket in buckets {
1016 let id = bucket["bucketId"].as_str().unwrap_or_default();
1017 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
1018 if let Some(slot) = CANONICAL_IDS.iter().position(|c| *c == id) {
1019 exact[slot].push(bucket);
1020 continue;
1021 }
1022 let window = bucket["window"].as_str().unwrap_or_default();
1023 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
1024 true
1025 } else if id.ends_with("5h") || window == "5h" {
1026 false
1027 } else {
1028 continue;
1031 };
1032 if own_pool {
1033 pools.push((group_name, bucket, is_weekly));
1034 continue;
1035 }
1036 if !renamed {
1039 continue;
1040 }
1041 let is_gemini = if id.starts_with("gemini") {
1042 true
1043 } else if id.starts_with("3p") {
1044 false
1045 } else if group_name.contains("Gemini") {
1046 true
1047 } else if group_name.contains("Claude") || group_name.contains("GPT") {
1048 false
1049 } else {
1050 pools.push((group_name, bucket, is_weekly));
1053 continue;
1054 };
1055
1056 let slot = usize::from(!is_gemini) * 2 + usize::from(is_weekly);
1057 inferred[slot].push(bucket);
1058 }
1059 }
1060
1061 let mut slots: [Option<UsageWindow>; 4] = Default::default();
1065 for slot in 0..4 {
1066 let candidates = if renamed {
1067 &inferred[slot]
1068 } else {
1069 &exact[slot]
1070 };
1071 match candidates.as_slice() {
1072 [] => {}
1073 [bucket] => slots[slot] = Some(usage_window(bucket, slot % 2 == 1)?),
1074 _ => {
1075 return Err(AppError::Schema(format!(
1076 "antigravity: duplicate {} bucket",
1077 SLOT_NAMES[slot]
1078 )));
1079 }
1080 }
1081 }
1082 let [gemini_5h, gemini_weekly, tp_5h, tp_weekly] = slots;
1083
1084 let model_pools: Vec<ModelPoolWindow> = pools
1088 .into_iter()
1089 .filter_map(|(group_name, bucket, is_weekly)| {
1090 let window = usage_window(bucket, is_weekly).ok()?;
1091 Some(ModelPoolWindow {
1092 name: pool_name(group_name, bucket),
1093 window,
1094 })
1095 })
1096 .collect();
1097
1098 if gemini_5h.is_none()
1102 && gemini_weekly.is_none()
1103 && tp_5h.is_none()
1104 && tp_weekly.is_none()
1105 && model_pools.is_empty()
1106 {
1107 return Err(no_usable_bucket(&seen));
1108 }
1109
1110 Ok(AntigravitySnapshot {
1111 plan,
1112 account: String::new(),
1115 source: AntigravitySource::Local,
1116 session: gemini_5h,
1117 weekly: gemini_weekly,
1118 third_party_session: tp_5h,
1119 third_party_weekly: tp_weekly,
1120 model_pools,
1121 })
1122}
1123
1124fn pool_name(group_name: &str, bucket: &serde_json::Value) -> String {
1128 [
1129 group_name,
1130 bucket["displayName"].as_str().unwrap_or_default(),
1131 bucket["bucketId"].as_str().unwrap_or_default(),
1132 ]
1133 .into_iter()
1134 .map(crate::display::sanitize_untrusted_field)
1135 .find(|name| !name.trim().is_empty())
1136 .unwrap_or_else(|| "Model pool".to_string())
1137}
1138
1139fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
1143 let remaining = bucket["remainingFraction"]
1144 .as_f64()
1145 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
1146 .ok_or_else(|| {
1147 AppError::Schema(format!(
1148 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
1149 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
1150 ))
1151 })?;
1152 Ok(UsageWindow {
1153 utilization_pct: pct_used(remaining),
1154 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
1155 window_duration: if is_weekly {
1156 chrono::Duration::days(7)
1157 } else {
1158 chrono::Duration::hours(5)
1159 },
1160 })
1161}
1162
1163fn pct_used(remaining_fraction: f64) -> i32 {
1166 let used = (1.0 - remaining_fraction) * 100.0;
1167 used.round() as i32
1168}
1169
1170fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
1171 match value {
1172 serde_json::Value::Null => Ok(None),
1173 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
1174 .map(|dt| Some(dt.with_timezone(&Utc)))
1175 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
1176 _ => Err(AppError::Schema(format!(
1177 "antigravity: {field} must be a timestamp or null"
1178 ))),
1179 }
1180}
1181
1182fn candidate_bases(ide_language_server: bool) -> Vec<Candidate> {
1190 let discovered = discover_ls_ports();
1191 let ide_pids = if ide_language_server {
1194 std::collections::HashSet::new()
1195 } else {
1196 ide_ls_pids(discovered.iter().map(|&(pid, _)| pid).collect())
1197 };
1198 candidate_bases_with(
1199 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
1200 discovered,
1201 &ide_pids,
1202 )
1203}
1204
1205fn candidate_bases_with(
1209 override_addr: Option<&str>,
1210 discovered: Vec<(u32, u16)>,
1211 ide_pids: &std::collections::HashSet<u32>,
1212) -> Vec<Candidate> {
1213 let mut candidates = Vec::new();
1214 if let Some(base) = override_addr.and_then(normalize_base) {
1215 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
1220 candidates.push(Candidate {
1221 base,
1222 pid,
1223 ide: false,
1224 });
1225 }
1226
1227 for (pid, port) in discovered {
1233 let base = format!("http://127.0.0.1:{port}");
1234 if !candidates.iter().any(|c| c.base == base) {
1235 candidates.push(Candidate {
1236 base,
1237 pid: Some(pid),
1238 ide: ide_pids.contains(&pid),
1239 });
1240 }
1241 }
1242
1243 candidates
1244}
1245
1246#[cfg(any(test, target_os = "linux", target_os = "macos"))]
1251fn is_ide_install(exe: &str) -> bool {
1252 let exe = exe.to_lowercase().replace('\\', "/");
1253 exe.contains("antigravity-ide") || exe.contains("antigravity ide")
1254}
1255
1256#[cfg(target_os = "linux")]
1258fn ide_ls_pids(pids: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1259 pids.into_iter()
1260 .filter(|pid| {
1261 std::fs::read_link(format!("/proc/{pid}/exe"))
1262 .map(|path| path.to_string_lossy().into_owned())
1263 .is_ok_and(|exe| is_ide_install(&exe))
1264 })
1265 .collect()
1266}
1267
1268#[cfg(target_os = "macos")]
1271fn ide_ls_pids(pids: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1272 if pids.is_empty() {
1273 return std::collections::HashSet::new();
1274 }
1275 let list = pids
1276 .iter()
1277 .map(u32::to_string)
1278 .collect::<Vec<_>>()
1279 .join(",");
1280 let Ok(output) = std::process::Command::new("/bin/ps")
1283 .args(["-o", "pid=,comm=", "-p", &list])
1284 .output()
1285 else {
1286 return std::collections::HashSet::new();
1287 };
1288 parse_ps_comm(&String::from_utf8_lossy(&output.stdout))
1289 .into_iter()
1290 .filter(|(pid, exe)| pids.contains(pid) && is_ide_install(exe))
1291 .map(|(pid, _)| pid)
1292 .collect()
1293}
1294
1295#[cfg(any(test, target_os = "macos"))]
1298fn parse_ps_comm(output: &str) -> Vec<(u32, String)> {
1299 output
1300 .lines()
1301 .filter_map(|line| {
1302 let (pid, exe) = line.trim().split_once(char::is_whitespace)?;
1303 Some((pid.parse().ok()?, exe.trim().to_string()))
1304 })
1305 .collect()
1306}
1307
1308#[cfg(target_os = "windows")]
1312fn ide_ls_pids(_: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1313 std::collections::HashSet::new()
1314}
1315
1316#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1317fn ide_ls_pids(_: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1318 std::collections::HashSet::new()
1319}
1320
1321fn base_port(base: &str) -> Option<u16> {
1323 let (_, authority) = base.split_once("://")?;
1324 authority.rsplit_once(':')?.1.parse().ok()
1325}
1326
1327fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
1330 let mut owners = discovered
1331 .iter()
1332 .filter(|(_, p)| *p == port)
1333 .map(|(pid, _)| *pid);
1334 let first = owners.next()?;
1335 owners.all(|pid| pid == first).then_some(first)
1336}
1337
1338fn normalize_base(addr: &str) -> Option<String> {
1347 let trimmed = addr.trim();
1348 let (scheme, authority) = match trimmed.split_once("://") {
1349 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
1350 _ => ("http", trimmed),
1351 };
1352 let authority = authority.trim_end_matches('/');
1353 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
1354}
1355
1356fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
1363 let comm = comm.trim().to_lowercase();
1364 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
1365 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
1366 return true;
1367 }
1368 exe.is_some_and(|p| {
1369 let p = p.to_lowercase().replace('\\', "/");
1370 let p = p.strip_suffix(".exe").unwrap_or(&p);
1371 p.contains("antigravity") || p.ends_with("/agy")
1372 })
1373}
1374
1375#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
1407fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
1408 let groups: Vec<(u32, Vec<u16>)> = per_pid
1409 .into_iter()
1410 .map(|(pid, mut group)| {
1411 group.sort_unstable_by(|a, b| b.cmp(a));
1412 group.dedup();
1417 (pid, group)
1418 })
1419 .collect();
1420
1421 let mut ports: Vec<(u32, u16)> = Vec::new();
1422 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
1423 for (pid, port) in groups
1424 .iter()
1425 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
1426 {
1427 if !ports.iter().any(|(_, p)| *p == port) {
1430 ports.push((pid, port));
1431 }
1432 }
1433 }
1434 ports
1435}
1436
1437#[cfg(target_os = "linux")]
1445pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1446 use std::collections::{BTreeMap, HashMap};
1447
1448 let mut owners: HashMap<u64, u32> = HashMap::new();
1451 let Ok(entries) = std::fs::read_dir("/proc") else {
1452 return Vec::new();
1453 };
1454 for entry in entries.flatten() {
1455 let pid_dir = entry.path();
1456 let Some(pid) = pid_dir
1457 .file_name()
1458 .and_then(|name| name.to_str())
1459 .and_then(|name| name.parse::<u32>().ok())
1460 else {
1461 continue;
1462 };
1463 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1464 continue;
1465 };
1466 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1467 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1468 continue;
1469 }
1470 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1471 continue;
1472 };
1473 for fd in fds.flatten() {
1474 let Ok(target) = std::fs::read_link(fd.path()) else {
1475 continue;
1476 };
1477 if let Some(ino) = target
1478 .to_str()
1479 .and_then(|s| s.strip_prefix("socket:["))
1480 .and_then(|s| s.strip_suffix(']'))
1481 .and_then(|s| s.parse::<u64>().ok())
1482 {
1483 owners.insert(ino, pid);
1484 }
1485 }
1486 }
1487
1488 if owners.is_empty() {
1489 return Vec::new();
1490 }
1491
1492 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1493 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1494 let Ok(contents) = std::fs::read_to_string(table) else {
1495 continue;
1496 };
1497 for line in contents.lines().skip(1) {
1498 if let Some((port, ino)) = parse_proc_net_line(line)
1499 && let Some(&pid) = owners.get(&ino)
1500 {
1501 per_pid.entry(pid).or_default().push(port);
1502 }
1503 }
1504 }
1505 probe_order(per_pid)
1506}
1507
1508#[cfg(target_os = "macos")]
1515pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1516 let Ok(output) = std::process::Command::new("/usr/sbin/lsof")
1518 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1519 .output()
1520 else {
1521 return Vec::new();
1522 };
1523 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1527}
1528
1529#[cfg(any(test, target_os = "macos"))]
1534fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1535 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1536 let mut pid = None;
1539 let mut owner = None;
1540 for line in output.lines() {
1541 let Some(rest) = line.get(1..) else { continue };
1542 match line.as_bytes().first() {
1543 Some(b'p') => {
1544 pid = rest.parse::<u32>().ok();
1545 owner = None;
1546 }
1547 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1548 Some(b'n') => {
1549 if let Some(pid) = owner
1550 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1551 {
1552 per_pid.entry(pid).or_default().push(port);
1553 }
1554 }
1555 _ => {}
1556 }
1557 }
1558 probe_order(per_pid)
1559}
1560
1561#[cfg(any(test, target_os = "windows"))]
1562#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1563struct WindowsTcpRow {
1564 local_addr: [u8; 4],
1565 local_port: u32,
1566 pid: u32,
1567}
1568
1569#[cfg(any(test, target_os = "windows"))]
1570fn decode_windows_process_name(raw: &[u16]) -> String {
1571 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1572 String::from_utf16_lossy(&raw[..end])
1573}
1574
1575#[cfg(any(test, target_os = "windows"))]
1576fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1577 processes
1578 .iter()
1579 .filter(|(_, name)| is_antigravity_process(name, None))
1580 .map(|(pid, _)| *pid)
1581 .collect()
1582}
1583
1584#[cfg(any(test, target_os = "windows"))]
1587fn matching_windows_ports(
1588 pids: &std::collections::HashSet<u32>,
1589 rows: &[WindowsTcpRow],
1590) -> Vec<(u32, u16)> {
1591 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1592 for row in rows {
1593 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1594 continue;
1595 }
1596 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1597 if port != 0 {
1598 per_pid.entry(row.pid).or_default().push(port);
1599 }
1600 }
1601 probe_order(per_pid)
1602}
1603
1604#[cfg(any(test, target_os = "windows"))]
1605fn checked_windows_row_count(
1606 buffer_len: usize,
1607 rows_offset: usize,
1608 row_size: usize,
1609 declared: usize,
1610) -> Option<usize> {
1611 let rows_len = row_size.checked_mul(declared)?;
1612 let end = rows_offset.checked_add(rows_len)?;
1613 (row_size != 0 && end <= buffer_len).then_some(declared)
1614}
1615
1616#[cfg(target_os = "windows")]
1617struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1618
1619#[cfg(target_os = "windows")]
1620impl Drop for WindowsHandle {
1621 fn drop(&mut self) {
1622 unsafe {
1623 windows_sys::Win32::Foundation::CloseHandle(self.0);
1624 }
1625 }
1626}
1627
1628#[cfg(target_os = "windows")]
1629fn windows_processes() -> Vec<(u32, String)> {
1630 use std::mem::size_of;
1631 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1632 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1633 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1634 TH32CS_SNAPPROCESS,
1635 };
1636
1637 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1638 if handle == INVALID_HANDLE_VALUE {
1639 return Vec::new();
1640 }
1641 let snapshot = WindowsHandle(handle);
1642 let mut entry = PROCESSENTRY32W::default();
1643 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1644 return Vec::new();
1645 };
1646 entry.dwSize = entry_size;
1647 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1648 return Vec::new();
1649 }
1650
1651 let mut processes = Vec::new();
1652 loop {
1653 processes.push((
1654 entry.th32ProcessID,
1655 decode_windows_process_name(&entry.szExeFile),
1656 ));
1657 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1658 break;
1659 }
1660 }
1661 processes
1662}
1663
1664#[cfg(target_os = "windows")]
1665fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1666 use std::mem::{offset_of, size_of, size_of_val};
1667 use windows_sys::Win32::NetworkManagement::IpHelper::{
1668 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1669 };
1670
1671 let available = used_bytes.min(size_of_val(buffer));
1672 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1673 if available < size_of::<u32>() || available < rows_offset {
1674 return Vec::new();
1675 }
1676 let base = buffer.as_ptr().cast::<u8>();
1677 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1678 if checked_windows_row_count(
1679 available,
1680 rows_offset,
1681 size_of::<MIB_TCPROW_OWNER_PID>(),
1682 declared,
1683 )
1684 .is_none()
1685 {
1686 return Vec::new();
1687 }
1688
1689 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1690 (0..declared)
1691 .map(|index| unsafe { rows.add(index).read_unaligned() })
1692 .map(|row| WindowsTcpRow {
1693 local_addr: row.dwLocalAddr.to_ne_bytes(),
1694 local_port: row.dwLocalPort,
1695 pid: row.dwOwningPid,
1696 })
1697 .collect()
1698}
1699
1700#[cfg(target_os = "windows")]
1701fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1702 use std::mem::size_of;
1703 use std::ptr::null_mut;
1704 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1705 use windows_sys::Win32::NetworkManagement::IpHelper::{
1706 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1707 };
1708 use windows_sys::Win32::Networking::WinSock::AF_INET;
1709
1710 let mut size = 0u32;
1711 let status = unsafe {
1712 GetExtendedTcpTable(
1713 null_mut(),
1714 &mut size,
1715 0,
1716 u32::from(AF_INET),
1717 TCP_TABLE_OWNER_PID_LISTENER,
1718 0,
1719 )
1720 };
1721 if status != ERROR_INSUFFICIENT_BUFFER {
1722 return Vec::new();
1723 }
1724
1725 for _ in 0..3 {
1726 let Some(words) = (size as usize)
1727 .checked_add(size_of::<u32>() - 1)
1728 .map(|bytes| bytes / size_of::<u32>())
1729 else {
1730 return Vec::new();
1731 };
1732 if words == 0 {
1733 return Vec::new();
1734 }
1735 let mut buffer = Vec::<u32>::new();
1736 if buffer.try_reserve_exact(words).is_err() {
1737 return Vec::new();
1738 }
1739 buffer.resize(words, 0);
1740 let mut used = size;
1741 let status = unsafe {
1742 GetExtendedTcpTable(
1743 buffer.as_mut_ptr().cast(),
1744 &mut used,
1745 0,
1746 u32::from(AF_INET),
1747 TCP_TABLE_OWNER_PID_LISTENER,
1748 0,
1749 )
1750 };
1751 if status == ERROR_INSUFFICIENT_BUFFER {
1752 size = used;
1753 continue;
1754 }
1755 if status != 0 {
1756 return Vec::new();
1757 }
1758 return parse_windows_tcp_rows(&buffer, used as usize);
1759 }
1760 Vec::new()
1761}
1762
1763#[cfg(target_os = "windows")]
1764pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1765 let pids = matching_windows_process_ids(&windows_processes());
1766 if pids.is_empty() {
1767 return Vec::new();
1768 }
1769 matching_windows_ports(&pids, &windows_tcp_rows())
1770}
1771
1772#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1773pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1774 Vec::new()
1775}
1776
1777#[cfg(target_os = "linux")]
1780fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1781 let cols: Vec<&str> = line.split_whitespace().collect();
1782 if cols.len() < 10 {
1783 return None;
1784 }
1785 if cols[3] != "0A" {
1787 return None;
1788 }
1789 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1790 let inode = cols[9].parse::<u64>().ok()?;
1791 Some((port, inode))
1792}
1793
1794fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1799 crate::outcome::fallback(cache, None, original, |bytes| {
1800 parse_cache_at(bytes, None, now)
1801 })
1802}
1803
1804fn fallback_with_error(
1809 cache: &Cache,
1810 last_error: Option<(u16, String)>,
1811 reason: AppError,
1812 now: DateTime<Utc>,
1813) -> Result<FetchOutcome> {
1814 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1815 parse_cache_at(bytes, None, now)
1816 })
1817}
1818
1819fn reuse_cache(
1820 bytes: Vec<u8>,
1821 cache: &Cache,
1822 stale: bool,
1823 account: Option<&str>,
1824 now: DateTime<Utc>,
1825) -> Result<FetchOutcome> {
1826 let snap = parse_cache_at(&bytes, account, now)?;
1827 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1828}
1829
1830pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1836 parse_cache_at(bytes, account, Utc::now())
1837}
1838
1839fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<String> {
1849 let shared = [
1850 ("Gemini 5h", snap.session.as_ref()),
1851 ("Gemini weekly", snap.weekly.as_ref()),
1852 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1853 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1854 ]
1855 .into_iter()
1856 .filter_map(|(name, w)| Some((name.to_string(), w?)));
1857 let pools = snap.model_pools.iter().map(|pool| {
1858 let name = format!("{} {}", pool.name, pool_cadence(pool.is_weekly()));
1859 (name, &pool.window)
1860 });
1861 shared
1862 .chain(pools)
1863 .find(|(_, w)| w.resets_at.is_some_and(|r| r <= now))
1864 .map(|(name, _)| name)
1865}
1866
1867fn pool_cadence(weekly: bool) -> &'static str {
1870 if weekly { "weekly" } else { "5h" }
1871}
1872
1873pub fn parse_cache_at(
1874 bytes: &[u8],
1875 account: Option<&str>,
1876 now: DateTime<Utc>,
1877) -> Result<AntigravitySnapshot> {
1878 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1879
1880 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1881 if let Some(expected) = account
1882 && cached_account != Some(expected)
1883 {
1884 return Err(AppError::Schema(
1885 "antigravity cache belongs to a different account; refetching".into(),
1886 ));
1887 }
1888
1889 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1899 match v.get(pct_key) {
1900 None => Err(AppError::Schema(format!(
1901 "antigravity: cached payload is missing {pct_key}"
1902 ))),
1903 Some(serde_json::Value::Null) => Ok(None),
1904 Some(value) => value
1905 .as_i64()
1906 .filter(|pct| (0..=100).contains(pct))
1907 .map(|pct| Some(pct as i32))
1908 .ok_or_else(|| {
1909 AppError::Schema(format!(
1910 "antigravity: cached {pct_key} must be an integer in 0..=100"
1911 ))
1912 }),
1913 }
1914 };
1915
1916 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1917 let Some(pct) = cached_pct(pct_key)? else {
1918 return Ok(None);
1919 };
1920 Ok::<_, AppError>(Some(UsageWindow {
1921 utilization_pct: pct,
1922 resets_at: parse_reset(&v[reset_key], reset_key)?,
1923 window_duration: if weekly {
1924 chrono::Duration::days(7)
1925 } else {
1926 chrono::Duration::hours(5)
1927 },
1928 }))
1929 };
1930
1931 let snap = AntigravitySnapshot {
1932 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1933 account: cached_account.unwrap_or_default().to_string(),
1934 source: v["source"]
1936 .as_str()
1937 .and_then(AntigravitySource::parse)
1938 .unwrap_or_default(),
1939 session: optional("session_pct", "session_reset", false)?,
1940 weekly: optional("weekly_pct", "weekly_reset", true)?,
1941 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1942 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1943 model_pools: cached_pools(&v)?,
1944 };
1945
1946 if snap.session.is_none()
1949 && snap.weekly.is_none()
1950 && snap.third_party_session.is_none()
1951 && snap.third_party_weekly.is_none()
1952 && snap.model_pools.is_empty()
1953 {
1954 return Err(AppError::Schema(
1955 "antigravity cache holds no usable window; refetching".into(),
1956 ));
1957 }
1958
1959 if let Some(window) = expired_window(&snap, now) {
1960 return Err(AppError::Schema(format!(
1961 "antigravity cache is past its {window} reset; refetching"
1962 )));
1963 }
1964 Ok(snap)
1965}
1966
1967fn cached_pools(v: &serde_json::Value) -> Result<Vec<ModelPoolWindow>> {
1971 let Some(entries) = v.get("model_pools") else {
1972 return Ok(Vec::new());
1973 };
1974 let malformed = || AppError::Schema("antigravity: cached model_pools is malformed".into());
1975 entries
1976 .as_array()
1977 .ok_or_else(malformed)?
1978 .iter()
1979 .map(|entry| {
1980 let name = entry["name"].as_str().ok_or_else(malformed)?;
1981 let weekly = match entry["window"].as_str() {
1982 Some(cadence) if cadence == pool_cadence(true) => true,
1983 Some(cadence) if cadence == pool_cadence(false) => false,
1984 _ => return Err(malformed()),
1985 };
1986 let pct = entry["pct"]
1987 .as_i64()
1988 .filter(|pct| (0..=100).contains(pct))
1989 .ok_or_else(malformed)?;
1990 Ok(ModelPoolWindow {
1991 name: crate::display::sanitize_untrusted_field(name),
1995 window: UsageWindow {
1996 utilization_pct: pct as i32,
1997 resets_at: parse_reset(&entry["reset"], "model_pools reset")?,
1998 window_duration: if weekly {
1999 chrono::Duration::days(7)
2000 } else {
2001 chrono::Duration::hours(5)
2002 },
2003 },
2004 })
2005 })
2006 .collect()
2007}
2008
2009pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
2010 let model_pools: Vec<serde_json::Value> = snap
2011 .model_pools
2012 .iter()
2013 .map(|pool| {
2014 serde_json::json!({
2015 "name": pool.name,
2016 "window": pool_cadence(pool.is_weekly()),
2017 "pct": pool.window.utilization_pct,
2018 "reset": pool.window.resets_at.map(|dt| dt.to_rfc3339()),
2019 })
2020 })
2021 .collect();
2022 serde_json::json!({
2023 "plan": snap.plan,
2024 "account": snap.account,
2025 "source": snap.source.as_str(),
2026 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
2027 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2028 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
2029 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2030 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
2031 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2032 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
2033 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2034 "model_pools": model_pools,
2035 })
2036}
2037
2038#[cfg(test)]
2039mod tests {
2040 use super::*;
2041
2042 const QUOTA_JSON: &str = r#"{
2046 "response": {
2047 "groups": [
2048 {
2049 "displayName": "Gemini Models",
2050 "buckets": [
2051 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
2052 "window": "weekly", "remainingFraction": 0.9191212,
2053 "resetTime": "2026-07-28T17:39:58Z"},
2054 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
2055 "window": "5h", "remainingFraction": 0.5672253,
2056 "resetTime": "2026-07-22T17:47:00Z"}
2057 ]
2058 },
2059 {
2060 "displayName": "Claude and GPT models",
2061 "buckets": [
2062 {"bucketId": "3p-weekly", "window": "weekly",
2063 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
2064 {"bucketId": "3p-5h", "window": "5h",
2065 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
2066 ]
2067 }
2068 ]
2069 }
2070 }"#;
2071
2072 fn now() -> DateTime<Utc> {
2075 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
2076 .unwrap()
2077 .with_timezone(&Utc)
2078 }
2079
2080 fn parsed() -> AntigravitySnapshot {
2081 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2082 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
2083 }
2084
2085 #[test]
2086 fn quota_summary_maps_four_distinct_windows() {
2087 let snap = parsed();
2088 assert_eq!(snap.plan, "Google AI Pro");
2089 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2091 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2092 assert_eq!(
2093 snap.third_party_session.as_ref().unwrap().utilization_pct,
2094 75
2095 );
2096 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
2097 }
2098
2099 #[test]
2100 fn each_window_keeps_its_own_reset_time() {
2101 let snap = parsed();
2102 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
2103 assert_eq!(
2104 snap.session.as_ref().unwrap().resets_at,
2105 at("2026-07-22T17:47:00Z")
2106 );
2107 assert_eq!(
2108 snap.weekly.as_ref().unwrap().resets_at,
2109 at("2026-07-28T17:39:58Z")
2110 );
2111 assert_eq!(
2112 snap.third_party_weekly.as_ref().unwrap().resets_at,
2113 at("2026-07-29T12:47:00Z")
2114 );
2115 assert_ne!(
2117 snap.session.as_ref().unwrap().resets_at,
2118 snap.weekly.as_ref().unwrap().resets_at
2119 );
2120 }
2121
2122 #[test]
2123 fn window_durations_match_their_bucket() {
2124 let snap = parsed();
2125 assert_eq!(
2126 snap.session.as_ref().unwrap().window_duration,
2127 chrono::Duration::hours(5)
2128 );
2129 assert_eq!(
2130 snap.weekly.as_ref().unwrap().window_duration,
2131 chrono::Duration::days(7)
2132 );
2133 assert_eq!(
2134 snap.third_party_weekly.as_ref().unwrap().window_duration,
2135 chrono::Duration::days(7)
2136 );
2137 }
2138
2139 #[test]
2140 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
2141 let v: serde_json::Value = serde_json::from_str(
2142 r#"{"response":{"groups":[
2143 {"displayName":"Gemini Models","buckets":[
2144 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
2145 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
2146 {"displayName":"Claude and GPT models","buckets":[
2147 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
2148 ]}}"#,
2149 )
2150 .unwrap();
2151 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2152 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
2153 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
2154 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
2155 assert!(snap.third_party_weekly.is_none());
2156 }
2157
2158 #[test]
2164 fn a_per_model_pool_is_reported_beside_the_canonical_buckets() {
2165 let v: serde_json::Value = serde_json::from_str(
2166 r#"{"response":{"groups":[
2167 {"displayName":"Gemini Models","buckets":[
2168 {"bucketId":"gemini-weekly","displayName":"Weekly Limit Remaining","window":"weekly",
2169 "remainingFraction":0.7839184,"resetTime":"2026-10-23T12:22:17Z"},
2170 {"bucketId":"gemini-5h","displayName":"Five Hour Limit Remaining","window":"5h",
2171 "remainingFraction":0.253501,"resetTime":"2026-10-16T17:22:17Z"}]},
2172 {"displayName":"Claude and GPT models","buckets":[
2173 {"bucketId":"3p-weekly","displayName":"Weekly Limit Remaining","window":"weekly",
2174 "remainingFraction":1,"resetTime":"2026-10-23T17:05:12Z"},
2175 {"bucketId":"3p-5h","displayName":"Five Hour Limit Remaining","window":"5h",
2176 "remainingFraction":1,"resetTime":"2026-10-16T22:05:12Z"}]},
2177 {"displayName":"Gemini 4 Argon","buckets":[
2178 {"bucketId":"gemini-argon-wtu-weekly","displayName":"Gemini 4 Argon Weekly Limit",
2179 "window":"weekly","remainingFraction":0.4,"resetTime":"2026-10-23T17:05:12Z"}]}
2180 ]}}"#,
2181 )
2182 .unwrap();
2183 let snap = parse_quota_summary(&v, "Google AI Pro".into()).unwrap();
2184 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
2185 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 22);
2186 assert_eq!(
2187 snap.third_party_session.as_ref().unwrap().utilization_pct,
2188 0
2189 );
2190 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
2191 let [argon] = snap.model_pools.as_slice() else {
2192 panic!("{:?}", snap.model_pools);
2193 };
2194 assert_eq!(argon.name, "Gemini 4 Argon");
2195 assert_eq!(argon.window.utilization_pct, 60);
2196 assert!(argon.is_weekly());
2197 }
2198
2199 #[test]
2202 fn a_malformed_model_pool_is_left_out() {
2203 let v: serde_json::Value = serde_json::from_str(
2204 r#"{"response":{"groups":[
2205 {"displayName":"Gemini Models","buckets":[
2206 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.5}]},
2207 {"displayName":"Gemini 4 Argon","buckets":[
2208 {"bucketId":"gemini-argon-wtu-weekly","window":"weekly","remainingFraction":"x"}]}
2209 ]}}"#,
2210 )
2211 .unwrap();
2212 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2213 assert_eq!(snap.weekly.unwrap().utilization_pct, 50);
2214 assert!(snap.model_pools.is_empty());
2215 }
2216
2217 #[test]
2219 fn a_model_pool_name_is_sanitized_and_falls_back_to_the_bucket() {
2220 let v: serde_json::Value = serde_json::from_str(
2221 "{\"response\":{\"groups\":[
2222 {\"displayName\":\"Gemini Models\",\"buckets\":[
2223 {\"bucketId\":\"gemini-weekly\",\"window\":\"weekly\",\"remainingFraction\":0.5}]},
2224 {\"displayName\":\"Argon\\u001b[31m\",\"buckets\":[
2225 {\"bucketId\":\"argon-weekly\",\"window\":\"weekly\",\"remainingFraction\":0.5}]},
2226 {\"buckets\":[
2227 {\"bucketId\":\"nova-5h\",\"displayName\":\"Nova Five Hour\",\"window\":\"5h\",\"remainingFraction\":0.5}]}
2228 ]}}",
2229 )
2230 .unwrap();
2231 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2232 let names: Vec<&str> = snap.model_pools.iter().map(|p| p.name.as_str()).collect();
2233 assert_eq!(names, ["Argon[31m", "Nova Five Hour"]);
2234 }
2235
2236 #[test]
2237 fn canonical_presence_disables_inference_in_a_shared_group() {
2238 for remaining in [serde_json::json!(0.1), serde_json::json!("unknown")] {
2239 let summary = serde_json::json!({"groups": [{
2240 "displayName": "Gemini Models",
2241 "buckets": [
2242 {"bucketId": "gemini-5h", "remainingFraction": 0.5},
2243 {"bucketId": "gemini-argon-wtu-weekly", "remainingFraction": remaining}
2244 ]
2245 }]});
2246 let snap = parse_quota_summary(&summary, "Pro".into()).unwrap();
2247 assert_eq!(snap.session.unwrap().utilization_pct, 50);
2248 assert!(snap.weekly.is_none());
2249 assert!(snap.model_pools.is_empty());
2250 }
2251 }
2252
2253 #[test]
2254 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
2255 let lookalike: serde_json::Value = serde_json::from_str(
2258 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
2259 {"bucketId":"gemini-5h-copy","window":"5h"},
2260 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2261 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
2262 ]}]}}"#,
2263 )
2264 .unwrap();
2265 let snap = parse_quota_summary(&lookalike, "Pro".into()).unwrap();
2266 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 10);
2267
2268 for buckets in [
2271 r#"{"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2272 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.1}"#,
2273 r#"{"bucketId":"x1","window":"5h","remainingFraction":0.9},
2274 {"bucketId":"x2","window":"5h","remainingFraction":0.1}"#,
2275 ] {
2276 let duplicate: serde_json::Value = serde_json::from_str(&format!(
2277 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[{buckets}]}}]}}}}"#
2278 ))
2279 .unwrap();
2280 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
2281 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
2282 }
2283
2284 let unrelated: serde_json::Value = serde_json::from_str(
2287 r#"{"response":{"groups":[
2288 {"displayName":"Gemini Models","buckets":[
2289 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2290 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
2291 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
2292 ]},
2293 {"displayName":"Future Models","buckets":[
2294 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
2295 ]}
2296 ]}}"#,
2297 )
2298 .unwrap();
2299 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
2300 assert!(snap.third_party_session.is_none());
2301 assert!(snap.third_party_weekly.is_none());
2302 let [future] = snap.model_pools.as_slice() else {
2303 panic!("{:?}", snap.model_pools);
2304 };
2305 assert_eq!(future.name, "Future Models");
2306 assert!(!future.is_weekly());
2307 }
2308
2309 #[test]
2312 fn a_bucket_without_a_usable_fraction_is_rejected() {
2313 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
2314 let v: serde_json::Value = serde_json::from_str(&format!(
2315 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
2316 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
2317 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
2318 ))
2319 .unwrap();
2320 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
2321 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
2322 }
2323 }
2324
2325 #[test]
2326 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
2327 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
2328 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2329 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
2330 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
2331 assert!(err.to_string().contains("resetTime"), "{err}");
2332 }
2333 }
2334
2335 #[test]
2339 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
2340 let mut snap = parsed();
2341 snap.session = None;
2342 snap.third_party_session = None;
2343
2344 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2345 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
2346
2347 assert!(back.session.is_none());
2348 assert_eq!(
2349 back.weekly.as_ref().unwrap().utilization_pct,
2350 snap.weekly.as_ref().unwrap().utilization_pct
2351 );
2352 }
2353
2354 #[test]
2359 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
2360 let summary = serde_json::json!({
2361 "groups": [
2362 {
2363 "displayName": "Gemini Models",
2364 "buckets": [{
2365 "bucketId": "gemini-weekly", "window": "weekly",
2366 "remainingFraction": 0.42,
2367 }],
2368 },
2369 {
2370 "displayName": "Claude and GPT models",
2371 "buckets": [{
2372 "bucketId": "3p-weekly", "window": "weekly",
2373 "remainingFraction": 0.9,
2374 }],
2375 },
2376 ],
2377 });
2378
2379 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
2380
2381 assert!(snap.session.is_none(), "no 5h bucket arrived");
2382 assert!(snap.third_party_session.is_none());
2383 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
2384 assert_eq!(
2385 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
2386 10
2387 );
2388 }
2389
2390 #[test]
2393 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
2394 let summary = serde_json::json!({
2395 "groups": [{
2396 "displayName": "Gemini Models",
2397 "buckets": [{
2398 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
2399 }],
2400 }],
2401 });
2402
2403 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
2404
2405 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
2406 assert!(snap.weekly.is_none());
2407 }
2408
2409 #[test]
2412 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
2413 let summary = serde_json::json!({
2414 "groups": [{
2415 "displayName": "Gemini Models",
2416 "buckets": [{
2417 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
2418 }],
2419 }],
2420 });
2421
2422 let rendered = parse_quota_summary(&summary, "Pro".into())
2423 .expect_err("an unrecognised cadence alone is not a snapshot")
2424 .to_string();
2425
2426 assert!(
2427 rendered.contains("no bucket in a window we recognise"),
2428 "{rendered}"
2429 );
2430 assert!(rendered.contains("gemini-daily"), "{rendered}");
2431 assert!(rendered.contains("window daily"), "{rendered}");
2432 }
2433
2434 #[test]
2437 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
2438 let summary = serde_json::json!({
2439 "groups": [{"displayName": "Gemini", "buckets": []}],
2440 });
2441
2442 let rendered = parse_quota_summary(&summary, "Pro".into())
2443 .expect_err("no buckets is an error")
2444 .to_string();
2445
2446 assert!(rendered.contains("no buckets at all"), "{rendered}");
2447 assert!(!rendered.contains("it offered:"), "{rendered}");
2448 }
2449
2450 #[test]
2453 fn buckets_without_an_id_are_still_named_in_the_error() {
2454 let summary = serde_json::json!({
2455 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
2456 });
2457
2458 let rendered = parse_quota_summary(&summary, "Pro".into())
2459 .expect_err("an unusable summary is an error")
2460 .to_string();
2461
2462 assert!(rendered.contains("<unnamed>"), "{rendered}");
2463 }
2464
2465 #[test]
2466 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
2467 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
2468 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
2469 }
2470
2471 #[test]
2472 fn cache_round_trip_preserves_every_window() {
2473 let snap = parsed();
2474 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2475 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2476 }
2477
2478 #[test]
2482 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
2483 let full = snap_to_json(&parsed());
2484 for missing in ["session_pct", "weekly_pct"] {
2485 let mut v = full.clone();
2486 v.as_object_mut().unwrap().remove(missing);
2487 let bytes = serde_json::to_vec(&v).unwrap();
2488 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2489 assert!(err.to_string().contains(missing), "{missing}: {err}");
2490 }
2491 assert!(parse_cache_at(b"{}", None, now()).is_err());
2493 }
2494
2495 #[test]
2496 fn cached_percentages_are_range_checked_before_narrowing() {
2497 let full = snap_to_json(&parsed());
2498 for (key, bad) in [
2499 ("session_pct", serde_json::json!(-1)),
2500 ("weekly_pct", serde_json::json!(101)),
2501 ("session_pct", serde_json::json!(i64::MAX)),
2502 ("tp_session_pct", serde_json::json!("75")),
2503 ] {
2504 let mut v = full.clone();
2505 v[key] = bad;
2506 let bytes = serde_json::to_vec(&v).unwrap();
2507 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2508 assert!(err.to_string().contains(key), "{key}: {err}");
2509 }
2510 }
2511
2512 #[test]
2513 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
2514 let mut v = snap_to_json(&parsed());
2515 v["session_reset"] = serde_json::json!("not-a-time");
2516 let bytes = serde_json::to_vec(&v).unwrap();
2517 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2518 assert!(err.to_string().contains("session_reset"), "{err}");
2519 }
2520
2521 #[test]
2526 fn a_cache_past_its_reset_is_refused() {
2527 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
2528 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2529
2530 assert!(parse_cache_at(&bytes, None, now()).is_ok());
2532 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
2534
2535 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
2537 assert!(err.to_string().contains("5h"), "{err}");
2538
2539 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
2541 }
2542
2543 #[test]
2546 fn expiry_names_the_window_that_rolled_over() {
2547 let mut snap = parsed();
2548 snap.session.as_mut().unwrap().resets_at = None;
2550 snap.third_party_session = None;
2551 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2552 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2553
2554 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
2556
2557 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
2559 assert!(err.to_string().contains("Gemini weekly"), "{err}");
2560 }
2561
2562 #[test]
2564 fn a_window_without_a_reset_never_expires() {
2565 let mut snap = parsed();
2566 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
2567 w.resets_at = None;
2568 }
2569 snap.third_party_session = None;
2570 snap.third_party_weekly = None;
2571 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2572 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
2573 .unwrap()
2574 .with_timezone(&Utc);
2575 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
2576 }
2577
2578 #[test]
2580 fn a_cache_from_another_account_is_rejected() {
2581 let mut snap = parsed();
2582 snap.account = "acct:aaaa".into();
2583 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2584
2585 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
2586 assert_eq!(
2587 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
2588 snap
2589 );
2590
2591 let mut legacy = snap_to_json(&snap);
2593 legacy.as_object_mut().unwrap().remove("account");
2594 let legacy = serde_json::to_vec(&legacy).unwrap();
2595 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
2596 }
2597
2598 #[test]
2601 fn an_unverifiable_cache_is_served_rather_than_discarded() {
2602 let mut snap = parsed();
2603 snap.account = "acct:aaaa".into();
2604 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2605 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2606 }
2607
2608 #[test]
2609 fn account_key_fingerprints_rather_than_storing_the_address() {
2610 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
2611 let a = with("someone@example.com");
2612 assert!(!a.contains("someone"), "{a}");
2613 assert!(!a.contains('@'), "{a}");
2614 assert_eq!(a, with("someone@example.com"), "must be stable");
2615 assert_ne!(a, with("other@example.com"));
2616 let unknown = account_key(&serde_json::json!({}));
2618 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
2619 assert_ne!(unknown, a);
2620 }
2621
2622 #[test]
2625 fn absent_third_party_windows_are_not_treated_as_corruption() {
2626 let mut snap = parsed();
2627 snap.third_party_session = None;
2628 snap.third_party_weekly = None;
2629 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2630 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2631 }
2632
2633 #[test]
2634 fn cache_round_trip_preserves_absent_third_party_windows() {
2635 let mut snap = parsed();
2636 snap.third_party_session = None;
2637 snap.third_party_weekly = None;
2638 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2639 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2640 }
2641
2642 #[test]
2643 fn pct_used_inverts_valid_fractions() {
2644 assert_eq!(pct_used(1.0), 0);
2645 assert_eq!(pct_used(0.0), 100);
2646 assert_eq!(pct_used(0.5), 50);
2647 }
2648
2649 #[test]
2650 fn plan_falls_back_through_the_status_payload() {
2651 let tier: serde_json::Value =
2652 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2653 .unwrap();
2654 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2655
2656 let plan_only: serde_json::Value = serde_json::from_str(
2657 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2658 )
2659 .unwrap();
2660 assert_eq!(plan_from_status(&plan_only), "Pro");
2661
2662 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2663 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2664 }
2665
2666 #[cfg(target_os = "linux")]
2667 #[test]
2668 fn proc_net_parser_keeps_only_listening_rows() {
2669 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2670 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2671 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2672
2673 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2674 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2675 assert_eq!(parse_proc_net_line(established), None);
2676
2677 assert_eq!(parse_proc_net_line("garbage"), None);
2678 }
2679
2680 #[test]
2681 fn explicit_address_comes_first_and_gets_a_scheme() {
2682 let candidate = |base: &str, pid| Candidate {
2683 base: base.to_string(),
2684 pid,
2685 ide: false,
2686 };
2687 assert_eq!(
2688 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)], &no_ide()),
2689 vec![
2690 candidate("http://127.0.0.1:1234", None),
2691 candidate("http://127.0.0.1:5678", Some(10)),
2692 ]
2693 );
2694 assert_eq!(
2696 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)], &no_ide()),
2697 vec![
2698 candidate("http://127.0.0.1:1234", None),
2699 candidate("http://127.0.0.1:5678", Some(10)),
2700 ]
2701 );
2702 assert_eq!(
2705 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)], &no_ide()),
2706 vec![candidate("http://127.0.0.1:5678", Some(10))]
2707 );
2708 assert_eq!(
2710 candidate_bases_with(None, vec![(10, 5678), (10, 5678)], &no_ide()),
2711 vec![candidate("http://127.0.0.1:5678", Some(10))]
2712 );
2713 assert_eq!(
2715 candidate_bases_with(Some("https://host:9"), vec![], &no_ide()),
2716 vec![candidate("https://host:9", None)]
2717 );
2718 }
2719
2720 fn no_ide() -> std::collections::HashSet<u32> {
2723 std::collections::HashSet::new()
2724 }
2725
2726 fn http(status: u16) -> AppError {
2727 AppError::Http {
2728 status,
2729 body: String::new(),
2730 }
2731 }
2732
2733 fn missing_csrf() -> AppError {
2734 missing_csrf_with("missing CSRF token")
2735 }
2736
2737 fn missing_csrf_with(message: &str) -> AppError {
2738 AppError::Http {
2739 status: 401,
2740 body: serde_json::json!({
2741 "code": "unauthenticated",
2742 "message": message,
2743 })
2744 .to_string(),
2745 }
2746 }
2747
2748 #[test]
2749 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2750 assert!(is_missing_csrf(&missing_csrf()));
2751 assert!(is_missing_csrf(&missing_csrf_with(
2752 "unauthenticated: missing CSRF token"
2753 )));
2754 assert!(is_missing_csrf(&missing_csrf_with(
2755 " unauthenticated: missing CSRF token "
2756 )));
2757 assert!(!is_missing_csrf(&missing_csrf_with(
2758 "unauthenticated: invalid token"
2759 )));
2760 assert!(!is_missing_csrf(&AppError::Http {
2761 status: 403,
2762 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2763 }));
2764 assert!(!is_missing_csrf(&AppError::Http {
2765 status: 401,
2766 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2767 }));
2768 assert!(!is_missing_csrf(&AppError::Http {
2769 status: 401,
2770 body: "prefix: missing CSRF token".into(),
2771 }));
2772 }
2773
2774 #[test]
2775 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2776 for errors in [
2777 vec![missing_csrf(), http(403)],
2778 vec![http(401), missing_csrf()],
2779 ] {
2780 let err = select_probe_error(errors);
2781 assert!(is_actionable(&err), "{err}");
2782 assert!(!is_missing_csrf(&err), "{err}");
2783 }
2784 }
2785
2786 #[test]
2787 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2788 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2789 assert!(is_missing_csrf(&err), "{err}");
2790 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2791 }
2792
2793 #[test]
2797 fn an_auth_failure_outranks_later_transport_noise() {
2798 let err = select_probe_error(vec![
2799 http(401),
2800 AppError::Transport("connection refused".into()),
2801 ]);
2802 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2803
2804 let err = select_probe_error(vec![
2805 AppError::Transport("connection refused".into()),
2806 http(403),
2807 ]);
2808 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2809 }
2810
2811 #[test]
2814 fn the_first_auth_failure_wins() {
2815 let err = select_probe_error(vec![http(401), http(403)]);
2816 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2817 }
2818
2819 #[test]
2823 fn without_an_auth_failure_the_last_error_stands() {
2824 let err = select_probe_error(vec![
2825 AppError::Transport("first".into()),
2826 http(500),
2827 AppError::Transport("last".into()),
2828 ]);
2829 assert!(
2830 matches!(&err, AppError::Transport(m) if m == "last"),
2831 "{err}"
2832 );
2833 assert!(err.is_transient());
2834 }
2835
2836 #[test]
2839 fn a_server_error_is_not_treated_as_actionable() {
2840 let err = select_probe_error(vec![http(500), http(401)]);
2841 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2842 }
2843
2844 #[test]
2845 fn no_candidates_at_all_yields_a_generic_error() {
2846 let err = select_probe_error(Vec::new());
2847 assert!(
2848 err.to_string().contains("no local server answered"),
2849 "{err}"
2850 );
2851 }
2852
2853 fn tls_echo() -> AppError {
2857 AppError::Http {
2858 status: 400,
2859 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2860 }
2861 }
2862
2863 #[test]
2868 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2869 let err = select_probe_error(vec![
2870 AppError::Http {
2871 status: 500,
2872 body: "GetUserStatus: internal".into(),
2873 },
2874 tls_echo(),
2875 ]);
2876 assert!(
2877 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2878 "{err}"
2879 );
2880 }
2881
2882 #[test]
2887 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2888 let err = select_probe_error(vec![
2889 AppError::Transport("connection refused".into()),
2890 tls_echo(),
2891 ]);
2892 assert!(
2893 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2894 "{err}"
2895 );
2896 assert!(
2897 err.is_transient(),
2898 "the echo must not make the run non-transient: {err}"
2899 );
2900 }
2901
2902 #[test]
2905 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2906 let err = select_probe_error(vec![tls_echo()]);
2907 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2908 }
2909
2910 #[test]
2913 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2914 let err = select_probe_error(vec![
2915 AppError::Http {
2916 status: 400,
2917 body: "unknown method GetUserStatus".into(),
2918 },
2919 tls_echo(),
2920 ]);
2921 assert!(
2922 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2923 "{err}"
2924 );
2925 }
2926
2927 #[test]
2929 fn an_auth_failure_still_outranks_a_tls_echo() {
2930 let err = select_probe_error(vec![tls_echo(), http(401)]);
2931 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2932 }
2933
2934 #[test]
2935 fn every_discovered_port_is_probed_in_order() {
2936 assert_eq!(
2937 candidate_bases_with(None, vec![(10, 33875), (10, 37435)], &no_ide()),
2938 vec![
2939 Candidate {
2940 base: "http://127.0.0.1:33875".to_string(),
2941 pid: Some(10),
2942 ide: false,
2943 },
2944 Candidate {
2945 base: "http://127.0.0.1:37435".to_string(),
2946 pid: Some(10),
2947 ide: false,
2948 },
2949 ]
2950 );
2951 }
2952
2953 #[test]
2958 fn empty_discovery_yields_no_candidates() {
2959 assert!(candidate_bases_with(None, vec![], &no_ide()).is_empty());
2960 assert!(candidate_bases_with(Some(""), vec![], &no_ide()).is_empty());
2961 }
2962
2963 #[test]
2964 fn every_antigravity_product_is_recognised() {
2965 assert!(is_antigravity_process(
2967 "language_server\n",
2968 Some("/opt/antigravity/resources/bin/language_server")
2969 ));
2970 assert!(is_antigravity_process(
2972 "agy\n",
2973 Some("/home/u/.local/bin/agy")
2974 ));
2975 assert!(is_antigravity_process(
2977 "node",
2978 Some("/opt/antigravity/bin/helper")
2979 ));
2980 assert!(is_antigravity_process("antigravity", None));
2981 assert!(is_antigravity_process("agy.exe", None));
2982 assert!(is_antigravity_process("Antigravity.exe", None));
2983 assert!(is_antigravity_process("language_server.exe", None));
2984 assert!(is_antigravity_process(
2985 "language_server_windows_x64.exe",
2986 None
2987 ));
2988 assert!(is_antigravity_process(
2989 "node.exe",
2990 Some(r"C:\Users\u\AppData\Local\agy.exe")
2991 ));
2992 }
2993
2994 #[test]
2995 fn unrelated_processes_are_not_probed() {
2996 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2997 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2998 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
3000 assert!(!is_antigravity_process("legacy.exe", None));
3001 assert!(!is_antigravity_process("not-agy.exe", None));
3002 assert!(!is_antigravity_process("", None));
3003 }
3004
3005 #[test]
3006 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
3007 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
3008 raw.extend([0, b'x' as u16]);
3009 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
3010 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
3011 assert_eq!(decode_windows_process_name(&[]), "");
3012 }
3013
3014 #[test]
3015 fn windows_process_filter_keeps_only_antigravity_pids() {
3016 let processes = vec![
3017 (10, "agy.exe".to_string()),
3018 (20, "language_server_windows_x64.exe".to_string()),
3019 (30, "sshd.exe".to_string()),
3020 ];
3021 let pids = matching_windows_process_ids(&processes);
3022 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
3023 }
3024
3025 #[test]
3026 fn windows_listener_filter_joins_pid_loopback_and_port() {
3027 let pids = std::collections::HashSet::from([10]);
3028 let rows = [
3029 WindowsTcpRow {
3030 local_addr: [127, 0, 0, 1],
3031 local_port: u32::from(59870u16.to_be()),
3032 pid: 10,
3033 },
3034 WindowsTcpRow {
3035 local_addr: [127, 0, 0, 1],
3036 local_port: u32::from(59868u16.to_be()),
3037 pid: 10,
3038 },
3039 WindowsTcpRow {
3040 local_addr: [127, 0, 0, 1],
3041 local_port: u32::from(59870u16.to_be()),
3042 pid: 10,
3043 },
3044 WindowsTcpRow {
3045 local_addr: [0, 0, 0, 0],
3046 local_port: u32::from(50000u16.to_be()),
3047 pid: 10,
3048 },
3049 WindowsTcpRow {
3050 local_addr: [127, 0, 0, 1],
3051 local_port: u32::from(50001u16.to_be()),
3052 pid: 99,
3053 },
3054 WindowsTcpRow {
3055 local_addr: [127, 0, 0, 1],
3056 local_port: 0,
3057 pid: 10,
3058 },
3059 ];
3060 assert_eq!(
3061 matching_windows_ports(&pids, &rows),
3062 vec![(10, 59870), (10, 59868)]
3063 );
3064 }
3065
3066 #[test]
3070 fn windows_ports_from_two_products_keep_tls_listeners_last() {
3071 let pids = std::collections::HashSet::from([10, 20]);
3072 let row = |port: u16, pid: u32| WindowsTcpRow {
3073 local_addr: [127, 0, 0, 1],
3074 local_port: u32::from(port.to_be()),
3075 pid,
3076 };
3077 let rows = [
3078 row(40000, 10),
3079 row(40001, 10),
3080 row(50000, 20),
3081 row(50001, 20),
3082 ];
3083 assert_eq!(
3084 matching_windows_ports(&pids, &rows),
3085 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
3086 );
3087 }
3088
3089 #[test]
3093 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
3094 use std::collections::BTreeMap;
3095
3096 assert_eq!(
3098 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
3099 vec![(10, 59870), (10, 59868)]
3100 );
3101 assert_eq!(
3106 probe_order(BTreeMap::from([
3107 (10, vec![40000, 40001]),
3108 (20, vec![50000, 50001]),
3109 ])),
3110 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
3111 );
3112 assert_eq!(
3115 probe_order(BTreeMap::from([
3116 (10, vec![6000, 5000, 4000]),
3117 (20, vec![6000, 7000]),
3118 ])),
3119 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
3120 );
3121 assert!(probe_order(BTreeMap::new()).is_empty());
3122 }
3123
3124 #[test]
3129 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
3130 use std::collections::BTreeMap;
3131
3132 assert_eq!(
3133 probe_order(BTreeMap::from([
3134 (10, vec![40001, 40001, 40000, 40000]),
3135 (20, vec![50001, 50000]),
3136 ])),
3137 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
3138 "duplicate rows must not reorder the ranks below them"
3139 );
3140 }
3141
3142 #[test]
3147 fn a_half_started_product_is_the_documented_exception() {
3148 use std::collections::BTreeMap;
3149
3150 assert_eq!(
3151 probe_order(BTreeMap::from([
3152 (10, vec![40000]),
3153 (20, vec![50001, 50000])
3154 ])),
3155 vec![(10, 40000), (20, 50001), (20, 50000)]
3156 );
3157 }
3158
3159 #[test]
3164 fn an_override_with_no_authority_is_dropped_not_probed() {
3165 for junk in ["/", "///", "http://", "https://", " / "] {
3166 assert_eq!(
3167 candidate_bases_with(Some(junk), vec![(10, 4242)], &no_ide()),
3168 vec![Candidate {
3169 base: "http://127.0.0.1:4242".to_string(),
3170 pid: Some(10),
3171 ide: false,
3172 }],
3173 "{junk:?} should not survive as a candidate"
3174 );
3175 }
3176 assert!(candidate_bases_with(Some("/"), vec![], &no_ide()).is_empty());
3177 }
3178
3179 #[test]
3180 fn windows_table_bounds_reject_truncation_and_overflow() {
3181 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
3182 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
3183 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
3184 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
3185 assert_eq!(
3186 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
3187 None
3188 );
3189 }
3190
3191 #[cfg(target_os = "windows")]
3192 #[test]
3193 fn windows_tcp_table_parser_copies_complete_rows_only() {
3194 use std::mem::{offset_of, size_of};
3195 use windows_sys::Win32::NetworkManagement::IpHelper::{
3196 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
3197 };
3198
3199 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
3200 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
3201 let words = used.div_ceil(size_of::<u32>());
3202 let mut buffer = vec![0u32; words];
3203 let first = MIB_TCPROW_OWNER_PID {
3204 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
3205 dwLocalPort: u32::from(59868u16.to_be()),
3206 dwOwningPid: 10,
3207 ..Default::default()
3208 };
3209 let second = MIB_TCPROW_OWNER_PID {
3210 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
3211 dwLocalPort: u32::from(59870u16.to_be()),
3212 dwOwningPid: 10,
3213 ..Default::default()
3214 };
3215 unsafe {
3216 buffer.as_mut_ptr().write_unaligned(2);
3217 let rows = buffer
3218 .as_mut_ptr()
3219 .cast::<u8>()
3220 .add(offset)
3221 .cast::<MIB_TCPROW_OWNER_PID>();
3222 rows.write_unaligned(first);
3223 rows.add(1).write_unaligned(second);
3224 }
3225
3226 assert_eq!(
3227 parse_windows_tcp_rows(&buffer, used),
3228 vec![
3229 WindowsTcpRow {
3230 local_addr: [127, 0, 0, 1],
3231 local_port: u32::from(59868u16.to_be()),
3232 pid: 10,
3233 },
3234 WindowsTcpRow {
3235 local_addr: [127, 0, 0, 1],
3236 local_port: u32::from(59870u16.to_be()),
3237 pid: 10,
3238 },
3239 ]
3240 );
3241 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
3242 }
3243
3244 #[test]
3245 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
3246 let output = "p74101\ncagy\nf10\nn127.0.0.1:8829\nf11\nn127.0.0.1:61289\nf12\nn127.0.0.1:61290\np200\ncsshd\nf5\nn*:22\n";
3249 assert_eq!(
3250 parse_lsof_pcn(output),
3251 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
3252 );
3253 }
3254
3255 #[test]
3259 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
3260 let output = concat!(
3261 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
3262 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
3263 );
3264 assert_eq!(
3265 parse_lsof_pcn(output),
3266 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
3267 "both HTTP listeners must precede both TLS listeners"
3268 );
3269 }
3270
3271 #[test]
3272 fn lsof_parser_matches_the_capitalised_macos_app_name() {
3273 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
3274 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
3275 }
3276
3277 #[test]
3278 fn lsof_parser_deduplicates_and_handles_empty_output() {
3279 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
3280 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
3281 assert!(parse_lsof_pcn("").is_empty());
3282 }
3283
3284 #[test]
3288 fn macos_system_tools_are_pinned_to_trusted_paths() {
3289 let sources = [
3290 crate::guard::production_code(include_str!("fetch.rs")),
3291 crate::guard::production_code(include_str!("../tray/tui_launch_macos.rs")),
3292 crate::guard::production_code(include_str!("../tray/browse.rs")),
3293 crate::guard::production_code(include_str!("../nous/oauth.rs")),
3294 crate::guard::production_code(include_str!("../../macos/ai-usagebar-menubar.swift")),
3295 ];
3296 let code = sources.join("\n");
3297 for bare in [
3298 r#"Command::new("lsof")"#,
3299 r#"Command::new("ps")"#,
3300 r#"Command::new("open")"#,
3301 r#"Command::new("osascript")"#,
3302 r#"do script \"bash "#,
3303 ] {
3304 assert!(!code.contains(bare), "{bare} looks the tool up on PATH");
3305 }
3306 for pinned in [
3307 r#"Command::new("/usr/sbin/lsof")"#,
3308 r#"Command::new("/bin/ps")"#,
3309 r#"Command::new("/usr/bin/open")"#,
3310 r#"Command::new("/usr/bin/osascript")"#,
3311 r#"do script \"/bin/bash "#,
3312 ] {
3313 assert!(
3314 code.contains(pinned),
3315 "missing trusted system path: {pinned}"
3316 );
3317 }
3318 }
3319
3320 #[test]
3323 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
3324 let dir = tempfile::tempdir().unwrap();
3325 let cache = Cache::at(dir.path().join("usage.json"));
3326 let reason = AppError::Credentials("Antigravity: no local language server found".into());
3327
3328 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
3329 let msg = err.to_string();
3330 assert!(msg.contains("no local language server found"), "{msg}");
3331 assert!(!msg.contains("no usable cache"), "{msg}");
3332 }
3333
3334 #[test]
3335 fn unusable_cache_does_not_replace_the_live_diagnosis() {
3336 let dir = tempfile::tempdir().unwrap();
3337 let cache = Cache::at(dir.path().join("antigravity"));
3338 cache.write_payload(b"{}").unwrap();
3339
3340 let reason = AppError::Credentials("Antigravity must be running".into());
3341 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
3342 assert!(err.to_string().contains("must be running"), "{err}");
3343
3344 let original = AppError::Transport("original loopback failure".into());
3345 let err = fallback_silent(&cache, now(), original).unwrap_err();
3346 assert!(
3347 err.to_string().contains("original loopback failure"),
3348 "{err}"
3349 );
3350 }
3351
3352 #[tokio::test]
3353 async fn rpc_error_bodies_are_bounded_too() {
3354 let mut server = mockito::Server::new_async().await;
3355 let path = format!("/{STATUS_RPC}");
3356 server
3357 .mock("POST", path.as_str())
3358 .with_status(500)
3359 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
3360 .create_async()
3361 .await;
3362
3363 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
3364 .await
3365 .unwrap_err();
3366 assert!(err.to_string().contains("exceeds"), "{err}");
3367 }
3368
3369 #[test]
3370 fn blank_override_falls_through_to_discovery() {
3371 assert_eq!(
3372 candidate_bases_with(Some(" "), vec![(10, 4242)], &no_ide()),
3373 vec![Candidate {
3374 base: "http://127.0.0.1:4242".to_string(),
3375 pid: Some(10),
3376 ide: false,
3377 }]
3378 );
3379 }
3380
3381 #[test]
3384 fn only_the_ide_install_is_recognised() {
3385 assert!(is_ide_install(
3386 "/opt/antigravity-ide/resources/app/extensions/antigravity/bin/language_server_linux_x64"
3387 ));
3388 assert!(is_ide_install(
3389 "/Applications/Antigravity IDE.app/Contents/Resources/app/extensions/antigravity/bin/language_server_darwin_arm64"
3390 ));
3391 assert!(is_ide_install(
3392 r"C:\Users\u\AppData\Local\Programs\Antigravity IDE\language_server_windows_x64.exe"
3393 ));
3394 assert!(!is_ide_install(
3395 "/opt/antigravity/resources/bin/language_server"
3396 ));
3397 assert!(!is_ide_install("/home/u/.local/bin/agy"));
3398 assert!(!is_ide_install("/usr/bin/node"));
3399 }
3400
3401 #[test]
3403 fn ps_output_pairs_pids_with_full_paths() {
3404 let output = " 101 /opt/antigravity-ide/bin/language_server_linux_x64\n\
3405 202 /Applications/Antigravity IDE.app/Contents/MacOS/Antigravity IDE\n";
3406 assert_eq!(
3407 parse_ps_comm(output),
3408 vec![
3409 (
3410 101,
3411 "/opt/antigravity-ide/bin/language_server_linux_x64".to_string()
3412 ),
3413 (
3414 202,
3415 "/Applications/Antigravity IDE.app/Contents/MacOS/Antigravity IDE".to_string()
3416 ),
3417 ]
3418 );
3419 assert!(parse_ps_comm("").is_empty());
3420 assert!(parse_ps_comm("not-a-pid /tmp/x").is_empty());
3421 }
3422
3423 #[test]
3427 fn ide_owned_listeners_are_marked_but_an_explicit_override_is_not() {
3428 let ide: std::collections::HashSet<u32> = [10].into_iter().collect();
3429
3430 let candidates = candidate_bases_with(None, vec![(10, 40001), (20, 50001)], &ide);
3431 assert_eq!(candidates.len(), 2);
3432 assert!(candidates[0].ide, "pid 10 is the IDE's language server");
3433 assert!(!candidates[1].ide);
3434
3435 let candidates = candidate_bases_with(Some("127.0.0.1:40001"), vec![(10, 40001)], &ide);
3436 assert_eq!(
3437 candidates,
3438 vec![Candidate {
3439 base: "http://127.0.0.1:40001".to_string(),
3440 pid: Some(10),
3441 ide: false,
3442 }],
3443 "ANTIGRAVITY_LS_ADDRESS is an explicit target, never skipped"
3444 );
3445 }
3446
3447 fn fixture() -> (tempfile::TempDir, Cache) {
3452 let td = tempfile::TempDir::new().unwrap();
3453 let cache = Cache::at(td.path().join("antigravity"));
3454 (td, cache)
3455 }
3456
3457 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
3458 let base = server.url();
3459 cloud::Endpoints {
3460 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
3461 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
3462 token: format!("{base}/token"),
3463 }
3464 }
3465
3466 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
3468 let mut token = serde_json::json!({
3469 "access_token": "KEYRING-AT",
3470 "expiry": expiry,
3471 });
3472 if with_refresh {
3473 token["refresh_token"] = serde_json::json!("KEYRING-RT");
3474 }
3475 serde_json::json!({ "token": token }).to_string()
3476 }
3477
3478 const EXPIRED: &str = "2026-07-22T11:00:00Z";
3480 const VALID: &str = "2026-07-22T13:00:00Z";
3482
3483 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
3485 RemoteOverride {
3486 credential: SavedCredential::Blob(blob),
3487 endpoints: Some(eps),
3488 local_bases: Some(vec![]),
3489 ..Default::default()
3490 }
3491 }
3492
3493 fn bare_summary() -> String {
3496 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
3497 v["response"].to_string()
3498 }
3499
3500 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
3501 server
3502 .mock("POST", "/daily/quota")
3503 .match_header("authorization", format!("Bearer {bearer}").as_str())
3504 .with_status(200)
3505 .with_body(bare_summary())
3506 }
3507
3508 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
3509 server
3510 .mock("POST", "/token")
3511 .with_status(200)
3512 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
3513 }
3514
3515 fn test_oauth() -> cloud::OauthClient {
3516 cloud::OauthClient {
3517 id: "test-client".into(),
3518 secret: "test-client-secret".into(),
3519 }
3520 }
3521
3522 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
3523 fetch_snapshot_at(
3524 &reqwest::Client::new(),
3525 cache,
3526 ttl,
3527 Some(&test_oauth()),
3528 remote,
3529 now(),
3530 )
3531 .await
3532 }
3533
3534 #[tokio::test]
3535 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
3536 let mut server = mockito::Server::new_async().await;
3537 let eps = endpoints(&server);
3538 let quota = quota_mock(&mut server, "KEYRING-AT")
3539 .expect(1)
3540 .create_async()
3541 .await;
3542 let plan = server
3543 .mock("POST", "/daily/plan")
3544 .match_header("authorization", "Bearer KEYRING-AT")
3545 .with_status(200)
3546 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3547 .expect(1)
3548 .create_async()
3549 .await;
3550 let token = token_mock(&mut server).expect(0).create_async().await;
3552 let blob = keyring_blob(VALID, true);
3553 let (_td, cache) = fixture();
3554
3555 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
3556 .await
3557 .expect("remote path yields a snapshot");
3558
3559 quota.assert_async().await;
3560 plan.assert_async().await;
3561 token.assert_async().await;
3562 let snap = outcome.snapshot;
3563 assert!(!outcome.stale);
3564 assert_eq!(snap.source, AntigravitySource::Remote);
3565 assert_eq!(snap.plan, "Pro");
3566 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
3567 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
3568 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
3569 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
3570 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3571 assert_eq!(snap.account, format!("acct:{fingerprint}"));
3572 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
3573
3574 let cached = parse_cache_at(
3576 &cache
3577 .fresh_payload(Duration::from_secs(60))
3578 .unwrap()
3579 .unwrap(),
3580 None,
3581 now(),
3582 )
3583 .unwrap();
3584 assert_eq!(cached.source, AntigravitySource::Remote);
3585 assert_eq!(cached.account, snap.account);
3586 }
3587
3588 fn cached_remote(cache: &Cache, blob: &str, expired: bool) -> AntigravitySnapshot {
3589 let mut snap = parsed();
3590 snap.source = AntigravitySource::Remote;
3591 snap.account = remote_account(&credential::parse_keyring_blob(blob).unwrap().fingerprint);
3592 if expired {
3593 snap.session.as_mut().unwrap().resets_at = Some(now());
3594 }
3595 cache
3596 .write_payload(&serde_json::to_vec(&snap_to_json(&snap)).unwrap())
3597 .unwrap();
3598 snap
3599 }
3600
3601 #[tokio::test]
3602 async fn an_active_backoff_blocks_cloud_requests_after_a_window_reset() {
3603 for expiry in [VALID, EXPIRED] {
3604 let mut server = mockito::Server::new_async().await;
3605 let eps = endpoints(&server);
3606 let network = server
3607 .mock("POST", mockito::Matcher::Any)
3608 .with_status(429)
3609 .expect(0)
3610 .create_async()
3611 .await;
3612 let blob = keyring_blob(expiry, true);
3613 let (_td, cache) = fixture();
3614 cached_remote(&cache, &blob, true);
3615 cache.write_last_error(429, "quota limited");
3616 let marker = std::fs::read(cache.retry_after_path()).unwrap();
3617
3618 let error = run(&cache, remote(&blob, &eps), Duration::ZERO)
3619 .await
3620 .unwrap_err();
3621
3622 network.assert_async().await;
3623 assert!(matches!(error, AppError::Http { status: 429, .. }));
3624 assert!(error.to_string().contains("next attempt in"));
3625 assert_eq!(std::fs::read(cache.retry_after_path()).unwrap(), marker);
3626 assert_eq!(cache.read_last_error(), Some((429, "quota limited".into())));
3627 }
3628 }
3629
3630 #[tokio::test]
3631 async fn an_active_backoff_keeps_an_unexpired_remote_cache_visible() {
3632 let mut server = mockito::Server::new_async().await;
3633 let eps = endpoints(&server);
3634 let network = server
3635 .mock("POST", mockito::Matcher::Any)
3636 .expect(0)
3637 .create_async()
3638 .await;
3639 let blob = keyring_blob(VALID, true);
3640 let (_td, cache) = fixture();
3641 let snap = cached_remote(&cache, &blob, false);
3642 cache.write_last_error(429, "quota limited");
3643
3644 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
3645 .await
3646 .unwrap();
3647
3648 network.assert_async().await;
3649 assert!(outcome.stale);
3650 assert_eq!(outcome.snapshot, snap);
3651 assert_eq!(outcome.last_error, Some((429, "quota limited".into())));
3652 }
3653
3654 #[tokio::test]
3655 async fn an_active_backoff_does_not_block_a_local_quota_after_a_window_reset() {
3656 let mut server = mockito::Server::new_async().await;
3657 let eps = endpoints(&server);
3658 let status_path = format!("/{STATUS_RPC}");
3659 server
3660 .mock("POST", status_path.as_str())
3661 .with_status(200)
3662 .with_body(r#"{"userStatus":{"email":"test@example.test"}}"#)
3663 .expect(1)
3664 .create_async()
3665 .await;
3666 let quota_path = format!("/{QUOTA_RPC}");
3667 let quota = server
3668 .mock("POST", quota_path.as_str())
3669 .with_status(200)
3670 .with_body(QUOTA_JSON)
3671 .expect(1)
3672 .create_async()
3673 .await;
3674 let blob = keyring_blob(VALID, true);
3675 let (_td, cache) = fixture();
3676 cached_remote(&cache, &blob, true);
3677 cache.write_last_error(429, "quota limited");
3678
3679 let outcome = run(
3680 &cache,
3681 RemoteOverride {
3682 local_bases: Some(vec![server.url().into()]),
3683 ..remote(&blob, &eps)
3684 },
3685 Duration::from_secs(60),
3686 )
3687 .await
3688 .unwrap();
3689
3690 quota.assert_async().await;
3691 assert!(!outcome.stale);
3692 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3693 }
3694
3695 #[cfg(unix)]
3696 #[tokio::test]
3697 async fn an_active_backoff_allows_agy_but_skips_its_cloud_plan_lookup() {
3698 let mut server = mockito::Server::new_async().await;
3699 let eps = endpoints(&server);
3700 let network = server
3701 .mock("POST", mockito::Matcher::Any)
3702 .expect(0)
3703 .create_async()
3704 .await;
3705 let blob = keyring_blob(VALID, true);
3706 let (td, cache) = fixture();
3707 cached_remote(&cache, &blob, true);
3708 cache.write_last_error(429, "quota limited");
3709 let runs = td.path().join("runs");
3710 let agy = agy_reading(td.path(), &runs);
3711
3712 let outcome = run(
3713 &cache,
3714 RemoteOverride {
3715 agy: super::super::agy::AgyCommand::At(agy),
3716 ..remote(&blob, &eps)
3717 },
3718 Duration::ZERO,
3719 )
3720 .await
3721 .unwrap();
3722
3723 network.assert_async().await;
3724 assert!(!outcome.stale);
3725 assert_eq!(outcome.snapshot.source, AntigravitySource::Agy);
3726 assert_eq!(outcome.snapshot.plan, DEFAULT_PLAN);
3727 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
3728 }
3729
3730 #[tokio::test]
3735 async fn agys_missing_csrf_response_uses_the_saved_session() {
3736 let mut server = mockito::Server::new_async().await;
3737 let eps = endpoints(&server);
3738 let root = server
3739 .mock("GET", "/")
3740 .with_status(404)
3741 .expect(1)
3742 .create_async()
3743 .await;
3744 let status_path = format!("/{STATUS_RPC}");
3745 let status = server
3746 .mock("POST", status_path.as_str())
3747 .with_status(401)
3748 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3749 .expect(1)
3750 .create_async()
3751 .await;
3752 let quota = quota_mock(&mut server, "KEYRING-AT")
3753 .expect(1)
3754 .create_async()
3755 .await;
3756 let plan = server
3757 .mock("POST", "/daily/plan")
3758 .match_header("authorization", "Bearer KEYRING-AT")
3759 .with_status(200)
3760 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3761 .expect(1)
3762 .create_async()
3763 .await;
3764 let blob = keyring_blob(VALID, true);
3765 let (_td, cache) = fixture();
3766
3767 let outcome = run(
3768 &cache,
3769 RemoteOverride {
3770 credential: SavedCredential::Blob(&blob),
3771 endpoints: Some(&eps),
3772 local_bases: Some(vec![server.url().into()]),
3773 ..Default::default()
3774 },
3775 Duration::ZERO,
3776 )
3777 .await
3778 .expect("saved session bypasses agy's unusable local RPC");
3779
3780 root.assert_async().await;
3781 status.assert_async().await;
3782 quota.assert_async().await;
3783 plan.assert_async().await;
3784 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3785 }
3786
3787 #[tokio::test]
3791 async fn agys_prefixed_missing_csrf_response_uses_the_saved_session() {
3792 let mut server = mockito::Server::new_async().await;
3793 let eps = endpoints(&server);
3794 let root = server
3795 .mock("GET", "/")
3796 .with_status(404)
3797 .expect(1)
3798 .create_async()
3799 .await;
3800 let status_path = format!("/{STATUS_RPC}");
3801 let status = server
3802 .mock("POST", status_path.as_str())
3803 .with_status(401)
3804 .with_body(
3805 r#"{"code":"unauthenticated","message":"unauthenticated: missing CSRF token"}"#,
3806 )
3807 .expect(1)
3808 .create_async()
3809 .await;
3810 let quota = quota_mock(&mut server, "KEYRING-AT")
3811 .expect(1)
3812 .create_async()
3813 .await;
3814 let plan = server
3815 .mock("POST", "/daily/plan")
3816 .match_header("authorization", "Bearer KEYRING-AT")
3817 .with_status(200)
3818 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3819 .expect(1)
3820 .create_async()
3821 .await;
3822 let blob = keyring_blob(VALID, true);
3823 let (_td, cache) = fixture();
3824
3825 let outcome = run(
3826 &cache,
3827 RemoteOverride {
3828 credential: SavedCredential::Blob(&blob),
3829 endpoints: Some(&eps),
3830 local_bases: Some(vec![server.url().into()]),
3831 ..Default::default()
3832 },
3833 Duration::ZERO,
3834 )
3835 .await
3836 .expect("saved session bypasses agy 1.3.1's prefixed missing-CSRF RPC");
3837
3838 root.assert_async().await;
3839 status.assert_async().await;
3840 quota.assert_async().await;
3841 plan.assert_async().await;
3842 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3843 }
3844
3845 #[tokio::test]
3846 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
3847 let mut server = mockito::Server::new_async().await;
3848 let eps = endpoints(&server);
3849 let root = server
3850 .mock("GET", "/")
3851 .with_status(404)
3852 .expect(1)
3853 .create_async()
3854 .await;
3855 let status_path = format!("/{STATUS_RPC}");
3856 let status = server
3857 .mock("POST", status_path.as_str())
3858 .with_status(401)
3859 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3860 .expect(1)
3861 .create_async()
3862 .await;
3863 let quota = quota_mock(&mut server, "KEYRING-AT")
3864 .expect(0)
3865 .create_async()
3866 .await;
3867 let (_td, cache) = fixture();
3868
3869 let error = run(
3870 &cache,
3871 RemoteOverride {
3872 credential: SavedCredential::Absent,
3873 endpoints: Some(&eps),
3874 local_bases: Some(vec![server.url().into()]),
3875 ..Default::default()
3876 },
3877 Duration::ZERO,
3878 )
3879 .await
3880 .expect_err("neither the local RPC nor a saved session is usable");
3881
3882 root.assert_async().await;
3883 status.assert_async().await;
3884 quota.assert_async().await;
3885 let message = error.to_string();
3886 assert!(message.contains("requires a CSRF token"), "{message}");
3887 assert!(message.contains("saved Google session"), "{message}");
3888 }
3889
3890 #[tokio::test]
3895 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
3896 let mut server1 = mockito::Server::new_async().await;
3897 let mut server2 = mockito::Server::new_async().await;
3898 let eps = endpoints(&server1);
3899 let root1 = server1
3900 .mock("GET", "/")
3901 .with_status(404)
3902 .expect(1)
3903 .create_async()
3904 .await;
3905 let status_path = format!("/{STATUS_RPC}");
3906 let status1 = server1
3907 .mock("POST", status_path.as_str())
3908 .with_status(401)
3909 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3910 .expect(1)
3911 .create_async()
3912 .await;
3913
3914 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
3916 let status2 = server2
3917 .mock("POST", status_path.as_str())
3918 .expect(0)
3919 .create_async()
3920 .await;
3921
3922 let (_td, cache) = fixture();
3923 let error = run(
3924 &cache,
3925 RemoteOverride {
3926 credential: SavedCredential::Absent,
3927 endpoints: Some(&eps),
3928 local_bases: Some(vec![
3929 Candidate {
3930 base: server1.url(),
3931 pid: Some(42),
3932 ide: false,
3933 },
3934 Candidate {
3935 base: server2.url(),
3936 pid: Some(42),
3937 ide: false,
3938 },
3939 ]),
3940 ..Default::default()
3941 },
3942 Duration::ZERO,
3943 )
3944 .await
3945 .expect_err("aborted probe returns agy missing csrf diagnosis");
3946
3947 root1.assert_async().await;
3948 status1.assert_async().await;
3949 root2.assert_async().await;
3950 status2.assert_async().await;
3951 let message = error.to_string();
3952 assert!(message.contains("requires a CSRF token"), "{message}");
3953 }
3954
3955 #[tokio::test]
3962 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
3963 let mut agy = mockito::Server::new_async().await;
3964 let mut ide = mockito::Server::new_async().await;
3965 let eps = endpoints(&agy);
3966 let agy_root = agy
3967 .mock("GET", "/")
3968 .with_status(404)
3969 .expect(1)
3970 .create_async()
3971 .await;
3972 let status_path = format!("/{STATUS_RPC}");
3973 let agy_status = agy
3974 .mock("POST", status_path.as_str())
3975 .with_status(401)
3976 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3977 .expect(1)
3978 .create_async()
3979 .await;
3980
3981 let ide_root = ide
3982 .mock("GET", "/")
3983 .with_status(404)
3984 .expect(1)
3985 .create_async()
3986 .await;
3987 let ide_status = ide
3988 .mock("POST", status_path.as_str())
3989 .with_status(200)
3990 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3991 .expect(1)
3992 .create_async()
3993 .await;
3994 let quota_path = format!("/{QUOTA_RPC}");
3995 let ide_quota = ide
3996 .mock("POST", quota_path.as_str())
3997 .with_status(200)
3998 .with_body(QUOTA_JSON)
3999 .expect(1)
4000 .create_async()
4001 .await;
4002
4003 let (_td, cache) = fixture();
4004 let outcome = run(
4005 &cache,
4006 RemoteOverride {
4007 credential: SavedCredential::Absent,
4008 endpoints: Some(&eps),
4009 local_bases: Some(vec![
4010 Candidate {
4011 base: agy.url(),
4012 pid: Some(42),
4013 ide: false,
4014 },
4015 Candidate {
4016 base: ide.url(),
4017 pid: Some(7),
4018 ide: false,
4019 },
4020 ]),
4021 ..Default::default()
4022 },
4023 Duration::ZERO,
4024 )
4025 .await
4026 .expect("the other product's server still answers");
4027
4028 agy_root.assert_async().await;
4029 agy_status.assert_async().await;
4030 ide_root.assert_async().await;
4031 ide_status.assert_async().await;
4032 ide_quota.assert_async().await;
4033 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
4034 assert_eq!(outcome.snapshot.plan, "Pro");
4035 }
4036
4037 #[tokio::test]
4040 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
4041 let mut server = mockito::Server::new_async().await;
4042 let eps = endpoints(&server);
4043 let _quota = quota_mock(&mut server, "NEW-AT")
4044 .expect(2)
4045 .create_async()
4046 .await;
4047 let token = token_mock(&mut server).expect(1).create_async().await;
4048 let blob = keyring_blob(EXPIRED, true);
4049 let (_td, cache) = fixture();
4050
4051 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
4052 .await
4053 .expect("refresh, then quota");
4054 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
4055 token.assert_async().await;
4056
4057 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
4058 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
4059 .expect("refreshed token persisted under the keyring session's fingerprint");
4060 assert_eq!(persisted.access_token, "NEW-AT");
4061
4062 run(&cache, remote(&blob, &eps), Duration::ZERO)
4064 .await
4065 .expect("persisted token reused");
4066 token.assert_async().await;
4067 }
4068
4069 #[tokio::test]
4072 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
4073 let mut server = mockito::Server::new_async().await;
4074 let eps = endpoints(&server);
4075 let rejected = server
4076 .mock("POST", "/daily/quota")
4077 .match_header("authorization", "Bearer KEYRING-AT")
4078 .with_status(401)
4079 .expect(1)
4080 .create_async()
4081 .await;
4082 let accepted = quota_mock(&mut server, "NEW-AT")
4083 .expect(1)
4084 .create_async()
4085 .await;
4086 let token = token_mock(&mut server).expect(1).create_async().await;
4087 let blob = keyring_blob(VALID, true);
4088 let (_td, cache) = fixture();
4089
4090 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
4091 .await
4092 .expect("retry with the refreshed token succeeds");
4093
4094 rejected.assert_async().await;
4095 accepted.assert_async().await;
4096 token.assert_async().await;
4097 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4098 }
4099
4100 #[tokio::test]
4103 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
4104 let mut server = mockito::Server::new_async().await;
4105 let eps = endpoints(&server);
4106 let _token = server
4107 .mock("POST", "/token")
4108 .with_status(400)
4109 .with_body(r#"{"error":"invalid_grant"}"#)
4110 .create_async()
4111 .await;
4112 let blob = keyring_blob(EXPIRED, true);
4113 let (_td, cache) = fixture();
4114
4115 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4116 .await
4117 .expect_err("no cache to fall back on");
4118
4119 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4120 let rendered = err.to_string();
4121 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
4122 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
4123 }
4124
4125 #[tokio::test]
4128 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
4129 let mut server = mockito::Server::new_async().await;
4130 let eps = endpoints(&server);
4131 let _quota = server
4132 .mock("POST", "/daily/quota")
4133 .with_status(401)
4134 .create_async()
4135 .await;
4136 let token = token_mock(&mut server).expect(1).create_async().await;
4137 let blob = keyring_blob(EXPIRED, true);
4138 let (_td, cache) = fixture();
4139
4140 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4141 .await
4142 .expect_err("rejected after refresh");
4143
4144 token.assert_async().await;
4145 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4146 assert!(err.to_string().contains("sign in again"), "{err}");
4147 }
4148
4149 #[cfg(unix)]
4153 #[tokio::test]
4154 async fn a_subscription_403_surfaces_the_plan_message_without_a_refresh_or_agy_run() {
4155 let mut server = mockito::Server::new_async().await;
4156 let eps = endpoints(&server);
4157 let quota = server
4158 .mock("POST", "/daily/quota")
4159 .match_header("authorization", "Bearer KEYRING-AT")
4160 .with_status(403)
4161 .with_body(
4162 r#"{"error":{"code":403,"status":"PERMISSION_DENIED","details":
4163 [{"reason":"SUBSCRIPTION_REQUIRED","domain":"cloud.google.com"}]}}"#,
4164 )
4165 .expect(1)
4166 .create_async()
4167 .await;
4168 let prod = server
4169 .mock("POST", "/prod/quota")
4170 .expect(0)
4171 .create_async()
4172 .await;
4173 let token = token_mock(&mut server).expect(0).create_async().await;
4174 let (td, cache) = fixture();
4175 let file = td.path().join("credential");
4176 std::fs::write(&file, keyring_blob(VALID, true)).unwrap();
4177 let runs = td.path().join("runs");
4178 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
4179
4180 let err = run_without_oauth(
4181 &cache,
4182 RemoteOverride {
4183 credential: SavedCredential::File(&file),
4184 endpoints: Some(&eps),
4185 local_bases: Some(vec![]),
4186 agy: super::super::agy::AgyCommand::At(agy),
4187 ..Default::default()
4188 },
4189 )
4190 .await
4191 .expect_err("a plan without quota is an error, not a snapshot");
4192
4193 quota.assert_async().await;
4194 prod.assert_async().await;
4195 token.assert_async().await;
4196 match err {
4197 AppError::Http { status, body } => {
4198 assert_eq!(status, 403);
4199 assert!(body.contains("paid subscription"), "{body}");
4200 }
4201 other => panic!("expected Http, got {other:?}"),
4202 }
4203 assert!(
4204 !runs.exists(),
4205 "a subscription verdict must not spawn `agy`"
4206 );
4207 }
4208
4209 #[tokio::test]
4213 async fn a_non_subscription_403_keeps_the_refresh_and_retry_behavior() {
4214 let mut server = mockito::Server::new_async().await;
4215 let eps = endpoints(&server);
4216 let rejected = server
4217 .mock("POST", "/daily/quota")
4218 .match_header("authorization", "Bearer KEYRING-AT")
4219 .with_status(403)
4220 .with_body(r#"{"error":{"message":"acl detail"}}"#)
4221 .expect(1)
4222 .create_async()
4223 .await;
4224 let retried = server
4225 .mock("POST", "/daily/quota")
4226 .match_header("authorization", "Bearer NEW-AT")
4227 .with_status(403)
4228 .with_body(r#"{"error":{"message":"acl detail"}}"#)
4229 .expect(1)
4230 .create_async()
4231 .await;
4232 let token = token_mock(&mut server).expect(1).create_async().await;
4233 let blob = keyring_blob(VALID, true);
4234 let (_td, cache) = fixture();
4235
4236 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4237 .await
4238 .expect_err("still a rejected session");
4239
4240 rejected.assert_async().await;
4241 retried.assert_async().await;
4242 token.assert_async().await;
4243 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4244 assert!(err.to_string().contains("sign in again"), "{err}");
4245 }
4246
4247 async fn run_without_oauth(cache: &Cache, remote: RemoteOverride<'_>) -> Result<FetchOutcome> {
4249 fetch_snapshot_at(
4250 &reqwest::Client::new(),
4251 cache,
4252 Duration::ZERO,
4253 None,
4254 remote,
4255 now(),
4256 )
4257 .await
4258 }
4259
4260 #[cfg(unix)]
4265 fn fake_agy(
4266 dir: &std::path::Path,
4267 runs: &std::path::Path,
4268 file: &std::path::Path,
4269 blob: &str,
4270 ) -> std::path::PathBuf {
4271 let path = dir.join("agy");
4272 super::super::agy::test_support::write_script(
4273 &path,
4274 &format!(
4275 "[ \"$1\" = -p ] && exit 1\necho \"$@\" >> '{}'\nprintf '%s' '{}' > '{}'",
4276 runs.display(),
4277 blob,
4278 file.display()
4279 ),
4280 );
4281 path
4282 }
4283
4284 #[cfg(unix)]
4287 #[tokio::test]
4288 async fn an_expired_session_is_renewed_by_running_agy() {
4289 let mut server = mockito::Server::new_async().await;
4290 let eps = endpoints(&server);
4291 let quota = quota_mock(&mut server, "KEYRING-AT").create_async().await;
4292 let _plan = server
4293 .mock("POST", "/daily/plan")
4294 .with_status(200)
4295 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
4296 .create_async()
4297 .await;
4298 let (td, cache) = fixture();
4299 let file = td.path().join("credential");
4300 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
4301 let runs = td.path().join("runs");
4302 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
4303
4304 let out = run_without_oauth(
4305 &cache,
4306 RemoteOverride {
4307 credential: SavedCredential::File(&file),
4308 endpoints: Some(&eps),
4309 local_bases: Some(vec![]),
4310 agy: super::super::agy::AgyCommand::At(agy),
4311 ..Default::default()
4312 },
4313 )
4314 .await
4315 .expect("agy renewed the session");
4316
4317 quota.assert_async().await;
4318 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4319 assert_eq!(std::fs::read_to_string(runs).unwrap().trim(), "models");
4320 }
4321
4322 #[cfg(unix)]
4325 #[tokio::test]
4326 async fn agy_that_does_not_renew_keeps_the_original_error_and_runs_once() {
4327 let server = mockito::Server::new_async().await;
4328 let eps = endpoints(&server);
4329 let (td, cache) = fixture();
4330 let file = td.path().join("credential");
4331 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
4332 let runs = td.path().join("runs");
4333 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(EXPIRED, true));
4334 let remote = || RemoteOverride {
4335 credential: SavedCredential::File(&file),
4336 endpoints: Some(&eps),
4337 local_bases: Some(vec![]),
4338 agy: super::super::agy::AgyCommand::At(agy.clone()),
4339 ..Default::default()
4340 };
4341
4342 for _ in 0..2 {
4343 let err = run_without_oauth(&cache, remote())
4344 .await
4345 .expect_err("still expired");
4346 assert!(err.to_string().contains("no OAuth client"), "{err}");
4347 }
4348 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4349 }
4350
4351 #[tokio::test]
4354 async fn without_an_agy_command_an_expired_session_is_left_alone() {
4355 let server = mockito::Server::new_async().await;
4356 let eps = endpoints(&server);
4357 let blob = keyring_blob(EXPIRED, true);
4358 let (_td, cache) = fixture();
4359
4360 let err = run_without_oauth(&cache, remote(&blob, &eps))
4361 .await
4362 .expect_err("nothing renews it");
4363 assert!(err.to_string().contains("no OAuth client"), "{err}");
4364 }
4365
4366 #[cfg(unix)]
4369 const AGY_QUOTA: &str = r#"{"status":"SUCCESS","num_turns":0,"command":{"name":"usage","data":{"groups":[
4370{"name":"Gemini Models","buckets":[
4371 {"id":"gemini-weekly","name":"Weekly Limit Remaining","window":"weekly","remaining_fraction":0.63,"reset_time":"2026-07-28T01:11:19Z"},
4372 {"id":"gemini-5h","name":"Five Hour Limit Remaining","window":"5h","remaining_fraction":0.49,"reset_time":"2026-07-22T15:25:45Z"}]},
4373{"name":"Claude and GPT models","buckets":[
4374 {"id":"3p-weekly","name":"Weekly Limit Remaining","window":"weekly","remaining_fraction":1,"reset_time":"2026-07-29T12:00:00Z"},
4375 {"id":"3p-5h","name":"Five Hour Limit Remaining","window":"5h","remaining_fraction":1,"reset_time":"2026-07-22T17:00:00Z"}]},
4376{"name":"Gemini 4 Argon","buckets":[
4377 {"id":"gemini-argon-wtu-weekly","name":"Gemini 4 Argon Weekly Limit","window":"weekly","remaining_fraction":0,"reset_time":"2026-07-28T01:10:47Z"}]}]}}}"#;
4378
4379 #[cfg(unix)]
4382 fn agy_reading(dir: &std::path::Path, runs: &std::path::Path) -> std::path::PathBuf {
4383 let path = dir.join("agy");
4384 super::super::agy::test_support::write_script(
4385 &path,
4386 &format!(
4387 "echo \"$@\" >> '{}'\n[ \"$1\" = -p ] || exit 1\ncat <<'EOF'\n{AGY_QUOTA}\nEOF",
4388 runs.display()
4389 ),
4390 );
4391 path
4392 }
4393
4394 #[cfg(unix)]
4398 #[tokio::test]
4399 async fn agys_own_reading_answers_before_the_cloud_code_api() {
4400 let mut server = mockito::Server::new_async().await;
4401 let eps = endpoints(&server);
4402 let daily = server
4403 .mock("POST", "/daily/quota")
4404 .expect(0)
4405 .create_async()
4406 .await;
4407 let prod = server
4408 .mock("POST", "/prod/quota")
4409 .expect(0)
4410 .create_async()
4411 .await;
4412 let plan = server
4413 .mock("POST", "/daily/plan")
4414 .match_header("authorization", "Bearer KEYRING-AT")
4415 .with_status(200)
4416 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4417 .expect(1)
4418 .create_async()
4419 .await;
4420 let blob = keyring_blob(VALID, true);
4421 let (td, cache) = fixture();
4422 let runs = td.path().join("runs");
4423 let agy = agy_reading(td.path(), &runs);
4424
4425 let out = run(
4426 &cache,
4427 RemoteOverride {
4428 agy: super::super::agy::AgyCommand::At(agy),
4429 ..remote(&blob, &eps)
4430 },
4431 Duration::from_secs(60),
4432 )
4433 .await
4434 .expect("agy answered");
4435
4436 daily.assert_async().await;
4437 prod.assert_async().await;
4438 plan.assert_async().await;
4439 let snap = out.snapshot;
4440 assert_eq!(snap.source, AntigravitySource::Agy);
4441 assert_eq!(snap.plan, "Pro");
4442 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 51);
4443 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 37);
4444 assert_eq!(
4445 snap.third_party_session.as_ref().unwrap().utilization_pct,
4446 0
4447 );
4448 assert_eq!(snap.model_pools[0].window.utilization_pct, 100);
4449 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
4450 assert_eq!(snap.account, format!("acct:{fingerprint}"));
4451 let cached = parse_cache_at(
4452 &cache
4453 .fresh_payload(Duration::from_secs(60))
4454 .unwrap()
4455 .unwrap(),
4456 None,
4457 now(),
4458 )
4459 .unwrap();
4460 assert_eq!(cached.source, AntigravitySource::Agy);
4461 assert_eq!(
4462 std::fs::read_to_string(runs).unwrap().trim(),
4463 "-p /quota --output-format json --log-file /dev/null"
4464 );
4465 }
4466
4467 #[cfg(unix)]
4470 #[tokio::test]
4471 async fn agys_reading_needs_no_good_token() {
4472 let mut server = mockito::Server::new_async().await;
4473 let eps = endpoints(&server);
4474 let plan = server
4475 .mock("POST", "/daily/plan")
4476 .expect(0)
4477 .create_async()
4478 .await;
4479 let token = token_mock(&mut server).expect(0).create_async().await;
4480 let blob = keyring_blob(EXPIRED, true);
4481 let (td, cache) = fixture();
4482 let runs = td.path().join("runs");
4483 let agy = agy_reading(td.path(), &runs);
4484
4485 let out = run_without_oauth(
4486 &cache,
4487 RemoteOverride {
4488 agy: super::super::agy::AgyCommand::At(agy),
4489 ..remote(&blob, &eps)
4490 },
4491 )
4492 .await
4493 .expect("agy answered");
4494
4495 plan.assert_async().await;
4496 token.assert_async().await;
4497 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4498 assert_eq!(out.snapshot.plan, "Antigravity");
4499 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4500 }
4501
4502 #[cfg(unix)]
4505 #[tokio::test]
4506 async fn when_agy_cannot_read_the_quota_the_cloud_code_api_answers() {
4507 let mut server = mockito::Server::new_async().await;
4508 let eps = endpoints(&server);
4509 let quota = quota_mock(&mut server, "KEYRING-AT")
4510 .expect(2)
4511 .create_async()
4512 .await;
4513 let blob = keyring_blob(VALID, true);
4514 let (td, cache) = fixture();
4515 let runs = td.path().join("runs");
4516 let agy = td.path().join("agy");
4517 super::super::agy::test_support::write_script(
4518 &agy,
4519 &format!("echo \"$@\" >> '{}'\nexit 1", runs.display()),
4520 );
4521 let remote = || RemoteOverride {
4522 agy: super::super::agy::AgyCommand::At(agy.clone()),
4523 ..remote(&blob, &eps)
4524 };
4525
4526 for _ in 0..2 {
4527 let out = run(&cache, remote(), Duration::ZERO)
4528 .await
4529 .expect("the Cloud Code API answered");
4530 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4531 std::fs::remove_file(cache.payload_path()).unwrap();
4533 }
4534 quota.assert_async().await;
4535 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4536 }
4537
4538 #[cfg(unix)]
4542 fn age_past_reading_floor(cache: &Cache) {
4543 let aged = std::time::SystemTime::now() - super::super::agy::QUOTA_INTERVAL * 2;
4544 std::fs::OpenOptions::new()
4545 .write(true)
4546 .open(cache.payload_path())
4547 .unwrap()
4548 .set_times(std::fs::FileTimes::new().set_modified(aged))
4549 .unwrap();
4550 }
4551
4552 #[cfg(unix)]
4557 #[tokio::test]
4558 async fn when_agy_stops_answering_its_last_reading_is_kept_over_the_cloud_code_api() {
4559 let mut server = mockito::Server::new_async().await;
4560 let eps = endpoints(&server);
4561 let quota = quota_mock(&mut server, "KEYRING-AT")
4562 .expect(0)
4563 .create_async()
4564 .await;
4565 let _plan = server
4566 .mock("POST", "/daily/plan")
4567 .with_status(200)
4568 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4569 .create_async()
4570 .await;
4571 let blob = keyring_blob(VALID, true);
4572 let (td, cache) = fixture();
4573 let runs = td.path().join("runs");
4574 let agy = agy_reading(td.path(), &runs);
4575 let remote = || RemoteOverride {
4576 agy: super::super::agy::AgyCommand::At(agy.clone()),
4577 ..remote(&blob, &eps)
4578 };
4579
4580 let first = run(&cache, remote(), Duration::ZERO)
4581 .await
4582 .expect("agy answered");
4583 assert_eq!(first.snapshot.source, AntigravitySource::Agy);
4584 assert!(!first.stale);
4585
4586 age_past_reading_floor(&cache);
4588 super::super::agy::test_support::write_script(
4589 &agy,
4590 &format!("echo \"$@\" >> '{}'\nexit 1", runs.display()),
4591 );
4592
4593 for _ in 0..2 {
4594 let out = run(&cache, remote(), Duration::ZERO)
4595 .await
4596 .expect("agy's reading is kept");
4597 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4598 assert_eq!(out.snapshot.weekly.as_ref().unwrap().utilization_pct, 37);
4599 assert_eq!(out.snapshot.plan, "Pro");
4600 assert!(out.stale, "a kept reading is not passed off as current");
4601 age_past_reading_floor(&cache);
4602 }
4603 quota.assert_async().await;
4604 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 2);
4607 }
4608
4609 #[cfg(unix)]
4612 #[tokio::test]
4613 async fn a_cloud_code_figure_is_not_kept_in_agys_place() {
4614 let mut server = mockito::Server::new_async().await;
4615 let eps = endpoints(&server);
4616 let quota = quota_mock(&mut server, "KEYRING-AT")
4617 .expect(2)
4618 .create_async()
4619 .await;
4620 let blob = keyring_blob(VALID, true);
4621 let (td, cache) = fixture();
4622 let agy = td.path().join("agy");
4623 super::super::agy::test_support::write_script(&agy, "exit 1");
4624 let remote = || RemoteOverride {
4625 agy: super::super::agy::AgyCommand::At(agy.clone()),
4626 ..remote(&blob, &eps)
4627 };
4628
4629 for _ in 0..2 {
4630 let out = run(&cache, remote(), Duration::ZERO)
4631 .await
4632 .expect("the Cloud Code API answered");
4633 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4634 assert!(!out.stale);
4635 age_past_reading_floor(&cache);
4636 }
4637 quota.assert_async().await;
4638 }
4639
4640 #[cfg(unix)]
4643 #[tokio::test]
4644 async fn agys_reading_is_taken_at_most_every_quota_interval() {
4645 let server = mockito::Server::new_async().await;
4646 let eps = endpoints(&server);
4647 let blob = keyring_blob(VALID, true);
4648 let (td, cache) = fixture();
4649 let runs = td.path().join("runs");
4650 let agy = agy_reading(td.path(), &runs);
4651 for _ in 0..2 {
4652 let out = run(
4653 &cache,
4654 RemoteOverride {
4655 agy: super::super::agy::AgyCommand::At(agy.clone()),
4656 ..remote(&blob, &eps)
4657 },
4658 Duration::ZERO,
4659 )
4660 .await
4661 .expect("a snapshot");
4662 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4663 }
4664 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4665 }
4666
4667 #[tokio::test]
4669 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
4670 let mut server = mockito::Server::new_async().await;
4671 let eps = endpoints(&server);
4672 let token = token_mock(&mut server).expect(0).create_async().await;
4673 let blob = keyring_blob(EXPIRED, false);
4674 let (_td, cache) = fixture();
4675
4676 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4677 .await
4678 .expect_err("nothing to refresh with");
4679
4680 token.assert_async().await;
4681 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4682 assert!(err.to_string().contains("sign in again"), "{err}");
4683 }
4684
4685 #[tokio::test]
4688 async fn no_saved_session_extends_the_no_local_server_error() {
4689 let mut server = mockito::Server::new_async().await;
4690 let eps = endpoints(&server);
4691 let quota = quota_mock(&mut server, "KEYRING-AT")
4692 .expect(0)
4693 .create_async()
4694 .await;
4695 let (_td, cache) = fixture();
4696
4697 let err = run(
4698 &cache,
4699 RemoteOverride {
4700 credential: SavedCredential::Absent,
4701 endpoints: Some(&eps),
4702 local_bases: Some(vec![]),
4703 ..Default::default()
4704 },
4705 Duration::ZERO,
4706 )
4707 .await
4708 .expect_err("no source at all");
4709
4710 quota.assert_async().await;
4711 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4712 let rendered = err.to_string();
4713 assert!(rendered.contains("no local server found"), "{rendered}");
4714 assert!(rendered.contains("saved Google session"), "{rendered}");
4715 }
4716
4717 #[tokio::test]
4720 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
4721 let mut server = mockito::Server::new_async().await;
4722 let eps = endpoints(&server);
4723 let status_path = format!("/{STATUS_RPC}");
4724 let _status = server
4725 .mock("POST", status_path.as_str())
4726 .with_status(401)
4727 .create_async()
4728 .await;
4729 let quota = quota_mock(&mut server, "KEYRING-AT")
4730 .expect(0)
4731 .create_async()
4732 .await;
4733 let blob = keyring_blob(VALID, true);
4734 let (_td, cache) = fixture();
4735
4736 let err = run(
4737 &cache,
4738 RemoteOverride {
4739 credential: SavedCredential::Blob(&blob),
4740 endpoints: Some(&eps),
4741 local_bases: Some(vec![server.url().into()]),
4742 ..Default::default()
4743 },
4744 Duration::ZERO,
4745 )
4746 .await
4747 .expect_err("signed-out local server");
4748
4749 quota.assert_async().await;
4750 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
4751 }
4752
4753 #[tokio::test]
4756 async fn a_403_auth_rejection_does_not_fall_back_to_a_warm_cache() {
4757 let mut server = mockito::Server::new_async().await;
4758 let eps = endpoints(&server);
4759 let status_path = format!("/{STATUS_RPC}");
4760 let _status = server
4761 .mock("POST", status_path.as_str())
4762 .with_status(403)
4763 .with_body(r#"{"error":"forbidden"}"#)
4764 .create_async()
4765 .await;
4766 let quota = quota_mock(&mut server, "KEYRING-AT")
4767 .expect(0)
4768 .create_async()
4769 .await;
4770 let blob = keyring_blob(VALID, true);
4771 let (_td, cache) = fixture();
4772
4773 let mut warm_snap = parsed();
4774 warm_snap.session = None;
4775 if let Some(w) = warm_snap.weekly.as_mut() {
4776 w.utilization_pct = 100;
4777 }
4778 if let Some(w) = warm_snap.third_party_weekly.as_mut() {
4779 w.utilization_pct = 100;
4780 }
4781 let bytes = serde_json::to_vec(&snap_to_json(&warm_snap)).unwrap();
4782 cache.write_payload(&bytes).unwrap();
4783
4784 let err = run(
4785 &cache,
4786 RemoteOverride {
4787 credential: SavedCredential::Blob(&blob),
4788 endpoints: Some(&eps),
4789 local_bases: Some(vec![server.url().into()]),
4790 ..Default::default()
4791 },
4792 Duration::ZERO,
4793 )
4794 .await
4795 .expect_err("HTTP 403 must not fall back to warm cache");
4796
4797 quota.assert_async().await;
4798 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
4799 assert_eq!(cache.read_last_error().map(|(code, _)| code), Some(403));
4800 }
4801
4802 #[tokio::test]
4808 async fn a_signed_out_ide_language_server_is_skipped_so_the_saved_session_answers() {
4809 let mut ide = mockito::Server::new_async().await;
4810 let mut agy = mockito::Server::new_async().await;
4811 let eps = endpoints(&agy);
4812 let ide_root = ide.mock("GET", "/").expect(0).create_async().await;
4814 let status_path = format!("/{STATUS_RPC}");
4815 let ide_status = ide
4816 .mock("POST", status_path.as_str())
4817 .with_status(403)
4818 .expect(0)
4819 .create_async()
4820 .await;
4821 let agy_root = agy
4823 .mock("GET", "/")
4824 .with_status(404)
4825 .expect(1)
4826 .create_async()
4827 .await;
4828 let agy_status = agy
4829 .mock("POST", status_path.as_str())
4830 .with_status(401)
4831 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
4832 .expect(1)
4833 .create_async()
4834 .await;
4835 let quota = quota_mock(&mut agy, "KEYRING-AT")
4836 .expect(1)
4837 .create_async()
4838 .await;
4839 let plan = agy
4840 .mock("POST", "/daily/plan")
4841 .match_header("authorization", "Bearer KEYRING-AT")
4842 .with_status(200)
4843 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4844 .expect(1)
4845 .create_async()
4846 .await;
4847 let blob = keyring_blob(VALID, true);
4848 let (_td, cache) = fixture();
4849
4850 let outcome = run(
4851 &cache,
4852 RemoteOverride {
4853 credential: SavedCredential::Blob(&blob),
4854 endpoints: Some(&eps),
4855 local_bases: Some(vec![
4856 Candidate {
4857 base: ide.url(),
4858 pid: Some(7),
4859 ide: true,
4860 },
4861 Candidate {
4862 base: agy.url(),
4863 pid: Some(42),
4864 ide: false,
4865 },
4866 ]),
4867 ide_language_server: false,
4868 ..Default::default()
4869 },
4870 Duration::ZERO,
4871 )
4872 .await
4873 .expect("with the IDE skipped, agy's missing-CSRF answer unlocks the saved session");
4874
4875 ide_root.assert_async().await;
4876 ide_status.assert_async().await;
4877 agy_root.assert_async().await;
4878 agy_status.assert_async().await;
4879 quota.assert_async().await;
4880 plan.assert_async().await;
4881 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4882 assert_eq!(outcome.snapshot.plan, "Pro");
4883 }
4884
4885 #[tokio::test]
4889 async fn skipping_the_only_ide_candidate_falls_through_to_the_saved_session() {
4890 let mut ide = mockito::Server::new_async().await;
4891 let eps = endpoints(&ide);
4892 let ide_root = ide.mock("GET", "/").expect(0).create_async().await;
4893 let status_path = format!("/{STATUS_RPC}");
4894 let ide_status = ide
4895 .mock("POST", status_path.as_str())
4896 .expect(0)
4897 .create_async()
4898 .await;
4899 let quota = quota_mock(&mut ide, "KEYRING-AT")
4900 .expect(1)
4901 .create_async()
4902 .await;
4903 let plan = ide
4904 .mock("POST", "/daily/plan")
4905 .with_status(200)
4906 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
4907 .expect(1)
4908 .create_async()
4909 .await;
4910 let blob = keyring_blob(VALID, true);
4911 let (_td, cache) = fixture();
4912
4913 let outcome = run(
4914 &cache,
4915 RemoteOverride {
4916 credential: SavedCredential::Blob(&blob),
4917 endpoints: Some(&eps),
4918 local_bases: Some(vec![Candidate {
4919 base: ide.url(),
4920 pid: Some(7),
4921 ide: true,
4922 }]),
4923 ide_language_server: false,
4924 ..Default::default()
4925 },
4926 Duration::ZERO,
4927 )
4928 .await
4929 .expect("an IDE-only machine with the IDE skipped uses the saved session");
4930
4931 ide_root.assert_async().await;
4932 ide_status.assert_async().await;
4933 quota.assert_async().await;
4934 plan.assert_async().await;
4935 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4936 }
4937
4938 #[tokio::test]
4941 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
4942 let mut server = mockito::Server::new_async().await;
4943 let eps = endpoints(&server);
4944 for path in ["/daily/quota", "/prod/quota"] {
4945 server
4946 .mock("POST", path)
4947 .with_status(500)
4948 .with_body("boom")
4949 .create_async()
4950 .await;
4951 }
4952 let blob = keyring_blob(VALID, true);
4953 let (_td, cache) = fixture();
4954 cache.ensure_dir().unwrap();
4955 let mut earlier = parsed();
4956 earlier.source = AntigravitySource::Remote;
4957 earlier.account = "acct:earlier".into();
4958 cache
4959 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
4960 .unwrap();
4961
4962 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
4963 .await
4964 .expect("stale cache stands in");
4965
4966 assert!(outcome.stale);
4967 assert_eq!(outcome.snapshot, earlier);
4968 assert!(
4969 matches!(outcome.last_error, Some((500, _))),
4970 "{:?}",
4971 outcome.last_error
4972 );
4973 }
4974
4975 fn argon(pct: i32, reset: &str) -> ModelPoolWindow {
4976 ModelPoolWindow {
4977 name: "Gemini 4 Argon".into(),
4978 window: UsageWindow {
4979 utilization_pct: pct,
4980 resets_at: Some(
4981 DateTime::parse_from_rfc3339(reset)
4982 .unwrap()
4983 .with_timezone(&Utc),
4984 ),
4985 window_duration: chrono::Duration::days(7),
4986 },
4987 }
4988 }
4989
4990 #[test]
4991 fn model_pools_round_trip_through_the_cache() {
4992 let mut snap = parsed();
4993 snap.model_pools.push(argon(60, "2026-07-28T17:00:00Z"));
4994 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4995 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
4996
4997 let mut legacy = snap_to_json(&parsed());
4999 legacy.as_object_mut().unwrap().remove("model_pools");
5000 let legacy = serde_json::to_vec(&legacy).unwrap();
5001 assert!(
5002 parse_cache_at(&legacy, None, now())
5003 .unwrap()
5004 .model_pools
5005 .is_empty()
5006 );
5007
5008 let mut torn = snap_to_json(&snap);
5010 torn["model_pools"][0]["window"] = "monthly".into();
5011 let torn = serde_json::to_vec(&torn).unwrap();
5012 assert!(parse_cache_at(&torn, None, now()).is_err());
5013
5014 let mut damaged = snap_to_json(&snap);
5017 damaged["model_pools"][0]["name"] = "Argon\u{001b}[31m".into();
5018 let damaged = serde_json::to_vec(&damaged).unwrap();
5019 let restored = parse_cache_at(&damaged, None, now()).unwrap();
5020 assert_eq!(restored.model_pools[0].name, "Argon[31m");
5021 }
5022
5023 #[test]
5025 fn a_cached_model_pool_past_its_reset_refetches() {
5026 let mut snap = parsed();
5027 snap.model_pools.push(argon(60, "2026-07-22T11:00:00Z"));
5028 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
5029 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
5030 assert!(err.to_string().contains("Gemini 4 Argon weekly"), "{err}");
5031 }
5032
5033 #[test]
5034 fn source_round_trips_through_the_cache_and_defaults_to_local() {
5035 let mut snap = parsed();
5036 snap.source = AntigravitySource::Remote;
5037 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
5038 assert_eq!(
5039 parse_cache_at(&bytes, None, now()).unwrap().source,
5040 AntigravitySource::Remote
5041 );
5042
5043 snap.source = AntigravitySource::Statusline;
5044 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
5045 assert_eq!(
5046 parse_cache_at(&bytes, None, now()).unwrap().source,
5047 AntigravitySource::Statusline
5048 );
5049
5050 snap.source = AntigravitySource::Agy;
5051 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
5052 assert_eq!(
5053 parse_cache_at(&bytes, None, now()).unwrap().source,
5054 AntigravitySource::Agy
5055 );
5056
5057 let mut legacy = snap_to_json(&snap);
5059 legacy.as_object_mut().unwrap().remove("source");
5060 let legacy = serde_json::to_vec(&legacy).unwrap();
5061 assert_eq!(
5062 parse_cache_at(&legacy, None, now()).unwrap().source,
5063 AntigravitySource::Local
5064 );
5065 }
5066}