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(
599 client: &reqwest::Client,
600 cache: &Cache,
601 oauth: Option<&cloud::OauthClient>,
602 endpoints: &cloud::Endpoints,
603 token: Result<StoredToken>,
604 renewal: Renewal<'_>,
605 now: DateTime<Utc>,
606) -> Result<AntigravitySnapshot> {
607 let mut token = token?;
608 let oauth_path = cloud::oauth_cache_path(cache);
609 let mut access =
610 match resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await {
611 Ok(access) => access,
612 Err(e) if oauth.is_none() => {
615 let Some(renewed) = renewal.renewed(cache, now).await else {
616 return Err(e);
617 };
618 token = renewed;
619 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
620 }
621 Err(e) => return Err(e),
622 };
623
624 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
625 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
626 access = if oauth.is_none() {
627 let Some(renewed) = renewal.renewed(cache, now).await else {
628 return Err(refresh_unconfigured());
629 };
630 token = renewed;
631 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?
632 } else {
633 refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?
634 };
635 cloud::fetch_quota(client, endpoints, &access.value)
636 .await
637 .map_err(|e| {
638 if is_auth_rejection(&e) {
639 session_expired()
640 } else {
641 e
642 }
643 })?
644 }
645 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
646 other => other?,
647 };
648
649 let plan = cloud::fetch_plan(client, endpoints, &access.value)
650 .await
651 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
652 let mut snap = parse_quota_summary("a, plan)?;
653 snap.account = remote_account(&token.fingerprint);
654 snap.source = AntigravitySource::Remote;
655 Ok(snap)
656}
657
658fn account_key(user_status: &serde_json::Value) -> String {
663 let email = user_status["userStatus"]["email"]
664 .as_str()
665 .filter(|s| !s.is_empty());
666 match email {
667 Some(e) => {
668 use std::hash::{Hash, Hasher};
669 let mut h = std::collections::hash_map::DefaultHasher::new();
670 e.hash(&mut h);
671 format!("acct:{:016x}", h.finish())
672 }
673 None => "acct:unknown".to_string(),
674 }
675}
676
677async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
681 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
682 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
685 .await
686 .ok()?;
687 let html = String::from_utf8_lossy(&bytes);
688 html.split("csrfToken\":\"")
689 .nth(1)
690 .and_then(|s| s.split('"').next())
691 .filter(|t| !t.is_empty())
692 .map(|t| t.to_string())
693}
694
695async fn post_rpc(
696 client: &reqwest::Client,
697 base: &str,
698 csrf: Option<&str>,
699 rpc: &str,
700) -> Result<serde_json::Value> {
701 let mut req = client
702 .post(format!("{base}/{rpc}"))
703 .header("Content-Type", "application/json")
704 .body("{}")
705 .timeout(HTTP_TIMEOUT);
706 if let Some(token) = csrf {
707 req = req.header("x-codeium-csrf-token", token);
708 }
709 let resp = req.send().await?;
710
711 let status = resp.status();
712 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
716 if !status.is_success() {
717 let body = String::from_utf8_lossy(&bytes).into_owned();
718 return Err(AppError::Http {
719 status: status.as_u16(),
720 body,
721 });
722 }
723
724 Ok(serde_json::from_slice(&bytes)?)
725}
726
727pub fn plan_from_status(v: &serde_json::Value) -> String {
728 v["userStatus"]["userTier"]["name"]
729 .as_str()
730 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
731 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
732 .filter(|s| !s.is_empty())
733 .unwrap_or(DEFAULT_PLAN)
734 .to_string()
735}
736
737fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
748 let id = if id.is_empty() { "<unnamed>" } else { id };
749 let group = if group_name.is_empty() {
750 String::new()
751 } else {
752 format!(" in {group_name:?}")
753 };
754 match window {
755 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
756 _ => format!("{id}{group}"),
757 }
758}
759
760fn no_usable_bucket(seen: &[String]) -> AppError {
766 if seen.is_empty() {
767 return AppError::Other(
768 "antigravity: quota summary has no buckets at all — the running product may \
769 not have a quota for this account yet"
770 .into(),
771 );
772 }
773 AppError::Other(format!(
774 "antigravity: quota summary has no bucket in a window we recognise (5h or \
775 weekly, Gemini or Claude/GPT); it offered: {}",
776 seen.join(", ")
777 ))
778}
779
780pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
781 let groups = v["response"]["groups"]
782 .as_array()
783 .or_else(|| v["groups"].as_array())
784 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
785
786 let mut gemini_5h = None;
787 let mut gemini_weekly = None;
788 let mut tp_5h = None;
789 let mut tp_weekly = None;
790 let mut seen: Vec<String> = Vec::new();
794
795 for group in groups {
796 let group_name = group["displayName"].as_str().unwrap_or_default();
797 let Some(buckets) = group["buckets"].as_array() else {
798 continue;
799 };
800 for bucket in buckets {
801 let id = bucket["bucketId"].as_str().unwrap_or_default();
802 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
803 let window = bucket["window"].as_str().unwrap_or_default();
804 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
805 true
806 } else if id.ends_with("5h") || window == "5h" {
807 false
808 } else {
809 continue;
812 };
813 let is_gemini = if id.starts_with("gemini") {
814 true
815 } else if id.starts_with("3p") {
816 false
817 } else if group_name.contains("Gemini") {
818 true
819 } else if group_name.contains("Claude") || group_name.contains("GPT") {
820 false
821 } else {
822 continue;
825 };
826
827 let (slot, slot_name) = match (is_gemini, is_weekly) {
828 (true, false) => (&mut gemini_5h, "Gemini 5h"),
829 (true, true) => (&mut gemini_weekly, "Gemini weekly"),
830 (false, false) => (&mut tp_5h, "Claude/GPT 5h"),
831 (false, true) => (&mut tp_weekly, "Claude/GPT weekly"),
832 };
833 let parsed = usage_window(bucket, is_weekly)?;
834 if slot.replace(parsed).is_some() {
835 return Err(AppError::Schema(format!(
836 "antigravity: duplicate {slot_name} bucket"
837 )));
838 }
839 }
840 }
841
842 if gemini_5h.is_none() && gemini_weekly.is_none() && tp_5h.is_none() && tp_weekly.is_none() {
846 return Err(no_usable_bucket(&seen));
847 }
848
849 Ok(AntigravitySnapshot {
850 plan,
851 account: String::new(),
854 source: AntigravitySource::Local,
855 session: gemini_5h,
856 weekly: gemini_weekly,
857 third_party_session: tp_5h,
858 third_party_weekly: tp_weekly,
859 })
860}
861
862fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
866 let remaining = bucket["remainingFraction"]
867 .as_f64()
868 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
869 .ok_or_else(|| {
870 AppError::Schema(format!(
871 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
872 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
873 ))
874 })?;
875 Ok(UsageWindow {
876 utilization_pct: pct_used(remaining),
877 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
878 window_duration: if is_weekly {
879 chrono::Duration::days(7)
880 } else {
881 chrono::Duration::hours(5)
882 },
883 })
884}
885
886fn pct_used(remaining_fraction: f64) -> i32 {
889 let used = (1.0 - remaining_fraction) * 100.0;
890 used.round() as i32
891}
892
893fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
894 match value {
895 serde_json::Value::Null => Ok(None),
896 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
897 .map(|dt| Some(dt.with_timezone(&Utc)))
898 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
899 _ => Err(AppError::Schema(format!(
900 "antigravity: {field} must be a timestamp or null"
901 ))),
902 }
903}
904
905fn candidate_bases() -> Vec<Candidate> {
911 candidate_bases_with(
912 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
913 discover_ls_ports(),
914 )
915}
916
917fn candidate_bases_with(
921 override_addr: Option<&str>,
922 discovered: Vec<(u32, u16)>,
923) -> Vec<Candidate> {
924 let mut candidates = Vec::new();
925 if let Some(base) = override_addr.and_then(normalize_base) {
926 let pid = base_port(&base).and_then(|port| unique_owner(&discovered, port));
930 candidates.push(Candidate { base, pid });
931 }
932
933 for (pid, port) in discovered {
939 let base = format!("http://127.0.0.1:{port}");
940 if !candidates.iter().any(|c| c.base == base) {
941 candidates.push(Candidate {
942 base,
943 pid: Some(pid),
944 });
945 }
946 }
947
948 candidates
949}
950
951fn base_port(base: &str) -> Option<u16> {
953 let (_, authority) = base.split_once("://")?;
954 authority.rsplit_once(':')?.1.parse().ok()
955}
956
957fn unique_owner(discovered: &[(u32, u16)], port: u16) -> Option<u32> {
960 let mut owners = discovered
961 .iter()
962 .filter(|(_, p)| *p == port)
963 .map(|(pid, _)| *pid);
964 let first = owners.next()?;
965 owners.all(|pid| pid == first).then_some(first)
966}
967
968fn normalize_base(addr: &str) -> Option<String> {
977 let trimmed = addr.trim();
978 let (scheme, authority) = match trimmed.split_once("://") {
979 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
980 _ => ("http", trimmed),
981 };
982 let authority = authority.trim_end_matches('/');
983 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
984}
985
986fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
993 let comm = comm.trim().to_lowercase();
994 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
995 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
996 return true;
997 }
998 exe.is_some_and(|p| {
999 let p = p.to_lowercase().replace('\\', "/");
1000 let p = p.strip_suffix(".exe").unwrap_or(&p);
1001 p.contains("antigravity") || p.ends_with("/agy")
1002 })
1003}
1004
1005#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
1037fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<(u32, u16)> {
1038 let groups: Vec<(u32, Vec<u16>)> = per_pid
1039 .into_iter()
1040 .map(|(pid, mut group)| {
1041 group.sort_unstable_by(|a, b| b.cmp(a));
1042 group.dedup();
1047 (pid, group)
1048 })
1049 .collect();
1050
1051 let mut ports: Vec<(u32, u16)> = Vec::new();
1052 for rank in 0..groups.iter().map(|(_, g)| g.len()).max().unwrap_or(0) {
1053 for (pid, port) in groups
1054 .iter()
1055 .filter_map(|(pid, group)| group.get(rank).map(|port| (*pid, *port)))
1056 {
1057 if !ports.iter().any(|(_, p)| *p == port) {
1060 ports.push((pid, port));
1061 }
1062 }
1063 }
1064 ports
1065}
1066
1067#[cfg(target_os = "linux")]
1075pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1076 use std::collections::{BTreeMap, HashMap};
1077
1078 let mut owners: HashMap<u64, u32> = HashMap::new();
1081 let Ok(entries) = std::fs::read_dir("/proc") else {
1082 return Vec::new();
1083 };
1084 for entry in entries.flatten() {
1085 let pid_dir = entry.path();
1086 let Some(pid) = pid_dir
1087 .file_name()
1088 .and_then(|name| name.to_str())
1089 .and_then(|name| name.parse::<u32>().ok())
1090 else {
1091 continue;
1092 };
1093 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
1094 continue;
1095 };
1096 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
1097 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
1098 continue;
1099 }
1100 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
1101 continue;
1102 };
1103 for fd in fds.flatten() {
1104 let Ok(target) = std::fs::read_link(fd.path()) else {
1105 continue;
1106 };
1107 if let Some(ino) = target
1108 .to_str()
1109 .and_then(|s| s.strip_prefix("socket:["))
1110 .and_then(|s| s.strip_suffix(']'))
1111 .and_then(|s| s.parse::<u64>().ok())
1112 {
1113 owners.insert(ino, pid);
1114 }
1115 }
1116 }
1117
1118 if owners.is_empty() {
1119 return Vec::new();
1120 }
1121
1122 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
1123 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
1124 let Ok(contents) = std::fs::read_to_string(table) else {
1125 continue;
1126 };
1127 for line in contents.lines().skip(1) {
1128 if let Some((port, ino)) = parse_proc_net_line(line)
1129 && let Some(&pid) = owners.get(&ino)
1130 {
1131 per_pid.entry(pid).or_default().push(port);
1132 }
1133 }
1134 }
1135 probe_order(per_pid)
1136}
1137
1138#[cfg(target_os = "macos")]
1145pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1146 let Ok(output) = std::process::Command::new("lsof")
1147 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
1148 .output()
1149 else {
1150 return Vec::new();
1151 };
1152 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
1156}
1157
1158#[cfg(any(test, target_os = "macos"))]
1163fn parse_lsof_pcn(output: &str) -> Vec<(u32, u16)> {
1164 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1165 let mut pid = None;
1168 let mut owner = None;
1169 for line in output.lines() {
1170 let Some(rest) = line.get(1..) else { continue };
1171 match line.as_bytes().first() {
1172 Some(b'p') => {
1173 pid = rest.parse::<u32>().ok();
1174 owner = None;
1175 }
1176 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
1177 Some(b'n') => {
1178 if let Some(pid) = owner
1179 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
1180 {
1181 per_pid.entry(pid).or_default().push(port);
1182 }
1183 }
1184 _ => {}
1185 }
1186 }
1187 probe_order(per_pid)
1188}
1189
1190#[cfg(any(test, target_os = "windows"))]
1191#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1192struct WindowsTcpRow {
1193 local_addr: [u8; 4],
1194 local_port: u32,
1195 pid: u32,
1196}
1197
1198#[cfg(any(test, target_os = "windows"))]
1199fn decode_windows_process_name(raw: &[u16]) -> String {
1200 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
1201 String::from_utf16_lossy(&raw[..end])
1202}
1203
1204#[cfg(any(test, target_os = "windows"))]
1205fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
1206 processes
1207 .iter()
1208 .filter(|(_, name)| is_antigravity_process(name, None))
1209 .map(|(pid, _)| *pid)
1210 .collect()
1211}
1212
1213#[cfg(any(test, target_os = "windows"))]
1216fn matching_windows_ports(
1217 pids: &std::collections::HashSet<u32>,
1218 rows: &[WindowsTcpRow],
1219) -> Vec<(u32, u16)> {
1220 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1221 for row in rows {
1222 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1223 continue;
1224 }
1225 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1226 if port != 0 {
1227 per_pid.entry(row.pid).or_default().push(port);
1228 }
1229 }
1230 probe_order(per_pid)
1231}
1232
1233#[cfg(any(test, target_os = "windows"))]
1234fn checked_windows_row_count(
1235 buffer_len: usize,
1236 rows_offset: usize,
1237 row_size: usize,
1238 declared: usize,
1239) -> Option<usize> {
1240 let rows_len = row_size.checked_mul(declared)?;
1241 let end = rows_offset.checked_add(rows_len)?;
1242 (row_size != 0 && end <= buffer_len).then_some(declared)
1243}
1244
1245#[cfg(target_os = "windows")]
1246struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1247
1248#[cfg(target_os = "windows")]
1249impl Drop for WindowsHandle {
1250 fn drop(&mut self) {
1251 unsafe {
1252 windows_sys::Win32::Foundation::CloseHandle(self.0);
1253 }
1254 }
1255}
1256
1257#[cfg(target_os = "windows")]
1258fn windows_processes() -> Vec<(u32, String)> {
1259 use std::mem::size_of;
1260 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1261 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1262 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1263 TH32CS_SNAPPROCESS,
1264 };
1265
1266 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1267 if handle == INVALID_HANDLE_VALUE {
1268 return Vec::new();
1269 }
1270 let snapshot = WindowsHandle(handle);
1271 let mut entry = PROCESSENTRY32W::default();
1272 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1273 return Vec::new();
1274 };
1275 entry.dwSize = entry_size;
1276 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1277 return Vec::new();
1278 }
1279
1280 let mut processes = Vec::new();
1281 loop {
1282 processes.push((
1283 entry.th32ProcessID,
1284 decode_windows_process_name(&entry.szExeFile),
1285 ));
1286 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1287 break;
1288 }
1289 }
1290 processes
1291}
1292
1293#[cfg(target_os = "windows")]
1294fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1295 use std::mem::{offset_of, size_of, size_of_val};
1296 use windows_sys::Win32::NetworkManagement::IpHelper::{
1297 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1298 };
1299
1300 let available = used_bytes.min(size_of_val(buffer));
1301 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1302 if available < size_of::<u32>() || available < rows_offset {
1303 return Vec::new();
1304 }
1305 let base = buffer.as_ptr().cast::<u8>();
1306 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1307 if checked_windows_row_count(
1308 available,
1309 rows_offset,
1310 size_of::<MIB_TCPROW_OWNER_PID>(),
1311 declared,
1312 )
1313 .is_none()
1314 {
1315 return Vec::new();
1316 }
1317
1318 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1319 (0..declared)
1320 .map(|index| unsafe { rows.add(index).read_unaligned() })
1321 .map(|row| WindowsTcpRow {
1322 local_addr: row.dwLocalAddr.to_ne_bytes(),
1323 local_port: row.dwLocalPort,
1324 pid: row.dwOwningPid,
1325 })
1326 .collect()
1327}
1328
1329#[cfg(target_os = "windows")]
1330fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1331 use std::mem::size_of;
1332 use std::ptr::null_mut;
1333 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1334 use windows_sys::Win32::NetworkManagement::IpHelper::{
1335 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1336 };
1337 use windows_sys::Win32::Networking::WinSock::AF_INET;
1338
1339 let mut size = 0u32;
1340 let status = unsafe {
1341 GetExtendedTcpTable(
1342 null_mut(),
1343 &mut size,
1344 0,
1345 u32::from(AF_INET),
1346 TCP_TABLE_OWNER_PID_LISTENER,
1347 0,
1348 )
1349 };
1350 if status != ERROR_INSUFFICIENT_BUFFER {
1351 return Vec::new();
1352 }
1353
1354 for _ in 0..3 {
1355 let Some(words) = (size as usize)
1356 .checked_add(size_of::<u32>() - 1)
1357 .map(|bytes| bytes / size_of::<u32>())
1358 else {
1359 return Vec::new();
1360 };
1361 if words == 0 {
1362 return Vec::new();
1363 }
1364 let mut buffer = Vec::<u32>::new();
1365 if buffer.try_reserve_exact(words).is_err() {
1366 return Vec::new();
1367 }
1368 buffer.resize(words, 0);
1369 let mut used = size;
1370 let status = unsafe {
1371 GetExtendedTcpTable(
1372 buffer.as_mut_ptr().cast(),
1373 &mut used,
1374 0,
1375 u32::from(AF_INET),
1376 TCP_TABLE_OWNER_PID_LISTENER,
1377 0,
1378 )
1379 };
1380 if status == ERROR_INSUFFICIENT_BUFFER {
1381 size = used;
1382 continue;
1383 }
1384 if status != 0 {
1385 return Vec::new();
1386 }
1387 return parse_windows_tcp_rows(&buffer, used as usize);
1388 }
1389 Vec::new()
1390}
1391
1392#[cfg(target_os = "windows")]
1393pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1394 let pids = matching_windows_process_ids(&windows_processes());
1395 if pids.is_empty() {
1396 return Vec::new();
1397 }
1398 matching_windows_ports(&pids, &windows_tcp_rows())
1399}
1400
1401#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1402pub(crate) fn discover_ls_ports() -> Vec<(u32, u16)> {
1403 Vec::new()
1404}
1405
1406#[cfg(target_os = "linux")]
1409fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1410 let cols: Vec<&str> = line.split_whitespace().collect();
1411 if cols.len() < 10 {
1412 return None;
1413 }
1414 if cols[3] != "0A" {
1416 return None;
1417 }
1418 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1419 let inode = cols[9].parse::<u64>().ok()?;
1420 Some((port, inode))
1421}
1422
1423fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1428 crate::outcome::fallback(cache, None, original, |bytes| {
1429 parse_cache_at(bytes, None, now)
1430 })
1431}
1432
1433fn fallback_with_error(
1438 cache: &Cache,
1439 last_error: Option<(u16, String)>,
1440 reason: AppError,
1441 now: DateTime<Utc>,
1442) -> Result<FetchOutcome> {
1443 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1444 parse_cache_at(bytes, None, now)
1445 })
1446}
1447
1448fn reuse_cache(
1449 bytes: Vec<u8>,
1450 cache: &Cache,
1451 stale: bool,
1452 account: Option<&str>,
1453 now: DateTime<Utc>,
1454) -> Result<FetchOutcome> {
1455 let snap = parse_cache_at(&bytes, account, now)?;
1456 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1457}
1458
1459pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1465 parse_cache_at(bytes, account, Utc::now())
1466}
1467
1468fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<&'static str> {
1478 [
1479 ("Gemini 5h", snap.session.as_ref()),
1480 ("Gemini weekly", snap.weekly.as_ref()),
1481 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1482 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1483 ]
1484 .into_iter()
1485 .find(|(_, w)| w.and_then(|w| w.resets_at).is_some_and(|r| r <= now))
1486 .map(|(name, _)| name)
1487}
1488
1489pub fn parse_cache_at(
1490 bytes: &[u8],
1491 account: Option<&str>,
1492 now: DateTime<Utc>,
1493) -> Result<AntigravitySnapshot> {
1494 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1495
1496 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1497 if let Some(expected) = account
1498 && cached_account != Some(expected)
1499 {
1500 return Err(AppError::Schema(
1501 "antigravity cache belongs to a different account; refetching".into(),
1502 ));
1503 }
1504
1505 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1515 match v.get(pct_key) {
1516 None => Err(AppError::Schema(format!(
1517 "antigravity: cached payload is missing {pct_key}"
1518 ))),
1519 Some(serde_json::Value::Null) => Ok(None),
1520 Some(value) => value
1521 .as_i64()
1522 .filter(|pct| (0..=100).contains(pct))
1523 .map(|pct| Some(pct as i32))
1524 .ok_or_else(|| {
1525 AppError::Schema(format!(
1526 "antigravity: cached {pct_key} must be an integer in 0..=100"
1527 ))
1528 }),
1529 }
1530 };
1531
1532 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1533 let Some(pct) = cached_pct(pct_key)? else {
1534 return Ok(None);
1535 };
1536 Ok::<_, AppError>(Some(UsageWindow {
1537 utilization_pct: pct,
1538 resets_at: parse_reset(&v[reset_key], reset_key)?,
1539 window_duration: if weekly {
1540 chrono::Duration::days(7)
1541 } else {
1542 chrono::Duration::hours(5)
1543 },
1544 }))
1545 };
1546
1547 let snap = AntigravitySnapshot {
1548 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1549 account: cached_account.unwrap_or_default().to_string(),
1550 source: v["source"]
1552 .as_str()
1553 .and_then(AntigravitySource::parse)
1554 .unwrap_or_default(),
1555 session: optional("session_pct", "session_reset", false)?,
1556 weekly: optional("weekly_pct", "weekly_reset", true)?,
1557 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1558 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1559 };
1560
1561 if snap.session.is_none()
1564 && snap.weekly.is_none()
1565 && snap.third_party_session.is_none()
1566 && snap.third_party_weekly.is_none()
1567 {
1568 return Err(AppError::Schema(
1569 "antigravity cache holds no usable window; refetching".into(),
1570 ));
1571 }
1572
1573 if let Some(window) = expired_window(&snap, now) {
1574 return Err(AppError::Schema(format!(
1575 "antigravity cache is past its {window} reset; refetching"
1576 )));
1577 }
1578 Ok(snap)
1579}
1580
1581pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1582 serde_json::json!({
1583 "plan": snap.plan,
1584 "account": snap.account,
1585 "source": snap.source.as_str(),
1586 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
1587 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1588 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
1589 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1590 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
1591 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1592 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
1593 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1594 })
1595}
1596
1597#[cfg(test)]
1598mod tests {
1599 use super::*;
1600
1601 const QUOTA_JSON: &str = r#"{
1605 "response": {
1606 "groups": [
1607 {
1608 "displayName": "Gemini Models",
1609 "buckets": [
1610 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
1611 "window": "weekly", "remainingFraction": 0.9191212,
1612 "resetTime": "2026-07-28T17:39:58Z"},
1613 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
1614 "window": "5h", "remainingFraction": 0.5672253,
1615 "resetTime": "2026-07-22T17:47:00Z"}
1616 ]
1617 },
1618 {
1619 "displayName": "Claude and GPT models",
1620 "buckets": [
1621 {"bucketId": "3p-weekly", "window": "weekly",
1622 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
1623 {"bucketId": "3p-5h", "window": "5h",
1624 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
1625 ]
1626 }
1627 ]
1628 }
1629 }"#;
1630
1631 fn now() -> DateTime<Utc> {
1634 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
1635 .unwrap()
1636 .with_timezone(&Utc)
1637 }
1638
1639 fn parsed() -> AntigravitySnapshot {
1640 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1641 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
1642 }
1643
1644 #[test]
1645 fn quota_summary_maps_four_distinct_windows() {
1646 let snap = parsed();
1647 assert_eq!(snap.plan, "Google AI Pro");
1648 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
1650 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
1651 assert_eq!(
1652 snap.third_party_session.as_ref().unwrap().utilization_pct,
1653 75
1654 );
1655 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
1656 }
1657
1658 #[test]
1659 fn each_window_keeps_its_own_reset_time() {
1660 let snap = parsed();
1661 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
1662 assert_eq!(
1663 snap.session.as_ref().unwrap().resets_at,
1664 at("2026-07-22T17:47:00Z")
1665 );
1666 assert_eq!(
1667 snap.weekly.as_ref().unwrap().resets_at,
1668 at("2026-07-28T17:39:58Z")
1669 );
1670 assert_eq!(
1671 snap.third_party_weekly.as_ref().unwrap().resets_at,
1672 at("2026-07-29T12:47:00Z")
1673 );
1674 assert_ne!(
1676 snap.session.as_ref().unwrap().resets_at,
1677 snap.weekly.as_ref().unwrap().resets_at
1678 );
1679 }
1680
1681 #[test]
1682 fn window_durations_match_their_bucket() {
1683 let snap = parsed();
1684 assert_eq!(
1685 snap.session.as_ref().unwrap().window_duration,
1686 chrono::Duration::hours(5)
1687 );
1688 assert_eq!(
1689 snap.weekly.as_ref().unwrap().window_duration,
1690 chrono::Duration::days(7)
1691 );
1692 assert_eq!(
1693 snap.third_party_weekly.as_ref().unwrap().window_duration,
1694 chrono::Duration::days(7)
1695 );
1696 }
1697
1698 #[test]
1699 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
1700 let v: serde_json::Value = serde_json::from_str(
1701 r#"{"response":{"groups":[
1702 {"displayName":"Gemini Models","buckets":[
1703 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
1704 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
1705 {"displayName":"Claude and GPT models","buckets":[
1706 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
1707 ]}}"#,
1708 )
1709 .unwrap();
1710 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
1711 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
1712 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
1713 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
1714 assert!(snap.third_party_weekly.is_none());
1715 }
1716
1717 #[test]
1718 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
1719 let duplicate: serde_json::Value = serde_json::from_str(
1720 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
1721 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1722 {"bucketId":"gemini-5h-copy","window":"5h","remainingFraction":0.1},
1723 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
1724 ]}]}}"#,
1725 )
1726 .unwrap();
1727 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
1728 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
1729
1730 let unrelated: serde_json::Value = serde_json::from_str(
1733 r#"{"response":{"groups":[
1734 {"displayName":"Gemini Models","buckets":[
1735 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1736 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
1737 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
1738 ]},
1739 {"displayName":"Future Models","buckets":[
1740 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
1741 ]}
1742 ]}}"#,
1743 )
1744 .unwrap();
1745 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
1746 assert!(snap.third_party_session.is_none());
1747 assert!(snap.third_party_weekly.is_none());
1748 }
1749
1750 #[test]
1753 fn a_bucket_without_a_usable_fraction_is_rejected() {
1754 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
1755 let v: serde_json::Value = serde_json::from_str(&format!(
1756 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
1757 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
1758 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
1759 ))
1760 .unwrap();
1761 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1762 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
1763 }
1764 }
1765
1766 #[test]
1767 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
1768 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
1769 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1770 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
1771 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1772 assert!(err.to_string().contains("resetTime"), "{err}");
1773 }
1774 }
1775
1776 #[test]
1780 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
1781 let mut snap = parsed();
1782 snap.session = None;
1783 snap.third_party_session = None;
1784
1785 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1786 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
1787
1788 assert!(back.session.is_none());
1789 assert_eq!(
1790 back.weekly.as_ref().unwrap().utilization_pct,
1791 snap.weekly.as_ref().unwrap().utilization_pct
1792 );
1793 }
1794
1795 #[test]
1800 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
1801 let summary = serde_json::json!({
1802 "groups": [
1803 {
1804 "displayName": "Gemini Models",
1805 "buckets": [{
1806 "bucketId": "gemini-weekly", "window": "weekly",
1807 "remainingFraction": 0.42,
1808 }],
1809 },
1810 {
1811 "displayName": "Claude and GPT models",
1812 "buckets": [{
1813 "bucketId": "3p-weekly", "window": "weekly",
1814 "remainingFraction": 0.9,
1815 }],
1816 },
1817 ],
1818 });
1819
1820 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
1821
1822 assert!(snap.session.is_none(), "no 5h bucket arrived");
1823 assert!(snap.third_party_session.is_none());
1824 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
1825 assert_eq!(
1826 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
1827 10
1828 );
1829 }
1830
1831 #[test]
1834 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
1835 let summary = serde_json::json!({
1836 "groups": [{
1837 "displayName": "Gemini Models",
1838 "buckets": [{
1839 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
1840 }],
1841 }],
1842 });
1843
1844 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
1845
1846 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
1847 assert!(snap.weekly.is_none());
1848 }
1849
1850 #[test]
1853 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
1854 let summary = serde_json::json!({
1855 "groups": [{
1856 "displayName": "Gemini Models",
1857 "buckets": [{
1858 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
1859 }],
1860 }],
1861 });
1862
1863 let rendered = parse_quota_summary(&summary, "Pro".into())
1864 .expect_err("an unrecognised cadence alone is not a snapshot")
1865 .to_string();
1866
1867 assert!(
1868 rendered.contains("no bucket in a window we recognise"),
1869 "{rendered}"
1870 );
1871 assert!(rendered.contains("gemini-daily"), "{rendered}");
1872 assert!(rendered.contains("window daily"), "{rendered}");
1873 }
1874
1875 #[test]
1878 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
1879 let summary = serde_json::json!({
1880 "groups": [{"displayName": "Gemini", "buckets": []}],
1881 });
1882
1883 let rendered = parse_quota_summary(&summary, "Pro".into())
1884 .expect_err("no buckets is an error")
1885 .to_string();
1886
1887 assert!(rendered.contains("no buckets at all"), "{rendered}");
1888 assert!(!rendered.contains("it offered:"), "{rendered}");
1889 }
1890
1891 #[test]
1894 fn buckets_without_an_id_are_still_named_in_the_error() {
1895 let summary = serde_json::json!({
1896 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
1897 });
1898
1899 let rendered = parse_quota_summary(&summary, "Pro".into())
1900 .expect_err("an unusable summary is an error")
1901 .to_string();
1902
1903 assert!(rendered.contains("<unnamed>"), "{rendered}");
1904 }
1905
1906 #[test]
1907 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
1908 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
1909 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
1910 }
1911
1912 #[test]
1913 fn cache_round_trip_preserves_every_window() {
1914 let snap = parsed();
1915 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1916 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1917 }
1918
1919 #[test]
1923 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
1924 let full = snap_to_json(&parsed());
1925 for missing in ["session_pct", "weekly_pct"] {
1926 let mut v = full.clone();
1927 v.as_object_mut().unwrap().remove(missing);
1928 let bytes = serde_json::to_vec(&v).unwrap();
1929 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1930 assert!(err.to_string().contains(missing), "{missing}: {err}");
1931 }
1932 assert!(parse_cache_at(b"{}", None, now()).is_err());
1934 }
1935
1936 #[test]
1937 fn cached_percentages_are_range_checked_before_narrowing() {
1938 let full = snap_to_json(&parsed());
1939 for (key, bad) in [
1940 ("session_pct", serde_json::json!(-1)),
1941 ("weekly_pct", serde_json::json!(101)),
1942 ("session_pct", serde_json::json!(i64::MAX)),
1943 ("tp_session_pct", serde_json::json!("75")),
1944 ] {
1945 let mut v = full.clone();
1946 v[key] = bad;
1947 let bytes = serde_json::to_vec(&v).unwrap();
1948 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1949 assert!(err.to_string().contains(key), "{key}: {err}");
1950 }
1951 }
1952
1953 #[test]
1954 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
1955 let mut v = snap_to_json(&parsed());
1956 v["session_reset"] = serde_json::json!("not-a-time");
1957 let bytes = serde_json::to_vec(&v).unwrap();
1958 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1959 assert!(err.to_string().contains("session_reset"), "{err}");
1960 }
1961
1962 #[test]
1967 fn a_cache_past_its_reset_is_refused() {
1968 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
1969 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1970
1971 assert!(parse_cache_at(&bytes, None, now()).is_ok());
1973 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
1975
1976 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
1978 assert!(err.to_string().contains("5h"), "{err}");
1979
1980 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
1982 }
1983
1984 #[test]
1987 fn expiry_names_the_window_that_rolled_over() {
1988 let mut snap = parsed();
1989 snap.session.as_mut().unwrap().resets_at = None;
1991 snap.third_party_session = None;
1992 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1993 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1994
1995 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
1997
1998 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
2000 assert!(err.to_string().contains("Gemini weekly"), "{err}");
2001 }
2002
2003 #[test]
2005 fn a_window_without_a_reset_never_expires() {
2006 let mut snap = parsed();
2007 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
2008 w.resets_at = None;
2009 }
2010 snap.third_party_session = None;
2011 snap.third_party_weekly = None;
2012 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2013 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
2014 .unwrap()
2015 .with_timezone(&Utc);
2016 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
2017 }
2018
2019 #[test]
2021 fn a_cache_from_another_account_is_rejected() {
2022 let mut snap = parsed();
2023 snap.account = "acct:aaaa".into();
2024 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2025
2026 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
2027 assert_eq!(
2028 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
2029 snap
2030 );
2031
2032 let mut legacy = snap_to_json(&snap);
2034 legacy.as_object_mut().unwrap().remove("account");
2035 let legacy = serde_json::to_vec(&legacy).unwrap();
2036 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
2037 }
2038
2039 #[test]
2042 fn an_unverifiable_cache_is_served_rather_than_discarded() {
2043 let mut snap = parsed();
2044 snap.account = "acct:aaaa".into();
2045 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2046 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2047 }
2048
2049 #[test]
2050 fn account_key_fingerprints_rather_than_storing_the_address() {
2051 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
2052 let a = with("someone@example.com");
2053 assert!(!a.contains("someone"), "{a}");
2054 assert!(!a.contains('@'), "{a}");
2055 assert_eq!(a, with("someone@example.com"), "must be stable");
2056 assert_ne!(a, with("other@example.com"));
2057 let unknown = account_key(&serde_json::json!({}));
2059 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
2060 assert_ne!(unknown, a);
2061 }
2062
2063 #[test]
2066 fn absent_third_party_windows_are_not_treated_as_corruption() {
2067 let mut snap = parsed();
2068 snap.third_party_session = None;
2069 snap.third_party_weekly = None;
2070 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2071 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2072 }
2073
2074 #[test]
2075 fn cache_round_trip_preserves_absent_third_party_windows() {
2076 let mut snap = parsed();
2077 snap.third_party_session = None;
2078 snap.third_party_weekly = None;
2079 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2080 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
2081 }
2082
2083 #[test]
2084 fn pct_used_inverts_valid_fractions() {
2085 assert_eq!(pct_used(1.0), 0);
2086 assert_eq!(pct_used(0.0), 100);
2087 assert_eq!(pct_used(0.5), 50);
2088 }
2089
2090 #[test]
2091 fn plan_falls_back_through_the_status_payload() {
2092 let tier: serde_json::Value =
2093 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
2094 .unwrap();
2095 assert_eq!(plan_from_status(&tier), "Google AI Pro");
2096
2097 let plan_only: serde_json::Value = serde_json::from_str(
2098 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
2099 )
2100 .unwrap();
2101 assert_eq!(plan_from_status(&plan_only), "Pro");
2102
2103 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
2104 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
2105 }
2106
2107 #[cfg(target_os = "linux")]
2108 #[test]
2109 fn proc_net_parser_keeps_only_listening_rows() {
2110 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
2111 00:00000000 00000000 1000 0 123456 1 0000 100 0";
2112 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
2113
2114 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
2115 00:00000000 00000000 1000 0 123457 1 0000 100 0";
2116 assert_eq!(parse_proc_net_line(established), None);
2117
2118 assert_eq!(parse_proc_net_line("garbage"), None);
2119 }
2120
2121 #[test]
2122 fn explicit_address_comes_first_and_gets_a_scheme() {
2123 let candidate = |base: &str, pid| Candidate {
2124 base: base.to_string(),
2125 pid,
2126 };
2127 assert_eq!(
2128 candidate_bases_with(Some("127.0.0.1:1234"), vec![(10, 5678)]),
2129 vec![
2130 candidate("http://127.0.0.1:1234", None),
2131 candidate("http://127.0.0.1:5678", Some(10)),
2132 ]
2133 );
2134 assert_eq!(
2136 candidate_bases_with(Some("127.0.0.1:1234/"), vec![(10, 5678)]),
2137 vec![
2138 candidate("http://127.0.0.1:1234", None),
2139 candidate("http://127.0.0.1:5678", Some(10)),
2140 ]
2141 );
2142 assert_eq!(
2145 candidate_bases_with(Some("127.0.0.1:5678"), vec![(10, 5678)]),
2146 vec![candidate("http://127.0.0.1:5678", Some(10))]
2147 );
2148 assert_eq!(
2150 candidate_bases_with(None, vec![(10, 5678), (10, 5678)]),
2151 vec![candidate("http://127.0.0.1:5678", Some(10))]
2152 );
2153 assert_eq!(
2155 candidate_bases_with(Some("https://host:9"), vec![]),
2156 vec![candidate("https://host:9", None)]
2157 );
2158 }
2159
2160 fn http(status: u16) -> AppError {
2161 AppError::Http {
2162 status,
2163 body: String::new(),
2164 }
2165 }
2166
2167 fn missing_csrf() -> AppError {
2168 missing_csrf_with("missing CSRF token")
2169 }
2170
2171 fn missing_csrf_with(message: &str) -> AppError {
2172 AppError::Http {
2173 status: 401,
2174 body: serde_json::json!({
2175 "code": "unauthenticated",
2176 "message": message,
2177 })
2178 .to_string(),
2179 }
2180 }
2181
2182 #[test]
2183 fn only_agys_exact_missing_csrf_response_enables_remote_fallback() {
2184 assert!(is_missing_csrf(&missing_csrf()));
2185 assert!(is_missing_csrf(&missing_csrf_with(
2186 "unauthenticated: missing CSRF token"
2187 )));
2188 assert!(is_missing_csrf(&missing_csrf_with(
2189 " unauthenticated: missing CSRF token "
2190 )));
2191 assert!(!is_missing_csrf(&missing_csrf_with(
2192 "unauthenticated: invalid token"
2193 )));
2194 assert!(!is_missing_csrf(&AppError::Http {
2195 status: 403,
2196 body: r#"{"code":"unauthenticated","message":"missing CSRF token"}"#.into(),
2197 }));
2198 assert!(!is_missing_csrf(&AppError::Http {
2199 status: 401,
2200 body: r#"{"code":"other","message":"missing CSRF token"}"#.into(),
2201 }));
2202 assert!(!is_missing_csrf(&AppError::Http {
2203 status: 401,
2204 body: "prefix: missing CSRF token".into(),
2205 }));
2206 }
2207
2208 #[test]
2209 fn a_real_auth_failure_outranks_agys_missing_csrf_response() {
2210 for errors in [
2211 vec![missing_csrf(), http(403)],
2212 vec![http(401), missing_csrf()],
2213 ] {
2214 let err = select_probe_error(errors);
2215 assert!(is_actionable(&err), "{err}");
2216 assert!(!is_missing_csrf(&err), "{err}");
2217 }
2218 }
2219
2220 #[test]
2221 fn several_agy_sessions_still_select_the_remote_fallback_reason() {
2222 let err = select_probe_error(vec![missing_csrf(), missing_csrf()]);
2223 assert!(is_missing_csrf(&err), "{err}");
2224 assert!(remote_fallback_reason(&err).is_some(), "{err}");
2225 }
2226
2227 #[test]
2231 fn an_auth_failure_outranks_later_transport_noise() {
2232 let err = select_probe_error(vec![
2233 http(401),
2234 AppError::Transport("connection refused".into()),
2235 ]);
2236 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2237
2238 let err = select_probe_error(vec![
2239 AppError::Transport("connection refused".into()),
2240 http(403),
2241 ]);
2242 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
2243 }
2244
2245 #[test]
2248 fn the_first_auth_failure_wins() {
2249 let err = select_probe_error(vec![http(401), http(403)]);
2250 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2251 }
2252
2253 #[test]
2257 fn without_an_auth_failure_the_last_error_stands() {
2258 let err = select_probe_error(vec![
2259 AppError::Transport("first".into()),
2260 http(500),
2261 AppError::Transport("last".into()),
2262 ]);
2263 assert!(
2264 matches!(&err, AppError::Transport(m) if m == "last"),
2265 "{err}"
2266 );
2267 assert!(err.is_transient());
2268 }
2269
2270 #[test]
2273 fn a_server_error_is_not_treated_as_actionable() {
2274 let err = select_probe_error(vec![http(500), http(401)]);
2275 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2276 }
2277
2278 #[test]
2279 fn no_candidates_at_all_yields_a_generic_error() {
2280 let err = select_probe_error(Vec::new());
2281 assert!(
2282 err.to_string().contains("no local server answered"),
2283 "{err}"
2284 );
2285 }
2286
2287 fn tls_echo() -> AppError {
2291 AppError::Http {
2292 status: 400,
2293 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2294 }
2295 }
2296
2297 #[test]
2302 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2303 let err = select_probe_error(vec![
2304 AppError::Http {
2305 status: 500,
2306 body: "GetUserStatus: internal".into(),
2307 },
2308 tls_echo(),
2309 ]);
2310 assert!(
2311 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2312 "{err}"
2313 );
2314 }
2315
2316 #[test]
2321 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2322 let err = select_probe_error(vec![
2323 AppError::Transport("connection refused".into()),
2324 tls_echo(),
2325 ]);
2326 assert!(
2327 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2328 "{err}"
2329 );
2330 assert!(
2331 err.is_transient(),
2332 "the echo must not make the run non-transient: {err}"
2333 );
2334 }
2335
2336 #[test]
2339 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2340 let err = select_probe_error(vec![tls_echo()]);
2341 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2342 }
2343
2344 #[test]
2347 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2348 let err = select_probe_error(vec![
2349 AppError::Http {
2350 status: 400,
2351 body: "unknown method GetUserStatus".into(),
2352 },
2353 tls_echo(),
2354 ]);
2355 assert!(
2356 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2357 "{err}"
2358 );
2359 }
2360
2361 #[test]
2363 fn an_auth_failure_still_outranks_a_tls_echo() {
2364 let err = select_probe_error(vec![tls_echo(), http(401)]);
2365 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2366 }
2367
2368 #[test]
2369 fn every_discovered_port_is_probed_in_order() {
2370 assert_eq!(
2371 candidate_bases_with(None, vec![(10, 33875), (10, 37435)]),
2372 vec![
2373 Candidate {
2374 base: "http://127.0.0.1:33875".to_string(),
2375 pid: Some(10),
2376 },
2377 Candidate {
2378 base: "http://127.0.0.1:37435".to_string(),
2379 pid: Some(10),
2380 },
2381 ]
2382 );
2383 }
2384
2385 #[test]
2390 fn empty_discovery_yields_no_candidates() {
2391 assert!(candidate_bases_with(None, vec![]).is_empty());
2392 assert!(candidate_bases_with(Some(""), vec![]).is_empty());
2393 }
2394
2395 #[test]
2396 fn every_antigravity_product_is_recognised() {
2397 assert!(is_antigravity_process(
2399 "language_server\n",
2400 Some("/opt/antigravity/resources/bin/language_server")
2401 ));
2402 assert!(is_antigravity_process(
2404 "agy\n",
2405 Some("/home/u/.local/bin/agy")
2406 ));
2407 assert!(is_antigravity_process(
2409 "node",
2410 Some("/opt/antigravity/bin/helper")
2411 ));
2412 assert!(is_antigravity_process("antigravity", None));
2413 assert!(is_antigravity_process("agy.exe", None));
2414 assert!(is_antigravity_process("Antigravity.exe", None));
2415 assert!(is_antigravity_process("language_server.exe", None));
2416 assert!(is_antigravity_process(
2417 "language_server_windows_x64.exe",
2418 None
2419 ));
2420 assert!(is_antigravity_process(
2421 "node.exe",
2422 Some(r"C:\Users\u\AppData\Local\agy.exe")
2423 ));
2424 }
2425
2426 #[test]
2427 fn unrelated_processes_are_not_probed() {
2428 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2429 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2430 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2432 assert!(!is_antigravity_process("legacy.exe", None));
2433 assert!(!is_antigravity_process("not-agy.exe", None));
2434 assert!(!is_antigravity_process("", None));
2435 }
2436
2437 #[test]
2438 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2439 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2440 raw.extend([0, b'x' as u16]);
2441 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2442 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2443 assert_eq!(decode_windows_process_name(&[]), "");
2444 }
2445
2446 #[test]
2447 fn windows_process_filter_keeps_only_antigravity_pids() {
2448 let processes = vec![
2449 (10, "agy.exe".to_string()),
2450 (20, "language_server_windows_x64.exe".to_string()),
2451 (30, "sshd.exe".to_string()),
2452 ];
2453 let pids = matching_windows_process_ids(&processes);
2454 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
2455 }
2456
2457 #[test]
2458 fn windows_listener_filter_joins_pid_loopback_and_port() {
2459 let pids = std::collections::HashSet::from([10]);
2460 let rows = [
2461 WindowsTcpRow {
2462 local_addr: [127, 0, 0, 1],
2463 local_port: u32::from(59870u16.to_be()),
2464 pid: 10,
2465 },
2466 WindowsTcpRow {
2467 local_addr: [127, 0, 0, 1],
2468 local_port: u32::from(59868u16.to_be()),
2469 pid: 10,
2470 },
2471 WindowsTcpRow {
2472 local_addr: [127, 0, 0, 1],
2473 local_port: u32::from(59870u16.to_be()),
2474 pid: 10,
2475 },
2476 WindowsTcpRow {
2477 local_addr: [0, 0, 0, 0],
2478 local_port: u32::from(50000u16.to_be()),
2479 pid: 10,
2480 },
2481 WindowsTcpRow {
2482 local_addr: [127, 0, 0, 1],
2483 local_port: u32::from(50001u16.to_be()),
2484 pid: 99,
2485 },
2486 WindowsTcpRow {
2487 local_addr: [127, 0, 0, 1],
2488 local_port: 0,
2489 pid: 10,
2490 },
2491 ];
2492 assert_eq!(
2493 matching_windows_ports(&pids, &rows),
2494 vec![(10, 59870), (10, 59868)]
2495 );
2496 }
2497
2498 #[test]
2502 fn windows_ports_from_two_products_keep_tls_listeners_last() {
2503 let pids = std::collections::HashSet::from([10, 20]);
2504 let row = |port: u16, pid: u32| WindowsTcpRow {
2505 local_addr: [127, 0, 0, 1],
2506 local_port: u32::from(port.to_be()),
2507 pid,
2508 };
2509 let rows = [
2510 row(40000, 10),
2511 row(40001, 10),
2512 row(50000, 20),
2513 row(50001, 20),
2514 ];
2515 assert_eq!(
2516 matching_windows_ports(&pids, &rows),
2517 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2518 );
2519 }
2520
2521 #[test]
2525 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
2526 use std::collections::BTreeMap;
2527
2528 assert_eq!(
2530 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
2531 vec![(10, 59870), (10, 59868)]
2532 );
2533 assert_eq!(
2538 probe_order(BTreeMap::from([
2539 (10, vec![40000, 40001]),
2540 (20, vec![50000, 50001]),
2541 ])),
2542 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)]
2543 );
2544 assert_eq!(
2547 probe_order(BTreeMap::from([
2548 (10, vec![6000, 5000, 4000]),
2549 (20, vec![6000, 7000]),
2550 ])),
2551 vec![(10, 6000), (20, 7000), (10, 5000), (10, 4000)]
2552 );
2553 assert!(probe_order(BTreeMap::new()).is_empty());
2554 }
2555
2556 #[test]
2561 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
2562 use std::collections::BTreeMap;
2563
2564 assert_eq!(
2565 probe_order(BTreeMap::from([
2566 (10, vec![40001, 40001, 40000, 40000]),
2567 (20, vec![50001, 50000]),
2568 ])),
2569 vec![(10, 40001), (20, 50001), (10, 40000), (20, 50000)],
2570 "duplicate rows must not reorder the ranks below them"
2571 );
2572 }
2573
2574 #[test]
2579 fn a_half_started_product_is_the_documented_exception() {
2580 use std::collections::BTreeMap;
2581
2582 assert_eq!(
2583 probe_order(BTreeMap::from([
2584 (10, vec![40000]),
2585 (20, vec![50001, 50000])
2586 ])),
2587 vec![(10, 40000), (20, 50001), (20, 50000)]
2588 );
2589 }
2590
2591 #[test]
2596 fn an_override_with_no_authority_is_dropped_not_probed() {
2597 for junk in ["/", "///", "http://", "https://", " / "] {
2598 assert_eq!(
2599 candidate_bases_with(Some(junk), vec![(10, 4242)]),
2600 vec![Candidate {
2601 base: "http://127.0.0.1:4242".to_string(),
2602 pid: Some(10),
2603 }],
2604 "{junk:?} should not survive as a candidate"
2605 );
2606 }
2607 assert!(candidate_bases_with(Some("/"), vec![]).is_empty());
2608 }
2609
2610 #[test]
2611 fn windows_table_bounds_reject_truncation_and_overflow() {
2612 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
2613 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
2614 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
2615 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
2616 assert_eq!(
2617 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
2618 None
2619 );
2620 }
2621
2622 #[cfg(target_os = "windows")]
2623 #[test]
2624 fn windows_tcp_table_parser_copies_complete_rows_only() {
2625 use std::mem::{offset_of, size_of};
2626 use windows_sys::Win32::NetworkManagement::IpHelper::{
2627 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
2628 };
2629
2630 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
2631 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
2632 let words = used.div_ceil(size_of::<u32>());
2633 let mut buffer = vec![0u32; words];
2634 let first = MIB_TCPROW_OWNER_PID {
2635 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2636 dwLocalPort: u32::from(59868u16.to_be()),
2637 dwOwningPid: 10,
2638 ..Default::default()
2639 };
2640 let second = MIB_TCPROW_OWNER_PID {
2641 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2642 dwLocalPort: u32::from(59870u16.to_be()),
2643 dwOwningPid: 10,
2644 ..Default::default()
2645 };
2646 unsafe {
2647 buffer.as_mut_ptr().write_unaligned(2);
2648 let rows = buffer
2649 .as_mut_ptr()
2650 .cast::<u8>()
2651 .add(offset)
2652 .cast::<MIB_TCPROW_OWNER_PID>();
2653 rows.write_unaligned(first);
2654 rows.add(1).write_unaligned(second);
2655 }
2656
2657 assert_eq!(
2658 parse_windows_tcp_rows(&buffer, used),
2659 vec![
2660 WindowsTcpRow {
2661 local_addr: [127, 0, 0, 1],
2662 local_port: u32::from(59868u16.to_be()),
2663 pid: 10,
2664 },
2665 WindowsTcpRow {
2666 local_addr: [127, 0, 0, 1],
2667 local_port: u32::from(59870u16.to_be()),
2668 pid: 10,
2669 },
2670 ]
2671 );
2672 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
2673 }
2674
2675 #[test]
2676 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
2677 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";
2680 assert_eq!(
2681 parse_lsof_pcn(output),
2682 vec![(74101, 61290), (74101, 61289), (74101, 8829)]
2683 );
2684 }
2685
2686 #[test]
2690 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
2691 let output = concat!(
2692 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
2693 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
2694 );
2695 assert_eq!(
2696 parse_lsof_pcn(output),
2697 vec![(100, 40001), (200, 50001), (100, 40000), (200, 50000)],
2698 "both HTTP listeners must precede both TLS listeners"
2699 );
2700 }
2701
2702 #[test]
2703 fn lsof_parser_matches_the_capitalised_macos_app_name() {
2704 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
2705 assert_eq!(parse_lsof_pcn(output), vec![(900, 54321)]);
2706 }
2707
2708 #[test]
2709 fn lsof_parser_deduplicates_and_handles_empty_output() {
2710 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
2711 assert_eq!(parse_lsof_pcn(output), vec![(1, 9000)]);
2712 assert!(parse_lsof_pcn("").is_empty());
2713 }
2714
2715 #[test]
2718 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
2719 let dir = tempfile::tempdir().unwrap();
2720 let cache = Cache::at(dir.path().join("usage.json"));
2721 let reason = AppError::Credentials("Antigravity: no local language server found".into());
2722
2723 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2724 let msg = err.to_string();
2725 assert!(msg.contains("no local language server found"), "{msg}");
2726 assert!(!msg.contains("no usable cache"), "{msg}");
2727 }
2728
2729 #[test]
2730 fn unusable_cache_does_not_replace_the_live_diagnosis() {
2731 let dir = tempfile::tempdir().unwrap();
2732 let cache = Cache::at(dir.path().join("antigravity"));
2733 cache.write_payload(b"{}").unwrap();
2734
2735 let reason = AppError::Credentials("Antigravity must be running".into());
2736 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2737 assert!(err.to_string().contains("must be running"), "{err}");
2738
2739 let original = AppError::Transport("original loopback failure".into());
2740 let err = fallback_silent(&cache, now(), original).unwrap_err();
2741 assert!(
2742 err.to_string().contains("original loopback failure"),
2743 "{err}"
2744 );
2745 }
2746
2747 #[tokio::test]
2748 async fn rpc_error_bodies_are_bounded_too() {
2749 let mut server = mockito::Server::new_async().await;
2750 let path = format!("/{STATUS_RPC}");
2751 server
2752 .mock("POST", path.as_str())
2753 .with_status(500)
2754 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
2755 .create_async()
2756 .await;
2757
2758 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
2759 .await
2760 .unwrap_err();
2761 assert!(err.to_string().contains("exceeds"), "{err}");
2762 }
2763
2764 #[test]
2765 fn blank_override_falls_through_to_discovery() {
2766 assert_eq!(
2767 candidate_bases_with(Some(" "), vec![(10, 4242)]),
2768 vec![Candidate {
2769 base: "http://127.0.0.1:4242".to_string(),
2770 pid: Some(10),
2771 }]
2772 );
2773 }
2774
2775 fn fixture() -> (tempfile::TempDir, Cache) {
2780 let td = tempfile::TempDir::new().unwrap();
2781 let cache = Cache::at(td.path().join("antigravity"));
2782 (td, cache)
2783 }
2784
2785 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
2786 let base = server.url();
2787 cloud::Endpoints {
2788 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
2789 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
2790 token: format!("{base}/token"),
2791 }
2792 }
2793
2794 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
2796 let mut token = serde_json::json!({
2797 "access_token": "KEYRING-AT",
2798 "expiry": expiry,
2799 });
2800 if with_refresh {
2801 token["refresh_token"] = serde_json::json!("KEYRING-RT");
2802 }
2803 serde_json::json!({ "token": token }).to_string()
2804 }
2805
2806 const EXPIRED: &str = "2026-07-22T11:00:00Z";
2808 const VALID: &str = "2026-07-22T13:00:00Z";
2810
2811 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
2813 RemoteOverride {
2814 credential: SavedCredential::Blob(blob),
2815 endpoints: Some(eps),
2816 local_bases: Some(vec![]),
2817 ..Default::default()
2818 }
2819 }
2820
2821 fn bare_summary() -> String {
2824 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2825 v["response"].to_string()
2826 }
2827
2828 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
2829 server
2830 .mock("POST", "/daily/quota")
2831 .match_header("authorization", format!("Bearer {bearer}").as_str())
2832 .with_status(200)
2833 .with_body(bare_summary())
2834 }
2835
2836 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
2837 server
2838 .mock("POST", "/token")
2839 .with_status(200)
2840 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
2841 }
2842
2843 fn test_oauth() -> cloud::OauthClient {
2844 cloud::OauthClient {
2845 id: "test-client".into(),
2846 secret: "test-client-secret".into(),
2847 }
2848 }
2849
2850 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
2851 fetch_snapshot_at(
2852 &reqwest::Client::new(),
2853 cache,
2854 ttl,
2855 Some(&test_oauth()),
2856 remote,
2857 now(),
2858 )
2859 .await
2860 }
2861
2862 #[tokio::test]
2863 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
2864 let mut server = mockito::Server::new_async().await;
2865 let eps = endpoints(&server);
2866 let quota = quota_mock(&mut server, "KEYRING-AT")
2867 .expect(1)
2868 .create_async()
2869 .await;
2870 let plan = server
2871 .mock("POST", "/daily/plan")
2872 .match_header("authorization", "Bearer KEYRING-AT")
2873 .with_status(200)
2874 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2875 .expect(1)
2876 .create_async()
2877 .await;
2878 let token = token_mock(&mut server).expect(0).create_async().await;
2880 let blob = keyring_blob(VALID, true);
2881 let (_td, cache) = fixture();
2882
2883 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
2884 .await
2885 .expect("remote path yields a snapshot");
2886
2887 quota.assert_async().await;
2888 plan.assert_async().await;
2889 token.assert_async().await;
2890 let snap = outcome.snapshot;
2891 assert!(!outcome.stale);
2892 assert_eq!(snap.source, AntigravitySource::Remote);
2893 assert_eq!(snap.plan, "Pro");
2894 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2895 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2896 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
2897 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
2898 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
2899 assert_eq!(snap.account, format!("acct:{fingerprint}"));
2900 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
2901
2902 let cached = parse_cache_at(
2904 &cache
2905 .fresh_payload(Duration::from_secs(60))
2906 .unwrap()
2907 .unwrap(),
2908 None,
2909 now(),
2910 )
2911 .unwrap();
2912 assert_eq!(cached.source, AntigravitySource::Remote);
2913 assert_eq!(cached.account, snap.account);
2914 }
2915
2916 #[tokio::test]
2921 async fn agys_missing_csrf_response_uses_the_saved_session() {
2922 let mut server = mockito::Server::new_async().await;
2923 let eps = endpoints(&server);
2924 let root = server
2925 .mock("GET", "/")
2926 .with_status(404)
2927 .expect(1)
2928 .create_async()
2929 .await;
2930 let status_path = format!("/{STATUS_RPC}");
2931 let status = server
2932 .mock("POST", status_path.as_str())
2933 .with_status(401)
2934 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
2935 .expect(1)
2936 .create_async()
2937 .await;
2938 let quota = quota_mock(&mut server, "KEYRING-AT")
2939 .expect(1)
2940 .create_async()
2941 .await;
2942 let plan = server
2943 .mock("POST", "/daily/plan")
2944 .match_header("authorization", "Bearer KEYRING-AT")
2945 .with_status(200)
2946 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2947 .expect(1)
2948 .create_async()
2949 .await;
2950 let blob = keyring_blob(VALID, true);
2951 let (_td, cache) = fixture();
2952
2953 let outcome = run(
2954 &cache,
2955 RemoteOverride {
2956 credential: SavedCredential::Blob(&blob),
2957 endpoints: Some(&eps),
2958 local_bases: Some(vec![server.url().into()]),
2959 ..Default::default()
2960 },
2961 Duration::ZERO,
2962 )
2963 .await
2964 .expect("saved session bypasses agy's unusable local RPC");
2965
2966 root.assert_async().await;
2967 status.assert_async().await;
2968 quota.assert_async().await;
2969 plan.assert_async().await;
2970 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
2971 }
2972
2973 #[tokio::test]
2977 async fn agys_prefixed_missing_csrf_response_uses_the_saved_session() {
2978 let mut server = mockito::Server::new_async().await;
2979 let eps = endpoints(&server);
2980 let root = server
2981 .mock("GET", "/")
2982 .with_status(404)
2983 .expect(1)
2984 .create_async()
2985 .await;
2986 let status_path = format!("/{STATUS_RPC}");
2987 let status = server
2988 .mock("POST", status_path.as_str())
2989 .with_status(401)
2990 .with_body(
2991 r#"{"code":"unauthenticated","message":"unauthenticated: missing CSRF token"}"#,
2992 )
2993 .expect(1)
2994 .create_async()
2995 .await;
2996 let quota = quota_mock(&mut server, "KEYRING-AT")
2997 .expect(1)
2998 .create_async()
2999 .await;
3000 let plan = server
3001 .mock("POST", "/daily/plan")
3002 .match_header("authorization", "Bearer KEYRING-AT")
3003 .with_status(200)
3004 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
3005 .expect(1)
3006 .create_async()
3007 .await;
3008 let blob = keyring_blob(VALID, true);
3009 let (_td, cache) = fixture();
3010
3011 let outcome = run(
3012 &cache,
3013 RemoteOverride {
3014 credential: SavedCredential::Blob(&blob),
3015 endpoints: Some(&eps),
3016 local_bases: Some(vec![server.url().into()]),
3017 ..Default::default()
3018 },
3019 Duration::ZERO,
3020 )
3021 .await
3022 .expect("saved session bypasses agy 1.3.1's prefixed missing-CSRF RPC");
3023
3024 root.assert_async().await;
3025 status.assert_async().await;
3026 quota.assert_async().await;
3027 plan.assert_async().await;
3028 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3029 }
3030
3031 #[tokio::test]
3032 async fn agys_missing_csrf_without_a_saved_session_explains_both_options() {
3033 let mut server = mockito::Server::new_async().await;
3034 let eps = endpoints(&server);
3035 let root = server
3036 .mock("GET", "/")
3037 .with_status(404)
3038 .expect(1)
3039 .create_async()
3040 .await;
3041 let status_path = format!("/{STATUS_RPC}");
3042 let status = server
3043 .mock("POST", status_path.as_str())
3044 .with_status(401)
3045 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3046 .expect(1)
3047 .create_async()
3048 .await;
3049 let quota = quota_mock(&mut server, "KEYRING-AT")
3050 .expect(0)
3051 .create_async()
3052 .await;
3053 let (_td, cache) = fixture();
3054
3055 let error = run(
3056 &cache,
3057 RemoteOverride {
3058 credential: SavedCredential::Absent,
3059 endpoints: Some(&eps),
3060 local_bases: Some(vec![server.url().into()]),
3061 ..Default::default()
3062 },
3063 Duration::ZERO,
3064 )
3065 .await
3066 .expect_err("neither the local RPC nor a saved session is usable");
3067
3068 root.assert_async().await;
3069 status.assert_async().await;
3070 quota.assert_async().await;
3071 let message = error.to_string();
3072 assert!(message.contains("requires a CSRF token"), "{message}");
3073 assert!(message.contains("saved Google session"), "{message}");
3074 }
3075
3076 #[tokio::test]
3081 async fn agys_missing_csrf_aborts_further_candidate_probes_to_spare_tls_listeners() {
3082 let mut server1 = mockito::Server::new_async().await;
3083 let mut server2 = mockito::Server::new_async().await;
3084 let eps = endpoints(&server1);
3085 let root1 = server1
3086 .mock("GET", "/")
3087 .with_status(404)
3088 .expect(1)
3089 .create_async()
3090 .await;
3091 let status_path = format!("/{STATUS_RPC}");
3092 let status1 = server1
3093 .mock("POST", status_path.as_str())
3094 .with_status(401)
3095 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3096 .expect(1)
3097 .create_async()
3098 .await;
3099
3100 let root2 = server2.mock("GET", "/").expect(0).create_async().await;
3102 let status2 = server2
3103 .mock("POST", status_path.as_str())
3104 .expect(0)
3105 .create_async()
3106 .await;
3107
3108 let (_td, cache) = fixture();
3109 let error = run(
3110 &cache,
3111 RemoteOverride {
3112 credential: SavedCredential::Absent,
3113 endpoints: Some(&eps),
3114 local_bases: Some(vec![
3115 Candidate {
3116 base: server1.url(),
3117 pid: Some(42),
3118 },
3119 Candidate {
3120 base: server2.url(),
3121 pid: Some(42),
3122 },
3123 ]),
3124 ..Default::default()
3125 },
3126 Duration::ZERO,
3127 )
3128 .await
3129 .expect_err("aborted probe returns agy missing csrf diagnosis");
3130
3131 root1.assert_async().await;
3132 status1.assert_async().await;
3133 root2.assert_async().await;
3134 status2.assert_async().await;
3135 let message = error.to_string();
3136 assert!(message.contains("requires a CSRF token"), "{message}");
3137 }
3138
3139 #[tokio::test]
3146 async fn agys_missing_csrf_does_not_skip_other_products_servers() {
3147 let mut agy = mockito::Server::new_async().await;
3148 let mut ide = mockito::Server::new_async().await;
3149 let eps = endpoints(&agy);
3150 let agy_root = agy
3151 .mock("GET", "/")
3152 .with_status(404)
3153 .expect(1)
3154 .create_async()
3155 .await;
3156 let status_path = format!("/{STATUS_RPC}");
3157 let agy_status = agy
3158 .mock("POST", status_path.as_str())
3159 .with_status(401)
3160 .with_body(r#"{"code":"unauthenticated","message":"missing CSRF token"}"#)
3161 .expect(1)
3162 .create_async()
3163 .await;
3164
3165 let ide_root = ide
3166 .mock("GET", "/")
3167 .with_status(404)
3168 .expect(1)
3169 .create_async()
3170 .await;
3171 let ide_status = ide
3172 .mock("POST", status_path.as_str())
3173 .with_status(200)
3174 .with_body(r#"{"userStatus":{"userTier":{"name":"Pro"},"email":"u@example.com"}}"#)
3175 .expect(1)
3176 .create_async()
3177 .await;
3178 let quota_path = format!("/{QUOTA_RPC}");
3179 let ide_quota = ide
3180 .mock("POST", quota_path.as_str())
3181 .with_status(200)
3182 .with_body(QUOTA_JSON)
3183 .expect(1)
3184 .create_async()
3185 .await;
3186
3187 let (_td, cache) = fixture();
3188 let outcome = run(
3189 &cache,
3190 RemoteOverride {
3191 credential: SavedCredential::Absent,
3192 endpoints: Some(&eps),
3193 local_bases: Some(vec![
3194 Candidate {
3195 base: agy.url(),
3196 pid: Some(42),
3197 },
3198 Candidate {
3199 base: ide.url(),
3200 pid: Some(7),
3201 },
3202 ]),
3203 ..Default::default()
3204 },
3205 Duration::ZERO,
3206 )
3207 .await
3208 .expect("the other product's server still answers");
3209
3210 agy_root.assert_async().await;
3211 agy_status.assert_async().await;
3212 ide_root.assert_async().await;
3213 ide_status.assert_async().await;
3214 ide_quota.assert_async().await;
3215 assert_eq!(outcome.snapshot.source, AntigravitySource::Local);
3216 assert_eq!(outcome.snapshot.plan, "Pro");
3217 }
3218
3219 #[tokio::test]
3222 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
3223 let mut server = mockito::Server::new_async().await;
3224 let eps = endpoints(&server);
3225 let _quota = quota_mock(&mut server, "NEW-AT")
3226 .expect(2)
3227 .create_async()
3228 .await;
3229 let token = token_mock(&mut server).expect(1).create_async().await;
3230 let blob = keyring_blob(EXPIRED, true);
3231 let (_td, cache) = fixture();
3232
3233 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
3234 .await
3235 .expect("refresh, then quota");
3236 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
3237 token.assert_async().await;
3238
3239 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
3240 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
3241 .expect("refreshed token persisted under the keyring session's fingerprint");
3242 assert_eq!(persisted.access_token, "NEW-AT");
3243
3244 run(&cache, remote(&blob, &eps), Duration::ZERO)
3246 .await
3247 .expect("persisted token reused");
3248 token.assert_async().await;
3249 }
3250
3251 #[tokio::test]
3254 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
3255 let mut server = mockito::Server::new_async().await;
3256 let eps = endpoints(&server);
3257 let rejected = server
3258 .mock("POST", "/daily/quota")
3259 .match_header("authorization", "Bearer KEYRING-AT")
3260 .with_status(401)
3261 .expect(1)
3262 .create_async()
3263 .await;
3264 let accepted = quota_mock(&mut server, "NEW-AT")
3265 .expect(1)
3266 .create_async()
3267 .await;
3268 let token = token_mock(&mut server).expect(1).create_async().await;
3269 let blob = keyring_blob(VALID, true);
3270 let (_td, cache) = fixture();
3271
3272 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3273 .await
3274 .expect("retry with the refreshed token succeeds");
3275
3276 rejected.assert_async().await;
3277 accepted.assert_async().await;
3278 token.assert_async().await;
3279 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
3280 }
3281
3282 #[tokio::test]
3285 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
3286 let mut server = mockito::Server::new_async().await;
3287 let eps = endpoints(&server);
3288 let _token = server
3289 .mock("POST", "/token")
3290 .with_status(400)
3291 .with_body(r#"{"error":"invalid_grant"}"#)
3292 .create_async()
3293 .await;
3294 let blob = keyring_blob(EXPIRED, true);
3295 let (_td, cache) = fixture();
3296
3297 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3298 .await
3299 .expect_err("no cache to fall back on");
3300
3301 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3302 let rendered = err.to_string();
3303 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
3304 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
3305 }
3306
3307 #[tokio::test]
3310 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
3311 let mut server = mockito::Server::new_async().await;
3312 let eps = endpoints(&server);
3313 let _quota = server
3314 .mock("POST", "/daily/quota")
3315 .with_status(401)
3316 .create_async()
3317 .await;
3318 let token = token_mock(&mut server).expect(1).create_async().await;
3319 let blob = keyring_blob(EXPIRED, true);
3320 let (_td, cache) = fixture();
3321
3322 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3323 .await
3324 .expect_err("rejected after refresh");
3325
3326 token.assert_async().await;
3327 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3328 assert!(err.to_string().contains("sign in again"), "{err}");
3329 }
3330
3331 async fn run_without_oauth(cache: &Cache, remote: RemoteOverride<'_>) -> Result<FetchOutcome> {
3333 fetch_snapshot_at(
3334 &reqwest::Client::new(),
3335 cache,
3336 Duration::ZERO,
3337 None,
3338 remote,
3339 now(),
3340 )
3341 .await
3342 }
3343
3344 #[cfg(unix)]
3347 fn fake_agy(
3348 dir: &std::path::Path,
3349 runs: &std::path::Path,
3350 file: &std::path::Path,
3351 blob: &str,
3352 ) -> std::path::PathBuf {
3353 use std::os::unix::fs::PermissionsExt;
3354 let path = dir.join("agy");
3355 std::fs::write(
3356 &path,
3357 format!(
3358 "#!/bin/sh\necho \"$@\" >> '{}'\nprintf '%s' '{}' > '{}'\n",
3359 runs.display(),
3360 blob,
3361 file.display()
3362 ),
3363 )
3364 .unwrap();
3365 std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o755)).unwrap();
3366 path
3367 }
3368
3369 #[cfg(unix)]
3372 #[tokio::test]
3373 async fn an_expired_session_is_renewed_by_running_agy() {
3374 let mut server = mockito::Server::new_async().await;
3375 let eps = endpoints(&server);
3376 let quota = quota_mock(&mut server, "KEYRING-AT").create_async().await;
3377 let _plan = server
3378 .mock("POST", "/daily/plan")
3379 .with_status(200)
3380 .with_body(r#"{"currentTier":{"name":"Pro"}}"#)
3381 .create_async()
3382 .await;
3383 let (td, cache) = fixture();
3384 let file = td.path().join("credential");
3385 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
3386 let runs = td.path().join("runs");
3387 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(VALID, true));
3388
3389 let out = run_without_oauth(
3390 &cache,
3391 RemoteOverride {
3392 credential: SavedCredential::File(&file),
3393 endpoints: Some(&eps),
3394 local_bases: Some(vec![]),
3395 agy: super::super::agy::AgyCommand::At(agy),
3396 },
3397 )
3398 .await
3399 .expect("agy renewed the session");
3400
3401 quota.assert_async().await;
3402 assert_eq!(out.snapshot.source, AntigravitySource::Remote);
3403 assert_eq!(std::fs::read_to_string(runs).unwrap().trim(), "models");
3404 }
3405
3406 #[cfg(unix)]
3409 #[tokio::test]
3410 async fn agy_that_does_not_renew_keeps_the_original_error_and_runs_once() {
3411 let server = mockito::Server::new_async().await;
3412 let eps = endpoints(&server);
3413 let (td, cache) = fixture();
3414 let file = td.path().join("credential");
3415 std::fs::write(&file, keyring_blob(EXPIRED, true)).unwrap();
3416 let runs = td.path().join("runs");
3417 let agy = fake_agy(td.path(), &runs, &file, &keyring_blob(EXPIRED, true));
3418 let remote = || RemoteOverride {
3419 credential: SavedCredential::File(&file),
3420 endpoints: Some(&eps),
3421 local_bases: Some(vec![]),
3422 agy: super::super::agy::AgyCommand::At(agy.clone()),
3423 };
3424
3425 for _ in 0..2 {
3426 let err = run_without_oauth(&cache, remote())
3427 .await
3428 .expect_err("still expired");
3429 assert!(err.to_string().contains("no OAuth client"), "{err}");
3430 }
3431 assert_eq!(std::fs::read_to_string(runs).unwrap().lines().count(), 1);
3432 }
3433
3434 #[tokio::test]
3437 async fn without_an_agy_command_an_expired_session_is_left_alone() {
3438 let server = mockito::Server::new_async().await;
3439 let eps = endpoints(&server);
3440 let blob = keyring_blob(EXPIRED, true);
3441 let (_td, cache) = fixture();
3442
3443 let err = run_without_oauth(&cache, remote(&blob, &eps))
3444 .await
3445 .expect_err("nothing renews it");
3446 assert!(err.to_string().contains("no OAuth client"), "{err}");
3447 }
3448
3449 #[tokio::test]
3451 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
3452 let mut server = mockito::Server::new_async().await;
3453 let eps = endpoints(&server);
3454 let token = token_mock(&mut server).expect(0).create_async().await;
3455 let blob = keyring_blob(EXPIRED, false);
3456 let (_td, cache) = fixture();
3457
3458 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
3459 .await
3460 .expect_err("nothing to refresh with");
3461
3462 token.assert_async().await;
3463 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3464 assert!(err.to_string().contains("sign in again"), "{err}");
3465 }
3466
3467 #[tokio::test]
3470 async fn no_saved_session_extends_the_no_local_server_error() {
3471 let mut server = mockito::Server::new_async().await;
3472 let eps = endpoints(&server);
3473 let quota = quota_mock(&mut server, "KEYRING-AT")
3474 .expect(0)
3475 .create_async()
3476 .await;
3477 let (_td, cache) = fixture();
3478
3479 let err = run(
3480 &cache,
3481 RemoteOverride {
3482 credential: SavedCredential::Absent,
3483 endpoints: Some(&eps),
3484 local_bases: Some(vec![]),
3485 ..Default::default()
3486 },
3487 Duration::ZERO,
3488 )
3489 .await
3490 .expect_err("no source at all");
3491
3492 quota.assert_async().await;
3493 assert!(matches!(err, AppError::Credentials(_)), "{err}");
3494 let rendered = err.to_string();
3495 assert!(rendered.contains("no local server found"), "{rendered}");
3496 assert!(rendered.contains("saved Google session"), "{rendered}");
3497 }
3498
3499 #[tokio::test]
3502 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
3503 let mut server = mockito::Server::new_async().await;
3504 let eps = endpoints(&server);
3505 let status_path = format!("/{STATUS_RPC}");
3506 let _status = server
3507 .mock("POST", status_path.as_str())
3508 .with_status(401)
3509 .create_async()
3510 .await;
3511 let quota = quota_mock(&mut server, "KEYRING-AT")
3512 .expect(0)
3513 .create_async()
3514 .await;
3515 let blob = keyring_blob(VALID, true);
3516 let (_td, cache) = fixture();
3517
3518 let err = run(
3519 &cache,
3520 RemoteOverride {
3521 credential: SavedCredential::Blob(&blob),
3522 endpoints: Some(&eps),
3523 local_bases: Some(vec![server.url().into()]),
3524 ..Default::default()
3525 },
3526 Duration::ZERO,
3527 )
3528 .await
3529 .expect_err("signed-out local server");
3530
3531 quota.assert_async().await;
3532 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
3533 }
3534
3535 #[tokio::test]
3538 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
3539 let mut server = mockito::Server::new_async().await;
3540 let eps = endpoints(&server);
3541 for path in ["/daily/quota", "/prod/quota"] {
3542 server
3543 .mock("POST", path)
3544 .with_status(500)
3545 .with_body("boom")
3546 .create_async()
3547 .await;
3548 }
3549 let blob = keyring_blob(VALID, true);
3550 let (_td, cache) = fixture();
3551 cache.ensure_dir().unwrap();
3552 let mut earlier = parsed();
3553 earlier.source = AntigravitySource::Remote;
3554 earlier.account = "acct:earlier".into();
3555 cache
3556 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
3557 .unwrap();
3558
3559 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
3560 .await
3561 .expect("stale cache stands in");
3562
3563 assert!(outcome.stale);
3564 assert_eq!(outcome.snapshot, earlier);
3565 assert!(
3566 matches!(outcome.last_error, Some((500, _))),
3567 "{:?}",
3568 outcome.last_error
3569 );
3570 }
3571
3572 #[test]
3573 fn source_round_trips_through_the_cache_and_defaults_to_local() {
3574 let mut snap = parsed();
3575 snap.source = AntigravitySource::Remote;
3576 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3577 assert_eq!(
3578 parse_cache_at(&bytes, None, now()).unwrap().source,
3579 AntigravitySource::Remote
3580 );
3581
3582 snap.source = AntigravitySource::Statusline;
3583 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
3584 assert_eq!(
3585 parse_cache_at(&bytes, None, now()).unwrap().source,
3586 AntigravitySource::Statusline
3587 );
3588
3589 let mut legacy = snap_to_json(&snap);
3591 legacy.as_object_mut().unwrap().remove("source");
3592 let legacy = serde_json::to_vec(&legacy).unwrap();
3593 assert_eq!(
3594 parse_cache_at(&legacy, None, now()).unwrap().source,
3595 AntigravitySource::Local
3596 );
3597 }
3598}