1use std::time::Duration;
40
41use chrono::{DateTime, Utc};
42
43use super::cloud;
44use super::credential::{self, StoredToken};
45use crate::cache::{Cache, acquire_lock_async};
46use crate::error::{AppError, Result};
47use crate::usage::{AntigravitySnapshot, AntigravitySource, ModelPoolWindow, UsageWindow};
48
49const HTTP_TIMEOUT: Duration = Duration::from_secs(5);
50const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
51
52const QUOTA_RPC: &str = "exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary";
53const STATUS_RPC: &str = "exa.language_server_pb.LanguageServerService/GetUserStatus";
54
55const DEFAULT_PLAN: &str = "Antigravity";
56
57const NO_LOCAL_SERVER: &str = "Antigravity: no local server found. Quota is only served while \
58 Antigravity is running — open the Antigravity app, or an interactive \
59 `agy` session, or point ANTIGRAVITY_LS_ADDRESS at a host:port.";
60
61const AGY_CSRF_UNAVAILABLE: &str = "Antigravity: the running `agy` server requires a CSRF token \
62 that it does not publish.";
63
64const NO_SAVED_SESSION: &str = "Or sign in to Antigravity once, so its saved Google session can \
68 be used while it is closed.";
69
70const SESSION_EXPIRED: &str =
71 "Antigravity's saved Google session expired; open Antigravity to sign in again";
72
73pub type FetchOutcome = crate::outcome::Outcome<AntigravitySnapshot>;
76
77impl From<FetchOutcome> for crate::vendor::VendorOutcome {
78 fn from(o: FetchOutcome) -> Self {
79 o.map(crate::usage::VendorSnapshot::Antigravity)
80 }
81}
82
83#[derive(Debug, Clone, Copy, Default)]
87pub enum SavedCredential<'a> {
88 #[default]
90 Keyring,
91 Blob(&'a str),
93 File(&'a std::path::Path),
96 Absent,
98}
99
100pub struct RemoteOverride<'a> {
105 pub credential: SavedCredential<'a>,
107 pub endpoints: Option<&'a cloud::Endpoints>,
109 pub local_bases: Option<Vec<Candidate>>,
114 pub agy: super::agy::AgyCommand,
118 pub ide_language_server: bool,
122}
123
124impl Default for RemoteOverride<'_> {
125 fn default() -> Self {
126 Self {
127 credential: SavedCredential::default(),
128 endpoints: None,
129 local_bases: None,
130 agy: super::agy::AgyCommand::default(),
131 ide_language_server: true,
132 }
133 }
134}
135
136#[derive(Debug, Clone, PartialEq, Eq)]
141pub struct Candidate {
142 pub base: String,
144 pub pid: Option<u32>,
147 pub ide: bool,
154}
155
156impl From<String> for Candidate {
157 fn from(base: String) -> Self {
158 Self {
159 base,
160 pid: None,
161 ide: false,
162 }
163 }
164}
165
166pub async fn fetch_snapshot(
169 client: &reqwest::Client,
170 cache: &Cache,
171 cache_ttl: Duration,
172 oauth: Option<&cloud::OauthClient>,
173 agy_binary: Option<&std::path::Path>,
174 ide_language_server: bool,
175) -> Result<FetchOutcome> {
176 fetch_snapshot_at(
177 client,
178 cache,
179 cache_ttl,
180 oauth,
181 RemoteOverride {
182 agy: agy_binary.map_or(super::agy::AgyCommand::Discover, |path| {
183 super::agy::AgyCommand::At(path.to_path_buf())
184 }),
185 ide_language_server,
186 ..Default::default()
187 },
188 Utc::now(),
189 )
190 .await
191}
192
193pub async fn fetch_snapshot_at(
197 client: &reqwest::Client,
198 cache: &Cache,
199 cache_ttl: Duration,
200 oauth: Option<&cloud::OauthClient>,
201 remote: RemoteOverride<'_>,
202 now: DateTime<Utc>,
203) -> Result<FetchOutcome> {
204 cache.ensure_dir()?;
205 let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
206
207 let origin = match open_session(
217 client,
218 remote.local_bases.as_deref(),
219 remote.ide_language_server,
220 )
221 .await
222 {
223 Ok(session) => Origin::Local(Ok(session)),
224 Err(error) => match remote_fallback_reason(&error) {
225 Some(reason) => Origin::Remote(saved_session(remote.credential, reason)),
226 None => Origin::Local(Err(error)),
227 },
228 };
229 let account = origin.account();
230
231 let cache_ttl = match origin {
234 Origin::Remote(Ok(_)) if remote.agy.is_available() => {
235 cache_ttl.max(super::agy::QUOTA_INTERVAL)
236 }
237 _ => cache_ttl,
238 };
239 if let Some(bytes) = cache.fresh_payload(cache_ttl)?
240 && let Ok(outcome) = reuse_cache(bytes, cache, false, account.as_deref(), now)
241 {
242 return Ok(outcome);
243 }
244
245 let default_endpoints = cloud::Endpoints::default();
246 let endpoints = remote.endpoints.unwrap_or(&default_endpoints);
247 let live = match origin {
248 Origin::Local(session) => fetch_live(client, session).await,
249 Origin::Remote(token) => {
250 let renewal = Renewal {
251 agy: &remote.agy,
252 credential: remote.credential,
253 };
254 fetch_remote(client, cache, oauth, endpoints, token, renewal, now).await
255 }
256 };
257
258 match live {
259 Ok(snap) => {
260 let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
261 cache.write_payload(&bytes)?;
262 Ok(crate::outcome::Outcome::fresh(snap))
263 }
264 Err(e) if e.is_transient() => fallback_silent(cache, now, e),
265 Err(AppError::Http { status, body }) => {
266 cache.mark_stale();
267 let last_error = Some(cache.write_last_error(status, &body));
268 let reason = AppError::Http { status, body };
269 fallback_with_error(cache, last_error, reason, now)
270 }
271 Err(e) => {
272 cache.mark_stale();
273 let last_error = Some(cache.write_last_error(0, &e.to_string()));
274 fallback_with_error(cache, last_error, e, now)
275 }
276 }
277}
278
279struct Session {
282 base: String,
283 csrf: Option<String>,
284 plan: String,
285 account: String,
286}
287
288enum Origin {
291 Local(Result<Session>),
292 Remote(Result<StoredToken>),
293}
294
295impl Origin {
296 fn account(&self) -> Option<String> {
298 match self {
299 Origin::Local(Ok(session)) => Some(session.account.clone()),
300 Origin::Remote(Ok(token)) => Some(remote_account(&token.fingerprint)),
301 Origin::Local(Err(_)) | Origin::Remote(Err(_)) => None,
302 }
303 }
304}
305
306fn no_local_server() -> AppError {
307 AppError::Credentials(NO_LOCAL_SERVER.into())
308}
309
310#[derive(Clone, Copy)]
311enum RemoteFallbackReason {
312 NoLocalServer,
313 AgyMissingCsrf,
314}
315
316impl RemoteFallbackReason {
317 fn message(self) -> &'static str {
318 match self {
319 Self::NoLocalServer => NO_LOCAL_SERVER,
320 Self::AgyMissingCsrf => AGY_CSRF_UNAVAILABLE,
321 }
322 }
323}
324
325fn remote_fallback_reason(error: &AppError) -> Option<RemoteFallbackReason> {
331 if matches!(error, AppError::Credentials(message) if message == NO_LOCAL_SERVER) {
332 Some(RemoteFallbackReason::NoLocalServer)
333 } else if is_missing_csrf(error) {
334 Some(RemoteFallbackReason::AgyMissingCsrf)
335 } else {
336 None
337 }
338}
339
340async fn open_session(
351 client: &reqwest::Client,
352 bases: Option<&[Candidate]>,
353 ide_language_server: bool,
354) -> Result<Session> {
355 let candidates = bases.map_or_else(
356 || candidate_bases(ide_language_server),
357 <[Candidate]>::to_vec,
358 );
359 let candidates: Vec<Candidate> = candidates.into_iter().filter(|c| !c.ide).collect();
360 if candidates.is_empty() {
361 return Err(no_local_server());
362 }
363
364 let mut errors = Vec::new();
365 let mut blocked_pids = std::collections::HashSet::new();
370 for candidate in candidates {
371 if candidate.pid.is_some_and(|pid| blocked_pids.contains(&pid)) {
372 continue;
373 }
374 let csrf = fetch_csrf(client, &candidate.base).await;
375 match post_rpc(client, &candidate.base, csrf.as_deref(), STATUS_RPC).await {
376 Ok(v) => {
377 return Ok(Session {
378 base: candidate.base,
379 csrf,
380 plan: plan_from_status(&v),
381 account: account_key(&v),
382 });
383 }
384 Err(e) => {
385 if is_missing_csrf(&e)
389 && let Some(pid) = candidate.pid
390 {
391 blocked_pids.insert(pid);
392 }
393 errors.push(e);
394 }
395 }
396 }
397 Err(select_probe_error(errors))
398}
399
400fn select_probe_error(errors: Vec<AppError>) -> AppError {
417 let mut actionable = None;
418 let mut missing_csrf = None;
419 let mut last = None;
420 let mut echo = None;
421 for e in errors {
422 if missing_csrf.is_none() && is_missing_csrf(&e) {
423 missing_csrf = Some(e);
424 } else if actionable.is_none() && is_actionable(&e) {
425 actionable = Some(e);
426 } else if is_tls_echo(&e) {
427 echo = Some(e);
428 } else {
429 last = Some(e);
430 }
431 }
432 actionable
433 .or(missing_csrf)
434 .or(last)
435 .or(echo)
436 .unwrap_or_else(|| {
437 AppError::Other("antigravity: no local server answered GetUserStatus".into())
438 })
439}
440
441fn is_missing_csrf(error: &AppError) -> bool {
447 let AppError::Http { status: 401, body } = error else {
448 return false;
449 };
450 let Ok(body) = serde_json::from_str::<serde_json::Value>(body) else {
451 return false;
452 };
453 matches!(
454 (body["code"].as_str(), body["message"].as_str()),
455 (Some(code), Some(message))
456 if code.eq_ignore_ascii_case("unauthenticated")
457 && is_missing_csrf_message(message)
458 )
459}
460
461fn is_missing_csrf_message(message: &str) -> bool {
464 message
465 .rsplit(':')
466 .next()
467 .is_some_and(|part| part.trim().eq_ignore_ascii_case("missing CSRF token"))
468}
469
470fn is_actionable(e: &AppError) -> bool {
475 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
476 && !is_missing_csrf(e)
477}
478
479fn is_tls_echo(e: &AppError) -> bool {
492 matches!(
493 e,
494 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
495 )
496}
497
498async fn fetch_live(
499 client: &reqwest::Client,
500 session: Result<Session>,
501) -> Result<AntigravitySnapshot> {
502 let session = session?;
503 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
504 let mut snap = parse_quota_summary("a, session.plan)?;
505 snap.account = session.account;
506 Ok(snap)
507}
508
509fn saved_session(
519 credential: SavedCredential<'_>,
520 reason: RemoteFallbackReason,
521) -> Result<StoredToken> {
522 let raw = match credential {
523 SavedCredential::Keyring => credential::read()?,
524 SavedCredential::Blob(blob) => Some(blob.to_string()),
525 SavedCredential::File(path) => std::fs::read_to_string(path).ok(),
526 SavedCredential::Absent => None,
527 };
528 let Some(raw) = raw else {
529 return Err(AppError::Credentials(format!(
530 "{} {NO_SAVED_SESSION}",
531 reason.message()
532 )));
533 };
534 credential::parse_keyring_blob(&raw)
535}
536
537fn remote_account(fingerprint: &str) -> String {
541 format!("acct:{fingerprint}")
542}
543
544fn session_expired() -> AppError {
545 AppError::Credentials(SESSION_EXPIRED.into())
546}
547
548fn refresh_unconfigured() -> AppError {
552 AppError::Credentials(
553 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
554 refresh it; open Antigravity (or run `agy`) to sign in again, or set [antigravity] oauth_client_id \
555 and oauth_client_secret in config.toml"
556 .into(),
557 )
558}
559
560fn is_auth_rejection(e: &AppError) -> bool {
561 matches!(
562 e,
563 AppError::Http {
564 status: 401 | 403,
565 ..
566 }
567 )
568}
569
570struct AccessToken {
574 value: String,
575 just_refreshed: bool,
576}
577
578fn current_access_token(
583 oauth_path: &std::path::Path,
584 token: &StoredToken,
585 now: DateTime<Utc>,
586) -> Option<String> {
587 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
588 && token
589 .expires_at
590 .is_none_or(|keyring| persisted.expires_at > keyring)
591 && !cloud::needs_refresh(Some(persisted.expires_at), now)
592 {
593 return Some(persisted.access_token);
594 }
595 (!cloud::needs_refresh(token.expires_at, now)).then(|| token.access_token.clone())
596}
597
598async fn resolve_access_token(
600 client: &reqwest::Client,
601 oauth: Option<&cloud::OauthClient>,
602 endpoints: &cloud::Endpoints,
603 oauth_path: &std::path::Path,
604 token: &StoredToken,
605 now: DateTime<Utc>,
606) -> Result<AccessToken> {
607 if let Some(value) = current_access_token(oauth_path, token, now) {
608 return Ok(AccessToken {
609 value,
610 just_refreshed: false,
611 });
612 }
613 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
614}
615
616async fn refresh_and_persist(
620 client: &reqwest::Client,
621 oauth: Option<&cloud::OauthClient>,
622 endpoints: &cloud::Endpoints,
623 oauth_path: &std::path::Path,
624 token: &StoredToken,
625) -> Result<AccessToken> {
626 let Some(refresh_token) = token.refresh_token.as_deref() else {
627 return Err(session_expired());
628 };
629 let Some(oauth) = oauth else {
630 return Err(refresh_unconfigured());
631 };
632 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
633 cloud::write_persisted(
634 oauth_path,
635 &cloud::PersistedOAuth {
636 fingerprint: token.fingerprint.clone(),
637 access_token: refreshed.access_token.clone(),
638 expires_at: refreshed.expires_at,
639 },
640 )?;
641 Ok(AccessToken {
642 value: refreshed.access_token,
643 just_refreshed: true,
644 })
645}
646
647struct Renewal<'a> {
650 agy: &'a super::agy::AgyCommand,
651 credential: SavedCredential<'a>,
652}
653
654impl Renewal<'_> {
655 async fn renewed(&self, cache: &Cache, now: DateTime<Utc>) -> Option<StoredToken> {
659 if !super::agy::renew_session(self.agy, cache, now).await {
660 return None;
661 }
662 let raw = match self.credential {
663 SavedCredential::Keyring => credential::read().ok().flatten()?,
664 SavedCredential::Blob(blob) => blob.to_string(),
665 SavedCredential::File(path) => std::fs::read_to_string(path).ok()?,
666 SavedCredential::Absent => return None,
667 };
668 credential::parse_keyring_blob(&raw).ok()
669 }
670}
671
672async fn fetch_remote(
688 client: &reqwest::Client,
689 cache: &Cache,
690 oauth: Option<&cloud::OauthClient>,
691 endpoints: &cloud::Endpoints,
692 token: Result<StoredToken>,
693 renewal: Renewal<'_>,
694 now: DateTime<Utc>,
695) -> Result<AntigravitySnapshot> {
696 let mut token = token?;
697 let oauth_path = cloud::oauth_cache_path(cache);
698 if let Some(mut snap) = super::agy::quota_snapshot(renewal.agy, cache, now).await {
699 snap.plan = match current_access_token(&oauth_path, &token, now) {
700 Some(access) => cloud::fetch_plan(client, endpoints, &access).await,
701 None => None,
702 }
703 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
704 snap.account = remote_account(&token.fingerprint);
705 snap.source = AntigravitySource::Agy;
706 return Ok(snap);
707 }
708 let mut access =
709 match resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await {
710 Ok(access) => access,
711 Err(e) if oauth.is_none() => {
714 let Some(renewed) = renewal.renewed(cache, now).await else {
715 return Err(e);
716 };
717 token = renewed;
718 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
719 }
720 Err(e) => return Err(e),
721 };
722
723 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
724 Err(e) if cloud::is_no_plan_quota(&e) => return Err(e),
727 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
728 access = if oauth.is_none() {
729 let Some(renewed) = renewal.renewed(cache, now).await else {
730 return Err(refresh_unconfigured());
731 };
732 token = renewed;
733 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
734 } else {
735 refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?
736 };
737 cloud::fetch_quota(client, endpoints, &access.value)
738 .await
739 .map_err(|e| {
740 if cloud::is_no_plan_quota(&e) {
741 e
742 } else if is_auth_rejection(&e) {
743 session_expired()
744 } else {
745 e
746 }
747 })?
748 }
749 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
750 other => other?,
751 };
752
753 let plan = cloud::fetch_plan(client, endpoints, &access.value)
754 .await
755 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
756 let mut snap = parse_quota_summary("a, plan)?;
757 snap.account = remote_account(&token.fingerprint);
758 snap.source = AntigravitySource::Remote;
759 Ok(snap)
760}
761
762fn account_key(user_status: &serde_json::Value) -> String {
767 let email = user_status["userStatus"]["email"]
768 .as_str()
769 .filter(|s| !s.is_empty());
770 match email {
771 Some(e) => {
772 use std::hash::{Hash, Hasher};
773 let mut h = std::collections::hash_map::DefaultHasher::new();
774 e.hash(&mut h);
775 format!("acct:{:016x}", h.finish())
776 }
777 None => "acct:unknown".to_string(),
778 }
779}
780
781async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
785 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
786 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
789 .await
790 .ok()?;
791 let html = String::from_utf8_lossy(&bytes);
792 html.split("csrfToken\":\"")
793 .nth(1)
794 .and_then(|s| s.split('"').next())
795 .filter(|t| !t.is_empty())
796 .map(|t| t.to_string())
797}
798
799async fn post_rpc(
800 client: &reqwest::Client,
801 base: &str,
802 csrf: Option<&str>,
803 rpc: &str,
804) -> Result<serde_json::Value> {
805 let mut req = client
806 .post(format!("{base}/{rpc}"))
807 .header("Content-Type", "application/json")
808 .body("{}")
809 .timeout(HTTP_TIMEOUT);
810 if let Some(token) = csrf {
811 req = req.header("x-codeium-csrf-token", token);
812 }
813 let resp = req.send().await?;
814
815 let status = resp.status();
816 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
820 if !status.is_success() {
821 let body = String::from_utf8_lossy(&bytes).into_owned();
822 return Err(AppError::Http {
823 status: status.as_u16(),
824 body,
825 });
826 }
827
828 Ok(serde_json::from_slice(&bytes)?)
829}
830
831pub fn plan_from_status(v: &serde_json::Value) -> String {
832 v["userStatus"]["userTier"]["name"]
833 .as_str()
834 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
835 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
836 .filter(|s| !s.is_empty())
837 .unwrap_or(DEFAULT_PLAN)
838 .to_string()
839}
840
841fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
853 let id = if id.is_empty() { "<unnamed>" } else { id };
854 let group = if group_name.is_empty() {
855 String::new()
856 } else {
857 format!(" in {group_name:?}")
858 };
859 match window {
860 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
861 _ => format!("{id}{group}"),
862 }
863}
864
865fn no_usable_bucket(seen: &[String]) -> AppError {
871 if seen.is_empty() {
872 return AppError::Other(
873 "antigravity: quota summary has no buckets at all — the running product may \
874 not have a quota for this account yet"
875 .into(),
876 );
877 }
878 AppError::Other(format!(
879 "antigravity: quota summary has no bucket in a window we recognise (5h or \
880 weekly, Gemini or Claude/GPT); it offered: {}",
881 seen.join(", ")
882 ))
883}
884
885pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
886 let groups = v["response"]["groups"]
887 .as_array()
888 .or_else(|| v["groups"].as_array())
889 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
890
891 const SLOT_NAMES: [&str; 4] = [
893 "Gemini 5h",
894 "Gemini weekly",
895 "Claude/GPT 5h",
896 "Claude/GPT weekly",
897 ];
898 const CANONICAL_IDS: [&str; 4] = ["gemini-5h", "gemini-weekly", "3p-5h", "3p-weekly"];
899 let is_canonical = |bucket: &serde_json::Value| {
900 CANONICAL_IDS.contains(&bucket["bucketId"].as_str().unwrap_or_default())
901 };
902 let shares_a_slot = |group: &serde_json::Value| {
903 group["buckets"]
904 .as_array()
905 .is_some_and(|buckets| buckets.iter().any(is_canonical))
906 };
907 let renamed = !groups.iter().any(shares_a_slot);
914
915 let mut exact: [Vec<&serde_json::Value>; 4] = Default::default();
918 let mut inferred: [Vec<&serde_json::Value>; 4] = Default::default();
919 let mut pools: Vec<(&str, &serde_json::Value, bool)> = Vec::new();
920 let mut seen: Vec<String> = Vec::new();
924
925 for group in groups {
926 let group_name = group["displayName"].as_str().unwrap_or_default();
927 let Some(buckets) = group["buckets"].as_array() else {
928 continue;
929 };
930 let own_pool = !renamed && !shares_a_slot(group);
931 for bucket in buckets {
932 let id = bucket["bucketId"].as_str().unwrap_or_default();
933 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
934 if let Some(slot) = CANONICAL_IDS.iter().position(|c| *c == id) {
935 exact[slot].push(bucket);
936 continue;
937 }
938 let window = bucket["window"].as_str().unwrap_or_default();
939 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
940 true
941 } else if id.ends_with("5h") || window == "5h" {
942 false
943 } else {
944 continue;
947 };
948 if own_pool {
949 pools.push((group_name, bucket, is_weekly));
950 continue;
951 }
952 if !renamed {
955 continue;
956 }
957 let is_gemini = if id.starts_with("gemini") {
958 true
959 } else if id.starts_with("3p") {
960 false
961 } else if group_name.contains("Gemini") {
962 true
963 } else if group_name.contains("Claude") || group_name.contains("GPT") {
964 false
965 } else {
966 pools.push((group_name, bucket, is_weekly));
969 continue;
970 };
971
972 let slot = usize::from(!is_gemini) * 2 + usize::from(is_weekly);
973 inferred[slot].push(bucket);
974 }
975 }
976
977 let mut slots: [Option<UsageWindow>; 4] = Default::default();
981 for slot in 0..4 {
982 let candidates = if renamed {
983 &inferred[slot]
984 } else {
985 &exact[slot]
986 };
987 match candidates.as_slice() {
988 [] => {}
989 [bucket] => slots[slot] = Some(usage_window(bucket, slot % 2 == 1)?),
990 _ => {
991 return Err(AppError::Schema(format!(
992 "antigravity: duplicate {} bucket",
993 SLOT_NAMES[slot]
994 )));
995 }
996 }
997 }
998 let [gemini_5h, gemini_weekly, tp_5h, tp_weekly] = slots;
999
1000 let model_pools: Vec<ModelPoolWindow> = pools
1004 .into_iter()
1005 .filter_map(|(group_name, bucket, is_weekly)| {
1006 let window = usage_window(bucket, is_weekly).ok()?;
1007 Some(ModelPoolWindow {
1008 name: pool_name(group_name, bucket),
1009 window,
1010 })
1011 })
1012 .collect();
1013
1014 if gemini_5h.is_none()
1018 && gemini_weekly.is_none()
1019 && tp_5h.is_none()
1020 && tp_weekly.is_none()
1021 && model_pools.is_empty()
1022 {
1023 return Err(no_usable_bucket(&seen));
1024 }
1025
1026 Ok(AntigravitySnapshot {
1027 plan,
1028 account: String::new(),
1031 source: AntigravitySource::Local,
1032 session: gemini_5h,
1033 weekly: gemini_weekly,
1034 third_party_session: tp_5h,
1035 third_party_weekly: tp_weekly,
1036 model_pools,
1037 })
1038}
1039
1040fn pool_name(group_name: &str, bucket: &serde_json::Value) -> String {
1044 [
1045 group_name,
1046 bucket["displayName"].as_str().unwrap_or_default(),
1047 bucket["bucketId"].as_str().unwrap_or_default(),
1048 ]
1049 .into_iter()
1050 .map(crate::display::sanitize_untrusted_field)
1051 .find(|name| !name.trim().is_empty())
1052 .unwrap_or_else(|| "Model pool".to_string())
1053}
1054
1055fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
1059 let remaining = bucket["remainingFraction"]
1060 .as_f64()
1061 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
1062 .ok_or_else(|| {
1063 AppError::Schema(format!(
1064 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
1065 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
1066 ))
1067 })?;
1068 Ok(UsageWindow {
1069 utilization_pct: pct_used(remaining),
1070 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
1071 window_duration: if is_weekly {
1072 chrono::Duration::days(7)
1073 } else {
1074 chrono::Duration::hours(5)
1075 },
1076 })
1077}
1078
1079fn pct_used(remaining_fraction: f64) -> i32 {
1082 let used = (1.0 - remaining_fraction) * 100.0;
1083 used.round() as i32
1084}
1085
1086fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
1087 match value {
1088 serde_json::Value::Null => Ok(None),
1089 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
1090 .map(|dt| Some(dt.with_timezone(&Utc)))
1091 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
1092 _ => Err(AppError::Schema(format!(
1093 "antigravity: {field} must be a timestamp or null"
1094 ))),
1095 }
1096}
1097
1098fn candidate_bases(ide_language_server: bool) -> Vec<Candidate> {
1106 let discovered = discover_ls_ports();
1107 let ide_pids = if ide_language_server {
1110 std::collections::HashSet::new()
1111 } else {
1112 ide_ls_pids(discovered.iter().map(|&(pid, _)| pid).collect())
1113 };
1114 candidate_bases_with(
1115 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
1116 discovered,
1117 &ide_pids,
1118 )
1119}
1120
1121fn candidate_bases_with(
1125 override_addr: Option<&str>,
1126 discovered: Vec<(u32, u16)>,
1127 ide_pids: &std::collections::HashSet<u32>,
1128) -> Vec<Candidate> {
1129 let mut candidates = Vec::new();
1130 if let Some(base) = override_addr.and_then(normalize_base) {
1131 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
1136 candidates.push(Candidate {
1137 base,
1138 pid,
1139 ide: false,
1140 });
1141 }
1142
1143 for (pid, port) in discovered {
1149 let base = format!("http://127.0.0.1:{port}");
1150 if !candidates.iter().any(|c| c.base == base) {
1151 candidates.push(Candidate {
1152 base,
1153 pid: Some(pid),
1154 ide: ide_pids.contains(&pid),
1155 });
1156 }
1157 }
1158
1159 candidates
1160}
1161
1162#[cfg(any(test, target_os = "linux", target_os = "macos"))]
1167fn is_ide_install(exe: &str) -> bool {
1168 let exe = exe.to_lowercase().replace('\\', "/");
1169 exe.contains("antigravity-ide") || exe.contains("antigravity ide")
1170}
1171
1172#[cfg(target_os = "linux")]
1174fn ide_ls_pids(pids: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1175 pids.into_iter()
1176 .filter(|pid| {
1177 std::fs::read_link(format!("/proc/{pid}/exe"))
1178 .map(|path| path.to_string_lossy().into_owned())
1179 .is_ok_and(|exe| is_ide_install(&exe))
1180 })
1181 .collect()
1182}
1183
1184#[cfg(target_os = "macos")]
1187fn ide_ls_pids(pids: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1188 if pids.is_empty() {
1189 return std::collections::HashSet::new();
1190 }
1191 let list = pids
1192 .iter()
1193 .map(u32::to_string)
1194 .collect::<Vec<_>>()
1195 .join(",");
1196 let Ok(output) = std::process::Command::new("ps")
1197 .args(["-o", "pid=,comm=", "-p", &list])
1198 .output()
1199 else {
1200 return std::collections::HashSet::new();
1201 };
1202 parse_ps_comm(&String::from_utf8_lossy(&output.stdout))
1203 .into_iter()
1204 .filter(|(pid, exe)| pids.contains(pid) && is_ide_install(exe))
1205 .map(|(pid, _)| pid)
1206 .collect()
1207}
1208
1209#[cfg(any(test, target_os = "macos"))]
1212fn parse_ps_comm(output: &str) -> Vec<(u32, String)> {
1213 output
1214 .lines()
1215 .filter_map(|line| {
1216 let (pid, exe) = line.trim().split_once(char::is_whitespace)?;
1217 Some((pid.parse().ok()?, exe.trim().to_string()))
1218 })
1219 .collect()
1220}
1221
1222#[cfg(target_os = "windows")]
1226fn ide_ls_pids(_: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1227 std::collections::HashSet::new()
1228}
1229
1230#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1231fn ide_ls_pids(_: std::collections::HashSet<u32>) -> std::collections::HashSet<u32> {
1232 std::collections::HashSet::new()
1233}
1234
1235fn base_port(base: &str) -> Option<u16> {
1237 let (_, authority) = base.split_once("://")?;
1238 authority.rsplit_once(':')?.1.parse().ok()
1239}
1240
1241fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
1244 let mut owners = discovered
1245 .iter()
1246 .filter(|(_, p)| *p == port)
1247 .map(|(pid, _)| *pid);
1248 let first = owners.next()?;
1249 owners.all(|pid| pid == first).then_some(first)
1250}
1251
1252fn normalize_base(addr: &str) -> Option<String> {
1261 let trimmed = addr.trim();
1262 let (scheme, authority) = match trimmed.split_once("://") {
1263 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
1264 _ => ("http", trimmed),
1265 };
1266 let authority = authority.trim_end_matches('/');
1267 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
1268}
1269
1270fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
1277 let comm = comm.trim().to_lowercase();
1278 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
1279 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
1280 return true;
1281 }
1282 exe.is_some_and(|p| {
1283 let p = p.to_lowercase().replace('\\', "/");
1284 let p = p.strip_suffix(".exe").unwrap_or(&p);
1285 p.contains("antigravity") || p.ends_with("/agy")
1286 })
1287}
1288
1289#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
1321fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
1322 let groups: Vec<(u32, Vec<u16>)> = per_pid
1323 .into_iter()
1324 .map(|(pid, mut group)| {
1325 group.sort_unstable_by(|a, b| b.cmp(a));
1326 group.dedup();
1331 (pid, group)
1332 })
1333 .collect();
1334
1335 let mut ports: Vec<(u32, u16)> = Vec::new();
1336 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
1337 for (pid, port) in groups
1338 .iter()
1339 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
1340 {
1341 if !ports.iter().any(|(_, p)| *p == port) {
1344 ports.push((pid, port));
1345 }
1346 }
1347 }
1348 ports
1349}
1350
1351#[cfg(target_os = "linux")]
1359pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1360 use std::collections::{BTreeMap, HashMap};
1361
1362 let mut owners: HashMap<u64, u32> = HashMap::new();
1365 let Ok(entries) = std::fs::read_dir("/proc") else {
1366 return Vec::new();
1367 };
1368 for entry in entries.flatten() {
1369 let pid_dir = entry.path();
1370 let Some(pid) = pid_dir
1371 .file_name()
1372 .and_then(|name| name.to_str())
1373 .and_then(|name| name.parse::<u32>().ok())
1374 else {
1375 continue;
1376 };
1377 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1378 continue;
1379 };
1380 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1381 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1382 continue;
1383 }
1384 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1385 continue;
1386 };
1387 for fd in fds.flatten() {
1388 let Ok(target) = std::fs::read_link(fd.path()) else {
1389 continue;
1390 };
1391 if let Some(ino) = target
1392 .to_str()
1393 .and_then(|s| s.strip_prefix("socket:["))
1394 .and_then(|s| s.strip_suffix(']'))
1395 .and_then(|s| s.parse::<u64>().ok())
1396 {
1397 owners.insert(ino, pid);
1398 }
1399 }
1400 }
1401
1402 if owners.is_empty() {
1403 return Vec::new();
1404 }
1405
1406 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1407 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1408 let Ok(contents) = std::fs::read_to_string(table) else {
1409 continue;
1410 };
1411 for line in contents.lines().skip(1) {
1412 if let Some((port, ino)) = parse_proc_net_line(line)
1413 && let Some(&pid) = owners.get(&ino)
1414 {
1415 per_pid.entry(pid).or_default().push(port);
1416 }
1417 }
1418 }
1419 probe_order(per_pid)
1420}
1421
1422#[cfg(target_os = "macos")]
1429pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1430 let Ok(output) = std::process::Command::new("lsof")
1431 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1432 .output()
1433 else {
1434 return Vec::new();
1435 };
1436 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1440}
1441
1442#[cfg(any(test, target_os = "macos"))]
1447fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1448 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1449 let mut pid = None;
1452 let mut owner = None;
1453 for line in output.lines() {
1454 let Some(rest) = line.get(1..) else { continue };
1455 match line.as_bytes().first() {
1456 Some(b'p') => {
1457 pid = rest.parse::<u32>().ok();
1458 owner = None;
1459 }
1460 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1461 Some(b'n') => {
1462 if let Some(pid) = owner
1463 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1464 {
1465 per_pid.entry(pid).or_default().push(port);
1466 }
1467 }
1468 _ => {}
1469 }
1470 }
1471 probe_order(per_pid)
1472}
1473
1474#[cfg(any(test, target_os = "windows"))]
1475#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1476struct WindowsTcpRow {
1477 local_addr: [u8; 4],
1478 local_port: u32,
1479 pid: u32,
1480}
1481
1482#[cfg(any(test, target_os = "windows"))]
1483fn decode_windows_process_name(raw: &[u16]) -> String {
1484 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1485 String::from_utf16_lossy(&raw[..end])
1486}
1487
1488#[cfg(any(test, target_os = "windows"))]
1489fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1490 processes
1491 .iter()
1492 .filter(|(_, name)| is_antigravity_process(name, None))
1493 .map(|(pid, _)| *pid)
1494 .collect()
1495}
1496
1497#[cfg(any(test, target_os = "windows"))]
1500fn matching_windows_ports(
1501 pids: &std::collections::HashSet<u32>,
1502 rows: &[WindowsTcpRow],
1503) -> Vec<(u32, u16)> {
1504 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1505 for row in rows {
1506 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1507 continue;
1508 }
1509 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1510 if port != 0 {
1511 per_pid.entry(row.pid).or_default().push(port);
1512 }
1513 }
1514 probe_order(per_pid)
1515}
1516
1517#[cfg(any(test, target_os = "windows"))]
1518fn checked_windows_row_count(
1519 buffer_len: usize,
1520 rows_offset: usize,
1521 row_size: usize,
1522 declared: usize,
1523) -> Option<usize> {
1524 let rows_len = row_size.checked_mul(declared)?;
1525 let end = rows_offset.checked_add(rows_len)?;
1526 (row_size != 0 && end <= buffer_len).then_some(declared)
1527}
1528
1529#[cfg(target_os = "windows")]
1530struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1531
1532#[cfg(target_os = "windows")]
1533impl Drop for WindowsHandle {
1534 fn drop(&mut self) {
1535 unsafe {
1536 windows_sys::Win32::Foundation::CloseHandle(self.0);
1537 }
1538 }
1539}
1540
1541#[cfg(target_os = "windows")]
1542fn windows_processes() -> Vec<(u32, String)> {
1543 use std::mem::size_of;
1544 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1545 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1546 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1547 TH32CS_SNAPPROCESS,
1548 };
1549
1550 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1551 if handle == INVALID_HANDLE_VALUE {
1552 return Vec::new();
1553 }
1554 let snapshot = WindowsHandle(handle);
1555 let mut entry = PROCESSENTRY32W::default();
1556 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1557 return Vec::new();
1558 };
1559 entry.dwSize = entry_size;
1560 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1561 return Vec::new();
1562 }
1563
1564 let mut processes = Vec::new();
1565 loop {
1566 processes.push((
1567 entry.th32ProcessID,
1568 decode_windows_process_name(&entry.szExeFile),
1569 ));
1570 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1571 break;
1572 }
1573 }
1574 processes
1575}
1576
1577#[cfg(target_os = "windows")]
1578fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1579 use std::mem::{offset_of, size_of, size_of_val};
1580 use windows_sys::Win32::NetworkManagement::IpHelper::{
1581 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1582 };
1583
1584 let available = used_bytes.min(size_of_val(buffer));
1585 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1586 if available < size_of::<u32>() || available < rows_offset {
1587 return Vec::new();
1588 }
1589 let base = buffer.as_ptr().cast::<u8>();
1590 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1591 if checked_windows_row_count(
1592 available,
1593 rows_offset,
1594 size_of::<MIB_TCPROW_OWNER_PID>(),
1595 declared,
1596 )
1597 .is_none()
1598 {
1599 return Vec::new();
1600 }
1601
1602 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1603 (0..declared)
1604 .map(|index| unsafe { rows.add(index).read_unaligned() })
1605 .map(|row| WindowsTcpRow {
1606 local_addr: row.dwLocalAddr.to_ne_bytes(),
1607 local_port: row.dwLocalPort,
1608 pid: row.dwOwningPid,
1609 })
1610 .collect()
1611}
1612
1613#[cfg(target_os = "windows")]
1614fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1615 use std::mem::size_of;
1616 use std::ptr::null_mut;
1617 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1618 use windows_sys::Win32::NetworkManagement::IpHelper::{
1619 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1620 };
1621 use windows_sys::Win32::Networking::WinSock::AF_INET;
1622
1623 let mut size = 0u32;
1624 let status = unsafe {
1625 GetExtendedTcpTable(
1626 null_mut(),
1627 &mut size,
1628 0,
1629 u32::from(AF_INET),
1630 TCP_TABLE_OWNER_PID_LISTENER,
1631 0,
1632 )
1633 };
1634 if status != ERROR_INSUFFICIENT_BUFFER {
1635 return Vec::new();
1636 }
1637
1638 for _ in 0..3 {
1639 let Some(words) = (size as usize)
1640 .checked_add(size_of::<u32>() - 1)
1641 .map(|bytes| bytes / size_of::<u32>())
1642 else {
1643 return Vec::new();
1644 };
1645 if words == 0 {
1646 return Vec::new();
1647 }
1648 let mut buffer = Vec::<u32>::new();
1649 if buffer.try_reserve_exact(words).is_err() {
1650 return Vec::new();
1651 }
1652 buffer.resize(words, 0);
1653 let mut used = size;
1654 let status = unsafe {
1655 GetExtendedTcpTable(
1656 buffer.as_mut_ptr().cast(),
1657 &mut used,
1658 0,
1659 u32::from(AF_INET),
1660 TCP_TABLE_OWNER_PID_LISTENER,
1661 0,
1662 )
1663 };
1664 if status == ERROR_INSUFFICIENT_BUFFER {
1665 size = used;
1666 continue;
1667 }
1668 if status != 0 {
1669 return Vec::new();
1670 }
1671 return parse_windows_tcp_rows(&buffer, used as usize);
1672 }
1673 Vec::new()
1674}
1675
1676#[cfg(target_os = "windows")]
1677pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1678 let pids = matching_windows_process_ids(&windows_processes());
1679 if pids.is_empty() {
1680 return Vec::new();
1681 }
1682 matching_windows_ports(&pids, &windows_tcp_rows())
1683}
1684
1685#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1686pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1687 Vec::new()
1688}
1689
1690#[cfg(target_os = "linux")]
1693fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1694 let cols: Vec<&str> = line.split_whitespace().collect();
1695 if cols.len() < 10 {
1696 return None;
1697 }
1698 if cols[3] != "0A" {
1700 return None;
1701 }
1702 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1703 let inode = cols[9].parse::<u64>().ok()?;
1704 Some((port, inode))
1705}
1706
1707fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1712 crate::outcome::fallback(cache, None, original, |bytes| {
1713 parse_cache_at(bytes, None, now)
1714 })
1715}
1716
1717fn fallback_with_error(
1722 cache: &Cache,
1723 last_error: Option<(u16, String)>,
1724 reason: AppError,
1725 now: DateTime<Utc>,
1726) -> Result<FetchOutcome> {
1727 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1728 parse_cache_at(bytes, None, now)
1729 })
1730}
1731
1732fn reuse_cache(
1733 bytes: Vec<u8>,
1734 cache: &Cache,
1735 stale: bool,
1736 account: Option<&str>,
1737 now: DateTime<Utc>,
1738) -> Result<FetchOutcome> {
1739 let snap = parse_cache_at(&bytes, account, now)?;
1740 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1741}
1742
1743pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1749 parse_cache_at(bytes, account, Utc::now())
1750}
1751
1752fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<String> {
1762 let shared = [
1763 ("Gemini 5h", snap.session.as_ref()),
1764 ("Gemini weekly", snap.weekly.as_ref()),
1765 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1766 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1767 ]
1768 .into_iter()
1769 .filter_map(|(name, w)| Some((name.to_string(), w?)));
1770 let pools = snap.model_pools.iter().map(|pool| {
1771 let name = format!("{} {}", pool.name, pool_cadence(pool.is_weekly()));
1772 (name, &pool.window)
1773 });
1774 shared
1775 .chain(pools)
1776 .find(|(_, w)| w.resets_at.is_some_and(|r| r <= now))
1777 .map(|(name, _)| name)
1778}
1779
1780fn pool_cadence(weekly: bool) -> &'static str {
1783 if weekly { "weekly" } else { "5h" }
1784}
1785
1786pub fn parse_cache_at(
1787 bytes: &[u8],
1788 account: Option<&str>,
1789 now: DateTime<Utc>,
1790) -> Result<AntigravitySnapshot> {
1791 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1792
1793 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1794 if let Some(expected) = account
1795 && cached_account != Some(expected)
1796 {
1797 return Err(AppError::Schema(
1798 "antigravity cache belongs to a different account; refetching".into(),
1799 ));
1800 }
1801
1802 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1812 match v.get(pct_key) {
1813 None => Err(AppError::Schema(format!(
1814 "antigravity: cached payload is missing {pct_key}"
1815 ))),
1816 Some(serde_json::Value::Null) => Ok(None),
1817 Some(value) => value
1818 .as_i64()
1819 .filter(|pct| (0..=100).contains(pct))
1820 .map(|pct| Some(pct as i32))
1821 .ok_or_else(|| {
1822 AppError::Schema(format!(
1823 "antigravity: cached {pct_key} must be an integer in 0..=100"
1824 ))
1825 }),
1826 }
1827 };
1828
1829 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1830 let Some(pct) = cached_pct(pct_key)? else {
1831 return Ok(None);
1832 };
1833 Ok::<_, AppError>(Some(UsageWindow {
1834 utilization_pct: pct,
1835 resets_at: parse_reset(&v[reset_key], reset_key)?,
1836 window_duration: if weekly {
1837 chrono::Duration::days(7)
1838 } else {
1839 chrono::Duration::hours(5)
1840 },
1841 }))
1842 };
1843
1844 let snap = AntigravitySnapshot {
1845 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1846 account: cached_account.unwrap_or_default().to_string(),
1847 source: v["source"]
1849 .as_str()
1850 .and_then(AntigravitySource::parse)
1851 .unwrap_or_default(),
1852 session: optional("session_pct", "session_reset", false)?,
1853 weekly: optional("weekly_pct", "weekly_reset", true)?,
1854 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1855 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1856 model_pools: cached_pools(&v)?,
1857 };
1858
1859 if snap.session.is_none()
1862 && snap.weekly.is_none()
1863 && snap.third_party_session.is_none()
1864 && snap.third_party_weekly.is_none()
1865 && snap.model_pools.is_empty()
1866 {
1867 return Err(AppError::Schema(
1868 "antigravity cache holds no usable window; refetching".into(),
1869 ));
1870 }
1871
1872 if let Some(window) = expired_window(&snap, now) {
1873 return Err(AppError::Schema(format!(
1874 "antigravity cache is past its {window} reset; refetching"
1875 )));
1876 }
1877 Ok(snap)
1878}
1879
1880fn cached_pools(v: &serde_json::Value) -> Result<Vec<ModelPoolWindow>> {
1884 let Some(entries) = v.get("model_pools") else {
1885 return Ok(Vec::new());
1886 };
1887 let malformed = || AppError::Schema("antigravity: cached model_pools is malformed".into());
1888 entries
1889 .as_array()
1890 .ok_or_else(malformed)?
1891 .iter()
1892 .map(|entry| {
1893 let name = entry["name"].as_str().ok_or_else(malformed)?;
1894 let weekly = match entry["window"].as_str() {
1895 Some(cadence) if cadence == pool_cadence(true) => true,
1896 Some(cadence) if cadence == pool_cadence(false) => false,
1897 _ => return Err(malformed()),
1898 };
1899 let pct = entry["pct"]
1900 .as_i64()
1901 .filter(|pct| (0..=100).contains(pct))
1902 .ok_or_else(malformed)?;
1903 Ok(ModelPoolWindow {
1904 name: crate::display::sanitize_untrusted_field(name),
1908 window: UsageWindow {
1909 utilization_pct: pct as i32,
1910 resets_at: parse_reset(&entry["reset"], "model_pools reset")?,
1911 window_duration: if weekly {
1912 chrono::Duration::days(7)
1913 } else {
1914 chrono::Duration::hours(5)
1915 },
1916 },
1917 })
1918 })
1919 .collect()
1920}
1921
1922pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1923 let model_pools: Vec<serde_json::Value> = snap
1924 .model_pools
1925 .iter()
1926 .map(|pool| {
1927 serde_json::json!({
1928 "name": pool.name,
1929 "window": pool_cadence(pool.is_weekly()),
1930 "pct": pool.window.utilization_pct,
1931 "reset": pool.window.resets_at.map(|dt| dt.to_rfc3339()),
1932 })
1933 })
1934 .collect();
1935 serde_json::json!({
1936 "plan": snap.plan,
1937 "account": snap.account,
1938 "source": snap.source.as_str(),
1939 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
1940 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1941 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
1942 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1943 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
1944 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1945 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
1946 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1947 "model_pools": model_pools,
1948 })
1949}
1950
1951#[cfg(test)]
1952mod tests {
1953 use super::*;
1954
1955 const QUOTA_JSON: &str = r#"{
1959 "response": {
1960 "groups": [
1961 {
1962 "displayName": "Gemini Models",
1963 "buckets": [
1964 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
1965 "window": "weekly", "remainingFraction": 0.9191212,
1966 "resetTime": "2026-07-28T17:39:58Z"},
1967 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
1968 "window": "5h", "remainingFraction": 0.5672253,
1969 "resetTime": "2026-07-22T17:47:00Z"}
1970 ]
1971 },
1972 {
1973 "displayName": "Claude and GPT models",
1974 "buckets": [
1975 {"bucketId": "3p-weekly", "window": "weekly",
1976 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
1977 {"bucketId": "3p-5h", "window": "5h",
1978 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
1979 ]
1980 }
1981 ]
1982 }
1983 }"#;
1984
1985 fn now() -> DateTime<Utc> {
1988 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
1989 .unwrap()
1990 .with_timezone(&Utc)
1991 }
1992
1993 fn parsed() -> AntigravitySnapshot {
1994 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1995 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
1996 }
1997
1998 #[test]
1999 fn quota_summary_maps_four_distinct_windows() {
2000 let snap = parsed();
2001 assert_eq!(snap.plan, "Google AI Pro");
2002 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2004 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2005 assert_eq!(
2006 snap.third_party_session.as_ref().unwrap().utilization_pct,
2007 75
2008 );
2009 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
2010 }
2011
2012 #[test]
2013 fn each_window_keeps_its_own_reset_time() {
2014 let snap = parsed();
2015 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
2016 assert_eq!(
2017 snap.session.as_ref().unwrap().resets_at,
2018 at("2026-07-22T17:47:00Z")
2019 );
2020 assert_eq!(
2021 snap.weekly.as_ref().unwrap().resets_at,
2022 at("2026-07-28T17:39:58Z")
2023 );
2024 assert_eq!(
2025 snap.third_party_weekly.as_ref().unwrap().resets_at,
2026 at("2026-07-29T12:47:00Z")
2027 );
2028 assert_ne!(
2030 snap.session.as_ref().unwrap().resets_at,
2031 snap.weekly.as_ref().unwrap().resets_at
2032 );
2033 }
2034
2035 #[test]
2036 fn window_durations_match_their_bucket() {
2037 let snap = parsed();
2038 assert_eq!(
2039 snap.session.as_ref().unwrap().window_duration,
2040 chrono::Duration::hours(5)
2041 );
2042 assert_eq!(
2043 snap.weekly.as_ref().unwrap().window_duration,
2044 chrono::Duration::days(7)
2045 );
2046 assert_eq!(
2047 snap.third_party_weekly.as_ref().unwrap().window_duration,
2048 chrono::Duration::days(7)
2049 );
2050 }
2051
2052 #[test]
2053 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
2054 let v: serde_json::Value = serde_json::from_str(
2055 r#"{"response":{"groups":[
2056 {"displayName":"Gemini Models","buckets":[
2057 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
2058 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
2059 {"displayName":"Claude and GPT models","buckets":[
2060 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
2061 ]}}"#,
2062 )
2063 .unwrap();
2064 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2065 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
2066 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
2067 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
2068 assert!(snap.third_party_weekly.is_none());
2069 }
2070
2071 #[test]
2077 fn a_per_model_pool_is_reported_beside_the_canonical_buckets() {
2078 let v: serde_json::Value = serde_json::from_str(
2079 r#"{"response":{"groups":[
2080 {"displayName":"Gemini Models","buckets":[
2081 {"bucketId":"gemini-weekly","displayName":"Weekly Limit Remaining","window":"weekly",
2082 "remainingFraction":0.7839184,"resetTime":"2026-10-23T12:22:17Z"},
2083 {"bucketId":"gemini-5h","displayName":"Five Hour Limit Remaining","window":"5h",
2084 "remainingFraction":0.253501,"resetTime":"2026-10-16T17:22:17Z"}]},
2085 {"displayName":"Claude and GPT models","buckets":[
2086 {"bucketId":"3p-weekly","displayName":"Weekly Limit Remaining","window":"weekly",
2087 "remainingFraction":1,"resetTime":"2026-10-23T17:05:12Z"},
2088 {"bucketId":"3p-5h","displayName":"Five Hour Limit Remaining","window":"5h",
2089 "remainingFraction":1,"resetTime":"2026-10-16T22:05:12Z"}]},
2090 {"displayName":"Gemini 4 Argon","buckets":[
2091 {"bucketId":"gemini-argon-wtu-weekly","displayName":"Gemini 4 Argon Weekly Limit",
2092 "window":"weekly","remainingFraction":0.4,"resetTime":"2026-10-23T17:05:12Z"}]}
2093 ]}}"#,
2094 )
2095 .unwrap();
2096 let snap = parse_quota_summary(&v, "Google AI Pro".into()).unwrap();
2097 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
2098 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 22);
2099 assert_eq!(
2100 snap.third_party_session.as_ref().unwrap().utilization_pct,
2101 0
2102 );
2103 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
2104 let [argon] = snap.model_pools.as_slice() else {
2105 panic!("{:?}", snap.model_pools);
2106 };
2107 assert_eq!(argon.name, "Gemini 4 Argon");
2108 assert_eq!(argon.window.utilization_pct, 60);
2109 assert!(argon.is_weekly());
2110 }
2111
2112 #[test]
2115 fn a_malformed_model_pool_is_left_out() {
2116 let v: serde_json::Value = serde_json::from_str(
2117 r#"{"response":{"groups":[
2118 {"displayName":"Gemini Models","buckets":[
2119 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.5}]},
2120 {"displayName":"Gemini 4 Argon","buckets":[
2121 {"bucketId":"gemini-argon-wtu-weekly","window":"weekly","remainingFraction":"x"}]}
2122 ]}}"#,
2123 )
2124 .unwrap();
2125 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2126 assert_eq!(snap.weekly.unwrap().utilization_pct, 50);
2127 assert!(snap.model_pools.is_empty());
2128 }
2129
2130 #[test]
2132 fn a_model_pool_name_is_sanitized_and_falls_back_to_the_bucket() {
2133 let v: serde_json::Value = serde_json::from_str(
2134 "{\"response\":{\"groups\":[
2135 {\"displayName\":\"Gemini Models\",\"buckets\":[
2136 {\"bucketId\":\"gemini-weekly\",\"window\":\"weekly\",\"remainingFraction\":0.5}]},
2137 {\"displayName\":\"Argon\\u001b[31m\",\"buckets\":[
2138 {\"bucketId\":\"argon-weekly\",\"window\":\"weekly\",\"remainingFraction\":0.5}]},
2139 {\"buckets\":[
2140 {\"bucketId\":\"nova-5h\",\"displayName\":\"Nova Five Hour\",\"window\":\"5h\",\"remainingFraction\":0.5}]}
2141 ]}}",
2142 )
2143 .unwrap();
2144 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
2145 let names: Vec<&str> = snap.model_pools.iter().map(|p| p.name.as_str()).collect();
2146 assert_eq!(names, ["Argon[31m", "Nova Five Hour"]);
2147 }
2148
2149 #[test]
2150 fn canonical_presence_disables_inference_in_a_shared_group() {
2151 for remaining in [serde_json::json!(0.1), serde_json::json!("unknown")] {
2152 let summary = serde_json::json!({"groups": [{
2153 "displayName": "Gemini Models",
2154 "buckets": [
2155 {"bucketId": "gemini-5h", "remainingFraction": 0.5},
2156 {"bucketId": "gemini-argon-wtu-weekly", "remainingFraction": remaining}
2157 ]
2158 }]});
2159 let snap = parse_quota_summary(&summary, "Pro".into()).unwrap();
2160 assert_eq!(snap.session.unwrap().utilization_pct, 50);
2161 assert!(snap.weekly.is_none());
2162 assert!(snap.model_pools.is_empty());
2163 }
2164 }
2165
2166 #[test]
2167 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
2168 let lookalike: serde_json::Value = serde_json::from_str(
2171 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
2172 {"bucketId":"gemini-5h-copy","window":"5h"},
2173 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2174 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
2175 ]}]}}"#,
2176 )
2177 .unwrap();
2178 let snap = parse_quota_summary(&lookalike, "Pro".into()).unwrap();
2179 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 10);
2180
2181 for buckets in [
2184 r#"{"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2185 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.1}"#,
2186 r#"{"bucketId":"x1","window":"5h","remainingFraction":0.9},
2187 {"bucketId":"x2","window":"5h","remainingFraction":0.1}"#,
2188 ] {
2189 let duplicate: serde_json::Value = serde_json::from_str(&format!(
2190 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[{buckets}]}}]}}}}"#
2191 ))
2192 .unwrap();
2193 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
2194 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
2195 }
2196
2197 let unrelated: serde_json::Value = serde_json::from_str(
2200 r#"{"response":{"groups":[
2201 {"displayName":"Gemini Models","buckets":[
2202 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
2203 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
2204 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
2205 ]},
2206 {"displayName":"Future Models","buckets":[
2207 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
2208 ]}
2209 ]}}"#,
2210 )
2211 .unwrap();
2212 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
2213 assert!(snap.third_party_session.is_none());
2214 assert!(snap.third_party_weekly.is_none());
2215 let [future] = snap.model_pools.as_slice() else {
2216 panic!("{:?}", snap.model_pools);
2217 };
2218 assert_eq!(future.name, "Future Models");
2219 assert!(!future.is_weekly());
2220 }
2221
2222 #[test]
2225 fn a_bucket_without_a_usable_fraction_is_rejected() {
2226 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
2227 let v: serde_json::Value = serde_json::from_str(&format!(
2228 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
2229 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
2230 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
2231 ))
2232 .unwrap();
2233 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
2234 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
2235 }
2236 }
2237
2238 #[test]
2239 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
2240 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
2241 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2242 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
2243 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
2244 assert!(err.to_string().contains("resetTime"), "{err}");
2245 }
2246 }
2247
2248 #[test]
2252 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
2253 let mut snap = parsed();
2254 snap.session = None;
2255 snap.third_party_session = None;
2256
2257 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2258 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
2259
2260 assert!(back.session.is_none());
2261 assert_eq!(
2262 back.weekly.as_ref().unwrap().utilization_pct,
2263 snap.weekly.as_ref().unwrap().utilization_pct
2264 );
2265 }
2266
2267 #[test]
2272 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
2273 let summary = serde_json::json!({
2274 "groups": [
2275 {
2276 "displayName": "Gemini Models",
2277 "buckets": [{
2278 "bucketId": "gemini-weekly", "window": "weekly",
2279 "remainingFraction": 0.42,
2280 }],
2281 },
2282 {
2283 "displayName": "Claude and GPT models",
2284 "buckets": [{
2285 "bucketId": "3p-weekly", "window": "weekly",
2286 "remainingFraction": 0.9,
2287 }],
2288 },
2289 ],
2290 });
2291
2292 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
2293
2294 assert!(snap.session.is_none(), "no 5h bucket arrived");
2295 assert!(snap.third_party_session.is_none());
2296 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
2297 assert_eq!(
2298 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
2299 10
2300 );
2301 }
2302
2303 #[test]
2306 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
2307 let summary = serde_json::json!({
2308 "groups": [{
2309 "displayName": "Gemini Models",
2310 "buckets": [{
2311 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
2312 }],
2313 }],
2314 });
2315
2316 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
2317
2318 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
2319 assert!(snap.weekly.is_none());
2320 }
2321
2322 #[test]
2325 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
2326 let summary = serde_json::json!({
2327 "groups": [{
2328 "displayName": "Gemini Models",
2329 "buckets": [{
2330 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
2331 }],
2332 }],
2333 });
2334
2335 let rendered = parse_quota_summary(&summary, "Pro".into())
2336 .expect_err("an unrecognised cadence alone is not a snapshot")
2337 .to_string();
2338
2339 assert!(
2340 rendered.contains("no bucket in a window we recognise"),
2341 "{rendered}"
2342 );
2343 assert!(rendered.contains("gemini-daily"), "{rendered}");
2344 assert!(rendered.contains("window daily"), "{rendered}");
2345 }
2346
2347 #[test]
2350 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
2351 let summary = serde_json::json!({
2352 "groups": [{"displayName": "Gemini", "buckets": []}],
2353 });
2354
2355 let rendered = parse_quota_summary(&summary, "Pro".into())
2356 .expect_err("no buckets is an error")
2357 .to_string();
2358
2359 assert!(rendered.contains("no buckets at all"), "{rendered}");
2360 assert!(!rendered.contains("it offered:"), "{rendered}");
2361 }
2362
2363 #[test]
2366 fn buckets_without_an_id_are_still_named_in_the_error() {
2367 let summary = serde_json::json!({
2368 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
2369 });
2370
2371 let rendered = parse_quota_summary(&summary, "Pro".into())
2372 .expect_err("an unusable summary is an error")
2373 .to_string();
2374
2375 assert!(rendered.contains("<unnamed>"), "{rendered}");
2376 }
2377
2378 #[test]
2379 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
2380 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
2381 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
2382 }
2383
2384 #[test]
2385 fn cache_round_trip_preserves_every_window() {
2386 let snap = parsed();
2387 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2388 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2389 }
2390
2391 #[test]
2395 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
2396 let full = snap_to_json(&parsed());
2397 for missing in ["session_pct", "weekly_pct"] {
2398 let mut v = full.clone();
2399 v.as_object_mut().unwrap().remove(missing);
2400 let bytes = serde_json::to_vec(&v).unwrap();
2401 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2402 assert!(err.to_string().contains(missing), "{missing}: {err}");
2403 }
2404 assert!(parse_cache_at(b"{}", None, now()).is_err());
2406 }
2407
2408 #[test]
2409 fn cached_percentages_are_range_checked_before_narrowing() {
2410 let full = snap_to_json(&parsed());
2411 for (key, bad) in [
2412 ("session_pct", serde_json::json!(-1)),
2413 ("weekly_pct", serde_json::json!(101)),
2414 ("session_pct", serde_json::json!(i64::MAX)),
2415 ("tp_session_pct", serde_json::json!("75")),
2416 ] {
2417 let mut v = full.clone();
2418 v[key] = bad;
2419 let bytes = serde_json::to_vec(&v).unwrap();
2420 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2421 assert!(err.to_string().contains(key), "{key}: {err}");
2422 }
2423 }
2424
2425 #[test]
2426 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
2427 let mut v = snap_to_json(&parsed());
2428 v["session_reset"] = serde_json::json!("not-a-time");
2429 let bytes = serde_json::to_vec(&v).unwrap();
2430 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
2431 assert!(err.to_string().contains("session_reset"), "{err}");
2432 }
2433
2434 #[test]
2439 fn a_cache_past_its_reset_is_refused() {
2440 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
2441 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2442
2443 assert!(parse_cache_at(&bytes, None, now()).is_ok());
2445 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
2447
2448 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
2450 assert!(err.to_string().contains("5h"), "{err}");
2451
2452 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
2454 }
2455
2456 #[test]
2459 fn expiry_names_the_window_that_rolled_over() {
2460 let mut snap = parsed();
2461 snap.session.as_mut().unwrap().resets_at = None;
2463 snap.third_party_session = None;
2464 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2465 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2466
2467 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
2469
2470 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
2472 assert!(err.to_string().contains("Gemini weekly"), "{err}");
2473 }
2474
2475 #[test]
2477 fn a_window_without_a_reset_never_expires() {
2478 let mut snap = parsed();
2479 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
2480 w.resets_at = None;
2481 }
2482 snap.third_party_session = None;
2483 snap.third_party_weekly = None;
2484 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2485 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
2486 .unwrap()
2487 .with_timezone(&Utc);
2488 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
2489 }
2490
2491 #[test]
2493 fn a_cache_from_another_account_is_rejected() {
2494 let mut snap = parsed();
2495 snap.account = "acct:aaaa".into();
2496 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2497
2498 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
2499 assert_eq!(
2500 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
2501 snap
2502 );
2503
2504 let mut legacy = snap_to_json(&snap);
2506 legacy.as_object_mut().unwrap().remove("account");
2507 let legacy = serde_json::to_vec(&legacy).unwrap();
2508 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
2509 }
2510
2511 #[test]
2514 fn an_unverifiable_cache_is_served_rather_than_discarded() {
2515 let mut snap = parsed();
2516 snap.account = "acct:aaaa".into();
2517 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2518 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2519 }
2520
2521 #[test]
2522 fn account_key_fingerprints_rather_than_storing_the_address() {
2523 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
2524 let a = with("someone@example.com");
2525 assert!(!a.contains("someone"), "{a}");
2526 assert!(!a.contains('@'), "{a}");
2527 assert_eq!(a, with("someone@example.com"), "must be stable");
2528 assert_ne!(a, with("other@example.com"));
2529 let unknown = account_key(&serde_json::json!({}));
2531 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
2532 assert_ne!(unknown, a);
2533 }
2534
2535 #[test]
2538 fn absent_third_party_windows_are_not_treated_as_corruption() {
2539 let mut snap = parsed();
2540 snap.third_party_session = None;
2541 snap.third_party_weekly = None;
2542 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2543 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2544 }
2545
2546 #[test]
2547 fn cache_round_trip_preserves_absent_third_party_windows() {
2548 let mut snap = parsed();
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 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2553 }
2554
2555 #[test]
2556 fn pct_used_inverts_valid_fractions() {
2557 assert_eq!(pct_used(1.0), 0);
2558 assert_eq!(pct_used(0.0), 100);
2559 assert_eq!(pct_used(0.5), 50);
2560 }
2561
2562 #[test]
2563 fn plan_falls_back_through_the_status_payload() {
2564 let tier: serde_json::Value =
2565 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2566 .unwrap();
2567 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2568
2569 let plan_only: serde_json::Value = serde_json::from_str(
2570 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2571 )
2572 .unwrap();
2573 assert_eq!(plan_from_status(&plan_only), "Pro");
2574
2575 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2576 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2577 }
2578
2579 #[cfg(target_os = "linux")]
2580 #[test]
2581 fn proc_net_parser_keeps_only_listening_rows() {
2582 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2583 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2584 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2585
2586 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2587 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2588 assert_eq!(parse_proc_net_line(established), None);
2589
2590 assert_eq!(parse_proc_net_line("garbage"), None);
2591 }
2592
2593 #[test]
2594 fn explicit_address_comes_first_and_gets_a_scheme() {
2595 let candidate = |base: &str, pid| Candidate {
2596 base: base.to_string(),
2597 pid,
2598 ide: false,
2599 };
2600 assert_eq!(
2601 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)], &no_ide()),
2602 vec![
2603 candidate("http://127.0.0.1:1234", None),
2604 candidate("http://127.0.0.1:5678", Some(10)),
2605 ]
2606 );
2607 assert_eq!(
2609 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)], &no_ide()),
2610 vec![
2611 candidate("http://127.0.0.1:1234", None),
2612 candidate("http://127.0.0.1:5678", Some(10)),
2613 ]
2614 );
2615 assert_eq!(
2618 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)], &no_ide()),
2619 vec![candidate("http://127.0.0.1:5678", Some(10))]
2620 );
2621 assert_eq!(
2623 candidate_bases_with(None, vec![(10, 5678), (10, 5678)], &no_ide()),
2624 vec![candidate("http://127.0.0.1:5678", Some(10))]
2625 );
2626 assert_eq!(
2628 candidate_bases_with(Some("https://host:9"), vec![], &no_ide()),
2629 vec![candidate("https://host:9", None)]
2630 );
2631 }
2632
2633 fn no_ide() -> std::collections::HashSet<u32> {
2636 std::collections::HashSet::new()
2637 }
2638
2639 fn http(status: u16) -> AppError {
2640 AppError::Http {
2641 status,
2642 body: String::new(),
2643 }
2644 }
2645
2646 fn missing_csrf() -> AppError {
2647 missing_csrf_with("missing CSRF token")
2648 }
2649
2650 fn missing_csrf_with(message: &str) -> AppError {
2651 AppError::Http {
2652 status: 401,
2653 body: serde_json::json!({
2654 "code": "unauthenticated",
2655 "message": message,
2656 })
2657 .to_string(),
2658 }
2659 }
2660
2661 #[test]
2662 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2663 assert!(is_missing_csrf(&missing_csrf()));
2664 assert!(is_missing_csrf(&missing_csrf_with(
2665 "unauthenticated: missing CSRF token"
2666 )));
2667 assert!(is_missing_csrf(&missing_csrf_with(
2668 " unauthenticated: missing CSRF token "
2669 )));
2670 assert!(!is_missing_csrf(&missing_csrf_with(
2671 "unauthenticated: invalid token"
2672 )));
2673 assert!(!is_missing_csrf(&AppError::Http {
2674 status: 403,
2675 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2676 }));
2677 assert!(!is_missing_csrf(&AppError::Http {
2678 status: 401,
2679 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2680 }));
2681 assert!(!is_missing_csrf(&AppError::Http {
2682 status: 401,
2683 body: "prefix: missing CSRF token".into(),
2684 }));
2685 }
2686
2687 #[test]
2688 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2689 for errors in [
2690 vec![missing_csrf(), http(403)],
2691 vec![http(401), missing_csrf()],
2692 ] {
2693 let err = select_probe_error(errors);
2694 assert!(is_actionable(&err), "{err}");
2695 assert!(!is_missing_csrf(&err), "{err}");
2696 }
2697 }
2698
2699 #[test]
2700 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2701 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2702 assert!(is_missing_csrf(&err), "{err}");
2703 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2704 }
2705
2706 #[test]
2710 fn an_auth_failure_outranks_later_transport_noise() {
2711 let err = select_probe_error(vec![
2712 http(401),
2713 AppError::Transport("connection refused".into()),
2714 ]);
2715 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2716
2717 let err = select_probe_error(vec![
2718 AppError::Transport("connection refused".into()),
2719 http(403),
2720 ]);
2721 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2722 }
2723
2724 #[test]
2727 fn the_first_auth_failure_wins() {
2728 let err = select_probe_error(vec![http(401), http(403)]);
2729 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2730 }
2731
2732 #[test]
2736 fn without_an_auth_failure_the_last_error_stands() {
2737 let err = select_probe_error(vec![
2738 AppError::Transport("first".into()),
2739 http(500),
2740 AppError::Transport("last".into()),
2741 ]);
2742 assert!(
2743 matches!(&err, AppError::Transport(m) if m == "last"),
2744 "{err}"
2745 );
2746 assert!(err.is_transient());
2747 }
2748
2749 #[test]
2752 fn a_server_error_is_not_treated_as_actionable() {
2753 let err = select_probe_error(vec![http(500), http(401)]);
2754 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2755 }
2756
2757 #[test]
2758 fn no_candidates_at_all_yields_a_generic_error() {
2759 let err = select_probe_error(Vec::new());
2760 assert!(
2761 err.to_string().contains("no local server answered"),
2762 "{err}"
2763 );
2764 }
2765
2766 fn tls_echo() -> AppError {
2770 AppError::Http {
2771 status: 400,
2772 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2773 }
2774 }
2775
2776 #[test]
2781 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2782 let err = select_probe_error(vec![
2783 AppError::Http {
2784 status: 500,
2785 body: "GetUserStatus: internal".into(),
2786 },
2787 tls_echo(),
2788 ]);
2789 assert!(
2790 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2791 "{err}"
2792 );
2793 }
2794
2795 #[test]
2800 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2801 let err = select_probe_error(vec![
2802 AppError::Transport("connection refused".into()),
2803 tls_echo(),
2804 ]);
2805 assert!(
2806 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2807 "{err}"
2808 );
2809 assert!(
2810 err.is_transient(),
2811 "the echo must not make the run non-transient: {err}"
2812 );
2813 }
2814
2815 #[test]
2818 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2819 let err = select_probe_error(vec![tls_echo()]);
2820 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2821 }
2822
2823 #[test]
2826 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2827 let err = select_probe_error(vec![
2828 AppError::Http {
2829 status: 400,
2830 body: "unknown method GetUserStatus".into(),
2831 },
2832 tls_echo(),
2833 ]);
2834 assert!(
2835 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2836 "{err}"
2837 );
2838 }
2839
2840 #[test]
2842 fn an_auth_failure_still_outranks_a_tls_echo() {
2843 let err = select_probe_error(vec![tls_echo(), http(401)]);
2844 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2845 }
2846
2847 #[test]
2848 fn every_discovered_port_is_probed_in_order() {
2849 assert_eq!(
2850 candidate_bases_with(None, vec![(10, 33875), (10, 37435)], &no_ide()),
2851 vec![
2852 Candidate {
2853 base: "http://127.0.0.1:33875".to_string(),
2854 pid: Some(10),
2855 ide: false,
2856 },
2857 Candidate {
2858 base: "http://127.0.0.1:37435".to_string(),
2859 pid: Some(10),
2860 ide: false,
2861 },
2862 ]
2863 );
2864 }
2865
2866 #[test]
2871 fn empty_discovery_yields_no_candidates() {
2872 assert!(candidate_bases_with(None, vec![], &no_ide()).is_empty());
2873 assert!(candidate_bases_with(Some(""), vec![], &no_ide()).is_empty());
2874 }
2875
2876 #[test]
2877 fn every_antigravity_product_is_recognised() {
2878 assert!(is_antigravity_process(
2880 "language_server\n",
2881 Some("/opt/antigravity/resources/bin/language_server")
2882 ));
2883 assert!(is_antigravity_process(
2885 "agy\n",
2886 Some("/home/u/.local/bin/agy")
2887 ));
2888 assert!(is_antigravity_process(
2890 "node",
2891 Some("/opt/antigravity/bin/helper")
2892 ));
2893 assert!(is_antigravity_process("antigravity", None));
2894 assert!(is_antigravity_process("agy.exe", None));
2895 assert!(is_antigravity_process("Antigravity.exe", None));
2896 assert!(is_antigravity_process("language_server.exe", None));
2897 assert!(is_antigravity_process(
2898 "language_server_windows_x64.exe",
2899 None
2900 ));
2901 assert!(is_antigravity_process(
2902 "node.exe",
2903 Some(r"C:\Users\u\AppData\Local\agy.exe")
2904 ));
2905 }
2906
2907 #[test]
2908 fn unrelated_processes_are_not_probed() {
2909 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2910 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2911 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2913 assert!(!is_antigravity_process("legacy.exe", None));
2914 assert!(!is_antigravity_process("not-agy.exe", None));
2915 assert!(!is_antigravity_process("", None));
2916 }
2917
2918 #[test]
2919 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2920 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2921 raw.extend([0, b'x' as u16]);
2922 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2923 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2924 assert_eq!(decode_windows_process_name(&[]), "");
2925 }
2926
2927 #[test]
2928 fn windows_process_filter_keeps_only_antigravity_pids() {
2929 let processes = vec![
2930 (10, "agy.exe".to_string()),
2931 (20, "language_server_windows_x64.exe".to_string()),
2932 (30, "sshd.exe".to_string()),
2933 ];
2934 let pids = matching_windows_process_ids(&processes);
2935 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
2936 }
2937
2938 #[test]
2939 fn windows_listener_filter_joins_pid_loopback_and_port() {
2940 let pids = std::collections::HashSet::from([10]);
2941 let rows = [
2942 WindowsTcpRow {
2943 local_addr: [127, 0, 0, 1],
2944 local_port: u32::from(59870u16.to_be()),
2945 pid: 10,
2946 },
2947 WindowsTcpRow {
2948 local_addr: [127, 0, 0, 1],
2949 local_port: u32::from(59868u16.to_be()),
2950 pid: 10,
2951 },
2952 WindowsTcpRow {
2953 local_addr: [127, 0, 0, 1],
2954 local_port: u32::from(59870u16.to_be()),
2955 pid: 10,
2956 },
2957 WindowsTcpRow {
2958 local_addr: [0, 0, 0, 0],
2959 local_port: u32::from(50000u16.to_be()),
2960 pid: 10,
2961 },
2962 WindowsTcpRow {
2963 local_addr: [127, 0, 0, 1],
2964 local_port: u32::from(50001u16.to_be()),
2965 pid: 99,
2966 },
2967 WindowsTcpRow {
2968 local_addr: [127, 0, 0, 1],
2969 local_port: 0,
2970 pid: 10,
2971 },
2972 ];
2973 assert_eq!(
2974 matching_windows_ports(&pids, &rows),
2975 vec![(10, 59870), (10, 59868)]
2976 );
2977 }
2978
2979 #[test]
2983 fn windows_ports_from_two_products_keep_tls_listeners_last() {
2984 let pids = std::collections::HashSet::from([10, 20]);
2985 let row = |port: u16, pid: u32| WindowsTcpRow {
2986 local_addr: [127, 0, 0, 1],
2987 local_port: u32::from(port.to_be()),
2988 pid,
2989 };
2990 let rows = [
2991 row(40000, 10),
2992 row(40001, 10),
2993 row(50000, 20),
2994 row(50001, 20),
2995 ];
2996 assert_eq!(
2997 matching_windows_ports(&pids, &rows),
2998 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2999 );
3000 }
3001
3002 #[test]
3006 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
3007 use std::collections::BTreeMap;
3008
3009 assert_eq!(
3011 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
3012 vec![(10, 59870), (10, 59868)]
3013 );
3014 assert_eq!(
3019 probe_order(BTreeMap::from([
3020 (10, vec![40000, 40001]),
3021 (20, vec![50000, 50001]),
3022 ])),
3023 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
3024 );
3025 assert_eq!(
3028 probe_order(BTreeMap::from([
3029 (10, vec![6000, 5000, 4000]),
3030 (20, vec![6000, 7000]),
3031 ])),
3032 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
3033 );
3034 assert!(probe_order(BTreeMap::new()).is_empty());
3035 }
3036
3037 #[test]
3042 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
3043 use std::collections::BTreeMap;
3044
3045 assert_eq!(
3046 probe_order(BTreeMap::from([
3047 (10, vec![40001, 40001, 40000, 40000]),
3048 (20, vec![50001, 50000]),
3049 ])),
3050 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
3051 "duplicate rows must not reorder the ranks below them"
3052 );
3053 }
3054
3055 #[test]
3060 fn a_half_started_product_is_the_documented_exception() {
3061 use std::collections::BTreeMap;
3062
3063 assert_eq!(
3064 probe_order(BTreeMap::from([
3065 (10, vec![40000]),
3066 (20, vec![50001, 50000])
3067 ])),
3068 vec![(10, 40000), (20, 50001), (20, 50000)]
3069 );
3070 }
3071
3072 #[test]
3077 fn an_override_with_no_authority_is_dropped_not_probed() {
3078 for junk in ["/", "///", "http://", "https://", " / "] {
3079 assert_eq!(
3080 candidate_bases_with(Some(junk), vec![(10, 4242)], &no_ide()),
3081 vec![Candidate {
3082 base: "http://127.0.0.1:4242".to_string(),
3083 pid: Some(10),
3084 ide: false,
3085 }],
3086 "{junk:?} should not survive as a candidate"
3087 );
3088 }
3089 assert!(candidate_bases_with(Some("/"), vec![], &no_ide()).is_empty());
3090 }
3091
3092 #[test]
3093 fn windows_table_bounds_reject_truncation_and_overflow() {
3094 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
3095 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
3096 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
3097 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
3098 assert_eq!(
3099 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
3100 None
3101 );
3102 }
3103
3104 #[cfg(target_os = "windows")]
3105 #[test]
3106 fn windows_tcp_table_parser_copies_complete_rows_only() {
3107 use std::mem::{offset_of, size_of};
3108 use windows_sys::Win32::NetworkManagement::IpHelper::{
3109 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
3110 };
3111
3112 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
3113 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
3114 let words = used.div_ceil(size_of::<u32>());
3115 let mut buffer = vec![0u32; words];
3116 let first = MIB_TCPROW_OWNER_PID {
3117 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
3118 dwLocalPort: u32::from(59868u16.to_be()),
3119 dwOwningPid: 10,
3120 ..Default::default()
3121 };
3122 let second = MIB_TCPROW_OWNER_PID {
3123 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
3124 dwLocalPort: u32::from(59870u16.to_be()),
3125 dwOwningPid: 10,
3126 ..Default::default()
3127 };
3128 unsafe {
3129 buffer.as_mut_ptr().write_unaligned(2);
3130 let rows = buffer
3131 .as_mut_ptr()
3132 .cast::<u8>()
3133 .add(offset)
3134 .cast::<MIB_TCPROW_OWNER_PID>();
3135 rows.write_unaligned(first);
3136 rows.add(1).write_unaligned(second);
3137 }
3138
3139 assert_eq!(
3140 parse_windows_tcp_rows(&buffer, used),
3141 vec![
3142 WindowsTcpRow {
3143 local_addr: [127, 0, 0, 1],
3144 local_port: u32::from(59868u16.to_be()),
3145 pid: 10,
3146 },
3147 WindowsTcpRow {
3148 local_addr: [127, 0, 0, 1],
3149 local_port: u32::from(59870u16.to_be()),
3150 pid: 10,
3151 },
3152 ]
3153 );
3154 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
3155 }
3156
3157 #[test]
3158 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
3159 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";
3162 assert_eq!(
3163 parse_lsof_pcn(output),
3164 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
3165 );
3166 }
3167
3168 #[test]
3172 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
3173 let output = concat!(
3174 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
3175 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
3176 );
3177 assert_eq!(
3178 parse_lsof_pcn(output),
3179 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
3180 "both HTTP listeners must precede both TLS listeners"
3181 );
3182 }
3183
3184 #[test]
3185 fn lsof_parser_matches_the_capitalised_macos_app_name() {
3186 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
3187 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
3188 }
3189
3190 #[test]
3191 fn lsof_parser_deduplicates_and_handles_empty_output() {
3192 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
3193 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
3194 assert!(parse_lsof_pcn("").is_empty());
3195 }
3196
3197 #[test]
3200 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
3201 let dir = tempfile::tempdir().unwrap();
3202 let cache = Cache::at(dir.path().join("usage.json"));
3203 let reason = AppError::Credentials("Antigravity: no local language server found".into());
3204
3205 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
3206 let msg = err.to_string();
3207 assert!(msg.contains("no local language server found"), "{msg}");
3208 assert!(!msg.contains("no usable cache"), "{msg}");
3209 }
3210
3211 #[test]
3212 fn unusable_cache_does_not_replace_the_live_diagnosis() {
3213 let dir = tempfile::tempdir().unwrap();
3214 let cache = Cache::at(dir.path().join("antigravity"));
3215 cache.write_payload(b"{}").unwrap();
3216
3217 let reason = AppError::Credentials("Antigravity must be running".into());
3218 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
3219 assert!(err.to_string().contains("must be running"), "{err}");
3220
3221 let original = AppError::Transport("original loopback failure".into());
3222 let err = fallback_silent(&cache, now(), original).unwrap_err();
3223 assert!(
3224 err.to_string().contains("original loopback failure"),
3225 "{err}"
3226 );
3227 }
3228
3229 #[tokio::test]
3230 async fn rpc_error_bodies_are_bounded_too() {
3231 let mut server = mockito::Server::new_async().await;
3232 let path = format!("/{STATUS_RPC}");
3233 server
3234 .mock("POST", path.as_str())
3235 .with_status(500)
3236 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
3237 .create_async()
3238 .await;
3239
3240 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
3241 .await
3242 .unwrap_err();
3243 assert!(err.to_string().contains("exceeds"), "{err}");
3244 }
3245
3246 #[test]
3247 fn blank_override_falls_through_to_discovery() {
3248 assert_eq!(
3249 candidate_bases_with(Some(" "), vec![(10, 4242)], &no_ide()),
3250 vec![Candidate {
3251 base: "http://127.0.0.1:4242".to_string(),
3252 pid: Some(10),
3253 ide: false,
3254 }]
3255 );
3256 }
3257
3258 #[test]
3261 fn only_the_ide_install_is_recognised() {
3262 assert!(is_ide_install(
3263 "/opt/antigravity-ide/resources/app/extensions/antigravity/bin/language_server_linux_x64"
3264 ));
3265 assert!(is_ide_install(
3266 "/Applications/Antigravity IDE.app/Contents/Resources/app/extensions/antigravity/bin/language_server_darwin_arm64"
3267 ));
3268 assert!(is_ide_install(
3269 r"C:\Users\u\AppData\Local\Programs\Antigravity IDE\language_server_windows_x64.exe"
3270 ));
3271 assert!(!is_ide_install(
3272 "/opt/antigravity/resources/bin/language_server"
3273 ));
3274 assert!(!is_ide_install("/home/u/.local/bin/agy"));
3275 assert!(!is_ide_install("/usr/bin/node"));
3276 }
3277
3278 #[test]
3280 fn ps_output_pairs_pids_with_full_paths() {
3281 let output = " 101 /opt/antigravity-ide/bin/language_server_linux_x64\n\
3282 202 /Applications/Antigravity IDE.app/Contents/MacOS/Antigravity IDE\n";
3283 assert_eq!(
3284 parse_ps_comm(output),
3285 vec![
3286 (
3287 101,
3288 "/opt/antigravity-ide/bin/language_server_linux_x64".to_string()
3289 ),
3290 (
3291 202,
3292 "/Applications/Antigravity IDE.app/Contents/MacOS/Antigravity IDE".to_string()
3293 ),
3294 ]
3295 );
3296 assert!(parse_ps_comm("").is_empty());
3297 assert!(parse_ps_comm("not-a-pid /tmp/x").is_empty());
3298 }
3299
3300 #[test]
3304 fn ide_owned_listeners_are_marked_but_an_explicit_override_is_not() {
3305 let ide: std::collections::HashSet<u32> = [10].into_iter().collect();
3306
3307 let candidates = candidate_bases_with(None, vec![(10, 40001), (20, 50001)], &ide);
3308 assert_eq!(candidates.len(), 2);
3309 assert!(candidates[0].ide, "pid 10 is the IDE's language server");
3310 assert!(!candidates[1].ide);
3311
3312 let candidates = candidate_bases_with(Some("127.0.0.1:40001"), vec![(10, 40001)], &ide);
3313 assert_eq!(
3314 candidates,
3315 vec![Candidate {
3316 base: "http://127.0.0.1:40001".to_string(),
3317 pid: Some(10),
3318 ide: false,
3319 }],
3320 "ANTIGRAVITY_LS_ADDRESS is an explicit target, never skipped"
3321 );
3322 }
3323
3324 fn fixture() -> (tempfile::TempDir, Cache) {
3329 let td = tempfile::TempDir::new().unwrap();
3330 let cache = Cache::at(td.path().join("antigravity"));
3331 (td, cache)
3332 }
3333
3334 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
3335 let base = server.url();
3336 cloud::Endpoints {
3337 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
3338 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
3339 token: format!("{base}/token"),
3340 }
3341 }
3342
3343 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
3345 let mut token = serde_json::json!({
3346 "access_token": "KEYRING-AT",
3347 "expiry": expiry,
3348 });
3349 if with_refresh {
3350 token["refresh_token"] = serde_json::json!("KEYRING-RT");
3351 }
3352 serde_json::json!({ "token": token }).to_string()
3353 }
3354
3355 const EXPIRED: &str = "2026-07-22T11:00:00Z";
3357 const VALID: &str = "2026-07-22T13:00:00Z";
3359
3360 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
3362 RemoteOverride {
3363 credential: SavedCredential::Blob(blob),
3364 endpoints: Some(eps),
3365 local_bases: Some(vec![]),
3366 ..Default::default()
3367 }
3368 }
3369
3370 fn bare_summary() -> String {
3373 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
3374 v["response"].to_string()
3375 }
3376
3377 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
3378 server
3379 .mock("POST", "/daily/quota")
3380 .match_header("authorization", format!("Bearer {bearer}").as_str())
3381 .with_status(200)
3382 .with_body(bare_summary())
3383 }
3384
3385 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
3386 server
3387 .mock("POST", "/token")
3388 .with_status(200)
3389 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
3390 }
3391
3392 fn test_oauth() -> cloud::OauthClient {
3393 cloud::OauthClient {
3394 id: "test-client".into(),
3395 secret: "test-client-secret".into(),
3396 }
3397 }
3398
3399 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
3400 fetch_snapshot_at(
3401 &reqwest::Client::new(),
3402 cache,
3403 ttl,
3404 Some(&test_oauth()),
3405 remote,
3406 now(),
3407 )
3408 .await
3409 }
3410
3411 #[tokio::test]
3412 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
3413 let mut server = mockito::Server::new_async().await;
3414 let eps = endpoints(&server);
3415 let quota = quota_mock(&mut server, "KEYRING-AT")
3416 .expect(1)
3417 .create_async()
3418 .await;
3419 let plan = server
3420 .mock("POST", "/daily/plan")
3421 .match_header("authorization", "Bearer KEYRING-AT")
3422 .with_status(200)
3423 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3424 .expect(1)
3425 .create_async()
3426 .await;
3427 let token = token_mock(&mut server).expect(0).create_async().await;
3429 let blob = keyring_blob(VALID, true);
3430 let (_td, cache) = fixture();
3431
3432 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
3433 .await
3434 .expect("remote path yields a snapshot");
3435
3436 quota.assert_async().await;
3437 plan.assert_async().await;
3438 token.assert_async().await;
3439 let snap = outcome.snapshot;
3440 assert!(!outcome.stale);
3441 assert_eq!(snap.source, AntigravitySource::Remote);
3442 assert_eq!(snap.plan, "Pro");
3443 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
3444 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
3445 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
3446 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
3447 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3448 assert_eq!(snap.account, format!("acct:{fingerprint}"));
3449 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
3450
3451 let cached = parse_cache_at(
3453 &cache
3454 .fresh_payload(Duration::from_secs(60))
3455 .unwrap()
3456 .unwrap(),
3457 None,
3458 now(),
3459 )
3460 .unwrap();
3461 assert_eq!(cached.source, AntigravitySource::Remote);
3462 assert_eq!(cached.account, snap.account);
3463 }
3464
3465 #[tokio::test]
3470 async fn agys_missing_csrf_response_uses_the_saved_session() {
3471 let mut server = mockito::Server::new_async().await;
3472 let eps = endpoints(&server);
3473 let root = server
3474 .mock("GET", "/")
3475 .with_status(404)
3476 .expect(1)
3477 .create_async()
3478 .await;
3479 let status_path = format!("/{STATUS_RPC}");
3480 let status = server
3481 .mock("POST", status_path.as_str())
3482 .with_status(401)
3483 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3484 .expect(1)
3485 .create_async()
3486 .await;
3487 let quota = quota_mock(&mut server, "KEYRING-AT")
3488 .expect(1)
3489 .create_async()
3490 .await;
3491 let plan = server
3492 .mock("POST", "/daily/plan")
3493 .match_header("authorization", "Bearer KEYRING-AT")
3494 .with_status(200)
3495 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3496 .expect(1)
3497 .create_async()
3498 .await;
3499 let blob = keyring_blob(VALID, true);
3500 let (_td, cache) = fixture();
3501
3502 let outcome = run(
3503 &cache,
3504 RemoteOverride {
3505 credential: SavedCredential::Blob(&blob),
3506 endpoints: Some(&eps),
3507 local_bases: Some(vec![server.url().into()]),
3508 ..Default::default()
3509 },
3510 Duration::ZERO,
3511 )
3512 .await
3513 .expect("saved session bypasses agy's unusable local RPC");
3514
3515 root.assert_async().await;
3516 status.assert_async().await;
3517 quota.assert_async().await;
3518 plan.assert_async().await;
3519 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3520 }
3521
3522 #[tokio::test]
3526 async fn agys_prefixed_missing_csrf_response_uses_the_saved_session() {
3527 let mut server = mockito::Server::new_async().await;
3528 let eps = endpoints(&server);
3529 let root = server
3530 .mock("GET", "/")
3531 .with_status(404)
3532 .expect(1)
3533 .create_async()
3534 .await;
3535 let status_path = format!("/{STATUS_RPC}");
3536 let status = server
3537 .mock("POST", status_path.as_str())
3538 .with_status(401)
3539 .with_body(
3540 r#"{"code":"unauthenticated","message":"unauthenticated: missing CSRF token"}"#,
3541 )
3542 .expect(1)
3543 .create_async()
3544 .await;
3545 let quota = quota_mock(&mut server, "KEYRING-AT")
3546 .expect(1)
3547 .create_async()
3548 .await;
3549 let plan = server
3550 .mock("POST", "/daily/plan")
3551 .match_header("authorization", "Bearer KEYRING-AT")
3552 .with_status(200)
3553 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3554 .expect(1)
3555 .create_async()
3556 .await;
3557 let blob = keyring_blob(VALID, true);
3558 let (_td, cache) = fixture();
3559
3560 let outcome = run(
3561 &cache,
3562 RemoteOverride {
3563 credential: SavedCredential::Blob(&blob),
3564 endpoints: Some(&eps),
3565 local_bases: Some(vec![server.url().into()]),
3566 ..Default::default()
3567 },
3568 Duration::ZERO,
3569 )
3570 .await
3571 .expect("saved session bypasses agy 1.3.1's prefixed missing-CSRF RPC");
3572
3573 root.assert_async().await;
3574 status.assert_async().await;
3575 quota.assert_async().await;
3576 plan.assert_async().await;
3577 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3578 }
3579
3580 #[tokio::test]
3581 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
3582 let mut server = mockito::Server::new_async().await;
3583 let eps = endpoints(&server);
3584 let root = server
3585 .mock("GET", "/")
3586 .with_status(404)
3587 .expect(1)
3588 .create_async()
3589 .await;
3590 let status_path = format!("/{STATUS_RPC}");
3591 let status = server
3592 .mock("POST", status_path.as_str())
3593 .with_status(401)
3594 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3595 .expect(1)
3596 .create_async()
3597 .await;
3598 let quota = quota_mock(&mut server, "KEYRING-AT")
3599 .expect(0)
3600 .create_async()
3601 .await;
3602 let (_td, cache) = fixture();
3603
3604 let error = run(
3605 &cache,
3606 RemoteOverride {
3607 credential: SavedCredential::Absent,
3608 endpoints: Some(&eps),
3609 local_bases: Some(vec![server.url().into()]),
3610 ..Default::default()
3611 },
3612 Duration::ZERO,
3613 )
3614 .await
3615 .expect_err("neither the local RPC nor a saved session is usable");
3616
3617 root.assert_async().await;
3618 status.assert_async().await;
3619 quota.assert_async().await;
3620 let message = error.to_string();
3621 assert!(message.contains("requires a CSRF token"), "{message}");
3622 assert!(message.contains("saved Google session"), "{message}");
3623 }
3624
3625 #[tokio::test]
3630 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
3631 let mut server1 = mockito::Server::new_async().await;
3632 let mut server2 = mockito::Server::new_async().await;
3633 let eps = endpoints(&server1);
3634 let root1 = server1
3635 .mock("GET", "/")
3636 .with_status(404)
3637 .expect(1)
3638 .create_async()
3639 .await;
3640 let status_path = format!("/{STATUS_RPC}");
3641 let status1 = server1
3642 .mock("POST", status_path.as_str())
3643 .with_status(401)
3644 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3645 .expect(1)
3646 .create_async()
3647 .await;
3648
3649 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
3651 let status2 = server2
3652 .mock("POST", status_path.as_str())
3653 .expect(0)
3654 .create_async()
3655 .await;
3656
3657 let (_td, cache) = fixture();
3658 let error = run(
3659 &cache,
3660 RemoteOverride {
3661 credential: SavedCredential::Absent,
3662 endpoints: Some(&eps),
3663 local_bases: Some(vec![
3664 Candidate {
3665 base: server1.url(),
3666 pid: Some(42),
3667 ide: false,
3668 },
3669 Candidate {
3670 base: server2.url(),
3671 pid: Some(42),
3672 ide: false,
3673 },
3674 ]),
3675 ..Default::default()
3676 },
3677 Duration::ZERO,
3678 )
3679 .await
3680 .expect_err("aborted probe returns agy missing csrf diagnosis");
3681
3682 root1.assert_async().await;
3683 status1.assert_async().await;
3684 root2.assert_async().await;
3685 status2.assert_async().await;
3686 let message = error.to_string();
3687 assert!(message.contains("requires a CSRF token"), "{message}");
3688 }
3689
3690 #[tokio::test]
3697 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
3698 let mut agy = mockito::Server::new_async().await;
3699 let mut ide = mockito::Server::new_async().await;
3700 let eps = endpoints(&agy);
3701 let agy_root = agy
3702 .mock("GET", "/")
3703 .with_status(404)
3704 .expect(1)
3705 .create_async()
3706 .await;
3707 let status_path = format!("/{STATUS_RPC}");
3708 let agy_status = agy
3709 .mock("POST", status_path.as_str())
3710 .with_status(401)
3711 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3712 .expect(1)
3713 .create_async()
3714 .await;
3715
3716 let ide_root = ide
3717 .mock("GET", "/")
3718 .with_status(404)
3719 .expect(1)
3720 .create_async()
3721 .await;
3722 let ide_status = ide
3723 .mock("POST", status_path.as_str())
3724 .with_status(200)
3725 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3726 .expect(1)
3727 .create_async()
3728 .await;
3729 let quota_path = format!("/{QUOTA_RPC}");
3730 let ide_quota = ide
3731 .mock("POST", quota_path.as_str())
3732 .with_status(200)
3733 .with_body(QUOTA_JSON)
3734 .expect(1)
3735 .create_async()
3736 .await;
3737
3738 let (_td, cache) = fixture();
3739 let outcome = run(
3740 &cache,
3741 RemoteOverride {
3742 credential: SavedCredential::Absent,
3743 endpoints: Some(&eps),
3744 local_bases: Some(vec![
3745 Candidate {
3746 base: agy.url(),
3747 pid: Some(42),
3748 ide: false,
3749 },
3750 Candidate {
3751 base: ide.url(),
3752 pid: Some(7),
3753 ide: false,
3754 },
3755 ]),
3756 ..Default::default()
3757 },
3758 Duration::ZERO,
3759 )
3760 .await
3761 .expect("the other product's server still answers");
3762
3763 agy_root.assert_async().await;
3764 agy_status.assert_async().await;
3765 ide_root.assert_async().await;
3766 ide_status.assert_async().await;
3767 ide_quota.assert_async().await;
3768 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3769 assert_eq!(outcome.snapshot.plan, "Pro");
3770 }
3771
3772 #[tokio::test]
3775 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
3776 let mut server = mockito::Server::new_async().await;
3777 let eps = endpoints(&server);
3778 let _quota = quota_mock(&mut server, "NEW-AT")
3779 .expect(2)
3780 .create_async()
3781 .await;
3782 let token = token_mock(&mut server).expect(1).create_async().await;
3783 let blob = keyring_blob(EXPIRED, true);
3784 let (_td, cache) = fixture();
3785
3786 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
3787 .await
3788 .expect("refresh, then quota");
3789 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
3790 token.assert_async().await;
3791
3792 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3793 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
3794 .expect("refreshed token persisted under the keyring session's fingerprint");
3795 assert_eq!(persisted.access_token, "NEW-AT");
3796
3797 run(&cache, remote(&blob, &eps), Duration::ZERO)
3799 .await
3800 .expect("persisted token reused");
3801 token.assert_async().await;
3802 }
3803
3804 #[tokio::test]
3807 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
3808 let mut server = mockito::Server::new_async().await;
3809 let eps = endpoints(&server);
3810 let rejected = server
3811 .mock("POST", "/daily/quota")
3812 .match_header("authorization", "Bearer KEYRING-AT")
3813 .with_status(401)
3814 .expect(1)
3815 .create_async()
3816 .await;
3817 let accepted = quota_mock(&mut server, "NEW-AT")
3818 .expect(1)
3819 .create_async()
3820 .await;
3821 let token = token_mock(&mut server).expect(1).create_async().await;
3822 let blob = keyring_blob(VALID, true);
3823 let (_td, cache) = fixture();
3824
3825 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3826 .await
3827 .expect("retry with the refreshed token succeeds");
3828
3829 rejected.assert_async().await;
3830 accepted.assert_async().await;
3831 token.assert_async().await;
3832 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3833 }
3834
3835 #[tokio::test]
3838 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
3839 let mut server = mockito::Server::new_async().await;
3840 let eps = endpoints(&server);
3841 let _token = server
3842 .mock("POST", "/token")
3843 .with_status(400)
3844 .with_body(r#"{"error":"invalid_grant"}"#)
3845 .create_async()
3846 .await;
3847 let blob = keyring_blob(EXPIRED, true);
3848 let (_td, cache) = fixture();
3849
3850 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3851 .await
3852 .expect_err("no cache to fall back on");
3853
3854 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3855 let rendered = err.to_string();
3856 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
3857 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
3858 }
3859
3860 #[tokio::test]
3863 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
3864 let mut server = mockito::Server::new_async().await;
3865 let eps = endpoints(&server);
3866 let _quota = server
3867 .mock("POST", "/daily/quota")
3868 .with_status(401)
3869 .create_async()
3870 .await;
3871 let token = token_mock(&mut server).expect(1).create_async().await;
3872 let blob = keyring_blob(EXPIRED, true);
3873 let (_td, cache) = fixture();
3874
3875 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3876 .await
3877 .expect_err("rejected after refresh");
3878
3879 token.assert_async().await;
3880 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3881 assert!(err.to_string().contains("sign in again"), "{err}");
3882 }
3883
3884 #[cfg(unix)]
3888 #[tokio::test]
3889 async fn a_subscription_403_surfaces_the_plan_message_without_a_refresh_or_agy_run() {
3890 let mut server = mockito::Server::new_async().await;
3891 let eps = endpoints(&server);
3892 let quota = server
3893 .mock("POST", "/daily/quota")
3894 .match_header("authorization", "Bearer KEYRING-AT")
3895 .with_status(403)
3896 .with_body(
3897 r#"{"error":{"code":403,"status":"PERMISSION_DENIED","details":
3898 [{"reason":"SUBSCRIPTION_REQUIRED","domain":"cloud.google.com"}]}}"#,
3899 )
3900 .expect(1)
3901 .create_async()
3902 .await;
3903 let prod = server
3904 .mock("POST", "/prod/quota")
3905 .expect(0)
3906 .create_async()
3907 .await;
3908 let token = token_mock(&mut server).expect(0).create_async().await;
3909 let (td, cache) = fixture();
3910 let file = td.path().join("credential");
3911 std::fs::write(&file, keyring_blob(VALID, true)).unwrap();
3912 let runs = td.path().join("runs");
3913 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
3914
3915 let err = run_without_oauth(
3916 &cache,
3917 RemoteOverride {
3918 credential: SavedCredential::File(&file),
3919 endpoints: Some(&eps),
3920 local_bases: Some(vec![]),
3921 agy: super::super::agy::AgyCommand::At(agy),
3922 ..Default::default()
3923 },
3924 )
3925 .await
3926 .expect_err("a plan without quota is an error, not a snapshot");
3927
3928 quota.assert_async().await;
3929 prod.assert_async().await;
3930 token.assert_async().await;
3931 match err {
3932 AppError::Http { status, body } => {
3933 assert_eq!(status, 403);
3934 assert!(body.contains("paid subscription"), "{body}");
3935 }
3936 other => panic!("expected Http, got {other:?}"),
3937 }
3938 assert!(
3939 !runs.exists(),
3940 "a subscription verdict must not spawn `agy`"
3941 );
3942 }
3943
3944 #[tokio::test]
3948 async fn a_non_subscription_403_keeps_the_refresh_and_retry_behavior() {
3949 let mut server = mockito::Server::new_async().await;
3950 let eps = endpoints(&server);
3951 let rejected = server
3952 .mock("POST", "/daily/quota")
3953 .match_header("authorization", "Bearer KEYRING-AT")
3954 .with_status(403)
3955 .with_body(r#"{"error":{"message":"acl detail"}}"#)
3956 .expect(1)
3957 .create_async()
3958 .await;
3959 let retried = server
3960 .mock("POST", "/daily/quota")
3961 .match_header("authorization", "Bearer NEW-AT")
3962 .with_status(403)
3963 .with_body(r#"{"error":{"message":"acl detail"}}"#)
3964 .expect(1)
3965 .create_async()
3966 .await;
3967 let token = token_mock(&mut server).expect(1).create_async().await;
3968 let blob = keyring_blob(VALID, true);
3969 let (_td, cache) = fixture();
3970
3971 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3972 .await
3973 .expect_err("still a rejected session");
3974
3975 rejected.assert_async().await;
3976 retried.assert_async().await;
3977 token.assert_async().await;
3978 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3979 assert!(err.to_string().contains("sign in again"), "{err}");
3980 }
3981
3982 async fn run_without_oauth(cache: &Cache, remote: RemoteOverride<'_>) -> Result<FetchOutcome> {
3984 fetch_snapshot_at(
3985 &reqwest::Client::new(),
3986 cache,
3987 Duration::ZERO,
3988 None,
3989 remote,
3990 now(),
3991 )
3992 .await
3993 }
3994
3995 #[cfg(unix)]
4000 fn fake_agy(
4001 dir: &std::path::Path,
4002 runs: &std::path::Path,
4003 file: &std::path::Path,
4004 blob: &str,
4005 ) -> std::path::PathBuf {
4006 let path = dir.join("agy");
4007 super::super::agy::test_support::write_script(
4008 &path,
4009 &format!(
4010 "[ \"$1\" = -p ] && exit 1\necho \"$@\" >> '{}'\nprintf '%s' '{}' > '{}'",
4011 runs.display(),
4012 blob,
4013 file.display()
4014 ),
4015 );
4016 path
4017 }
4018
4019 #[cfg(unix)]
4022 #[tokio::test]
4023 async fn an_expired_session_is_renewed_by_running_agy() {
4024 let mut server = mockito::Server::new_async().await;
4025 let eps = endpoints(&server);
4026 let quota = quota_mock(&mut server, "KEYRING-AT").create_async().await;
4027 let _plan = server
4028 .mock("POST", "/daily/plan")
4029 .with_status(200)
4030 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
4031 .create_async()
4032 .await;
4033 let (td, cache) = fixture();
4034 let file = td.path().join("credential");
4035 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
4036 let runs = td.path().join("runs");
4037 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
4038
4039 let out = run_without_oauth(
4040 &cache,
4041 RemoteOverride {
4042 credential: SavedCredential::File(&file),
4043 endpoints: Some(&eps),
4044 local_bases: Some(vec![]),
4045 agy: super::super::agy::AgyCommand::At(agy),
4046 ..Default::default()
4047 },
4048 )
4049 .await
4050 .expect("agy renewed the session");
4051
4052 quota.assert_async().await;
4053 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4054 assert_eq!(std::fs::read_to_string(runs).unwrap().trim(), "models");
4055 }
4056
4057 #[cfg(unix)]
4060 #[tokio::test]
4061 async fn agy_that_does_not_renew_keeps_the_original_error_and_runs_once() {
4062 let server = mockito::Server::new_async().await;
4063 let eps = endpoints(&server);
4064 let (td, cache) = fixture();
4065 let file = td.path().join("credential");
4066 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
4067 let runs = td.path().join("runs");
4068 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(EXPIRED, true));
4069 let remote = || RemoteOverride {
4070 credential: SavedCredential::File(&file),
4071 endpoints: Some(&eps),
4072 local_bases: Some(vec![]),
4073 agy: super::super::agy::AgyCommand::At(agy.clone()),
4074 ..Default::default()
4075 };
4076
4077 for _ in 0..2 {
4078 let err = run_without_oauth(&cache, remote())
4079 .await
4080 .expect_err("still expired");
4081 assert!(err.to_string().contains("no OAuth client"), "{err}");
4082 }
4083 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4084 }
4085
4086 #[tokio::test]
4089 async fn without_an_agy_command_an_expired_session_is_left_alone() {
4090 let server = mockito::Server::new_async().await;
4091 let eps = endpoints(&server);
4092 let blob = keyring_blob(EXPIRED, true);
4093 let (_td, cache) = fixture();
4094
4095 let err = run_without_oauth(&cache, remote(&blob, &eps))
4096 .await
4097 .expect_err("nothing renews it");
4098 assert!(err.to_string().contains("no OAuth client"), "{err}");
4099 }
4100
4101 #[cfg(unix)]
4104 const AGY_QUOTA: &str = r#"{"status":"SUCCESS","num_turns":0,"command":{"name":"usage","data":{"groups":[
4105{"name":"Gemini Models","buckets":[
4106 {"id":"gemini-weekly","name":"Weekly Limit Remaining","window":"weekly","remaining_fraction":0.63,"reset_time":"2026-07-28T01:11:19Z"},
4107 {"id":"gemini-5h","name":"Five Hour Limit Remaining","window":"5h","remaining_fraction":0.49,"reset_time":"2026-07-22T15:25:45Z"}]},
4108{"name":"Claude and GPT models","buckets":[
4109 {"id":"3p-weekly","name":"Weekly Limit Remaining","window":"weekly","remaining_fraction":1,"reset_time":"2026-07-29T12:00:00Z"},
4110 {"id":"3p-5h","name":"Five Hour Limit Remaining","window":"5h","remaining_fraction":1,"reset_time":"2026-07-22T17:00:00Z"}]},
4111{"name":"Gemini 4 Argon","buckets":[
4112 {"id":"gemini-argon-wtu-weekly","name":"Gemini 4 Argon Weekly Limit","window":"weekly","remaining_fraction":0,"reset_time":"2026-07-28T01:10:47Z"}]}]}}}"#;
4113
4114 #[cfg(unix)]
4117 fn agy_reading(dir: &std::path::Path, runs: &std::path::Path) -> std::path::PathBuf {
4118 let path = dir.join("agy");
4119 super::super::agy::test_support::write_script(
4120 &path,
4121 &format!(
4122 "echo \"$@\" >> '{}'\n[ \"$1\" = -p ] || exit 1\ncat <<'EOF'\n{AGY_QUOTA}\nEOF",
4123 runs.display()
4124 ),
4125 );
4126 path
4127 }
4128
4129 #[cfg(unix)]
4133 #[tokio::test]
4134 async fn agys_own_reading_answers_before_the_cloud_code_api() {
4135 let mut server = mockito::Server::new_async().await;
4136 let eps = endpoints(&server);
4137 let daily = server
4138 .mock("POST", "/daily/quota")
4139 .expect(0)
4140 .create_async()
4141 .await;
4142 let prod = server
4143 .mock("POST", "/prod/quota")
4144 .expect(0)
4145 .create_async()
4146 .await;
4147 let plan = server
4148 .mock("POST", "/daily/plan")
4149 .match_header("authorization", "Bearer KEYRING-AT")
4150 .with_status(200)
4151 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4152 .expect(1)
4153 .create_async()
4154 .await;
4155 let blob = keyring_blob(VALID, true);
4156 let (td, cache) = fixture();
4157 let runs = td.path().join("runs");
4158 let agy = agy_reading(td.path(), &runs);
4159
4160 let out = run(
4161 &cache,
4162 RemoteOverride {
4163 agy: super::super::agy::AgyCommand::At(agy),
4164 ..remote(&blob, &eps)
4165 },
4166 Duration::from_secs(60),
4167 )
4168 .await
4169 .expect("agy answered");
4170
4171 daily.assert_async().await;
4172 prod.assert_async().await;
4173 plan.assert_async().await;
4174 let snap = out.snapshot;
4175 assert_eq!(snap.source, AntigravitySource::Agy);
4176 assert_eq!(snap.plan, "Pro");
4177 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 51);
4178 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 37);
4179 assert_eq!(
4180 snap.third_party_session.as_ref().unwrap().utilization_pct,
4181 0
4182 );
4183 assert_eq!(snap.model_pools[0].window.utilization_pct, 100);
4184 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
4185 assert_eq!(snap.account, format!("acct:{fingerprint}"));
4186 let cached = parse_cache_at(
4187 &cache
4188 .fresh_payload(Duration::from_secs(60))
4189 .unwrap()
4190 .unwrap(),
4191 None,
4192 now(),
4193 )
4194 .unwrap();
4195 assert_eq!(cached.source, AntigravitySource::Agy);
4196 assert_eq!(
4197 std::fs::read_to_string(runs).unwrap().trim(),
4198 "-p /quota --output-format json --log-file /dev/null"
4199 );
4200 }
4201
4202 #[cfg(unix)]
4205 #[tokio::test]
4206 async fn agys_reading_needs_no_good_token() {
4207 let mut server = mockito::Server::new_async().await;
4208 let eps = endpoints(&server);
4209 let plan = server
4210 .mock("POST", "/daily/plan")
4211 .expect(0)
4212 .create_async()
4213 .await;
4214 let token = token_mock(&mut server).expect(0).create_async().await;
4215 let blob = keyring_blob(EXPIRED, true);
4216 let (td, cache) = fixture();
4217 let runs = td.path().join("runs");
4218 let agy = agy_reading(td.path(), &runs);
4219
4220 let out = run_without_oauth(
4221 &cache,
4222 RemoteOverride {
4223 agy: super::super::agy::AgyCommand::At(agy),
4224 ..remote(&blob, &eps)
4225 },
4226 )
4227 .await
4228 .expect("agy answered");
4229
4230 plan.assert_async().await;
4231 token.assert_async().await;
4232 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4233 assert_eq!(out.snapshot.plan, "Antigravity");
4234 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4235 }
4236
4237 #[cfg(unix)]
4240 #[tokio::test]
4241 async fn when_agy_cannot_read_the_quota_the_cloud_code_api_answers() {
4242 let mut server = mockito::Server::new_async().await;
4243 let eps = endpoints(&server);
4244 let quota = quota_mock(&mut server, "KEYRING-AT")
4245 .expect(2)
4246 .create_async()
4247 .await;
4248 let blob = keyring_blob(VALID, true);
4249 let (td, cache) = fixture();
4250 let runs = td.path().join("runs");
4251 let agy = td.path().join("agy");
4252 super::super::agy::test_support::write_script(
4253 &agy,
4254 &format!("echo \"$@\" >> '{}'\nexit 1", runs.display()),
4255 );
4256 let remote = || RemoteOverride {
4257 agy: super::super::agy::AgyCommand::At(agy.clone()),
4258 ..remote(&blob, &eps)
4259 };
4260
4261 for _ in 0..2 {
4262 let out = run(&cache, remote(), Duration::ZERO)
4263 .await
4264 .expect("the Cloud Code API answered");
4265 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
4266 std::fs::remove_file(cache.payload_path()).unwrap();
4268 }
4269 quota.assert_async().await;
4270 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4271 }
4272
4273 #[cfg(unix)]
4276 #[tokio::test]
4277 async fn agys_reading_is_taken_at_most_every_quota_interval() {
4278 let server = mockito::Server::new_async().await;
4279 let eps = endpoints(&server);
4280 let blob = keyring_blob(VALID, true);
4281 let (td, cache) = fixture();
4282 let runs = td.path().join("runs");
4283 let agy = agy_reading(td.path(), &runs);
4284 for _ in 0..2 {
4285 let out = run(
4286 &cache,
4287 RemoteOverride {
4288 agy: super::super::agy::AgyCommand::At(agy.clone()),
4289 ..remote(&blob, &eps)
4290 },
4291 Duration::ZERO,
4292 )
4293 .await
4294 .expect("a snapshot");
4295 assert_eq!(out.snapshot.source, AntigravitySource::Agy);
4296 }
4297 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
4298 }
4299
4300 #[tokio::test]
4302 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
4303 let mut server = mockito::Server::new_async().await;
4304 let eps = endpoints(&server);
4305 let token = token_mock(&mut server).expect(0).create_async().await;
4306 let blob = keyring_blob(EXPIRED, false);
4307 let (_td, cache) = fixture();
4308
4309 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
4310 .await
4311 .expect_err("nothing to refresh with");
4312
4313 token.assert_async().await;
4314 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4315 assert!(err.to_string().contains("sign in again"), "{err}");
4316 }
4317
4318 #[tokio::test]
4321 async fn no_saved_session_extends_the_no_local_server_error() {
4322 let mut server = mockito::Server::new_async().await;
4323 let eps = endpoints(&server);
4324 let quota = quota_mock(&mut server, "KEYRING-AT")
4325 .expect(0)
4326 .create_async()
4327 .await;
4328 let (_td, cache) = fixture();
4329
4330 let err = run(
4331 &cache,
4332 RemoteOverride {
4333 credential: SavedCredential::Absent,
4334 endpoints: Some(&eps),
4335 local_bases: Some(vec![]),
4336 ..Default::default()
4337 },
4338 Duration::ZERO,
4339 )
4340 .await
4341 .expect_err("no source at all");
4342
4343 quota.assert_async().await;
4344 assert!(matches!(err, AppError::Credentials(_)), "{err}");
4345 let rendered = err.to_string();
4346 assert!(rendered.contains("no local server found"), "{rendered}");
4347 assert!(rendered.contains("saved Google session"), "{rendered}");
4348 }
4349
4350 #[tokio::test]
4353 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
4354 let mut server = mockito::Server::new_async().await;
4355 let eps = endpoints(&server);
4356 let status_path = format!("/{STATUS_RPC}");
4357 let _status = server
4358 .mock("POST", status_path.as_str())
4359 .with_status(401)
4360 .create_async()
4361 .await;
4362 let quota = quota_mock(&mut server, "KEYRING-AT")
4363 .expect(0)
4364 .create_async()
4365 .await;
4366 let blob = keyring_blob(VALID, true);
4367 let (_td, cache) = fixture();
4368
4369 let err = run(
4370 &cache,
4371 RemoteOverride {
4372 credential: SavedCredential::Blob(&blob),
4373 endpoints: Some(&eps),
4374 local_bases: Some(vec![server.url().into()]),
4375 ..Default::default()
4376 },
4377 Duration::ZERO,
4378 )
4379 .await
4380 .expect_err("signed-out local server");
4381
4382 quota.assert_async().await;
4383 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
4384 }
4385
4386 #[tokio::test]
4392 async fn a_signed_out_ide_language_server_is_skipped_so_the_saved_session_answers() {
4393 let mut ide = mockito::Server::new_async().await;
4394 let mut agy = mockito::Server::new_async().await;
4395 let eps = endpoints(&agy);
4396 let ide_root = ide.mock("GET", "/").expect(0).create_async().await;
4398 let status_path = format!("/{STATUS_RPC}");
4399 let ide_status = ide
4400 .mock("POST", status_path.as_str())
4401 .with_status(403)
4402 .expect(0)
4403 .create_async()
4404 .await;
4405 let agy_root = agy
4407 .mock("GET", "/")
4408 .with_status(404)
4409 .expect(1)
4410 .create_async()
4411 .await;
4412 let agy_status = agy
4413 .mock("POST", status_path.as_str())
4414 .with_status(401)
4415 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
4416 .expect(1)
4417 .create_async()
4418 .await;
4419 let quota = quota_mock(&mut agy, "KEYRING-AT")
4420 .expect(1)
4421 .create_async()
4422 .await;
4423 let plan = agy
4424 .mock("POST", "/daily/plan")
4425 .match_header("authorization", "Bearer KEYRING-AT")
4426 .with_status(200)
4427 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
4428 .expect(1)
4429 .create_async()
4430 .await;
4431 let blob = keyring_blob(VALID, true);
4432 let (_td, cache) = fixture();
4433
4434 let outcome = run(
4435 &cache,
4436 RemoteOverride {
4437 credential: SavedCredential::Blob(&blob),
4438 endpoints: Some(&eps),
4439 local_bases: Some(vec![
4440 Candidate {
4441 base: ide.url(),
4442 pid: Some(7),
4443 ide: true,
4444 },
4445 Candidate {
4446 base: agy.url(),
4447 pid: Some(42),
4448 ide: false,
4449 },
4450 ]),
4451 ide_language_server: false,
4452 ..Default::default()
4453 },
4454 Duration::ZERO,
4455 )
4456 .await
4457 .expect("with the IDE skipped, agy's missing-CSRF answer unlocks the saved session");
4458
4459 ide_root.assert_async().await;
4460 ide_status.assert_async().await;
4461 agy_root.assert_async().await;
4462 agy_status.assert_async().await;
4463 quota.assert_async().await;
4464 plan.assert_async().await;
4465 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4466 assert_eq!(outcome.snapshot.plan, "Pro");
4467 }
4468
4469 #[tokio::test]
4473 async fn skipping_the_only_ide_candidate_falls_through_to_the_saved_session() {
4474 let mut ide = mockito::Server::new_async().await;
4475 let eps = endpoints(&ide);
4476 let ide_root = ide.mock("GET", "/").expect(0).create_async().await;
4477 let status_path = format!("/{STATUS_RPC}");
4478 let ide_status = ide
4479 .mock("POST", status_path.as_str())
4480 .expect(0)
4481 .create_async()
4482 .await;
4483 let quota = quota_mock(&mut ide, "KEYRING-AT")
4484 .expect(1)
4485 .create_async()
4486 .await;
4487 let plan = ide
4488 .mock("POST", "/daily/plan")
4489 .with_status(200)
4490 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
4491 .expect(1)
4492 .create_async()
4493 .await;
4494 let blob = keyring_blob(VALID, true);
4495 let (_td, cache) = fixture();
4496
4497 let outcome = run(
4498 &cache,
4499 RemoteOverride {
4500 credential: SavedCredential::Blob(&blob),
4501 endpoints: Some(&eps),
4502 local_bases: Some(vec![Candidate {
4503 base: ide.url(),
4504 pid: Some(7),
4505 ide: true,
4506 }]),
4507 ide_language_server: false,
4508 ..Default::default()
4509 },
4510 Duration::ZERO,
4511 )
4512 .await
4513 .expect("an IDE-only machine with the IDE skipped uses the saved session");
4514
4515 ide_root.assert_async().await;
4516 ide_status.assert_async().await;
4517 quota.assert_async().await;
4518 plan.assert_async().await;
4519 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
4520 }
4521
4522 #[tokio::test]
4525 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
4526 let mut server = mockito::Server::new_async().await;
4527 let eps = endpoints(&server);
4528 for path in ["/daily/quota", "/prod/quota"] {
4529 server
4530 .mock("POST", path)
4531 .with_status(500)
4532 .with_body("boom")
4533 .create_async()
4534 .await;
4535 }
4536 let blob = keyring_blob(VALID, true);
4537 let (_td, cache) = fixture();
4538 cache.ensure_dir().unwrap();
4539 let mut earlier = parsed();
4540 earlier.source = AntigravitySource::Remote;
4541 earlier.account = "acct:earlier".into();
4542 cache
4543 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
4544 .unwrap();
4545
4546 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
4547 .await
4548 .expect("stale cache stands in");
4549
4550 assert!(outcome.stale);
4551 assert_eq!(outcome.snapshot, earlier);
4552 assert!(
4553 matches!(outcome.last_error, Some((500, _))),
4554 "{:?}",
4555 outcome.last_error
4556 );
4557 }
4558
4559 fn argon(pct: i32, reset: &str) -> ModelPoolWindow {
4560 ModelPoolWindow {
4561 name: "Gemini 4 Argon".into(),
4562 window: UsageWindow {
4563 utilization_pct: pct,
4564 resets_at: Some(
4565 DateTime::parse_from_rfc3339(reset)
4566 .unwrap()
4567 .with_timezone(&Utc),
4568 ),
4569 window_duration: chrono::Duration::days(7),
4570 },
4571 }
4572 }
4573
4574 #[test]
4575 fn model_pools_round_trip_through_the_cache() {
4576 let mut snap = parsed();
4577 snap.model_pools.push(argon(60, "2026-07-28T17:00:00Z"));
4578 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4579 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
4580
4581 let mut legacy = snap_to_json(&parsed());
4583 legacy.as_object_mut().unwrap().remove("model_pools");
4584 let legacy = serde_json::to_vec(&legacy).unwrap();
4585 assert!(
4586 parse_cache_at(&legacy, None, now())
4587 .unwrap()
4588 .model_pools
4589 .is_empty()
4590 );
4591
4592 let mut torn = snap_to_json(&snap);
4594 torn["model_pools"][0]["window"] = "monthly".into();
4595 let torn = serde_json::to_vec(&torn).unwrap();
4596 assert!(parse_cache_at(&torn, None, now()).is_err());
4597
4598 let mut damaged = snap_to_json(&snap);
4601 damaged["model_pools"][0]["name"] = "Argon\u{001b}[31m".into();
4602 let damaged = serde_json::to_vec(&damaged).unwrap();
4603 let restored = parse_cache_at(&damaged, None, now()).unwrap();
4604 assert_eq!(restored.model_pools[0].name, "Argon[31m");
4605 }
4606
4607 #[test]
4609 fn a_cached_model_pool_past_its_reset_refetches() {
4610 let mut snap = parsed();
4611 snap.model_pools.push(argon(60, "2026-07-22T11:00:00Z"));
4612 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4613 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
4614 assert!(err.to_string().contains("Gemini 4 Argon weekly"), "{err}");
4615 }
4616
4617 #[test]
4618 fn source_round_trips_through_the_cache_and_defaults_to_local() {
4619 let mut snap = parsed();
4620 snap.source = AntigravitySource::Remote;
4621 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4622 assert_eq!(
4623 parse_cache_at(&bytes, None, now()).unwrap().source,
4624 AntigravitySource::Remote
4625 );
4626
4627 snap.source = AntigravitySource::Statusline;
4628 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4629 assert_eq!(
4630 parse_cache_at(&bytes, None, now()).unwrap().source,
4631 AntigravitySource::Statusline
4632 );
4633
4634 snap.source = AntigravitySource::Agy;
4635 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
4636 assert_eq!(
4637 parse_cache_at(&bytes, None, now()).unwrap().source,
4638 AntigravitySource::Agy
4639 );
4640
4641 let mut legacy = snap_to_json(&snap);
4643 legacy.as_object_mut().unwrap().remove("source");
4644 let legacy = serde_json::to_vec(&legacy).unwrap();
4645 assert_eq!(
4646 parse_cache_at(&legacy, None, now()).unwrap().source,
4647 AntigravitySource::Local
4648 );
4649 }
4650}