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 Err(e) => Err(e),
268 }
269 }
270 };
271
272 match live {
273 Ok(snap) => {
274 let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
275 cache.write_payload(&bytes)?;
276 Ok(crate::outcome::Outcome::fresh(snap))
277 }
278 Err(e) if e.is_transient() => fallback_silent(cache, now, e),
279 Err(AppError::Http { status, body }) if matches!(status, 401 | 403) => {
280 cache.mark_stale();
281 let _ = cache.write_last_error(status, &body);
282 Err(AppError::Http { status, body })
283 }
284 Err(AppError::Http { status, body }) => {
285 cache.mark_stale();
286 let last_error = Some(cache.write_last_error(status, &body));
287 let reason = AppError::Http { status, body };
288 fallback_with_error(cache, last_error, reason, now)
289 }
290 Err(e) => {
291 cache.mark_stale();
292 let last_error = Some(cache.write_last_error(0, &e.to_string()));
293 fallback_with_error(cache, last_error, e, now)
294 }
295 }
296}
297
298struct Session {
301 base: String,
302 csrf: Option<String>,
303 plan: String,
304 account: String,
305}
306
307enum Origin {
310 Local(Result<Session>),
311 Remote(Result<StoredToken>),
312}
313
314impl Origin {
315 fn account(&self) -> Option<String> {
317 match self {
318 Origin::Local(Ok(session)) => Some(session.account.clone()),
319 Origin::Remote(Ok(token)) => Some(remote_account(&token.fingerprint)),
320 Origin::Local(Err(_)) | Origin::Remote(Err(_)) => None,
321 }
322 }
323}
324
325fn no_local_server() -> AppError {
326 AppError::Credentials(NO_LOCAL_SERVER.into())
327}
328
329#[derive(Clone, Copy)]
330enum RemoteFallbackReason {
331 NoLocalServer,
332 AgyMissingCsrf,
333}
334
335impl RemoteFallbackReason {
336 fn message(self) -> &'static str {
337 match self {
338 Self::NoLocalServer => NO_LOCAL_SERVER,
339 Self::AgyMissingCsrf => AGY_CSRF_UNAVAILABLE,
340 }
341 }
342}
343
344fn remote_fallback_reason(error: &AppError) -> Option<RemoteFallbackReason> {
350 if matches!(error, AppError::Credentials(message) if message == NO_LOCAL_SERVER) {
351 Some(RemoteFallbackReason::NoLocalServer)
352 } else if is_missing_csrf(error) {
353 Some(RemoteFallbackReason::AgyMissingCsrf)
354 } else {
355 None
356 }
357}
358
359async fn open_session(
370 client: &reqwest::Client,
371 bases: Option<&[Candidate]>,
372 ide_language_server: bool,
373) -> Result<Session> {
374 let candidates = bases.map_or_else(
375 || candidate_bases(ide_language_server),
376 <[Candidate]>::to_vec,
377 );
378 let candidates: Vec<Candidate> = candidates.into_iter().filter(|c| !c.ide).collect();
379 if candidates.is_empty() {
380 return Err(no_local_server());
381 }
382
383 let mut errors = Vec::new();
384 let mut blocked_pids = std::collections::HashSet::new();
389 for candidate in candidates {
390 if candidate.pid.is_some_and(|pid| blocked_pids.contains(&pid)) {
391 continue;
392 }
393 let csrf = fetch_csrf(client, &candidate.base).await;
394 match post_rpc(client, &candidate.base, csrf.as_deref(), STATUS_RPC).await {
395 Ok(v) => {
396 return Ok(Session {
397 base: candidate.base,
398 csrf,
399 plan: plan_from_status(&v),
400 account: account_key(&v),
401 });
402 }
403 Err(e) => {
404 if is_missing_csrf(&e)
408 && let Some(pid) = candidate.pid
409 {
410 blocked_pids.insert(pid);
411 }
412 errors.push(e);
413 }
414 }
415 }
416 Err(select_probe_error(errors))
417}
418
419fn select_probe_error(errors: Vec<AppError>) -> AppError {
436 let mut actionable = None;
437 let mut missing_csrf = None;
438 let mut last = None;
439 let mut echo = None;
440 for e in errors {
441 if missing_csrf.is_none() && is_missing_csrf(&e) {
442 missing_csrf = Some(e);
443 } else if actionable.is_none() && is_actionable(&e) {
444 actionable = Some(e);
445 } else if is_tls_echo(&e) {
446 echo = Some(e);
447 } else {
448 last = Some(e);
449 }
450 }
451 actionable
452 .or(missing_csrf)
453 .or(last)
454 .or(echo)
455 .unwrap_or_else(|| {
456 AppError::Other("antigravity: no local server answered GetUserStatus".into())
457 })
458}
459
460fn is_missing_csrf(error: &AppError) -> bool {
466 let AppError::Http { status: 401, body } = error else {
467 return false;
468 };
469 let Ok(body) = serde_json::from_str::<serde_json::Value>(body) else {
470 return false;
471 };
472 matches!(
473 (body["code"].as_str(), body["message"].as_str()),
474 (Some(code), Some(message))
475 if code.eq_ignore_ascii_case("unauthenticated")
476 && is_missing_csrf_message(message)
477 )
478}
479
480fn is_missing_csrf_message(message: &str) -> bool {
483 message
484 .rsplit(':')
485 .next()
486 .is_some_and(|part| part.trim().eq_ignore_ascii_case("missing CSRF token"))
487}
488
489fn is_actionable(e: &AppError) -> bool {
494 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
495 && !is_missing_csrf(e)
496}
497
498fn is_tls_echo(e: &AppError) -> bool {
511 matches!(
512 e,
513 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
514 )
515}
516
517async fn fetch_live(
518 client: &reqwest::Client,
519 session: Result<Session>,
520) -> Result<AntigravitySnapshot> {
521 let session = session?;
522 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
523 let mut snap = parse_quota_summary("a, session.plan)?;
524 snap.account = session.account;
525 Ok(snap)
526}
527
528fn saved_session(
538 credential: SavedCredential<'_>,
539 reason: RemoteFallbackReason,
540) -> Result<StoredToken> {
541 let raw = match credential {
542 SavedCredential::Keyring => credential::read()?,
543 SavedCredential::Blob(blob) => Some(blob.to_string()),
544 SavedCredential::File(path) => std::fs::read_to_string(path).ok(),
545 SavedCredential::Absent => None,
546 };
547 let Some(raw) = raw else {
548 return Err(AppError::Credentials(format!(
549 "{} {NO_SAVED_SESSION}",
550 reason.message()
551 )));
552 };
553 credential::parse_keyring_blob(&raw)
554}
555
556fn remote_account(fingerprint: &str) -> String {
560 format!("acct:{fingerprint}")
561}
562
563fn session_expired() -> AppError {
564 AppError::Credentials(SESSION_EXPIRED.into())
565}
566
567fn refresh_unconfigured() -> AppError {
571 AppError::Credentials(
572 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
573 refresh it; open Antigravity (or run `agy`) to sign in again, or set [antigravity] oauth_client_id \
574 and oauth_client_secret in config.toml"
575 .into(),
576 )
577}
578
579fn is_auth_rejection(e: &AppError) -> bool {
580 matches!(
581 e,
582 AppError::Http {
583 status: 401 | 403,
584 ..
585 }
586 )
587}
588
589struct AccessToken {
593 value: String,
594 just_refreshed: bool,
595}
596
597fn current_access_token(
602 oauth_path: &std::path::Path,
603 token: &StoredToken,
604 now: DateTime<Utc>,
605) -> Option<String> {
606 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
607 && token
608 .expires_at
609 .is_none_or(|keyring| persisted.expires_at > keyring)
610 && !cloud::needs_refresh(Some(persisted.expires_at), now)
611 {
612 return Some(persisted.access_token);
613 }
614 (!cloud::needs_refresh(token.expires_at, now)).then(|| token.access_token.clone())
615}
616
617async fn resolve_access_token(
619 client: &reqwest::Client,
620 oauth: Option<&cloud::OauthClient>,
621 endpoints: &cloud::Endpoints,
622 oauth_path: &std::path::Path,
623 token: &StoredToken,
624 now: DateTime<Utc>,
625) -> Result<AccessToken> {
626 if let Some(value) = current_access_token(oauth_path, token, now) {
627 return Ok(AccessToken {
628 value,
629 just_refreshed: false,
630 });
631 }
632 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
633}
634
635async fn refresh_and_persist(
639 client: &reqwest::Client,
640 oauth: Option<&cloud::OauthClient>,
641 endpoints: &cloud::Endpoints,
642 oauth_path: &std::path::Path,
643 token: &StoredToken,
644) -> Result<AccessToken> {
645 let Some(refresh_token) = token.refresh_token.as_deref() else {
646 return Err(session_expired());
647 };
648 let Some(oauth) = oauth else {
649 return Err(refresh_unconfigured());
650 };
651 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
652 cloud::write_persisted(
653 oauth_path,
654 &cloud::PersistedOAuth {
655 fingerprint: token.fingerprint.clone(),
656 access_token: refreshed.access_token.clone(),
657 expires_at: refreshed.expires_at,
658 },
659 )?;
660 Ok(AccessToken {
661 value: refreshed.access_token,
662 just_refreshed: true,
663 })
664}
665
666struct Renewal<'a> {
669 agy: &'a super::agy::AgyCommand,
670 credential: SavedCredential<'a>,
671}
672
673impl Renewal<'_> {
674 async fn renewed(&self, cache: &Cache, now: DateTime<Utc>) -> Option<StoredToken> {
678 if !super::agy::renew_session(self.agy, cache, now).await {
679 return None;
680 }
681 let raw = match self.credential {
682 SavedCredential::Keyring => credential::read().ok().flatten()?,
683 SavedCredential::Blob(blob) => blob.to_string(),
684 SavedCredential::File(path) => std::fs::read_to_string(path).ok()?,
685 SavedCredential::Absent => return None,
686 };
687 credential::parse_keyring_blob(&raw).ok()
688 }
689}
690
691enum RemoteReading {
695 Live(AntigravitySnapshot),
696 Kept(FetchOutcome),
697}
698
699async fn fetch_remote(
725 client: &reqwest::Client,
726 cache: &Cache,
727 oauth: Option<&cloud::OauthClient>,
728 endpoints: &cloud::Endpoints,
729 token: Result<StoredToken>,
730 renewal: Renewal<'_>,
731 now: DateTime<Utc>,
732) -> Result<RemoteReading> {
733 let mut token = token?;
734 let oauth_path = cloud::oauth_cache_path(cache);
735 match super::agy::quota_snapshot(renewal.agy, cache, now).await {
736 Some(mut snap) => {
737 snap.plan = match current_access_token(&oauth_path, &token, now) {
738 Some(access) => cloud::fetch_plan(client, endpoints, &access).await,
739 None => None,
740 }
741 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
742 snap.account = remote_account(&token.fingerprint);
743 snap.source = AntigravitySource::Agy;
744 return Ok(RemoteReading::Live(snap));
745 }
746 None if renewal.agy.is_available() => {
747 if let Some(kept) = kept_agy_reading(cache, &remote_account(&token.fingerprint), now) {
748 return Ok(RemoteReading::Kept(kept));
749 }
750 }
751 None => {}
752 }
753 let mut access =
754 match resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await {
755 Ok(access) => access,
756 Err(e) if oauth.is_none() => {
759 let Some(renewed) = renewal.renewed(cache, now).await else {
760 return Err(e);
761 };
762 token = renewed;
763 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
764 }
765 Err(e) => return Err(e),
766 };
767
768 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
769 Err(e) if cloud::is_no_plan_quota(&e) => return Err(e),
772 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
773 access = if oauth.is_none() {
774 let Some(renewed) = renewal.renewed(cache, now).await else {
775 return Err(refresh_unconfigured());
776 };
777 token = renewed;
778 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
779 } else {
780 refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?
781 };
782 cloud::fetch_quota(client, endpoints, &access.value)
783 .await
784 .map_err(|e| {
785 if cloud::is_no_plan_quota(&e) {
786 e
787 } else if is_auth_rejection(&e) {
788 session_expired()
789 } else {
790 e
791 }
792 })?
793 }
794 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
795 other => other?,
796 };
797
798 let plan = cloud::fetch_plan(client, endpoints, &access.value)
799 .await
800 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
801 let mut snap = parse_quota_summary("a, plan)?;
802 snap.account = remote_account(&token.fingerprint);
803 snap.source = AntigravitySource::Remote;
804 Ok(RemoteReading::Live(snap))
805}
806
807fn kept_agy_reading(cache: &Cache, account: &str, now: DateTime<Utc>) -> Option<FetchOutcome> {
816 let not_agys = || AppError::Other("Antigravity: no reading of `agy`'s own to keep".into());
817 crate::outcome::fallback(cache, None, not_agys(), |bytes| {
818 let snap = parse_cache_at(bytes, Some(account), now)?;
819 (snap.source == AntigravitySource::Agy)
820 .then_some(snap)
821 .ok_or_else(not_agys)
822 })
823 .ok()
824}
825
826fn account_key(user_status: &serde_json::Value) -> String {
831 let email = user_status["userStatus"]["email"]
832 .as_str()
833 .filter(|s| !s.is_empty());
834 match email {
835 Some(e) => {
836 use std::hash::{Hash, Hasher};
837 let mut h = std::collections::hash_map::DefaultHasher::new();
838 e.hash(&mut h);
839 format!("acct:{:016x}", h.finish())
840 }
841 None => "acct:unknown".to_string(),
842 }
843}
844
845async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
849 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
850 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
853 .await
854 .ok()?;
855 let html = String::from_utf8_lossy(&bytes);
856 html.split("csrfToken\":\"")
857 .nth(1)
858 .and_then(|s| s.split('"').next())
859 .filter(|t| !t.is_empty())
860 .map(|t| t.to_string())
861}
862
863async fn post_rpc(
864 client: &reqwest::Client,
865 base: &str,
866 csrf: Option<&str>,
867 rpc: &str,
868) -> Result<serde_json::Value> {
869 let mut req = client
870 .post(format!("{base}/{rpc}"))
871 .header("Content-Type", "application/json")
872 .body("{}")
873 .timeout(HTTP_TIMEOUT);
874 if let Some(token) = csrf {
875 req = req.header("x-codeium-csrf-token", token);
876 }
877 let resp = req.send().await?;
878
879 let status = resp.status();
880 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
884 if !status.is_success() {
885 let body = String::from_utf8_lossy(&bytes).into_owned();
886 return Err(AppError::Http {
887 status: status.as_u16(),
888 body,
889 });
890 }
891
892 Ok(serde_json::from_slice(&bytes)?)
893}
894
895pub fn plan_from_status(v: &serde_json::Value) -> String {
896 v["userStatus"]["userTier"]["name"]
897 .as_str()
898 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
899 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
900 .filter(|s| !s.is_empty())
901 .unwrap_or(DEFAULT_PLAN)
902 .to_string()
903}
904
905fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
917 let id = if id.is_empty() { "<unnamed>" } else { id };
918 let group = if group_name.is_empty() {
919 String::new()
920 } else {
921 format!(" in {group_name:?}")
922 };
923 match window {
924 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
925 _ => format!("{id}{group}"),
926 }
927}
928
929fn no_usable_bucket(seen: &[String]) -> AppError {
935 if seen.is_empty() {
936 return AppError::Other(
937 "antigravity: quota summary has no buckets at all — the running product may \
938 not have a quota for this account yet"
939 .into(),
940 );
941 }
942 AppError::Other(format!(
943 "antigravity: quota summary has no bucket in a window we recognise (5h or \
944 weekly, Gemini or Claude/GPT); it offered: {}",
945 seen.join(", ")
946 ))
947}
948
949pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
950 let groups = v["response"]["groups"]
951 .as_array()
952 .or_else(|| v["groups"].as_array())
953 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
954
955 const SLOT_NAMES: [&str; 4] = [
957 "Gemini 5h",
958 "Gemini weekly",
959 "Claude/GPT 5h",
960 "Claude/GPT weekly",
961 ];
962 const CANONICAL_IDS: [&str; 4] = ["gemini-5h", "gemini-weekly", "3p-5h", "3p-weekly"];
963 let is_canonical = |bucket: &serde_json::Value| {
964 CANONICAL_IDS.contains(&bucket["bucketId"].as_str().unwrap_or_default())
965 };
966 let shares_a_slot = |group: &serde_json::Value| {
967 group["buckets"]
968 .as_array()
969 .is_some_and(|buckets| buckets.iter().any(is_canonical))
970 };
971 let renamed = !groups.iter().any(shares_a_slot);
978
979 let mut exact: [Vec<&serde_json::Value>; 4] = Default::default();
982 let mut inferred: [Vec<&serde_json::Value>; 4] = Default::default();
983 let mut pools: Vec<(&str, &serde_json::Value, bool)> = Vec::new();
984 let mut seen: Vec<String> = Vec::new();
988
989 for group in groups {
990 let group_name = group["displayName"].as_str().unwrap_or_default();
991 let Some(buckets) = group["buckets"].as_array() else {
992 continue;
993 };
994 let own_pool = !renamed && !shares_a_slot(group);
995 for bucket in buckets {
996 let id = bucket["bucketId"].as_str().unwrap_or_default();
997 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
998 if let Some(slot) = CANONICAL_IDS.iter().position(|c| *c == id) {
999 exact[slot].push(bucket);
1000 continue;
1001 }
1002 let window = bucket["window"].as_str().unwrap_or_default();
1003 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
1004 true
1005 } else if id.ends_with("5h") || window == "5h" {
1006 false
1007 } else {
1008 continue;
1011 };
1012 if own_pool {
1013 pools.push((group_name, bucket, is_weekly));
1014 continue;
1015 }
1016 if !renamed {
1019 continue;
1020 }
1021 let is_gemini = if id.starts_with("gemini") {
1022 true
1023 } else if id.starts_with("3p") {
1024 false
1025 } else if group_name.contains("Gemini") {
1026 true
1027 } else if group_name.contains("Claude") || group_name.contains("GPT") {
1028 false
1029 } else {
1030 pools.push((group_name, bucket, is_weekly));
1033 continue;
1034 };
1035
1036 let slot = usize::from(!is_gemini) * 2 + usize::from(is_weekly);
1037 inferred[slot].push(bucket);
1038 }
1039 }
1040
1041 let mut slots: [Option<UsageWindow>; 4] = Default::default();
1045 for slot in 0..4 {
1046 let candidates = if renamed {
1047 &inferred[slot]
1048 } else {
1049 &exact[slot]
1050 };
1051 match candidates.as_slice() {
1052 [] => {}
1053 [bucket] => slots[slot] = Some(usage_window(bucket, slot % 2 == 1)?),
1054 _ => {
1055 return Err(AppError::Schema(format!(
1056 "antigravity: duplicate {} bucket",
1057 SLOT_NAMES[slot]
1058 )));
1059 }
1060 }
1061 }
1062 let [gemini_5h, gemini_weekly, tp_5h, tp_weekly] = slots;
1063
1064 let model_pools: Vec<ModelPoolWindow> = pools
1068 .into_iter()
1069 .filter_map(|(group_name, bucket, is_weekly)| {
1070 let window = usage_window(bucket, is_weekly).ok()?;
1071 Some(ModelPoolWindow {
1072 name: pool_name(group_name, bucket),
1073 window,
1074 })
1075 })
1076 .collect();
1077
1078 if gemini_5h.is_none()
1082 && gemini_weekly.is_none()
1083 && tp_5h.is_none()
1084 && tp_weekly.is_none()
1085 && model_pools.is_empty()
1086 {
1087 return Err(no_usable_bucket(&seen));
1088 }
1089
1090 Ok(AntigravitySnapshot {
1091 plan,
1092 account: String::new(),
1095 source: AntigravitySource::Local,
1096 session: gemini_5h,
1097 weekly: gemini_weekly,
1098 third_party_session: tp_5h,
1099 third_party_weekly: tp_weekly,
1100 model_pools,
1101 })
1102}
1103
1104fn pool_name(group_name: &str, bucket: &serde_json::Value) -> String {
1108 [
1109 group_name,
1110 bucket["displayName"].as_str().unwrap_or_default(),
1111 bucket["bucketId"].as_str().unwrap_or_default(),
1112 ]
1113 .into_iter()
1114 .map(crate::display::sanitize_untrusted_field)
1115 .find(|name| !name.trim().is_empty())
1116 .unwrap_or_else(|| "Model pool".to_string())
1117}
1118
1119fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
1123 let remaining = bucket["remainingFraction"]
1124 .as_f64()
1125 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
1126 .ok_or_else(|| {
1127 AppError::Schema(format!(
1128 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
1129 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
1130 ))
1131 })?;
1132 Ok(UsageWindow {
1133 utilization_pct: pct_used(remaining),
1134 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
1135 window_duration: if is_weekly {
1136 chrono::Duration::days(7)
1137 } else {
1138 chrono::Duration::hours(5)
1139 },
1140 })
1141}
1142
1143fn pct_used(remaining_fraction: f64) -> i32 {
1146 let used = (1.0 - remaining_fraction) * 100.0;
1147 used.round() as i32
1148}
1149
1150fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
1151 match value {
1152 serde_json::Value::Null => Ok(None),
1153 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
1154 .map(|dt| Some(dt.with_timezone(&Utc)))
1155 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
1156 _ => Err(AppError::Schema(format!(
1157 "antigravity: {field} must be a timestamp or null"
1158 ))),
1159 }
1160}
1161
1162fn candidate_bases(ide_language_server: bool) -> Vec<Candidate> {
1170 let discovered = discover_ls_ports();
1171 let ide_pids = if ide_language_server {
1174 std::collections::HashSet::new()
1175 } else {
1176 ide_ls_pids(discovered.iter().map(|&(pid, _)| pid).collect())
1177 };
1178 candidate_bases_with(
1179 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
1180 discovered,
1181 &ide_pids,
1182 )
1183}
1184
1185fn candidate_bases_with(
1189 override_addr: Option<&str>,
1190 discovered: Vec<(u32, u16)>,
1191 ide_pids: &std::collections::HashSet<u32>,
1192) -> Vec<Candidate> {
1193 let mut candidates = Vec::new();
1194 if let Some(base) = override_addr.and_then(normalize_base) {
1195 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
1200 candidates.push(Candidate {
1201 base,
1202 pid,
1203 ide: false,
1204 });
1205 }
1206
1207 for (pid, port) in discovered {
1213 let base = format!("http://127.0.0.1:{port}");
1214 if !candidates.iter().any(|c| c.base == base) {
1215 candidates.push(Candidate {
1216 base,
1217 pid: Some(pid),
1218 ide: ide_pids.contains(&pid),
1219 });
1220 }
1221 }
1222
1223 candidates
1224}
1225
1226#[cfg(any(test, target_os = "linux", target_os = "macos"))]
1231fn is_ide_install(exe: &str) -> bool {
1232 let exe = exe.to_lowercase().replace('\\', "/");
1233 exe.contains("antigravity-ide") || exe.contains("antigravity ide")
1234}
1235
1236#[cfg(target_os = "linux")]
1238fn ide_ls_pids(pids: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1239 pids.into_iter()
1240 .filter(|pid| {
1241 std::fs::read_link(format!("/proc/{pid}/exe"))
1242 .map(|path| path.to_string_lossy().into_owned())
1243 .is_ok_and(|exe| is_ide_install(&exe))
1244 })
1245 .collect()
1246}
1247
1248#[cfg(target_os = "macos")]
1251fn ide_ls_pids(pids: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1252 if pids.is_empty() {
1253 return std::collections::HashSet::new();
1254 }
1255 let list = pids
1256 .iter()
1257 .map(u32::to_string)
1258 .collect::<Vec<_>>()
1259 .join(",");
1260 let Ok(output) = std::process::Command::new("/bin/ps")
1263 .args(["-o", "pid=,comm=", "-p", &list])
1264 .output()
1265 else {
1266 return std::collections::HashSet::new();
1267 };
1268 parse_ps_comm(&String::from_utf8_lossy(&output.stdout))
1269 .into_iter()
1270 .filter(|(pid, exe)| pids.contains(pid) && is_ide_install(exe))
1271 .map(|(pid, _)| pid)
1272 .collect()
1273}
1274
1275#[cfg(any(test, target_os = "macos"))]
1278fn parse_ps_comm(output: &str) -> Vec<(u32, String)> {
1279 output
1280 .lines()
1281 .filter_map(|line| {
1282 let (pid, exe) = line.trim().split_once(char::is_whitespace)?;
1283 Some((pid.parse().ok()?, exe.trim().to_string()))
1284 })
1285 .collect()
1286}
1287
1288#[cfg(target_os = "windows")]
1292fn ide_ls_pids(_: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1293 std::collections::HashSet::new()
1294}
1295
1296#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1297fn ide_ls_pids(_: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1298 std::collections::HashSet::new()
1299}
1300
1301fn base_port(base: &str) -> Option<u16> {
1303 let (_, authority) = base.split_once("://")?;
1304 authority.rsplit_once(':')?.1.parse().ok()
1305}
1306
1307fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
1310 let mut owners = discovered
1311 .iter()
1312 .filter(|(_, p)| *p == port)
1313 .map(|(pid, _)| *pid);
1314 let first = owners.next()?;
1315 owners.all(|pid| pid == first).then_some(first)
1316}
1317
1318fn normalize_base(addr: &str) -> Option<String> {
1327 let trimmed = addr.trim();
1328 let (scheme, authority) = match trimmed.split_once("://") {
1329 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
1330 _ => ("http", trimmed),
1331 };
1332 let authority = authority.trim_end_matches('/');
1333 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
1334}
1335
1336fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
1343 let comm = comm.trim().to_lowercase();
1344 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
1345 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
1346 return true;
1347 }
1348 exe.is_some_and(|p| {
1349 let p = p.to_lowercase().replace('\\', "/");
1350 let p = p.strip_suffix(".exe").unwrap_or(&p);
1351 p.contains("antigravity") || p.ends_with("/agy")
1352 })
1353}
1354
1355#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
1387fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
1388 let groups: Vec<(u32, Vec<u16>)> = per_pid
1389 .into_iter()
1390 .map(|(pid, mut group)| {
1391 group.sort_unstable_by(|a, b| b.cmp(a));
1392 group.dedup();
1397 (pid, group)
1398 })
1399 .collect();
1400
1401 let mut ports: Vec<(u32, u16)> = Vec::new();
1402 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
1403 for (pid, port) in groups
1404 .iter()
1405 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
1406 {
1407 if !ports.iter().any(|(_, p)| *p == port) {
1410 ports.push((pid, port));
1411 }
1412 }
1413 }
1414 ports
1415}
1416
1417#[cfg(target_os = "linux")]
1425pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1426 use std::collections::{BTreeMap, HashMap};
1427
1428 let mut owners: HashMap<u64, u32> = HashMap::new();
1431 let Ok(entries) = std::fs::read_dir("/proc") else {
1432 return Vec::new();
1433 };
1434 for entry in entries.flatten() {
1435 let pid_dir = entry.path();
1436 let Some(pid) = pid_dir
1437 .file_name()
1438 .and_then(|name| name.to_str())
1439 .and_then(|name| name.parse::<u32>().ok())
1440 else {
1441 continue;
1442 };
1443 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1444 continue;
1445 };
1446 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1447 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1448 continue;
1449 }
1450 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1451 continue;
1452 };
1453 for fd in fds.flatten() {
1454 let Ok(target) = std::fs::read_link(fd.path()) else {
1455 continue;
1456 };
1457 if let Some(ino) = target
1458 .to_str()
1459 .and_then(|s| s.strip_prefix("socket:["))
1460 .and_then(|s| s.strip_suffix(']'))
1461 .and_then(|s| s.parse::<u64>().ok())
1462 {
1463 owners.insert(ino, pid);
1464 }
1465 }
1466 }
1467
1468 if owners.is_empty() {
1469 return Vec::new();
1470 }
1471
1472 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1473 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1474 let Ok(contents) = std::fs::read_to_string(table) else {
1475 continue;
1476 };
1477 for line in contents.lines().skip(1) {
1478 if let Some((port, ino)) = parse_proc_net_line(line)
1479 && let Some(&pid) = owners.get(&ino)
1480 {
1481 per_pid.entry(pid).or_default().push(port);
1482 }
1483 }
1484 }
1485 probe_order(per_pid)
1486}
1487
1488#[cfg(target_os = "macos")]
1495pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1496 let Ok(output) = std::process::Command::new("/usr/sbin/lsof")
1498 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1499 .output()
1500 else {
1501 return Vec::new();
1502 };
1503 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1507}
1508
1509#[cfg(any(test, target_os = "macos"))]
1514fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1515 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1516 let mut pid = None;
1519 let mut owner = None;
1520 for line in output.lines() {
1521 let Some(rest) = line.get(1..) else { continue };
1522 match line.as_bytes().first() {
1523 Some(b'p') => {
1524 pid = rest.parse::<u32>().ok();
1525 owner = None;
1526 }
1527 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1528 Some(b'n') => {
1529 if let Some(pid) = owner
1530 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1531 {
1532 per_pid.entry(pid).or_default().push(port);
1533 }
1534 }
1535 _ => {}
1536 }
1537 }
1538 probe_order(per_pid)
1539}
1540
1541#[cfg(any(test, target_os = "windows"))]
1542#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1543struct WindowsTcpRow {
1544 local_addr: [u8; 4],
1545 local_port: u32,
1546 pid: u32,
1547}
1548
1549#[cfg(any(test, target_os = "windows"))]
1550fn decode_windows_process_name(raw: &[u16]) -> String {
1551 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1552 String::from_utf16_lossy(&raw[..end])
1553}
1554
1555#[cfg(any(test, target_os = "windows"))]
1556fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1557 processes
1558 .iter()
1559 .filter(|(_, name)| is_antigravity_process(name, None))
1560 .map(|(pid, _)| *pid)
1561 .collect()
1562}
1563
1564#[cfg(any(test, target_os = "windows"))]
1567fn matching_windows_ports(
1568 pids: &std::collections::HashSet<u32>,
1569 rows: &[WindowsTcpRow],
1570) -> Vec<(u32, u16)> {
1571 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1572 for row in rows {
1573 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1574 continue;
1575 }
1576 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1577 if port != 0 {
1578 per_pid.entry(row.pid).or_default().push(port);
1579 }
1580 }
1581 probe_order(per_pid)
1582}
1583
1584#[cfg(any(test, target_os = "windows"))]
1585fn checked_windows_row_count(
1586 buffer_len: usize,
1587 rows_offset: usize,
1588 row_size: usize,
1589 declared: usize,
1590) -> Option<usize> {
1591 let rows_len = row_size.checked_mul(declared)?;
1592 let end = rows_offset.checked_add(rows_len)?;
1593 (row_size != 0 && end <= buffer_len).then_some(declared)
1594}
1595
1596#[cfg(target_os = "windows")]
1597struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1598
1599#[cfg(target_os = "windows")]
1600impl Drop for WindowsHandle {
1601 fn drop(&mut self) {
1602 unsafe {
1603 windows_sys::Win32::Foundation::CloseHandle(self.0);
1604 }
1605 }
1606}
1607
1608#[cfg(target_os = "windows")]
1609fn windows_processes() -> Vec<(u32, String)> {
1610 use std::mem::size_of;
1611 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1612 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1613 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1614 TH32CS_SNAPPROCESS,
1615 };
1616
1617 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1618 if handle == INVALID_HANDLE_VALUE {
1619 return Vec::new();
1620 }
1621 let snapshot = WindowsHandle(handle);
1622 let mut entry = PROCESSENTRY32W::default();
1623 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1624 return Vec::new();
1625 };
1626 entry.dwSize = entry_size;
1627 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1628 return Vec::new();
1629 }
1630
1631 let mut processes = Vec::new();
1632 loop {
1633 processes.push((
1634 entry.th32ProcessID,
1635 decode_windows_process_name(&entry.szExeFile),
1636 ));
1637 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1638 break;
1639 }
1640 }
1641 processes
1642}
1643
1644#[cfg(target_os = "windows")]
1645fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1646 use std::mem::{offset_of, size_of, size_of_val};
1647 use windows_sys::Win32::NetworkManagement::IpHelper::{
1648 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1649 };
1650
1651 let available = used_bytes.min(size_of_val(buffer));
1652 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1653 if available < size_of::<u32>() || available < rows_offset {
1654 return Vec::new();
1655 }
1656 let base = buffer.as_ptr().cast::<u8>();
1657 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1658 if checked_windows_row_count(
1659 available,
1660 rows_offset,
1661 size_of::<MIB_TCPROW_OWNER_PID>(),
1662 declared,
1663 )
1664 .is_none()
1665 {
1666 return Vec::new();
1667 }
1668
1669 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1670 (0..declared)
1671 .map(|index| unsafe { rows.add(index).read_unaligned() })
1672 .map(|row| WindowsTcpRow {
1673 local_addr: row.dwLocalAddr.to_ne_bytes(),
1674 local_port: row.dwLocalPort,
1675 pid: row.dwOwningPid,
1676 })
1677 .collect()
1678}
1679
1680#[cfg(target_os = "windows")]
1681fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1682 use std::mem::size_of;
1683 use std::ptr::null_mut;
1684 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1685 use windows_sys::Win32::NetworkManagement::IpHelper::{
1686 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1687 };
1688 use windows_sys::Win32::Networking::WinSock::AF_INET;
1689
1690 let mut size = 0u32;
1691 let status = unsafe {
1692 GetExtendedTcpTable(
1693 null_mut(),
1694 &mut size,
1695 0,
1696 u32::from(AF_INET),
1697 TCP_TABLE_OWNER_PID_LISTENER,
1698 0,
1699 )
1700 };
1701 if status != ERROR_INSUFFICIENT_BUFFER {
1702 return Vec::new();
1703 }
1704
1705 for _ in 0..3 {
1706 let Some(words) = (size as usize)
1707 .checked_add(size_of::<u32>() - 1)
1708 .map(|bytes| bytes / size_of::<u32>())
1709 else {
1710 return Vec::new();
1711 };
1712 if words == 0 {
1713 return Vec::new();
1714 }
1715 let mut buffer = Vec::<u32>::new();
1716 if buffer.try_reserve_exact(words).is_err() {
1717 return Vec::new();
1718 }
1719 buffer.resize(words, 0);
1720 let mut used = size;
1721 let status = unsafe {
1722 GetExtendedTcpTable(
1723 buffer.as_mut_ptr().cast(),
1724 &mut used,
1725 0,
1726 u32::from(AF_INET),
1727 TCP_TABLE_OWNER_PID_LISTENER,
1728 0,
1729 )
1730 };
1731 if status == ERROR_INSUFFICIENT_BUFFER {
1732 size = used;
1733 continue;
1734 }
1735 if status != 0 {
1736 return Vec::new();
1737 }
1738 return parse_windows_tcp_rows(&buffer, used as usize);
1739 }
1740 Vec::new()
1741}
1742
1743#[cfg(target_os = "windows")]
1744pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1745 let pids = matching_windows_process_ids(&windows_processes());
1746 if pids.is_empty() {
1747 return Vec::new();
1748 }
1749 matching_windows_ports(&pids, &windows_tcp_rows())
1750}
1751
1752#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1753pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1754 Vec::new()
1755}
1756
1757#[cfg(target_os = "linux")]
1760fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1761 let cols: Vec<&str> = line.split_whitespace().collect();
1762 if cols.len() < 10 {
1763 return None;
1764 }
1765 if cols[3] != "0A" {
1767 return None;
1768 }
1769 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1770 let inode = cols[9].parse::<u64>().ok()?;
1771 Some((port, inode))
1772}
1773
1774fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1779 crate::outcome::fallback(cache, None, original, |bytes| {
1780 parse_cache_at(bytes, None, now)
1781 })
1782}
1783
1784fn fallback_with_error(
1789 cache: &Cache,
1790 last_error: Option<(u16, String)>,
1791 reason: AppError,
1792 now: DateTime<Utc>,
1793) -> Result<FetchOutcome> {
1794 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1795 parse_cache_at(bytes, None, now)
1796 })
1797}
1798
1799fn reuse_cache(
1800 bytes: Vec<u8>,
1801 cache: &Cache,
1802 stale: bool,
1803 account: Option<&str>,
1804 now: DateTime<Utc>,
1805) -> Result<FetchOutcome> {
1806 let snap = parse_cache_at(&bytes, account, now)?;
1807 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1808}
1809
1810pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1816 parse_cache_at(bytes, account, Utc::now())
1817}
1818
1819fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<String> {
1829 let shared = [
1830 ("Gemini 5h", snap.session.as_ref()),
1831 ("Gemini weekly", snap.weekly.as_ref()),
1832 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1833 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1834 ]
1835 .into_iter()
1836 .filter_map(|(name, w)| Some((name.to_string(), w?)));
1837 let pools = snap.model_pools.iter().map(|pool| {
1838 let name = format!("{} {}", pool.name, pool_cadence(pool.is_weekly()));
1839 (name, &pool.window)
1840 });
1841 shared
1842 .chain(pools)
1843 .find(|(_, w)| w.resets_at.is_some_and(|r| r <= now))
1844 .map(|(name, _)| name)
1845}
1846
1847fn pool_cadence(weekly: bool) -> &'static str {
1850 if weekly { "weekly" } else { "5h" }
1851}
1852
1853pub fn parse_cache_at(
1854 bytes: &[u8],
1855 account: Option<&str>,
1856 now: DateTime<Utc>,
1857) -> Result<AntigravitySnapshot> {
1858 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1859
1860 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1861 if let Some(expected) = account
1862 && cached_account != Some(expected)
1863 {
1864 return Err(AppError::Schema(
1865 "antigravity cache belongs to a different account; refetching".into(),
1866 ));
1867 }
1868
1869 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1879 match v.get(pct_key) {
1880 None => Err(AppError::Schema(format!(
1881 "antigravity: cached payload is missing {pct_key}"
1882 ))),
1883 Some(serde_json::Value::Null) => Ok(None),
1884 Some(value) => value
1885 .as_i64()
1886 .filter(|pct| (0..=100).contains(pct))
1887 .map(|pct| Some(pct as i32))
1888 .ok_or_else(|| {
1889 AppError::Schema(format!(
1890 "antigravity: cached {pct_key} must be an integer in 0..=100"
1891 ))
1892 }),
1893 }
1894 };
1895
1896 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1897 let Some(pct) = cached_pct(pct_key)? else {
1898 return Ok(None);
1899 };
1900 Ok::<_, AppError>(Some(UsageWindow {
1901 utilization_pct: pct,
1902 resets_at: parse_reset(&v[reset_key], reset_key)?,
1903 window_duration: if weekly {
1904 chrono::Duration::days(7)
1905 } else {
1906 chrono::Duration::hours(5)
1907 },
1908 }))
1909 };
1910
1911 let snap = AntigravitySnapshot {
1912 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1913 account: cached_account.unwrap_or_default().to_string(),
1914 source: v["source"]
1916 .as_str()
1917 .and_then(AntigravitySource::parse)
1918 .unwrap_or_default(),
1919 session: optional("session_pct", "session_reset", false)?,
1920 weekly: optional("weekly_pct", "weekly_reset", true)?,
1921 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1922 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1923 model_pools: cached_pools(&v)?,
1924 };
1925
1926 if snap.session.is_none()
1929 && snap.weekly.is_none()
1930 && snap.third_party_session.is_none()
1931 && snap.third_party_weekly.is_none()
1932 && snap.model_pools.is_empty()
1933 {
1934 return Err(AppError::Schema(
1935 "antigravity cache holds no usable window; refetching".into(),
1936 ));
1937 }
1938
1939 if let Some(window) = expired_window(&snap, now) {
1940 return Err(AppError::Schema(format!(
1941 "antigravity cache is past its {window} reset; refetching"
1942 )));
1943 }
1944 Ok(snap)
1945}
1946
1947fn cached_pools(v: &serde_json::Value) -> Result<Vec<ModelPoolWindow>> {
1951 let Some(entries) = v.get("model_pools") else {
1952 return Ok(Vec::new());
1953 };
1954 let malformed = || AppError::Schema("antigravity: cached model_pools is malformed".into());
1955 entries
1956 .as_array()
1957 .ok_or_else(malformed)?
1958 .iter()
1959 .map(|entry| {
1960 let name = entry["name"].as_str().ok_or_else(malformed)?;
1961 let weekly = match entry["window"].as_str() {
1962 Some(cadence) if cadence == pool_cadence(true) => true,
1963 Some(cadence) if cadence == pool_cadence(false) => false,
1964 _ => return Err(malformed()),
1965 };
1966 let pct = entry["pct"]
1967 .as_i64()
1968 .filter(|pct| (0..=100).contains(pct))
1969 .ok_or_else(malformed)?;
1970 Ok(ModelPoolWindow {
1971 name: crate::display::sanitize_untrusted_field(name),
1975 window: UsageWindow {
1976 utilization_pct: pct as i32,
1977 resets_at: parse_reset(&entry["reset"], "model_pools reset")?,
1978 window_duration: if weekly {
1979 chrono::Duration::days(7)
1980 } else {
1981 chrono::Duration::hours(5)
1982 },
1983 },
1984 })
1985 })
1986 .collect()
1987}
1988
1989pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1990 let model_pools: Vec<serde_json::Value> = snap
1991 .model_pools
1992 .iter()
1993 .map(|pool| {
1994 serde_json::json!({
1995 "name": pool.name,
1996 "window": pool_cadence(pool.is_weekly()),
1997 "pct": pool.window.utilization_pct,
1998 "reset": pool.window.resets_at.map(|dt| dt.to_rfc3339()),
1999 })
2000 })
2001 .collect();
2002 serde_json::json!({
2003 "plan": snap.plan,
2004 "account": snap.account,
2005 "source": snap.source.as_str(),
2006 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
2007 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2008 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
2009 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2010 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
2011 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2012 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
2013 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
2014 "model_pools": model_pools,
2015 })
2016}
2017
2018#[cfg(test)]
2019mod tests {
2020 use super::*;
2021
2022 const QUOTA_JSON: &str = r#"{
2026 "response": {
2027 "groups": [
2028 {
2029 "displayName": "Gemini Models",
2030 "buckets": [
2031 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
2032 "window": "weekly", "remainingFraction": 0.9191212,
2033 "resetTime": "2026-07-28T17:39:58Z"},
2034 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
2035 "window": "5h", "remainingFraction": 0.5672253,
2036 "resetTime": "2026-07-22T17:47:00Z"}
2037 ]
2038 },
2039 {
2040 "displayName": "Claude and GPT models",
2041 "buckets": [
2042 {"bucketId": "3p-weekly", "window": "weekly",
2043 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
2044 {"bucketId": "3p-5h", "window": "5h",
2045 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
2046 ]
2047 }
2048 ]
2049 }
2050 }"#;
2051
2052 fn now() -> DateTime<Utc> {
2055 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
2056 .unwrap()
2057 .with_timezone(&Utc)
2058 }
2059
2060 fn parsed() -> AntigravitySnapshot {
2061 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2062 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
2063 }
2064
2065 #[test]
2066 fn quota_summary_maps_four_distinct_windows() {
2067 let snap = parsed();
2068 assert_eq!(snap.plan, "Google AI Pro");
2069 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2071 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2072 assert_eq!(
2073 snap.third_party_session.as_ref().unwrap().utilization_pct,
2074 75
2075 );
2076 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
2077 }
2078
2079 #[test]
2080 fn each_window_keeps_its_own_reset_time() {
2081 let snap = parsed();
2082 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
2083 assert_eq!(
2084 snap.session.as_ref().unwrap().resets_at,
2085 at("2026-07-22T17:47:00Z")
2086 );
2087 assert_eq!(
2088 snap.weekly.as_ref().unwrap().resets_at,
2089 at("2026-07-28T17:39:58Z")
2090 );
2091 assert_eq!(
2092 snap.third_party_weekly.as_ref().unwrap().resets_at,
2093 at("2026-07-29T12:47:00Z")
2094 );
2095 assert_ne!(
2097 snap.session.as_ref().unwrap().resets_at,
2098 snap.weekly.as_ref().unwrap().resets_at
2099 );
2100 }
2101
2102 #[test]
2103 fn window_durations_match_their_bucket() {
2104 let snap = parsed();
2105 assert_eq!(
2106 snap.session.as_ref().unwrap().window_duration,
2107 chrono::Duration::hours(5)
2108 );
2109 assert_eq!(
2110 snap.weekly.as_ref().unwrap().window_duration,
2111 chrono::Duration::days(7)
2112 );
2113 assert_eq!(
2114 snap.third_party_weekly.as_ref().unwrap().window_duration,
2115 chrono::Duration::days(7)
2116 );
2117 }
2118
2119 #[test]
2120 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
2121 let v: serde_json::Value = serde_json::from_str(
2122 r#"{"response":{"groups":[
2123 {"displayName":"Gemini Models","buckets":[
2124 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
2125 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
2126 {"displayName":"Claude and GPT models","buckets":[
2127 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
2128 ]}}"#,
2129 )
2130 .unwrap();
2131 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2132 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
2133 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
2134 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
2135 assert!(snap.third_party_weekly.is_none());
2136 }
2137
2138 #[test]
2144 fn a_per_model_pool_is_reported_beside_the_canonical_buckets() {
2145 let v: serde_json::Value = serde_json::from_str(
2146 r#"{"response":{"groups":[
2147 {"displayName":"Gemini Models","buckets":[
2148 {"bucketId":"gemini-weekly","displayName":"Weekly Limit Remaining","window":"weekly",
2149 "remainingFraction":0.7839184,"resetTime":"2026-10-23T12:22:17Z"},
2150 {"bucketId":"gemini-5h","displayName":"Five Hour Limit Remaining","window":"5h",
2151 "remainingFraction":0.253501,"resetTime":"2026-10-16T17:22:17Z"}]},
2152 {"displayName":"Claude and GPT models","buckets":[
2153 {"bucketId":"3p-weekly","displayName":"Weekly Limit Remaining","window":"weekly",
2154 "remainingFraction":1,"resetTime":"2026-10-23T17:05:12Z"},
2155 {"bucketId":"3p-5h","displayName":"Five Hour Limit Remaining","window":"5h",
2156 "remainingFraction":1,"resetTime":"2026-10-16T22:05:12Z"}]},
2157 {"displayName":"Gemini 4 Argon","buckets":[
2158 {"bucketId":"gemini-argon-wtu-weekly","displayName":"Gemini 4 Argon Weekly Limit",
2159 "window":"weekly","remainingFraction":0.4,"resetTime":"2026-10-23T17:05:12Z"}]}
2160 ]}}"#,
2161 )
2162 .unwrap();
2163 let snap = parse_quota_summary(&v, "Google AI Pro".into()).unwrap();
2164 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
2165 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 22);
2166 assert_eq!(
2167 snap.third_party_session.as_ref().unwrap().utilization_pct,
2168 0
2169 );
2170 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
2171 let [argon] = snap.model_pools.as_slice() else {
2172 panic!("{:?}", snap.model_pools);
2173 };
2174 assert_eq!(argon.name, "Gemini 4 Argon");
2175 assert_eq!(argon.window.utilization_pct, 60);
2176 assert!(argon.is_weekly());
2177 }
2178
2179 #[test]
2182 fn a_malformed_model_pool_is_left_out() {
2183 let v: serde_json::Value = serde_json::from_str(
2184 r#"{"response":{"groups":[
2185 {"displayName":"Gemini Models","buckets":[
2186 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.5}]},
2187 {"displayName":"Gemini 4 Argon","buckets":[
2188 {"bucketId":"gemini-argon-wtu-weekly","window":"weekly","remainingFraction":"x"}]}
2189 ]}}"#,
2190 )
2191 .unwrap();
2192 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2193 assert_eq!(snap.weekly.unwrap().utilization_pct, 50);
2194 assert!(snap.model_pools.is_empty());
2195 }
2196
2197 #[test]
2199 fn a_model_pool_name_is_sanitized_and_falls_back_to_the_bucket() {
2200 let v: serde_json::Value = serde_json::from_str(
2201 "{\"response\":{\"groups\":[
2202 {\"displayName\":\"Gemini Models\",\"buckets\":[
2203 {\"bucketId\":\"gemini-weekly\",\"window\":\"weekly\",\"remainingFraction\":0.5}]},
2204 {\"displayName\":\"Argon\\u001b[31m\",\"buckets\":[
2205 {\"bucketId\":\"argon-weekly\",\"window\":\"weekly\",\"remainingFraction\":0.5}]},
2206 {\"buckets\":[
2207 {\"bucketId\":\"nova-5h\",\"displayName\":\"Nova Five Hour\",\"window\":\"5h\",\"remainingFraction\":0.5}]}
2208 ]}}",
2209 )
2210 .unwrap();
2211 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2212 let names: Vec<&str> = snap.model_pools.iter().map(|p| p.name.as_str()).collect();
2213 assert_eq!(names, ["Argon[31m", "Nova Five Hour"]);
2214 }
2215
2216 #[test]
2217 fn canonical_presence_disables_inference_in_a_shared_group() {
2218 for remaining in [serde_json::json!(0.1), serde_json::json!("unknown")] {
2219 let summary = serde_json::json!({"groups": [{
2220 "displayName": "Gemini Models",
2221 "buckets": [
2222 {"bucketId": "gemini-5h", "remainingFraction": 0.5},
2223 {"bucketId": "gemini-argon-wtu-weekly", "remainingFraction": remaining}
2224 ]
2225 }]});
2226 let snap = parse_quota_summary(&summary, "Pro".into()).unwrap();
2227 assert_eq!(snap.session.unwrap().utilization_pct, 50);
2228 assert!(snap.weekly.is_none());
2229 assert!(snap.model_pools.is_empty());
2230 }
2231 }
2232
2233 #[test]
2234 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
2235 let lookalike: serde_json::Value = serde_json::from_str(
2238 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
2239 {"bucketId":"gemini-5h-copy","window":"5h"},
2240 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2241 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
2242 ]}]}}"#,
2243 )
2244 .unwrap();
2245 let snap = parse_quota_summary(&lookalike, "Pro".into()).unwrap();
2246 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 10);
2247
2248 for buckets in [
2251 r#"{"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2252 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.1}"#,
2253 r#"{"bucketId":"x1","window":"5h","remainingFraction":0.9},
2254 {"bucketId":"x2","window":"5h","remainingFraction":0.1}"#,
2255 ] {
2256 let duplicate: serde_json::Value = serde_json::from_str(&format!(
2257 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[{buckets}]}}]}}}}"#
2258 ))
2259 .unwrap();
2260 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
2261 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
2262 }
2263
2264 let unrelated: serde_json::Value = serde_json::from_str(
2267 r#"{"response":{"groups":[
2268 {"displayName":"Gemini Models","buckets":[
2269 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2270 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
2271 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
2272 ]},
2273 {"displayName":"Future Models","buckets":[
2274 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
2275 ]}
2276 ]}}"#,
2277 )
2278 .unwrap();
2279 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
2280 assert!(snap.third_party_session.is_none());
2281 assert!(snap.third_party_weekly.is_none());
2282 let [future] = snap.model_pools.as_slice() else {
2283 panic!("{:?}", snap.model_pools);
2284 };
2285 assert_eq!(future.name, "Future Models");
2286 assert!(!future.is_weekly());
2287 }
2288
2289 #[test]
2292 fn a_bucket_without_a_usable_fraction_is_rejected() {
2293 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
2294 let v: serde_json::Value = serde_json::from_str(&format!(
2295 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
2296 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
2297 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
2298 ))
2299 .unwrap();
2300 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
2301 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
2302 }
2303 }
2304
2305 #[test]
2306 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
2307 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
2308 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2309 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
2310 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
2311 assert!(err.to_string().contains("resetTime"), "{err}");
2312 }
2313 }
2314
2315 #[test]
2319 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
2320 let mut snap = parsed();
2321 snap.session = None;
2322 snap.third_party_session = None;
2323
2324 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2325 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
2326
2327 assert!(back.session.is_none());
2328 assert_eq!(
2329 back.weekly.as_ref().unwrap().utilization_pct,
2330 snap.weekly.as_ref().unwrap().utilization_pct
2331 );
2332 }
2333
2334 #[test]
2339 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
2340 let summary = serde_json::json!({
2341 "groups": [
2342 {
2343 "displayName": "Gemini Models",
2344 "buckets": [{
2345 "bucketId": "gemini-weekly", "window": "weekly",
2346 "remainingFraction": 0.42,
2347 }],
2348 },
2349 {
2350 "displayName": "Claude and GPT models",
2351 "buckets": [{
2352 "bucketId": "3p-weekly", "window": "weekly",
2353 "remainingFraction": 0.9,
2354 }],
2355 },
2356 ],
2357 });
2358
2359 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
2360
2361 assert!(snap.session.is_none(), "no 5h bucket arrived");
2362 assert!(snap.third_party_session.is_none());
2363 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
2364 assert_eq!(
2365 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
2366 10
2367 );
2368 }
2369
2370 #[test]
2373 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
2374 let summary = serde_json::json!({
2375 "groups": [{
2376 "displayName": "Gemini Models",
2377 "buckets": [{
2378 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
2379 }],
2380 }],
2381 });
2382
2383 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
2384
2385 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
2386 assert!(snap.weekly.is_none());
2387 }
2388
2389 #[test]
2392 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
2393 let summary = serde_json::json!({
2394 "groups": [{
2395 "displayName": "Gemini Models",
2396 "buckets": [{
2397 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
2398 }],
2399 }],
2400 });
2401
2402 let rendered = parse_quota_summary(&summary, "Pro".into())
2403 .expect_err("an unrecognised cadence alone is not a snapshot")
2404 .to_string();
2405
2406 assert!(
2407 rendered.contains("no bucket in a window we recognise"),
2408 "{rendered}"
2409 );
2410 assert!(rendered.contains("gemini-daily"), "{rendered}");
2411 assert!(rendered.contains("window daily"), "{rendered}");
2412 }
2413
2414 #[test]
2417 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
2418 let summary = serde_json::json!({
2419 "groups": [{"displayName": "Gemini", "buckets": []}],
2420 });
2421
2422 let rendered = parse_quota_summary(&summary, "Pro".into())
2423 .expect_err("no buckets is an error")
2424 .to_string();
2425
2426 assert!(rendered.contains("no buckets at all"), "{rendered}");
2427 assert!(!rendered.contains("it offered:"), "{rendered}");
2428 }
2429
2430 #[test]
2433 fn buckets_without_an_id_are_still_named_in_the_error() {
2434 let summary = serde_json::json!({
2435 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
2436 });
2437
2438 let rendered = parse_quota_summary(&summary, "Pro".into())
2439 .expect_err("an unusable summary is an error")
2440 .to_string();
2441
2442 assert!(rendered.contains("<unnamed>"), "{rendered}");
2443 }
2444
2445 #[test]
2446 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
2447 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
2448 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
2449 }
2450
2451 #[test]
2452 fn cache_round_trip_preserves_every_window() {
2453 let snap = parsed();
2454 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2455 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2456 }
2457
2458 #[test]
2462 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
2463 let full = snap_to_json(&parsed());
2464 for missing in ["session_pct", "weekly_pct"] {
2465 let mut v = full.clone();
2466 v.as_object_mut().unwrap().remove(missing);
2467 let bytes = serde_json::to_vec(&v).unwrap();
2468 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2469 assert!(err.to_string().contains(missing), "{missing}: {err}");
2470 }
2471 assert!(parse_cache_at(b"{}", None, now()).is_err());
2473 }
2474
2475 #[test]
2476 fn cached_percentages_are_range_checked_before_narrowing() {
2477 let full = snap_to_json(&parsed());
2478 for (key, bad) in [
2479 ("session_pct", serde_json::json!(-1)),
2480 ("weekly_pct", serde_json::json!(101)),
2481 ("session_pct", serde_json::json!(i64::MAX)),
2482 ("tp_session_pct", serde_json::json!("75")),
2483 ] {
2484 let mut v = full.clone();
2485 v[key] = bad;
2486 let bytes = serde_json::to_vec(&v).unwrap();
2487 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2488 assert!(err.to_string().contains(key), "{key}: {err}");
2489 }
2490 }
2491
2492 #[test]
2493 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
2494 let mut v = snap_to_json(&parsed());
2495 v["session_reset"] = serde_json::json!("not-a-time");
2496 let bytes = serde_json::to_vec(&v).unwrap();
2497 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2498 assert!(err.to_string().contains("session_reset"), "{err}");
2499 }
2500
2501 #[test]
2506 fn a_cache_past_its_reset_is_refused() {
2507 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
2508 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2509
2510 assert!(parse_cache_at(&bytes, None, now()).is_ok());
2512 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
2514
2515 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
2517 assert!(err.to_string().contains("5h"), "{err}");
2518
2519 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
2521 }
2522
2523 #[test]
2526 fn expiry_names_the_window_that_rolled_over() {
2527 let mut snap = parsed();
2528 snap.session.as_mut().unwrap().resets_at = None;
2530 snap.third_party_session = None;
2531 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2532 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2533
2534 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
2536
2537 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
2539 assert!(err.to_string().contains("Gemini weekly"), "{err}");
2540 }
2541
2542 #[test]
2544 fn a_window_without_a_reset_never_expires() {
2545 let mut snap = parsed();
2546 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
2547 w.resets_at = None;
2548 }
2549 snap.third_party_session = None;
2550 snap.third_party_weekly = None;
2551 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2552 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
2553 .unwrap()
2554 .with_timezone(&Utc);
2555 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
2556 }
2557
2558 #[test]
2560 fn a_cache_from_another_account_is_rejected() {
2561 let mut snap = parsed();
2562 snap.account = "acct:aaaa".into();
2563 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2564
2565 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
2566 assert_eq!(
2567 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
2568 snap
2569 );
2570
2571 let mut legacy = snap_to_json(&snap);
2573 legacy.as_object_mut().unwrap().remove("account");
2574 let legacy = serde_json::to_vec(&legacy).unwrap();
2575 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
2576 }
2577
2578 #[test]
2581 fn an_unverifiable_cache_is_served_rather_than_discarded() {
2582 let mut snap = parsed();
2583 snap.account = "acct:aaaa".into();
2584 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2585 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2586 }
2587
2588 #[test]
2589 fn account_key_fingerprints_rather_than_storing_the_address() {
2590 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
2591 let a = with("someone@example.com");
2592 assert!(!a.contains("someone"), "{a}");
2593 assert!(!a.contains('@'), "{a}");
2594 assert_eq!(a, with("someone@example.com"), "must be stable");
2595 assert_ne!(a, with("other@example.com"));
2596 let unknown = account_key(&serde_json::json!({}));
2598 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
2599 assert_ne!(unknown, a);
2600 }
2601
2602 #[test]
2605 fn absent_third_party_windows_are_not_treated_as_corruption() {
2606 let mut snap = parsed();
2607 snap.third_party_session = None;
2608 snap.third_party_weekly = None;
2609 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2610 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2611 }
2612
2613 #[test]
2614 fn cache_round_trip_preserves_absent_third_party_windows() {
2615 let mut snap = parsed();
2616 snap.third_party_session = None;
2617 snap.third_party_weekly = None;
2618 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2619 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2620 }
2621
2622 #[test]
2623 fn pct_used_inverts_valid_fractions() {
2624 assert_eq!(pct_used(1.0), 0);
2625 assert_eq!(pct_used(0.0), 100);
2626 assert_eq!(pct_used(0.5), 50);
2627 }
2628
2629 #[test]
2630 fn plan_falls_back_through_the_status_payload() {
2631 let tier: serde_json::Value =
2632 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2633 .unwrap();
2634 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2635
2636 let plan_only: serde_json::Value = serde_json::from_str(
2637 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2638 )
2639 .unwrap();
2640 assert_eq!(plan_from_status(&plan_only), "Pro");
2641
2642 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2643 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2644 }
2645
2646 #[cfg(target_os = "linux")]
2647 #[test]
2648 fn proc_net_parser_keeps_only_listening_rows() {
2649 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2650 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2651 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2652
2653 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2654 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2655 assert_eq!(parse_proc_net_line(established), None);
2656
2657 assert_eq!(parse_proc_net_line("garbage"), None);
2658 }
2659
2660 #[test]
2661 fn explicit_address_comes_first_and_gets_a_scheme() {
2662 let candidate = |base: &str, pid| Candidate {
2663 base: base.to_string(),
2664 pid,
2665 ide: false,
2666 };
2667 assert_eq!(
2668 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)], &no_ide()),
2669 vec![
2670 candidate("http://127.0.0.1:1234", None),
2671 candidate("http://127.0.0.1:5678", Some(10)),
2672 ]
2673 );
2674 assert_eq!(
2676 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)], &no_ide()),
2677 vec![
2678 candidate("http://127.0.0.1:1234", None),
2679 candidate("http://127.0.0.1:5678", Some(10)),
2680 ]
2681 );
2682 assert_eq!(
2685 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)], &no_ide()),
2686 vec![candidate("http://127.0.0.1:5678", Some(10))]
2687 );
2688 assert_eq!(
2690 candidate_bases_with(None, vec![(10, 5678), (10, 5678)], &no_ide()),
2691 vec![candidate("http://127.0.0.1:5678", Some(10))]
2692 );
2693 assert_eq!(
2695 candidate_bases_with(Some("https://host:9"), vec![], &no_ide()),
2696 vec![candidate("https://host:9", None)]
2697 );
2698 }
2699
2700 fn no_ide() -> std::collections::HashSet<u32> {
2703 std::collections::HashSet::new()
2704 }
2705
2706 fn http(status: u16) -> AppError {
2707 AppError::Http {
2708 status,
2709 body: String::new(),
2710 }
2711 }
2712
2713 fn missing_csrf() -> AppError {
2714 missing_csrf_with("missing CSRF token")
2715 }
2716
2717 fn missing_csrf_with(message: &str) -> AppError {
2718 AppError::Http {
2719 status: 401,
2720 body: serde_json::json!({
2721 "code": "unauthenticated",
2722 "message": message,
2723 })
2724 .to_string(),
2725 }
2726 }
2727
2728 #[test]
2729 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2730 assert!(is_missing_csrf(&missing_csrf()));
2731 assert!(is_missing_csrf(&missing_csrf_with(
2732 "unauthenticated: missing CSRF token"
2733 )));
2734 assert!(is_missing_csrf(&missing_csrf_with(
2735 " unauthenticated: missing CSRF token "
2736 )));
2737 assert!(!is_missing_csrf(&missing_csrf_with(
2738 "unauthenticated: invalid token"
2739 )));
2740 assert!(!is_missing_csrf(&AppError::Http {
2741 status: 403,
2742 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2743 }));
2744 assert!(!is_missing_csrf(&AppError::Http {
2745 status: 401,
2746 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2747 }));
2748 assert!(!is_missing_csrf(&AppError::Http {
2749 status: 401,
2750 body: "prefix: missing CSRF token".into(),
2751 }));
2752 }
2753
2754 #[test]
2755 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2756 for errors in [
2757 vec![missing_csrf(), http(403)],
2758 vec![http(401), missing_csrf()],
2759 ] {
2760 let err = select_probe_error(errors);
2761 assert!(is_actionable(&err), "{err}");
2762 assert!(!is_missing_csrf(&err), "{err}");
2763 }
2764 }
2765
2766 #[test]
2767 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2768 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2769 assert!(is_missing_csrf(&err), "{err}");
2770 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2771 }
2772
2773 #[test]
2777 fn an_auth_failure_outranks_later_transport_noise() {
2778 let err = select_probe_error(vec![
2779 http(401),
2780 AppError::Transport("connection refused".into()),
2781 ]);
2782 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2783
2784 let err = select_probe_error(vec![
2785 AppError::Transport("connection refused".into()),
2786 http(403),
2787 ]);
2788 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2789 }
2790
2791 #[test]
2794 fn the_first_auth_failure_wins() {
2795 let err = select_probe_error(vec![http(401), http(403)]);
2796 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2797 }
2798
2799 #[test]
2803 fn without_an_auth_failure_the_last_error_stands() {
2804 let err = select_probe_error(vec![
2805 AppError::Transport("first".into()),
2806 http(500),
2807 AppError::Transport("last".into()),
2808 ]);
2809 assert!(
2810 matches!(&err, AppError::Transport(m) if m == "last"),
2811 "{err}"
2812 );
2813 assert!(err.is_transient());
2814 }
2815
2816 #[test]
2819 fn a_server_error_is_not_treated_as_actionable() {
2820 let err = select_probe_error(vec![http(500), http(401)]);
2821 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2822 }
2823
2824 #[test]
2825 fn no_candidates_at_all_yields_a_generic_error() {
2826 let err = select_probe_error(Vec::new());
2827 assert!(
2828 err.to_string().contains("no local server answered"),
2829 "{err}"
2830 );
2831 }
2832
2833 fn tls_echo() -> AppError {
2837 AppError::Http {
2838 status: 400,
2839 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2840 }
2841 }
2842
2843 #[test]
2848 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2849 let err = select_probe_error(vec![
2850 AppError::Http {
2851 status: 500,
2852 body: "GetUserStatus: internal".into(),
2853 },
2854 tls_echo(),
2855 ]);
2856 assert!(
2857 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2858 "{err}"
2859 );
2860 }
2861
2862 #[test]
2867 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2868 let err = select_probe_error(vec![
2869 AppError::Transport("connection refused".into()),
2870 tls_echo(),
2871 ]);
2872 assert!(
2873 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2874 "{err}"
2875 );
2876 assert!(
2877 err.is_transient(),
2878 "the echo must not make the run non-transient: {err}"
2879 );
2880 }
2881
2882 #[test]
2885 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2886 let err = select_probe_error(vec![tls_echo()]);
2887 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2888 }
2889
2890 #[test]
2893 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2894 let err = select_probe_error(vec![
2895 AppError::Http {
2896 status: 400,
2897 body: "unknown method GetUserStatus".into(),
2898 },
2899 tls_echo(),
2900 ]);
2901 assert!(
2902 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2903 "{err}"
2904 );
2905 }
2906
2907 #[test]
2909 fn an_auth_failure_still_outranks_a_tls_echo() {
2910 let err = select_probe_error(vec![tls_echo(), http(401)]);
2911 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2912 }
2913
2914 #[test]
2915 fn every_discovered_port_is_probed_in_order() {
2916 assert_eq!(
2917 candidate_bases_with(None, vec![(10, 33875), (10, 37435)], &no_ide()),
2918 vec![
2919 Candidate {
2920 base: "http://127.0.0.1:33875".to_string(),
2921 pid: Some(10),
2922 ide: false,
2923 },
2924 Candidate {
2925 base: "http://127.0.0.1:37435".to_string(),
2926 pid: Some(10),
2927 ide: false,
2928 },
2929 ]
2930 );
2931 }
2932
2933 #[test]
2938 fn empty_discovery_yields_no_candidates() {
2939 assert!(candidate_bases_with(None, vec![], &no_ide()).is_empty());
2940 assert!(candidate_bases_with(Some(""), vec![], &no_ide()).is_empty());
2941 }
2942
2943 #[test]
2944 fn every_antigravity_product_is_recognised() {
2945 assert!(is_antigravity_process(
2947 "language_server\n",
2948 Some("/opt/antigravity/resources/bin/language_server")
2949 ));
2950 assert!(is_antigravity_process(
2952 "agy\n",
2953 Some("/home/u/.local/bin/agy")
2954 ));
2955 assert!(is_antigravity_process(
2957 "node",
2958 Some("/opt/antigravity/bin/helper")
2959 ));
2960 assert!(is_antigravity_process("antigravity", None));
2961 assert!(is_antigravity_process("agy.exe", None));
2962 assert!(is_antigravity_process("Antigravity.exe", None));
2963 assert!(is_antigravity_process("language_server.exe", None));
2964 assert!(is_antigravity_process(
2965 "language_server_windows_x64.exe",
2966 None
2967 ));
2968 assert!(is_antigravity_process(
2969 "node.exe",
2970 Some(r"C:\Users\u\AppData\Local\agy.exe")
2971 ));
2972 }
2973
2974 #[test]
2975 fn unrelated_processes_are_not_probed() {
2976 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2977 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2978 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2980 assert!(!is_antigravity_process("legacy.exe", None));
2981 assert!(!is_antigravity_process("not-agy.exe", None));
2982 assert!(!is_antigravity_process("", None));
2983 }
2984
2985 #[test]
2986 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2987 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2988 raw.extend([0, b'x' as u16]);
2989 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2990 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2991 assert_eq!(decode_windows_process_name(&[]), "");
2992 }
2993
2994 #[test]
2995 fn windows_process_filter_keeps_only_antigravity_pids() {
2996 let processes = vec![
2997 (10, "agy.exe".to_string()),
2998 (20, "language_server_windows_x64.exe".to_string()),
2999 (30, "sshd.exe".to_string()),
3000 ];
3001 let pids = matching_windows_process_ids(&processes);
3002 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
3003 }
3004
3005 #[test]
3006 fn windows_listener_filter_joins_pid_loopback_and_port() {
3007 let pids = std::collections::HashSet::from([10]);
3008 let rows = [
3009 WindowsTcpRow {
3010 local_addr: [127, 0, 0, 1],
3011 local_port: u32::from(59870u16.to_be()),
3012 pid: 10,
3013 },
3014 WindowsTcpRow {
3015 local_addr: [127, 0, 0, 1],
3016 local_port: u32::from(59868u16.to_be()),
3017 pid: 10,
3018 },
3019 WindowsTcpRow {
3020 local_addr: [127, 0, 0, 1],
3021 local_port: u32::from(59870u16.to_be()),
3022 pid: 10,
3023 },
3024 WindowsTcpRow {
3025 local_addr: [0, 0, 0, 0],
3026 local_port: u32::from(50000u16.to_be()),
3027 pid: 10,
3028 },
3029 WindowsTcpRow {
3030 local_addr: [127, 0, 0, 1],
3031 local_port: u32::from(50001u16.to_be()),
3032 pid: 99,
3033 },
3034 WindowsTcpRow {
3035 local_addr: [127, 0, 0, 1],
3036 local_port: 0,
3037 pid: 10,
3038 },
3039 ];
3040 assert_eq!(
3041 matching_windows_ports(&pids, &rows),
3042 vec![(10, 59870), (10, 59868)]
3043 );
3044 }
3045
3046 #[test]
3050 fn windows_ports_from_two_products_keep_tls_listeners_last() {
3051 let pids = std::collections::HashSet::from([10, 20]);
3052 let row = |port: u16, pid: u32| WindowsTcpRow {
3053 local_addr: [127, 0, 0, 1],
3054 local_port: u32::from(port.to_be()),
3055 pid,
3056 };
3057 let rows = [
3058 row(40000, 10),
3059 row(40001, 10),
3060 row(50000, 20),
3061 row(50001, 20),
3062 ];
3063 assert_eq!(
3064 matching_windows_ports(&pids, &rows),
3065 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
3066 );
3067 }
3068
3069 #[test]
3073 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
3074 use std::collections::BTreeMap;
3075
3076 assert_eq!(
3078 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
3079 vec![(10, 59870), (10, 59868)]
3080 );
3081 assert_eq!(
3086 probe_order(BTreeMap::from([
3087 (10, vec![40000, 40001]),
3088 (20, vec![50000, 50001]),
3089 ])),
3090 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
3091 );
3092 assert_eq!(
3095 probe_order(BTreeMap::from([
3096 (10, vec![6000, 5000, 4000]),
3097 (20, vec![6000, 7000]),
3098 ])),
3099 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
3100 );
3101 assert!(probe_order(BTreeMap::new()).is_empty());
3102 }
3103
3104 #[test]
3109 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
3110 use std::collections::BTreeMap;
3111
3112 assert_eq!(
3113 probe_order(BTreeMap::from([
3114 (10, vec![40001, 40001, 40000, 40000]),
3115 (20, vec![50001, 50000]),
3116 ])),
3117 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
3118 "duplicate rows must not reorder the ranks below them"
3119 );
3120 }
3121
3122 #[test]
3127 fn a_half_started_product_is_the_documented_exception() {
3128 use std::collections::BTreeMap;
3129
3130 assert_eq!(
3131 probe_order(BTreeMap::from([
3132 (10, vec![40000]),
3133 (20, vec![50001, 50000])
3134 ])),
3135 vec![(10, 40000), (20, 50001), (20, 50000)]
3136 );
3137 }
3138
3139 #[test]
3144 fn an_override_with_no_authority_is_dropped_not_probed() {
3145 for junk in ["/", "///", "http://", "https://", " / "] {
3146 assert_eq!(
3147 candidate_bases_with(Some(junk), vec![(10, 4242)], &no_ide()),
3148 vec![Candidate {
3149 base: "http://127.0.0.1:4242".to_string(),
3150 pid: Some(10),
3151 ide: false,
3152 }],
3153 "{junk:?} should not survive as a candidate"
3154 );
3155 }
3156 assert!(candidate_bases_with(Some("/"), vec![], &no_ide()).is_empty());
3157 }
3158
3159 #[test]
3160 fn windows_table_bounds_reject_truncation_and_overflow() {
3161 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
3162 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
3163 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
3164 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
3165 assert_eq!(
3166 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
3167 None
3168 );
3169 }
3170
3171 #[cfg(target_os = "windows")]
3172 #[test]
3173 fn windows_tcp_table_parser_copies_complete_rows_only() {
3174 use std::mem::{offset_of, size_of};
3175 use windows_sys::Win32::NetworkManagement::IpHelper::{
3176 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
3177 };
3178
3179 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
3180 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
3181 let words = used.div_ceil(size_of::<u32>());
3182 let mut buffer = vec![0u32; words];
3183 let first = MIB_TCPROW_OWNER_PID {
3184 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
3185 dwLocalPort: u32::from(59868u16.to_be()),
3186 dwOwningPid: 10,
3187 ..Default::default()
3188 };
3189 let second = MIB_TCPROW_OWNER_PID {
3190 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
3191 dwLocalPort: u32::from(59870u16.to_be()),
3192 dwOwningPid: 10,
3193 ..Default::default()
3194 };
3195 unsafe {
3196 buffer.as_mut_ptr().write_unaligned(2);
3197 let rows = buffer
3198 .as_mut_ptr()
3199 .cast::<u8>()
3200 .add(offset)
3201 .cast::<MIB_TCPROW_OWNER_PID>();
3202 rows.write_unaligned(first);
3203 rows.add(1).write_unaligned(second);
3204 }
3205
3206 assert_eq!(
3207 parse_windows_tcp_rows(&buffer, used),
3208 vec![
3209 WindowsTcpRow {
3210 local_addr: [127, 0, 0, 1],
3211 local_port: u32::from(59868u16.to_be()),
3212 pid: 10,
3213 },
3214 WindowsTcpRow {
3215 local_addr: [127, 0, 0, 1],
3216 local_port: u32::from(59870u16.to_be()),
3217 pid: 10,
3218 },
3219 ]
3220 );
3221 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
3222 }
3223
3224 #[test]
3225 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
3226 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";
3229 assert_eq!(
3230 parse_lsof_pcn(output),
3231 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
3232 );
3233 }
3234
3235 #[test]
3239 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
3240 let output = concat!(
3241 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
3242 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
3243 );
3244 assert_eq!(
3245 parse_lsof_pcn(output),
3246 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
3247 "both HTTP listeners must precede both TLS listeners"
3248 );
3249 }
3250
3251 #[test]
3252 fn lsof_parser_matches_the_capitalised_macos_app_name() {
3253 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
3254 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
3255 }
3256
3257 #[test]
3258 fn lsof_parser_deduplicates_and_handles_empty_output() {
3259 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
3260 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
3261 assert!(parse_lsof_pcn("").is_empty());
3262 }
3263
3264 #[test]
3268 fn macos_system_tools_are_pinned_to_trusted_paths() {
3269 let sources = [
3270 crate::guard::production_code(include_str!("fetch.rs")),
3271 crate::guard::production_code(include_str!("../tray/tui_launch_macos.rs")),
3272 crate::guard::production_code(include_str!("../tray/browse.rs")),
3273 crate::guard::production_code(include_str!("../nous/oauth.rs")),
3274 crate::guard::production_code(include_str!("../../macos/ai-usagebar-menubar.swift")),
3275 ];
3276 let code = sources.join("\n");
3277 for bare in [
3278 r#"Command::new("lsof")"#,
3279 r#"Command::new("ps")"#,
3280 r#"Command::new("open")"#,
3281 r#"Command::new("osascript")"#,
3282 r#"do script \"bash "#,
3283 ] {
3284 assert!(!code.contains(bare), "{bare} looks the tool up on PATH");
3285 }
3286 for pinned in [
3287 r#"Command::new("/usr/sbin/lsof")"#,
3288 r#"Command::new("/bin/ps")"#,
3289 r#"Command::new("/usr/bin/open")"#,
3290 r#"Command::new("/usr/bin/osascript")"#,
3291 r#"do script \"/bin/bash "#,
3292 ] {
3293 assert!(
3294 code.contains(pinned),
3295 "missing trusted system path: {pinned}"
3296 );
3297 }
3298 }
3299
3300 #[test]
3303 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
3304 let dir = tempfile::tempdir().unwrap();
3305 let cache = Cache::at(dir.path().join("usage.json"));
3306 let reason = AppError::Credentials("Antigravity: no local language server found".into());
3307
3308 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
3309 let msg = err.to_string();
3310 assert!(msg.contains("no local language server found"), "{msg}");
3311 assert!(!msg.contains("no usable cache"), "{msg}");
3312 }
3313
3314 #[test]
3315 fn unusable_cache_does_not_replace_the_live_diagnosis() {
3316 let dir = tempfile::tempdir().unwrap();
3317 let cache = Cache::at(dir.path().join("antigravity"));
3318 cache.write_payload(b"{}").unwrap();
3319
3320 let reason = AppError::Credentials("Antigravity must be running".into());
3321 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
3322 assert!(err.to_string().contains("must be running"), "{err}");
3323
3324 let original = AppError::Transport("original loopback failure".into());
3325 let err = fallback_silent(&cache, now(), original).unwrap_err();
3326 assert!(
3327 err.to_string().contains("original loopback failure"),
3328 "{err}"
3329 );
3330 }
3331
3332 #[tokio::test]
3333 async fn rpc_error_bodies_are_bounded_too() {
3334 let mut server = mockito::Server::new_async().await;
3335 let path = format!("/{STATUS_RPC}");
3336 server
3337 .mock("POST", path.as_str())
3338 .with_status(500)
3339 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
3340 .create_async()
3341 .await;
3342
3343 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
3344 .await
3345 .unwrap_err();
3346 assert!(err.to_string().contains("exceeds"), "{err}");
3347 }
3348
3349 #[test]
3350 fn blank_override_falls_through_to_discovery() {
3351 assert_eq!(
3352 candidate_bases_with(Some(" "), vec![(10, 4242)], &no_ide()),
3353 vec![Candidate {
3354 base: "http://127.0.0.1:4242".to_string(),
3355 pid: Some(10),
3356 ide: false,
3357 }]
3358 );
3359 }
3360
3361 #[test]
3364 fn only_the_ide_install_is_recognised() {
3365 assert!(is_ide_install(
3366 "/opt/antigravity-ide/resources/app/extensions/antigravity/bin/language_server_linux_x64"
3367 ));
3368 assert!(is_ide_install(
3369 "/Applications/Antigravity IDE.app/Contents/Resources/app/extensions/antigravity/bin/language_server_darwin_arm64"
3370 ));
3371 assert!(is_ide_install(
3372 r"C:\Users\u\AppData\Local\Programs\Antigravity IDE\language_server_windows_x64.exe"
3373 ));
3374 assert!(!is_ide_install(
3375 "/opt/antigravity/resources/bin/language_server"
3376 ));
3377 assert!(!is_ide_install("/home/u/.local/bin/agy"));
3378 assert!(!is_ide_install("/usr/bin/node"));
3379 }
3380
3381 #[test]
3383 fn ps_output_pairs_pids_with_full_paths() {
3384 let output = " 101 /opt/antigravity-ide/bin/language_server_linux_x64\n\
3385 202 /Applications/Antigravity IDE.app/Contents/MacOS/Antigravity IDE\n";
3386 assert_eq!(
3387 parse_ps_comm(output),
3388 vec![
3389 (
3390 101,
3391 "/opt/antigravity-ide/bin/language_server_linux_x64".to_string()
3392 ),
3393 (
3394 202,
3395 "/Applications/Antigravity IDE.app/Contents/MacOS/Antigravity IDE".to_string()
3396 ),
3397 ]
3398 );
3399 assert!(parse_ps_comm("").is_empty());
3400 assert!(parse_ps_comm("not-a-pid /tmp/x").is_empty());
3401 }
3402
3403 #[test]
3407 fn ide_owned_listeners_are_marked_but_an_explicit_override_is_not() {
3408 let ide: std::collections::HashSet<u32> = [10].into_iter().collect();
3409
3410 let candidates = candidate_bases_with(None, vec![(10, 40001), (20, 50001)], &ide);
3411 assert_eq!(candidates.len(), 2);
3412 assert!(candidates[0].ide, "pid 10 is the IDE's language server");
3413 assert!(!candidates[1].ide);
3414
3415 let candidates = candidate_bases_with(Some("127.0.0.1:40001"), vec![(10, 40001)], &ide);
3416 assert_eq!(
3417 candidates,
3418 vec![Candidate {
3419 base: "http://127.0.0.1:40001".to_string(),
3420 pid: Some(10),
3421 ide: false,
3422 }],
3423 "ANTIGRAVITY_LS_ADDRESS is an explicit target, never skipped"
3424 );
3425 }
3426
3427 fn fixture() -> (tempfile::TempDir, Cache) {
3432 let td = tempfile::TempDir::new().unwrap();
3433 let cache = Cache::at(td.path().join("antigravity"));
3434 (td, cache)
3435 }
3436
3437 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
3438 let base = server.url();
3439 cloud::Endpoints {
3440 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
3441 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
3442 token: format!("{base}/token"),
3443 }
3444 }
3445
3446 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
3448 let mut token = serde_json::json!({
3449 "access_token": "KEYRING-AT",
3450 "expiry": expiry,
3451 });
3452 if with_refresh {
3453 token["refresh_token"] = serde_json::json!("KEYRING-RT");
3454 }
3455 serde_json::json!({ "token": token }).to_string()
3456 }
3457
3458 const EXPIRED: &str = "2026-07-22T11:00:00Z";
3460 const VALID: &str = "2026-07-22T13:00:00Z";
3462
3463 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
3465 RemoteOverride {
3466 credential: SavedCredential::Blob(blob),
3467 endpoints: Some(eps),
3468 local_bases: Some(vec![]),
3469 ..Default::default()
3470 }
3471 }
3472
3473 fn bare_summary() -> String {
3476 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
3477 v["response"].to_string()
3478 }
3479
3480 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
3481 server
3482 .mock("POST", "/daily/quota")
3483 .match_header("authorization", format!("Bearer {bearer}").as_str())
3484 .with_status(200)
3485 .with_body(bare_summary())
3486 }
3487
3488 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
3489 server
3490 .mock("POST", "/token")
3491 .with_status(200)
3492 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
3493 }
3494
3495 fn test_oauth() -> cloud::OauthClient {
3496 cloud::OauthClient {
3497 id: "test-client".into(),
3498 secret: "test-client-secret".into(),
3499 }
3500 }
3501
3502 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
3503 fetch_snapshot_at(
3504 &reqwest::Client::new(),
3505 cache,
3506 ttl,
3507 Some(&test_oauth()),
3508 remote,
3509 now(),
3510 )
3511 .await
3512 }
3513
3514 #[tokio::test]
3515 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
3516 let mut server = mockito::Server::new_async().await;
3517 let eps = endpoints(&server);
3518 let quota = quota_mock(&mut server, "KEYRING-AT")
3519 .expect(1)
3520 .create_async()
3521 .await;
3522 let plan = server
3523 .mock("POST", "/daily/plan")
3524 .match_header("authorization", "Bearer KEYRING-AT")
3525 .with_status(200)
3526 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3527 .expect(1)
3528 .create_async()
3529 .await;
3530 let token = token_mock(&mut server).expect(0).create_async().await;
3532 let blob = keyring_blob(VALID, true);
3533 let (_td, cache) = fixture();
3534
3535 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
3536 .await
3537 .expect("remote path yields a snapshot");
3538
3539 quota.assert_async().await;
3540 plan.assert_async().await;
3541 token.assert_async().await;
3542 let snap = outcome.snapshot;
3543 assert!(!outcome.stale);
3544 assert_eq!(snap.source, AntigravitySource::Remote);
3545 assert_eq!(snap.plan, "Pro");
3546 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
3547 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
3548 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
3549 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
3550 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3551 assert_eq!(snap.account, format!("acct:{fingerprint}"));
3552 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
3553
3554 let cached = parse_cache_at(
3556 &cache
3557 .fresh_payload(Duration::from_secs(60))
3558 .unwrap()
3559 .unwrap(),
3560 None,
3561 now(),
3562 )
3563 .unwrap();
3564 assert_eq!(cached.source, AntigravitySource::Remote);
3565 assert_eq!(cached.account, snap.account);
3566 }
3567
3568 #[tokio::test]
3573 async fn agys_missing_csrf_response_uses_the_saved_session() {
3574 let mut server = mockito::Server::new_async().await;
3575 let eps = endpoints(&server);
3576 let root = server
3577 .mock("GET", "/")
3578 .with_status(404)
3579 .expect(1)
3580 .create_async()
3581 .await;
3582 let status_path = format!("/{STATUS_RPC}");
3583 let status = server
3584 .mock("POST", status_path.as_str())
3585 .with_status(401)
3586 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3587 .expect(1)
3588 .create_async()
3589 .await;
3590 let quota = quota_mock(&mut server, "KEYRING-AT")
3591 .expect(1)
3592 .create_async()
3593 .await;
3594 let plan = server
3595 .mock("POST", "/daily/plan")
3596 .match_header("authorization", "Bearer KEYRING-AT")
3597 .with_status(200)
3598 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3599 .expect(1)
3600 .create_async()
3601 .await;
3602 let blob = keyring_blob(VALID, true);
3603 let (_td, cache) = fixture();
3604
3605 let outcome = run(
3606 &cache,
3607 RemoteOverride {
3608 credential: SavedCredential::Blob(&blob),
3609 endpoints: Some(&eps),
3610 local_bases: Some(vec![server.url().into()]),
3611 ..Default::default()
3612 },
3613 Duration::ZERO,
3614 )
3615 .await
3616 .expect("saved session bypasses agy's unusable local RPC");
3617
3618 root.assert_async().await;
3619 status.assert_async().await;
3620 quota.assert_async().await;
3621 plan.assert_async().await;
3622 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3623 }
3624
3625 #[tokio::test]
3629 async fn agys_prefixed_missing_csrf_response_uses_the_saved_session() {
3630 let mut server = mockito::Server::new_async().await;
3631 let eps = endpoints(&server);
3632 let root = server
3633 .mock("GET", "/")
3634 .with_status(404)
3635 .expect(1)
3636 .create_async()
3637 .await;
3638 let status_path = format!("/{STATUS_RPC}");
3639 let status = server
3640 .mock("POST", status_path.as_str())
3641 .with_status(401)
3642 .with_body(
3643 r#"{"code":"unauthenticated","message":"unauthenticated: missing CSRF token"}"#,
3644 )
3645 .expect(1)
3646 .create_async()
3647 .await;
3648 let quota = quota_mock(&mut server, "KEYRING-AT")
3649 .expect(1)
3650 .create_async()
3651 .await;
3652 let plan = server
3653 .mock("POST", "/daily/plan")
3654 .match_header("authorization", "Bearer KEYRING-AT")
3655 .with_status(200)
3656 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3657 .expect(1)
3658 .create_async()
3659 .await;
3660 let blob = keyring_blob(VALID, true);
3661 let (_td, cache) = fixture();
3662
3663 let outcome = run(
3664 &cache,
3665 RemoteOverride {
3666 credential: SavedCredential::Blob(&blob),
3667 endpoints: Some(&eps),
3668 local_bases: Some(vec![server.url().into()]),
3669 ..Default::default()
3670 },
3671 Duration::ZERO,
3672 )
3673 .await
3674 .expect("saved session bypasses agy 1.3.1's prefixed missing-CSRF RPC");
3675
3676 root.assert_async().await;
3677 status.assert_async().await;
3678 quota.assert_async().await;
3679 plan.assert_async().await;
3680 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3681 }
3682
3683 #[tokio::test]
3684 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
3685 let mut server = mockito::Server::new_async().await;
3686 let eps = endpoints(&server);
3687 let root = server
3688 .mock("GET", "/")
3689 .with_status(404)
3690 .expect(1)
3691 .create_async()
3692 .await;
3693 let status_path = format!("/{STATUS_RPC}");
3694 let status = server
3695 .mock("POST", status_path.as_str())
3696 .with_status(401)
3697 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3698 .expect(1)
3699 .create_async()
3700 .await;
3701 let quota = quota_mock(&mut server, "KEYRING-AT")
3702 .expect(0)
3703 .create_async()
3704 .await;
3705 let (_td, cache) = fixture();
3706
3707 let error = run(
3708 &cache,
3709 RemoteOverride {
3710 credential: SavedCredential::Absent,
3711 endpoints: Some(&eps),
3712 local_bases: Some(vec![server.url().into()]),
3713 ..Default::default()
3714 },
3715 Duration::ZERO,
3716 )
3717 .await
3718 .expect_err("neither the local RPC nor a saved session is usable");
3719
3720 root.assert_async().await;
3721 status.assert_async().await;
3722 quota.assert_async().await;
3723 let message = error.to_string();
3724 assert!(message.contains("requires a CSRF token"), "{message}");
3725 assert!(message.contains("saved Google session"), "{message}");
3726 }
3727
3728 #[tokio::test]
3733 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
3734 let mut server1 = mockito::Server::new_async().await;
3735 let mut server2 = mockito::Server::new_async().await;
3736 let eps = endpoints(&server1);
3737 let root1 = server1
3738 .mock("GET", "/")
3739 .with_status(404)
3740 .expect(1)
3741 .create_async()
3742 .await;
3743 let status_path = format!("/{STATUS_RPC}");
3744 let status1 = server1
3745 .mock("POST", status_path.as_str())
3746 .with_status(401)
3747 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3748 .expect(1)
3749 .create_async()
3750 .await;
3751
3752 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
3754 let status2 = server2
3755 .mock("POST", status_path.as_str())
3756 .expect(0)
3757 .create_async()
3758 .await;
3759
3760 let (_td, cache) = fixture();
3761 let error = run(
3762 &cache,
3763 RemoteOverride {
3764 credential: SavedCredential::Absent,
3765 endpoints: Some(&eps),
3766 local_bases: Some(vec![
3767 Candidate {
3768 base: server1.url(),
3769 pid: Some(42),
3770 ide: false,
3771 },
3772 Candidate {
3773 base: server2.url(),
3774 pid: Some(42),
3775 ide: false,
3776 },
3777 ]),
3778 ..Default::default()
3779 },
3780 Duration::ZERO,
3781 )
3782 .await
3783 .expect_err("aborted probe returns agy missing csrf diagnosis");
3784
3785 root1.assert_async().await;
3786 status1.assert_async().await;
3787 root2.assert_async().await;
3788 status2.assert_async().await;
3789 let message = error.to_string();
3790 assert!(message.contains("requires a CSRF token"), "{message}");
3791 }
3792
3793 #[tokio::test]
3800 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
3801 let mut agy = mockito::Server::new_async().await;
3802 let mut ide = mockito::Server::new_async().await;
3803 let eps = endpoints(&agy);
3804 let agy_root = agy
3805 .mock("GET", "/")
3806 .with_status(404)
3807 .expect(1)
3808 .create_async()
3809 .await;
3810 let status_path = format!("/{STATUS_RPC}");
3811 let agy_status = agy
3812 .mock("POST", status_path.as_str())
3813 .with_status(401)
3814 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3815 .expect(1)
3816 .create_async()
3817 .await;
3818
3819 let ide_root = ide
3820 .mock("GET", "/")
3821 .with_status(404)
3822 .expect(1)
3823 .create_async()
3824 .await;
3825 let ide_status = ide
3826 .mock("POST", status_path.as_str())
3827 .with_status(200)
3828 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3829 .expect(1)
3830 .create_async()
3831 .await;
3832 let quota_path = format!("/{QUOTA_RPC}");
3833 let ide_quota = ide
3834 .mock("POST", quota_path.as_str())
3835 .with_status(200)
3836 .with_body(QUOTA_JSON)
3837 .expect(1)
3838 .create_async()
3839 .await;
3840
3841 let (_td, cache) = fixture();
3842 let outcome = run(
3843 &cache,
3844 RemoteOverride {
3845 credential: SavedCredential::Absent,
3846 endpoints: Some(&eps),
3847 local_bases: Some(vec![
3848 Candidate {
3849 base: agy.url(),
3850 pid: Some(42),
3851 ide: false,
3852 },
3853 Candidate {
3854 base: ide.url(),
3855 pid: Some(7),
3856 ide: false,
3857 },
3858 ]),
3859 ..Default::default()
3860 },
3861 Duration::ZERO,
3862 )
3863 .await
3864 .expect("the other product's server still answers");
3865
3866 agy_root.assert_async().await;
3867 agy_status.assert_async().await;
3868 ide_root.assert_async().await;
3869 ide_status.assert_async().await;
3870 ide_quota.assert_async().await;
3871 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3872 assert_eq!(outcome.snapshot.plan, "Pro");
3873 }
3874
3875 #[tokio::test]
3878 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
3879 let mut server = mockito::Server::new_async().await;
3880 let eps = endpoints(&server);
3881 let _quota = quota_mock(&mut server, "NEW-AT")
3882 .expect(2)
3883 .create_async()
3884 .await;
3885 let token = token_mock(&mut server).expect(1).create_async().await;
3886 let blob = keyring_blob(EXPIRED, true);
3887 let (_td, cache) = fixture();
3888
3889 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
3890 .await
3891 .expect("refresh, then quota");
3892 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
3893 token.assert_async().await;
3894
3895 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3896 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
3897 .expect("refreshed token persisted under the keyring session's fingerprint");
3898 assert_eq!(persisted.access_token, "NEW-AT");
3899
3900 run(&cache, remote(&blob, &eps), Duration::ZERO)
3902 .await
3903 .expect("persisted token reused");
3904 token.assert_async().await;
3905 }
3906
3907 #[tokio::test]
3910 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
3911 let mut server = mockito::Server::new_async().await;
3912 let eps = endpoints(&server);
3913 let rejected = server
3914 .mock("POST", "/daily/quota")
3915 .match_header("authorization", "Bearer KEYRING-AT")
3916 .with_status(401)
3917 .expect(1)
3918 .create_async()
3919 .await;
3920 let accepted = quota_mock(&mut server, "NEW-AT")
3921 .expect(1)
3922 .create_async()
3923 .await;
3924 let token = token_mock(&mut server).expect(1).create_async().await;
3925 let blob = keyring_blob(VALID, true);
3926 let (_td, cache) = fixture();
3927
3928 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3929 .await
3930 .expect("retry with the refreshed token succeeds");
3931
3932 rejected.assert_async().await;
3933 accepted.assert_async().await;
3934 token.assert_async().await;
3935 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3936 }
3937
3938 #[tokio::test]
3941 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
3942 let mut server = mockito::Server::new_async().await;
3943 let eps = endpoints(&server);
3944 let _token = server
3945 .mock("POST", "/token")
3946 .with_status(400)
3947 .with_body(r#"{"error":"invalid_grant"}"#)
3948 .create_async()
3949 .await;
3950 let blob = keyring_blob(EXPIRED, true);
3951 let (_td, cache) = fixture();
3952
3953 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3954 .await
3955 .expect_err("no cache to fall back on");
3956
3957 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3958 let rendered = err.to_string();
3959 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
3960 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
3961 }
3962
3963 #[tokio::test]
3966 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
3967 let mut server = mockito::Server::new_async().await;
3968 let eps = endpoints(&server);
3969 let _quota = server
3970 .mock("POST", "/daily/quota")
3971 .with_status(401)
3972 .create_async()
3973 .await;
3974 let token = token_mock(&mut server).expect(1).create_async().await;
3975 let blob = keyring_blob(EXPIRED, true);
3976 let (_td, cache) = fixture();
3977
3978 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3979 .await
3980 .expect_err("rejected after refresh");
3981
3982 token.assert_async().await;
3983 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3984 assert!(err.to_string().contains("sign in again"), "{err}");
3985 }
3986
3987 #[cfg(unix)]
3991 #[tokio::test]
3992 async fn a_subscription_403_surfaces_the_plan_message_without_a_refresh_or_agy_run() {
3993 let mut server = mockito::Server::new_async().await;
3994 let eps = endpoints(&server);
3995 let quota = server
3996 .mock("POST", "/daily/quota")
3997 .match_header("authorization", "Bearer KEYRING-AT")
3998 .with_status(403)
3999 .with_body(
4000 r#"{"error":{"code":403,"status":"PERMISSION_DENIED","details":
4001 [{"reason":"SUBSCRIPTION_REQUIRED","domain":"cloud.google.com"}]}}"#,
4002 )
4003 .expect(1)
4004 .create_async()
4005 .await;
4006 let prod = server
4007 .mock("POST", "/prod/quota")
4008 .expect(0)
4009 .create_async()
4010 .await;
4011 let token = token_mock(&mut server).expect(0).create_async().await;
4012 let (td, cache) = fixture();
4013 let file = td.path().join("credential");
4014 std::fs::write(&file, keyring_blob(VALID, true)).unwrap();
4015 let runs = td.path().join("runs");
4016 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
4017
4018 let err = run_without_oauth(
4019 &cache,
4020 RemoteOverride {
4021 credential: SavedCredential::File(&file),
4022 endpoints: Some(&eps),
4023 local_bases: Some(vec![]),
4024 agy: super::super::agy::AgyCommand::At(agy),
4025 ..Default::default()
4026 },
4027 )
4028 .await
4029 .expect_err("a plan without quota is an error, not a snapshot");
4030
4031 quota.assert_async().await;
4032 prod.assert_async().await;
4033 token.assert_async().await;
4034 match err {
4035 AppError::Http { status, body } => {
4036 assert_eq!(status, 403);
4037 assert!(body.contains("paid subscription"), "{body}");
4038 }
4039 other => panic!("expected Http, got {other:?}"),
4040 }
4041 assert!(
4042 !runs.exists(),
4043 "a subscription verdict must not spawn `agy`"
4044 );
4045 }
4046
4047 #[tokio::test]
4051 async fn a_non_subscription_403_keeps_the_refresh_and_retry_behavior() {
4052 let mut server = mockito::Server::new_async().await;
4053 let eps = endpoints(&server);
4054 let rejected = server
4055 .mock("POST", "/daily/quota")
4056 .match_header("authorization", "Bearer KEYRING-AT")
4057 .with_status(403)
4058 .with_body(r#"{"error":{"message":"acl detail"}}"#)
4059 .expect(1)
4060 .create_async()
4061 .await;
4062 let retried = server
4063 .mock("POST", "/daily/quota")
4064 .match_header("authorization", "Bearer NEW-AT")
4065 .with_status(403)
4066 .with_body(r#"{"error":{"message":"acl detail"}}"#)
4067 .expect(1)
4068 .create_async()
4069 .await;
4070 let token = token_mock(&mut server).expect(1).create_async().await;
4071 let blob = keyring_blob(VALID, true);
4072 let (_td, cache) = fixture();
4073
4074 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4075 .await
4076 .expect_err("still a rejected session");
4077
4078 rejected.assert_async().await;
4079 retried.assert_async().await;
4080 token.assert_async().await;
4081 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4082 assert!(err.to_string().contains("sign in again"), "{err}");
4083 }
4084
4085 async fn run_without_oauth(cache: &Cache, remote: RemoteOverride<'_>) -> Result<FetchOutcome> {
4087 fetch_snapshot_at(
4088 &reqwest::Client::new(),
4089 cache,
4090 Duration::ZERO,
4091 None,
4092 remote,
4093 now(),
4094 )
4095 .await
4096 }
4097
4098 #[cfg(unix)]
4103 fn fake_agy(
4104 dir: &std::path::Path,
4105 runs: &std::path::Path,
4106 file: &std::path::Path,
4107 blob: &str,
4108 ) -> std::path::PathBuf {
4109 let path = dir.join("agy");
4110 super::super::agy::test_support::write_script(
4111 &path,
4112 &format!(
4113 "[ \"$1\" = -p ] && exit 1\necho \"$@\" >> '{}'\nprintf '%s' '{}' > '{}'",
4114 runs.display(),
4115 blob,
4116 file.display()
4117 ),
4118 );
4119 path
4120 }
4121
4122 #[cfg(unix)]
4125 #[tokio::test]
4126 async fn an_expired_session_is_renewed_by_running_agy() {
4127 let mut server = mockito::Server::new_async().await;
4128 let eps = endpoints(&server);
4129 let quota = quota_mock(&mut server, "KEYRING-AT").create_async().await;
4130 let _plan = server
4131 .mock("POST", "/daily/plan")
4132 .with_status(200)
4133 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
4134 .create_async()
4135 .await;
4136 let (td, cache) = fixture();
4137 let file = td.path().join("credential");
4138 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
4139 let runs = td.path().join("runs");
4140 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
4141
4142 let out = run_without_oauth(
4143 &cache,
4144 RemoteOverride {
4145 credential: SavedCredential::File(&file),
4146 endpoints: Some(&eps),
4147 local_bases: Some(vec![]),
4148 agy: super::super::agy::AgyCommand::At(agy),
4149 ..Default::default()
4150 },
4151 )
4152 .await
4153 .expect("agy renewed the session");
4154
4155 quota.assert_async().await;
4156 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4157 assert_eq!(std::fs::read_to_string(runs).unwrap().trim(), "models");
4158 }
4159
4160 #[cfg(unix)]
4163 #[tokio::test]
4164 async fn agy_that_does_not_renew_keeps_the_original_error_and_runs_once() {
4165 let server = mockito::Server::new_async().await;
4166 let eps = endpoints(&server);
4167 let (td, cache) = fixture();
4168 let file = td.path().join("credential");
4169 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
4170 let runs = td.path().join("runs");
4171 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(EXPIRED, true));
4172 let remote = || RemoteOverride {
4173 credential: SavedCredential::File(&file),
4174 endpoints: Some(&eps),
4175 local_bases: Some(vec![]),
4176 agy: super::super::agy::AgyCommand::At(agy.clone()),
4177 ..Default::default()
4178 };
4179
4180 for _ in 0..2 {
4181 let err = run_without_oauth(&cache, remote())
4182 .await
4183 .expect_err("still expired");
4184 assert!(err.to_string().contains("no OAuth client"), "{err}");
4185 }
4186 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4187 }
4188
4189 #[tokio::test]
4192 async fn without_an_agy_command_an_expired_session_is_left_alone() {
4193 let server = mockito::Server::new_async().await;
4194 let eps = endpoints(&server);
4195 let blob = keyring_blob(EXPIRED, true);
4196 let (_td, cache) = fixture();
4197
4198 let err = run_without_oauth(&cache, remote(&blob, &eps))
4199 .await
4200 .expect_err("nothing renews it");
4201 assert!(err.to_string().contains("no OAuth client"), "{err}");
4202 }
4203
4204 #[cfg(unix)]
4207 const AGY_QUOTA: &str = r#"{"status":"SUCCESS","num_turns":0,"command":{"name":"usage","data":{"groups":[
4208{"name":"Gemini Models","buckets":[
4209 {"id":"gemini-weekly","name":"Weekly Limit Remaining","window":"weekly","remaining_fraction":0.63,"reset_time":"2026-07-28T01:11:19Z"},
4210 {"id":"gemini-5h","name":"Five Hour Limit Remaining","window":"5h","remaining_fraction":0.49,"reset_time":"2026-07-22T15:25:45Z"}]},
4211{"name":"Claude and GPT models","buckets":[
4212 {"id":"3p-weekly","name":"Weekly Limit Remaining","window":"weekly","remaining_fraction":1,"reset_time":"2026-07-29T12:00:00Z"},
4213 {"id":"3p-5h","name":"Five Hour Limit Remaining","window":"5h","remaining_fraction":1,"reset_time":"2026-07-22T17:00:00Z"}]},
4214{"name":"Gemini 4 Argon","buckets":[
4215 {"id":"gemini-argon-wtu-weekly","name":"Gemini 4 Argon Weekly Limit","window":"weekly","remaining_fraction":0,"reset_time":"2026-07-28T01:10:47Z"}]}]}}}"#;
4216
4217 #[cfg(unix)]
4220 fn agy_reading(dir: &std::path::Path, runs: &std::path::Path) -> std::path::PathBuf {
4221 let path = dir.join("agy");
4222 super::super::agy::test_support::write_script(
4223 &path,
4224 &format!(
4225 "echo \"$@\" >> '{}'\n[ \"$1\" = -p ] || exit 1\ncat <<'EOF'\n{AGY_QUOTA}\nEOF",
4226 runs.display()
4227 ),
4228 );
4229 path
4230 }
4231
4232 #[cfg(unix)]
4236 #[tokio::test]
4237 async fn agys_own_reading_answers_before_the_cloud_code_api() {
4238 let mut server = mockito::Server::new_async().await;
4239 let eps = endpoints(&server);
4240 let daily = server
4241 .mock("POST", "/daily/quota")
4242 .expect(0)
4243 .create_async()
4244 .await;
4245 let prod = server
4246 .mock("POST", "/prod/quota")
4247 .expect(0)
4248 .create_async()
4249 .await;
4250 let plan = server
4251 .mock("POST", "/daily/plan")
4252 .match_header("authorization", "Bearer KEYRING-AT")
4253 .with_status(200)
4254 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4255 .expect(1)
4256 .create_async()
4257 .await;
4258 let blob = keyring_blob(VALID, true);
4259 let (td, cache) = fixture();
4260 let runs = td.path().join("runs");
4261 let agy = agy_reading(td.path(), &runs);
4262
4263 let out = run(
4264 &cache,
4265 RemoteOverride {
4266 agy: super::super::agy::AgyCommand::At(agy),
4267 ..remote(&blob, &eps)
4268 },
4269 Duration::from_secs(60),
4270 )
4271 .await
4272 .expect("agy answered");
4273
4274 daily.assert_async().await;
4275 prod.assert_async().await;
4276 plan.assert_async().await;
4277 let snap = out.snapshot;
4278 assert_eq!(snap.source, AntigravitySource::Agy);
4279 assert_eq!(snap.plan, "Pro");
4280 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 51);
4281 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 37);
4282 assert_eq!(
4283 snap.third_party_session.as_ref().unwrap().utilization_pct,
4284 0
4285 );
4286 assert_eq!(snap.model_pools[0].window.utilization_pct, 100);
4287 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
4288 assert_eq!(snap.account, format!("acct:{fingerprint}"));
4289 let cached = parse_cache_at(
4290 &cache
4291 .fresh_payload(Duration::from_secs(60))
4292 .unwrap()
4293 .unwrap(),
4294 None,
4295 now(),
4296 )
4297 .unwrap();
4298 assert_eq!(cached.source, AntigravitySource::Agy);
4299 assert_eq!(
4300 std::fs::read_to_string(runs).unwrap().trim(),
4301 "-p /quota --output-format json --log-file /dev/null"
4302 );
4303 }
4304
4305 #[cfg(unix)]
4308 #[tokio::test]
4309 async fn agys_reading_needs_no_good_token() {
4310 let mut server = mockito::Server::new_async().await;
4311 let eps = endpoints(&server);
4312 let plan = server
4313 .mock("POST", "/daily/plan")
4314 .expect(0)
4315 .create_async()
4316 .await;
4317 let token = token_mock(&mut server).expect(0).create_async().await;
4318 let blob = keyring_blob(EXPIRED, true);
4319 let (td, cache) = fixture();
4320 let runs = td.path().join("runs");
4321 let agy = agy_reading(td.path(), &runs);
4322
4323 let out = run_without_oauth(
4324 &cache,
4325 RemoteOverride {
4326 agy: super::super::agy::AgyCommand::At(agy),
4327 ..remote(&blob, &eps)
4328 },
4329 )
4330 .await
4331 .expect("agy answered");
4332
4333 plan.assert_async().await;
4334 token.assert_async().await;
4335 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4336 assert_eq!(out.snapshot.plan, "Antigravity");
4337 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4338 }
4339
4340 #[cfg(unix)]
4343 #[tokio::test]
4344 async fn when_agy_cannot_read_the_quota_the_cloud_code_api_answers() {
4345 let mut server = mockito::Server::new_async().await;
4346 let eps = endpoints(&server);
4347 let quota = quota_mock(&mut server, "KEYRING-AT")
4348 .expect(2)
4349 .create_async()
4350 .await;
4351 let blob = keyring_blob(VALID, true);
4352 let (td, cache) = fixture();
4353 let runs = td.path().join("runs");
4354 let agy = td.path().join("agy");
4355 super::super::agy::test_support::write_script(
4356 &agy,
4357 &format!("echo \"$@\" >> '{}'\nexit 1", runs.display()),
4358 );
4359 let remote = || RemoteOverride {
4360 agy: super::super::agy::AgyCommand::At(agy.clone()),
4361 ..remote(&blob, &eps)
4362 };
4363
4364 for _ in 0..2 {
4365 let out = run(&cache, remote(), Duration::ZERO)
4366 .await
4367 .expect("the Cloud Code API answered");
4368 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4369 std::fs::remove_file(cache.payload_path()).unwrap();
4371 }
4372 quota.assert_async().await;
4373 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4374 }
4375
4376 #[cfg(unix)]
4380 fn age_past_reading_floor(cache: &Cache) {
4381 let aged = std::time::SystemTime::now() - super::super::agy::QUOTA_INTERVAL * 2;
4382 std::fs::OpenOptions::new()
4383 .write(true)
4384 .open(cache.payload_path())
4385 .unwrap()
4386 .set_times(std::fs::FileTimes::new().set_modified(aged))
4387 .unwrap();
4388 }
4389
4390 #[cfg(unix)]
4395 #[tokio::test]
4396 async fn when_agy_stops_answering_its_last_reading_is_kept_over_the_cloud_code_api() {
4397 let mut server = mockito::Server::new_async().await;
4398 let eps = endpoints(&server);
4399 let quota = quota_mock(&mut server, "KEYRING-AT")
4400 .expect(0)
4401 .create_async()
4402 .await;
4403 let _plan = server
4404 .mock("POST", "/daily/plan")
4405 .with_status(200)
4406 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4407 .create_async()
4408 .await;
4409 let blob = keyring_blob(VALID, true);
4410 let (td, cache) = fixture();
4411 let runs = td.path().join("runs");
4412 let agy = agy_reading(td.path(), &runs);
4413 let remote = || RemoteOverride {
4414 agy: super::super::agy::AgyCommand::At(agy.clone()),
4415 ..remote(&blob, &eps)
4416 };
4417
4418 let first = run(&cache, remote(), Duration::ZERO)
4419 .await
4420 .expect("agy answered");
4421 assert_eq!(first.snapshot.source, AntigravitySource::Agy);
4422 assert!(!first.stale);
4423
4424 age_past_reading_floor(&cache);
4426 super::super::agy::test_support::write_script(
4427 &agy,
4428 &format!("echo \"$@\" >> '{}'\nexit 1", runs.display()),
4429 );
4430
4431 for _ in 0..2 {
4432 let out = run(&cache, remote(), Duration::ZERO)
4433 .await
4434 .expect("agy's reading is kept");
4435 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4436 assert_eq!(out.snapshot.weekly.as_ref().unwrap().utilization_pct, 37);
4437 assert_eq!(out.snapshot.plan, "Pro");
4438 assert!(out.stale, "a kept reading is not passed off as current");
4439 age_past_reading_floor(&cache);
4440 }
4441 quota.assert_async().await;
4442 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 2);
4445 }
4446
4447 #[cfg(unix)]
4450 #[tokio::test]
4451 async fn a_cloud_code_figure_is_not_kept_in_agys_place() {
4452 let mut server = mockito::Server::new_async().await;
4453 let eps = endpoints(&server);
4454 let quota = quota_mock(&mut server, "KEYRING-AT")
4455 .expect(2)
4456 .create_async()
4457 .await;
4458 let blob = keyring_blob(VALID, true);
4459 let (td, cache) = fixture();
4460 let agy = td.path().join("agy");
4461 super::super::agy::test_support::write_script(&agy, "exit 1");
4462 let remote = || RemoteOverride {
4463 agy: super::super::agy::AgyCommand::At(agy.clone()),
4464 ..remote(&blob, &eps)
4465 };
4466
4467 for _ in 0..2 {
4468 let out = run(&cache, remote(), Duration::ZERO)
4469 .await
4470 .expect("the Cloud Code API answered");
4471 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4472 assert!(!out.stale);
4473 age_past_reading_floor(&cache);
4474 }
4475 quota.assert_async().await;
4476 }
4477
4478 #[cfg(unix)]
4481 #[tokio::test]
4482 async fn agys_reading_is_taken_at_most_every_quota_interval() {
4483 let server = mockito::Server::new_async().await;
4484 let eps = endpoints(&server);
4485 let blob = keyring_blob(VALID, true);
4486 let (td, cache) = fixture();
4487 let runs = td.path().join("runs");
4488 let agy = agy_reading(td.path(), &runs);
4489 for _ in 0..2 {
4490 let out = run(
4491 &cache,
4492 RemoteOverride {
4493 agy: super::super::agy::AgyCommand::At(agy.clone()),
4494 ..remote(&blob, &eps)
4495 },
4496 Duration::ZERO,
4497 )
4498 .await
4499 .expect("a snapshot");
4500 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4501 }
4502 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4503 }
4504
4505 #[tokio::test]
4507 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
4508 let mut server = mockito::Server::new_async().await;
4509 let eps = endpoints(&server);
4510 let token = token_mock(&mut server).expect(0).create_async().await;
4511 let blob = keyring_blob(EXPIRED, false);
4512 let (_td, cache) = fixture();
4513
4514 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4515 .await
4516 .expect_err("nothing to refresh with");
4517
4518 token.assert_async().await;
4519 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4520 assert!(err.to_string().contains("sign in again"), "{err}");
4521 }
4522
4523 #[tokio::test]
4526 async fn no_saved_session_extends_the_no_local_server_error() {
4527 let mut server = mockito::Server::new_async().await;
4528 let eps = endpoints(&server);
4529 let quota = quota_mock(&mut server, "KEYRING-AT")
4530 .expect(0)
4531 .create_async()
4532 .await;
4533 let (_td, cache) = fixture();
4534
4535 let err = run(
4536 &cache,
4537 RemoteOverride {
4538 credential: SavedCredential::Absent,
4539 endpoints: Some(&eps),
4540 local_bases: Some(vec![]),
4541 ..Default::default()
4542 },
4543 Duration::ZERO,
4544 )
4545 .await
4546 .expect_err("no source at all");
4547
4548 quota.assert_async().await;
4549 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4550 let rendered = err.to_string();
4551 assert!(rendered.contains("no local server found"), "{rendered}");
4552 assert!(rendered.contains("saved Google session"), "{rendered}");
4553 }
4554
4555 #[tokio::test]
4558 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
4559 let mut server = mockito::Server::new_async().await;
4560 let eps = endpoints(&server);
4561 let status_path = format!("/{STATUS_RPC}");
4562 let _status = server
4563 .mock("POST", status_path.as_str())
4564 .with_status(401)
4565 .create_async()
4566 .await;
4567 let quota = quota_mock(&mut server, "KEYRING-AT")
4568 .expect(0)
4569 .create_async()
4570 .await;
4571 let blob = keyring_blob(VALID, true);
4572 let (_td, cache) = fixture();
4573
4574 let err = run(
4575 &cache,
4576 RemoteOverride {
4577 credential: SavedCredential::Blob(&blob),
4578 endpoints: Some(&eps),
4579 local_bases: Some(vec![server.url().into()]),
4580 ..Default::default()
4581 },
4582 Duration::ZERO,
4583 )
4584 .await
4585 .expect_err("signed-out local server");
4586
4587 quota.assert_async().await;
4588 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
4589 }
4590
4591 #[tokio::test]
4594 async fn a_403_auth_rejection_does_not_fall_back_to_a_warm_cache() {
4595 let mut server = mockito::Server::new_async().await;
4596 let eps = endpoints(&server);
4597 let status_path = format!("/{STATUS_RPC}");
4598 let _status = server
4599 .mock("POST", status_path.as_str())
4600 .with_status(403)
4601 .with_body(r#"{"error":"forbidden"}"#)
4602 .create_async()
4603 .await;
4604 let quota = quota_mock(&mut server, "KEYRING-AT")
4605 .expect(0)
4606 .create_async()
4607 .await;
4608 let blob = keyring_blob(VALID, true);
4609 let (_td, cache) = fixture();
4610
4611 let mut warm_snap = parsed();
4612 warm_snap.session = None;
4613 if let Some(w) = warm_snap.weekly.as_mut() {
4614 w.utilization_pct = 100;
4615 }
4616 if let Some(w) = warm_snap.third_party_weekly.as_mut() {
4617 w.utilization_pct = 100;
4618 }
4619 let bytes = serde_json::to_vec(&snap_to_json(&warm_snap)).unwrap();
4620 cache.write_payload(&bytes).unwrap();
4621
4622 let err = run(
4623 &cache,
4624 RemoteOverride {
4625 credential: SavedCredential::Blob(&blob),
4626 endpoints: Some(&eps),
4627 local_bases: Some(vec![server.url().into()]),
4628 ..Default::default()
4629 },
4630 Duration::ZERO,
4631 )
4632 .await
4633 .expect_err("HTTP 403 must not fall back to warm cache");
4634
4635 quota.assert_async().await;
4636 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
4637 assert_eq!(cache.read_last_error().map(|(code, _)| code), Some(403));
4638 }
4639
4640 #[tokio::test]
4646 async fn a_signed_out_ide_language_server_is_skipped_so_the_saved_session_answers() {
4647 let mut ide = mockito::Server::new_async().await;
4648 let mut agy = mockito::Server::new_async().await;
4649 let eps = endpoints(&agy);
4650 let ide_root = ide.mock("GET", "/").expect(0).create_async().await;
4652 let status_path = format!("/{STATUS_RPC}");
4653 let ide_status = ide
4654 .mock("POST", status_path.as_str())
4655 .with_status(403)
4656 .expect(0)
4657 .create_async()
4658 .await;
4659 let agy_root = agy
4661 .mock("GET", "/")
4662 .with_status(404)
4663 .expect(1)
4664 .create_async()
4665 .await;
4666 let agy_status = agy
4667 .mock("POST", status_path.as_str())
4668 .with_status(401)
4669 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
4670 .expect(1)
4671 .create_async()
4672 .await;
4673 let quota = quota_mock(&mut agy, "KEYRING-AT")
4674 .expect(1)
4675 .create_async()
4676 .await;
4677 let plan = agy
4678 .mock("POST", "/daily/plan")
4679 .match_header("authorization", "Bearer KEYRING-AT")
4680 .with_status(200)
4681 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4682 .expect(1)
4683 .create_async()
4684 .await;
4685 let blob = keyring_blob(VALID, true);
4686 let (_td, cache) = fixture();
4687
4688 let outcome = run(
4689 &cache,
4690 RemoteOverride {
4691 credential: SavedCredential::Blob(&blob),
4692 endpoints: Some(&eps),
4693 local_bases: Some(vec![
4694 Candidate {
4695 base: ide.url(),
4696 pid: Some(7),
4697 ide: true,
4698 },
4699 Candidate {
4700 base: agy.url(),
4701 pid: Some(42),
4702 ide: false,
4703 },
4704 ]),
4705 ide_language_server: false,
4706 ..Default::default()
4707 },
4708 Duration::ZERO,
4709 )
4710 .await
4711 .expect("with the IDE skipped, agy's missing-CSRF answer unlocks the saved session");
4712
4713 ide_root.assert_async().await;
4714 ide_status.assert_async().await;
4715 agy_root.assert_async().await;
4716 agy_status.assert_async().await;
4717 quota.assert_async().await;
4718 plan.assert_async().await;
4719 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4720 assert_eq!(outcome.snapshot.plan, "Pro");
4721 }
4722
4723 #[tokio::test]
4727 async fn skipping_the_only_ide_candidate_falls_through_to_the_saved_session() {
4728 let mut ide = mockito::Server::new_async().await;
4729 let eps = endpoints(&ide);
4730 let ide_root = ide.mock("GET", "/").expect(0).create_async().await;
4731 let status_path = format!("/{STATUS_RPC}");
4732 let ide_status = ide
4733 .mock("POST", status_path.as_str())
4734 .expect(0)
4735 .create_async()
4736 .await;
4737 let quota = quota_mock(&mut ide, "KEYRING-AT")
4738 .expect(1)
4739 .create_async()
4740 .await;
4741 let plan = ide
4742 .mock("POST", "/daily/plan")
4743 .with_status(200)
4744 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
4745 .expect(1)
4746 .create_async()
4747 .await;
4748 let blob = keyring_blob(VALID, true);
4749 let (_td, cache) = fixture();
4750
4751 let outcome = run(
4752 &cache,
4753 RemoteOverride {
4754 credential: SavedCredential::Blob(&blob),
4755 endpoints: Some(&eps),
4756 local_bases: Some(vec![Candidate {
4757 base: ide.url(),
4758 pid: Some(7),
4759 ide: true,
4760 }]),
4761 ide_language_server: false,
4762 ..Default::default()
4763 },
4764 Duration::ZERO,
4765 )
4766 .await
4767 .expect("an IDE-only machine with the IDE skipped uses the saved session");
4768
4769 ide_root.assert_async().await;
4770 ide_status.assert_async().await;
4771 quota.assert_async().await;
4772 plan.assert_async().await;
4773 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4774 }
4775
4776 #[tokio::test]
4779 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
4780 let mut server = mockito::Server::new_async().await;
4781 let eps = endpoints(&server);
4782 for path in ["/daily/quota", "/prod/quota"] {
4783 server
4784 .mock("POST", path)
4785 .with_status(500)
4786 .with_body("boom")
4787 .create_async()
4788 .await;
4789 }
4790 let blob = keyring_blob(VALID, true);
4791 let (_td, cache) = fixture();
4792 cache.ensure_dir().unwrap();
4793 let mut earlier = parsed();
4794 earlier.source = AntigravitySource::Remote;
4795 earlier.account = "acct:earlier".into();
4796 cache
4797 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
4798 .unwrap();
4799
4800 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
4801 .await
4802 .expect("stale cache stands in");
4803
4804 assert!(outcome.stale);
4805 assert_eq!(outcome.snapshot, earlier);
4806 assert!(
4807 matches!(outcome.last_error, Some((500, _))),
4808 "{:?}",
4809 outcome.last_error
4810 );
4811 }
4812
4813 fn argon(pct: i32, reset: &str) -> ModelPoolWindow {
4814 ModelPoolWindow {
4815 name: "Gemini 4 Argon".into(),
4816 window: UsageWindow {
4817 utilization_pct: pct,
4818 resets_at: Some(
4819 DateTime::parse_from_rfc3339(reset)
4820 .unwrap()
4821 .with_timezone(&Utc),
4822 ),
4823 window_duration: chrono::Duration::days(7),
4824 },
4825 }
4826 }
4827
4828 #[test]
4829 fn model_pools_round_trip_through_the_cache() {
4830 let mut snap = parsed();
4831 snap.model_pools.push(argon(60, "2026-07-28T17:00:00Z"));
4832 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4833 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
4834
4835 let mut legacy = snap_to_json(&parsed());
4837 legacy.as_object_mut().unwrap().remove("model_pools");
4838 let legacy = serde_json::to_vec(&legacy).unwrap();
4839 assert!(
4840 parse_cache_at(&legacy, None, now())
4841 .unwrap()
4842 .model_pools
4843 .is_empty()
4844 );
4845
4846 let mut torn = snap_to_json(&snap);
4848 torn["model_pools"][0]["window"] = "monthly".into();
4849 let torn = serde_json::to_vec(&torn).unwrap();
4850 assert!(parse_cache_at(&torn, None, now()).is_err());
4851
4852 let mut damaged = snap_to_json(&snap);
4855 damaged["model_pools"][0]["name"] = "Argon\u{001b}[31m".into();
4856 let damaged = serde_json::to_vec(&damaged).unwrap();
4857 let restored = parse_cache_at(&damaged, None, now()).unwrap();
4858 assert_eq!(restored.model_pools[0].name, "Argon[31m");
4859 }
4860
4861 #[test]
4863 fn a_cached_model_pool_past_its_reset_refetches() {
4864 let mut snap = parsed();
4865 snap.model_pools.push(argon(60, "2026-07-22T11:00:00Z"));
4866 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4867 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
4868 assert!(err.to_string().contains("Gemini 4 Argon weekly"), "{err}");
4869 }
4870
4871 #[test]
4872 fn source_round_trips_through_the_cache_and_defaults_to_local() {
4873 let mut snap = parsed();
4874 snap.source = AntigravitySource::Remote;
4875 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4876 assert_eq!(
4877 parse_cache_at(&bytes, None, now()).unwrap().source,
4878 AntigravitySource::Remote
4879 );
4880
4881 snap.source = AntigravitySource::Statusline;
4882 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4883 assert_eq!(
4884 parse_cache_at(&bytes, None, now()).unwrap().source,
4885 AntigravitySource::Statusline
4886 );
4887
4888 snap.source = AntigravitySource::Agy;
4889 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4890 assert_eq!(
4891 parse_cache_at(&bytes, None, now()).unwrap().source,
4892 AntigravitySource::Agy
4893 );
4894
4895 let mut legacy = snap_to_json(&snap);
4897 legacy.as_object_mut().unwrap().remove("source");
4898 let legacy = serde_json::to_vec(&legacy).unwrap();
4899 assert_eq!(
4900 parse_cache_at(&legacy, None, now()).unwrap().source,
4901 AntigravitySource::Local
4902 );
4903 }
4904}