1use std::time::Duration;
26
27use chrono::{DateTime, Utc};
28
29use super::cloud;
30use super::credential::{self, StoredToken};
31use crate::cache::{Cache, acquire_lock_async};
32use crate::error::{AppError, Result};
33use crate::usage::{AntigravitySnapshot, AntigravitySource, UsageWindow};
34
35const HTTP_TIMEOUT: Duration = Duration::from_secs(5);
36const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
37
38const QUOTA_RPC: &str = "exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary";
39const STATUS_RPC: &str = "exa.language_server_pb.LanguageServerService/GetUserStatus";
40
41const DEFAULT_PLAN: &str = "Antigravity";
42
43const NO_LOCAL_SERVER: &str = "Antigravity: no local server found. Quota is only served while \
44 Antigravity is running — open the Antigravity app, or an interactive \
45 `agy` session, or point ANTIGRAVITY_LS_ADDRESS at a host:port.";
46
47const AGY_CSRF_UNAVAILABLE: &str = "Antigravity: the running `agy` server requires a CSRF token \
48 that it does not publish.";
49
50const NO_SAVED_SESSION: &str = "Or sign in to Antigravity once, so its saved Google session can \
54 be used while it is closed.";
55
56const SESSION_EXPIRED: &str =
57 "Antigravity's saved Google session expired; open Antigravity to sign in again";
58
59pub type FetchOutcome = crate::outcome::Outcome<AntigravitySnapshot>;
62
63impl From<FetchOutcome> for crate::vendor::VendorOutcome {
64 fn from(o: FetchOutcome) -> Self {
65 o.map(crate::usage::VendorSnapshot::Antigravity)
66 }
67}
68
69#[derive(Debug, Clone, Copy, Default)]
73pub enum SavedCredential<'a> {
74 #[default]
76 Keyring,
77 Blob(&'a str),
79 Absent,
81}
82
83#[derive(Default)]
88pub struct RemoteOverride<'a> {
89 pub credential: SavedCredential<'a>,
91 pub endpoints: Option<&'a cloud::Endpoints>,
93 pub local_bases: Option<Vec<Candidate>>,
98}
99
100#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct Candidate {
106 pub base: String,
108 pub pid: Option<u32>,
111}
112
113impl From<String> for Candidate {
114 fn from(base: String) -> Self {
115 Self { base, pid: None }
116 }
117}
118
119pub async fn fetch_snapshot(
120 client: &reqwest::Client,
121 cache: &Cache,
122 cache_ttl: Duration,
123 oauth: Option<&cloud::OauthClient>,
124) -> Result<FetchOutcome> {
125 fetch_snapshot_at(
126 client,
127 cache,
128 cache_ttl,
129 oauth,
130 RemoteOverride::default(),
131 Utc::now(),
132 )
133 .await
134}
135
136pub async fn fetch_snapshot_at(
140 client: &reqwest::Client,
141 cache: &Cache,
142 cache_ttl: Duration,
143 oauth: Option<&cloud::OauthClient>,
144 remote: RemoteOverride<'_>,
145 now: DateTime<Utc>,
146) -> Result<FetchOutcome> {
147 cache.ensure_dir()?;
148 let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
149
150 let origin = match open_session(client, remote.local_bases.as_deref()).await {
160 Ok(session) => Origin::Local(Ok(session)),
161 Err(error) => match remote_fallback_reason(&error) {
162 Some(reason) => Origin::Remote(saved_session(remote.credential, reason)),
163 None => Origin::Local(Err(error)),
164 },
165 };
166 let account = origin.account();
167
168 if let Some(bytes) = cache.fresh_payload(cache_ttl)?
169 && let Ok(outcome) = reuse_cache(bytes, cache, false, account.as_deref(), now)
170 {
171 return Ok(outcome);
172 }
173
174 let default_endpoints = cloud::Endpoints::default();
175 let endpoints = remote.endpoints.unwrap_or(&default_endpoints);
176 let live = match origin {
177 Origin::Local(session) => fetch_live(client, session).await,
178 Origin::Remote(token) => fetch_remote(client, cache, oauth, endpoints, token, now).await,
179 };
180
181 match live {
182 Ok(snap) => {
183 let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
184 cache.write_payload(&bytes)?;
185 Ok(crate::outcome::Outcome::fresh(snap))
186 }
187 Err(e) if e.is_transient() => fallback_silent(cache, now, e),
188 Err(AppError::Http { status, body }) => {
189 cache.mark_stale();
190 let last_error = Some(cache.write_last_error(status, &body));
191 let reason = AppError::Http { status, body };
192 fallback_with_error(cache, last_error, reason, now)
193 }
194 Err(e) => {
195 cache.mark_stale();
196 let last_error = Some(cache.write_last_error(0, &e.to_string()));
197 fallback_with_error(cache, last_error, e, now)
198 }
199 }
200}
201
202struct Session {
205 base: String,
206 csrf: Option<String>,
207 plan: String,
208 account: String,
209}
210
211enum Origin {
214 Local(Result<Session>),
215 Remote(Result<StoredToken>),
216}
217
218impl Origin {
219 fn account(&self) -> Option<String> {
221 match self {
222 Origin::Local(Ok(session)) => Some(session.account.clone()),
223 Origin::Remote(Ok(token)) => Some(remote_account(&token.fingerprint)),
224 Origin::Local(Err(_)) | Origin::Remote(Err(_)) => None,
225 }
226 }
227}
228
229fn no_local_server() -> AppError {
230 AppError::Credentials(NO_LOCAL_SERVER.into())
231}
232
233#[derive(Clone, Copy)]
234enum RemoteFallbackReason {
235 NoLocalServer,
236 AgyMissingCsrf,
237}
238
239impl RemoteFallbackReason {
240 fn message(self) -> &'static str {
241 match self {
242 Self::NoLocalServer => NO_LOCAL_SERVER,
243 Self::AgyMissingCsrf => AGY_CSRF_UNAVAILABLE,
244 }
245 }
246}
247
248fn remote_fallback_reason(error: &AppError) -> Option<RemoteFallbackReason> {
252 if matches!(error, AppError::Credentials(message) if message == NO_LOCAL_SERVER) {
253 Some(RemoteFallbackReason::NoLocalServer)
254 } else if is_missing_csrf(error) {
255 Some(RemoteFallbackReason::AgyMissingCsrf)
256 } else {
257 None
258 }
259}
260
261async fn open_session(client: &reqwest::Client, bases: Option<&[Candidate]>) -> Result<Session> {
267 let candidates = bases.map_or_else(candidate_bases, <[Candidate]>::to_vec);
268 if candidates.is_empty() {
269 return Err(no_local_server());
270 }
271
272 let mut errors = Vec::new();
273 let mut blocked_pids = std::collections::HashSet::new();
278 for candidate in candidates {
279 if candidate.pid.is_some_and(|pid| blocked_pids.contains(&pid)) {
280 continue;
281 }
282 let csrf = fetch_csrf(client, &candidate.base).await;
283 match post_rpc(client, &candidate.base, csrf.as_deref(), STATUS_RPC).await {
284 Ok(v) => {
285 return Ok(Session {
286 base: candidate.base,
287 csrf,
288 plan: plan_from_status(&v),
289 account: account_key(&v),
290 });
291 }
292 Err(e) => {
293 if is_missing_csrf(&e)
297 && let Some(pid) = candidate.pid
298 {
299 blocked_pids.insert(pid);
300 }
301 errors.push(e);
302 }
303 }
304 }
305 Err(select_probe_error(errors))
306}
307
308fn select_probe_error(errors: Vec<AppError>) -> AppError {
325 let mut actionable = None;
326 let mut missing_csrf = None;
327 let mut last = None;
328 let mut echo = None;
329 for e in errors {
330 if missing_csrf.is_none() && is_missing_csrf(&e) {
331 missing_csrf = Some(e);
332 } else if actionable.is_none() && is_actionable(&e) {
333 actionable = Some(e);
334 } else if is_tls_echo(&e) {
335 echo = Some(e);
336 } else {
337 last = Some(e);
338 }
339 }
340 actionable
341 .or(missing_csrf)
342 .or(last)
343 .or(echo)
344 .unwrap_or_else(|| {
345 AppError::Other("antigravity: no local server answered GetUserStatus".into())
346 })
347}
348
349fn is_missing_csrf(error: &AppError) -> bool {
354 let AppError::Http { status: 401, body } = error else {
355 return false;
356 };
357 let Ok(body) = serde_json::from_str::<serde_json::Value>(body) else {
358 return false;
359 };
360 matches!(
361 (body["code"].as_str(), body["message"].as_str()),
362 (Some(code), Some(message))
363 if code.eq_ignore_ascii_case("unauthenticated")
364 && message.trim().eq_ignore_ascii_case("missing CSRF token")
365 )
366}
367
368fn is_actionable(e: &AppError) -> bool {
373 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
374 && !is_missing_csrf(e)
375}
376
377fn is_tls_echo(e: &AppError) -> bool {
390 matches!(
391 e,
392 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
393 )
394}
395
396async fn fetch_live(
397 client: &reqwest::Client,
398 session: Result<Session>,
399) -> Result<AntigravitySnapshot> {
400 let session = session?;
401 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
402 let mut snap = parse_quota_summary("a, session.plan)?;
403 snap.account = session.account;
404 Ok(snap)
405}
406
407fn saved_session(
417 credential: SavedCredential<'_>,
418 reason: RemoteFallbackReason,
419) -> Result<StoredToken> {
420 let raw = match credential {
421 SavedCredential::Keyring => credential::read()?,
422 SavedCredential::Blob(blob) => Some(blob.to_string()),
423 SavedCredential::Absent => None,
424 };
425 let Some(raw) = raw else {
426 return Err(AppError::Credentials(format!(
427 "{} {NO_SAVED_SESSION}",
428 reason.message()
429 )));
430 };
431 credential::parse_keyring_blob(&raw)
432}
433
434fn remote_account(fingerprint: &str) -> String {
438 format!("acct:{fingerprint}")
439}
440
441fn session_expired() -> AppError {
442 AppError::Credentials(SESSION_EXPIRED.into())
443}
444
445fn refresh_unconfigured() -> AppError {
449 AppError::Credentials(
450 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
451 refresh it; open Antigravity to sign in again, or set [antigravity] oauth_client_id \
452 and oauth_client_secret in config.toml"
453 .into(),
454 )
455}
456
457fn is_auth_rejection(e: &AppError) -> bool {
458 matches!(
459 e,
460 AppError::Http {
461 status: 401 | 403,
462 ..
463 }
464 )
465}
466
467struct AccessToken {
471 value: String,
472 just_refreshed: bool,
473}
474
475async fn resolve_access_token(
480 client: &reqwest::Client,
481 oauth: Option<&cloud::OauthClient>,
482 endpoints: &cloud::Endpoints,
483 oauth_path: &std::path::Path,
484 token: &StoredToken,
485 now: DateTime<Utc>,
486) -> Result<AccessToken> {
487 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
488 && token
489 .expires_at
490 .is_none_or(|keyring| persisted.expires_at > keyring)
491 && !cloud::needs_refresh(Some(persisted.expires_at), now)
492 {
493 return Ok(AccessToken {
494 value: persisted.access_token,
495 just_refreshed: false,
496 });
497 }
498 if !cloud::needs_refresh(token.expires_at, now) {
499 return Ok(AccessToken {
500 value: token.access_token.clone(),
501 just_refreshed: false,
502 });
503 }
504 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
505}
506
507async fn refresh_and_persist(
511 client: &reqwest::Client,
512 oauth: Option<&cloud::OauthClient>,
513 endpoints: &cloud::Endpoints,
514 oauth_path: &std::path::Path,
515 token: &StoredToken,
516) -> Result<AccessToken> {
517 let Some(refresh_token) = token.refresh_token.as_deref() else {
518 return Err(session_expired());
519 };
520 let Some(oauth) = oauth else {
521 return Err(refresh_unconfigured());
522 };
523 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
524 cloud::write_persisted(
525 oauth_path,
526 &cloud::PersistedOAuth {
527 fingerprint: token.fingerprint.clone(),
528 access_token: refreshed.access_token.clone(),
529 expires_at: refreshed.expires_at,
530 },
531 )?;
532 Ok(AccessToken {
533 value: refreshed.access_token,
534 just_refreshed: true,
535 })
536}
537
538async fn fetch_remote(
545 client: &reqwest::Client,
546 cache: &Cache,
547 oauth: Option<&cloud::OauthClient>,
548 endpoints: &cloud::Endpoints,
549 token: Result<StoredToken>,
550 now: DateTime<Utc>,
551) -> Result<AntigravitySnapshot> {
552 let token = token?;
553 let oauth_path = cloud::oauth_cache_path(cache);
554 let mut access =
555 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?;
556
557 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
558 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
559 access = refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?;
560 cloud::fetch_quota(client, endpoints, &access.value)
561 .await
562 .map_err(|e| {
563 if is_auth_rejection(&e) {
564 session_expired()
565 } else {
566 e
567 }
568 })?
569 }
570 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
571 other => other?,
572 };
573
574 let plan = cloud::fetch_plan(client, endpoints, &access.value)
575 .await
576 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
577 let mut snap = parse_quota_summary("a, plan)?;
578 snap.account = remote_account(&token.fingerprint);
579 snap.source = AntigravitySource::Remote;
580 Ok(snap)
581}
582
583fn account_key(user_status: &serde_json::Value) -> String {
588 let email = user_status["userStatus"]["email"]
589 .as_str()
590 .filter(|s| !s.is_empty());
591 match email {
592 Some(e) => {
593 use std::hash::{Hash, Hasher};
594 let mut h = std::collections::hash_map::DefaultHasher::new();
595 e.hash(&mut h);
596 format!("acct:{:016x}", h.finish())
597 }
598 None => "acct:unknown".to_string(),
599 }
600}
601
602async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
606 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
607 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
610 .await
611 .ok()?;
612 let html = String::from_utf8_lossy(&bytes);
613 html.split("csrfToken\":\"")
614 .nth(1)
615 .and_then(|s| s.split('"').next())
616 .filter(|t| !t.is_empty())
617 .map(|t| t.to_string())
618}
619
620async fn post_rpc(
621 client: &reqwest::Client,
622 base: &str,
623 csrf: Option<&str>,
624 rpc: &str,
625) -> Result<serde_json::Value> {
626 let mut req = client
627 .post(format!("{base}/{rpc}"))
628 .header("Content-Type", "application/json")
629 .body("{}")
630 .timeout(HTTP_TIMEOUT);
631 if let Some(token) = csrf {
632 req = req.header("x-codeium-csrf-token", token);
633 }
634 let resp = req.send().await?;
635
636 let status = resp.status();
637 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
641 if !status.is_success() {
642 let body = String::from_utf8_lossy(&bytes).into_owned();
643 return Err(AppError::Http {
644 status: status.as_u16(),
645 body,
646 });
647 }
648
649 Ok(serde_json::from_slice(&bytes)?)
650}
651
652pub fn plan_from_status(v: &serde_json::Value) -> String {
653 v["userStatus"]["userTier"]["name"]
654 .as_str()
655 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
656 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
657 .filter(|s| !s.is_empty())
658 .unwrap_or(DEFAULT_PLAN)
659 .to_string()
660}
661
662fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
673 let id = if id.is_empty() { "<unnamed>" } else { id };
674 let group = if group_name.is_empty() {
675 String::new()
676 } else {
677 format!(" in {group_name:?}")
678 };
679 match window {
680 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
681 _ => format!("{id}{group}"),
682 }
683}
684
685fn no_usable_bucket(seen: &[String]) -> AppError {
691 if seen.is_empty() {
692 return AppError::Other(
693 "antigravity: quota summary has no buckets at all — the running product may \
694 not have a quota for this account yet"
695 .into(),
696 );
697 }
698 AppError::Other(format!(
699 "antigravity: quota summary has no bucket in a window we recognise (5h or \
700 weekly, Gemini or Claude/GPT); it offered: {}",
701 seen.join(", ")
702 ))
703}
704
705pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
706 let groups = v["response"]["groups"]
707 .as_array()
708 .or_else(|| v["groups"].as_array())
709 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
710
711 let mut gemini_5h = None;
712 let mut gemini_weekly = None;
713 let mut tp_5h = None;
714 let mut tp_weekly = None;
715 let mut seen: Vec<String> = Vec::new();
719
720 for group in groups {
721 let group_name = group["displayName"].as_str().unwrap_or_default();
722 let Some(buckets) = group["buckets"].as_array() else {
723 continue;
724 };
725 for bucket in buckets {
726 let id = bucket["bucketId"].as_str().unwrap_or_default();
727 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
728 let window = bucket["window"].as_str().unwrap_or_default();
729 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
730 true
731 } else if id.ends_with("5h") || window == "5h" {
732 false
733 } else {
734 continue;
737 };
738 let is_gemini = if id.starts_with("gemini") {
739 true
740 } else if id.starts_with("3p") {
741 false
742 } else if group_name.contains("Gemini") {
743 true
744 } else if group_name.contains("Claude") || group_name.contains("GPT") {
745 false
746 } else {
747 continue;
750 };
751
752 let (slot, slot_name) = match (is_gemini, is_weekly) {
753 (true, false) => (&mut gemini_5h, "Gemini 5h"),
754 (true, true) => (&mut gemini_weekly, "Gemini weekly"),
755 (false, false) => (&mut tp_5h, "Claude/GPT 5h"),
756 (false, true) => (&mut tp_weekly, "Claude/GPT weekly"),
757 };
758 let parsed = usage_window(bucket, is_weekly)?;
759 if slot.replace(parsed).is_some() {
760 return Err(AppError::Schema(format!(
761 "antigravity: duplicate {slot_name} bucket"
762 )));
763 }
764 }
765 }
766
767 if gemini_5h.is_none() && gemini_weekly.is_none() && tp_5h.is_none() && tp_weekly.is_none() {
771 return Err(no_usable_bucket(&seen));
772 }
773
774 Ok(AntigravitySnapshot {
775 plan,
776 account: String::new(),
779 source: AntigravitySource::Local,
780 session: gemini_5h,
781 weekly: gemini_weekly,
782 third_party_session: tp_5h,
783 third_party_weekly: tp_weekly,
784 })
785}
786
787fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
791 let remaining = bucket["remainingFraction"]
792 .as_f64()
793 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
794 .ok_or_else(|| {
795 AppError::Schema(format!(
796 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
797 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
798 ))
799 })?;
800 Ok(UsageWindow {
801 utilization_pct: pct_used(remaining),
802 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
803 window_duration: if is_weekly {
804 chrono::Duration::days(7)
805 } else {
806 chrono::Duration::hours(5)
807 },
808 })
809}
810
811fn pct_used(remaining_fraction: f64) -> i32 {
814 let used = (1.0 - remaining_fraction) * 100.0;
815 used.round() as i32
816}
817
818fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
819 match value {
820 serde_json::Value::Null => Ok(None),
821 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
822 .map(|dt| Some(dt.with_timezone(&Utc)))
823 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
824 _ => Err(AppError::Schema(format!(
825 "antigravity: {field} must be a timestamp or null"
826 ))),
827 }
828}
829
830fn candidate_bases() -> Vec<Candidate> {
836 candidate_bases_with(
837 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
838 discover_ls_ports(),
839 )
840}
841
842fn candidate_bases_with(
846 override_addr: Option<&str>,
847 discovered: Vec<(u32, u16)>,
848) -> Vec<Candidate> {
849 let mut candidates = Vec::new();
850 if let Some(base) = override_addr.and_then(normalize_base) {
851 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
855 candidates.push(Candidate { base, pid });
856 }
857
858 for (pid, port) in discovered {
864 let base = format!("http://127.0.0.1:{port}");
865 if !candidates.iter().any(|c| c.base == base) {
866 candidates.push(Candidate {
867 base,
868 pid: Some(pid),
869 });
870 }
871 }
872
873 candidates
874}
875
876fn base_port(base: &str) -> Option<u16> {
878 let (_, authority) = base.split_once("://")?;
879 authority.rsplit_once(':')?.1.parse().ok()
880}
881
882fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
885 let mut owners = discovered
886 .iter()
887 .filter(|(_, p)| *p == port)
888 .map(|(pid, _)| *pid);
889 let first = owners.next()?;
890 owners.all(|pid| pid == first).then_some(first)
891}
892
893fn normalize_base(addr: &str) -> Option<String> {
902 let trimmed = addr.trim();
903 let (scheme, authority) = match trimmed.split_once("://") {
904 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
905 _ => ("http", trimmed),
906 };
907 let authority = authority.trim_end_matches('/');
908 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
909}
910
911fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
918 let comm = comm.trim().to_lowercase();
919 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
920 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
921 return true;
922 }
923 exe.is_some_and(|p| {
924 let p = p.to_lowercase().replace('\\', "/");
925 let p = p.strip_suffix(".exe").unwrap_or(&p);
926 p.contains("antigravity") || p.ends_with("/agy")
927 })
928}
929
930#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
962fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
963 let groups: Vec<(u32, Vec<u16>)> = per_pid
964 .into_iter()
965 .map(|(pid, mut group)| {
966 group.sort_unstable_by(|a, b| b.cmp(a));
967 group.dedup();
972 (pid, group)
973 })
974 .collect();
975
976 let mut ports: Vec<(u32, u16)> = Vec::new();
977 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
978 for (pid, port) in groups
979 .iter()
980 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
981 {
982 if !ports.iter().any(|(_, p)| *p == port) {
985 ports.push((pid, port));
986 }
987 }
988 }
989 ports
990}
991
992#[cfg(target_os = "linux")]
1000pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1001 use std::collections::{BTreeMap, HashMap};
1002
1003 let mut owners: HashMap<u64, u32> = HashMap::new();
1006 let Ok(entries) = std::fs::read_dir("/proc") else {
1007 return Vec::new();
1008 };
1009 for entry in entries.flatten() {
1010 let pid_dir = entry.path();
1011 let Some(pid) = pid_dir
1012 .file_name()
1013 .and_then(|name| name.to_str())
1014 .and_then(|name| name.parse::<u32>().ok())
1015 else {
1016 continue;
1017 };
1018 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1019 continue;
1020 };
1021 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1022 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1023 continue;
1024 }
1025 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1026 continue;
1027 };
1028 for fd in fds.flatten() {
1029 let Ok(target) = std::fs::read_link(fd.path()) else {
1030 continue;
1031 };
1032 if let Some(ino) = target
1033 .to_str()
1034 .and_then(|s| s.strip_prefix("socket:["))
1035 .and_then(|s| s.strip_suffix(']'))
1036 .and_then(|s| s.parse::<u64>().ok())
1037 {
1038 owners.insert(ino, pid);
1039 }
1040 }
1041 }
1042
1043 if owners.is_empty() {
1044 return Vec::new();
1045 }
1046
1047 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1048 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1049 let Ok(contents) = std::fs::read_to_string(table) else {
1050 continue;
1051 };
1052 for line in contents.lines().skip(1) {
1053 if let Some((port, ino)) = parse_proc_net_line(line)
1054 && let Some(&pid) = owners.get(&ino)
1055 {
1056 per_pid.entry(pid).or_default().push(port);
1057 }
1058 }
1059 }
1060 probe_order(per_pid)
1061}
1062
1063#[cfg(target_os = "macos")]
1070pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1071 let Ok(output) = std::process::Command::new("lsof")
1072 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1073 .output()
1074 else {
1075 return Vec::new();
1076 };
1077 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1081}
1082
1083#[cfg(any(test, target_os = "macos"))]
1088fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1089 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1090 let mut pid = None;
1093 let mut owner = None;
1094 for line in output.lines() {
1095 let Some(rest) = line.get(1..) else { continue };
1096 match line.as_bytes().first() {
1097 Some(b'p') => {
1098 pid = rest.parse::<u32>().ok();
1099 owner = None;
1100 }
1101 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1102 Some(b'n') => {
1103 if let Some(pid) = owner
1104 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1105 {
1106 per_pid.entry(pid).or_default().push(port);
1107 }
1108 }
1109 _ => {}
1110 }
1111 }
1112 probe_order(per_pid)
1113}
1114
1115#[cfg(any(test, target_os = "windows"))]
1116#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1117struct WindowsTcpRow {
1118 local_addr: [u8; 4],
1119 local_port: u32,
1120 pid: u32,
1121}
1122
1123#[cfg(any(test, target_os = "windows"))]
1124fn decode_windows_process_name(raw: &[u16]) -> String {
1125 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1126 String::from_utf16_lossy(&raw[..end])
1127}
1128
1129#[cfg(any(test, target_os = "windows"))]
1130fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1131 processes
1132 .iter()
1133 .filter(|(_, name)| is_antigravity_process(name, None))
1134 .map(|(pid, _)| *pid)
1135 .collect()
1136}
1137
1138#[cfg(any(test, target_os = "windows"))]
1141fn matching_windows_ports(
1142 pids: &std::collections::HashSet<u32>,
1143 rows: &[WindowsTcpRow],
1144) -> Vec<(u32, u16)> {
1145 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1146 for row in rows {
1147 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1148 continue;
1149 }
1150 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1151 if port != 0 {
1152 per_pid.entry(row.pid).or_default().push(port);
1153 }
1154 }
1155 probe_order(per_pid)
1156}
1157
1158#[cfg(any(test, target_os = "windows"))]
1159fn checked_windows_row_count(
1160 buffer_len: usize,
1161 rows_offset: usize,
1162 row_size: usize,
1163 declared: usize,
1164) -> Option<usize> {
1165 let rows_len = row_size.checked_mul(declared)?;
1166 let end = rows_offset.checked_add(rows_len)?;
1167 (row_size != 0 && end <= buffer_len).then_some(declared)
1168}
1169
1170#[cfg(target_os = "windows")]
1171struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1172
1173#[cfg(target_os = "windows")]
1174impl Drop for WindowsHandle {
1175 fn drop(&mut self) {
1176 unsafe {
1177 windows_sys::Win32::Foundation::CloseHandle(self.0);
1178 }
1179 }
1180}
1181
1182#[cfg(target_os = "windows")]
1183fn windows_processes() -> Vec<(u32, String)> {
1184 use std::mem::size_of;
1185 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1186 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1187 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1188 TH32CS_SNAPPROCESS,
1189 };
1190
1191 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1192 if handle == INVALID_HANDLE_VALUE {
1193 return Vec::new();
1194 }
1195 let snapshot = WindowsHandle(handle);
1196 let mut entry = PROCESSENTRY32W::default();
1197 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1198 return Vec::new();
1199 };
1200 entry.dwSize = entry_size;
1201 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1202 return Vec::new();
1203 }
1204
1205 let mut processes = Vec::new();
1206 loop {
1207 processes.push((
1208 entry.th32ProcessID,
1209 decode_windows_process_name(&entry.szExeFile),
1210 ));
1211 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1212 break;
1213 }
1214 }
1215 processes
1216}
1217
1218#[cfg(target_os = "windows")]
1219fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1220 use std::mem::{offset_of, size_of, size_of_val};
1221 use windows_sys::Win32::NetworkManagement::IpHelper::{
1222 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1223 };
1224
1225 let available = used_bytes.min(size_of_val(buffer));
1226 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1227 if available < size_of::<u32>() || available < rows_offset {
1228 return Vec::new();
1229 }
1230 let base = buffer.as_ptr().cast::<u8>();
1231 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1232 if checked_windows_row_count(
1233 available,
1234 rows_offset,
1235 size_of::<MIB_TCPROW_OWNER_PID>(),
1236 declared,
1237 )
1238 .is_none()
1239 {
1240 return Vec::new();
1241 }
1242
1243 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1244 (0..declared)
1245 .map(|index| unsafe { rows.add(index).read_unaligned() })
1246 .map(|row| WindowsTcpRow {
1247 local_addr: row.dwLocalAddr.to_ne_bytes(),
1248 local_port: row.dwLocalPort,
1249 pid: row.dwOwningPid,
1250 })
1251 .collect()
1252}
1253
1254#[cfg(target_os = "windows")]
1255fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1256 use std::mem::size_of;
1257 use std::ptr::null_mut;
1258 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1259 use windows_sys::Win32::NetworkManagement::IpHelper::{
1260 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1261 };
1262 use windows_sys::Win32::Networking::WinSock::AF_INET;
1263
1264 let mut size = 0u32;
1265 let status = unsafe {
1266 GetExtendedTcpTable(
1267 null_mut(),
1268 &mut size,
1269 0,
1270 u32::from(AF_INET),
1271 TCP_TABLE_OWNER_PID_LISTENER,
1272 0,
1273 )
1274 };
1275 if status != ERROR_INSUFFICIENT_BUFFER {
1276 return Vec::new();
1277 }
1278
1279 for _ in 0..3 {
1280 let Some(words) = (size as usize)
1281 .checked_add(size_of::<u32>() - 1)
1282 .map(|bytes| bytes / size_of::<u32>())
1283 else {
1284 return Vec::new();
1285 };
1286 if words == 0 {
1287 return Vec::new();
1288 }
1289 let mut buffer = Vec::<u32>::new();
1290 if buffer.try_reserve_exact(words).is_err() {
1291 return Vec::new();
1292 }
1293 buffer.resize(words, 0);
1294 let mut used = size;
1295 let status = unsafe {
1296 GetExtendedTcpTable(
1297 buffer.as_mut_ptr().cast(),
1298 &mut used,
1299 0,
1300 u32::from(AF_INET),
1301 TCP_TABLE_OWNER_PID_LISTENER,
1302 0,
1303 )
1304 };
1305 if status == ERROR_INSUFFICIENT_BUFFER {
1306 size = used;
1307 continue;
1308 }
1309 if status != 0 {
1310 return Vec::new();
1311 }
1312 return parse_windows_tcp_rows(&buffer, used as usize);
1313 }
1314 Vec::new()
1315}
1316
1317#[cfg(target_os = "windows")]
1318pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1319 let pids = matching_windows_process_ids(&windows_processes());
1320 if pids.is_empty() {
1321 return Vec::new();
1322 }
1323 matching_windows_ports(&pids, &windows_tcp_rows())
1324}
1325
1326#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1327pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1328 Vec::new()
1329}
1330
1331#[cfg(target_os = "linux")]
1334fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1335 let cols: Vec<&str> = line.split_whitespace().collect();
1336 if cols.len() < 10 {
1337 return None;
1338 }
1339 if cols[3] != "0A" {
1341 return None;
1342 }
1343 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1344 let inode = cols[9].parse::<u64>().ok()?;
1345 Some((port, inode))
1346}
1347
1348fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1353 crate::outcome::fallback(cache, None, original, |bytes| {
1354 parse_cache_at(bytes, None, now)
1355 })
1356}
1357
1358fn fallback_with_error(
1363 cache: &Cache,
1364 last_error: Option<(u16, String)>,
1365 reason: AppError,
1366 now: DateTime<Utc>,
1367) -> Result<FetchOutcome> {
1368 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1369 parse_cache_at(bytes, None, now)
1370 })
1371}
1372
1373fn reuse_cache(
1374 bytes: Vec<u8>,
1375 cache: &Cache,
1376 stale: bool,
1377 account: Option<&str>,
1378 now: DateTime<Utc>,
1379) -> Result<FetchOutcome> {
1380 let snap = parse_cache_at(&bytes, account, now)?;
1381 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1382}
1383
1384pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1390 parse_cache_at(bytes, account, Utc::now())
1391}
1392
1393fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<&'static str> {
1403 [
1404 ("Gemini 5h", snap.session.as_ref()),
1405 ("Gemini weekly", snap.weekly.as_ref()),
1406 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1407 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1408 ]
1409 .into_iter()
1410 .find(|(_, w)| w.and_then(|w| w.resets_at).is_some_and(|r| r <= now))
1411 .map(|(name, _)| name)
1412}
1413
1414pub fn parse_cache_at(
1415 bytes: &[u8],
1416 account: Option<&str>,
1417 now: DateTime<Utc>,
1418) -> Result<AntigravitySnapshot> {
1419 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1420
1421 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1422 if let Some(expected) = account
1423 && cached_account != Some(expected)
1424 {
1425 return Err(AppError::Schema(
1426 "antigravity cache belongs to a different account; refetching".into(),
1427 ));
1428 }
1429
1430 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1440 match v.get(pct_key) {
1441 None => Err(AppError::Schema(format!(
1442 "antigravity: cached payload is missing {pct_key}"
1443 ))),
1444 Some(serde_json::Value::Null) => Ok(None),
1445 Some(value) => value
1446 .as_i64()
1447 .filter(|pct| (0..=100).contains(pct))
1448 .map(|pct| Some(pct as i32))
1449 .ok_or_else(|| {
1450 AppError::Schema(format!(
1451 "antigravity: cached {pct_key} must be an integer in 0..=100"
1452 ))
1453 }),
1454 }
1455 };
1456
1457 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1458 let Some(pct) = cached_pct(pct_key)? else {
1459 return Ok(None);
1460 };
1461 Ok::<_, AppError>(Some(UsageWindow {
1462 utilization_pct: pct,
1463 resets_at: parse_reset(&v[reset_key], reset_key)?,
1464 window_duration: if weekly {
1465 chrono::Duration::days(7)
1466 } else {
1467 chrono::Duration::hours(5)
1468 },
1469 }))
1470 };
1471
1472 let snap = AntigravitySnapshot {
1473 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1474 account: cached_account.unwrap_or_default().to_string(),
1475 source: v["source"]
1477 .as_str()
1478 .and_then(AntigravitySource::parse)
1479 .unwrap_or_default(),
1480 session: optional("session_pct", "session_reset", false)?,
1481 weekly: optional("weekly_pct", "weekly_reset", true)?,
1482 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1483 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1484 };
1485
1486 if snap.session.is_none()
1489 && snap.weekly.is_none()
1490 && snap.third_party_session.is_none()
1491 && snap.third_party_weekly.is_none()
1492 {
1493 return Err(AppError::Schema(
1494 "antigravity cache holds no usable window; refetching".into(),
1495 ));
1496 }
1497
1498 if let Some(window) = expired_window(&snap, now) {
1499 return Err(AppError::Schema(format!(
1500 "antigravity cache is past its {window} reset; refetching"
1501 )));
1502 }
1503 Ok(snap)
1504}
1505
1506pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1507 serde_json::json!({
1508 "plan": snap.plan,
1509 "account": snap.account,
1510 "source": snap.source.as_str(),
1511 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
1512 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1513 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
1514 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1515 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
1516 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1517 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
1518 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1519 })
1520}
1521
1522#[cfg(test)]
1523mod tests {
1524 use super::*;
1525
1526 const QUOTA_JSON: &str = r#"{
1530 "response": {
1531 "groups": [
1532 {
1533 "displayName": "Gemini Models",
1534 "buckets": [
1535 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
1536 "window": "weekly", "remainingFraction": 0.9191212,
1537 "resetTime": "2026-07-28T17:39:58Z"},
1538 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
1539 "window": "5h", "remainingFraction": 0.5672253,
1540 "resetTime": "2026-07-22T17:47:00Z"}
1541 ]
1542 },
1543 {
1544 "displayName": "Claude and GPT models",
1545 "buckets": [
1546 {"bucketId": "3p-weekly", "window": "weekly",
1547 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
1548 {"bucketId": "3p-5h", "window": "5h",
1549 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
1550 ]
1551 }
1552 ]
1553 }
1554 }"#;
1555
1556 fn now() -> DateTime<Utc> {
1559 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
1560 .unwrap()
1561 .with_timezone(&Utc)
1562 }
1563
1564 fn parsed() -> AntigravitySnapshot {
1565 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1566 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
1567 }
1568
1569 #[test]
1570 fn quota_summary_maps_four_distinct_windows() {
1571 let snap = parsed();
1572 assert_eq!(snap.plan, "Google AI Pro");
1573 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
1575 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
1576 assert_eq!(
1577 snap.third_party_session.as_ref().unwrap().utilization_pct,
1578 75
1579 );
1580 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
1581 }
1582
1583 #[test]
1584 fn each_window_keeps_its_own_reset_time() {
1585 let snap = parsed();
1586 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
1587 assert_eq!(
1588 snap.session.as_ref().unwrap().resets_at,
1589 at("2026-07-22T17:47:00Z")
1590 );
1591 assert_eq!(
1592 snap.weekly.as_ref().unwrap().resets_at,
1593 at("2026-07-28T17:39:58Z")
1594 );
1595 assert_eq!(
1596 snap.third_party_weekly.as_ref().unwrap().resets_at,
1597 at("2026-07-29T12:47:00Z")
1598 );
1599 assert_ne!(
1601 snap.session.as_ref().unwrap().resets_at,
1602 snap.weekly.as_ref().unwrap().resets_at
1603 );
1604 }
1605
1606 #[test]
1607 fn window_durations_match_their_bucket() {
1608 let snap = parsed();
1609 assert_eq!(
1610 snap.session.as_ref().unwrap().window_duration,
1611 chrono::Duration::hours(5)
1612 );
1613 assert_eq!(
1614 snap.weekly.as_ref().unwrap().window_duration,
1615 chrono::Duration::days(7)
1616 );
1617 assert_eq!(
1618 snap.third_party_weekly.as_ref().unwrap().window_duration,
1619 chrono::Duration::days(7)
1620 );
1621 }
1622
1623 #[test]
1624 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
1625 let v: serde_json::Value = serde_json::from_str(
1626 r#"{"response":{"groups":[
1627 {"displayName":"Gemini Models","buckets":[
1628 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
1629 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
1630 {"displayName":"Claude and GPT models","buckets":[
1631 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
1632 ]}}"#,
1633 )
1634 .unwrap();
1635 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
1636 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
1637 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
1638 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
1639 assert!(snap.third_party_weekly.is_none());
1640 }
1641
1642 #[test]
1643 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
1644 let duplicate: serde_json::Value = serde_json::from_str(
1645 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
1646 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1647 {"bucketId":"gemini-5h-copy","window":"5h","remainingFraction":0.1},
1648 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
1649 ]}]}}"#,
1650 )
1651 .unwrap();
1652 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
1653 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
1654
1655 let unrelated: serde_json::Value = serde_json::from_str(
1658 r#"{"response":{"groups":[
1659 {"displayName":"Gemini Models","buckets":[
1660 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1661 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
1662 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
1663 ]},
1664 {"displayName":"Future Models","buckets":[
1665 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
1666 ]}
1667 ]}}"#,
1668 )
1669 .unwrap();
1670 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
1671 assert!(snap.third_party_session.is_none());
1672 assert!(snap.third_party_weekly.is_none());
1673 }
1674
1675 #[test]
1678 fn a_bucket_without_a_usable_fraction_is_rejected() {
1679 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
1680 let v: serde_json::Value = serde_json::from_str(&format!(
1681 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
1682 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
1683 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
1684 ))
1685 .unwrap();
1686 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1687 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
1688 }
1689 }
1690
1691 #[test]
1692 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
1693 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
1694 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1695 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
1696 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1697 assert!(err.to_string().contains("resetTime"), "{err}");
1698 }
1699 }
1700
1701 #[test]
1705 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
1706 let mut snap = parsed();
1707 snap.session = None;
1708 snap.third_party_session = None;
1709
1710 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1711 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
1712
1713 assert!(back.session.is_none());
1714 assert_eq!(
1715 back.weekly.as_ref().unwrap().utilization_pct,
1716 snap.weekly.as_ref().unwrap().utilization_pct
1717 );
1718 }
1719
1720 #[test]
1725 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
1726 let summary = serde_json::json!({
1727 "groups": [
1728 {
1729 "displayName": "Gemini Models",
1730 "buckets": [{
1731 "bucketId": "gemini-weekly", "window": "weekly",
1732 "remainingFraction": 0.42,
1733 }],
1734 },
1735 {
1736 "displayName": "Claude and GPT models",
1737 "buckets": [{
1738 "bucketId": "3p-weekly", "window": "weekly",
1739 "remainingFraction": 0.9,
1740 }],
1741 },
1742 ],
1743 });
1744
1745 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
1746
1747 assert!(snap.session.is_none(), "no 5h bucket arrived");
1748 assert!(snap.third_party_session.is_none());
1749 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
1750 assert_eq!(
1751 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
1752 10
1753 );
1754 }
1755
1756 #[test]
1759 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
1760 let summary = serde_json::json!({
1761 "groups": [{
1762 "displayName": "Gemini Models",
1763 "buckets": [{
1764 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
1765 }],
1766 }],
1767 });
1768
1769 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
1770
1771 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
1772 assert!(snap.weekly.is_none());
1773 }
1774
1775 #[test]
1778 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
1779 let summary = serde_json::json!({
1780 "groups": [{
1781 "displayName": "Gemini Models",
1782 "buckets": [{
1783 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
1784 }],
1785 }],
1786 });
1787
1788 let rendered = parse_quota_summary(&summary, "Pro".into())
1789 .expect_err("an unrecognised cadence alone is not a snapshot")
1790 .to_string();
1791
1792 assert!(
1793 rendered.contains("no bucket in a window we recognise"),
1794 "{rendered}"
1795 );
1796 assert!(rendered.contains("gemini-daily"), "{rendered}");
1797 assert!(rendered.contains("window daily"), "{rendered}");
1798 }
1799
1800 #[test]
1803 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
1804 let summary = serde_json::json!({
1805 "groups": [{"displayName": "Gemini", "buckets": []}],
1806 });
1807
1808 let rendered = parse_quota_summary(&summary, "Pro".into())
1809 .expect_err("no buckets is an error")
1810 .to_string();
1811
1812 assert!(rendered.contains("no buckets at all"), "{rendered}");
1813 assert!(!rendered.contains("it offered:"), "{rendered}");
1814 }
1815
1816 #[test]
1819 fn buckets_without_an_id_are_still_named_in_the_error() {
1820 let summary = serde_json::json!({
1821 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
1822 });
1823
1824 let rendered = parse_quota_summary(&summary, "Pro".into())
1825 .expect_err("an unusable summary is an error")
1826 .to_string();
1827
1828 assert!(rendered.contains("<unnamed>"), "{rendered}");
1829 }
1830
1831 #[test]
1832 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
1833 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
1834 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
1835 }
1836
1837 #[test]
1838 fn cache_round_trip_preserves_every_window() {
1839 let snap = parsed();
1840 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1841 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1842 }
1843
1844 #[test]
1848 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
1849 let full = snap_to_json(&parsed());
1850 for missing in ["session_pct", "weekly_pct"] {
1851 let mut v = full.clone();
1852 v.as_object_mut().unwrap().remove(missing);
1853 let bytes = serde_json::to_vec(&v).unwrap();
1854 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1855 assert!(err.to_string().contains(missing), "{missing}: {err}");
1856 }
1857 assert!(parse_cache_at(b"{}", None, now()).is_err());
1859 }
1860
1861 #[test]
1862 fn cached_percentages_are_range_checked_before_narrowing() {
1863 let full = snap_to_json(&parsed());
1864 for (key, bad) in [
1865 ("session_pct", serde_json::json!(-1)),
1866 ("weekly_pct", serde_json::json!(101)),
1867 ("session_pct", serde_json::json!(i64::MAX)),
1868 ("tp_session_pct", serde_json::json!("75")),
1869 ] {
1870 let mut v = full.clone();
1871 v[key] = bad;
1872 let bytes = serde_json::to_vec(&v).unwrap();
1873 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1874 assert!(err.to_string().contains(key), "{key}: {err}");
1875 }
1876 }
1877
1878 #[test]
1879 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
1880 let mut v = snap_to_json(&parsed());
1881 v["session_reset"] = serde_json::json!("not-a-time");
1882 let bytes = serde_json::to_vec(&v).unwrap();
1883 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1884 assert!(err.to_string().contains("session_reset"), "{err}");
1885 }
1886
1887 #[test]
1892 fn a_cache_past_its_reset_is_refused() {
1893 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
1894 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1895
1896 assert!(parse_cache_at(&bytes, None, now()).is_ok());
1898 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
1900
1901 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
1903 assert!(err.to_string().contains("5h"), "{err}");
1904
1905 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
1907 }
1908
1909 #[test]
1912 fn expiry_names_the_window_that_rolled_over() {
1913 let mut snap = parsed();
1914 snap.session.as_mut().unwrap().resets_at = None;
1916 snap.third_party_session = None;
1917 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1918 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1919
1920 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
1922
1923 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
1925 assert!(err.to_string().contains("Gemini weekly"), "{err}");
1926 }
1927
1928 #[test]
1930 fn a_window_without_a_reset_never_expires() {
1931 let mut snap = parsed();
1932 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
1933 w.resets_at = None;
1934 }
1935 snap.third_party_session = None;
1936 snap.third_party_weekly = None;
1937 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1938 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
1939 .unwrap()
1940 .with_timezone(&Utc);
1941 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
1942 }
1943
1944 #[test]
1946 fn a_cache_from_another_account_is_rejected() {
1947 let mut snap = parsed();
1948 snap.account = "acct:aaaa".into();
1949 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1950
1951 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
1952 assert_eq!(
1953 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
1954 snap
1955 );
1956
1957 let mut legacy = snap_to_json(&snap);
1959 legacy.as_object_mut().unwrap().remove("account");
1960 let legacy = serde_json::to_vec(&legacy).unwrap();
1961 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
1962 }
1963
1964 #[test]
1967 fn an_unverifiable_cache_is_served_rather_than_discarded() {
1968 let mut snap = parsed();
1969 snap.account = "acct:aaaa".into();
1970 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1971 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1972 }
1973
1974 #[test]
1975 fn account_key_fingerprints_rather_than_storing_the_address() {
1976 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
1977 let a = with("someone@example.com");
1978 assert!(!a.contains("someone"), "{a}");
1979 assert!(!a.contains('@'), "{a}");
1980 assert_eq!(a, with("someone@example.com"), "must be stable");
1981 assert_ne!(a, with("other@example.com"));
1982 let unknown = account_key(&serde_json::json!({}));
1984 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
1985 assert_ne!(unknown, a);
1986 }
1987
1988 #[test]
1991 fn absent_third_party_windows_are_not_treated_as_corruption() {
1992 let mut snap = parsed();
1993 snap.third_party_session = None;
1994 snap.third_party_weekly = None;
1995 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1996 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1997 }
1998
1999 #[test]
2000 fn cache_round_trip_preserves_absent_third_party_windows() {
2001 let mut snap = parsed();
2002 snap.third_party_session = None;
2003 snap.third_party_weekly = None;
2004 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2005 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2006 }
2007
2008 #[test]
2009 fn pct_used_inverts_valid_fractions() {
2010 assert_eq!(pct_used(1.0), 0);
2011 assert_eq!(pct_used(0.0), 100);
2012 assert_eq!(pct_used(0.5), 50);
2013 }
2014
2015 #[test]
2016 fn plan_falls_back_through_the_status_payload() {
2017 let tier: serde_json::Value =
2018 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2019 .unwrap();
2020 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2021
2022 let plan_only: serde_json::Value = serde_json::from_str(
2023 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2024 )
2025 .unwrap();
2026 assert_eq!(plan_from_status(&plan_only), "Pro");
2027
2028 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2029 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2030 }
2031
2032 #[cfg(target_os = "linux")]
2033 #[test]
2034 fn proc_net_parser_keeps_only_listening_rows() {
2035 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2036 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2037 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2038
2039 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2040 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2041 assert_eq!(parse_proc_net_line(established), None);
2042
2043 assert_eq!(parse_proc_net_line("garbage"), None);
2044 }
2045
2046 #[test]
2047 fn explicit_address_comes_first_and_gets_a_scheme() {
2048 let candidate = |base: &str, pid| Candidate {
2049 base: base.to_string(),
2050 pid,
2051 };
2052 assert_eq!(
2053 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)]),
2054 vec![
2055 candidate("http://127.0.0.1:1234", None),
2056 candidate("http://127.0.0.1:5678", Some(10)),
2057 ]
2058 );
2059 assert_eq!(
2061 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)]),
2062 vec![
2063 candidate("http://127.0.0.1:1234", None),
2064 candidate("http://127.0.0.1:5678", Some(10)),
2065 ]
2066 );
2067 assert_eq!(
2070 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)]),
2071 vec![candidate("http://127.0.0.1:5678", Some(10))]
2072 );
2073 assert_eq!(
2075 candidate_bases_with(None, vec![(10, 5678), (10, 5678)]),
2076 vec![candidate("http://127.0.0.1:5678", Some(10))]
2077 );
2078 assert_eq!(
2080 candidate_bases_with(Some("https://host:9"), vec![]),
2081 vec![candidate("https://host:9", None)]
2082 );
2083 }
2084
2085 fn http(status: u16) -> AppError {
2086 AppError::Http {
2087 status,
2088 body: String::new(),
2089 }
2090 }
2091
2092 fn missing_csrf() -> AppError {
2093 AppError::Http {
2094 status: 401,
2095 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2096 }
2097 }
2098
2099 #[test]
2100 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2101 assert!(is_missing_csrf(&missing_csrf()));
2102 assert!(!is_missing_csrf(&AppError::Http {
2103 status: 403,
2104 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2105 }));
2106 assert!(!is_missing_csrf(&AppError::Http {
2107 status: 401,
2108 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2109 }));
2110 assert!(!is_missing_csrf(&AppError::Http {
2111 status: 401,
2112 body: "prefix: missing CSRF token".into(),
2113 }));
2114 }
2115
2116 #[test]
2117 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2118 for errors in [
2119 vec![missing_csrf(), http(403)],
2120 vec![http(401), missing_csrf()],
2121 ] {
2122 let err = select_probe_error(errors);
2123 assert!(is_actionable(&err), "{err}");
2124 assert!(!is_missing_csrf(&err), "{err}");
2125 }
2126 }
2127
2128 #[test]
2129 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2130 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2131 assert!(is_missing_csrf(&err), "{err}");
2132 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2133 }
2134
2135 #[test]
2139 fn an_auth_failure_outranks_later_transport_noise() {
2140 let err = select_probe_error(vec![
2141 http(401),
2142 AppError::Transport("connection refused".into()),
2143 ]);
2144 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2145
2146 let err = select_probe_error(vec![
2147 AppError::Transport("connection refused".into()),
2148 http(403),
2149 ]);
2150 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2151 }
2152
2153 #[test]
2156 fn the_first_auth_failure_wins() {
2157 let err = select_probe_error(vec![http(401), http(403)]);
2158 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2159 }
2160
2161 #[test]
2165 fn without_an_auth_failure_the_last_error_stands() {
2166 let err = select_probe_error(vec![
2167 AppError::Transport("first".into()),
2168 http(500),
2169 AppError::Transport("last".into()),
2170 ]);
2171 assert!(
2172 matches!(&err, AppError::Transport(m) if m == "last"),
2173 "{err}"
2174 );
2175 assert!(err.is_transient());
2176 }
2177
2178 #[test]
2181 fn a_server_error_is_not_treated_as_actionable() {
2182 let err = select_probe_error(vec![http(500), http(401)]);
2183 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2184 }
2185
2186 #[test]
2187 fn no_candidates_at_all_yields_a_generic_error() {
2188 let err = select_probe_error(Vec::new());
2189 assert!(
2190 err.to_string().contains("no local server answered"),
2191 "{err}"
2192 );
2193 }
2194
2195 fn tls_echo() -> AppError {
2199 AppError::Http {
2200 status: 400,
2201 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2202 }
2203 }
2204
2205 #[test]
2210 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2211 let err = select_probe_error(vec![
2212 AppError::Http {
2213 status: 500,
2214 body: "GetUserStatus: internal".into(),
2215 },
2216 tls_echo(),
2217 ]);
2218 assert!(
2219 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2220 "{err}"
2221 );
2222 }
2223
2224 #[test]
2229 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2230 let err = select_probe_error(vec![
2231 AppError::Transport("connection refused".into()),
2232 tls_echo(),
2233 ]);
2234 assert!(
2235 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2236 "{err}"
2237 );
2238 assert!(
2239 err.is_transient(),
2240 "the echo must not make the run non-transient: {err}"
2241 );
2242 }
2243
2244 #[test]
2247 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2248 let err = select_probe_error(vec![tls_echo()]);
2249 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2250 }
2251
2252 #[test]
2255 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2256 let err = select_probe_error(vec![
2257 AppError::Http {
2258 status: 400,
2259 body: "unknown method GetUserStatus".into(),
2260 },
2261 tls_echo(),
2262 ]);
2263 assert!(
2264 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2265 "{err}"
2266 );
2267 }
2268
2269 #[test]
2271 fn an_auth_failure_still_outranks_a_tls_echo() {
2272 let err = select_probe_error(vec![tls_echo(), http(401)]);
2273 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2274 }
2275
2276 #[test]
2277 fn every_discovered_port_is_probed_in_order() {
2278 assert_eq!(
2279 candidate_bases_with(None, vec![(10, 33875), (10, 37435)]),
2280 vec![
2281 Candidate {
2282 base: "http://127.0.0.1:33875".to_string(),
2283 pid: Some(10),
2284 },
2285 Candidate {
2286 base: "http://127.0.0.1:37435".to_string(),
2287 pid: Some(10),
2288 },
2289 ]
2290 );
2291 }
2292
2293 #[test]
2298 fn empty_discovery_yields_no_candidates() {
2299 assert!(candidate_bases_with(None, vec![]).is_empty());
2300 assert!(candidate_bases_with(Some(""), vec![]).is_empty());
2301 }
2302
2303 #[test]
2304 fn every_antigravity_product_is_recognised() {
2305 assert!(is_antigravity_process(
2307 "language_server\n",
2308 Some("/opt/antigravity/resources/bin/language_server")
2309 ));
2310 assert!(is_antigravity_process(
2312 "agy\n",
2313 Some("/home/u/.local/bin/agy")
2314 ));
2315 assert!(is_antigravity_process(
2317 "node",
2318 Some("/opt/antigravity/bin/helper")
2319 ));
2320 assert!(is_antigravity_process("antigravity", None));
2321 assert!(is_antigravity_process("agy.exe", None));
2322 assert!(is_antigravity_process("Antigravity.exe", None));
2323 assert!(is_antigravity_process("language_server.exe", None));
2324 assert!(is_antigravity_process(
2325 "language_server_windows_x64.exe",
2326 None
2327 ));
2328 assert!(is_antigravity_process(
2329 "node.exe",
2330 Some(r"C:\Users\u\AppData\Local\agy.exe")
2331 ));
2332 }
2333
2334 #[test]
2335 fn unrelated_processes_are_not_probed() {
2336 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2337 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2338 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2340 assert!(!is_antigravity_process("legacy.exe", None));
2341 assert!(!is_antigravity_process("not-agy.exe", None));
2342 assert!(!is_antigravity_process("", None));
2343 }
2344
2345 #[test]
2346 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2347 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2348 raw.extend([0, b'x' as u16]);
2349 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2350 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2351 assert_eq!(decode_windows_process_name(&[]), "");
2352 }
2353
2354 #[test]
2355 fn windows_process_filter_keeps_only_antigravity_pids() {
2356 let processes = vec![
2357 (10, "agy.exe".to_string()),
2358 (20, "language_server_windows_x64.exe".to_string()),
2359 (30, "sshd.exe".to_string()),
2360 ];
2361 let pids = matching_windows_process_ids(&processes);
2362 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
2363 }
2364
2365 #[test]
2366 fn windows_listener_filter_joins_pid_loopback_and_port() {
2367 let pids = std::collections::HashSet::from([10]);
2368 let rows = [
2369 WindowsTcpRow {
2370 local_addr: [127, 0, 0, 1],
2371 local_port: u32::from(59870u16.to_be()),
2372 pid: 10,
2373 },
2374 WindowsTcpRow {
2375 local_addr: [127, 0, 0, 1],
2376 local_port: u32::from(59868u16.to_be()),
2377 pid: 10,
2378 },
2379 WindowsTcpRow {
2380 local_addr: [127, 0, 0, 1],
2381 local_port: u32::from(59870u16.to_be()),
2382 pid: 10,
2383 },
2384 WindowsTcpRow {
2385 local_addr: [0, 0, 0, 0],
2386 local_port: u32::from(50000u16.to_be()),
2387 pid: 10,
2388 },
2389 WindowsTcpRow {
2390 local_addr: [127, 0, 0, 1],
2391 local_port: u32::from(50001u16.to_be()),
2392 pid: 99,
2393 },
2394 WindowsTcpRow {
2395 local_addr: [127, 0, 0, 1],
2396 local_port: 0,
2397 pid: 10,
2398 },
2399 ];
2400 assert_eq!(
2401 matching_windows_ports(&pids, &rows),
2402 vec![(10, 59870), (10, 59868)]
2403 );
2404 }
2405
2406 #[test]
2410 fn windows_ports_from_two_products_keep_tls_listeners_last() {
2411 let pids = std::collections::HashSet::from([10, 20]);
2412 let row = |port: u16, pid: u32| WindowsTcpRow {
2413 local_addr: [127, 0, 0, 1],
2414 local_port: u32::from(port.to_be()),
2415 pid,
2416 };
2417 let rows = [
2418 row(40000, 10),
2419 row(40001, 10),
2420 row(50000, 20),
2421 row(50001, 20),
2422 ];
2423 assert_eq!(
2424 matching_windows_ports(&pids, &rows),
2425 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2426 );
2427 }
2428
2429 #[test]
2433 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
2434 use std::collections::BTreeMap;
2435
2436 assert_eq!(
2438 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
2439 vec![(10, 59870), (10, 59868)]
2440 );
2441 assert_eq!(
2446 probe_order(BTreeMap::from([
2447 (10, vec![40000, 40001]),
2448 (20, vec![50000, 50001]),
2449 ])),
2450 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2451 );
2452 assert_eq!(
2455 probe_order(BTreeMap::from([
2456 (10, vec![6000, 5000, 4000]),
2457 (20, vec![6000, 7000]),
2458 ])),
2459 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
2460 );
2461 assert!(probe_order(BTreeMap::new()).is_empty());
2462 }
2463
2464 #[test]
2469 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
2470 use std::collections::BTreeMap;
2471
2472 assert_eq!(
2473 probe_order(BTreeMap::from([
2474 (10, vec![40001, 40001, 40000, 40000]),
2475 (20, vec![50001, 50000]),
2476 ])),
2477 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
2478 "duplicate rows must not reorder the ranks below them"
2479 );
2480 }
2481
2482 #[test]
2487 fn a_half_started_product_is_the_documented_exception() {
2488 use std::collections::BTreeMap;
2489
2490 assert_eq!(
2491 probe_order(BTreeMap::from([
2492 (10, vec![40000]),
2493 (20, vec![50001, 50000])
2494 ])),
2495 vec![(10, 40000), (20, 50001), (20, 50000)]
2496 );
2497 }
2498
2499 #[test]
2504 fn an_override_with_no_authority_is_dropped_not_probed() {
2505 for junk in ["/", "///", "http://", "https://", " / "] {
2506 assert_eq!(
2507 candidate_bases_with(Some(junk), vec![(10, 4242)]),
2508 vec![Candidate {
2509 base: "http://127.0.0.1:4242".to_string(),
2510 pid: Some(10),
2511 }],
2512 "{junk:?} should not survive as a candidate"
2513 );
2514 }
2515 assert!(candidate_bases_with(Some("/"), vec![]).is_empty());
2516 }
2517
2518 #[test]
2519 fn windows_table_bounds_reject_truncation_and_overflow() {
2520 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
2521 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
2522 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
2523 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
2524 assert_eq!(
2525 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
2526 None
2527 );
2528 }
2529
2530 #[cfg(target_os = "windows")]
2531 #[test]
2532 fn windows_tcp_table_parser_copies_complete_rows_only() {
2533 use std::mem::{offset_of, size_of};
2534 use windows_sys::Win32::NetworkManagement::IpHelper::{
2535 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
2536 };
2537
2538 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
2539 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
2540 let words = used.div_ceil(size_of::<u32>());
2541 let mut buffer = vec![0u32; words];
2542 let first = MIB_TCPROW_OWNER_PID {
2543 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2544 dwLocalPort: u32::from(59868u16.to_be()),
2545 dwOwningPid: 10,
2546 ..Default::default()
2547 };
2548 let second = MIB_TCPROW_OWNER_PID {
2549 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2550 dwLocalPort: u32::from(59870u16.to_be()),
2551 dwOwningPid: 10,
2552 ..Default::default()
2553 };
2554 unsafe {
2555 buffer.as_mut_ptr().write_unaligned(2);
2556 let rows = buffer
2557 .as_mut_ptr()
2558 .cast::<u8>()
2559 .add(offset)
2560 .cast::<MIB_TCPROW_OWNER_PID>();
2561 rows.write_unaligned(first);
2562 rows.add(1).write_unaligned(second);
2563 }
2564
2565 assert_eq!(
2566 parse_windows_tcp_rows(&buffer, used),
2567 vec![
2568 WindowsTcpRow {
2569 local_addr: [127, 0, 0, 1],
2570 local_port: u32::from(59868u16.to_be()),
2571 pid: 10,
2572 },
2573 WindowsTcpRow {
2574 local_addr: [127, 0, 0, 1],
2575 local_port: u32::from(59870u16.to_be()),
2576 pid: 10,
2577 },
2578 ]
2579 );
2580 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
2581 }
2582
2583 #[test]
2584 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
2585 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";
2588 assert_eq!(
2589 parse_lsof_pcn(output),
2590 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
2591 );
2592 }
2593
2594 #[test]
2598 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
2599 let output = concat!(
2600 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
2601 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
2602 );
2603 assert_eq!(
2604 parse_lsof_pcn(output),
2605 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
2606 "both HTTP listeners must precede both TLS listeners"
2607 );
2608 }
2609
2610 #[test]
2611 fn lsof_parser_matches_the_capitalised_macos_app_name() {
2612 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
2613 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
2614 }
2615
2616 #[test]
2617 fn lsof_parser_deduplicates_and_handles_empty_output() {
2618 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
2619 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
2620 assert!(parse_lsof_pcn("").is_empty());
2621 }
2622
2623 #[test]
2626 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
2627 let dir = tempfile::tempdir().unwrap();
2628 let cache = Cache::at(dir.path().join("usage.json"));
2629 let reason = AppError::Credentials("Antigravity: no local language server found".into());
2630
2631 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2632 let msg = err.to_string();
2633 assert!(msg.contains("no local language server found"), "{msg}");
2634 assert!(!msg.contains("no usable cache"), "{msg}");
2635 }
2636
2637 #[test]
2638 fn unusable_cache_does_not_replace_the_live_diagnosis() {
2639 let dir = tempfile::tempdir().unwrap();
2640 let cache = Cache::at(dir.path().join("antigravity"));
2641 cache.write_payload(b"{}").unwrap();
2642
2643 let reason = AppError::Credentials("Antigravity must be running".into());
2644 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2645 assert!(err.to_string().contains("must be running"), "{err}");
2646
2647 let original = AppError::Transport("original loopback failure".into());
2648 let err = fallback_silent(&cache, now(), original).unwrap_err();
2649 assert!(
2650 err.to_string().contains("original loopback failure"),
2651 "{err}"
2652 );
2653 }
2654
2655 #[tokio::test]
2656 async fn rpc_error_bodies_are_bounded_too() {
2657 let mut server = mockito::Server::new_async().await;
2658 let path = format!("/{STATUS_RPC}");
2659 server
2660 .mock("POST", path.as_str())
2661 .with_status(500)
2662 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
2663 .create_async()
2664 .await;
2665
2666 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
2667 .await
2668 .unwrap_err();
2669 assert!(err.to_string().contains("exceeds"), "{err}");
2670 }
2671
2672 #[test]
2673 fn blank_override_falls_through_to_discovery() {
2674 assert_eq!(
2675 candidate_bases_with(Some(" "), vec![(10, 4242)]),
2676 vec![Candidate {
2677 base: "http://127.0.0.1:4242".to_string(),
2678 pid: Some(10),
2679 }]
2680 );
2681 }
2682
2683 fn fixture() -> (tempfile::TempDir, Cache) {
2688 let td = tempfile::TempDir::new().unwrap();
2689 let cache = Cache::at(td.path().join("antigravity"));
2690 (td, cache)
2691 }
2692
2693 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
2694 let base = server.url();
2695 cloud::Endpoints {
2696 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
2697 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
2698 token: format!("{base}/token"),
2699 }
2700 }
2701
2702 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
2704 let mut token = serde_json::json!({
2705 "access_token": "KEYRING-AT",
2706 "expiry": expiry,
2707 });
2708 if with_refresh {
2709 token["refresh_token"] = serde_json::json!("KEYRING-RT");
2710 }
2711 serde_json::json!({ "token": token }).to_string()
2712 }
2713
2714 const EXPIRED: &str = "2026-07-22T11:00:00Z";
2716 const VALID: &str = "2026-07-22T13:00:00Z";
2718
2719 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
2721 RemoteOverride {
2722 credential: SavedCredential::Blob(blob),
2723 endpoints: Some(eps),
2724 local_bases: Some(vec![]),
2725 }
2726 }
2727
2728 fn bare_summary() -> String {
2731 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2732 v["response"].to_string()
2733 }
2734
2735 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
2736 server
2737 .mock("POST", "/daily/quota")
2738 .match_header("authorization", format!("Bearer {bearer}").as_str())
2739 .with_status(200)
2740 .with_body(bare_summary())
2741 }
2742
2743 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
2744 server
2745 .mock("POST", "/token")
2746 .with_status(200)
2747 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
2748 }
2749
2750 fn test_oauth() -> cloud::OauthClient {
2751 cloud::OauthClient {
2752 id: "test-client".into(),
2753 secret: "test-client-secret".into(),
2754 }
2755 }
2756
2757 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
2758 fetch_snapshot_at(
2759 &reqwest::Client::new(),
2760 cache,
2761 ttl,
2762 Some(&test_oauth()),
2763 remote,
2764 now(),
2765 )
2766 .await
2767 }
2768
2769 #[tokio::test]
2770 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
2771 let mut server = mockito::Server::new_async().await;
2772 let eps = endpoints(&server);
2773 let quota = quota_mock(&mut server, "KEYRING-AT")
2774 .expect(1)
2775 .create_async()
2776 .await;
2777 let plan = server
2778 .mock("POST", "/daily/plan")
2779 .match_header("authorization", "Bearer KEYRING-AT")
2780 .with_status(200)
2781 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2782 .expect(1)
2783 .create_async()
2784 .await;
2785 let token = token_mock(&mut server).expect(0).create_async().await;
2787 let blob = keyring_blob(VALID, true);
2788 let (_td, cache) = fixture();
2789
2790 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
2791 .await
2792 .expect("remote path yields a snapshot");
2793
2794 quota.assert_async().await;
2795 plan.assert_async().await;
2796 token.assert_async().await;
2797 let snap = outcome.snapshot;
2798 assert!(!outcome.stale);
2799 assert_eq!(snap.source, AntigravitySource::Remote);
2800 assert_eq!(snap.plan, "Pro");
2801 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2802 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2803 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
2804 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
2805 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
2806 assert_eq!(snap.account, format!("acct:{fingerprint}"));
2807 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
2808
2809 let cached = parse_cache_at(
2811 &cache
2812 .fresh_payload(Duration::from_secs(60))
2813 .unwrap()
2814 .unwrap(),
2815 None,
2816 now(),
2817 )
2818 .unwrap();
2819 assert_eq!(cached.source, AntigravitySource::Remote);
2820 assert_eq!(cached.account, snap.account);
2821 }
2822
2823 #[tokio::test]
2828 async fn agys_missing_csrf_response_uses_the_saved_session() {
2829 let mut server = mockito::Server::new_async().await;
2830 let eps = endpoints(&server);
2831 let root = server
2832 .mock("GET", "/")
2833 .with_status(404)
2834 .expect(1)
2835 .create_async()
2836 .await;
2837 let status_path = format!("/{STATUS_RPC}");
2838 let status = server
2839 .mock("POST", status_path.as_str())
2840 .with_status(401)
2841 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2842 .expect(1)
2843 .create_async()
2844 .await;
2845 let quota = quota_mock(&mut server, "KEYRING-AT")
2846 .expect(1)
2847 .create_async()
2848 .await;
2849 let plan = server
2850 .mock("POST", "/daily/plan")
2851 .match_header("authorization", "Bearer KEYRING-AT")
2852 .with_status(200)
2853 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2854 .expect(1)
2855 .create_async()
2856 .await;
2857 let blob = keyring_blob(VALID, true);
2858 let (_td, cache) = fixture();
2859
2860 let outcome = run(
2861 &cache,
2862 RemoteOverride {
2863 credential: SavedCredential::Blob(&blob),
2864 endpoints: Some(&eps),
2865 local_bases: Some(vec![server.url().into()]),
2866 },
2867 Duration::ZERO,
2868 )
2869 .await
2870 .expect("saved session bypasses agy's unusable local RPC");
2871
2872 root.assert_async().await;
2873 status.assert_async().await;
2874 quota.assert_async().await;
2875 plan.assert_async().await;
2876 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
2877 }
2878
2879 #[tokio::test]
2880 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
2881 let mut server = mockito::Server::new_async().await;
2882 let eps = endpoints(&server);
2883 let root = server
2884 .mock("GET", "/")
2885 .with_status(404)
2886 .expect(1)
2887 .create_async()
2888 .await;
2889 let status_path = format!("/{STATUS_RPC}");
2890 let status = server
2891 .mock("POST", status_path.as_str())
2892 .with_status(401)
2893 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2894 .expect(1)
2895 .create_async()
2896 .await;
2897 let quota = quota_mock(&mut server, "KEYRING-AT")
2898 .expect(0)
2899 .create_async()
2900 .await;
2901 let (_td, cache) = fixture();
2902
2903 let error = run(
2904 &cache,
2905 RemoteOverride {
2906 credential: SavedCredential::Absent,
2907 endpoints: Some(&eps),
2908 local_bases: Some(vec![server.url().into()]),
2909 },
2910 Duration::ZERO,
2911 )
2912 .await
2913 .expect_err("neither the local RPC nor a saved session is usable");
2914
2915 root.assert_async().await;
2916 status.assert_async().await;
2917 quota.assert_async().await;
2918 let message = error.to_string();
2919 assert!(message.contains("requires a CSRF token"), "{message}");
2920 assert!(message.contains("saved Google session"), "{message}");
2921 }
2922
2923 #[tokio::test]
2928 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
2929 let mut server1 = mockito::Server::new_async().await;
2930 let mut server2 = mockito::Server::new_async().await;
2931 let eps = endpoints(&server1);
2932 let root1 = server1
2933 .mock("GET", "/")
2934 .with_status(404)
2935 .expect(1)
2936 .create_async()
2937 .await;
2938 let status_path = format!("/{STATUS_RPC}");
2939 let status1 = server1
2940 .mock("POST", status_path.as_str())
2941 .with_status(401)
2942 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2943 .expect(1)
2944 .create_async()
2945 .await;
2946
2947 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
2949 let status2 = server2
2950 .mock("POST", status_path.as_str())
2951 .expect(0)
2952 .create_async()
2953 .await;
2954
2955 let (_td, cache) = fixture();
2956 let error = run(
2957 &cache,
2958 RemoteOverride {
2959 credential: SavedCredential::Absent,
2960 endpoints: Some(&eps),
2961 local_bases: Some(vec![
2962 Candidate {
2963 base: server1.url(),
2964 pid: Some(42),
2965 },
2966 Candidate {
2967 base: server2.url(),
2968 pid: Some(42),
2969 },
2970 ]),
2971 },
2972 Duration::ZERO,
2973 )
2974 .await
2975 .expect_err("aborted probe returns agy missing csrf diagnosis");
2976
2977 root1.assert_async().await;
2978 status1.assert_async().await;
2979 root2.assert_async().await;
2980 status2.assert_async().await;
2981 let message = error.to_string();
2982 assert!(message.contains("requires a CSRF token"), "{message}");
2983 }
2984
2985 #[tokio::test]
2992 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
2993 let mut agy = mockito::Server::new_async().await;
2994 let mut ide = mockito::Server::new_async().await;
2995 let eps = endpoints(&agy);
2996 let agy_root = agy
2997 .mock("GET", "/")
2998 .with_status(404)
2999 .expect(1)
3000 .create_async()
3001 .await;
3002 let status_path = format!("/{STATUS_RPC}");
3003 let agy_status = agy
3004 .mock("POST", status_path.as_str())
3005 .with_status(401)
3006 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3007 .expect(1)
3008 .create_async()
3009 .await;
3010
3011 let ide_root = ide
3012 .mock("GET", "/")
3013 .with_status(404)
3014 .expect(1)
3015 .create_async()
3016 .await;
3017 let ide_status = ide
3018 .mock("POST", status_path.as_str())
3019 .with_status(200)
3020 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3021 .expect(1)
3022 .create_async()
3023 .await;
3024 let quota_path = format!("/{QUOTA_RPC}");
3025 let ide_quota = ide
3026 .mock("POST", quota_path.as_str())
3027 .with_status(200)
3028 .with_body(QUOTA_JSON)
3029 .expect(1)
3030 .create_async()
3031 .await;
3032
3033 let (_td, cache) = fixture();
3034 let outcome = run(
3035 &cache,
3036 RemoteOverride {
3037 credential: SavedCredential::Absent,
3038 endpoints: Some(&eps),
3039 local_bases: Some(vec![
3040 Candidate {
3041 base: agy.url(),
3042 pid: Some(42),
3043 },
3044 Candidate {
3045 base: ide.url(),
3046 pid: Some(7),
3047 },
3048 ]),
3049 },
3050 Duration::ZERO,
3051 )
3052 .await
3053 .expect("the other product's server still answers");
3054
3055 agy_root.assert_async().await;
3056 agy_status.assert_async().await;
3057 ide_root.assert_async().await;
3058 ide_status.assert_async().await;
3059 ide_quota.assert_async().await;
3060 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3061 assert_eq!(outcome.snapshot.plan, "Pro");
3062 }
3063
3064 #[tokio::test]
3067 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
3068 let mut server = mockito::Server::new_async().await;
3069 let eps = endpoints(&server);
3070 let _quota = quota_mock(&mut server, "NEW-AT")
3071 .expect(2)
3072 .create_async()
3073 .await;
3074 let token = token_mock(&mut server).expect(1).create_async().await;
3075 let blob = keyring_blob(EXPIRED, true);
3076 let (_td, cache) = fixture();
3077
3078 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
3079 .await
3080 .expect("refresh, then quota");
3081 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
3082 token.assert_async().await;
3083
3084 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3085 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
3086 .expect("refreshed token persisted under the keyring session's fingerprint");
3087 assert_eq!(persisted.access_token, "NEW-AT");
3088
3089 run(&cache, remote(&blob, &eps), Duration::ZERO)
3091 .await
3092 .expect("persisted token reused");
3093 token.assert_async().await;
3094 }
3095
3096 #[tokio::test]
3099 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
3100 let mut server = mockito::Server::new_async().await;
3101 let eps = endpoints(&server);
3102 let rejected = server
3103 .mock("POST", "/daily/quota")
3104 .match_header("authorization", "Bearer KEYRING-AT")
3105 .with_status(401)
3106 .expect(1)
3107 .create_async()
3108 .await;
3109 let accepted = quota_mock(&mut server, "NEW-AT")
3110 .expect(1)
3111 .create_async()
3112 .await;
3113 let token = token_mock(&mut server).expect(1).create_async().await;
3114 let blob = keyring_blob(VALID, true);
3115 let (_td, cache) = fixture();
3116
3117 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3118 .await
3119 .expect("retry with the refreshed token succeeds");
3120
3121 rejected.assert_async().await;
3122 accepted.assert_async().await;
3123 token.assert_async().await;
3124 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3125 }
3126
3127 #[tokio::test]
3130 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
3131 let mut server = mockito::Server::new_async().await;
3132 let eps = endpoints(&server);
3133 let _token = server
3134 .mock("POST", "/token")
3135 .with_status(400)
3136 .with_body(r#"{"error":"invalid_grant"}"#)
3137 .create_async()
3138 .await;
3139 let blob = keyring_blob(EXPIRED, true);
3140 let (_td, cache) = fixture();
3141
3142 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3143 .await
3144 .expect_err("no cache to fall back on");
3145
3146 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3147 let rendered = err.to_string();
3148 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
3149 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
3150 }
3151
3152 #[tokio::test]
3155 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
3156 let mut server = mockito::Server::new_async().await;
3157 let eps = endpoints(&server);
3158 let _quota = server
3159 .mock("POST", "/daily/quota")
3160 .with_status(401)
3161 .create_async()
3162 .await;
3163 let token = token_mock(&mut server).expect(1).create_async().await;
3164 let blob = keyring_blob(EXPIRED, true);
3165 let (_td, cache) = fixture();
3166
3167 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3168 .await
3169 .expect_err("rejected after refresh");
3170
3171 token.assert_async().await;
3172 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3173 assert!(err.to_string().contains("sign in again"), "{err}");
3174 }
3175
3176 #[tokio::test]
3178 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
3179 let mut server = mockito::Server::new_async().await;
3180 let eps = endpoints(&server);
3181 let token = token_mock(&mut server).expect(0).create_async().await;
3182 let blob = keyring_blob(EXPIRED, false);
3183 let (_td, cache) = fixture();
3184
3185 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3186 .await
3187 .expect_err("nothing to refresh with");
3188
3189 token.assert_async().await;
3190 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3191 assert!(err.to_string().contains("sign in again"), "{err}");
3192 }
3193
3194 #[tokio::test]
3197 async fn no_saved_session_extends_the_no_local_server_error() {
3198 let mut server = mockito::Server::new_async().await;
3199 let eps = endpoints(&server);
3200 let quota = quota_mock(&mut server, "KEYRING-AT")
3201 .expect(0)
3202 .create_async()
3203 .await;
3204 let (_td, cache) = fixture();
3205
3206 let err = run(
3207 &cache,
3208 RemoteOverride {
3209 credential: SavedCredential::Absent,
3210 endpoints: Some(&eps),
3211 local_bases: Some(vec![]),
3212 },
3213 Duration::ZERO,
3214 )
3215 .await
3216 .expect_err("no source at all");
3217
3218 quota.assert_async().await;
3219 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3220 let rendered = err.to_string();
3221 assert!(rendered.contains("no local server found"), "{rendered}");
3222 assert!(rendered.contains("saved Google session"), "{rendered}");
3223 }
3224
3225 #[tokio::test]
3228 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
3229 let mut server = mockito::Server::new_async().await;
3230 let eps = endpoints(&server);
3231 let status_path = format!("/{STATUS_RPC}");
3232 let _status = server
3233 .mock("POST", status_path.as_str())
3234 .with_status(401)
3235 .create_async()
3236 .await;
3237 let quota = quota_mock(&mut server, "KEYRING-AT")
3238 .expect(0)
3239 .create_async()
3240 .await;
3241 let blob = keyring_blob(VALID, true);
3242 let (_td, cache) = fixture();
3243
3244 let err = run(
3245 &cache,
3246 RemoteOverride {
3247 credential: SavedCredential::Blob(&blob),
3248 endpoints: Some(&eps),
3249 local_bases: Some(vec![server.url().into()]),
3250 },
3251 Duration::ZERO,
3252 )
3253 .await
3254 .expect_err("signed-out local server");
3255
3256 quota.assert_async().await;
3257 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
3258 }
3259
3260 #[tokio::test]
3263 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
3264 let mut server = mockito::Server::new_async().await;
3265 let eps = endpoints(&server);
3266 for path in ["/daily/quota", "/prod/quota"] {
3267 server
3268 .mock("POST", path)
3269 .with_status(500)
3270 .with_body("boom")
3271 .create_async()
3272 .await;
3273 }
3274 let blob = keyring_blob(VALID, true);
3275 let (_td, cache) = fixture();
3276 cache.ensure_dir().unwrap();
3277 let mut earlier = parsed();
3278 earlier.source = AntigravitySource::Remote;
3279 earlier.account = "acct:earlier".into();
3280 cache
3281 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
3282 .unwrap();
3283
3284 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3285 .await
3286 .expect("stale cache stands in");
3287
3288 assert!(outcome.stale);
3289 assert_eq!(outcome.snapshot, earlier);
3290 assert!(
3291 matches!(outcome.last_error, Some((500, _))),
3292 "{:?}",
3293 outcome.last_error
3294 );
3295 }
3296
3297 #[test]
3298 fn source_round_trips_through_the_cache_and_defaults_to_local() {
3299 let mut snap = parsed();
3300 snap.source = AntigravitySource::Remote;
3301 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3302 assert_eq!(
3303 parse_cache_at(&bytes, None, now()).unwrap().source,
3304 AntigravitySource::Remote
3305 );
3306
3307 let mut legacy = snap_to_json(&snap);
3309 legacy.as_object_mut().unwrap().remove("source");
3310 let legacy = serde_json::to_vec(&legacy).unwrap();
3311 assert_eq!(
3312 parse_cache_at(&legacy, None, now()).unwrap().source,
3313 AntigravitySource::Local
3314 );
3315 }
3316}