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 File(&'a std::path::Path),
82 Absent,
84}
85
86#[derive(Default)]
91pub struct RemoteOverride<'a> {
92 pub credential: SavedCredential<'a>,
94 pub endpoints: Option<&'a cloud::Endpoints>,
96 pub local_bases: Option<Vec<Candidate>>,
101 pub agy: super::agy::AgyCommand,
105}
106
107#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct Candidate {
113 pub base: String,
115 pub pid: Option<u32>,
118}
119
120impl From<String> for Candidate {
121 fn from(base: String) -> Self {
122 Self { base, pid: None }
123 }
124}
125
126pub async fn fetch_snapshot(
127 client: &reqwest::Client,
128 cache: &Cache,
129 cache_ttl: Duration,
130 oauth: Option<&cloud::OauthClient>,
131) -> Result<FetchOutcome> {
132 fetch_snapshot_at(
133 client,
134 cache,
135 cache_ttl,
136 oauth,
137 RemoteOverride {
138 agy: super::agy::AgyCommand::Discover,
139 ..Default::default()
140 },
141 Utc::now(),
142 )
143 .await
144}
145
146pub async fn fetch_snapshot_at(
150 client: &reqwest::Client,
151 cache: &Cache,
152 cache_ttl: Duration,
153 oauth: Option<&cloud::OauthClient>,
154 remote: RemoteOverride<'_>,
155 now: DateTime<Utc>,
156) -> Result<FetchOutcome> {
157 cache.ensure_dir()?;
158 let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
159
160 let origin = match open_session(client, remote.local_bases.as_deref()).await {
170 Ok(session) => Origin::Local(Ok(session)),
171 Err(error) => match remote_fallback_reason(&error) {
172 Some(reason) => Origin::Remote(saved_session(remote.credential, reason)),
173 None => Origin::Local(Err(error)),
174 },
175 };
176 let account = origin.account();
177
178 if let Some(bytes) = cache.fresh_payload(cache_ttl)?
179 && let Ok(outcome) = reuse_cache(bytes, cache, false, account.as_deref(), now)
180 {
181 return Ok(outcome);
182 }
183
184 let default_endpoints = cloud::Endpoints::default();
185 let endpoints = remote.endpoints.unwrap_or(&default_endpoints);
186 let live = match origin {
187 Origin::Local(session) => fetch_live(client, session).await,
188 Origin::Remote(token) => {
189 let renewal = Renewal {
190 agy: &remote.agy,
191 credential: remote.credential,
192 };
193 fetch_remote(client, cache, oauth, endpoints, token, renewal, now).await
194 }
195 };
196
197 match live {
198 Ok(snap) => {
199 let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
200 cache.write_payload(&bytes)?;
201 Ok(crate::outcome::Outcome::fresh(snap))
202 }
203 Err(e) if e.is_transient() => fallback_silent(cache, now, e),
204 Err(AppError::Http { status, body }) => {
205 cache.mark_stale();
206 let last_error = Some(cache.write_last_error(status, &body));
207 let reason = AppError::Http { status, body };
208 fallback_with_error(cache, last_error, reason, now)
209 }
210 Err(e) => {
211 cache.mark_stale();
212 let last_error = Some(cache.write_last_error(0, &e.to_string()));
213 fallback_with_error(cache, last_error, e, now)
214 }
215 }
216}
217
218struct Session {
221 base: String,
222 csrf: Option<String>,
223 plan: String,
224 account: String,
225}
226
227enum Origin {
230 Local(Result<Session>),
231 Remote(Result<StoredToken>),
232}
233
234impl Origin {
235 fn account(&self) -> Option<String> {
237 match self {
238 Origin::Local(Ok(session)) => Some(session.account.clone()),
239 Origin::Remote(Ok(token)) => Some(remote_account(&token.fingerprint)),
240 Origin::Local(Err(_)) | Origin::Remote(Err(_)) => None,
241 }
242 }
243}
244
245fn no_local_server() -> AppError {
246 AppError::Credentials(NO_LOCAL_SERVER.into())
247}
248
249#[derive(Clone, Copy)]
250enum RemoteFallbackReason {
251 NoLocalServer,
252 AgyMissingCsrf,
253}
254
255impl RemoteFallbackReason {
256 fn message(self) -> &'static str {
257 match self {
258 Self::NoLocalServer => NO_LOCAL_SERVER,
259 Self::AgyMissingCsrf => AGY_CSRF_UNAVAILABLE,
260 }
261 }
262}
263
264fn remote_fallback_reason(error: &AppError) -> Option<RemoteFallbackReason> {
268 if matches!(error, AppError::Credentials(message) if message == NO_LOCAL_SERVER) {
269 Some(RemoteFallbackReason::NoLocalServer)
270 } else if is_missing_csrf(error) {
271 Some(RemoteFallbackReason::AgyMissingCsrf)
272 } else {
273 None
274 }
275}
276
277async fn open_session(client: &reqwest::Client, bases: Option<&[Candidate]>) -> Result<Session> {
283 let candidates = bases.map_or_else(candidate_bases, <[Candidate]>::to_vec);
284 if candidates.is_empty() {
285 return Err(no_local_server());
286 }
287
288 let mut errors = Vec::new();
289 let mut blocked_pids = std::collections::HashSet::new();
294 for candidate in candidates {
295 if candidate.pid.is_some_and(|pid| blocked_pids.contains(&pid)) {
296 continue;
297 }
298 let csrf = fetch_csrf(client, &candidate.base).await;
299 match post_rpc(client, &candidate.base, csrf.as_deref(), STATUS_RPC).await {
300 Ok(v) => {
301 return Ok(Session {
302 base: candidate.base,
303 csrf,
304 plan: plan_from_status(&v),
305 account: account_key(&v),
306 });
307 }
308 Err(e) => {
309 if is_missing_csrf(&e)
313 && let Some(pid) = candidate.pid
314 {
315 blocked_pids.insert(pid);
316 }
317 errors.push(e);
318 }
319 }
320 }
321 Err(select_probe_error(errors))
322}
323
324fn select_probe_error(errors: Vec<AppError>) -> AppError {
341 let mut actionable = None;
342 let mut missing_csrf = None;
343 let mut last = None;
344 let mut echo = None;
345 for e in errors {
346 if missing_csrf.is_none() && is_missing_csrf(&e) {
347 missing_csrf = Some(e);
348 } else if actionable.is_none() && is_actionable(&e) {
349 actionable = Some(e);
350 } else if is_tls_echo(&e) {
351 echo = Some(e);
352 } else {
353 last = Some(e);
354 }
355 }
356 actionable
357 .or(missing_csrf)
358 .or(last)
359 .or(echo)
360 .unwrap_or_else(|| {
361 AppError::Other("antigravity: no local server answered GetUserStatus".into())
362 })
363}
364
365fn is_missing_csrf(error: &AppError) -> bool {
370 let AppError::Http { status: 401, body } = error else {
371 return false;
372 };
373 let Ok(body) = serde_json::from_str::<serde_json::Value>(body) else {
374 return false;
375 };
376 matches!(
377 (body["code"].as_str(), body["message"].as_str()),
378 (Some(code), Some(message))
379 if code.eq_ignore_ascii_case("unauthenticated")
380 && message.trim().eq_ignore_ascii_case("missing CSRF token")
381 )
382}
383
384fn is_actionable(e: &AppError) -> bool {
389 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
390 && !is_missing_csrf(e)
391}
392
393fn is_tls_echo(e: &AppError) -> bool {
406 matches!(
407 e,
408 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
409 )
410}
411
412async fn fetch_live(
413 client: &reqwest::Client,
414 session: Result<Session>,
415) -> Result<AntigravitySnapshot> {
416 let session = session?;
417 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
418 let mut snap = parse_quota_summary("a, session.plan)?;
419 snap.account = session.account;
420 Ok(snap)
421}
422
423fn saved_session(
433 credential: SavedCredential<'_>,
434 reason: RemoteFallbackReason,
435) -> Result<StoredToken> {
436 let raw = match credential {
437 SavedCredential::Keyring => credential::read()?,
438 SavedCredential::Blob(blob) => Some(blob.to_string()),
439 SavedCredential::File(path) => std::fs::read_to_string(path).ok(),
440 SavedCredential::Absent => None,
441 };
442 let Some(raw) = raw else {
443 return Err(AppError::Credentials(format!(
444 "{} {NO_SAVED_SESSION}",
445 reason.message()
446 )));
447 };
448 credential::parse_keyring_blob(&raw)
449}
450
451fn remote_account(fingerprint: &str) -> String {
455 format!("acct:{fingerprint}")
456}
457
458fn session_expired() -> AppError {
459 AppError::Credentials(SESSION_EXPIRED.into())
460}
461
462fn refresh_unconfigured() -> AppError {
466 AppError::Credentials(
467 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
468 refresh it; open Antigravity (or run `agy`) to sign in again, or set [antigravity] oauth_client_id \
469 and oauth_client_secret in config.toml"
470 .into(),
471 )
472}
473
474fn is_auth_rejection(e: &AppError) -> bool {
475 matches!(
476 e,
477 AppError::Http {
478 status: 401 | 403,
479 ..
480 }
481 )
482}
483
484struct AccessToken {
488 value: String,
489 just_refreshed: bool,
490}
491
492async fn resolve_access_token(
497 client: &reqwest::Client,
498 oauth: Option<&cloud::OauthClient>,
499 endpoints: &cloud::Endpoints,
500 oauth_path: &std::path::Path,
501 token: &StoredToken,
502 now: DateTime<Utc>,
503) -> Result<AccessToken> {
504 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
505 && token
506 .expires_at
507 .is_none_or(|keyring| persisted.expires_at > keyring)
508 && !cloud::needs_refresh(Some(persisted.expires_at), now)
509 {
510 return Ok(AccessToken {
511 value: persisted.access_token,
512 just_refreshed: false,
513 });
514 }
515 if !cloud::needs_refresh(token.expires_at, now) {
516 return Ok(AccessToken {
517 value: token.access_token.clone(),
518 just_refreshed: false,
519 });
520 }
521 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
522}
523
524async fn refresh_and_persist(
528 client: &reqwest::Client,
529 oauth: Option<&cloud::OauthClient>,
530 endpoints: &cloud::Endpoints,
531 oauth_path: &std::path::Path,
532 token: &StoredToken,
533) -> Result<AccessToken> {
534 let Some(refresh_token) = token.refresh_token.as_deref() else {
535 return Err(session_expired());
536 };
537 let Some(oauth) = oauth else {
538 return Err(refresh_unconfigured());
539 };
540 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
541 cloud::write_persisted(
542 oauth_path,
543 &cloud::PersistedOAuth {
544 fingerprint: token.fingerprint.clone(),
545 access_token: refreshed.access_token.clone(),
546 expires_at: refreshed.expires_at,
547 },
548 )?;
549 Ok(AccessToken {
550 value: refreshed.access_token,
551 just_refreshed: true,
552 })
553}
554
555struct Renewal<'a> {
558 agy: &'a super::agy::AgyCommand,
559 credential: SavedCredential<'a>,
560}
561
562impl Renewal<'_> {
563 async fn renewed(&self, cache: &Cache, now: DateTime<Utc>) -> Option<StoredToken> {
567 if !super::agy::renew_session(self.agy, cache, now).await {
568 return None;
569 }
570 let raw = match self.credential {
571 SavedCredential::Keyring => credential::read().ok().flatten()?,
572 SavedCredential::Blob(blob) => blob.to_string(),
573 SavedCredential::File(path) => std::fs::read_to_string(path).ok()?,
574 SavedCredential::Absent => return None,
575 };
576 credential::parse_keyring_blob(&raw).ok()
577 }
578}
579
580async fn fetch_remote(
587 client: &reqwest::Client,
588 cache: &Cache,
589 oauth: Option<&cloud::OauthClient>,
590 endpoints: &cloud::Endpoints,
591 token: Result<StoredToken>,
592 renewal: Renewal<'_>,
593 now: DateTime<Utc>,
594) -> Result<AntigravitySnapshot> {
595 let mut token = token?;
596 let oauth_path = cloud::oauth_cache_path(cache);
597 let mut access =
598 match resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await {
599 Ok(access) => access,
600 Err(e) if oauth.is_none() => {
603 let Some(renewed) = renewal.renewed(cache, now).await else {
604 return Err(e);
605 };
606 token = renewed;
607 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
608 }
609 Err(e) => return Err(e),
610 };
611
612 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
613 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
614 access = if oauth.is_none() {
615 let Some(renewed) = renewal.renewed(cache, now).await else {
616 return Err(refresh_unconfigured());
617 };
618 token = renewed;
619 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
620 } else {
621 refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?
622 };
623 cloud::fetch_quota(client, endpoints, &access.value)
624 .await
625 .map_err(|e| {
626 if is_auth_rejection(&e) {
627 session_expired()
628 } else {
629 e
630 }
631 })?
632 }
633 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
634 other => other?,
635 };
636
637 let plan = cloud::fetch_plan(client, endpoints, &access.value)
638 .await
639 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
640 let mut snap = parse_quota_summary("a, plan)?;
641 snap.account = remote_account(&token.fingerprint);
642 snap.source = AntigravitySource::Remote;
643 Ok(snap)
644}
645
646fn account_key(user_status: &serde_json::Value) -> String {
651 let email = user_status["userStatus"]["email"]
652 .as_str()
653 .filter(|s| !s.is_empty());
654 match email {
655 Some(e) => {
656 use std::hash::{Hash, Hasher};
657 let mut h = std::collections::hash_map::DefaultHasher::new();
658 e.hash(&mut h);
659 format!("acct:{:016x}", h.finish())
660 }
661 None => "acct:unknown".to_string(),
662 }
663}
664
665async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
669 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
670 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
673 .await
674 .ok()?;
675 let html = String::from_utf8_lossy(&bytes);
676 html.split("csrfToken\":\"")
677 .nth(1)
678 .and_then(|s| s.split('"').next())
679 .filter(|t| !t.is_empty())
680 .map(|t| t.to_string())
681}
682
683async fn post_rpc(
684 client: &reqwest::Client,
685 base: &str,
686 csrf: Option<&str>,
687 rpc: &str,
688) -> Result<serde_json::Value> {
689 let mut req = client
690 .post(format!("{base}/{rpc}"))
691 .header("Content-Type", "application/json")
692 .body("{}")
693 .timeout(HTTP_TIMEOUT);
694 if let Some(token) = csrf {
695 req = req.header("x-codeium-csrf-token", token);
696 }
697 let resp = req.send().await?;
698
699 let status = resp.status();
700 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
704 if !status.is_success() {
705 let body = String::from_utf8_lossy(&bytes).into_owned();
706 return Err(AppError::Http {
707 status: status.as_u16(),
708 body,
709 });
710 }
711
712 Ok(serde_json::from_slice(&bytes)?)
713}
714
715pub fn plan_from_status(v: &serde_json::Value) -> String {
716 v["userStatus"]["userTier"]["name"]
717 .as_str()
718 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
719 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
720 .filter(|s| !s.is_empty())
721 .unwrap_or(DEFAULT_PLAN)
722 .to_string()
723}
724
725fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
736 let id = if id.is_empty() { "<unnamed>" } else { id };
737 let group = if group_name.is_empty() {
738 String::new()
739 } else {
740 format!(" in {group_name:?}")
741 };
742 match window {
743 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
744 _ => format!("{id}{group}"),
745 }
746}
747
748fn no_usable_bucket(seen: &[String]) -> AppError {
754 if seen.is_empty() {
755 return AppError::Other(
756 "antigravity: quota summary has no buckets at all — the running product may \
757 not have a quota for this account yet"
758 .into(),
759 );
760 }
761 AppError::Other(format!(
762 "antigravity: quota summary has no bucket in a window we recognise (5h or \
763 weekly, Gemini or Claude/GPT); it offered: {}",
764 seen.join(", ")
765 ))
766}
767
768pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
769 let groups = v["response"]["groups"]
770 .as_array()
771 .or_else(|| v["groups"].as_array())
772 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
773
774 let mut gemini_5h = None;
775 let mut gemini_weekly = None;
776 let mut tp_5h = None;
777 let mut tp_weekly = None;
778 let mut seen: Vec<String> = Vec::new();
782
783 for group in groups {
784 let group_name = group["displayName"].as_str().unwrap_or_default();
785 let Some(buckets) = group["buckets"].as_array() else {
786 continue;
787 };
788 for bucket in buckets {
789 let id = bucket["bucketId"].as_str().unwrap_or_default();
790 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
791 let window = bucket["window"].as_str().unwrap_or_default();
792 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
793 true
794 } else if id.ends_with("5h") || window == "5h" {
795 false
796 } else {
797 continue;
800 };
801 let is_gemini = if id.starts_with("gemini") {
802 true
803 } else if id.starts_with("3p") {
804 false
805 } else if group_name.contains("Gemini") {
806 true
807 } else if group_name.contains("Claude") || group_name.contains("GPT") {
808 false
809 } else {
810 continue;
813 };
814
815 let (slot, slot_name) = match (is_gemini, is_weekly) {
816 (true, false) => (&mut gemini_5h, "Gemini 5h"),
817 (true, true) => (&mut gemini_weekly, "Gemini weekly"),
818 (false, false) => (&mut tp_5h, "Claude/GPT 5h"),
819 (false, true) => (&mut tp_weekly, "Claude/GPT weekly"),
820 };
821 let parsed = usage_window(bucket, is_weekly)?;
822 if slot.replace(parsed).is_some() {
823 return Err(AppError::Schema(format!(
824 "antigravity: duplicate {slot_name} bucket"
825 )));
826 }
827 }
828 }
829
830 if gemini_5h.is_none() && gemini_weekly.is_none() && tp_5h.is_none() && tp_weekly.is_none() {
834 return Err(no_usable_bucket(&seen));
835 }
836
837 Ok(AntigravitySnapshot {
838 plan,
839 account: String::new(),
842 source: AntigravitySource::Local,
843 session: gemini_5h,
844 weekly: gemini_weekly,
845 third_party_session: tp_5h,
846 third_party_weekly: tp_weekly,
847 })
848}
849
850fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
854 let remaining = bucket["remainingFraction"]
855 .as_f64()
856 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
857 .ok_or_else(|| {
858 AppError::Schema(format!(
859 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
860 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
861 ))
862 })?;
863 Ok(UsageWindow {
864 utilization_pct: pct_used(remaining),
865 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
866 window_duration: if is_weekly {
867 chrono::Duration::days(7)
868 } else {
869 chrono::Duration::hours(5)
870 },
871 })
872}
873
874fn pct_used(remaining_fraction: f64) -> i32 {
877 let used = (1.0 - remaining_fraction) * 100.0;
878 used.round() as i32
879}
880
881fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
882 match value {
883 serde_json::Value::Null => Ok(None),
884 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
885 .map(|dt| Some(dt.with_timezone(&Utc)))
886 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
887 _ => Err(AppError::Schema(format!(
888 "antigravity: {field} must be a timestamp or null"
889 ))),
890 }
891}
892
893fn candidate_bases() -> Vec<Candidate> {
899 candidate_bases_with(
900 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
901 discover_ls_ports(),
902 )
903}
904
905fn candidate_bases_with(
909 override_addr: Option<&str>,
910 discovered: Vec<(u32, u16)>,
911) -> Vec<Candidate> {
912 let mut candidates = Vec::new();
913 if let Some(base) = override_addr.and_then(normalize_base) {
914 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
918 candidates.push(Candidate { base, pid });
919 }
920
921 for (pid, port) in discovered {
927 let base = format!("http://127.0.0.1:{port}");
928 if !candidates.iter().any(|c| c.base == base) {
929 candidates.push(Candidate {
930 base,
931 pid: Some(pid),
932 });
933 }
934 }
935
936 candidates
937}
938
939fn base_port(base: &str) -> Option<u16> {
941 let (_, authority) = base.split_once("://")?;
942 authority.rsplit_once(':')?.1.parse().ok()
943}
944
945fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
948 let mut owners = discovered
949 .iter()
950 .filter(|(_, p)| *p == port)
951 .map(|(pid, _)| *pid);
952 let first = owners.next()?;
953 owners.all(|pid| pid == first).then_some(first)
954}
955
956fn normalize_base(addr: &str) -> Option<String> {
965 let trimmed = addr.trim();
966 let (scheme, authority) = match trimmed.split_once("://") {
967 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
968 _ => ("http", trimmed),
969 };
970 let authority = authority.trim_end_matches('/');
971 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
972}
973
974fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
981 let comm = comm.trim().to_lowercase();
982 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
983 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
984 return true;
985 }
986 exe.is_some_and(|p| {
987 let p = p.to_lowercase().replace('\\', "/");
988 let p = p.strip_suffix(".exe").unwrap_or(&p);
989 p.contains("antigravity") || p.ends_with("/agy")
990 })
991}
992
993#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
1025fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
1026 let groups: Vec<(u32, Vec<u16>)> = per_pid
1027 .into_iter()
1028 .map(|(pid, mut group)| {
1029 group.sort_unstable_by(|a, b| b.cmp(a));
1030 group.dedup();
1035 (pid, group)
1036 })
1037 .collect();
1038
1039 let mut ports: Vec<(u32, u16)> = Vec::new();
1040 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
1041 for (pid, port) in groups
1042 .iter()
1043 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
1044 {
1045 if !ports.iter().any(|(_, p)| *p == port) {
1048 ports.push((pid, port));
1049 }
1050 }
1051 }
1052 ports
1053}
1054
1055#[cfg(target_os = "linux")]
1063pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1064 use std::collections::{BTreeMap, HashMap};
1065
1066 let mut owners: HashMap<u64, u32> = HashMap::new();
1069 let Ok(entries) = std::fs::read_dir("/proc") else {
1070 return Vec::new();
1071 };
1072 for entry in entries.flatten() {
1073 let pid_dir = entry.path();
1074 let Some(pid) = pid_dir
1075 .file_name()
1076 .and_then(|name| name.to_str())
1077 .and_then(|name| name.parse::<u32>().ok())
1078 else {
1079 continue;
1080 };
1081 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1082 continue;
1083 };
1084 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1085 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1086 continue;
1087 }
1088 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1089 continue;
1090 };
1091 for fd in fds.flatten() {
1092 let Ok(target) = std::fs::read_link(fd.path()) else {
1093 continue;
1094 };
1095 if let Some(ino) = target
1096 .to_str()
1097 .and_then(|s| s.strip_prefix("socket:["))
1098 .and_then(|s| s.strip_suffix(']'))
1099 .and_then(|s| s.parse::<u64>().ok())
1100 {
1101 owners.insert(ino, pid);
1102 }
1103 }
1104 }
1105
1106 if owners.is_empty() {
1107 return Vec::new();
1108 }
1109
1110 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1111 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1112 let Ok(contents) = std::fs::read_to_string(table) else {
1113 continue;
1114 };
1115 for line in contents.lines().skip(1) {
1116 if let Some((port, ino)) = parse_proc_net_line(line)
1117 && let Some(&pid) = owners.get(&ino)
1118 {
1119 per_pid.entry(pid).or_default().push(port);
1120 }
1121 }
1122 }
1123 probe_order(per_pid)
1124}
1125
1126#[cfg(target_os = "macos")]
1133pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1134 let Ok(output) = std::process::Command::new("lsof")
1135 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1136 .output()
1137 else {
1138 return Vec::new();
1139 };
1140 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1144}
1145
1146#[cfg(any(test, target_os = "macos"))]
1151fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1152 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1153 let mut pid = None;
1156 let mut owner = None;
1157 for line in output.lines() {
1158 let Some(rest) = line.get(1..) else { continue };
1159 match line.as_bytes().first() {
1160 Some(b'p') => {
1161 pid = rest.parse::<u32>().ok();
1162 owner = None;
1163 }
1164 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1165 Some(b'n') => {
1166 if let Some(pid) = owner
1167 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1168 {
1169 per_pid.entry(pid).or_default().push(port);
1170 }
1171 }
1172 _ => {}
1173 }
1174 }
1175 probe_order(per_pid)
1176}
1177
1178#[cfg(any(test, target_os = "windows"))]
1179#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1180struct WindowsTcpRow {
1181 local_addr: [u8; 4],
1182 local_port: u32,
1183 pid: u32,
1184}
1185
1186#[cfg(any(test, target_os = "windows"))]
1187fn decode_windows_process_name(raw: &[u16]) -> String {
1188 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1189 String::from_utf16_lossy(&raw[..end])
1190}
1191
1192#[cfg(any(test, target_os = "windows"))]
1193fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1194 processes
1195 .iter()
1196 .filter(|(_, name)| is_antigravity_process(name, None))
1197 .map(|(pid, _)| *pid)
1198 .collect()
1199}
1200
1201#[cfg(any(test, target_os = "windows"))]
1204fn matching_windows_ports(
1205 pids: &std::collections::HashSet<u32>,
1206 rows: &[WindowsTcpRow],
1207) -> Vec<(u32, u16)> {
1208 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1209 for row in rows {
1210 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1211 continue;
1212 }
1213 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1214 if port != 0 {
1215 per_pid.entry(row.pid).or_default().push(port);
1216 }
1217 }
1218 probe_order(per_pid)
1219}
1220
1221#[cfg(any(test, target_os = "windows"))]
1222fn checked_windows_row_count(
1223 buffer_len: usize,
1224 rows_offset: usize,
1225 row_size: usize,
1226 declared: usize,
1227) -> Option<usize> {
1228 let rows_len = row_size.checked_mul(declared)?;
1229 let end = rows_offset.checked_add(rows_len)?;
1230 (row_size != 0 && end <= buffer_len).then_some(declared)
1231}
1232
1233#[cfg(target_os = "windows")]
1234struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1235
1236#[cfg(target_os = "windows")]
1237impl Drop for WindowsHandle {
1238 fn drop(&mut self) {
1239 unsafe {
1240 windows_sys::Win32::Foundation::CloseHandle(self.0);
1241 }
1242 }
1243}
1244
1245#[cfg(target_os = "windows")]
1246fn windows_processes() -> Vec<(u32, String)> {
1247 use std::mem::size_of;
1248 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1249 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1250 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1251 TH32CS_SNAPPROCESS,
1252 };
1253
1254 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1255 if handle == INVALID_HANDLE_VALUE {
1256 return Vec::new();
1257 }
1258 let snapshot = WindowsHandle(handle);
1259 let mut entry = PROCESSENTRY32W::default();
1260 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1261 return Vec::new();
1262 };
1263 entry.dwSize = entry_size;
1264 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1265 return Vec::new();
1266 }
1267
1268 let mut processes = Vec::new();
1269 loop {
1270 processes.push((
1271 entry.th32ProcessID,
1272 decode_windows_process_name(&entry.szExeFile),
1273 ));
1274 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1275 break;
1276 }
1277 }
1278 processes
1279}
1280
1281#[cfg(target_os = "windows")]
1282fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1283 use std::mem::{offset_of, size_of, size_of_val};
1284 use windows_sys::Win32::NetworkManagement::IpHelper::{
1285 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1286 };
1287
1288 let available = used_bytes.min(size_of_val(buffer));
1289 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1290 if available < size_of::<u32>() || available < rows_offset {
1291 return Vec::new();
1292 }
1293 let base = buffer.as_ptr().cast::<u8>();
1294 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1295 if checked_windows_row_count(
1296 available,
1297 rows_offset,
1298 size_of::<MIB_TCPROW_OWNER_PID>(),
1299 declared,
1300 )
1301 .is_none()
1302 {
1303 return Vec::new();
1304 }
1305
1306 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1307 (0..declared)
1308 .map(|index| unsafe { rows.add(index).read_unaligned() })
1309 .map(|row| WindowsTcpRow {
1310 local_addr: row.dwLocalAddr.to_ne_bytes(),
1311 local_port: row.dwLocalPort,
1312 pid: row.dwOwningPid,
1313 })
1314 .collect()
1315}
1316
1317#[cfg(target_os = "windows")]
1318fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1319 use std::mem::size_of;
1320 use std::ptr::null_mut;
1321 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1322 use windows_sys::Win32::NetworkManagement::IpHelper::{
1323 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1324 };
1325 use windows_sys::Win32::Networking::WinSock::AF_INET;
1326
1327 let mut size = 0u32;
1328 let status = unsafe {
1329 GetExtendedTcpTable(
1330 null_mut(),
1331 &mut size,
1332 0,
1333 u32::from(AF_INET),
1334 TCP_TABLE_OWNER_PID_LISTENER,
1335 0,
1336 )
1337 };
1338 if status != ERROR_INSUFFICIENT_BUFFER {
1339 return Vec::new();
1340 }
1341
1342 for _ in 0..3 {
1343 let Some(words) = (size as usize)
1344 .checked_add(size_of::<u32>() - 1)
1345 .map(|bytes| bytes / size_of::<u32>())
1346 else {
1347 return Vec::new();
1348 };
1349 if words == 0 {
1350 return Vec::new();
1351 }
1352 let mut buffer = Vec::<u32>::new();
1353 if buffer.try_reserve_exact(words).is_err() {
1354 return Vec::new();
1355 }
1356 buffer.resize(words, 0);
1357 let mut used = size;
1358 let status = unsafe {
1359 GetExtendedTcpTable(
1360 buffer.as_mut_ptr().cast(),
1361 &mut used,
1362 0,
1363 u32::from(AF_INET),
1364 TCP_TABLE_OWNER_PID_LISTENER,
1365 0,
1366 )
1367 };
1368 if status == ERROR_INSUFFICIENT_BUFFER {
1369 size = used;
1370 continue;
1371 }
1372 if status != 0 {
1373 return Vec::new();
1374 }
1375 return parse_windows_tcp_rows(&buffer, used as usize);
1376 }
1377 Vec::new()
1378}
1379
1380#[cfg(target_os = "windows")]
1381pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1382 let pids = matching_windows_process_ids(&windows_processes());
1383 if pids.is_empty() {
1384 return Vec::new();
1385 }
1386 matching_windows_ports(&pids, &windows_tcp_rows())
1387}
1388
1389#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1390pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1391 Vec::new()
1392}
1393
1394#[cfg(target_os = "linux")]
1397fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1398 let cols: Vec<&str> = line.split_whitespace().collect();
1399 if cols.len() < 10 {
1400 return None;
1401 }
1402 if cols[3] != "0A" {
1404 return None;
1405 }
1406 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1407 let inode = cols[9].parse::<u64>().ok()?;
1408 Some((port, inode))
1409}
1410
1411fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1416 crate::outcome::fallback(cache, None, original, |bytes| {
1417 parse_cache_at(bytes, None, now)
1418 })
1419}
1420
1421fn fallback_with_error(
1426 cache: &Cache,
1427 last_error: Option<(u16, String)>,
1428 reason: AppError,
1429 now: DateTime<Utc>,
1430) -> Result<FetchOutcome> {
1431 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1432 parse_cache_at(bytes, None, now)
1433 })
1434}
1435
1436fn reuse_cache(
1437 bytes: Vec<u8>,
1438 cache: &Cache,
1439 stale: bool,
1440 account: Option<&str>,
1441 now: DateTime<Utc>,
1442) -> Result<FetchOutcome> {
1443 let snap = parse_cache_at(&bytes, account, now)?;
1444 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1445}
1446
1447pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1453 parse_cache_at(bytes, account, Utc::now())
1454}
1455
1456fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<&'static str> {
1466 [
1467 ("Gemini 5h", snap.session.as_ref()),
1468 ("Gemini weekly", snap.weekly.as_ref()),
1469 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1470 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1471 ]
1472 .into_iter()
1473 .find(|(_, w)| w.and_then(|w| w.resets_at).is_some_and(|r| r <= now))
1474 .map(|(name, _)| name)
1475}
1476
1477pub fn parse_cache_at(
1478 bytes: &[u8],
1479 account: Option<&str>,
1480 now: DateTime<Utc>,
1481) -> Result<AntigravitySnapshot> {
1482 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1483
1484 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1485 if let Some(expected) = account
1486 && cached_account != Some(expected)
1487 {
1488 return Err(AppError::Schema(
1489 "antigravity cache belongs to a different account; refetching".into(),
1490 ));
1491 }
1492
1493 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1503 match v.get(pct_key) {
1504 None => Err(AppError::Schema(format!(
1505 "antigravity: cached payload is missing {pct_key}"
1506 ))),
1507 Some(serde_json::Value::Null) => Ok(None),
1508 Some(value) => value
1509 .as_i64()
1510 .filter(|pct| (0..=100).contains(pct))
1511 .map(|pct| Some(pct as i32))
1512 .ok_or_else(|| {
1513 AppError::Schema(format!(
1514 "antigravity: cached {pct_key} must be an integer in 0..=100"
1515 ))
1516 }),
1517 }
1518 };
1519
1520 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1521 let Some(pct) = cached_pct(pct_key)? else {
1522 return Ok(None);
1523 };
1524 Ok::<_, AppError>(Some(UsageWindow {
1525 utilization_pct: pct,
1526 resets_at: parse_reset(&v[reset_key], reset_key)?,
1527 window_duration: if weekly {
1528 chrono::Duration::days(7)
1529 } else {
1530 chrono::Duration::hours(5)
1531 },
1532 }))
1533 };
1534
1535 let snap = AntigravitySnapshot {
1536 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1537 account: cached_account.unwrap_or_default().to_string(),
1538 source: v["source"]
1540 .as_str()
1541 .and_then(AntigravitySource::parse)
1542 .unwrap_or_default(),
1543 session: optional("session_pct", "session_reset", false)?,
1544 weekly: optional("weekly_pct", "weekly_reset", true)?,
1545 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1546 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1547 };
1548
1549 if snap.session.is_none()
1552 && snap.weekly.is_none()
1553 && snap.third_party_session.is_none()
1554 && snap.third_party_weekly.is_none()
1555 {
1556 return Err(AppError::Schema(
1557 "antigravity cache holds no usable window; refetching".into(),
1558 ));
1559 }
1560
1561 if let Some(window) = expired_window(&snap, now) {
1562 return Err(AppError::Schema(format!(
1563 "antigravity cache is past its {window} reset; refetching"
1564 )));
1565 }
1566 Ok(snap)
1567}
1568
1569pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1570 serde_json::json!({
1571 "plan": snap.plan,
1572 "account": snap.account,
1573 "source": snap.source.as_str(),
1574 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
1575 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1576 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
1577 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1578 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
1579 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1580 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
1581 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1582 })
1583}
1584
1585#[cfg(test)]
1586mod tests {
1587 use super::*;
1588
1589 const QUOTA_JSON: &str = r#"{
1593 "response": {
1594 "groups": [
1595 {
1596 "displayName": "Gemini Models",
1597 "buckets": [
1598 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
1599 "window": "weekly", "remainingFraction": 0.9191212,
1600 "resetTime": "2026-07-28T17:39:58Z"},
1601 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
1602 "window": "5h", "remainingFraction": 0.5672253,
1603 "resetTime": "2026-07-22T17:47:00Z"}
1604 ]
1605 },
1606 {
1607 "displayName": "Claude and GPT models",
1608 "buckets": [
1609 {"bucketId": "3p-weekly", "window": "weekly",
1610 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
1611 {"bucketId": "3p-5h", "window": "5h",
1612 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
1613 ]
1614 }
1615 ]
1616 }
1617 }"#;
1618
1619 fn now() -> DateTime<Utc> {
1622 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
1623 .unwrap()
1624 .with_timezone(&Utc)
1625 }
1626
1627 fn parsed() -> AntigravitySnapshot {
1628 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1629 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
1630 }
1631
1632 #[test]
1633 fn quota_summary_maps_four_distinct_windows() {
1634 let snap = parsed();
1635 assert_eq!(snap.plan, "Google AI Pro");
1636 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
1638 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
1639 assert_eq!(
1640 snap.third_party_session.as_ref().unwrap().utilization_pct,
1641 75
1642 );
1643 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
1644 }
1645
1646 #[test]
1647 fn each_window_keeps_its_own_reset_time() {
1648 let snap = parsed();
1649 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
1650 assert_eq!(
1651 snap.session.as_ref().unwrap().resets_at,
1652 at("2026-07-22T17:47:00Z")
1653 );
1654 assert_eq!(
1655 snap.weekly.as_ref().unwrap().resets_at,
1656 at("2026-07-28T17:39:58Z")
1657 );
1658 assert_eq!(
1659 snap.third_party_weekly.as_ref().unwrap().resets_at,
1660 at("2026-07-29T12:47:00Z")
1661 );
1662 assert_ne!(
1664 snap.session.as_ref().unwrap().resets_at,
1665 snap.weekly.as_ref().unwrap().resets_at
1666 );
1667 }
1668
1669 #[test]
1670 fn window_durations_match_their_bucket() {
1671 let snap = parsed();
1672 assert_eq!(
1673 snap.session.as_ref().unwrap().window_duration,
1674 chrono::Duration::hours(5)
1675 );
1676 assert_eq!(
1677 snap.weekly.as_ref().unwrap().window_duration,
1678 chrono::Duration::days(7)
1679 );
1680 assert_eq!(
1681 snap.third_party_weekly.as_ref().unwrap().window_duration,
1682 chrono::Duration::days(7)
1683 );
1684 }
1685
1686 #[test]
1687 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
1688 let v: serde_json::Value = serde_json::from_str(
1689 r#"{"response":{"groups":[
1690 {"displayName":"Gemini Models","buckets":[
1691 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
1692 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
1693 {"displayName":"Claude and GPT models","buckets":[
1694 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
1695 ]}}"#,
1696 )
1697 .unwrap();
1698 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
1699 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
1700 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
1701 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
1702 assert!(snap.third_party_weekly.is_none());
1703 }
1704
1705 #[test]
1706 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
1707 let duplicate: serde_json::Value = serde_json::from_str(
1708 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
1709 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1710 {"bucketId":"gemini-5h-copy","window":"5h","remainingFraction":0.1},
1711 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
1712 ]}]}}"#,
1713 )
1714 .unwrap();
1715 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
1716 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
1717
1718 let unrelated: serde_json::Value = serde_json::from_str(
1721 r#"{"response":{"groups":[
1722 {"displayName":"Gemini Models","buckets":[
1723 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1724 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
1725 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
1726 ]},
1727 {"displayName":"Future Models","buckets":[
1728 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
1729 ]}
1730 ]}}"#,
1731 )
1732 .unwrap();
1733 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
1734 assert!(snap.third_party_session.is_none());
1735 assert!(snap.third_party_weekly.is_none());
1736 }
1737
1738 #[test]
1741 fn a_bucket_without_a_usable_fraction_is_rejected() {
1742 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
1743 let v: serde_json::Value = serde_json::from_str(&format!(
1744 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
1745 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
1746 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
1747 ))
1748 .unwrap();
1749 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1750 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
1751 }
1752 }
1753
1754 #[test]
1755 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
1756 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
1757 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1758 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
1759 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1760 assert!(err.to_string().contains("resetTime"), "{err}");
1761 }
1762 }
1763
1764 #[test]
1768 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
1769 let mut snap = parsed();
1770 snap.session = None;
1771 snap.third_party_session = None;
1772
1773 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1774 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
1775
1776 assert!(back.session.is_none());
1777 assert_eq!(
1778 back.weekly.as_ref().unwrap().utilization_pct,
1779 snap.weekly.as_ref().unwrap().utilization_pct
1780 );
1781 }
1782
1783 #[test]
1788 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
1789 let summary = serde_json::json!({
1790 "groups": [
1791 {
1792 "displayName": "Gemini Models",
1793 "buckets": [{
1794 "bucketId": "gemini-weekly", "window": "weekly",
1795 "remainingFraction": 0.42,
1796 }],
1797 },
1798 {
1799 "displayName": "Claude and GPT models",
1800 "buckets": [{
1801 "bucketId": "3p-weekly", "window": "weekly",
1802 "remainingFraction": 0.9,
1803 }],
1804 },
1805 ],
1806 });
1807
1808 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
1809
1810 assert!(snap.session.is_none(), "no 5h bucket arrived");
1811 assert!(snap.third_party_session.is_none());
1812 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
1813 assert_eq!(
1814 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
1815 10
1816 );
1817 }
1818
1819 #[test]
1822 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
1823 let summary = serde_json::json!({
1824 "groups": [{
1825 "displayName": "Gemini Models",
1826 "buckets": [{
1827 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
1828 }],
1829 }],
1830 });
1831
1832 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
1833
1834 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
1835 assert!(snap.weekly.is_none());
1836 }
1837
1838 #[test]
1841 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
1842 let summary = serde_json::json!({
1843 "groups": [{
1844 "displayName": "Gemini Models",
1845 "buckets": [{
1846 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
1847 }],
1848 }],
1849 });
1850
1851 let rendered = parse_quota_summary(&summary, "Pro".into())
1852 .expect_err("an unrecognised cadence alone is not a snapshot")
1853 .to_string();
1854
1855 assert!(
1856 rendered.contains("no bucket in a window we recognise"),
1857 "{rendered}"
1858 );
1859 assert!(rendered.contains("gemini-daily"), "{rendered}");
1860 assert!(rendered.contains("window daily"), "{rendered}");
1861 }
1862
1863 #[test]
1866 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
1867 let summary = serde_json::json!({
1868 "groups": [{"displayName": "Gemini", "buckets": []}],
1869 });
1870
1871 let rendered = parse_quota_summary(&summary, "Pro".into())
1872 .expect_err("no buckets is an error")
1873 .to_string();
1874
1875 assert!(rendered.contains("no buckets at all"), "{rendered}");
1876 assert!(!rendered.contains("it offered:"), "{rendered}");
1877 }
1878
1879 #[test]
1882 fn buckets_without_an_id_are_still_named_in_the_error() {
1883 let summary = serde_json::json!({
1884 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
1885 });
1886
1887 let rendered = parse_quota_summary(&summary, "Pro".into())
1888 .expect_err("an unusable summary is an error")
1889 .to_string();
1890
1891 assert!(rendered.contains("<unnamed>"), "{rendered}");
1892 }
1893
1894 #[test]
1895 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
1896 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
1897 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
1898 }
1899
1900 #[test]
1901 fn cache_round_trip_preserves_every_window() {
1902 let snap = parsed();
1903 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1904 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1905 }
1906
1907 #[test]
1911 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
1912 let full = snap_to_json(&parsed());
1913 for missing in ["session_pct", "weekly_pct"] {
1914 let mut v = full.clone();
1915 v.as_object_mut().unwrap().remove(missing);
1916 let bytes = serde_json::to_vec(&v).unwrap();
1917 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1918 assert!(err.to_string().contains(missing), "{missing}: {err}");
1919 }
1920 assert!(parse_cache_at(b"{}", None, now()).is_err());
1922 }
1923
1924 #[test]
1925 fn cached_percentages_are_range_checked_before_narrowing() {
1926 let full = snap_to_json(&parsed());
1927 for (key, bad) in [
1928 ("session_pct", serde_json::json!(-1)),
1929 ("weekly_pct", serde_json::json!(101)),
1930 ("session_pct", serde_json::json!(i64::MAX)),
1931 ("tp_session_pct", serde_json::json!("75")),
1932 ] {
1933 let mut v = full.clone();
1934 v[key] = bad;
1935 let bytes = serde_json::to_vec(&v).unwrap();
1936 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1937 assert!(err.to_string().contains(key), "{key}: {err}");
1938 }
1939 }
1940
1941 #[test]
1942 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
1943 let mut v = snap_to_json(&parsed());
1944 v["session_reset"] = serde_json::json!("not-a-time");
1945 let bytes = serde_json::to_vec(&v).unwrap();
1946 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1947 assert!(err.to_string().contains("session_reset"), "{err}");
1948 }
1949
1950 #[test]
1955 fn a_cache_past_its_reset_is_refused() {
1956 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
1957 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1958
1959 assert!(parse_cache_at(&bytes, None, now()).is_ok());
1961 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
1963
1964 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
1966 assert!(err.to_string().contains("5h"), "{err}");
1967
1968 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
1970 }
1971
1972 #[test]
1975 fn expiry_names_the_window_that_rolled_over() {
1976 let mut snap = parsed();
1977 snap.session.as_mut().unwrap().resets_at = None;
1979 snap.third_party_session = None;
1980 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1981 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1982
1983 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
1985
1986 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
1988 assert!(err.to_string().contains("Gemini weekly"), "{err}");
1989 }
1990
1991 #[test]
1993 fn a_window_without_a_reset_never_expires() {
1994 let mut snap = parsed();
1995 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
1996 w.resets_at = None;
1997 }
1998 snap.third_party_session = None;
1999 snap.third_party_weekly = None;
2000 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2001 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
2002 .unwrap()
2003 .with_timezone(&Utc);
2004 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
2005 }
2006
2007 #[test]
2009 fn a_cache_from_another_account_is_rejected() {
2010 let mut snap = parsed();
2011 snap.account = "acct:aaaa".into();
2012 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2013
2014 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
2015 assert_eq!(
2016 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
2017 snap
2018 );
2019
2020 let mut legacy = snap_to_json(&snap);
2022 legacy.as_object_mut().unwrap().remove("account");
2023 let legacy = serde_json::to_vec(&legacy).unwrap();
2024 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
2025 }
2026
2027 #[test]
2030 fn an_unverifiable_cache_is_served_rather_than_discarded() {
2031 let mut snap = parsed();
2032 snap.account = "acct:aaaa".into();
2033 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2034 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2035 }
2036
2037 #[test]
2038 fn account_key_fingerprints_rather_than_storing_the_address() {
2039 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
2040 let a = with("someone@example.com");
2041 assert!(!a.contains("someone"), "{a}");
2042 assert!(!a.contains('@'), "{a}");
2043 assert_eq!(a, with("someone@example.com"), "must be stable");
2044 assert_ne!(a, with("other@example.com"));
2045 let unknown = account_key(&serde_json::json!({}));
2047 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
2048 assert_ne!(unknown, a);
2049 }
2050
2051 #[test]
2054 fn absent_third_party_windows_are_not_treated_as_corruption() {
2055 let mut snap = parsed();
2056 snap.third_party_session = None;
2057 snap.third_party_weekly = None;
2058 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2059 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2060 }
2061
2062 #[test]
2063 fn cache_round_trip_preserves_absent_third_party_windows() {
2064 let mut snap = parsed();
2065 snap.third_party_session = None;
2066 snap.third_party_weekly = None;
2067 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2068 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2069 }
2070
2071 #[test]
2072 fn pct_used_inverts_valid_fractions() {
2073 assert_eq!(pct_used(1.0), 0);
2074 assert_eq!(pct_used(0.0), 100);
2075 assert_eq!(pct_used(0.5), 50);
2076 }
2077
2078 #[test]
2079 fn plan_falls_back_through_the_status_payload() {
2080 let tier: serde_json::Value =
2081 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2082 .unwrap();
2083 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2084
2085 let plan_only: serde_json::Value = serde_json::from_str(
2086 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2087 )
2088 .unwrap();
2089 assert_eq!(plan_from_status(&plan_only), "Pro");
2090
2091 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2092 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2093 }
2094
2095 #[cfg(target_os = "linux")]
2096 #[test]
2097 fn proc_net_parser_keeps_only_listening_rows() {
2098 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2099 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2100 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2101
2102 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2103 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2104 assert_eq!(parse_proc_net_line(established), None);
2105
2106 assert_eq!(parse_proc_net_line("garbage"), None);
2107 }
2108
2109 #[test]
2110 fn explicit_address_comes_first_and_gets_a_scheme() {
2111 let candidate = |base: &str, pid| Candidate {
2112 base: base.to_string(),
2113 pid,
2114 };
2115 assert_eq!(
2116 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)]),
2117 vec![
2118 candidate("http://127.0.0.1:1234", None),
2119 candidate("http://127.0.0.1:5678", Some(10)),
2120 ]
2121 );
2122 assert_eq!(
2124 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)]),
2125 vec![
2126 candidate("http://127.0.0.1:1234", None),
2127 candidate("http://127.0.0.1:5678", Some(10)),
2128 ]
2129 );
2130 assert_eq!(
2133 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)]),
2134 vec![candidate("http://127.0.0.1:5678", Some(10))]
2135 );
2136 assert_eq!(
2138 candidate_bases_with(None, vec![(10, 5678), (10, 5678)]),
2139 vec![candidate("http://127.0.0.1:5678", Some(10))]
2140 );
2141 assert_eq!(
2143 candidate_bases_with(Some("https://host:9"), vec![]),
2144 vec![candidate("https://host:9", None)]
2145 );
2146 }
2147
2148 fn http(status: u16) -> AppError {
2149 AppError::Http {
2150 status,
2151 body: String::new(),
2152 }
2153 }
2154
2155 fn missing_csrf() -> AppError {
2156 AppError::Http {
2157 status: 401,
2158 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2159 }
2160 }
2161
2162 #[test]
2163 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2164 assert!(is_missing_csrf(&missing_csrf()));
2165 assert!(!is_missing_csrf(&AppError::Http {
2166 status: 403,
2167 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2168 }));
2169 assert!(!is_missing_csrf(&AppError::Http {
2170 status: 401,
2171 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2172 }));
2173 assert!(!is_missing_csrf(&AppError::Http {
2174 status: 401,
2175 body: "prefix: missing CSRF token".into(),
2176 }));
2177 }
2178
2179 #[test]
2180 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2181 for errors in [
2182 vec![missing_csrf(), http(403)],
2183 vec![http(401), missing_csrf()],
2184 ] {
2185 let err = select_probe_error(errors);
2186 assert!(is_actionable(&err), "{err}");
2187 assert!(!is_missing_csrf(&err), "{err}");
2188 }
2189 }
2190
2191 #[test]
2192 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2193 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2194 assert!(is_missing_csrf(&err), "{err}");
2195 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2196 }
2197
2198 #[test]
2202 fn an_auth_failure_outranks_later_transport_noise() {
2203 let err = select_probe_error(vec![
2204 http(401),
2205 AppError::Transport("connection refused".into()),
2206 ]);
2207 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2208
2209 let err = select_probe_error(vec![
2210 AppError::Transport("connection refused".into()),
2211 http(403),
2212 ]);
2213 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2214 }
2215
2216 #[test]
2219 fn the_first_auth_failure_wins() {
2220 let err = select_probe_error(vec![http(401), http(403)]);
2221 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2222 }
2223
2224 #[test]
2228 fn without_an_auth_failure_the_last_error_stands() {
2229 let err = select_probe_error(vec![
2230 AppError::Transport("first".into()),
2231 http(500),
2232 AppError::Transport("last".into()),
2233 ]);
2234 assert!(
2235 matches!(&err, AppError::Transport(m) if m == "last"),
2236 "{err}"
2237 );
2238 assert!(err.is_transient());
2239 }
2240
2241 #[test]
2244 fn a_server_error_is_not_treated_as_actionable() {
2245 let err = select_probe_error(vec![http(500), http(401)]);
2246 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2247 }
2248
2249 #[test]
2250 fn no_candidates_at_all_yields_a_generic_error() {
2251 let err = select_probe_error(Vec::new());
2252 assert!(
2253 err.to_string().contains("no local server answered"),
2254 "{err}"
2255 );
2256 }
2257
2258 fn tls_echo() -> AppError {
2262 AppError::Http {
2263 status: 400,
2264 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2265 }
2266 }
2267
2268 #[test]
2273 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2274 let err = select_probe_error(vec![
2275 AppError::Http {
2276 status: 500,
2277 body: "GetUserStatus: internal".into(),
2278 },
2279 tls_echo(),
2280 ]);
2281 assert!(
2282 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2283 "{err}"
2284 );
2285 }
2286
2287 #[test]
2292 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2293 let err = select_probe_error(vec![
2294 AppError::Transport("connection refused".into()),
2295 tls_echo(),
2296 ]);
2297 assert!(
2298 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2299 "{err}"
2300 );
2301 assert!(
2302 err.is_transient(),
2303 "the echo must not make the run non-transient: {err}"
2304 );
2305 }
2306
2307 #[test]
2310 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2311 let err = select_probe_error(vec![tls_echo()]);
2312 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2313 }
2314
2315 #[test]
2318 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2319 let err = select_probe_error(vec![
2320 AppError::Http {
2321 status: 400,
2322 body: "unknown method GetUserStatus".into(),
2323 },
2324 tls_echo(),
2325 ]);
2326 assert!(
2327 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2328 "{err}"
2329 );
2330 }
2331
2332 #[test]
2334 fn an_auth_failure_still_outranks_a_tls_echo() {
2335 let err = select_probe_error(vec![tls_echo(), http(401)]);
2336 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2337 }
2338
2339 #[test]
2340 fn every_discovered_port_is_probed_in_order() {
2341 assert_eq!(
2342 candidate_bases_with(None, vec![(10, 33875), (10, 37435)]),
2343 vec![
2344 Candidate {
2345 base: "http://127.0.0.1:33875".to_string(),
2346 pid: Some(10),
2347 },
2348 Candidate {
2349 base: "http://127.0.0.1:37435".to_string(),
2350 pid: Some(10),
2351 },
2352 ]
2353 );
2354 }
2355
2356 #[test]
2361 fn empty_discovery_yields_no_candidates() {
2362 assert!(candidate_bases_with(None, vec![]).is_empty());
2363 assert!(candidate_bases_with(Some(""), vec![]).is_empty());
2364 }
2365
2366 #[test]
2367 fn every_antigravity_product_is_recognised() {
2368 assert!(is_antigravity_process(
2370 "language_server\n",
2371 Some("/opt/antigravity/resources/bin/language_server")
2372 ));
2373 assert!(is_antigravity_process(
2375 "agy\n",
2376 Some("/home/u/.local/bin/agy")
2377 ));
2378 assert!(is_antigravity_process(
2380 "node",
2381 Some("/opt/antigravity/bin/helper")
2382 ));
2383 assert!(is_antigravity_process("antigravity", None));
2384 assert!(is_antigravity_process("agy.exe", None));
2385 assert!(is_antigravity_process("Antigravity.exe", None));
2386 assert!(is_antigravity_process("language_server.exe", None));
2387 assert!(is_antigravity_process(
2388 "language_server_windows_x64.exe",
2389 None
2390 ));
2391 assert!(is_antigravity_process(
2392 "node.exe",
2393 Some(r"C:\Users\u\AppData\Local\agy.exe")
2394 ));
2395 }
2396
2397 #[test]
2398 fn unrelated_processes_are_not_probed() {
2399 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2400 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2401 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2403 assert!(!is_antigravity_process("legacy.exe", None));
2404 assert!(!is_antigravity_process("not-agy.exe", None));
2405 assert!(!is_antigravity_process("", None));
2406 }
2407
2408 #[test]
2409 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2410 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2411 raw.extend([0, b'x' as u16]);
2412 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2413 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2414 assert_eq!(decode_windows_process_name(&[]), "");
2415 }
2416
2417 #[test]
2418 fn windows_process_filter_keeps_only_antigravity_pids() {
2419 let processes = vec![
2420 (10, "agy.exe".to_string()),
2421 (20, "language_server_windows_x64.exe".to_string()),
2422 (30, "sshd.exe".to_string()),
2423 ];
2424 let pids = matching_windows_process_ids(&processes);
2425 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
2426 }
2427
2428 #[test]
2429 fn windows_listener_filter_joins_pid_loopback_and_port() {
2430 let pids = std::collections::HashSet::from([10]);
2431 let rows = [
2432 WindowsTcpRow {
2433 local_addr: [127, 0, 0, 1],
2434 local_port: u32::from(59870u16.to_be()),
2435 pid: 10,
2436 },
2437 WindowsTcpRow {
2438 local_addr: [127, 0, 0, 1],
2439 local_port: u32::from(59868u16.to_be()),
2440 pid: 10,
2441 },
2442 WindowsTcpRow {
2443 local_addr: [127, 0, 0, 1],
2444 local_port: u32::from(59870u16.to_be()),
2445 pid: 10,
2446 },
2447 WindowsTcpRow {
2448 local_addr: [0, 0, 0, 0],
2449 local_port: u32::from(50000u16.to_be()),
2450 pid: 10,
2451 },
2452 WindowsTcpRow {
2453 local_addr: [127, 0, 0, 1],
2454 local_port: u32::from(50001u16.to_be()),
2455 pid: 99,
2456 },
2457 WindowsTcpRow {
2458 local_addr: [127, 0, 0, 1],
2459 local_port: 0,
2460 pid: 10,
2461 },
2462 ];
2463 assert_eq!(
2464 matching_windows_ports(&pids, &rows),
2465 vec![(10, 59870), (10, 59868)]
2466 );
2467 }
2468
2469 #[test]
2473 fn windows_ports_from_two_products_keep_tls_listeners_last() {
2474 let pids = std::collections::HashSet::from([10, 20]);
2475 let row = |port: u16, pid: u32| WindowsTcpRow {
2476 local_addr: [127, 0, 0, 1],
2477 local_port: u32::from(port.to_be()),
2478 pid,
2479 };
2480 let rows = [
2481 row(40000, 10),
2482 row(40001, 10),
2483 row(50000, 20),
2484 row(50001, 20),
2485 ];
2486 assert_eq!(
2487 matching_windows_ports(&pids, &rows),
2488 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2489 );
2490 }
2491
2492 #[test]
2496 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
2497 use std::collections::BTreeMap;
2498
2499 assert_eq!(
2501 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
2502 vec![(10, 59870), (10, 59868)]
2503 );
2504 assert_eq!(
2509 probe_order(BTreeMap::from([
2510 (10, vec![40000, 40001]),
2511 (20, vec![50000, 50001]),
2512 ])),
2513 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2514 );
2515 assert_eq!(
2518 probe_order(BTreeMap::from([
2519 (10, vec![6000, 5000, 4000]),
2520 (20, vec![6000, 7000]),
2521 ])),
2522 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
2523 );
2524 assert!(probe_order(BTreeMap::new()).is_empty());
2525 }
2526
2527 #[test]
2532 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
2533 use std::collections::BTreeMap;
2534
2535 assert_eq!(
2536 probe_order(BTreeMap::from([
2537 (10, vec![40001, 40001, 40000, 40000]),
2538 (20, vec![50001, 50000]),
2539 ])),
2540 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
2541 "duplicate rows must not reorder the ranks below them"
2542 );
2543 }
2544
2545 #[test]
2550 fn a_half_started_product_is_the_documented_exception() {
2551 use std::collections::BTreeMap;
2552
2553 assert_eq!(
2554 probe_order(BTreeMap::from([
2555 (10, vec![40000]),
2556 (20, vec![50001, 50000])
2557 ])),
2558 vec![(10, 40000), (20, 50001), (20, 50000)]
2559 );
2560 }
2561
2562 #[test]
2567 fn an_override_with_no_authority_is_dropped_not_probed() {
2568 for junk in ["/", "///", "http://", "https://", " / "] {
2569 assert_eq!(
2570 candidate_bases_with(Some(junk), vec![(10, 4242)]),
2571 vec![Candidate {
2572 base: "http://127.0.0.1:4242".to_string(),
2573 pid: Some(10),
2574 }],
2575 "{junk:?} should not survive as a candidate"
2576 );
2577 }
2578 assert!(candidate_bases_with(Some("/"), vec![]).is_empty());
2579 }
2580
2581 #[test]
2582 fn windows_table_bounds_reject_truncation_and_overflow() {
2583 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
2584 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
2585 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
2586 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
2587 assert_eq!(
2588 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
2589 None
2590 );
2591 }
2592
2593 #[cfg(target_os = "windows")]
2594 #[test]
2595 fn windows_tcp_table_parser_copies_complete_rows_only() {
2596 use std::mem::{offset_of, size_of};
2597 use windows_sys::Win32::NetworkManagement::IpHelper::{
2598 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
2599 };
2600
2601 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
2602 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
2603 let words = used.div_ceil(size_of::<u32>());
2604 let mut buffer = vec![0u32; words];
2605 let first = MIB_TCPROW_OWNER_PID {
2606 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2607 dwLocalPort: u32::from(59868u16.to_be()),
2608 dwOwningPid: 10,
2609 ..Default::default()
2610 };
2611 let second = MIB_TCPROW_OWNER_PID {
2612 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2613 dwLocalPort: u32::from(59870u16.to_be()),
2614 dwOwningPid: 10,
2615 ..Default::default()
2616 };
2617 unsafe {
2618 buffer.as_mut_ptr().write_unaligned(2);
2619 let rows = buffer
2620 .as_mut_ptr()
2621 .cast::<u8>()
2622 .add(offset)
2623 .cast::<MIB_TCPROW_OWNER_PID>();
2624 rows.write_unaligned(first);
2625 rows.add(1).write_unaligned(second);
2626 }
2627
2628 assert_eq!(
2629 parse_windows_tcp_rows(&buffer, used),
2630 vec![
2631 WindowsTcpRow {
2632 local_addr: [127, 0, 0, 1],
2633 local_port: u32::from(59868u16.to_be()),
2634 pid: 10,
2635 },
2636 WindowsTcpRow {
2637 local_addr: [127, 0, 0, 1],
2638 local_port: u32::from(59870u16.to_be()),
2639 pid: 10,
2640 },
2641 ]
2642 );
2643 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
2644 }
2645
2646 #[test]
2647 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
2648 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";
2651 assert_eq!(
2652 parse_lsof_pcn(output),
2653 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
2654 );
2655 }
2656
2657 #[test]
2661 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
2662 let output = concat!(
2663 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
2664 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
2665 );
2666 assert_eq!(
2667 parse_lsof_pcn(output),
2668 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
2669 "both HTTP listeners must precede both TLS listeners"
2670 );
2671 }
2672
2673 #[test]
2674 fn lsof_parser_matches_the_capitalised_macos_app_name() {
2675 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
2676 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
2677 }
2678
2679 #[test]
2680 fn lsof_parser_deduplicates_and_handles_empty_output() {
2681 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
2682 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
2683 assert!(parse_lsof_pcn("").is_empty());
2684 }
2685
2686 #[test]
2689 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
2690 let dir = tempfile::tempdir().unwrap();
2691 let cache = Cache::at(dir.path().join("usage.json"));
2692 let reason = AppError::Credentials("Antigravity: no local language server found".into());
2693
2694 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2695 let msg = err.to_string();
2696 assert!(msg.contains("no local language server found"), "{msg}");
2697 assert!(!msg.contains("no usable cache"), "{msg}");
2698 }
2699
2700 #[test]
2701 fn unusable_cache_does_not_replace_the_live_diagnosis() {
2702 let dir = tempfile::tempdir().unwrap();
2703 let cache = Cache::at(dir.path().join("antigravity"));
2704 cache.write_payload(b"{}").unwrap();
2705
2706 let reason = AppError::Credentials("Antigravity must be running".into());
2707 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2708 assert!(err.to_string().contains("must be running"), "{err}");
2709
2710 let original = AppError::Transport("original loopback failure".into());
2711 let err = fallback_silent(&cache, now(), original).unwrap_err();
2712 assert!(
2713 err.to_string().contains("original loopback failure"),
2714 "{err}"
2715 );
2716 }
2717
2718 #[tokio::test]
2719 async fn rpc_error_bodies_are_bounded_too() {
2720 let mut server = mockito::Server::new_async().await;
2721 let path = format!("/{STATUS_RPC}");
2722 server
2723 .mock("POST", path.as_str())
2724 .with_status(500)
2725 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
2726 .create_async()
2727 .await;
2728
2729 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
2730 .await
2731 .unwrap_err();
2732 assert!(err.to_string().contains("exceeds"), "{err}");
2733 }
2734
2735 #[test]
2736 fn blank_override_falls_through_to_discovery() {
2737 assert_eq!(
2738 candidate_bases_with(Some(" "), vec![(10, 4242)]),
2739 vec![Candidate {
2740 base: "http://127.0.0.1:4242".to_string(),
2741 pid: Some(10),
2742 }]
2743 );
2744 }
2745
2746 fn fixture() -> (tempfile::TempDir, Cache) {
2751 let td = tempfile::TempDir::new().unwrap();
2752 let cache = Cache::at(td.path().join("antigravity"));
2753 (td, cache)
2754 }
2755
2756 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
2757 let base = server.url();
2758 cloud::Endpoints {
2759 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
2760 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
2761 token: format!("{base}/token"),
2762 }
2763 }
2764
2765 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
2767 let mut token = serde_json::json!({
2768 "access_token": "KEYRING-AT",
2769 "expiry": expiry,
2770 });
2771 if with_refresh {
2772 token["refresh_token"] = serde_json::json!("KEYRING-RT");
2773 }
2774 serde_json::json!({ "token": token }).to_string()
2775 }
2776
2777 const EXPIRED: &str = "2026-07-22T11:00:00Z";
2779 const VALID: &str = "2026-07-22T13:00:00Z";
2781
2782 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
2784 RemoteOverride {
2785 credential: SavedCredential::Blob(blob),
2786 endpoints: Some(eps),
2787 local_bases: Some(vec![]),
2788 ..Default::default()
2789 }
2790 }
2791
2792 fn bare_summary() -> String {
2795 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2796 v["response"].to_string()
2797 }
2798
2799 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
2800 server
2801 .mock("POST", "/daily/quota")
2802 .match_header("authorization", format!("Bearer {bearer}").as_str())
2803 .with_status(200)
2804 .with_body(bare_summary())
2805 }
2806
2807 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
2808 server
2809 .mock("POST", "/token")
2810 .with_status(200)
2811 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
2812 }
2813
2814 fn test_oauth() -> cloud::OauthClient {
2815 cloud::OauthClient {
2816 id: "test-client".into(),
2817 secret: "test-client-secret".into(),
2818 }
2819 }
2820
2821 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
2822 fetch_snapshot_at(
2823 &reqwest::Client::new(),
2824 cache,
2825 ttl,
2826 Some(&test_oauth()),
2827 remote,
2828 now(),
2829 )
2830 .await
2831 }
2832
2833 #[tokio::test]
2834 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
2835 let mut server = mockito::Server::new_async().await;
2836 let eps = endpoints(&server);
2837 let quota = quota_mock(&mut server, "KEYRING-AT")
2838 .expect(1)
2839 .create_async()
2840 .await;
2841 let plan = server
2842 .mock("POST", "/daily/plan")
2843 .match_header("authorization", "Bearer KEYRING-AT")
2844 .with_status(200)
2845 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2846 .expect(1)
2847 .create_async()
2848 .await;
2849 let token = token_mock(&mut server).expect(0).create_async().await;
2851 let blob = keyring_blob(VALID, true);
2852 let (_td, cache) = fixture();
2853
2854 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
2855 .await
2856 .expect("remote path yields a snapshot");
2857
2858 quota.assert_async().await;
2859 plan.assert_async().await;
2860 token.assert_async().await;
2861 let snap = outcome.snapshot;
2862 assert!(!outcome.stale);
2863 assert_eq!(snap.source, AntigravitySource::Remote);
2864 assert_eq!(snap.plan, "Pro");
2865 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2866 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2867 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
2868 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
2869 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
2870 assert_eq!(snap.account, format!("acct:{fingerprint}"));
2871 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
2872
2873 let cached = parse_cache_at(
2875 &cache
2876 .fresh_payload(Duration::from_secs(60))
2877 .unwrap()
2878 .unwrap(),
2879 None,
2880 now(),
2881 )
2882 .unwrap();
2883 assert_eq!(cached.source, AntigravitySource::Remote);
2884 assert_eq!(cached.account, snap.account);
2885 }
2886
2887 #[tokio::test]
2892 async fn agys_missing_csrf_response_uses_the_saved_session() {
2893 let mut server = mockito::Server::new_async().await;
2894 let eps = endpoints(&server);
2895 let root = server
2896 .mock("GET", "/")
2897 .with_status(404)
2898 .expect(1)
2899 .create_async()
2900 .await;
2901 let status_path = format!("/{STATUS_RPC}");
2902 let status = server
2903 .mock("POST", status_path.as_str())
2904 .with_status(401)
2905 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2906 .expect(1)
2907 .create_async()
2908 .await;
2909 let quota = quota_mock(&mut server, "KEYRING-AT")
2910 .expect(1)
2911 .create_async()
2912 .await;
2913 let plan = server
2914 .mock("POST", "/daily/plan")
2915 .match_header("authorization", "Bearer KEYRING-AT")
2916 .with_status(200)
2917 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2918 .expect(1)
2919 .create_async()
2920 .await;
2921 let blob = keyring_blob(VALID, true);
2922 let (_td, cache) = fixture();
2923
2924 let outcome = run(
2925 &cache,
2926 RemoteOverride {
2927 credential: SavedCredential::Blob(&blob),
2928 endpoints: Some(&eps),
2929 local_bases: Some(vec![server.url().into()]),
2930 ..Default::default()
2931 },
2932 Duration::ZERO,
2933 )
2934 .await
2935 .expect("saved session bypasses agy's unusable local RPC");
2936
2937 root.assert_async().await;
2938 status.assert_async().await;
2939 quota.assert_async().await;
2940 plan.assert_async().await;
2941 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
2942 }
2943
2944 #[tokio::test]
2945 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
2946 let mut server = mockito::Server::new_async().await;
2947 let eps = endpoints(&server);
2948 let root = server
2949 .mock("GET", "/")
2950 .with_status(404)
2951 .expect(1)
2952 .create_async()
2953 .await;
2954 let status_path = format!("/{STATUS_RPC}");
2955 let status = server
2956 .mock("POST", status_path.as_str())
2957 .with_status(401)
2958 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2959 .expect(1)
2960 .create_async()
2961 .await;
2962 let quota = quota_mock(&mut server, "KEYRING-AT")
2963 .expect(0)
2964 .create_async()
2965 .await;
2966 let (_td, cache) = fixture();
2967
2968 let error = run(
2969 &cache,
2970 RemoteOverride {
2971 credential: SavedCredential::Absent,
2972 endpoints: Some(&eps),
2973 local_bases: Some(vec![server.url().into()]),
2974 ..Default::default()
2975 },
2976 Duration::ZERO,
2977 )
2978 .await
2979 .expect_err("neither the local RPC nor a saved session is usable");
2980
2981 root.assert_async().await;
2982 status.assert_async().await;
2983 quota.assert_async().await;
2984 let message = error.to_string();
2985 assert!(message.contains("requires a CSRF token"), "{message}");
2986 assert!(message.contains("saved Google session"), "{message}");
2987 }
2988
2989 #[tokio::test]
2994 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
2995 let mut server1 = mockito::Server::new_async().await;
2996 let mut server2 = mockito::Server::new_async().await;
2997 let eps = endpoints(&server1);
2998 let root1 = server1
2999 .mock("GET", "/")
3000 .with_status(404)
3001 .expect(1)
3002 .create_async()
3003 .await;
3004 let status_path = format!("/{STATUS_RPC}");
3005 let status1 = server1
3006 .mock("POST", status_path.as_str())
3007 .with_status(401)
3008 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3009 .expect(1)
3010 .create_async()
3011 .await;
3012
3013 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
3015 let status2 = server2
3016 .mock("POST", status_path.as_str())
3017 .expect(0)
3018 .create_async()
3019 .await;
3020
3021 let (_td, cache) = fixture();
3022 let error = run(
3023 &cache,
3024 RemoteOverride {
3025 credential: SavedCredential::Absent,
3026 endpoints: Some(&eps),
3027 local_bases: Some(vec![
3028 Candidate {
3029 base: server1.url(),
3030 pid: Some(42),
3031 },
3032 Candidate {
3033 base: server2.url(),
3034 pid: Some(42),
3035 },
3036 ]),
3037 ..Default::default()
3038 },
3039 Duration::ZERO,
3040 )
3041 .await
3042 .expect_err("aborted probe returns agy missing csrf diagnosis");
3043
3044 root1.assert_async().await;
3045 status1.assert_async().await;
3046 root2.assert_async().await;
3047 status2.assert_async().await;
3048 let message = error.to_string();
3049 assert!(message.contains("requires a CSRF token"), "{message}");
3050 }
3051
3052 #[tokio::test]
3059 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
3060 let mut agy = mockito::Server::new_async().await;
3061 let mut ide = mockito::Server::new_async().await;
3062 let eps = endpoints(&agy);
3063 let agy_root = agy
3064 .mock("GET", "/")
3065 .with_status(404)
3066 .expect(1)
3067 .create_async()
3068 .await;
3069 let status_path = format!("/{STATUS_RPC}");
3070 let agy_status = agy
3071 .mock("POST", status_path.as_str())
3072 .with_status(401)
3073 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3074 .expect(1)
3075 .create_async()
3076 .await;
3077
3078 let ide_root = ide
3079 .mock("GET", "/")
3080 .with_status(404)
3081 .expect(1)
3082 .create_async()
3083 .await;
3084 let ide_status = ide
3085 .mock("POST", status_path.as_str())
3086 .with_status(200)
3087 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3088 .expect(1)
3089 .create_async()
3090 .await;
3091 let quota_path = format!("/{QUOTA_RPC}");
3092 let ide_quota = ide
3093 .mock("POST", quota_path.as_str())
3094 .with_status(200)
3095 .with_body(QUOTA_JSON)
3096 .expect(1)
3097 .create_async()
3098 .await;
3099
3100 let (_td, cache) = fixture();
3101 let outcome = run(
3102 &cache,
3103 RemoteOverride {
3104 credential: SavedCredential::Absent,
3105 endpoints: Some(&eps),
3106 local_bases: Some(vec![
3107 Candidate {
3108 base: agy.url(),
3109 pid: Some(42),
3110 },
3111 Candidate {
3112 base: ide.url(),
3113 pid: Some(7),
3114 },
3115 ]),
3116 ..Default::default()
3117 },
3118 Duration::ZERO,
3119 )
3120 .await
3121 .expect("the other product's server still answers");
3122
3123 agy_root.assert_async().await;
3124 agy_status.assert_async().await;
3125 ide_root.assert_async().await;
3126 ide_status.assert_async().await;
3127 ide_quota.assert_async().await;
3128 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3129 assert_eq!(outcome.snapshot.plan, "Pro");
3130 }
3131
3132 #[tokio::test]
3135 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
3136 let mut server = mockito::Server::new_async().await;
3137 let eps = endpoints(&server);
3138 let _quota = quota_mock(&mut server, "NEW-AT")
3139 .expect(2)
3140 .create_async()
3141 .await;
3142 let token = token_mock(&mut server).expect(1).create_async().await;
3143 let blob = keyring_blob(EXPIRED, true);
3144 let (_td, cache) = fixture();
3145
3146 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
3147 .await
3148 .expect("refresh, then quota");
3149 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
3150 token.assert_async().await;
3151
3152 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3153 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
3154 .expect("refreshed token persisted under the keyring session's fingerprint");
3155 assert_eq!(persisted.access_token, "NEW-AT");
3156
3157 run(&cache, remote(&blob, &eps), Duration::ZERO)
3159 .await
3160 .expect("persisted token reused");
3161 token.assert_async().await;
3162 }
3163
3164 #[tokio::test]
3167 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
3168 let mut server = mockito::Server::new_async().await;
3169 let eps = endpoints(&server);
3170 let rejected = server
3171 .mock("POST", "/daily/quota")
3172 .match_header("authorization", "Bearer KEYRING-AT")
3173 .with_status(401)
3174 .expect(1)
3175 .create_async()
3176 .await;
3177 let accepted = quota_mock(&mut server, "NEW-AT")
3178 .expect(1)
3179 .create_async()
3180 .await;
3181 let token = token_mock(&mut server).expect(1).create_async().await;
3182 let blob = keyring_blob(VALID, true);
3183 let (_td, cache) = fixture();
3184
3185 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3186 .await
3187 .expect("retry with the refreshed token succeeds");
3188
3189 rejected.assert_async().await;
3190 accepted.assert_async().await;
3191 token.assert_async().await;
3192 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3193 }
3194
3195 #[tokio::test]
3198 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
3199 let mut server = mockito::Server::new_async().await;
3200 let eps = endpoints(&server);
3201 let _token = server
3202 .mock("POST", "/token")
3203 .with_status(400)
3204 .with_body(r#"{"error":"invalid_grant"}"#)
3205 .create_async()
3206 .await;
3207 let blob = keyring_blob(EXPIRED, true);
3208 let (_td, cache) = fixture();
3209
3210 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3211 .await
3212 .expect_err("no cache to fall back on");
3213
3214 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3215 let rendered = err.to_string();
3216 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
3217 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
3218 }
3219
3220 #[tokio::test]
3223 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
3224 let mut server = mockito::Server::new_async().await;
3225 let eps = endpoints(&server);
3226 let _quota = server
3227 .mock("POST", "/daily/quota")
3228 .with_status(401)
3229 .create_async()
3230 .await;
3231 let token = token_mock(&mut server).expect(1).create_async().await;
3232 let blob = keyring_blob(EXPIRED, true);
3233 let (_td, cache) = fixture();
3234
3235 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3236 .await
3237 .expect_err("rejected after refresh");
3238
3239 token.assert_async().await;
3240 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3241 assert!(err.to_string().contains("sign in again"), "{err}");
3242 }
3243
3244 async fn run_without_oauth(cache: &Cache, remote: RemoteOverride<'_>) -> Result<FetchOutcome> {
3246 fetch_snapshot_at(
3247 &reqwest::Client::new(),
3248 cache,
3249 Duration::ZERO,
3250 None,
3251 remote,
3252 now(),
3253 )
3254 .await
3255 }
3256
3257 #[cfg(unix)]
3260 fn fake_agy(
3261 dir: &std::path::Path,
3262 runs: &std::path::Path,
3263 file: &std::path::Path,
3264 blob: &str,
3265 ) -> std::path::PathBuf {
3266 use std::os::unix::fs::PermissionsExt;
3267 let path = dir.join("agy");
3268 std::fs::write(
3269 &path,
3270 format!(
3271 "#!/bin/sh\necho \"$@\" >> '{}'\nprintf '%s' '{}' > '{}'\n",
3272 runs.display(),
3273 blob,
3274 file.display()
3275 ),
3276 )
3277 .unwrap();
3278 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o755)).unwrap();
3279 path
3280 }
3281
3282 #[cfg(unix)]
3285 #[tokio::test]
3286 async fn an_expired_session_is_renewed_by_running_agy() {
3287 let mut server = mockito::Server::new_async().await;
3288 let eps = endpoints(&server);
3289 let quota = quota_mock(&mut server, "KEYRING-AT").create_async().await;
3290 let _plan = server
3291 .mock("POST", "/daily/plan")
3292 .with_status(200)
3293 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
3294 .create_async()
3295 .await;
3296 let (td, cache) = fixture();
3297 let file = td.path().join("credential");
3298 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
3299 let runs = td.path().join("runs");
3300 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
3301
3302 let out = run_without_oauth(
3303 &cache,
3304 RemoteOverride {
3305 credential: SavedCredential::File(&file),
3306 endpoints: Some(&eps),
3307 local_bases: Some(vec![]),
3308 agy: super::super::agy::AgyCommand::At(agy),
3309 },
3310 )
3311 .await
3312 .expect("agy renewed the session");
3313
3314 quota.assert_async().await;
3315 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
3316 assert_eq!(std::fs::read_to_string(runs).unwrap().trim(), "models");
3317 }
3318
3319 #[cfg(unix)]
3322 #[tokio::test]
3323 async fn agy_that_does_not_renew_keeps_the_original_error_and_runs_once() {
3324 let server = mockito::Server::new_async().await;
3325 let eps = endpoints(&server);
3326 let (td, cache) = fixture();
3327 let file = td.path().join("credential");
3328 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
3329 let runs = td.path().join("runs");
3330 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(EXPIRED, true));
3331 let remote = || RemoteOverride {
3332 credential: SavedCredential::File(&file),
3333 endpoints: Some(&eps),
3334 local_bases: Some(vec![]),
3335 agy: super::super::agy::AgyCommand::At(agy.clone()),
3336 };
3337
3338 for _ in 0..2 {
3339 let err = run_without_oauth(&cache, remote())
3340 .await
3341 .expect_err("still expired");
3342 assert!(err.to_string().contains("no OAuth client"), "{err}");
3343 }
3344 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
3345 }
3346
3347 #[tokio::test]
3350 async fn without_an_agy_command_an_expired_session_is_left_alone() {
3351 let server = mockito::Server::new_async().await;
3352 let eps = endpoints(&server);
3353 let blob = keyring_blob(EXPIRED, true);
3354 let (_td, cache) = fixture();
3355
3356 let err = run_without_oauth(&cache, remote(&blob, &eps))
3357 .await
3358 .expect_err("nothing renews it");
3359 assert!(err.to_string().contains("no OAuth client"), "{err}");
3360 }
3361
3362 #[tokio::test]
3364 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
3365 let mut server = mockito::Server::new_async().await;
3366 let eps = endpoints(&server);
3367 let token = token_mock(&mut server).expect(0).create_async().await;
3368 let blob = keyring_blob(EXPIRED, false);
3369 let (_td, cache) = fixture();
3370
3371 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3372 .await
3373 .expect_err("nothing to refresh with");
3374
3375 token.assert_async().await;
3376 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3377 assert!(err.to_string().contains("sign in again"), "{err}");
3378 }
3379
3380 #[tokio::test]
3383 async fn no_saved_session_extends_the_no_local_server_error() {
3384 let mut server = mockito::Server::new_async().await;
3385 let eps = endpoints(&server);
3386 let quota = quota_mock(&mut server, "KEYRING-AT")
3387 .expect(0)
3388 .create_async()
3389 .await;
3390 let (_td, cache) = fixture();
3391
3392 let err = run(
3393 &cache,
3394 RemoteOverride {
3395 credential: SavedCredential::Absent,
3396 endpoints: Some(&eps),
3397 local_bases: Some(vec![]),
3398 ..Default::default()
3399 },
3400 Duration::ZERO,
3401 )
3402 .await
3403 .expect_err("no source at all");
3404
3405 quota.assert_async().await;
3406 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3407 let rendered = err.to_string();
3408 assert!(rendered.contains("no local server found"), "{rendered}");
3409 assert!(rendered.contains("saved Google session"), "{rendered}");
3410 }
3411
3412 #[tokio::test]
3415 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
3416 let mut server = mockito::Server::new_async().await;
3417 let eps = endpoints(&server);
3418 let status_path = format!("/{STATUS_RPC}");
3419 let _status = server
3420 .mock("POST", status_path.as_str())
3421 .with_status(401)
3422 .create_async()
3423 .await;
3424 let quota = quota_mock(&mut server, "KEYRING-AT")
3425 .expect(0)
3426 .create_async()
3427 .await;
3428 let blob = keyring_blob(VALID, true);
3429 let (_td, cache) = fixture();
3430
3431 let err = run(
3432 &cache,
3433 RemoteOverride {
3434 credential: SavedCredential::Blob(&blob),
3435 endpoints: Some(&eps),
3436 local_bases: Some(vec![server.url().into()]),
3437 ..Default::default()
3438 },
3439 Duration::ZERO,
3440 )
3441 .await
3442 .expect_err("signed-out local server");
3443
3444 quota.assert_async().await;
3445 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
3446 }
3447
3448 #[tokio::test]
3451 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
3452 let mut server = mockito::Server::new_async().await;
3453 let eps = endpoints(&server);
3454 for path in ["/daily/quota", "/prod/quota"] {
3455 server
3456 .mock("POST", path)
3457 .with_status(500)
3458 .with_body("boom")
3459 .create_async()
3460 .await;
3461 }
3462 let blob = keyring_blob(VALID, true);
3463 let (_td, cache) = fixture();
3464 cache.ensure_dir().unwrap();
3465 let mut earlier = parsed();
3466 earlier.source = AntigravitySource::Remote;
3467 earlier.account = "acct:earlier".into();
3468 cache
3469 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
3470 .unwrap();
3471
3472 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3473 .await
3474 .expect("stale cache stands in");
3475
3476 assert!(outcome.stale);
3477 assert_eq!(outcome.snapshot, earlier);
3478 assert!(
3479 matches!(outcome.last_error, Some((500, _))),
3480 "{:?}",
3481 outcome.last_error
3482 );
3483 }
3484
3485 #[test]
3486 fn source_round_trips_through_the_cache_and_defaults_to_local() {
3487 let mut snap = parsed();
3488 snap.source = AntigravitySource::Remote;
3489 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3490 assert_eq!(
3491 parse_cache_at(&bytes, None, now()).unwrap().source,
3492 AntigravitySource::Remote
3493 );
3494
3495 snap.source = AntigravitySource::Statusline;
3496 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3497 assert_eq!(
3498 parse_cache_at(&bytes, None, now()).unwrap().source,
3499 AntigravitySource::Statusline
3500 );
3501
3502 let mut legacy = snap_to_json(&snap);
3504 legacy.as_object_mut().unwrap().remove("source");
3505 let legacy = serde_json::to_vec(&legacy).unwrap();
3506 assert_eq!(
3507 parse_cache_at(&legacy, None, now()).unwrap().source,
3508 AntigravitySource::Local
3509 );
3510 }
3511}