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> {
270 if matches!(error, AppError::Credentials(message) if message == NO_LOCAL_SERVER) {
271 Some(RemoteFallbackReason::NoLocalServer)
272 } else if is_missing_csrf(error) {
273 Some(RemoteFallbackReason::AgyMissingCsrf)
274 } else {
275 None
276 }
277}
278
279async fn open_session(client: &reqwest::Client, bases: Option<&[Candidate]>) -> Result<Session> {
285 let candidates = bases.map_or_else(candidate_bases, <[Candidate]>::to_vec);
286 if candidates.is_empty() {
287 return Err(no_local_server());
288 }
289
290 let mut errors = Vec::new();
291 let mut blocked_pids = std::collections::HashSet::new();
296 for candidate in candidates {
297 if candidate.pid.is_some_and(|pid| blocked_pids.contains(&pid)) {
298 continue;
299 }
300 let csrf = fetch_csrf(client, &candidate.base).await;
301 match post_rpc(client, &candidate.base, csrf.as_deref(), STATUS_RPC).await {
302 Ok(v) => {
303 return Ok(Session {
304 base: candidate.base,
305 csrf,
306 plan: plan_from_status(&v),
307 account: account_key(&v),
308 });
309 }
310 Err(e) => {
311 if is_missing_csrf(&e)
315 && let Some(pid) = candidate.pid
316 {
317 blocked_pids.insert(pid);
318 }
319 errors.push(e);
320 }
321 }
322 }
323 Err(select_probe_error(errors))
324}
325
326fn select_probe_error(errors: Vec<AppError>) -> AppError {
343 let mut actionable = None;
344 let mut missing_csrf = None;
345 let mut last = None;
346 let mut echo = None;
347 for e in errors {
348 if missing_csrf.is_none() && is_missing_csrf(&e) {
349 missing_csrf = Some(e);
350 } else if actionable.is_none() && is_actionable(&e) {
351 actionable = Some(e);
352 } else if is_tls_echo(&e) {
353 echo = Some(e);
354 } else {
355 last = Some(e);
356 }
357 }
358 actionable
359 .or(missing_csrf)
360 .or(last)
361 .or(echo)
362 .unwrap_or_else(|| {
363 AppError::Other("antigravity: no local server answered GetUserStatus".into())
364 })
365}
366
367fn is_missing_csrf(error: &AppError) -> bool {
373 let AppError::Http { status: 401, body } = error else {
374 return false;
375 };
376 let Ok(body) = serde_json::from_str::<serde_json::Value>(body) else {
377 return false;
378 };
379 matches!(
380 (body["code"].as_str(), body["message"].as_str()),
381 (Some(code), Some(message))
382 if code.eq_ignore_ascii_case("unauthenticated")
383 && is_missing_csrf_message(message)
384 )
385}
386
387fn is_missing_csrf_message(message: &str) -> bool {
390 message
391 .rsplit(':')
392 .next()
393 .is_some_and(|part| part.trim().eq_ignore_ascii_case("missing CSRF token"))
394}
395
396fn is_actionable(e: &AppError) -> bool {
401 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
402 && !is_missing_csrf(e)
403}
404
405fn is_tls_echo(e: &AppError) -> bool {
418 matches!(
419 e,
420 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
421 )
422}
423
424async fn fetch_live(
425 client: &reqwest::Client,
426 session: Result<Session>,
427) -> Result<AntigravitySnapshot> {
428 let session = session?;
429 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
430 let mut snap = parse_quota_summary("a, session.plan)?;
431 snap.account = session.account;
432 Ok(snap)
433}
434
435fn saved_session(
445 credential: SavedCredential<'_>,
446 reason: RemoteFallbackReason,
447) -> Result<StoredToken> {
448 let raw = match credential {
449 SavedCredential::Keyring => credential::read()?,
450 SavedCredential::Blob(blob) => Some(blob.to_string()),
451 SavedCredential::File(path) => std::fs::read_to_string(path).ok(),
452 SavedCredential::Absent => None,
453 };
454 let Some(raw) = raw else {
455 return Err(AppError::Credentials(format!(
456 "{} {NO_SAVED_SESSION}",
457 reason.message()
458 )));
459 };
460 credential::parse_keyring_blob(&raw)
461}
462
463fn remote_account(fingerprint: &str) -> String {
467 format!("acct:{fingerprint}")
468}
469
470fn session_expired() -> AppError {
471 AppError::Credentials(SESSION_EXPIRED.into())
472}
473
474fn refresh_unconfigured() -> AppError {
478 AppError::Credentials(
479 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
480 refresh it; open Antigravity (or run `agy`) to sign in again, or set [antigravity] oauth_client_id \
481 and oauth_client_secret in config.toml"
482 .into(),
483 )
484}
485
486fn is_auth_rejection(e: &AppError) -> bool {
487 matches!(
488 e,
489 AppError::Http {
490 status: 401 | 403,
491 ..
492 }
493 )
494}
495
496struct AccessToken {
500 value: String,
501 just_refreshed: bool,
502}
503
504async fn resolve_access_token(
509 client: &reqwest::Client,
510 oauth: Option<&cloud::OauthClient>,
511 endpoints: &cloud::Endpoints,
512 oauth_path: &std::path::Path,
513 token: &StoredToken,
514 now: DateTime<Utc>,
515) -> Result<AccessToken> {
516 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
517 && token
518 .expires_at
519 .is_none_or(|keyring| persisted.expires_at > keyring)
520 && !cloud::needs_refresh(Some(persisted.expires_at), now)
521 {
522 return Ok(AccessToken {
523 value: persisted.access_token,
524 just_refreshed: false,
525 });
526 }
527 if !cloud::needs_refresh(token.expires_at, now) {
528 return Ok(AccessToken {
529 value: token.access_token.clone(),
530 just_refreshed: false,
531 });
532 }
533 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
534}
535
536async fn refresh_and_persist(
540 client: &reqwest::Client,
541 oauth: Option<&cloud::OauthClient>,
542 endpoints: &cloud::Endpoints,
543 oauth_path: &std::path::Path,
544 token: &StoredToken,
545) -> Result<AccessToken> {
546 let Some(refresh_token) = token.refresh_token.as_deref() else {
547 return Err(session_expired());
548 };
549 let Some(oauth) = oauth else {
550 return Err(refresh_unconfigured());
551 };
552 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
553 cloud::write_persisted(
554 oauth_path,
555 &cloud::PersistedOAuth {
556 fingerprint: token.fingerprint.clone(),
557 access_token: refreshed.access_token.clone(),
558 expires_at: refreshed.expires_at,
559 },
560 )?;
561 Ok(AccessToken {
562 value: refreshed.access_token,
563 just_refreshed: true,
564 })
565}
566
567struct Renewal<'a> {
570 agy: &'a super::agy::AgyCommand,
571 credential: SavedCredential<'a>,
572}
573
574impl Renewal<'_> {
575 async fn renewed(&self, cache: &Cache, now: DateTime<Utc>) -> Option<StoredToken> {
579 if !super::agy::renew_session(self.agy, cache, now).await {
580 return None;
581 }
582 let raw = match self.credential {
583 SavedCredential::Keyring => credential::read().ok().flatten()?,
584 SavedCredential::Blob(blob) => blob.to_string(),
585 SavedCredential::File(path) => std::fs::read_to_string(path).ok()?,
586 SavedCredential::Absent => return None,
587 };
588 credential::parse_keyring_blob(&raw).ok()
589 }
590}
591
592async fn fetch_remote(
601 client: &reqwest::Client,
602 cache: &Cache,
603 oauth: Option<&cloud::OauthClient>,
604 endpoints: &cloud::Endpoints,
605 token: Result<StoredToken>,
606 renewal: Renewal<'_>,
607 now: DateTime<Utc>,
608) -> Result<AntigravitySnapshot> {
609 let mut token = token?;
610 let oauth_path = cloud::oauth_cache_path(cache);
611 let mut access =
612 match resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await {
613 Ok(access) => access,
614 Err(e) if oauth.is_none() => {
617 let Some(renewed) = renewal.renewed(cache, now).await else {
618 return Err(e);
619 };
620 token = renewed;
621 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
622 }
623 Err(e) => return Err(e),
624 };
625
626 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
627 Err(e) if cloud::is_no_plan_quota(&e) => return Err(e),
630 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
631 access = if oauth.is_none() {
632 let Some(renewed) = renewal.renewed(cache, now).await else {
633 return Err(refresh_unconfigured());
634 };
635 token = renewed;
636 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
637 } else {
638 refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?
639 };
640 cloud::fetch_quota(client, endpoints, &access.value)
641 .await
642 .map_err(|e| {
643 if cloud::is_no_plan_quota(&e) {
644 e
645 } else if is_auth_rejection(&e) {
646 session_expired()
647 } else {
648 e
649 }
650 })?
651 }
652 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
653 other => other?,
654 };
655
656 let plan = cloud::fetch_plan(client, endpoints, &access.value)
657 .await
658 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
659 let mut snap = parse_quota_summary("a, plan)?;
660 snap.account = remote_account(&token.fingerprint);
661 snap.source = AntigravitySource::Remote;
662 Ok(snap)
663}
664
665fn account_key(user_status: &serde_json::Value) -> String {
670 let email = user_status["userStatus"]["email"]
671 .as_str()
672 .filter(|s| !s.is_empty());
673 match email {
674 Some(e) => {
675 use std::hash::{Hash, Hasher};
676 let mut h = std::collections::hash_map::DefaultHasher::new();
677 e.hash(&mut h);
678 format!("acct:{:016x}", h.finish())
679 }
680 None => "acct:unknown".to_string(),
681 }
682}
683
684async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
688 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
689 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
692 .await
693 .ok()?;
694 let html = String::from_utf8_lossy(&bytes);
695 html.split("csrfToken\":\"")
696 .nth(1)
697 .and_then(|s| s.split('"').next())
698 .filter(|t| !t.is_empty())
699 .map(|t| t.to_string())
700}
701
702async fn post_rpc(
703 client: &reqwest::Client,
704 base: &str,
705 csrf: Option<&str>,
706 rpc: &str,
707) -> Result<serde_json::Value> {
708 let mut req = client
709 .post(format!("{base}/{rpc}"))
710 .header("Content-Type", "application/json")
711 .body("{}")
712 .timeout(HTTP_TIMEOUT);
713 if let Some(token) = csrf {
714 req = req.header("x-codeium-csrf-token", token);
715 }
716 let resp = req.send().await?;
717
718 let status = resp.status();
719 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
723 if !status.is_success() {
724 let body = String::from_utf8_lossy(&bytes).into_owned();
725 return Err(AppError::Http {
726 status: status.as_u16(),
727 body,
728 });
729 }
730
731 Ok(serde_json::from_slice(&bytes)?)
732}
733
734pub fn plan_from_status(v: &serde_json::Value) -> String {
735 v["userStatus"]["userTier"]["name"]
736 .as_str()
737 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
738 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
739 .filter(|s| !s.is_empty())
740 .unwrap_or(DEFAULT_PLAN)
741 .to_string()
742}
743
744fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
755 let id = if id.is_empty() { "<unnamed>" } else { id };
756 let group = if group_name.is_empty() {
757 String::new()
758 } else {
759 format!(" in {group_name:?}")
760 };
761 match window {
762 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
763 _ => format!("{id}{group}"),
764 }
765}
766
767fn no_usable_bucket(seen: &[String]) -> AppError {
773 if seen.is_empty() {
774 return AppError::Other(
775 "antigravity: quota summary has no buckets at all — the running product may \
776 not have a quota for this account yet"
777 .into(),
778 );
779 }
780 AppError::Other(format!(
781 "antigravity: quota summary has no bucket in a window we recognise (5h or \
782 weekly, Gemini or Claude/GPT); it offered: {}",
783 seen.join(", ")
784 ))
785}
786
787pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
788 let groups = v["response"]["groups"]
789 .as_array()
790 .or_else(|| v["groups"].as_array())
791 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
792
793 let mut gemini_5h = None;
794 let mut gemini_weekly = None;
795 let mut tp_5h = None;
796 let mut tp_weekly = None;
797 let mut seen: Vec<String> = Vec::new();
801
802 for group in groups {
803 let group_name = group["displayName"].as_str().unwrap_or_default();
804 let Some(buckets) = group["buckets"].as_array() else {
805 continue;
806 };
807 for bucket in buckets {
808 let id = bucket["bucketId"].as_str().unwrap_or_default();
809 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
810 let window = bucket["window"].as_str().unwrap_or_default();
811 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
812 true
813 } else if id.ends_with("5h") || window == "5h" {
814 false
815 } else {
816 continue;
819 };
820 let is_gemini = if id.starts_with("gemini") {
821 true
822 } else if id.starts_with("3p") {
823 false
824 } else if group_name.contains("Gemini") {
825 true
826 } else if group_name.contains("Claude") || group_name.contains("GPT") {
827 false
828 } else {
829 continue;
832 };
833
834 let (slot, slot_name) = match (is_gemini, is_weekly) {
835 (true, false) => (&mut gemini_5h, "Gemini 5h"),
836 (true, true) => (&mut gemini_weekly, "Gemini weekly"),
837 (false, false) => (&mut tp_5h, "Claude/GPT 5h"),
838 (false, true) => (&mut tp_weekly, "Claude/GPT weekly"),
839 };
840 let parsed = usage_window(bucket, is_weekly)?;
841 if slot.replace(parsed).is_some() {
842 return Err(AppError::Schema(format!(
843 "antigravity: duplicate {slot_name} bucket"
844 )));
845 }
846 }
847 }
848
849 if gemini_5h.is_none() && gemini_weekly.is_none() && tp_5h.is_none() && tp_weekly.is_none() {
853 return Err(no_usable_bucket(&seen));
854 }
855
856 Ok(AntigravitySnapshot {
857 plan,
858 account: String::new(),
861 source: AntigravitySource::Local,
862 session: gemini_5h,
863 weekly: gemini_weekly,
864 third_party_session: tp_5h,
865 third_party_weekly: tp_weekly,
866 })
867}
868
869fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
873 let remaining = bucket["remainingFraction"]
874 .as_f64()
875 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
876 .ok_or_else(|| {
877 AppError::Schema(format!(
878 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
879 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
880 ))
881 })?;
882 Ok(UsageWindow {
883 utilization_pct: pct_used(remaining),
884 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
885 window_duration: if is_weekly {
886 chrono::Duration::days(7)
887 } else {
888 chrono::Duration::hours(5)
889 },
890 })
891}
892
893fn pct_used(remaining_fraction: f64) -> i32 {
896 let used = (1.0 - remaining_fraction) * 100.0;
897 used.round() as i32
898}
899
900fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
901 match value {
902 serde_json::Value::Null => Ok(None),
903 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
904 .map(|dt| Some(dt.with_timezone(&Utc)))
905 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
906 _ => Err(AppError::Schema(format!(
907 "antigravity: {field} must be a timestamp or null"
908 ))),
909 }
910}
911
912fn candidate_bases() -> Vec<Candidate> {
918 candidate_bases_with(
919 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
920 discover_ls_ports(),
921 )
922}
923
924fn candidate_bases_with(
928 override_addr: Option<&str>,
929 discovered: Vec<(u32, u16)>,
930) -> Vec<Candidate> {
931 let mut candidates = Vec::new();
932 if let Some(base) = override_addr.and_then(normalize_base) {
933 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
937 candidates.push(Candidate { base, pid });
938 }
939
940 for (pid, port) in discovered {
946 let base = format!("http://127.0.0.1:{port}");
947 if !candidates.iter().any(|c| c.base == base) {
948 candidates.push(Candidate {
949 base,
950 pid: Some(pid),
951 });
952 }
953 }
954
955 candidates
956}
957
958fn base_port(base: &str) -> Option<u16> {
960 let (_, authority) = base.split_once("://")?;
961 authority.rsplit_once(':')?.1.parse().ok()
962}
963
964fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
967 let mut owners = discovered
968 .iter()
969 .filter(|(_, p)| *p == port)
970 .map(|(pid, _)| *pid);
971 let first = owners.next()?;
972 owners.all(|pid| pid == first).then_some(first)
973}
974
975fn normalize_base(addr: &str) -> Option<String> {
984 let trimmed = addr.trim();
985 let (scheme, authority) = match trimmed.split_once("://") {
986 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
987 _ => ("http", trimmed),
988 };
989 let authority = authority.trim_end_matches('/');
990 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
991}
992
993fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
1000 let comm = comm.trim().to_lowercase();
1001 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
1002 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
1003 return true;
1004 }
1005 exe.is_some_and(|p| {
1006 let p = p.to_lowercase().replace('\\', "/");
1007 let p = p.strip_suffix(".exe").unwrap_or(&p);
1008 p.contains("antigravity") || p.ends_with("/agy")
1009 })
1010}
1011
1012#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
1044fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
1045 let groups: Vec<(u32, Vec<u16>)> = per_pid
1046 .into_iter()
1047 .map(|(pid, mut group)| {
1048 group.sort_unstable_by(|a, b| b.cmp(a));
1049 group.dedup();
1054 (pid, group)
1055 })
1056 .collect();
1057
1058 let mut ports: Vec<(u32, u16)> = Vec::new();
1059 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
1060 for (pid, port) in groups
1061 .iter()
1062 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
1063 {
1064 if !ports.iter().any(|(_, p)| *p == port) {
1067 ports.push((pid, port));
1068 }
1069 }
1070 }
1071 ports
1072}
1073
1074#[cfg(target_os = "linux")]
1082pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1083 use std::collections::{BTreeMap, HashMap};
1084
1085 let mut owners: HashMap<u64, u32> = HashMap::new();
1088 let Ok(entries) = std::fs::read_dir("/proc") else {
1089 return Vec::new();
1090 };
1091 for entry in entries.flatten() {
1092 let pid_dir = entry.path();
1093 let Some(pid) = pid_dir
1094 .file_name()
1095 .and_then(|name| name.to_str())
1096 .and_then(|name| name.parse::<u32>().ok())
1097 else {
1098 continue;
1099 };
1100 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1101 continue;
1102 };
1103 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1104 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1105 continue;
1106 }
1107 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1108 continue;
1109 };
1110 for fd in fds.flatten() {
1111 let Ok(target) = std::fs::read_link(fd.path()) else {
1112 continue;
1113 };
1114 if let Some(ino) = target
1115 .to_str()
1116 .and_then(|s| s.strip_prefix("socket:["))
1117 .and_then(|s| s.strip_suffix(']'))
1118 .and_then(|s| s.parse::<u64>().ok())
1119 {
1120 owners.insert(ino, pid);
1121 }
1122 }
1123 }
1124
1125 if owners.is_empty() {
1126 return Vec::new();
1127 }
1128
1129 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1130 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1131 let Ok(contents) = std::fs::read_to_string(table) else {
1132 continue;
1133 };
1134 for line in contents.lines().skip(1) {
1135 if let Some((port, ino)) = parse_proc_net_line(line)
1136 && let Some(&pid) = owners.get(&ino)
1137 {
1138 per_pid.entry(pid).or_default().push(port);
1139 }
1140 }
1141 }
1142 probe_order(per_pid)
1143}
1144
1145#[cfg(target_os = "macos")]
1152pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1153 let Ok(output) = std::process::Command::new("lsof")
1154 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1155 .output()
1156 else {
1157 return Vec::new();
1158 };
1159 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1163}
1164
1165#[cfg(any(test, target_os = "macos"))]
1170fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1171 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1172 let mut pid = None;
1175 let mut owner = None;
1176 for line in output.lines() {
1177 let Some(rest) = line.get(1..) else { continue };
1178 match line.as_bytes().first() {
1179 Some(b'p') => {
1180 pid = rest.parse::<u32>().ok();
1181 owner = None;
1182 }
1183 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1184 Some(b'n') => {
1185 if let Some(pid) = owner
1186 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1187 {
1188 per_pid.entry(pid).or_default().push(port);
1189 }
1190 }
1191 _ => {}
1192 }
1193 }
1194 probe_order(per_pid)
1195}
1196
1197#[cfg(any(test, target_os = "windows"))]
1198#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1199struct WindowsTcpRow {
1200 local_addr: [u8; 4],
1201 local_port: u32,
1202 pid: u32,
1203}
1204
1205#[cfg(any(test, target_os = "windows"))]
1206fn decode_windows_process_name(raw: &[u16]) -> String {
1207 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1208 String::from_utf16_lossy(&raw[..end])
1209}
1210
1211#[cfg(any(test, target_os = "windows"))]
1212fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1213 processes
1214 .iter()
1215 .filter(|(_, name)| is_antigravity_process(name, None))
1216 .map(|(pid, _)| *pid)
1217 .collect()
1218}
1219
1220#[cfg(any(test, target_os = "windows"))]
1223fn matching_windows_ports(
1224 pids: &std::collections::HashSet<u32>,
1225 rows: &[WindowsTcpRow],
1226) -> Vec<(u32, u16)> {
1227 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1228 for row in rows {
1229 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1230 continue;
1231 }
1232 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1233 if port != 0 {
1234 per_pid.entry(row.pid).or_default().push(port);
1235 }
1236 }
1237 probe_order(per_pid)
1238}
1239
1240#[cfg(any(test, target_os = "windows"))]
1241fn checked_windows_row_count(
1242 buffer_len: usize,
1243 rows_offset: usize,
1244 row_size: usize,
1245 declared: usize,
1246) -> Option<usize> {
1247 let rows_len = row_size.checked_mul(declared)?;
1248 let end = rows_offset.checked_add(rows_len)?;
1249 (row_size != 0 && end <= buffer_len).then_some(declared)
1250}
1251
1252#[cfg(target_os = "windows")]
1253struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1254
1255#[cfg(target_os = "windows")]
1256impl Drop for WindowsHandle {
1257 fn drop(&mut self) {
1258 unsafe {
1259 windows_sys::Win32::Foundation::CloseHandle(self.0);
1260 }
1261 }
1262}
1263
1264#[cfg(target_os = "windows")]
1265fn windows_processes() -> Vec<(u32, String)> {
1266 use std::mem::size_of;
1267 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1268 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1269 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1270 TH32CS_SNAPPROCESS,
1271 };
1272
1273 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1274 if handle == INVALID_HANDLE_VALUE {
1275 return Vec::new();
1276 }
1277 let snapshot = WindowsHandle(handle);
1278 let mut entry = PROCESSENTRY32W::default();
1279 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1280 return Vec::new();
1281 };
1282 entry.dwSize = entry_size;
1283 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1284 return Vec::new();
1285 }
1286
1287 let mut processes = Vec::new();
1288 loop {
1289 processes.push((
1290 entry.th32ProcessID,
1291 decode_windows_process_name(&entry.szExeFile),
1292 ));
1293 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1294 break;
1295 }
1296 }
1297 processes
1298}
1299
1300#[cfg(target_os = "windows")]
1301fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1302 use std::mem::{offset_of, size_of, size_of_val};
1303 use windows_sys::Win32::NetworkManagement::IpHelper::{
1304 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1305 };
1306
1307 let available = used_bytes.min(size_of_val(buffer));
1308 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1309 if available < size_of::<u32>() || available < rows_offset {
1310 return Vec::new();
1311 }
1312 let base = buffer.as_ptr().cast::<u8>();
1313 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1314 if checked_windows_row_count(
1315 available,
1316 rows_offset,
1317 size_of::<MIB_TCPROW_OWNER_PID>(),
1318 declared,
1319 )
1320 .is_none()
1321 {
1322 return Vec::new();
1323 }
1324
1325 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1326 (0..declared)
1327 .map(|index| unsafe { rows.add(index).read_unaligned() })
1328 .map(|row| WindowsTcpRow {
1329 local_addr: row.dwLocalAddr.to_ne_bytes(),
1330 local_port: row.dwLocalPort,
1331 pid: row.dwOwningPid,
1332 })
1333 .collect()
1334}
1335
1336#[cfg(target_os = "windows")]
1337fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1338 use std::mem::size_of;
1339 use std::ptr::null_mut;
1340 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1341 use windows_sys::Win32::NetworkManagement::IpHelper::{
1342 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1343 };
1344 use windows_sys::Win32::Networking::WinSock::AF_INET;
1345
1346 let mut size = 0u32;
1347 let status = unsafe {
1348 GetExtendedTcpTable(
1349 null_mut(),
1350 &mut size,
1351 0,
1352 u32::from(AF_INET),
1353 TCP_TABLE_OWNER_PID_LISTENER,
1354 0,
1355 )
1356 };
1357 if status != ERROR_INSUFFICIENT_BUFFER {
1358 return Vec::new();
1359 }
1360
1361 for _ in 0..3 {
1362 let Some(words) = (size as usize)
1363 .checked_add(size_of::<u32>() - 1)
1364 .map(|bytes| bytes / size_of::<u32>())
1365 else {
1366 return Vec::new();
1367 };
1368 if words == 0 {
1369 return Vec::new();
1370 }
1371 let mut buffer = Vec::<u32>::new();
1372 if buffer.try_reserve_exact(words).is_err() {
1373 return Vec::new();
1374 }
1375 buffer.resize(words, 0);
1376 let mut used = size;
1377 let status = unsafe {
1378 GetExtendedTcpTable(
1379 buffer.as_mut_ptr().cast(),
1380 &mut used,
1381 0,
1382 u32::from(AF_INET),
1383 TCP_TABLE_OWNER_PID_LISTENER,
1384 0,
1385 )
1386 };
1387 if status == ERROR_INSUFFICIENT_BUFFER {
1388 size = used;
1389 continue;
1390 }
1391 if status != 0 {
1392 return Vec::new();
1393 }
1394 return parse_windows_tcp_rows(&buffer, used as usize);
1395 }
1396 Vec::new()
1397}
1398
1399#[cfg(target_os = "windows")]
1400pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1401 let pids = matching_windows_process_ids(&windows_processes());
1402 if pids.is_empty() {
1403 return Vec::new();
1404 }
1405 matching_windows_ports(&pids, &windows_tcp_rows())
1406}
1407
1408#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1409pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1410 Vec::new()
1411}
1412
1413#[cfg(target_os = "linux")]
1416fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1417 let cols: Vec<&str> = line.split_whitespace().collect();
1418 if cols.len() < 10 {
1419 return None;
1420 }
1421 if cols[3] != "0A" {
1423 return None;
1424 }
1425 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1426 let inode = cols[9].parse::<u64>().ok()?;
1427 Some((port, inode))
1428}
1429
1430fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1435 crate::outcome::fallback(cache, None, original, |bytes| {
1436 parse_cache_at(bytes, None, now)
1437 })
1438}
1439
1440fn fallback_with_error(
1445 cache: &Cache,
1446 last_error: Option<(u16, String)>,
1447 reason: AppError,
1448 now: DateTime<Utc>,
1449) -> Result<FetchOutcome> {
1450 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1451 parse_cache_at(bytes, None, now)
1452 })
1453}
1454
1455fn reuse_cache(
1456 bytes: Vec<u8>,
1457 cache: &Cache,
1458 stale: bool,
1459 account: Option<&str>,
1460 now: DateTime<Utc>,
1461) -> Result<FetchOutcome> {
1462 let snap = parse_cache_at(&bytes, account, now)?;
1463 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1464}
1465
1466pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1472 parse_cache_at(bytes, account, Utc::now())
1473}
1474
1475fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<&'static str> {
1485 [
1486 ("Gemini 5h", snap.session.as_ref()),
1487 ("Gemini weekly", snap.weekly.as_ref()),
1488 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1489 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1490 ]
1491 .into_iter()
1492 .find(|(_, w)| w.and_then(|w| w.resets_at).is_some_and(|r| r <= now))
1493 .map(|(name, _)| name)
1494}
1495
1496pub fn parse_cache_at(
1497 bytes: &[u8],
1498 account: Option<&str>,
1499 now: DateTime<Utc>,
1500) -> Result<AntigravitySnapshot> {
1501 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1502
1503 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1504 if let Some(expected) = account
1505 && cached_account != Some(expected)
1506 {
1507 return Err(AppError::Schema(
1508 "antigravity cache belongs to a different account; refetching".into(),
1509 ));
1510 }
1511
1512 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1522 match v.get(pct_key) {
1523 None => Err(AppError::Schema(format!(
1524 "antigravity: cached payload is missing {pct_key}"
1525 ))),
1526 Some(serde_json::Value::Null) => Ok(None),
1527 Some(value) => value
1528 .as_i64()
1529 .filter(|pct| (0..=100).contains(pct))
1530 .map(|pct| Some(pct as i32))
1531 .ok_or_else(|| {
1532 AppError::Schema(format!(
1533 "antigravity: cached {pct_key} must be an integer in 0..=100"
1534 ))
1535 }),
1536 }
1537 };
1538
1539 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1540 let Some(pct) = cached_pct(pct_key)? else {
1541 return Ok(None);
1542 };
1543 Ok::<_, AppError>(Some(UsageWindow {
1544 utilization_pct: pct,
1545 resets_at: parse_reset(&v[reset_key], reset_key)?,
1546 window_duration: if weekly {
1547 chrono::Duration::days(7)
1548 } else {
1549 chrono::Duration::hours(5)
1550 },
1551 }))
1552 };
1553
1554 let snap = AntigravitySnapshot {
1555 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1556 account: cached_account.unwrap_or_default().to_string(),
1557 source: v["source"]
1559 .as_str()
1560 .and_then(AntigravitySource::parse)
1561 .unwrap_or_default(),
1562 session: optional("session_pct", "session_reset", false)?,
1563 weekly: optional("weekly_pct", "weekly_reset", true)?,
1564 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1565 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1566 };
1567
1568 if snap.session.is_none()
1571 && snap.weekly.is_none()
1572 && snap.third_party_session.is_none()
1573 && snap.third_party_weekly.is_none()
1574 {
1575 return Err(AppError::Schema(
1576 "antigravity cache holds no usable window; refetching".into(),
1577 ));
1578 }
1579
1580 if let Some(window) = expired_window(&snap, now) {
1581 return Err(AppError::Schema(format!(
1582 "antigravity cache is past its {window} reset; refetching"
1583 )));
1584 }
1585 Ok(snap)
1586}
1587
1588pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1589 serde_json::json!({
1590 "plan": snap.plan,
1591 "account": snap.account,
1592 "source": snap.source.as_str(),
1593 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
1594 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1595 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
1596 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1597 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
1598 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1599 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
1600 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1601 })
1602}
1603
1604#[cfg(test)]
1605mod tests {
1606 use super::*;
1607
1608 const QUOTA_JSON: &str = r#"{
1612 "response": {
1613 "groups": [
1614 {
1615 "displayName": "Gemini Models",
1616 "buckets": [
1617 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
1618 "window": "weekly", "remainingFraction": 0.9191212,
1619 "resetTime": "2026-07-28T17:39:58Z"},
1620 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
1621 "window": "5h", "remainingFraction": 0.5672253,
1622 "resetTime": "2026-07-22T17:47:00Z"}
1623 ]
1624 },
1625 {
1626 "displayName": "Claude and GPT models",
1627 "buckets": [
1628 {"bucketId": "3p-weekly", "window": "weekly",
1629 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
1630 {"bucketId": "3p-5h", "window": "5h",
1631 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
1632 ]
1633 }
1634 ]
1635 }
1636 }"#;
1637
1638 fn now() -> DateTime<Utc> {
1641 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
1642 .unwrap()
1643 .with_timezone(&Utc)
1644 }
1645
1646 fn parsed() -> AntigravitySnapshot {
1647 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1648 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
1649 }
1650
1651 #[test]
1652 fn quota_summary_maps_four_distinct_windows() {
1653 let snap = parsed();
1654 assert_eq!(snap.plan, "Google AI Pro");
1655 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
1657 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
1658 assert_eq!(
1659 snap.third_party_session.as_ref().unwrap().utilization_pct,
1660 75
1661 );
1662 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
1663 }
1664
1665 #[test]
1666 fn each_window_keeps_its_own_reset_time() {
1667 let snap = parsed();
1668 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
1669 assert_eq!(
1670 snap.session.as_ref().unwrap().resets_at,
1671 at("2026-07-22T17:47:00Z")
1672 );
1673 assert_eq!(
1674 snap.weekly.as_ref().unwrap().resets_at,
1675 at("2026-07-28T17:39:58Z")
1676 );
1677 assert_eq!(
1678 snap.third_party_weekly.as_ref().unwrap().resets_at,
1679 at("2026-07-29T12:47:00Z")
1680 );
1681 assert_ne!(
1683 snap.session.as_ref().unwrap().resets_at,
1684 snap.weekly.as_ref().unwrap().resets_at
1685 );
1686 }
1687
1688 #[test]
1689 fn window_durations_match_their_bucket() {
1690 let snap = parsed();
1691 assert_eq!(
1692 snap.session.as_ref().unwrap().window_duration,
1693 chrono::Duration::hours(5)
1694 );
1695 assert_eq!(
1696 snap.weekly.as_ref().unwrap().window_duration,
1697 chrono::Duration::days(7)
1698 );
1699 assert_eq!(
1700 snap.third_party_weekly.as_ref().unwrap().window_duration,
1701 chrono::Duration::days(7)
1702 );
1703 }
1704
1705 #[test]
1706 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
1707 let v: serde_json::Value = serde_json::from_str(
1708 r#"{"response":{"groups":[
1709 {"displayName":"Gemini Models","buckets":[
1710 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
1711 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
1712 {"displayName":"Claude and GPT models","buckets":[
1713 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
1714 ]}}"#,
1715 )
1716 .unwrap();
1717 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
1718 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
1719 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
1720 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
1721 assert!(snap.third_party_weekly.is_none());
1722 }
1723
1724 #[test]
1725 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
1726 let duplicate: serde_json::Value = serde_json::from_str(
1727 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
1728 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1729 {"bucketId":"gemini-5h-copy","window":"5h","remainingFraction":0.1},
1730 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
1731 ]}]}}"#,
1732 )
1733 .unwrap();
1734 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
1735 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
1736
1737 let unrelated: serde_json::Value = serde_json::from_str(
1740 r#"{"response":{"groups":[
1741 {"displayName":"Gemini Models","buckets":[
1742 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1743 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
1744 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
1745 ]},
1746 {"displayName":"Future Models","buckets":[
1747 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
1748 ]}
1749 ]}}"#,
1750 )
1751 .unwrap();
1752 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
1753 assert!(snap.third_party_session.is_none());
1754 assert!(snap.third_party_weekly.is_none());
1755 }
1756
1757 #[test]
1760 fn a_bucket_without_a_usable_fraction_is_rejected() {
1761 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
1762 let v: serde_json::Value = serde_json::from_str(&format!(
1763 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
1764 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
1765 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
1766 ))
1767 .unwrap();
1768 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1769 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
1770 }
1771 }
1772
1773 #[test]
1774 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
1775 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
1776 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1777 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
1778 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1779 assert!(err.to_string().contains("resetTime"), "{err}");
1780 }
1781 }
1782
1783 #[test]
1787 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
1788 let mut snap = parsed();
1789 snap.session = None;
1790 snap.third_party_session = None;
1791
1792 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1793 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
1794
1795 assert!(back.session.is_none());
1796 assert_eq!(
1797 back.weekly.as_ref().unwrap().utilization_pct,
1798 snap.weekly.as_ref().unwrap().utilization_pct
1799 );
1800 }
1801
1802 #[test]
1807 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
1808 let summary = serde_json::json!({
1809 "groups": [
1810 {
1811 "displayName": "Gemini Models",
1812 "buckets": [{
1813 "bucketId": "gemini-weekly", "window": "weekly",
1814 "remainingFraction": 0.42,
1815 }],
1816 },
1817 {
1818 "displayName": "Claude and GPT models",
1819 "buckets": [{
1820 "bucketId": "3p-weekly", "window": "weekly",
1821 "remainingFraction": 0.9,
1822 }],
1823 },
1824 ],
1825 });
1826
1827 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
1828
1829 assert!(snap.session.is_none(), "no 5h bucket arrived");
1830 assert!(snap.third_party_session.is_none());
1831 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
1832 assert_eq!(
1833 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
1834 10
1835 );
1836 }
1837
1838 #[test]
1841 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
1842 let summary = serde_json::json!({
1843 "groups": [{
1844 "displayName": "Gemini Models",
1845 "buckets": [{
1846 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
1847 }],
1848 }],
1849 });
1850
1851 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
1852
1853 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
1854 assert!(snap.weekly.is_none());
1855 }
1856
1857 #[test]
1860 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
1861 let summary = serde_json::json!({
1862 "groups": [{
1863 "displayName": "Gemini Models",
1864 "buckets": [{
1865 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
1866 }],
1867 }],
1868 });
1869
1870 let rendered = parse_quota_summary(&summary, "Pro".into())
1871 .expect_err("an unrecognised cadence alone is not a snapshot")
1872 .to_string();
1873
1874 assert!(
1875 rendered.contains("no bucket in a window we recognise"),
1876 "{rendered}"
1877 );
1878 assert!(rendered.contains("gemini-daily"), "{rendered}");
1879 assert!(rendered.contains("window daily"), "{rendered}");
1880 }
1881
1882 #[test]
1885 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
1886 let summary = serde_json::json!({
1887 "groups": [{"displayName": "Gemini", "buckets": []}],
1888 });
1889
1890 let rendered = parse_quota_summary(&summary, "Pro".into())
1891 .expect_err("no buckets is an error")
1892 .to_string();
1893
1894 assert!(rendered.contains("no buckets at all"), "{rendered}");
1895 assert!(!rendered.contains("it offered:"), "{rendered}");
1896 }
1897
1898 #[test]
1901 fn buckets_without_an_id_are_still_named_in_the_error() {
1902 let summary = serde_json::json!({
1903 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
1904 });
1905
1906 let rendered = parse_quota_summary(&summary, "Pro".into())
1907 .expect_err("an unusable summary is an error")
1908 .to_string();
1909
1910 assert!(rendered.contains("<unnamed>"), "{rendered}");
1911 }
1912
1913 #[test]
1914 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
1915 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
1916 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
1917 }
1918
1919 #[test]
1920 fn cache_round_trip_preserves_every_window() {
1921 let snap = parsed();
1922 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1923 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1924 }
1925
1926 #[test]
1930 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
1931 let full = snap_to_json(&parsed());
1932 for missing in ["session_pct", "weekly_pct"] {
1933 let mut v = full.clone();
1934 v.as_object_mut().unwrap().remove(missing);
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(missing), "{missing}: {err}");
1938 }
1939 assert!(parse_cache_at(b"{}", None, now()).is_err());
1941 }
1942
1943 #[test]
1944 fn cached_percentages_are_range_checked_before_narrowing() {
1945 let full = snap_to_json(&parsed());
1946 for (key, bad) in [
1947 ("session_pct", serde_json::json!(-1)),
1948 ("weekly_pct", serde_json::json!(101)),
1949 ("session_pct", serde_json::json!(i64::MAX)),
1950 ("tp_session_pct", serde_json::json!("75")),
1951 ] {
1952 let mut v = full.clone();
1953 v[key] = bad;
1954 let bytes = serde_json::to_vec(&v).unwrap();
1955 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1956 assert!(err.to_string().contains(key), "{key}: {err}");
1957 }
1958 }
1959
1960 #[test]
1961 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
1962 let mut v = snap_to_json(&parsed());
1963 v["session_reset"] = serde_json::json!("not-a-time");
1964 let bytes = serde_json::to_vec(&v).unwrap();
1965 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1966 assert!(err.to_string().contains("session_reset"), "{err}");
1967 }
1968
1969 #[test]
1974 fn a_cache_past_its_reset_is_refused() {
1975 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
1976 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1977
1978 assert!(parse_cache_at(&bytes, None, now()).is_ok());
1980 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
1982
1983 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
1985 assert!(err.to_string().contains("5h"), "{err}");
1986
1987 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
1989 }
1990
1991 #[test]
1994 fn expiry_names_the_window_that_rolled_over() {
1995 let mut snap = parsed();
1996 snap.session.as_mut().unwrap().resets_at = None;
1998 snap.third_party_session = None;
1999 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2000 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
2001
2002 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
2004
2005 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
2007 assert!(err.to_string().contains("Gemini weekly"), "{err}");
2008 }
2009
2010 #[test]
2012 fn a_window_without_a_reset_never_expires() {
2013 let mut snap = parsed();
2014 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
2015 w.resets_at = None;
2016 }
2017 snap.third_party_session = None;
2018 snap.third_party_weekly = None;
2019 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2020 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
2021 .unwrap()
2022 .with_timezone(&Utc);
2023 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
2024 }
2025
2026 #[test]
2028 fn a_cache_from_another_account_is_rejected() {
2029 let mut snap = parsed();
2030 snap.account = "acct:aaaa".into();
2031 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2032
2033 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
2034 assert_eq!(
2035 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
2036 snap
2037 );
2038
2039 let mut legacy = snap_to_json(&snap);
2041 legacy.as_object_mut().unwrap().remove("account");
2042 let legacy = serde_json::to_vec(&legacy).unwrap();
2043 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
2044 }
2045
2046 #[test]
2049 fn an_unverifiable_cache_is_served_rather_than_discarded() {
2050 let mut snap = parsed();
2051 snap.account = "acct:aaaa".into();
2052 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2053 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2054 }
2055
2056 #[test]
2057 fn account_key_fingerprints_rather_than_storing_the_address() {
2058 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
2059 let a = with("someone@example.com");
2060 assert!(!a.contains("someone"), "{a}");
2061 assert!(!a.contains('@'), "{a}");
2062 assert_eq!(a, with("someone@example.com"), "must be stable");
2063 assert_ne!(a, with("other@example.com"));
2064 let unknown = account_key(&serde_json::json!({}));
2066 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
2067 assert_ne!(unknown, a);
2068 }
2069
2070 #[test]
2073 fn absent_third_party_windows_are_not_treated_as_corruption() {
2074 let mut snap = parsed();
2075 snap.third_party_session = None;
2076 snap.third_party_weekly = None;
2077 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2078 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2079 }
2080
2081 #[test]
2082 fn cache_round_trip_preserves_absent_third_party_windows() {
2083 let mut snap = parsed();
2084 snap.third_party_session = None;
2085 snap.third_party_weekly = None;
2086 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2087 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2088 }
2089
2090 #[test]
2091 fn pct_used_inverts_valid_fractions() {
2092 assert_eq!(pct_used(1.0), 0);
2093 assert_eq!(pct_used(0.0), 100);
2094 assert_eq!(pct_used(0.5), 50);
2095 }
2096
2097 #[test]
2098 fn plan_falls_back_through_the_status_payload() {
2099 let tier: serde_json::Value =
2100 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2101 .unwrap();
2102 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2103
2104 let plan_only: serde_json::Value = serde_json::from_str(
2105 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2106 )
2107 .unwrap();
2108 assert_eq!(plan_from_status(&plan_only), "Pro");
2109
2110 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2111 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2112 }
2113
2114 #[cfg(target_os = "linux")]
2115 #[test]
2116 fn proc_net_parser_keeps_only_listening_rows() {
2117 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2118 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2119 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2120
2121 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2122 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2123 assert_eq!(parse_proc_net_line(established), None);
2124
2125 assert_eq!(parse_proc_net_line("garbage"), None);
2126 }
2127
2128 #[test]
2129 fn explicit_address_comes_first_and_gets_a_scheme() {
2130 let candidate = |base: &str, pid| Candidate {
2131 base: base.to_string(),
2132 pid,
2133 };
2134 assert_eq!(
2135 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)]),
2136 vec![
2137 candidate("http://127.0.0.1:1234", None),
2138 candidate("http://127.0.0.1:5678", Some(10)),
2139 ]
2140 );
2141 assert_eq!(
2143 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)]),
2144 vec![
2145 candidate("http://127.0.0.1:1234", None),
2146 candidate("http://127.0.0.1:5678", Some(10)),
2147 ]
2148 );
2149 assert_eq!(
2152 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)]),
2153 vec![candidate("http://127.0.0.1:5678", Some(10))]
2154 );
2155 assert_eq!(
2157 candidate_bases_with(None, vec![(10, 5678), (10, 5678)]),
2158 vec![candidate("http://127.0.0.1:5678", Some(10))]
2159 );
2160 assert_eq!(
2162 candidate_bases_with(Some("https://host:9"), vec![]),
2163 vec![candidate("https://host:9", None)]
2164 );
2165 }
2166
2167 fn http(status: u16) -> AppError {
2168 AppError::Http {
2169 status,
2170 body: String::new(),
2171 }
2172 }
2173
2174 fn missing_csrf() -> AppError {
2175 missing_csrf_with("missing CSRF token")
2176 }
2177
2178 fn missing_csrf_with(message: &str) -> AppError {
2179 AppError::Http {
2180 status: 401,
2181 body: serde_json::json!({
2182 "code": "unauthenticated",
2183 "message": message,
2184 })
2185 .to_string(),
2186 }
2187 }
2188
2189 #[test]
2190 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2191 assert!(is_missing_csrf(&missing_csrf()));
2192 assert!(is_missing_csrf(&missing_csrf_with(
2193 "unauthenticated: missing CSRF token"
2194 )));
2195 assert!(is_missing_csrf(&missing_csrf_with(
2196 " unauthenticated: missing CSRF token "
2197 )));
2198 assert!(!is_missing_csrf(&missing_csrf_with(
2199 "unauthenticated: invalid token"
2200 )));
2201 assert!(!is_missing_csrf(&AppError::Http {
2202 status: 403,
2203 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2204 }));
2205 assert!(!is_missing_csrf(&AppError::Http {
2206 status: 401,
2207 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2208 }));
2209 assert!(!is_missing_csrf(&AppError::Http {
2210 status: 401,
2211 body: "prefix: missing CSRF token".into(),
2212 }));
2213 }
2214
2215 #[test]
2216 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2217 for errors in [
2218 vec![missing_csrf(), http(403)],
2219 vec![http(401), missing_csrf()],
2220 ] {
2221 let err = select_probe_error(errors);
2222 assert!(is_actionable(&err), "{err}");
2223 assert!(!is_missing_csrf(&err), "{err}");
2224 }
2225 }
2226
2227 #[test]
2228 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2229 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2230 assert!(is_missing_csrf(&err), "{err}");
2231 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2232 }
2233
2234 #[test]
2238 fn an_auth_failure_outranks_later_transport_noise() {
2239 let err = select_probe_error(vec![
2240 http(401),
2241 AppError::Transport("connection refused".into()),
2242 ]);
2243 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2244
2245 let err = select_probe_error(vec![
2246 AppError::Transport("connection refused".into()),
2247 http(403),
2248 ]);
2249 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2250 }
2251
2252 #[test]
2255 fn the_first_auth_failure_wins() {
2256 let err = select_probe_error(vec![http(401), http(403)]);
2257 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2258 }
2259
2260 #[test]
2264 fn without_an_auth_failure_the_last_error_stands() {
2265 let err = select_probe_error(vec![
2266 AppError::Transport("first".into()),
2267 http(500),
2268 AppError::Transport("last".into()),
2269 ]);
2270 assert!(
2271 matches!(&err, AppError::Transport(m) if m == "last"),
2272 "{err}"
2273 );
2274 assert!(err.is_transient());
2275 }
2276
2277 #[test]
2280 fn a_server_error_is_not_treated_as_actionable() {
2281 let err = select_probe_error(vec![http(500), http(401)]);
2282 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2283 }
2284
2285 #[test]
2286 fn no_candidates_at_all_yields_a_generic_error() {
2287 let err = select_probe_error(Vec::new());
2288 assert!(
2289 err.to_string().contains("no local server answered"),
2290 "{err}"
2291 );
2292 }
2293
2294 fn tls_echo() -> AppError {
2298 AppError::Http {
2299 status: 400,
2300 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2301 }
2302 }
2303
2304 #[test]
2309 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2310 let err = select_probe_error(vec![
2311 AppError::Http {
2312 status: 500,
2313 body: "GetUserStatus: internal".into(),
2314 },
2315 tls_echo(),
2316 ]);
2317 assert!(
2318 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2319 "{err}"
2320 );
2321 }
2322
2323 #[test]
2328 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2329 let err = select_probe_error(vec![
2330 AppError::Transport("connection refused".into()),
2331 tls_echo(),
2332 ]);
2333 assert!(
2334 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2335 "{err}"
2336 );
2337 assert!(
2338 err.is_transient(),
2339 "the echo must not make the run non-transient: {err}"
2340 );
2341 }
2342
2343 #[test]
2346 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2347 let err = select_probe_error(vec![tls_echo()]);
2348 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2349 }
2350
2351 #[test]
2354 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2355 let err = select_probe_error(vec![
2356 AppError::Http {
2357 status: 400,
2358 body: "unknown method GetUserStatus".into(),
2359 },
2360 tls_echo(),
2361 ]);
2362 assert!(
2363 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2364 "{err}"
2365 );
2366 }
2367
2368 #[test]
2370 fn an_auth_failure_still_outranks_a_tls_echo() {
2371 let err = select_probe_error(vec![tls_echo(), http(401)]);
2372 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2373 }
2374
2375 #[test]
2376 fn every_discovered_port_is_probed_in_order() {
2377 assert_eq!(
2378 candidate_bases_with(None, vec![(10, 33875), (10, 37435)]),
2379 vec![
2380 Candidate {
2381 base: "http://127.0.0.1:33875".to_string(),
2382 pid: Some(10),
2383 },
2384 Candidate {
2385 base: "http://127.0.0.1:37435".to_string(),
2386 pid: Some(10),
2387 },
2388 ]
2389 );
2390 }
2391
2392 #[test]
2397 fn empty_discovery_yields_no_candidates() {
2398 assert!(candidate_bases_with(None, vec![]).is_empty());
2399 assert!(candidate_bases_with(Some(""), vec![]).is_empty());
2400 }
2401
2402 #[test]
2403 fn every_antigravity_product_is_recognised() {
2404 assert!(is_antigravity_process(
2406 "language_server\n",
2407 Some("/opt/antigravity/resources/bin/language_server")
2408 ));
2409 assert!(is_antigravity_process(
2411 "agy\n",
2412 Some("/home/u/.local/bin/agy")
2413 ));
2414 assert!(is_antigravity_process(
2416 "node",
2417 Some("/opt/antigravity/bin/helper")
2418 ));
2419 assert!(is_antigravity_process("antigravity", None));
2420 assert!(is_antigravity_process("agy.exe", None));
2421 assert!(is_antigravity_process("Antigravity.exe", None));
2422 assert!(is_antigravity_process("language_server.exe", None));
2423 assert!(is_antigravity_process(
2424 "language_server_windows_x64.exe",
2425 None
2426 ));
2427 assert!(is_antigravity_process(
2428 "node.exe",
2429 Some(r"C:\Users\u\AppData\Local\agy.exe")
2430 ));
2431 }
2432
2433 #[test]
2434 fn unrelated_processes_are_not_probed() {
2435 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2436 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2437 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2439 assert!(!is_antigravity_process("legacy.exe", None));
2440 assert!(!is_antigravity_process("not-agy.exe", None));
2441 assert!(!is_antigravity_process("", None));
2442 }
2443
2444 #[test]
2445 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2446 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2447 raw.extend([0, b'x' as u16]);
2448 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2449 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2450 assert_eq!(decode_windows_process_name(&[]), "");
2451 }
2452
2453 #[test]
2454 fn windows_process_filter_keeps_only_antigravity_pids() {
2455 let processes = vec![
2456 (10, "agy.exe".to_string()),
2457 (20, "language_server_windows_x64.exe".to_string()),
2458 (30, "sshd.exe".to_string()),
2459 ];
2460 let pids = matching_windows_process_ids(&processes);
2461 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
2462 }
2463
2464 #[test]
2465 fn windows_listener_filter_joins_pid_loopback_and_port() {
2466 let pids = std::collections::HashSet::from([10]);
2467 let rows = [
2468 WindowsTcpRow {
2469 local_addr: [127, 0, 0, 1],
2470 local_port: u32::from(59870u16.to_be()),
2471 pid: 10,
2472 },
2473 WindowsTcpRow {
2474 local_addr: [127, 0, 0, 1],
2475 local_port: u32::from(59868u16.to_be()),
2476 pid: 10,
2477 },
2478 WindowsTcpRow {
2479 local_addr: [127, 0, 0, 1],
2480 local_port: u32::from(59870u16.to_be()),
2481 pid: 10,
2482 },
2483 WindowsTcpRow {
2484 local_addr: [0, 0, 0, 0],
2485 local_port: u32::from(50000u16.to_be()),
2486 pid: 10,
2487 },
2488 WindowsTcpRow {
2489 local_addr: [127, 0, 0, 1],
2490 local_port: u32::from(50001u16.to_be()),
2491 pid: 99,
2492 },
2493 WindowsTcpRow {
2494 local_addr: [127, 0, 0, 1],
2495 local_port: 0,
2496 pid: 10,
2497 },
2498 ];
2499 assert_eq!(
2500 matching_windows_ports(&pids, &rows),
2501 vec![(10, 59870), (10, 59868)]
2502 );
2503 }
2504
2505 #[test]
2509 fn windows_ports_from_two_products_keep_tls_listeners_last() {
2510 let pids = std::collections::HashSet::from([10, 20]);
2511 let row = |port: u16, pid: u32| WindowsTcpRow {
2512 local_addr: [127, 0, 0, 1],
2513 local_port: u32::from(port.to_be()),
2514 pid,
2515 };
2516 let rows = [
2517 row(40000, 10),
2518 row(40001, 10),
2519 row(50000, 20),
2520 row(50001, 20),
2521 ];
2522 assert_eq!(
2523 matching_windows_ports(&pids, &rows),
2524 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2525 );
2526 }
2527
2528 #[test]
2532 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
2533 use std::collections::BTreeMap;
2534
2535 assert_eq!(
2537 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
2538 vec![(10, 59870), (10, 59868)]
2539 );
2540 assert_eq!(
2545 probe_order(BTreeMap::from([
2546 (10, vec![40000, 40001]),
2547 (20, vec![50000, 50001]),
2548 ])),
2549 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2550 );
2551 assert_eq!(
2554 probe_order(BTreeMap::from([
2555 (10, vec![6000, 5000, 4000]),
2556 (20, vec![6000, 7000]),
2557 ])),
2558 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
2559 );
2560 assert!(probe_order(BTreeMap::new()).is_empty());
2561 }
2562
2563 #[test]
2568 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
2569 use std::collections::BTreeMap;
2570
2571 assert_eq!(
2572 probe_order(BTreeMap::from([
2573 (10, vec![40001, 40001, 40000, 40000]),
2574 (20, vec![50001, 50000]),
2575 ])),
2576 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
2577 "duplicate rows must not reorder the ranks below them"
2578 );
2579 }
2580
2581 #[test]
2586 fn a_half_started_product_is_the_documented_exception() {
2587 use std::collections::BTreeMap;
2588
2589 assert_eq!(
2590 probe_order(BTreeMap::from([
2591 (10, vec![40000]),
2592 (20, vec![50001, 50000])
2593 ])),
2594 vec![(10, 40000), (20, 50001), (20, 50000)]
2595 );
2596 }
2597
2598 #[test]
2603 fn an_override_with_no_authority_is_dropped_not_probed() {
2604 for junk in ["/", "///", "http://", "https://", " / "] {
2605 assert_eq!(
2606 candidate_bases_with(Some(junk), vec![(10, 4242)]),
2607 vec![Candidate {
2608 base: "http://127.0.0.1:4242".to_string(),
2609 pid: Some(10),
2610 }],
2611 "{junk:?} should not survive as a candidate"
2612 );
2613 }
2614 assert!(candidate_bases_with(Some("/"), vec![]).is_empty());
2615 }
2616
2617 #[test]
2618 fn windows_table_bounds_reject_truncation_and_overflow() {
2619 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
2620 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
2621 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
2622 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
2623 assert_eq!(
2624 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
2625 None
2626 );
2627 }
2628
2629 #[cfg(target_os = "windows")]
2630 #[test]
2631 fn windows_tcp_table_parser_copies_complete_rows_only() {
2632 use std::mem::{offset_of, size_of};
2633 use windows_sys::Win32::NetworkManagement::IpHelper::{
2634 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
2635 };
2636
2637 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
2638 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
2639 let words = used.div_ceil(size_of::<u32>());
2640 let mut buffer = vec![0u32; words];
2641 let first = MIB_TCPROW_OWNER_PID {
2642 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2643 dwLocalPort: u32::from(59868u16.to_be()),
2644 dwOwningPid: 10,
2645 ..Default::default()
2646 };
2647 let second = MIB_TCPROW_OWNER_PID {
2648 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2649 dwLocalPort: u32::from(59870u16.to_be()),
2650 dwOwningPid: 10,
2651 ..Default::default()
2652 };
2653 unsafe {
2654 buffer.as_mut_ptr().write_unaligned(2);
2655 let rows = buffer
2656 .as_mut_ptr()
2657 .cast::<u8>()
2658 .add(offset)
2659 .cast::<MIB_TCPROW_OWNER_PID>();
2660 rows.write_unaligned(first);
2661 rows.add(1).write_unaligned(second);
2662 }
2663
2664 assert_eq!(
2665 parse_windows_tcp_rows(&buffer, used),
2666 vec![
2667 WindowsTcpRow {
2668 local_addr: [127, 0, 0, 1],
2669 local_port: u32::from(59868u16.to_be()),
2670 pid: 10,
2671 },
2672 WindowsTcpRow {
2673 local_addr: [127, 0, 0, 1],
2674 local_port: u32::from(59870u16.to_be()),
2675 pid: 10,
2676 },
2677 ]
2678 );
2679 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
2680 }
2681
2682 #[test]
2683 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
2684 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";
2687 assert_eq!(
2688 parse_lsof_pcn(output),
2689 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
2690 );
2691 }
2692
2693 #[test]
2697 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
2698 let output = concat!(
2699 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
2700 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
2701 );
2702 assert_eq!(
2703 parse_lsof_pcn(output),
2704 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
2705 "both HTTP listeners must precede both TLS listeners"
2706 );
2707 }
2708
2709 #[test]
2710 fn lsof_parser_matches_the_capitalised_macos_app_name() {
2711 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
2712 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
2713 }
2714
2715 #[test]
2716 fn lsof_parser_deduplicates_and_handles_empty_output() {
2717 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
2718 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
2719 assert!(parse_lsof_pcn("").is_empty());
2720 }
2721
2722 #[test]
2725 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
2726 let dir = tempfile::tempdir().unwrap();
2727 let cache = Cache::at(dir.path().join("usage.json"));
2728 let reason = AppError::Credentials("Antigravity: no local language server found".into());
2729
2730 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2731 let msg = err.to_string();
2732 assert!(msg.contains("no local language server found"), "{msg}");
2733 assert!(!msg.contains("no usable cache"), "{msg}");
2734 }
2735
2736 #[test]
2737 fn unusable_cache_does_not_replace_the_live_diagnosis() {
2738 let dir = tempfile::tempdir().unwrap();
2739 let cache = Cache::at(dir.path().join("antigravity"));
2740 cache.write_payload(b"{}").unwrap();
2741
2742 let reason = AppError::Credentials("Antigravity must be running".into());
2743 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2744 assert!(err.to_string().contains("must be running"), "{err}");
2745
2746 let original = AppError::Transport("original loopback failure".into());
2747 let err = fallback_silent(&cache, now(), original).unwrap_err();
2748 assert!(
2749 err.to_string().contains("original loopback failure"),
2750 "{err}"
2751 );
2752 }
2753
2754 #[tokio::test]
2755 async fn rpc_error_bodies_are_bounded_too() {
2756 let mut server = mockito::Server::new_async().await;
2757 let path = format!("/{STATUS_RPC}");
2758 server
2759 .mock("POST", path.as_str())
2760 .with_status(500)
2761 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
2762 .create_async()
2763 .await;
2764
2765 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
2766 .await
2767 .unwrap_err();
2768 assert!(err.to_string().contains("exceeds"), "{err}");
2769 }
2770
2771 #[test]
2772 fn blank_override_falls_through_to_discovery() {
2773 assert_eq!(
2774 candidate_bases_with(Some(" "), vec![(10, 4242)]),
2775 vec![Candidate {
2776 base: "http://127.0.0.1:4242".to_string(),
2777 pid: Some(10),
2778 }]
2779 );
2780 }
2781
2782 fn fixture() -> (tempfile::TempDir, Cache) {
2787 let td = tempfile::TempDir::new().unwrap();
2788 let cache = Cache::at(td.path().join("antigravity"));
2789 (td, cache)
2790 }
2791
2792 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
2793 let base = server.url();
2794 cloud::Endpoints {
2795 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
2796 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
2797 token: format!("{base}/token"),
2798 }
2799 }
2800
2801 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
2803 let mut token = serde_json::json!({
2804 "access_token": "KEYRING-AT",
2805 "expiry": expiry,
2806 });
2807 if with_refresh {
2808 token["refresh_token"] = serde_json::json!("KEYRING-RT");
2809 }
2810 serde_json::json!({ "token": token }).to_string()
2811 }
2812
2813 const EXPIRED: &str = "2026-07-22T11:00:00Z";
2815 const VALID: &str = "2026-07-22T13:00:00Z";
2817
2818 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
2820 RemoteOverride {
2821 credential: SavedCredential::Blob(blob),
2822 endpoints: Some(eps),
2823 local_bases: Some(vec![]),
2824 ..Default::default()
2825 }
2826 }
2827
2828 fn bare_summary() -> String {
2831 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2832 v["response"].to_string()
2833 }
2834
2835 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
2836 server
2837 .mock("POST", "/daily/quota")
2838 .match_header("authorization", format!("Bearer {bearer}").as_str())
2839 .with_status(200)
2840 .with_body(bare_summary())
2841 }
2842
2843 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
2844 server
2845 .mock("POST", "/token")
2846 .with_status(200)
2847 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
2848 }
2849
2850 fn test_oauth() -> cloud::OauthClient {
2851 cloud::OauthClient {
2852 id: "test-client".into(),
2853 secret: "test-client-secret".into(),
2854 }
2855 }
2856
2857 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
2858 fetch_snapshot_at(
2859 &reqwest::Client::new(),
2860 cache,
2861 ttl,
2862 Some(&test_oauth()),
2863 remote,
2864 now(),
2865 )
2866 .await
2867 }
2868
2869 #[tokio::test]
2870 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
2871 let mut server = mockito::Server::new_async().await;
2872 let eps = endpoints(&server);
2873 let quota = quota_mock(&mut server, "KEYRING-AT")
2874 .expect(1)
2875 .create_async()
2876 .await;
2877 let plan = server
2878 .mock("POST", "/daily/plan")
2879 .match_header("authorization", "Bearer KEYRING-AT")
2880 .with_status(200)
2881 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2882 .expect(1)
2883 .create_async()
2884 .await;
2885 let token = token_mock(&mut server).expect(0).create_async().await;
2887 let blob = keyring_blob(VALID, true);
2888 let (_td, cache) = fixture();
2889
2890 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
2891 .await
2892 .expect("remote path yields a snapshot");
2893
2894 quota.assert_async().await;
2895 plan.assert_async().await;
2896 token.assert_async().await;
2897 let snap = outcome.snapshot;
2898 assert!(!outcome.stale);
2899 assert_eq!(snap.source, AntigravitySource::Remote);
2900 assert_eq!(snap.plan, "Pro");
2901 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2902 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2903 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
2904 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
2905 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
2906 assert_eq!(snap.account, format!("acct:{fingerprint}"));
2907 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
2908
2909 let cached = parse_cache_at(
2911 &cache
2912 .fresh_payload(Duration::from_secs(60))
2913 .unwrap()
2914 .unwrap(),
2915 None,
2916 now(),
2917 )
2918 .unwrap();
2919 assert_eq!(cached.source, AntigravitySource::Remote);
2920 assert_eq!(cached.account, snap.account);
2921 }
2922
2923 #[tokio::test]
2928 async fn agys_missing_csrf_response_uses_the_saved_session() {
2929 let mut server = mockito::Server::new_async().await;
2930 let eps = endpoints(&server);
2931 let root = server
2932 .mock("GET", "/")
2933 .with_status(404)
2934 .expect(1)
2935 .create_async()
2936 .await;
2937 let status_path = format!("/{STATUS_RPC}");
2938 let status = server
2939 .mock("POST", status_path.as_str())
2940 .with_status(401)
2941 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2942 .expect(1)
2943 .create_async()
2944 .await;
2945 let quota = quota_mock(&mut server, "KEYRING-AT")
2946 .expect(1)
2947 .create_async()
2948 .await;
2949 let plan = server
2950 .mock("POST", "/daily/plan")
2951 .match_header("authorization", "Bearer KEYRING-AT")
2952 .with_status(200)
2953 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2954 .expect(1)
2955 .create_async()
2956 .await;
2957 let blob = keyring_blob(VALID, true);
2958 let (_td, cache) = fixture();
2959
2960 let outcome = run(
2961 &cache,
2962 RemoteOverride {
2963 credential: SavedCredential::Blob(&blob),
2964 endpoints: Some(&eps),
2965 local_bases: Some(vec![server.url().into()]),
2966 ..Default::default()
2967 },
2968 Duration::ZERO,
2969 )
2970 .await
2971 .expect("saved session bypasses agy's unusable local RPC");
2972
2973 root.assert_async().await;
2974 status.assert_async().await;
2975 quota.assert_async().await;
2976 plan.assert_async().await;
2977 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
2978 }
2979
2980 #[tokio::test]
2984 async fn agys_prefixed_missing_csrf_response_uses_the_saved_session() {
2985 let mut server = mockito::Server::new_async().await;
2986 let eps = endpoints(&server);
2987 let root = server
2988 .mock("GET", "/")
2989 .with_status(404)
2990 .expect(1)
2991 .create_async()
2992 .await;
2993 let status_path = format!("/{STATUS_RPC}");
2994 let status = server
2995 .mock("POST", status_path.as_str())
2996 .with_status(401)
2997 .with_body(
2998 r#"{"code":"unauthenticated","message":"unauthenticated: missing CSRF token"}"#,
2999 )
3000 .expect(1)
3001 .create_async()
3002 .await;
3003 let quota = quota_mock(&mut server, "KEYRING-AT")
3004 .expect(1)
3005 .create_async()
3006 .await;
3007 let plan = server
3008 .mock("POST", "/daily/plan")
3009 .match_header("authorization", "Bearer KEYRING-AT")
3010 .with_status(200)
3011 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3012 .expect(1)
3013 .create_async()
3014 .await;
3015 let blob = keyring_blob(VALID, true);
3016 let (_td, cache) = fixture();
3017
3018 let outcome = run(
3019 &cache,
3020 RemoteOverride {
3021 credential: SavedCredential::Blob(&blob),
3022 endpoints: Some(&eps),
3023 local_bases: Some(vec![server.url().into()]),
3024 ..Default::default()
3025 },
3026 Duration::ZERO,
3027 )
3028 .await
3029 .expect("saved session bypasses agy 1.3.1's prefixed missing-CSRF RPC");
3030
3031 root.assert_async().await;
3032 status.assert_async().await;
3033 quota.assert_async().await;
3034 plan.assert_async().await;
3035 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3036 }
3037
3038 #[tokio::test]
3039 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
3040 let mut server = mockito::Server::new_async().await;
3041 let eps = endpoints(&server);
3042 let root = server
3043 .mock("GET", "/")
3044 .with_status(404)
3045 .expect(1)
3046 .create_async()
3047 .await;
3048 let status_path = format!("/{STATUS_RPC}");
3049 let status = server
3050 .mock("POST", status_path.as_str())
3051 .with_status(401)
3052 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3053 .expect(1)
3054 .create_async()
3055 .await;
3056 let quota = quota_mock(&mut server, "KEYRING-AT")
3057 .expect(0)
3058 .create_async()
3059 .await;
3060 let (_td, cache) = fixture();
3061
3062 let error = run(
3063 &cache,
3064 RemoteOverride {
3065 credential: SavedCredential::Absent,
3066 endpoints: Some(&eps),
3067 local_bases: Some(vec![server.url().into()]),
3068 ..Default::default()
3069 },
3070 Duration::ZERO,
3071 )
3072 .await
3073 .expect_err("neither the local RPC nor a saved session is usable");
3074
3075 root.assert_async().await;
3076 status.assert_async().await;
3077 quota.assert_async().await;
3078 let message = error.to_string();
3079 assert!(message.contains("requires a CSRF token"), "{message}");
3080 assert!(message.contains("saved Google session"), "{message}");
3081 }
3082
3083 #[tokio::test]
3088 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
3089 let mut server1 = mockito::Server::new_async().await;
3090 let mut server2 = mockito::Server::new_async().await;
3091 let eps = endpoints(&server1);
3092 let root1 = server1
3093 .mock("GET", "/")
3094 .with_status(404)
3095 .expect(1)
3096 .create_async()
3097 .await;
3098 let status_path = format!("/{STATUS_RPC}");
3099 let status1 = server1
3100 .mock("POST", status_path.as_str())
3101 .with_status(401)
3102 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3103 .expect(1)
3104 .create_async()
3105 .await;
3106
3107 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
3109 let status2 = server2
3110 .mock("POST", status_path.as_str())
3111 .expect(0)
3112 .create_async()
3113 .await;
3114
3115 let (_td, cache) = fixture();
3116 let error = run(
3117 &cache,
3118 RemoteOverride {
3119 credential: SavedCredential::Absent,
3120 endpoints: Some(&eps),
3121 local_bases: Some(vec![
3122 Candidate {
3123 base: server1.url(),
3124 pid: Some(42),
3125 },
3126 Candidate {
3127 base: server2.url(),
3128 pid: Some(42),
3129 },
3130 ]),
3131 ..Default::default()
3132 },
3133 Duration::ZERO,
3134 )
3135 .await
3136 .expect_err("aborted probe returns agy missing csrf diagnosis");
3137
3138 root1.assert_async().await;
3139 status1.assert_async().await;
3140 root2.assert_async().await;
3141 status2.assert_async().await;
3142 let message = error.to_string();
3143 assert!(message.contains("requires a CSRF token"), "{message}");
3144 }
3145
3146 #[tokio::test]
3153 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
3154 let mut agy = mockito::Server::new_async().await;
3155 let mut ide = mockito::Server::new_async().await;
3156 let eps = endpoints(&agy);
3157 let agy_root = agy
3158 .mock("GET", "/")
3159 .with_status(404)
3160 .expect(1)
3161 .create_async()
3162 .await;
3163 let status_path = format!("/{STATUS_RPC}");
3164 let agy_status = agy
3165 .mock("POST", status_path.as_str())
3166 .with_status(401)
3167 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3168 .expect(1)
3169 .create_async()
3170 .await;
3171
3172 let ide_root = ide
3173 .mock("GET", "/")
3174 .with_status(404)
3175 .expect(1)
3176 .create_async()
3177 .await;
3178 let ide_status = ide
3179 .mock("POST", status_path.as_str())
3180 .with_status(200)
3181 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3182 .expect(1)
3183 .create_async()
3184 .await;
3185 let quota_path = format!("/{QUOTA_RPC}");
3186 let ide_quota = ide
3187 .mock("POST", quota_path.as_str())
3188 .with_status(200)
3189 .with_body(QUOTA_JSON)
3190 .expect(1)
3191 .create_async()
3192 .await;
3193
3194 let (_td, cache) = fixture();
3195 let outcome = run(
3196 &cache,
3197 RemoteOverride {
3198 credential: SavedCredential::Absent,
3199 endpoints: Some(&eps),
3200 local_bases: Some(vec![
3201 Candidate {
3202 base: agy.url(),
3203 pid: Some(42),
3204 },
3205 Candidate {
3206 base: ide.url(),
3207 pid: Some(7),
3208 },
3209 ]),
3210 ..Default::default()
3211 },
3212 Duration::ZERO,
3213 )
3214 .await
3215 .expect("the other product's server still answers");
3216
3217 agy_root.assert_async().await;
3218 agy_status.assert_async().await;
3219 ide_root.assert_async().await;
3220 ide_status.assert_async().await;
3221 ide_quota.assert_async().await;
3222 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3223 assert_eq!(outcome.snapshot.plan, "Pro");
3224 }
3225
3226 #[tokio::test]
3229 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
3230 let mut server = mockito::Server::new_async().await;
3231 let eps = endpoints(&server);
3232 let _quota = quota_mock(&mut server, "NEW-AT")
3233 .expect(2)
3234 .create_async()
3235 .await;
3236 let token = token_mock(&mut server).expect(1).create_async().await;
3237 let blob = keyring_blob(EXPIRED, true);
3238 let (_td, cache) = fixture();
3239
3240 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
3241 .await
3242 .expect("refresh, then quota");
3243 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
3244 token.assert_async().await;
3245
3246 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3247 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
3248 .expect("refreshed token persisted under the keyring session's fingerprint");
3249 assert_eq!(persisted.access_token, "NEW-AT");
3250
3251 run(&cache, remote(&blob, &eps), Duration::ZERO)
3253 .await
3254 .expect("persisted token reused");
3255 token.assert_async().await;
3256 }
3257
3258 #[tokio::test]
3261 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
3262 let mut server = mockito::Server::new_async().await;
3263 let eps = endpoints(&server);
3264 let rejected = server
3265 .mock("POST", "/daily/quota")
3266 .match_header("authorization", "Bearer KEYRING-AT")
3267 .with_status(401)
3268 .expect(1)
3269 .create_async()
3270 .await;
3271 let accepted = quota_mock(&mut server, "NEW-AT")
3272 .expect(1)
3273 .create_async()
3274 .await;
3275 let token = token_mock(&mut server).expect(1).create_async().await;
3276 let blob = keyring_blob(VALID, true);
3277 let (_td, cache) = fixture();
3278
3279 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3280 .await
3281 .expect("retry with the refreshed token succeeds");
3282
3283 rejected.assert_async().await;
3284 accepted.assert_async().await;
3285 token.assert_async().await;
3286 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3287 }
3288
3289 #[tokio::test]
3292 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
3293 let mut server = mockito::Server::new_async().await;
3294 let eps = endpoints(&server);
3295 let _token = server
3296 .mock("POST", "/token")
3297 .with_status(400)
3298 .with_body(r#"{"error":"invalid_grant"}"#)
3299 .create_async()
3300 .await;
3301 let blob = keyring_blob(EXPIRED, true);
3302 let (_td, cache) = fixture();
3303
3304 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3305 .await
3306 .expect_err("no cache to fall back on");
3307
3308 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3309 let rendered = err.to_string();
3310 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
3311 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
3312 }
3313
3314 #[tokio::test]
3317 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
3318 let mut server = mockito::Server::new_async().await;
3319 let eps = endpoints(&server);
3320 let _quota = server
3321 .mock("POST", "/daily/quota")
3322 .with_status(401)
3323 .create_async()
3324 .await;
3325 let token = token_mock(&mut server).expect(1).create_async().await;
3326 let blob = keyring_blob(EXPIRED, true);
3327 let (_td, cache) = fixture();
3328
3329 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3330 .await
3331 .expect_err("rejected after refresh");
3332
3333 token.assert_async().await;
3334 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3335 assert!(err.to_string().contains("sign in again"), "{err}");
3336 }
3337
3338 #[cfg(unix)]
3342 #[tokio::test]
3343 async fn a_subscription_403_surfaces_the_plan_message_without_a_refresh_or_agy_run() {
3344 let mut server = mockito::Server::new_async().await;
3345 let eps = endpoints(&server);
3346 let quota = server
3347 .mock("POST", "/daily/quota")
3348 .match_header("authorization", "Bearer KEYRING-AT")
3349 .with_status(403)
3350 .with_body(
3351 r#"{"error":{"code":403,"status":"PERMISSION_DENIED","details":
3352 [{"reason":"SUBSCRIPTION_REQUIRED","domain":"cloud.google.com"}]}}"#,
3353 )
3354 .expect(1)
3355 .create_async()
3356 .await;
3357 let prod = server
3358 .mock("POST", "/prod/quota")
3359 .expect(0)
3360 .create_async()
3361 .await;
3362 let token = token_mock(&mut server).expect(0).create_async().await;
3363 let (td, cache) = fixture();
3364 let file = td.path().join("credential");
3365 std::fs::write(&file, keyring_blob(VALID, true)).unwrap();
3366 let runs = td.path().join("runs");
3367 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
3368
3369 let err = run_without_oauth(
3370 &cache,
3371 RemoteOverride {
3372 credential: SavedCredential::File(&file),
3373 endpoints: Some(&eps),
3374 local_bases: Some(vec![]),
3375 agy: super::super::agy::AgyCommand::At(agy),
3376 },
3377 )
3378 .await
3379 .expect_err("a plan without quota is an error, not a snapshot");
3380
3381 quota.assert_async().await;
3382 prod.assert_async().await;
3383 token.assert_async().await;
3384 match err {
3385 AppError::Http { status, body } => {
3386 assert_eq!(status, 403);
3387 assert!(body.contains("paid subscription"), "{body}");
3388 }
3389 other => panic!("expected Http, got {other:?}"),
3390 }
3391 assert!(
3392 !runs.exists(),
3393 "a subscription verdict must not spawn `agy`"
3394 );
3395 }
3396
3397 #[tokio::test]
3401 async fn a_non_subscription_403_keeps_the_refresh_and_retry_behavior() {
3402 let mut server = mockito::Server::new_async().await;
3403 let eps = endpoints(&server);
3404 let rejected = server
3405 .mock("POST", "/daily/quota")
3406 .match_header("authorization", "Bearer KEYRING-AT")
3407 .with_status(403)
3408 .with_body(r#"{"error":{"message":"acl detail"}}"#)
3409 .expect(1)
3410 .create_async()
3411 .await;
3412 let retried = server
3413 .mock("POST", "/daily/quota")
3414 .match_header("authorization", "Bearer NEW-AT")
3415 .with_status(403)
3416 .with_body(r#"{"error":{"message":"acl detail"}}"#)
3417 .expect(1)
3418 .create_async()
3419 .await;
3420 let token = token_mock(&mut server).expect(1).create_async().await;
3421 let blob = keyring_blob(VALID, true);
3422 let (_td, cache) = fixture();
3423
3424 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3425 .await
3426 .expect_err("still a rejected session");
3427
3428 rejected.assert_async().await;
3429 retried.assert_async().await;
3430 token.assert_async().await;
3431 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3432 assert!(err.to_string().contains("sign in again"), "{err}");
3433 }
3434
3435 async fn run_without_oauth(cache: &Cache, remote: RemoteOverride<'_>) -> Result<FetchOutcome> {
3437 fetch_snapshot_at(
3438 &reqwest::Client::new(),
3439 cache,
3440 Duration::ZERO,
3441 None,
3442 remote,
3443 now(),
3444 )
3445 .await
3446 }
3447
3448 #[cfg(unix)]
3451 fn fake_agy(
3452 dir: &std::path::Path,
3453 runs: &std::path::Path,
3454 file: &std::path::Path,
3455 blob: &str,
3456 ) -> std::path::PathBuf {
3457 use std::os::unix::fs::PermissionsExt;
3458 let path = dir.join("agy");
3459 std::fs::write(
3460 &path,
3461 format!(
3462 "#!/bin/sh\necho \"$@\" >> '{}'\nprintf '%s' '{}' > '{}'\n",
3463 runs.display(),
3464 blob,
3465 file.display()
3466 ),
3467 )
3468 .unwrap();
3469 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o755)).unwrap();
3470 path
3471 }
3472
3473 #[cfg(unix)]
3476 #[tokio::test]
3477 async fn an_expired_session_is_renewed_by_running_agy() {
3478 let mut server = mockito::Server::new_async().await;
3479 let eps = endpoints(&server);
3480 let quota = quota_mock(&mut server, "KEYRING-AT").create_async().await;
3481 let _plan = server
3482 .mock("POST", "/daily/plan")
3483 .with_status(200)
3484 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
3485 .create_async()
3486 .await;
3487 let (td, cache) = fixture();
3488 let file = td.path().join("credential");
3489 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
3490 let runs = td.path().join("runs");
3491 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
3492
3493 let out = run_without_oauth(
3494 &cache,
3495 RemoteOverride {
3496 credential: SavedCredential::File(&file),
3497 endpoints: Some(&eps),
3498 local_bases: Some(vec![]),
3499 agy: super::super::agy::AgyCommand::At(agy),
3500 },
3501 )
3502 .await
3503 .expect("agy renewed the session");
3504
3505 quota.assert_async().await;
3506 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
3507 assert_eq!(std::fs::read_to_string(runs).unwrap().trim(), "models");
3508 }
3509
3510 #[cfg(unix)]
3513 #[tokio::test]
3514 async fn agy_that_does_not_renew_keeps_the_original_error_and_runs_once() {
3515 let server = mockito::Server::new_async().await;
3516 let eps = endpoints(&server);
3517 let (td, cache) = fixture();
3518 let file = td.path().join("credential");
3519 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
3520 let runs = td.path().join("runs");
3521 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(EXPIRED, true));
3522 let remote = || RemoteOverride {
3523 credential: SavedCredential::File(&file),
3524 endpoints: Some(&eps),
3525 local_bases: Some(vec![]),
3526 agy: super::super::agy::AgyCommand::At(agy.clone()),
3527 };
3528
3529 for _ in 0..2 {
3530 let err = run_without_oauth(&cache, remote())
3531 .await
3532 .expect_err("still expired");
3533 assert!(err.to_string().contains("no OAuth client"), "{err}");
3534 }
3535 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
3536 }
3537
3538 #[tokio::test]
3541 async fn without_an_agy_command_an_expired_session_is_left_alone() {
3542 let server = mockito::Server::new_async().await;
3543 let eps = endpoints(&server);
3544 let blob = keyring_blob(EXPIRED, true);
3545 let (_td, cache) = fixture();
3546
3547 let err = run_without_oauth(&cache, remote(&blob, &eps))
3548 .await
3549 .expect_err("nothing renews it");
3550 assert!(err.to_string().contains("no OAuth client"), "{err}");
3551 }
3552
3553 #[tokio::test]
3555 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
3556 let mut server = mockito::Server::new_async().await;
3557 let eps = endpoints(&server);
3558 let token = token_mock(&mut server).expect(0).create_async().await;
3559 let blob = keyring_blob(EXPIRED, false);
3560 let (_td, cache) = fixture();
3561
3562 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3563 .await
3564 .expect_err("nothing to refresh with");
3565
3566 token.assert_async().await;
3567 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3568 assert!(err.to_string().contains("sign in again"), "{err}");
3569 }
3570
3571 #[tokio::test]
3574 async fn no_saved_session_extends_the_no_local_server_error() {
3575 let mut server = mockito::Server::new_async().await;
3576 let eps = endpoints(&server);
3577 let quota = quota_mock(&mut server, "KEYRING-AT")
3578 .expect(0)
3579 .create_async()
3580 .await;
3581 let (_td, cache) = fixture();
3582
3583 let err = run(
3584 &cache,
3585 RemoteOverride {
3586 credential: SavedCredential::Absent,
3587 endpoints: Some(&eps),
3588 local_bases: Some(vec![]),
3589 ..Default::default()
3590 },
3591 Duration::ZERO,
3592 )
3593 .await
3594 .expect_err("no source at all");
3595
3596 quota.assert_async().await;
3597 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3598 let rendered = err.to_string();
3599 assert!(rendered.contains("no local server found"), "{rendered}");
3600 assert!(rendered.contains("saved Google session"), "{rendered}");
3601 }
3602
3603 #[tokio::test]
3606 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
3607 let mut server = mockito::Server::new_async().await;
3608 let eps = endpoints(&server);
3609 let status_path = format!("/{STATUS_RPC}");
3610 let _status = server
3611 .mock("POST", status_path.as_str())
3612 .with_status(401)
3613 .create_async()
3614 .await;
3615 let quota = quota_mock(&mut server, "KEYRING-AT")
3616 .expect(0)
3617 .create_async()
3618 .await;
3619 let blob = keyring_blob(VALID, true);
3620 let (_td, cache) = fixture();
3621
3622 let err = run(
3623 &cache,
3624 RemoteOverride {
3625 credential: SavedCredential::Blob(&blob),
3626 endpoints: Some(&eps),
3627 local_bases: Some(vec![server.url().into()]),
3628 ..Default::default()
3629 },
3630 Duration::ZERO,
3631 )
3632 .await
3633 .expect_err("signed-out local server");
3634
3635 quota.assert_async().await;
3636 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
3637 }
3638
3639 #[tokio::test]
3642 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
3643 let mut server = mockito::Server::new_async().await;
3644 let eps = endpoints(&server);
3645 for path in ["/daily/quota", "/prod/quota"] {
3646 server
3647 .mock("POST", path)
3648 .with_status(500)
3649 .with_body("boom")
3650 .create_async()
3651 .await;
3652 }
3653 let blob = keyring_blob(VALID, true);
3654 let (_td, cache) = fixture();
3655 cache.ensure_dir().unwrap();
3656 let mut earlier = parsed();
3657 earlier.source = AntigravitySource::Remote;
3658 earlier.account = "acct:earlier".into();
3659 cache
3660 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
3661 .unwrap();
3662
3663 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3664 .await
3665 .expect("stale cache stands in");
3666
3667 assert!(outcome.stale);
3668 assert_eq!(outcome.snapshot, earlier);
3669 assert!(
3670 matches!(outcome.last_error, Some((500, _))),
3671 "{:?}",
3672 outcome.last_error
3673 );
3674 }
3675
3676 #[test]
3677 fn source_round_trips_through_the_cache_and_defaults_to_local() {
3678 let mut snap = parsed();
3679 snap.source = AntigravitySource::Remote;
3680 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3681 assert_eq!(
3682 parse_cache_at(&bytes, None, now()).unwrap().source,
3683 AntigravitySource::Remote
3684 );
3685
3686 snap.source = AntigravitySource::Statusline;
3687 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3688 assert_eq!(
3689 parse_cache_at(&bytes, None, now()).unwrap().source,
3690 AntigravitySource::Statusline
3691 );
3692
3693 let mut legacy = snap_to_json(&snap);
3695 legacy.as_object_mut().unwrap().remove("source");
3696 let legacy = serde_json::to_vec(&legacy).unwrap();
3697 assert_eq!(
3698 parse_cache_at(&legacy, None, now()).unwrap().source,
3699 AntigravitySource::Local
3700 );
3701 }
3702}