1use std::time::Duration;
24
25use chrono::{DateTime, Utc};
26
27use super::cloud;
28use super::credential::{self, StoredToken};
29use crate::cache::{Cache, acquire_lock_async};
30use crate::error::{AppError, Result};
31use crate::usage::{AntigravitySnapshot, AntigravitySource, UsageWindow};
32
33const HTTP_TIMEOUT: Duration = Duration::from_secs(5);
34const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
35
36const QUOTA_RPC: &str = "exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary";
37const STATUS_RPC: &str = "exa.language_server_pb.LanguageServerService/GetUserStatus";
38
39const DEFAULT_PLAN: &str = "Antigravity";
40
41const NO_LOCAL_SERVER: &str = "Antigravity: no local server found. Quota is only served while \
42 Antigravity is running — open the Antigravity app, or an interactive \
43 `agy` session, or point ANTIGRAVITY_LS_ADDRESS at a host:port.";
44
45const NO_SAVED_SESSION: &str = "Or sign in to Antigravity once, so its saved Google session can \
49 be used while it is closed.";
50
51const SESSION_EXPIRED: &str =
52 "Antigravity's saved Google session expired; open Antigravity to sign in again";
53
54pub type FetchOutcome = crate::outcome::Outcome<AntigravitySnapshot>;
57
58impl From<FetchOutcome> for crate::vendor::VendorOutcome {
59 fn from(o: FetchOutcome) -> Self {
60 o.map(crate::usage::VendorSnapshot::Antigravity)
61 }
62}
63
64#[derive(Debug, Clone, Copy, Default)]
68pub enum SavedCredential<'a> {
69 #[default]
71 Keyring,
72 Blob(&'a str),
74 Absent,
76}
77
78#[derive(Default)]
83pub struct RemoteOverride<'a> {
84 pub credential: SavedCredential<'a>,
86 pub endpoints: Option<&'a cloud::Endpoints>,
88 pub local_bases: Option<Vec<String>>,
91}
92
93pub async fn fetch_snapshot(
94 client: &reqwest::Client,
95 cache: &Cache,
96 cache_ttl: Duration,
97 oauth: Option<&cloud::OauthClient>,
98) -> Result<FetchOutcome> {
99 fetch_snapshot_at(
100 client,
101 cache,
102 cache_ttl,
103 oauth,
104 RemoteOverride::default(),
105 Utc::now(),
106 )
107 .await
108}
109
110pub async fn fetch_snapshot_at(
114 client: &reqwest::Client,
115 cache: &Cache,
116 cache_ttl: Duration,
117 oauth: Option<&cloud::OauthClient>,
118 remote: RemoteOverride<'_>,
119 now: DateTime<Utc>,
120) -> Result<FetchOutcome> {
121 cache.ensure_dir()?;
122 let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
123
124 let origin = match open_session(client, remote.local_bases.as_deref()).await {
134 Err(e) if is_no_local_server(&e) => Origin::Remote(saved_session(remote.credential)),
135 session => Origin::Local(session),
136 };
137 let account = origin.account();
138
139 if let Some(bytes) = cache.fresh_payload(cache_ttl)?
140 && let Ok(outcome) = reuse_cache(bytes, cache, false, account.as_deref(), now)
141 {
142 return Ok(outcome);
143 }
144
145 let default_endpoints = cloud::Endpoints::default();
146 let endpoints = remote.endpoints.unwrap_or(&default_endpoints);
147 let live = match origin {
148 Origin::Local(session) => fetch_live(client, session).await,
149 Origin::Remote(token) => fetch_remote(client, cache, oauth, endpoints, token, now).await,
150 };
151
152 match live {
153 Ok(snap) => {
154 let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
155 cache.write_payload(&bytes)?;
156 Ok(crate::outcome::Outcome::fresh(snap))
157 }
158 Err(e) if e.is_transient() => fallback_silent(cache, now, e),
159 Err(AppError::Http { status, body }) => {
160 cache.mark_stale();
161 let last_error = Some(cache.write_last_error(status, &body));
162 let reason = AppError::Http { status, body };
163 fallback_with_error(cache, last_error, reason, now)
164 }
165 Err(e) => {
166 cache.mark_stale();
167 let last_error = Some(cache.write_last_error(0, &e.to_string()));
168 fallback_with_error(cache, last_error, e, now)
169 }
170 }
171}
172
173struct Session {
176 base: String,
177 csrf: Option<String>,
178 plan: String,
179 account: String,
180}
181
182enum Origin {
185 Local(Result<Session>),
186 Remote(Result<StoredToken>),
187}
188
189impl Origin {
190 fn account(&self) -> Option<String> {
192 match self {
193 Origin::Local(Ok(session)) => Some(session.account.clone()),
194 Origin::Remote(Ok(token)) => Some(remote_account(&token.fingerprint)),
195 Origin::Local(Err(_)) | Origin::Remote(Err(_)) => None,
196 }
197 }
198}
199
200fn no_local_server() -> AppError {
201 AppError::Credentials(NO_LOCAL_SERVER.into())
202}
203
204fn is_no_local_server(e: &AppError) -> bool {
208 matches!(e, AppError::Credentials(msg) if msg == NO_LOCAL_SERVER)
209}
210
211async fn open_session(client: &reqwest::Client, bases: Option<&[String]>) -> Result<Session> {
217 let bases = bases.map_or_else(candidate_bases, <[String]>::to_vec);
218 if bases.is_empty() {
219 return Err(no_local_server());
220 }
221
222 let mut errors = Vec::new();
223 for base in bases {
224 let csrf = fetch_csrf(client, &base).await;
225 match post_rpc(client, &base, csrf.as_deref(), STATUS_RPC).await {
226 Ok(v) => {
227 return Ok(Session {
228 base,
229 csrf,
230 plan: plan_from_status(&v),
231 account: account_key(&v),
232 });
233 }
234 Err(e) => errors.push(e),
235 }
236 }
237 Err(select_probe_error(errors))
238}
239
240fn select_probe_error(errors: Vec<AppError>) -> AppError {
255 let mut actionable = None;
256 let mut last = None;
257 let mut echo = None;
258 for e in errors {
259 if actionable.is_none() && is_actionable(&e) {
260 actionable = Some(e);
261 } else if is_tls_echo(&e) {
262 echo = Some(e);
263 } else {
264 last = Some(e);
265 }
266 }
267 actionable.or(last).or(echo).unwrap_or_else(|| {
268 AppError::Other("antigravity: no local server answered GetUserStatus".into())
269 })
270}
271
272fn is_actionable(e: &AppError) -> bool {
276 matches!(e, AppError::Http { status, .. } if *status == 401 || *status == 403)
277}
278
279fn is_tls_echo(e: &AppError) -> bool {
292 matches!(
293 e,
294 AppError::Http { status: 400, body } if body.contains("HTTP request to an HTTPS server")
295 )
296}
297
298async fn fetch_live(
299 client: &reqwest::Client,
300 session: Result<Session>,
301) -> Result<AntigravitySnapshot> {
302 let session = session?;
303 let quota = post_rpc(client, &session.base, session.csrf.as_deref(), QUOTA_RPC).await?;
304 let mut snap = parse_quota_summary("a, session.plan)?;
305 snap.account = session.account;
306 Ok(snap)
307}
308
309fn saved_session(credential: SavedCredential<'_>) -> Result<StoredToken> {
319 let raw = match credential {
320 SavedCredential::Keyring => credential::read()?,
321 SavedCredential::Blob(blob) => Some(blob.to_string()),
322 SavedCredential::Absent => None,
323 };
324 let Some(raw) = raw else {
325 return Err(AppError::Credentials(format!(
326 "{NO_LOCAL_SERVER} {NO_SAVED_SESSION}"
327 )));
328 };
329 credential::parse_keyring_blob(&raw)
330}
331
332fn remote_account(fingerprint: &str) -> String {
336 format!("acct:{fingerprint}")
337}
338
339fn session_expired() -> AppError {
340 AppError::Credentials(SESSION_EXPIRED.into())
341}
342
343fn refresh_unconfigured() -> AppError {
347 AppError::Credentials(
348 "Antigravity's saved Google session expired and ai-usagebar has no OAuth client to \
349 refresh it; open Antigravity to sign in again, or set [antigravity] oauth_client_id \
350 and oauth_client_secret in config.toml"
351 .into(),
352 )
353}
354
355fn is_auth_rejection(e: &AppError) -> bool {
356 matches!(
357 e,
358 AppError::Http {
359 status: 401 | 403,
360 ..
361 }
362 )
363}
364
365struct AccessToken {
369 value: String,
370 just_refreshed: bool,
371}
372
373async fn resolve_access_token(
378 client: &reqwest::Client,
379 oauth: Option<&cloud::OauthClient>,
380 endpoints: &cloud::Endpoints,
381 oauth_path: &std::path::Path,
382 token: &StoredToken,
383 now: DateTime<Utc>,
384) -> Result<AccessToken> {
385 if let Some(persisted) = cloud::read_persisted(oauth_path, &token.fingerprint)
386 && token
387 .expires_at
388 .is_none_or(|keyring| persisted.expires_at > keyring)
389 && !cloud::needs_refresh(Some(persisted.expires_at), now)
390 {
391 return Ok(AccessToken {
392 value: persisted.access_token,
393 just_refreshed: false,
394 });
395 }
396 if !cloud::needs_refresh(token.expires_at, now) {
397 return Ok(AccessToken {
398 value: token.access_token.clone(),
399 just_refreshed: false,
400 });
401 }
402 refresh_and_persist(client, oauth, endpoints, oauth_path, token).await
403}
404
405async fn refresh_and_persist(
409 client: &reqwest::Client,
410 oauth: Option<&cloud::OauthClient>,
411 endpoints: &cloud::Endpoints,
412 oauth_path: &std::path::Path,
413 token: &StoredToken,
414) -> Result<AccessToken> {
415 let Some(refresh_token) = token.refresh_token.as_deref() else {
416 return Err(session_expired());
417 };
418 let Some(oauth) = oauth else {
419 return Err(refresh_unconfigured());
420 };
421 let refreshed = cloud::refresh(client, &endpoints.token, oauth, refresh_token).await?;
422 cloud::write_persisted(
423 oauth_path,
424 &cloud::PersistedOAuth {
425 fingerprint: token.fingerprint.clone(),
426 access_token: refreshed.access_token.clone(),
427 expires_at: refreshed.expires_at,
428 },
429 )?;
430 Ok(AccessToken {
431 value: refreshed.access_token,
432 just_refreshed: true,
433 })
434}
435
436async fn fetch_remote(
443 client: &reqwest::Client,
444 cache: &Cache,
445 oauth: Option<&cloud::OauthClient>,
446 endpoints: &cloud::Endpoints,
447 token: Result<StoredToken>,
448 now: DateTime<Utc>,
449) -> Result<AntigravitySnapshot> {
450 let token = token?;
451 let oauth_path = cloud::oauth_cache_path(cache);
452 let mut access =
453 resolve_access_token(client, oauth, endpoints, &oauth_path, &token, now).await?;
454
455 let quota = match cloud::fetch_quota(client, endpoints, &access.value).await {
456 Err(e) if is_auth_rejection(&e) && !access.just_refreshed => {
457 access = refresh_and_persist(client, oauth, endpoints, &oauth_path, &token).await?;
458 cloud::fetch_quota(client, endpoints, &access.value)
459 .await
460 .map_err(|e| {
461 if is_auth_rejection(&e) {
462 session_expired()
463 } else {
464 e
465 }
466 })?
467 }
468 Err(e) if is_auth_rejection(&e) => return Err(session_expired()),
469 other => other?,
470 };
471
472 let plan = cloud::fetch_plan(client, endpoints, &access.value)
473 .await
474 .unwrap_or_else(|| DEFAULT_PLAN.to_string());
475 let mut snap = parse_quota_summary("a, plan)?;
476 snap.account = remote_account(&token.fingerprint);
477 snap.source = AntigravitySource::Remote;
478 Ok(snap)
479}
480
481fn account_key(user_status: &serde_json::Value) -> String {
486 let email = user_status["userStatus"]["email"]
487 .as_str()
488 .filter(|s| !s.is_empty());
489 match email {
490 Some(e) => {
491 use std::hash::{Hash, Hasher};
492 let mut h = std::collections::hash_map::DefaultHasher::new();
493 e.hash(&mut h);
494 format!("acct:{:016x}", h.finish())
495 }
496 None => "acct:unknown".to_string(),
497 }
498}
499
500async fn fetch_csrf(client: &reqwest::Client, base: &str) -> Option<String> {
505 let resp = client.get(base).timeout(HTTP_TIMEOUT).send().await.ok()?;
506 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES)
509 .await
510 .ok()?;
511 let html = String::from_utf8_lossy(&bytes);
512 html.split("csrfToken\":\"")
513 .nth(1)
514 .and_then(|s| s.split('"').next())
515 .filter(|t| !t.is_empty())
516 .map(|t| t.to_string())
517}
518
519async fn post_rpc(
520 client: &reqwest::Client,
521 base: &str,
522 csrf: Option<&str>,
523 rpc: &str,
524) -> Result<serde_json::Value> {
525 let mut req = client
526 .post(format!("{base}/{rpc}"))
527 .header("Content-Type", "application/json")
528 .body("{}")
529 .timeout(HTTP_TIMEOUT);
530 if let Some(token) = csrf {
531 req = req.header("x-codeium-csrf-token", token);
532 }
533 let resp = req.send().await?;
534
535 let status = resp.status();
536 let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
540 if !status.is_success() {
541 let body = String::from_utf8_lossy(&bytes).into_owned();
542 return Err(AppError::Http {
543 status: status.as_u16(),
544 body,
545 });
546 }
547
548 Ok(serde_json::from_slice(&bytes)?)
549}
550
551pub fn plan_from_status(v: &serde_json::Value) -> String {
552 v["userStatus"]["userTier"]["name"]
553 .as_str()
554 .or_else(|| v["userStatus"]["userTier"]["description"].as_str())
555 .or_else(|| v["userStatus"]["planStatus"]["planInfo"]["planName"].as_str())
556 .filter(|s| !s.is_empty())
557 .unwrap_or(DEFAULT_PLAN)
558 .to_string()
559}
560
561fn describe_bucket(group_name: &str, id: &str, window: Option<&str>) -> String {
572 let id = if id.is_empty() { "<unnamed>" } else { id };
573 let group = if group_name.is_empty() {
574 String::new()
575 } else {
576 format!(" in {group_name:?}")
577 };
578 match window {
579 Some(window) if !window.is_empty() => format!("{id} (window {window}){group}"),
580 _ => format!("{id}{group}"),
581 }
582}
583
584fn no_usable_bucket(seen: &[String]) -> AppError {
590 if seen.is_empty() {
591 return AppError::Other(
592 "antigravity: quota summary has no buckets at all — the running product may \
593 not have a quota for this account yet"
594 .into(),
595 );
596 }
597 AppError::Other(format!(
598 "antigravity: quota summary has no bucket in a window we recognise (5h or \
599 weekly, Gemini or Claude/GPT); it offered: {}",
600 seen.join(", ")
601 ))
602}
603
604pub fn parse_quota_summary(v: &serde_json::Value, plan: String) -> Result<AntigravitySnapshot> {
605 let groups = v["response"]["groups"]
606 .as_array()
607 .or_else(|| v["groups"].as_array())
608 .ok_or_else(|| AppError::Other("antigravity: quota summary has no groups".into()))?;
609
610 let mut gemini_5h = None;
611 let mut gemini_weekly = None;
612 let mut tp_5h = None;
613 let mut tp_weekly = None;
614 let mut seen: Vec<String> = Vec::new();
618
619 for group in groups {
620 let group_name = group["displayName"].as_str().unwrap_or_default();
621 let Some(buckets) = group["buckets"].as_array() else {
622 continue;
623 };
624 for bucket in buckets {
625 let id = bucket["bucketId"].as_str().unwrap_or_default();
626 seen.push(describe_bucket(group_name, id, bucket["window"].as_str()));
627 let window = bucket["window"].as_str().unwrap_or_default();
628 let is_weekly = if id.ends_with("weekly") || window == "weekly" {
629 true
630 } else if id.ends_with("5h") || window == "5h" {
631 false
632 } else {
633 continue;
636 };
637 let is_gemini = if id.starts_with("gemini") {
638 true
639 } else if id.starts_with("3p") {
640 false
641 } else if group_name.contains("Gemini") {
642 true
643 } else if group_name.contains("Claude") || group_name.contains("GPT") {
644 false
645 } else {
646 continue;
649 };
650
651 let (slot, slot_name) = match (is_gemini, is_weekly) {
652 (true, false) => (&mut gemini_5h, "Gemini 5h"),
653 (true, true) => (&mut gemini_weekly, "Gemini weekly"),
654 (false, false) => (&mut tp_5h, "Claude/GPT 5h"),
655 (false, true) => (&mut tp_weekly, "Claude/GPT weekly"),
656 };
657 let parsed = usage_window(bucket, is_weekly)?;
658 if slot.replace(parsed).is_some() {
659 return Err(AppError::Schema(format!(
660 "antigravity: duplicate {slot_name} bucket"
661 )));
662 }
663 }
664 }
665
666 if gemini_5h.is_none() && gemini_weekly.is_none() && tp_5h.is_none() && tp_weekly.is_none() {
670 return Err(no_usable_bucket(&seen));
671 }
672
673 Ok(AntigravitySnapshot {
674 plan,
675 account: String::new(),
678 source: AntigravitySource::Local,
679 session: gemini_5h,
680 weekly: gemini_weekly,
681 third_party_session: tp_5h,
682 third_party_weekly: tp_weekly,
683 })
684}
685
686fn usage_window(bucket: &serde_json::Value, is_weekly: bool) -> Result<UsageWindow> {
690 let remaining = bucket["remainingFraction"]
691 .as_f64()
692 .filter(|f| f.is_finite() && (0.0..=1.0).contains(f))
693 .ok_or_else(|| {
694 AppError::Schema(format!(
695 "antigravity: bucket {} has no valid remainingFraction in 0..=1",
696 bucket["bucketId"].as_str().unwrap_or("<unnamed>")
697 ))
698 })?;
699 Ok(UsageWindow {
700 utilization_pct: pct_used(remaining),
701 resets_at: parse_reset(&bucket["resetTime"], "quota resetTime")?,
702 window_duration: if is_weekly {
703 chrono::Duration::days(7)
704 } else {
705 chrono::Duration::hours(5)
706 },
707 })
708}
709
710fn pct_used(remaining_fraction: f64) -> i32 {
713 let used = (1.0 - remaining_fraction) * 100.0;
714 used.round() as i32
715}
716
717fn parse_reset(value: &serde_json::Value, field: &str) -> Result<Option<DateTime<Utc>>> {
718 match value {
719 serde_json::Value::Null => Ok(None),
720 serde_json::Value::String(s) => DateTime::parse_from_rfc3339(s)
721 .map(|dt| Some(dt.with_timezone(&Utc)))
722 .map_err(|_| AppError::Schema(format!("antigravity: invalid {field}"))),
723 _ => Err(AppError::Schema(format!(
724 "antigravity: {field} must be a timestamp or null"
725 ))),
726 }
727}
728
729fn candidate_bases() -> Vec<String> {
735 candidate_bases_with(
736 std::env::var("ANTIGRAVITY_LS_ADDRESS").ok().as_deref(),
737 discover_ls_ports(),
738 )
739}
740
741fn candidate_bases_with(override_addr: Option<&str>, discovered: Vec<u16>) -> Vec<String> {
744 let mut bases = Vec::new();
745 if let Some(base) = override_addr.and_then(normalize_base) {
746 bases.push(base);
747 }
748
749 for p in discovered {
755 let candidate = format!("http://127.0.0.1:{p}");
756 if !bases.contains(&candidate) {
757 bases.push(candidate);
758 }
759 }
760
761 bases
762}
763
764fn normalize_base(addr: &str) -> Option<String> {
773 let trimmed = addr.trim();
774 let (scheme, authority) = match trimmed.split_once("://") {
775 Some((scheme @ ("http" | "https"), rest)) => (scheme, rest),
776 _ => ("http", trimmed),
777 };
778 let authority = authority.trim_end_matches('/');
779 (!authority.is_empty()).then(|| format!("{scheme}://{authority}"))
780}
781
782fn is_antigravity_process(comm: &str, exe: Option<&str>) -> bool {
789 let comm = comm.trim().to_lowercase();
790 let comm = comm.strip_suffix(".exe").unwrap_or(&comm);
791 if comm.contains("language_server") || comm == "agy" || comm == "antigravity" {
792 return true;
793 }
794 exe.is_some_and(|p| {
795 let p = p.to_lowercase().replace('\\', "/");
796 let p = p.strip_suffix(".exe").unwrap_or(&p);
797 p.contains("antigravity") || p.ends_with("/agy")
798 })
799}
800
801#[cfg(any(test, target_os = "linux", target_os = "macos", target_os = "windows"))]
831fn probe_order(per_pid: std::collections::BTreeMap<u32, Vec<u16>>) -> Vec<u16> {
832 let groups: Vec<Vec<u16>> = per_pid
833 .into_values()
834 .map(|mut group| {
835 group.sort_unstable_by(|a, b| b.cmp(a));
836 group.dedup();
841 group
842 })
843 .collect();
844
845 let mut ports: Vec<u16> = Vec::new();
846 for rank in 0..groups.iter().map(Vec::len).max().unwrap_or(0) {
847 for port in groups.iter().filter_map(|group| group.get(rank)) {
848 if !ports.contains(port) {
849 ports.push(*port);
850 }
851 }
852 }
853 ports
854}
855
856#[cfg(target_os = "linux")]
863pub(crate) fn discover_ls_ports() -> Vec<u16> {
864 use std::collections::{BTreeMap, HashMap};
865
866 let mut owners: HashMap<u64, u32> = HashMap::new();
869 let Ok(entries) = std::fs::read_dir("/proc") else {
870 return Vec::new();
871 };
872 for entry in entries.flatten() {
873 let pid_dir = entry.path();
874 let Some(pid) = pid_dir
875 .file_name()
876 .and_then(|name| name.to_str())
877 .and_then(|name| name.parse::<u32>().ok())
878 else {
879 continue;
880 };
881 let Ok(comm) = std::fs::read_to_string(pid_dir.join("comm")) else {
882 continue;
883 };
884 let exe = std::fs::read_link(pid_dir.join("exe")).ok();
885 if !is_antigravity_process(&comm, exe.as_deref().and_then(|p| p.to_str())) {
886 continue;
887 }
888 let Ok(fds) = std::fs::read_dir(pid_dir.join("fd")) else {
889 continue;
890 };
891 for fd in fds.flatten() {
892 let Ok(target) = std::fs::read_link(fd.path()) else {
893 continue;
894 };
895 if let Some(ino) = target
896 .to_str()
897 .and_then(|s| s.strip_prefix("socket:["))
898 .and_then(|s| s.strip_suffix(']'))
899 .and_then(|s| s.parse::<u64>().ok())
900 {
901 owners.insert(ino, pid);
902 }
903 }
904 }
905
906 if owners.is_empty() {
907 return Vec::new();
908 }
909
910 let mut per_pid: BTreeMap<u32, Vec<u16>> = BTreeMap::new();
911 for table in ["/proc/net/tcp", "/proc/net/tcp6"] {
912 let Ok(contents) = std::fs::read_to_string(table) else {
913 continue;
914 };
915 for line in contents.lines().skip(1) {
916 if let Some((port, ino)) = parse_proc_net_line(line)
917 && let Some(&pid) = owners.get(&ino)
918 {
919 per_pid.entry(pid).or_default().push(port);
920 }
921 }
922 }
923 probe_order(per_pid)
924}
925
926#[cfg(target_os = "macos")]
933pub(crate) fn discover_ls_ports() -> Vec<u16> {
934 let Ok(output) = std::process::Command::new("lsof")
935 .args(["-nP", "-iTCP", "-sTCP:LISTEN", "-F", "pcn"])
936 .output()
937 else {
938 return Vec::new();
939 };
940 parse_lsof_pcn(&String::from_utf8_lossy(&output.stdout))
944}
945
946#[cfg(any(test, target_os = "macos"))]
951fn parse_lsof_pcn(output: &str) -> Vec<u16> {
952 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
953 let mut pid = None;
956 let mut owner = None;
957 for line in output.lines() {
958 let Some(rest) = line.get(1..) else { continue };
959 match line.as_bytes().first() {
960 Some(b'p') => {
961 pid = rest.parse::<u32>().ok();
962 owner = None;
963 }
964 Some(b'c') => owner = pid.filter(|_| is_antigravity_process(rest, None)),
965 Some(b'n') => {
966 if let Some(pid) = owner
967 && let Some(port) = rest.rsplit(':').next().and_then(|p| p.parse::<u16>().ok())
968 {
969 per_pid.entry(pid).or_default().push(port);
970 }
971 }
972 _ => {}
973 }
974 }
975 probe_order(per_pid)
976}
977
978#[cfg(any(test, target_os = "windows"))]
979#[derive(Debug, Clone, Copy, PartialEq, Eq)]
980struct WindowsTcpRow {
981 local_addr: [u8; 4],
982 local_port: u32,
983 pid: u32,
984}
985
986#[cfg(any(test, target_os = "windows"))]
987fn decode_windows_process_name(raw: &[u16]) -> String {
988 let end = raw.iter().position(|&unit| unit == 0).unwrap_or(raw.len());
989 String::from_utf16_lossy(&raw[..end])
990}
991
992#[cfg(any(test, target_os = "windows"))]
993fn matching_windows_process_ids(processes: &[(u32, String)]) -> std::collections::HashSet<u32> {
994 processes
995 .iter()
996 .filter(|(_, name)| is_antigravity_process(name, None))
997 .map(|(pid, _)| *pid)
998 .collect()
999}
1000
1001#[cfg(any(test, target_os = "windows"))]
1004fn matching_windows_ports(
1005 pids: &std::collections::HashSet<u32>,
1006 rows: &[WindowsTcpRow],
1007) -> Vec<u16> {
1008 let mut per_pid: std::collections::BTreeMap<u32, Vec<u16>> = std::collections::BTreeMap::new();
1009 for row in rows {
1010 if !pids.contains(&row.pid) || row.local_addr != [127, 0, 0, 1] {
1011 continue;
1012 }
1013 let port = u16::from_be((row.local_port & u32::from(u16::MAX)) as u16);
1014 if port != 0 {
1015 per_pid.entry(row.pid).or_default().push(port);
1016 }
1017 }
1018 probe_order(per_pid)
1019}
1020
1021#[cfg(any(test, target_os = "windows"))]
1022fn checked_windows_row_count(
1023 buffer_len: usize,
1024 rows_offset: usize,
1025 row_size: usize,
1026 declared: usize,
1027) -> Option<usize> {
1028 let rows_len = row_size.checked_mul(declared)?;
1029 let end = rows_offset.checked_add(rows_len)?;
1030 (row_size != 0 && end <= buffer_len).then_some(declared)
1031}
1032
1033#[cfg(target_os = "windows")]
1034struct WindowsHandle(windows_sys::Win32::Foundation::HANDLE);
1035
1036#[cfg(target_os = "windows")]
1037impl Drop for WindowsHandle {
1038 fn drop(&mut self) {
1039 unsafe {
1040 windows_sys::Win32::Foundation::CloseHandle(self.0);
1041 }
1042 }
1043}
1044
1045#[cfg(target_os = "windows")]
1046fn windows_processes() -> Vec<(u32, String)> {
1047 use std::mem::size_of;
1048 use windows_sys::Win32::Foundation::INVALID_HANDLE_VALUE;
1049 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
1050 CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
1051 TH32CS_SNAPPROCESS,
1052 };
1053
1054 let handle = unsafe { CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) };
1055 if handle == INVALID_HANDLE_VALUE {
1056 return Vec::new();
1057 }
1058 let snapshot = WindowsHandle(handle);
1059 let mut entry = PROCESSENTRY32W::default();
1060 let Ok(entry_size) = u32::try_from(size_of::<PROCESSENTRY32W>()) else {
1061 return Vec::new();
1062 };
1063 entry.dwSize = entry_size;
1064 if unsafe { Process32FirstW(snapshot.0, &mut entry) } == 0 {
1065 return Vec::new();
1066 }
1067
1068 let mut processes = Vec::new();
1069 loop {
1070 processes.push((
1071 entry.th32ProcessID,
1072 decode_windows_process_name(&entry.szExeFile),
1073 ));
1074 if unsafe { Process32NextW(snapshot.0, &mut entry) } == 0 {
1075 break;
1076 }
1077 }
1078 processes
1079}
1080
1081#[cfg(target_os = "windows")]
1082fn parse_windows_tcp_rows(buffer: &[u32], used_bytes: usize) -> Vec<WindowsTcpRow> {
1083 use std::mem::{offset_of, size_of, size_of_val};
1084 use windows_sys::Win32::NetworkManagement::IpHelper::{
1085 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
1086 };
1087
1088 let available = used_bytes.min(size_of_val(buffer));
1089 let rows_offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
1090 if available < size_of::<u32>() || available < rows_offset {
1091 return Vec::new();
1092 }
1093 let base = buffer.as_ptr().cast::<u8>();
1094 let declared = unsafe { base.cast::<u32>().read_unaligned() } as usize;
1095 if checked_windows_row_count(
1096 available,
1097 rows_offset,
1098 size_of::<MIB_TCPROW_OWNER_PID>(),
1099 declared,
1100 )
1101 .is_none()
1102 {
1103 return Vec::new();
1104 }
1105
1106 let rows = unsafe { base.add(rows_offset).cast::<MIB_TCPROW_OWNER_PID>() };
1107 (0..declared)
1108 .map(|index| unsafe { rows.add(index).read_unaligned() })
1109 .map(|row| WindowsTcpRow {
1110 local_addr: row.dwLocalAddr.to_ne_bytes(),
1111 local_port: row.dwLocalPort,
1112 pid: row.dwOwningPid,
1113 })
1114 .collect()
1115}
1116
1117#[cfg(target_os = "windows")]
1118fn windows_tcp_rows() -> Vec<WindowsTcpRow> {
1119 use std::mem::size_of;
1120 use std::ptr::null_mut;
1121 use windows_sys::Win32::Foundation::ERROR_INSUFFICIENT_BUFFER;
1122 use windows_sys::Win32::NetworkManagement::IpHelper::{
1123 GetExtendedTcpTable, TCP_TABLE_OWNER_PID_LISTENER,
1124 };
1125 use windows_sys::Win32::Networking::WinSock::AF_INET;
1126
1127 let mut size = 0u32;
1128 let status = unsafe {
1129 GetExtendedTcpTable(
1130 null_mut(),
1131 &mut size,
1132 0,
1133 u32::from(AF_INET),
1134 TCP_TABLE_OWNER_PID_LISTENER,
1135 0,
1136 )
1137 };
1138 if status != ERROR_INSUFFICIENT_BUFFER {
1139 return Vec::new();
1140 }
1141
1142 for _ in 0..3 {
1143 let Some(words) = (size as usize)
1144 .checked_add(size_of::<u32>() - 1)
1145 .map(|bytes| bytes / size_of::<u32>())
1146 else {
1147 return Vec::new();
1148 };
1149 if words == 0 {
1150 return Vec::new();
1151 }
1152 let mut buffer = Vec::<u32>::new();
1153 if buffer.try_reserve_exact(words).is_err() {
1154 return Vec::new();
1155 }
1156 buffer.resize(words, 0);
1157 let mut used = size;
1158 let status = unsafe {
1159 GetExtendedTcpTable(
1160 buffer.as_mut_ptr().cast(),
1161 &mut used,
1162 0,
1163 u32::from(AF_INET),
1164 TCP_TABLE_OWNER_PID_LISTENER,
1165 0,
1166 )
1167 };
1168 if status == ERROR_INSUFFICIENT_BUFFER {
1169 size = used;
1170 continue;
1171 }
1172 if status != 0 {
1173 return Vec::new();
1174 }
1175 return parse_windows_tcp_rows(&buffer, used as usize);
1176 }
1177 Vec::new()
1178}
1179
1180#[cfg(target_os = "windows")]
1181pub(crate) fn discover_ls_ports() -> Vec<u16> {
1182 let pids = matching_windows_process_ids(&windows_processes());
1183 if pids.is_empty() {
1184 return Vec::new();
1185 }
1186 matching_windows_ports(&pids, &windows_tcp_rows())
1187}
1188
1189#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))]
1190pub(crate) fn discover_ls_ports() -> Vec<u16> {
1191 Vec::new()
1192}
1193
1194#[cfg(target_os = "linux")]
1197fn parse_proc_net_line(line: &str) -> Option<(u16, u64)> {
1198 let cols: Vec<&str> = line.split_whitespace().collect();
1199 if cols.len() < 10 {
1200 return None;
1201 }
1202 if cols[3] != "0A" {
1204 return None;
1205 }
1206 let port = u16::from_str_radix(cols[1].split(':').nth(1)?, 16).ok()?;
1207 let inode = cols[9].parse::<u64>().ok()?;
1208 Some((port, inode))
1209}
1210
1211fn fallback_silent(cache: &Cache, now: DateTime<Utc>, original: AppError) -> Result<FetchOutcome> {
1216 crate::outcome::fallback(cache, None, original, |bytes| {
1217 parse_cache_at(bytes, None, now)
1218 })
1219}
1220
1221fn fallback_with_error(
1226 cache: &Cache,
1227 last_error: Option<(u16, String)>,
1228 reason: AppError,
1229 now: DateTime<Utc>,
1230) -> Result<FetchOutcome> {
1231 crate::outcome::fallback(cache, last_error, reason, |bytes| {
1232 parse_cache_at(bytes, None, now)
1233 })
1234}
1235
1236fn reuse_cache(
1237 bytes: Vec<u8>,
1238 cache: &Cache,
1239 stale: bool,
1240 account: Option<&str>,
1241 now: DateTime<Utc>,
1242) -> Result<FetchOutcome> {
1243 let snap = parse_cache_at(&bytes, account, now)?;
1244 Ok(crate::outcome::Outcome::cached(snap, cache, stale))
1245}
1246
1247pub fn parse_cache(bytes: &[u8], account: Option<&str>) -> Result<AntigravitySnapshot> {
1253 parse_cache_at(bytes, account, Utc::now())
1254}
1255
1256fn expired_window(snap: &AntigravitySnapshot, now: DateTime<Utc>) -> Option<&'static str> {
1266 [
1267 ("Gemini 5h", snap.session.as_ref()),
1268 ("Gemini weekly", snap.weekly.as_ref()),
1269 ("Claude & GPT OSS 5h", snap.third_party_session.as_ref()),
1270 ("Claude & GPT OSS weekly", snap.third_party_weekly.as_ref()),
1271 ]
1272 .into_iter()
1273 .find(|(_, w)| w.and_then(|w| w.resets_at).is_some_and(|r| r <= now))
1274 .map(|(name, _)| name)
1275}
1276
1277pub fn parse_cache_at(
1278 bytes: &[u8],
1279 account: Option<&str>,
1280 now: DateTime<Utc>,
1281) -> Result<AntigravitySnapshot> {
1282 let v: serde_json::Value = serde_json::from_slice(bytes)?;
1283
1284 let cached_account = v.get("account").and_then(serde_json::Value::as_str);
1285 if let Some(expected) = account
1286 && cached_account != Some(expected)
1287 {
1288 return Err(AppError::Schema(
1289 "antigravity cache belongs to a different account; refetching".into(),
1290 ));
1291 }
1292
1293 let cached_pct = |pct_key: &'static str| -> Result<Option<i32>> {
1303 match v.get(pct_key) {
1304 None => Err(AppError::Schema(format!(
1305 "antigravity: cached payload is missing {pct_key}"
1306 ))),
1307 Some(serde_json::Value::Null) => Ok(None),
1308 Some(value) => value
1309 .as_i64()
1310 .filter(|pct| (0..=100).contains(pct))
1311 .map(|pct| Some(pct as i32))
1312 .ok_or_else(|| {
1313 AppError::Schema(format!(
1314 "antigravity: cached {pct_key} must be an integer in 0..=100"
1315 ))
1316 }),
1317 }
1318 };
1319
1320 let optional = |pct_key: &'static str, reset_key: &str, weekly: bool| {
1321 let Some(pct) = cached_pct(pct_key)? else {
1322 return Ok(None);
1323 };
1324 Ok::<_, AppError>(Some(UsageWindow {
1325 utilization_pct: pct,
1326 resets_at: parse_reset(&v[reset_key], reset_key)?,
1327 window_duration: if weekly {
1328 chrono::Duration::days(7)
1329 } else {
1330 chrono::Duration::hours(5)
1331 },
1332 }))
1333 };
1334
1335 let snap = AntigravitySnapshot {
1336 plan: v["plan"].as_str().unwrap_or(DEFAULT_PLAN).to_string(),
1337 account: cached_account.unwrap_or_default().to_string(),
1338 source: v["source"]
1340 .as_str()
1341 .and_then(AntigravitySource::parse)
1342 .unwrap_or_default(),
1343 session: optional("session_pct", "session_reset", false)?,
1344 weekly: optional("weekly_pct", "weekly_reset", true)?,
1345 third_party_session: optional("tp_session_pct", "tp_session_reset", false)?,
1346 third_party_weekly: optional("tp_weekly_pct", "tp_weekly_reset", true)?,
1347 };
1348
1349 if snap.session.is_none()
1352 && snap.weekly.is_none()
1353 && snap.third_party_session.is_none()
1354 && snap.third_party_weekly.is_none()
1355 {
1356 return Err(AppError::Schema(
1357 "antigravity cache holds no usable window; refetching".into(),
1358 ));
1359 }
1360
1361 if let Some(window) = expired_window(&snap, now) {
1362 return Err(AppError::Schema(format!(
1363 "antigravity cache is past its {window} reset; refetching"
1364 )));
1365 }
1366 Ok(snap)
1367}
1368
1369pub fn snap_to_json(snap: &AntigravitySnapshot) -> serde_json::Value {
1370 serde_json::json!({
1371 "plan": snap.plan,
1372 "account": snap.account,
1373 "source": snap.source.as_str(),
1374 "session_pct": snap.session.as_ref().map(|w| w.utilization_pct),
1375 "session_reset": snap.session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1376 "weekly_pct": snap.weekly.as_ref().map(|w| w.utilization_pct),
1377 "weekly_reset": snap.weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1378 "tp_session_pct": snap.third_party_session.as_ref().map(|w| w.utilization_pct),
1379 "tp_session_reset": snap.third_party_session.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1380 "tp_weekly_pct": snap.third_party_weekly.as_ref().map(|w| w.utilization_pct),
1381 "tp_weekly_reset": snap.third_party_weekly.as_ref().and_then(|w| w.resets_at.map(|dt| dt.to_rfc3339())),
1382 })
1383}
1384
1385#[cfg(test)]
1386mod tests {
1387 use super::*;
1388
1389 const QUOTA_JSON: &str = r#"{
1393 "response": {
1394 "groups": [
1395 {
1396 "displayName": "Gemini Models",
1397 "buckets": [
1398 {"bucketId": "gemini-weekly", "displayName": "Weekly Limit",
1399 "window": "weekly", "remainingFraction": 0.9191212,
1400 "resetTime": "2026-07-28T17:39:58Z"},
1401 {"bucketId": "gemini-5h", "displayName": "Five Hour Limit",
1402 "window": "5h", "remainingFraction": 0.5672253,
1403 "resetTime": "2026-07-22T17:47:00Z"}
1404 ]
1405 },
1406 {
1407 "displayName": "Claude and GPT models",
1408 "buckets": [
1409 {"bucketId": "3p-weekly", "window": "weekly",
1410 "remainingFraction": 1, "resetTime": "2026-07-29T12:47:00Z"},
1411 {"bucketId": "3p-5h", "window": "5h",
1412 "remainingFraction": 0.25, "resetTime": "2026-07-22T17:47:00Z"}
1413 ]
1414 }
1415 ]
1416 }
1417 }"#;
1418
1419 fn now() -> DateTime<Utc> {
1422 DateTime::parse_from_rfc3339("2026-07-22T12:00:00Z")
1423 .unwrap()
1424 .with_timezone(&Utc)
1425 }
1426
1427 fn parsed() -> AntigravitySnapshot {
1428 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1429 parse_quota_summary(&v, "Google AI Pro".into()).unwrap()
1430 }
1431
1432 #[test]
1433 fn quota_summary_maps_four_distinct_windows() {
1434 let snap = parsed();
1435 assert_eq!(snap.plan, "Google AI Pro");
1436 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
1438 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
1439 assert_eq!(
1440 snap.third_party_session.as_ref().unwrap().utilization_pct,
1441 75
1442 );
1443 assert_eq!(snap.third_party_weekly.as_ref().unwrap().utilization_pct, 0);
1444 }
1445
1446 #[test]
1447 fn each_window_keeps_its_own_reset_time() {
1448 let snap = parsed();
1449 let at = |s: &str| Some(DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc));
1450 assert_eq!(
1451 snap.session.as_ref().unwrap().resets_at,
1452 at("2026-07-22T17:47:00Z")
1453 );
1454 assert_eq!(
1455 snap.weekly.as_ref().unwrap().resets_at,
1456 at("2026-07-28T17:39:58Z")
1457 );
1458 assert_eq!(
1459 snap.third_party_weekly.as_ref().unwrap().resets_at,
1460 at("2026-07-29T12:47:00Z")
1461 );
1462 assert_ne!(
1464 snap.session.as_ref().unwrap().resets_at,
1465 snap.weekly.as_ref().unwrap().resets_at
1466 );
1467 }
1468
1469 #[test]
1470 fn window_durations_match_their_bucket() {
1471 let snap = parsed();
1472 assert_eq!(
1473 snap.session.as_ref().unwrap().window_duration,
1474 chrono::Duration::hours(5)
1475 );
1476 assert_eq!(
1477 snap.weekly.as_ref().unwrap().window_duration,
1478 chrono::Duration::days(7)
1479 );
1480 assert_eq!(
1481 snap.third_party_weekly.as_ref().unwrap().window_duration,
1482 chrono::Duration::days(7)
1483 );
1484 }
1485
1486 #[test]
1487 fn groups_are_matched_by_display_name_when_bucket_ids_change() {
1488 let v: serde_json::Value = serde_json::from_str(
1489 r#"{"response":{"groups":[
1490 {"displayName":"Gemini Models","buckets":[
1491 {"bucketId":"x1","window":"5h","remainingFraction":0.5,"resetTime":"2026-07-22T17:47:00Z"},
1492 {"bucketId":"x2","window":"weekly","remainingFraction":0.9,"resetTime":"2026-07-28T17:39:58Z"}]},
1493 {"displayName":"Claude and GPT models","buckets":[
1494 {"bucketId":"y1","window":"5h","remainingFraction":0.0,"resetTime":"2026-07-22T17:47:00Z"}]}
1495 ]}}"#,
1496 )
1497 .unwrap();
1498 let snap = parse_quota_summary(&v, "Pro".into()).unwrap();
1499 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 50);
1500 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 10);
1501 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 100);
1502 assert!(snap.third_party_weekly.is_none());
1503 }
1504
1505 #[test]
1506 fn duplicate_or_unclassified_buckets_cannot_overwrite_a_slot() {
1507 let duplicate: serde_json::Value = serde_json::from_str(
1508 r#"{"response":{"groups":[{"displayName":"Gemini Models","buckets":[
1509 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1510 {"bucketId":"gemini-5h-copy","window":"5h","remainingFraction":0.1},
1511 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8}
1512 ]}]}}"#,
1513 )
1514 .unwrap();
1515 let err = parse_quota_summary(&duplicate, "Pro".into()).unwrap_err();
1516 assert!(err.to_string().contains("duplicate Gemini 5h"), "{err}");
1517
1518 let unrelated: serde_json::Value = serde_json::from_str(
1521 r#"{"response":{"groups":[
1522 {"displayName":"Gemini Models","buckets":[
1523 {"bucketId":"gemini-5h","window":"5h","remainingFraction":0.9},
1524 {"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.8},
1525 {"bucketId":"gemini-monthly","window":"monthly","remainingFraction":0.7}
1526 ]},
1527 {"displayName":"Future Models","buckets":[
1528 {"bucketId":"future-5h","window":"5h","remainingFraction":0.1}
1529 ]}
1530 ]}}"#,
1531 )
1532 .unwrap();
1533 let snap = parse_quota_summary(&unrelated, "Pro".into()).unwrap();
1534 assert!(snap.third_party_session.is_none());
1535 assert!(snap.third_party_weekly.is_none());
1536 }
1537
1538 #[test]
1541 fn a_bucket_without_a_usable_fraction_is_rejected() {
1542 for bad in [r#""oops""#, "null", "-0.01", "1.01"] {
1543 let v: serde_json::Value = serde_json::from_str(&format!(
1544 r#"{{"response":{{"groups":[{{"displayName":"Gemini Models","buckets":[
1545 {{"bucketId":"gemini-5h","window":"5h","remainingFraction":{bad}}},
1546 {{"bucketId":"gemini-weekly","window":"weekly","remainingFraction":0.9}}]}}]}}}}"#
1547 ))
1548 .unwrap();
1549 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1550 assert!(err.to_string().contains("gemini-5h"), "{bad}: {err}");
1551 }
1552 }
1553
1554 #[test]
1555 fn malformed_present_reset_is_rejected_instead_of_disabling_expiry() {
1556 for bad in [serde_json::json!("not-a-time"), serde_json::json!(42)] {
1557 let mut v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
1558 v["response"]["groups"][0]["buckets"][0]["resetTime"] = bad;
1559 let err = parse_quota_summary(&v, "Pro".into()).unwrap_err();
1560 assert!(err.to_string().contains("resetTime"), "{err}");
1561 }
1562 }
1563
1564 #[test]
1568 fn a_weekly_only_snapshot_round_trips_through_the_cache() {
1569 let mut snap = parsed();
1570 snap.session = None;
1571 snap.third_party_session = None;
1572
1573 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1574 let back = parse_cache_at(&bytes, None, now()).expect("weekly-only cache is usable");
1575
1576 assert!(back.session.is_none());
1577 assert_eq!(
1578 back.weekly.as_ref().unwrap().utilization_pct,
1579 snap.weekly.as_ref().unwrap().utilization_pct
1580 );
1581 }
1582
1583 #[test]
1588 fn a_product_reporting_only_weekly_buckets_still_renders_them() {
1589 let summary = serde_json::json!({
1590 "groups": [
1591 {
1592 "displayName": "Gemini Models",
1593 "buckets": [{
1594 "bucketId": "gemini-weekly", "window": "weekly",
1595 "remainingFraction": 0.42,
1596 }],
1597 },
1598 {
1599 "displayName": "Claude and GPT models",
1600 "buckets": [{
1601 "bucketId": "3p-weekly", "window": "weekly",
1602 "remainingFraction": 0.9,
1603 }],
1604 },
1605 ],
1606 });
1607
1608 let snap = parse_quota_summary(&summary, "Pro".into()).expect("weekly-only is usable");
1609
1610 assert!(snap.session.is_none(), "no 5h bucket arrived");
1611 assert!(snap.third_party_session.is_none());
1612 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 58);
1613 assert_eq!(
1614 snap.third_party_weekly.as_ref().unwrap().utilization_pct,
1615 10
1616 );
1617 }
1618
1619 #[test]
1622 fn a_product_reporting_only_five_hour_buckets_still_renders_them() {
1623 let summary = serde_json::json!({
1624 "groups": [{
1625 "displayName": "Gemini Models",
1626 "buckets": [{
1627 "bucketId": "gemini-5h", "window": "5h", "remainingFraction": 0.25,
1628 }],
1629 }],
1630 });
1631
1632 let snap = parse_quota_summary(&summary, "Pro".into()).expect("5h-only is usable");
1633
1634 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 75);
1635 assert!(snap.weekly.is_none());
1636 }
1637
1638 #[test]
1641 fn a_summary_with_no_recognisable_bucket_errors_and_names_what_it_had() {
1642 let summary = serde_json::json!({
1643 "groups": [{
1644 "displayName": "Gemini Models",
1645 "buckets": [{
1646 "bucketId": "gemini-daily", "window": "daily", "remainingFraction": 0.9,
1647 }],
1648 }],
1649 });
1650
1651 let rendered = parse_quota_summary(&summary, "Pro".into())
1652 .expect_err("an unrecognised cadence alone is not a snapshot")
1653 .to_string();
1654
1655 assert!(
1656 rendered.contains("no bucket in a window we recognise"),
1657 "{rendered}"
1658 );
1659 assert!(rendered.contains("gemini-daily"), "{rendered}");
1660 assert!(rendered.contains("window daily"), "{rendered}");
1661 }
1662
1663 #[test]
1666 fn a_summary_with_no_buckets_at_all_says_that_rather_than_listing_nothing() {
1667 let summary = serde_json::json!({
1668 "groups": [{"displayName": "Gemini", "buckets": []}],
1669 });
1670
1671 let rendered = parse_quota_summary(&summary, "Pro".into())
1672 .expect_err("no buckets is an error")
1673 .to_string();
1674
1675 assert!(rendered.contains("no buckets at all"), "{rendered}");
1676 assert!(!rendered.contains("it offered:"), "{rendered}");
1677 }
1678
1679 #[test]
1682 fn buckets_without_an_id_are_still_named_in_the_error() {
1683 let summary = serde_json::json!({
1684 "groups": [{"displayName": "", "buckets": [{"remainingFraction": 0.5}]}],
1685 });
1686
1687 let rendered = parse_quota_summary(&summary, "Pro".into())
1688 .expect_err("an unusable summary is an error")
1689 .to_string();
1690
1691 assert!(rendered.contains("<unnamed>"), "{rendered}");
1692 }
1693
1694 #[test]
1695 fn missing_gemini_buckets_is_an_error_not_a_zero_bar() {
1696 let v: serde_json::Value = serde_json::from_str(r#"{"response":{"groups":[]}}"#).unwrap();
1697 assert!(parse_quota_summary(&v, "Pro".into()).is_err());
1698 }
1699
1700 #[test]
1701 fn cache_round_trip_preserves_every_window() {
1702 let snap = parsed();
1703 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1704 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1705 }
1706
1707 #[test]
1711 fn a_truncated_cached_payload_is_rejected_not_zeroed() {
1712 let full = snap_to_json(&parsed());
1713 for missing in ["session_pct", "weekly_pct"] {
1714 let mut v = full.clone();
1715 v.as_object_mut().unwrap().remove(missing);
1716 let bytes = serde_json::to_vec(&v).unwrap();
1717 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1718 assert!(err.to_string().contains(missing), "{missing}: {err}");
1719 }
1720 assert!(parse_cache_at(b"{}", None, now()).is_err());
1722 }
1723
1724 #[test]
1725 fn cached_percentages_are_range_checked_before_narrowing() {
1726 let full = snap_to_json(&parsed());
1727 for (key, bad) in [
1728 ("session_pct", serde_json::json!(-1)),
1729 ("weekly_pct", serde_json::json!(101)),
1730 ("session_pct", serde_json::json!(i64::MAX)),
1731 ("tp_session_pct", serde_json::json!("75")),
1732 ] {
1733 let mut v = full.clone();
1734 v[key] = bad;
1735 let bytes = serde_json::to_vec(&v).unwrap();
1736 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1737 assert!(err.to_string().contains(key), "{key}: {err}");
1738 }
1739 }
1740
1741 #[test]
1742 fn malformed_cached_reset_is_rejected_instead_of_served_for_a_week() {
1743 let mut v = snap_to_json(&parsed());
1744 v["session_reset"] = serde_json::json!("not-a-time");
1745 let bytes = serde_json::to_vec(&v).unwrap();
1746 let err = parse_cache_at(&bytes, None, now()).unwrap_err();
1747 assert!(err.to_string().contains("session_reset"), "{err}");
1748 }
1749
1750 #[test]
1755 fn a_cache_past_its_reset_is_refused() {
1756 let bytes = serde_json::to_vec(&snap_to_json(&parsed())).unwrap();
1757 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1758
1759 assert!(parse_cache_at(&bytes, None, now()).is_ok());
1761 assert!(parse_cache_at(&bytes, None, at("2026-07-22T17:46:59Z")).is_ok());
1763
1764 let err = parse_cache_at(&bytes, None, at("2026-07-22T17:47:00Z")).unwrap_err();
1766 assert!(err.to_string().contains("5h"), "{err}");
1767
1768 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_err());
1770 }
1771
1772 #[test]
1775 fn expiry_names_the_window_that_rolled_over() {
1776 let mut snap = parsed();
1777 snap.session.as_mut().unwrap().resets_at = None;
1779 snap.third_party_session = None;
1780 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1781 let at = |s: &str| DateTime::parse_from_rfc3339(s).unwrap().with_timezone(&Utc);
1782
1783 assert!(parse_cache_at(&bytes, None, at("2026-07-23T09:00:00Z")).is_ok());
1785
1786 let err = parse_cache_at(&bytes, None, at("2026-07-28T18:00:00Z")).unwrap_err();
1788 assert!(err.to_string().contains("Gemini weekly"), "{err}");
1789 }
1790
1791 #[test]
1793 fn a_window_without_a_reset_never_expires() {
1794 let mut snap = parsed();
1795 for w in [&mut snap.session, &mut snap.weekly].into_iter().flatten() {
1796 w.resets_at = None;
1797 }
1798 snap.third_party_session = None;
1799 snap.third_party_weekly = None;
1800 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1801 let far_future = DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
1802 .unwrap()
1803 .with_timezone(&Utc);
1804 assert!(parse_cache_at(&bytes, None, far_future).is_ok());
1805 }
1806
1807 #[test]
1809 fn a_cache_from_another_account_is_rejected() {
1810 let mut snap = parsed();
1811 snap.account = "acct:aaaa".into();
1812 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1813
1814 assert!(parse_cache_at(&bytes, Some("acct:bbbb"), now()).is_err());
1815 assert_eq!(
1816 parse_cache_at(&bytes, Some("acct:aaaa"), now()).unwrap(),
1817 snap
1818 );
1819
1820 let mut legacy = snap_to_json(&snap);
1822 legacy.as_object_mut().unwrap().remove("account");
1823 let legacy = serde_json::to_vec(&legacy).unwrap();
1824 assert!(parse_cache_at(&legacy, Some("acct:aaaa"), now()).is_err());
1825 }
1826
1827 #[test]
1830 fn an_unverifiable_cache_is_served_rather_than_discarded() {
1831 let mut snap = parsed();
1832 snap.account = "acct:aaaa".into();
1833 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1834 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1835 }
1836
1837 #[test]
1838 fn account_key_fingerprints_rather_than_storing_the_address() {
1839 let with = |email: &str| account_key(&serde_json::json!({"userStatus": {"email": email}}));
1840 let a = with("someone@example.com");
1841 assert!(!a.contains("someone"), "{a}");
1842 assert!(!a.contains('@'), "{a}");
1843 assert_eq!(a, with("someone@example.com"), "must be stable");
1844 assert_ne!(a, with("other@example.com"));
1845 let unknown = account_key(&serde_json::json!({}));
1847 assert_eq!(unknown, account_key(&serde_json::json!({"userStatus": {}})));
1848 assert_ne!(unknown, a);
1849 }
1850
1851 #[test]
1854 fn absent_third_party_windows_are_not_treated_as_corruption() {
1855 let mut snap = parsed();
1856 snap.third_party_session = None;
1857 snap.third_party_weekly = None;
1858 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1859 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1860 }
1861
1862 #[test]
1863 fn cache_round_trip_preserves_absent_third_party_windows() {
1864 let mut snap = parsed();
1865 snap.third_party_session = None;
1866 snap.third_party_weekly = None;
1867 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
1868 assert_eq!(parse_cache_at(&bytes, None, now()).unwrap(), snap);
1869 }
1870
1871 #[test]
1872 fn pct_used_inverts_valid_fractions() {
1873 assert_eq!(pct_used(1.0), 0);
1874 assert_eq!(pct_used(0.0), 100);
1875 assert_eq!(pct_used(0.5), 50);
1876 }
1877
1878 #[test]
1879 fn plan_falls_back_through_the_status_payload() {
1880 let tier: serde_json::Value =
1881 serde_json::from_str(r#"{"userStatus":{"userTier":{"name":"Google AI Pro"}}}"#)
1882 .unwrap();
1883 assert_eq!(plan_from_status(&tier), "Google AI Pro");
1884
1885 let plan_only: serde_json::Value = serde_json::from_str(
1886 r#"{"userStatus":{"planStatus":{"planInfo":{"planName":"Pro"}}}}"#,
1887 )
1888 .unwrap();
1889 assert_eq!(plan_from_status(&plan_only), "Pro");
1890
1891 let empty: serde_json::Value = serde_json::from_str("{}").unwrap();
1892 assert_eq!(plan_from_status(&empty), DEFAULT_PLAN);
1893 }
1894
1895 #[cfg(target_os = "linux")]
1896 #[test]
1897 fn proc_net_parser_keeps_only_listening_rows() {
1898 let listen = " 0: 0100007F:975B 00000000:0000 0A 00000000:00000000 \
1899 00:00000000 00000000 1000 0 123456 1 0000 100 0";
1900 assert_eq!(parse_proc_net_line(listen), Some((38747, 123456)));
1901
1902 let established = " 1: 0100007F:975B 0100007F:A1B2 01 00000000:00000000 \
1903 00:00000000 00000000 1000 0 123457 1 0000 100 0";
1904 assert_eq!(parse_proc_net_line(established), None);
1905
1906 assert_eq!(parse_proc_net_line("garbage"), None);
1907 }
1908
1909 #[test]
1910 fn explicit_address_comes_first_and_gets_a_scheme() {
1911 assert_eq!(
1912 candidate_bases_with(Some("127.0.0.1:1234"), vec![5678]),
1913 vec![
1914 "http://127.0.0.1:1234".to_string(),
1915 "http://127.0.0.1:5678".to_string(),
1916 ]
1917 );
1918 assert_eq!(
1920 candidate_bases_with(Some("127.0.0.1:1234/"), vec![5678]),
1921 vec![
1922 "http://127.0.0.1:1234".to_string(),
1923 "http://127.0.0.1:5678".to_string(),
1924 ]
1925 );
1926 assert_eq!(
1928 candidate_bases_with(Some("127.0.0.1:5678"), vec![5678]),
1929 vec!["http://127.0.0.1:5678".to_string()]
1930 );
1931 assert_eq!(
1933 candidate_bases_with(None, vec![5678, 5678]),
1934 vec!["http://127.0.0.1:5678".to_string()]
1935 );
1936 assert_eq!(
1938 candidate_bases_with(Some("https://host:9"), vec![]),
1939 vec!["https://host:9".to_string()]
1940 );
1941 }
1942
1943 fn http(status: u16) -> AppError {
1944 AppError::Http {
1945 status,
1946 body: String::new(),
1947 }
1948 }
1949
1950 #[test]
1954 fn an_auth_failure_outranks_later_transport_noise() {
1955 let err = select_probe_error(vec![
1956 http(401),
1957 AppError::Transport("connection refused".into()),
1958 ]);
1959 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
1960
1961 let err = select_probe_error(vec![
1962 AppError::Transport("connection refused".into()),
1963 http(403),
1964 ]);
1965 assert!(matches!(err, AppError::Http { status: 403, .. }), "{err}");
1966 }
1967
1968 #[test]
1971 fn the_first_auth_failure_wins() {
1972 let err = select_probe_error(vec![http(401), http(403)]);
1973 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
1974 }
1975
1976 #[test]
1980 fn without_an_auth_failure_the_last_error_stands() {
1981 let err = select_probe_error(vec![
1982 AppError::Transport("first".into()),
1983 http(500),
1984 AppError::Transport("last".into()),
1985 ]);
1986 assert!(
1987 matches!(&err, AppError::Transport(m) if m == "last"),
1988 "{err}"
1989 );
1990 assert!(err.is_transient());
1991 }
1992
1993 #[test]
1996 fn a_server_error_is_not_treated_as_actionable() {
1997 let err = select_probe_error(vec![http(500), http(401)]);
1998 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
1999 }
2000
2001 #[test]
2002 fn no_candidates_at_all_yields_a_generic_error() {
2003 let err = select_probe_error(Vec::new());
2004 assert!(
2005 err.to_string().contains("no local server answered"),
2006 "{err}"
2007 );
2008 }
2009
2010 fn tls_echo() -> AppError {
2014 AppError::Http {
2015 status: 400,
2016 body: "Client sent an HTTP request to an HTTPS server.\n".into(),
2017 }
2018 }
2019
2020 #[test]
2025 fn a_tls_echo_does_not_mask_what_the_rpc_listener_said() {
2026 let err = select_probe_error(vec![
2027 AppError::Http {
2028 status: 500,
2029 body: "GetUserStatus: internal".into(),
2030 },
2031 tls_echo(),
2032 ]);
2033 assert!(
2034 matches!(&err, AppError::Http { status: 500, body } if body.contains("internal")),
2035 "{err}"
2036 );
2037 }
2038
2039 #[test]
2044 fn a_tls_echo_does_not_cost_the_silent_fallback() {
2045 let err = select_probe_error(vec![
2046 AppError::Transport("connection refused".into()),
2047 tls_echo(),
2048 ]);
2049 assert!(
2050 matches!(&err, AppError::Transport(m) if m == "connection refused"),
2051 "{err}"
2052 );
2053 assert!(
2054 err.is_transient(),
2055 "the echo must not make the run non-transient: {err}"
2056 );
2057 }
2058
2059 #[test]
2062 fn a_tls_echo_still_stands_when_it_is_the_only_thing_that_answered() {
2063 let err = select_probe_error(vec![tls_echo()]);
2064 assert!(matches!(err, AppError::Http { status: 400, .. }), "{err}");
2065 }
2066
2067 #[test]
2070 fn a_genuine_bad_request_is_not_mistaken_for_a_tls_echo() {
2071 let err = select_probe_error(vec![
2072 AppError::Http {
2073 status: 400,
2074 body: "unknown method GetUserStatus".into(),
2075 },
2076 tls_echo(),
2077 ]);
2078 assert!(
2079 matches!(&err, AppError::Http { status: 400, body } if body.contains("unknown method")),
2080 "{err}"
2081 );
2082 }
2083
2084 #[test]
2086 fn an_auth_failure_still_outranks_a_tls_echo() {
2087 let err = select_probe_error(vec![tls_echo(), http(401)]);
2088 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2089 }
2090
2091 #[test]
2092 fn every_discovered_port_is_probed_in_order() {
2093 assert_eq!(
2094 candidate_bases_with(None, vec![33875, 37435]),
2095 vec![
2096 "http://127.0.0.1:33875".to_string(),
2097 "http://127.0.0.1:37435".to_string(),
2098 ]
2099 );
2100 }
2101
2102 #[test]
2107 fn empty_discovery_yields_no_candidates() {
2108 assert!(candidate_bases_with(None, vec![]).is_empty());
2109 assert!(candidate_bases_with(Some(""), vec![]).is_empty());
2110 }
2111
2112 #[test]
2113 fn every_antigravity_product_is_recognised() {
2114 assert!(is_antigravity_process(
2116 "language_server\n",
2117 Some("/opt/antigravity/resources/bin/language_server")
2118 ));
2119 assert!(is_antigravity_process(
2121 "agy\n",
2122 Some("/home/u/.local/bin/agy")
2123 ));
2124 assert!(is_antigravity_process(
2126 "node",
2127 Some("/opt/antigravity/bin/helper")
2128 ));
2129 assert!(is_antigravity_process("antigravity", None));
2130 assert!(is_antigravity_process("agy.exe", None));
2131 assert!(is_antigravity_process("Antigravity.exe", None));
2132 assert!(is_antigravity_process("language_server.exe", None));
2133 assert!(is_antigravity_process(
2134 "language_server_windows_x64.exe",
2135 None
2136 ));
2137 assert!(is_antigravity_process(
2138 "node.exe",
2139 Some(r"C:\Users\u\AppData\Local\agy.exe")
2140 ));
2141 }
2142
2143 #[test]
2144 fn unrelated_processes_are_not_probed() {
2145 assert!(!is_antigravity_process("sshd", Some("/usr/sbin/sshd")));
2146 assert!(!is_antigravity_process("node", Some("/usr/bin/node")));
2147 assert!(!is_antigravity_process("legacy", Some("/usr/bin/legacy")));
2149 assert!(!is_antigravity_process("legacy.exe", None));
2150 assert!(!is_antigravity_process("not-agy.exe", None));
2151 assert!(!is_antigravity_process("", None));
2152 }
2153
2154 #[test]
2155 fn windows_process_names_decode_until_nul_and_tolerate_invalid_utf16() {
2156 let mut raw: Vec<u16> = "agy.exe".encode_utf16().collect();
2157 raw.extend([0, b'x' as u16]);
2158 assert_eq!(decode_windows_process_name(&raw), "agy.exe");
2159 assert_eq!(decode_windows_process_name(&[0xd800]), "�");
2160 assert_eq!(decode_windows_process_name(&[]), "");
2161 }
2162
2163 #[test]
2164 fn windows_process_filter_keeps_only_antigravity_pids() {
2165 let processes = vec![
2166 (10, "agy.exe".to_string()),
2167 (20, "language_server_windows_x64.exe".to_string()),
2168 (30, "sshd.exe".to_string()),
2169 ];
2170 let pids = matching_windows_process_ids(&processes);
2171 assert_eq!(pids, std::collections::HashSet::from([10, 20]));
2172 }
2173
2174 #[test]
2175 fn windows_listener_filter_joins_pid_loopback_and_port() {
2176 let pids = std::collections::HashSet::from([10]);
2177 let rows = [
2178 WindowsTcpRow {
2179 local_addr: [127, 0, 0, 1],
2180 local_port: u32::from(59870u16.to_be()),
2181 pid: 10,
2182 },
2183 WindowsTcpRow {
2184 local_addr: [127, 0, 0, 1],
2185 local_port: u32::from(59868u16.to_be()),
2186 pid: 10,
2187 },
2188 WindowsTcpRow {
2189 local_addr: [127, 0, 0, 1],
2190 local_port: u32::from(59870u16.to_be()),
2191 pid: 10,
2192 },
2193 WindowsTcpRow {
2194 local_addr: [0, 0, 0, 0],
2195 local_port: u32::from(50000u16.to_be()),
2196 pid: 10,
2197 },
2198 WindowsTcpRow {
2199 local_addr: [127, 0, 0, 1],
2200 local_port: u32::from(50001u16.to_be()),
2201 pid: 99,
2202 },
2203 WindowsTcpRow {
2204 local_addr: [127, 0, 0, 1],
2205 local_port: 0,
2206 pid: 10,
2207 },
2208 ];
2209 assert_eq!(matching_windows_ports(&pids, &rows), vec![59870, 59868]);
2210 }
2211
2212 #[test]
2216 fn windows_ports_from_two_products_keep_tls_listeners_last() {
2217 let pids = std::collections::HashSet::from([10, 20]);
2218 let row = |port: u16, pid: u32| WindowsTcpRow {
2219 local_addr: [127, 0, 0, 1],
2220 local_port: u32::from(port.to_be()),
2221 pid,
2222 };
2223 let rows = [
2224 row(40000, 10),
2225 row(40001, 10),
2226 row(50000, 20),
2227 row(50001, 20),
2228 ];
2229 assert_eq!(
2230 matching_windows_ports(&pids, &rows),
2231 vec![40001, 50001, 40000, 50000]
2232 );
2233 }
2234
2235 #[test]
2239 fn probe_order_puts_every_rpc_listener_ahead_of_every_tls_listener() {
2240 use std::collections::BTreeMap;
2241
2242 assert_eq!(
2244 probe_order(BTreeMap::from([(10, vec![59868, 59870])])),
2245 vec![59870, 59868]
2246 );
2247 assert_eq!(
2252 probe_order(BTreeMap::from([
2253 (10, vec![40000, 40001]),
2254 (20, vec![50000, 50001]),
2255 ])),
2256 vec![40001, 50001, 40000, 50000]
2257 );
2258 assert_eq!(
2261 probe_order(BTreeMap::from([
2262 (10, vec![6000, 5000, 4000]),
2263 (20, vec![6000, 7000]),
2264 ])),
2265 vec![6000, 7000, 5000, 4000]
2266 );
2267 assert!(probe_order(BTreeMap::new()).is_empty());
2268 }
2269
2270 #[test]
2275 fn a_port_named_twice_by_one_product_still_occupies_one_rank() {
2276 use std::collections::BTreeMap;
2277
2278 assert_eq!(
2279 probe_order(BTreeMap::from([
2280 (10, vec![40001, 40001, 40000, 40000]),
2281 (20, vec![50001, 50000]),
2282 ])),
2283 vec![40001, 50001, 40000, 50000],
2284 "duplicate rows must not reorder the ranks below them"
2285 );
2286 }
2287
2288 #[test]
2293 fn a_half_started_product_is_the_documented_exception() {
2294 use std::collections::BTreeMap;
2295
2296 assert_eq!(
2297 probe_order(BTreeMap::from([
2298 (10, vec![40000]),
2299 (20, vec![50001, 50000])
2300 ])),
2301 vec![40000, 50001, 50000]
2302 );
2303 }
2304
2305 #[test]
2310 fn an_override_with_no_authority_is_dropped_not_probed() {
2311 for junk in ["/", "///", "http://", "https://", " / "] {
2312 assert_eq!(
2313 candidate_bases_with(Some(junk), vec![4242]),
2314 vec!["http://127.0.0.1:4242".to_string()],
2315 "{junk:?} should not survive as a candidate"
2316 );
2317 }
2318 assert!(candidate_bases_with(Some("/"), vec![]).is_empty());
2319 }
2320
2321 #[test]
2322 fn windows_table_bounds_reject_truncation_and_overflow() {
2323 assert_eq!(checked_windows_row_count(52, 4, 24, 2), Some(2));
2324 assert_eq!(checked_windows_row_count(51, 4, 24, 2), None);
2325 assert_eq!(checked_windows_row_count(4, 4, 24, 0), Some(0));
2326 assert_eq!(checked_windows_row_count(52, 4, 0, 2), None);
2327 assert_eq!(
2328 checked_windows_row_count(usize::MAX, 4, 24, usize::MAX),
2329 None
2330 );
2331 }
2332
2333 #[cfg(target_os = "windows")]
2334 #[test]
2335 fn windows_tcp_table_parser_copies_complete_rows_only() {
2336 use std::mem::{offset_of, size_of};
2337 use windows_sys::Win32::NetworkManagement::IpHelper::{
2338 MIB_TCPROW_OWNER_PID, MIB_TCPTABLE_OWNER_PID,
2339 };
2340
2341 let offset = offset_of!(MIB_TCPTABLE_OWNER_PID, table);
2342 let used = offset + 2 * size_of::<MIB_TCPROW_OWNER_PID>();
2343 let words = used.div_ceil(size_of::<u32>());
2344 let mut buffer = vec![0u32; words];
2345 let first = MIB_TCPROW_OWNER_PID {
2346 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2347 dwLocalPort: u32::from(59868u16.to_be()),
2348 dwOwningPid: 10,
2349 ..Default::default()
2350 };
2351 let second = MIB_TCPROW_OWNER_PID {
2352 dwLocalAddr: u32::from_ne_bytes([127, 0, 0, 1]),
2353 dwLocalPort: u32::from(59870u16.to_be()),
2354 dwOwningPid: 10,
2355 ..Default::default()
2356 };
2357 unsafe {
2358 buffer.as_mut_ptr().write_unaligned(2);
2359 let rows = buffer
2360 .as_mut_ptr()
2361 .cast::<u8>()
2362 .add(offset)
2363 .cast::<MIB_TCPROW_OWNER_PID>();
2364 rows.write_unaligned(first);
2365 rows.add(1).write_unaligned(second);
2366 }
2367
2368 assert_eq!(
2369 parse_windows_tcp_rows(&buffer, used),
2370 vec![
2371 WindowsTcpRow {
2372 local_addr: [127, 0, 0, 1],
2373 local_port: u32::from(59868u16.to_be()),
2374 pid: 10,
2375 },
2376 WindowsTcpRow {
2377 local_addr: [127, 0, 0, 1],
2378 local_port: u32::from(59870u16.to_be()),
2379 pid: 10,
2380 },
2381 ]
2382 );
2383 assert!(parse_windows_tcp_rows(&buffer, used - 1).is_empty());
2384 }
2385
2386 #[test]
2387 fn lsof_parser_keeps_only_ports_owned_by_antigravity_processes() {
2388 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";
2391 assert_eq!(parse_lsof_pcn(output), vec![61290, 61289, 8829]);
2392 }
2393
2394 #[test]
2398 fn lsof_parser_keeps_each_products_ports_in_its_own_group() {
2399 let output = concat!(
2400 "p100\ncagy\nf3\nn127.0.0.1:40000\nf4\nn127.0.0.1:40001\n",
2401 "p200\nclanguage_server\nf5\nn127.0.0.1:50000\nf6\nn127.0.0.1:50001\n",
2402 );
2403 assert_eq!(
2404 parse_lsof_pcn(output),
2405 vec![40001, 50001, 40000, 50000],
2406 "both HTTP listeners must precede both TLS listeners"
2407 );
2408 }
2409
2410 #[test]
2411 fn lsof_parser_matches_the_capitalised_macos_app_name() {
2412 let output = "p900\ncAntigravity\nf7\nn127.0.0.1:54321\n";
2413 assert_eq!(parse_lsof_pcn(output), vec![54321]);
2414 }
2415
2416 #[test]
2417 fn lsof_parser_deduplicates_and_handles_empty_output() {
2418 let output = "p1\ncagy\nf3\nn127.0.0.1:9000\nf4\nn127.0.0.1:9000\n";
2419 assert_eq!(parse_lsof_pcn(output), vec![9000]);
2420 assert!(parse_lsof_pcn("").is_empty());
2421 }
2422
2423 #[test]
2426 fn missing_cache_surfaces_the_diagnosis_not_a_cache_miss() {
2427 let dir = tempfile::tempdir().unwrap();
2428 let cache = Cache::at(dir.path().join("usage.json"));
2429 let reason = AppError::Credentials("Antigravity: no local language server found".into());
2430
2431 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2432 let msg = err.to_string();
2433 assert!(msg.contains("no local language server found"), "{msg}");
2434 assert!(!msg.contains("no usable cache"), "{msg}");
2435 }
2436
2437 #[test]
2438 fn unusable_cache_does_not_replace_the_live_diagnosis() {
2439 let dir = tempfile::tempdir().unwrap();
2440 let cache = Cache::at(dir.path().join("antigravity"));
2441 cache.write_payload(b"{}").unwrap();
2442
2443 let reason = AppError::Credentials("Antigravity must be running".into());
2444 let err = fallback_with_error(&cache, None, reason, now()).unwrap_err();
2445 assert!(err.to_string().contains("must be running"), "{err}");
2446
2447 let original = AppError::Transport("original loopback failure".into());
2448 let err = fallback_silent(&cache, now(), original).unwrap_err();
2449 assert!(
2450 err.to_string().contains("original loopback failure"),
2451 "{err}"
2452 );
2453 }
2454
2455 #[tokio::test]
2456 async fn rpc_error_bodies_are_bounded_too() {
2457 let mut server = mockito::Server::new_async().await;
2458 let path = format!("/{STATUS_RPC}");
2459 server
2460 .mock("POST", path.as_str())
2461 .with_status(500)
2462 .with_body("x".repeat(crate::vendor::MAX_BODY_BYTES + 1))
2463 .create_async()
2464 .await;
2465
2466 let err = post_rpc(&reqwest::Client::new(), &server.url(), None, STATUS_RPC)
2467 .await
2468 .unwrap_err();
2469 assert!(err.to_string().contains("exceeds"), "{err}");
2470 }
2471
2472 #[test]
2473 fn blank_override_falls_through_to_discovery() {
2474 assert_eq!(
2475 candidate_bases_with(Some(" "), vec![4242]),
2476 vec!["http://127.0.0.1:4242".to_string()]
2477 );
2478 }
2479
2480 fn fixture() -> (tempfile::TempDir, Cache) {
2485 let td = tempfile::TempDir::new().unwrap();
2486 let cache = Cache::at(td.path().join("antigravity"));
2487 (td, cache)
2488 }
2489
2490 fn endpoints(server: &mockito::Server) -> cloud::Endpoints {
2491 let base = server.url();
2492 cloud::Endpoints {
2493 quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
2494 load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
2495 token: format!("{base}/token"),
2496 }
2497 }
2498
2499 fn keyring_blob(expiry: &str, with_refresh: bool) -> String {
2501 let mut token = serde_json::json!({
2502 "access_token": "KEYRING-AT",
2503 "expiry": expiry,
2504 });
2505 if with_refresh {
2506 token["refresh_token"] = serde_json::json!("KEYRING-RT");
2507 }
2508 serde_json::json!({ "token": token }).to_string()
2509 }
2510
2511 const EXPIRED: &str = "2026-07-22T11:00:00Z";
2513 const VALID: &str = "2026-07-22T13:00:00Z";
2515
2516 fn remote<'a>(blob: &'a str, eps: &'a cloud::Endpoints) -> RemoteOverride<'a> {
2518 RemoteOverride {
2519 credential: SavedCredential::Blob(blob),
2520 endpoints: Some(eps),
2521 local_bases: Some(vec![]),
2522 }
2523 }
2524
2525 fn bare_summary() -> String {
2528 let v: serde_json::Value = serde_json::from_str(QUOTA_JSON).unwrap();
2529 v["response"].to_string()
2530 }
2531
2532 fn quota_mock(server: &mut mockito::Server, bearer: &str) -> mockito::Mock {
2533 server
2534 .mock("POST", "/daily/quota")
2535 .match_header("authorization", format!("Bearer {bearer}").as_str())
2536 .with_status(200)
2537 .with_body(bare_summary())
2538 }
2539
2540 fn token_mock(server: &mut mockito::Server) -> mockito::Mock {
2541 server
2542 .mock("POST", "/token")
2543 .with_status(200)
2544 .with_body(r#"{"access_token":"NEW-AT","expires_in":3600}"#)
2545 }
2546
2547 fn test_oauth() -> cloud::OauthClient {
2548 cloud::OauthClient {
2549 id: "test-client".into(),
2550 secret: "test-client-secret".into(),
2551 }
2552 }
2553
2554 async fn run(cache: &Cache, remote: RemoteOverride<'_>, ttl: Duration) -> Result<FetchOutcome> {
2555 fetch_snapshot_at(
2556 &reqwest::Client::new(),
2557 cache,
2558 ttl,
2559 Some(&test_oauth()),
2560 remote,
2561 now(),
2562 )
2563 .await
2564 }
2565
2566 #[tokio::test]
2567 async fn with_no_local_server_the_saved_session_answers_from_the_api() {
2568 let mut server = mockito::Server::new_async().await;
2569 let eps = endpoints(&server);
2570 let quota = quota_mock(&mut server, "KEYRING-AT")
2571 .expect(1)
2572 .create_async()
2573 .await;
2574 let plan = server
2575 .mock("POST", "/daily/plan")
2576 .match_header("authorization", "Bearer KEYRING-AT")
2577 .with_status(200)
2578 .with_body(r#"{"currentTier":{"name":"google_ai_pro"}}"#)
2579 .expect(1)
2580 .create_async()
2581 .await;
2582 let token = token_mock(&mut server).expect(0).create_async().await;
2584 let blob = keyring_blob(VALID, true);
2585 let (_td, cache) = fixture();
2586
2587 let outcome = run(&cache, remote(&blob, &eps), Duration::from_secs(60))
2588 .await
2589 .expect("remote path yields a snapshot");
2590
2591 quota.assert_async().await;
2592 plan.assert_async().await;
2593 token.assert_async().await;
2594 let snap = outcome.snapshot;
2595 assert!(!outcome.stale);
2596 assert_eq!(snap.source, AntigravitySource::Remote);
2597 assert_eq!(snap.plan, "Pro");
2598 assert_eq!(snap.session.as_ref().unwrap().utilization_pct, 43);
2599 assert_eq!(snap.weekly.as_ref().unwrap().utilization_pct, 8);
2600 assert_eq!(snap.third_party_session.unwrap().utilization_pct, 75);
2601 assert_eq!(snap.third_party_weekly.unwrap().utilization_pct, 0);
2602 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
2603 assert_eq!(snap.account, format!("acct:{fingerprint}"));
2604 assert!(!snap.account.contains("KEYRING"), "{}", snap.account);
2605
2606 let cached = parse_cache_at(
2608 &cache
2609 .fresh_payload(Duration::from_secs(60))
2610 .unwrap()
2611 .unwrap(),
2612 None,
2613 now(),
2614 )
2615 .unwrap();
2616 assert_eq!(cached.source, AntigravitySource::Remote);
2617 assert_eq!(cached.account, snap.account);
2618 }
2619
2620 #[tokio::test]
2623 async fn an_expired_keyring_token_is_refreshed_once_and_the_refresh_is_reused() {
2624 let mut server = mockito::Server::new_async().await;
2625 let eps = endpoints(&server);
2626 let _quota = quota_mock(&mut server, "NEW-AT")
2627 .expect(2)
2628 .create_async()
2629 .await;
2630 let token = token_mock(&mut server).expect(1).create_async().await;
2631 let blob = keyring_blob(EXPIRED, true);
2632 let (_td, cache) = fixture();
2633
2634 let first = run(&cache, remote(&blob, &eps), Duration::ZERO)
2635 .await
2636 .expect("refresh, then quota");
2637 assert_eq!(first.snapshot.source, AntigravitySource::Remote);
2638 token.assert_async().await;
2639
2640 let fingerprint = credential::parse_keyring_blob(&blob).unwrap().fingerprint;
2641 let persisted = cloud::read_persisted(&cloud::oauth_cache_path(&cache), &fingerprint)
2642 .expect("refreshed token persisted under the keyring session's fingerprint");
2643 assert_eq!(persisted.access_token, "NEW-AT");
2644
2645 run(&cache, remote(&blob, &eps), Duration::ZERO)
2647 .await
2648 .expect("persisted token reused");
2649 token.assert_async().await;
2650 }
2651
2652 #[tokio::test]
2655 async fn a_stale_token_the_api_rejects_gets_one_refresh_and_one_retry() {
2656 let mut server = mockito::Server::new_async().await;
2657 let eps = endpoints(&server);
2658 let rejected = server
2659 .mock("POST", "/daily/quota")
2660 .match_header("authorization", "Bearer KEYRING-AT")
2661 .with_status(401)
2662 .expect(1)
2663 .create_async()
2664 .await;
2665 let accepted = quota_mock(&mut server, "NEW-AT")
2666 .expect(1)
2667 .create_async()
2668 .await;
2669 let token = token_mock(&mut server).expect(1).create_async().await;
2670 let blob = keyring_blob(VALID, true);
2671 let (_td, cache) = fixture();
2672
2673 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
2674 .await
2675 .expect("retry with the refreshed token succeeds");
2676
2677 rejected.assert_async().await;
2678 accepted.assert_async().await;
2679 token.assert_async().await;
2680 assert_eq!(outcome.snapshot.source, AntigravitySource::Remote);
2681 }
2682
2683 #[tokio::test]
2686 async fn a_refused_refresh_is_a_credentials_error_without_the_token() {
2687 let mut server = mockito::Server::new_async().await;
2688 let eps = endpoints(&server);
2689 let _token = server
2690 .mock("POST", "/token")
2691 .with_status(400)
2692 .with_body(r#"{"error":"invalid_grant"}"#)
2693 .create_async()
2694 .await;
2695 let blob = keyring_blob(EXPIRED, true);
2696 let (_td, cache) = fixture();
2697
2698 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
2699 .await
2700 .expect_err("no cache to fall back on");
2701
2702 assert!(matches!(err, AppError::Credentials(_)), "{err}");
2703 let rendered = err.to_string();
2704 assert!(!rendered.contains("KEYRING-RT"), "{rendered}");
2705 assert!(!rendered.contains("KEYRING-AT"), "{rendered}");
2706 }
2707
2708 #[tokio::test]
2711 async fn a_rejection_right_after_a_refresh_asks_to_sign_in_again() {
2712 let mut server = mockito::Server::new_async().await;
2713 let eps = endpoints(&server);
2714 let _quota = server
2715 .mock("POST", "/daily/quota")
2716 .with_status(401)
2717 .create_async()
2718 .await;
2719 let token = token_mock(&mut server).expect(1).create_async().await;
2720 let blob = keyring_blob(EXPIRED, true);
2721 let (_td, cache) = fixture();
2722
2723 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
2724 .await
2725 .expect_err("rejected after refresh");
2726
2727 token.assert_async().await;
2728 assert!(matches!(err, AppError::Credentials(_)), "{err}");
2729 assert!(err.to_string().contains("sign in again"), "{err}");
2730 }
2731
2732 #[tokio::test]
2734 async fn an_expired_session_without_a_refresh_token_asks_to_sign_in_again() {
2735 let mut server = mockito::Server::new_async().await;
2736 let eps = endpoints(&server);
2737 let token = token_mock(&mut server).expect(0).create_async().await;
2738 let blob = keyring_blob(EXPIRED, false);
2739 let (_td, cache) = fixture();
2740
2741 let err = run(&cache, remote(&blob, &eps), Duration::ZERO)
2742 .await
2743 .expect_err("nothing to refresh with");
2744
2745 token.assert_async().await;
2746 assert!(matches!(err, AppError::Credentials(_)), "{err}");
2747 assert!(err.to_string().contains("sign in again"), "{err}");
2748 }
2749
2750 #[tokio::test]
2753 async fn no_saved_session_extends_the_no_local_server_error() {
2754 let mut server = mockito::Server::new_async().await;
2755 let eps = endpoints(&server);
2756 let quota = quota_mock(&mut server, "KEYRING-AT")
2757 .expect(0)
2758 .create_async()
2759 .await;
2760 let (_td, cache) = fixture();
2761
2762 let err = run(
2763 &cache,
2764 RemoteOverride {
2765 credential: SavedCredential::Absent,
2766 endpoints: Some(&eps),
2767 local_bases: Some(vec![]),
2768 },
2769 Duration::ZERO,
2770 )
2771 .await
2772 .expect_err("no source at all");
2773
2774 quota.assert_async().await;
2775 assert!(matches!(err, AppError::Credentials(_)), "{err}");
2776 let rendered = err.to_string();
2777 assert!(rendered.contains("no local server found"), "{rendered}");
2778 assert!(rendered.contains("saved Google session"), "{rendered}");
2779 }
2780
2781 #[tokio::test]
2784 async fn a_signed_out_local_server_is_reported_rather_than_bypassed() {
2785 let mut server = mockito::Server::new_async().await;
2786 let eps = endpoints(&server);
2787 let status_path = format!("/{STATUS_RPC}");
2788 let _status = server
2789 .mock("POST", status_path.as_str())
2790 .with_status(401)
2791 .create_async()
2792 .await;
2793 let quota = quota_mock(&mut server, "KEYRING-AT")
2794 .expect(0)
2795 .create_async()
2796 .await;
2797 let blob = keyring_blob(VALID, true);
2798 let (_td, cache) = fixture();
2799
2800 let err = run(
2801 &cache,
2802 RemoteOverride {
2803 credential: SavedCredential::Blob(&blob),
2804 endpoints: Some(&eps),
2805 local_bases: Some(vec![server.url()]),
2806 },
2807 Duration::ZERO,
2808 )
2809 .await
2810 .expect_err("signed-out local server");
2811
2812 quota.assert_async().await;
2813 assert!(matches!(err, AppError::Http { status: 401, .. }), "{err}");
2814 }
2815
2816 #[tokio::test]
2819 async fn a_cached_remote_snapshot_is_served_when_the_api_fails() {
2820 let mut server = mockito::Server::new_async().await;
2821 let eps = endpoints(&server);
2822 for path in ["/daily/quota", "/prod/quota"] {
2823 server
2824 .mock("POST", path)
2825 .with_status(500)
2826 .with_body("boom")
2827 .create_async()
2828 .await;
2829 }
2830 let blob = keyring_blob(VALID, true);
2831 let (_td, cache) = fixture();
2832 cache.ensure_dir().unwrap();
2833 let mut earlier = parsed();
2834 earlier.source = AntigravitySource::Remote;
2835 earlier.account = "acct:earlier".into();
2836 cache
2837 .write_payload(&serde_json::to_vec(&snap_to_json(&earlier)).unwrap())
2838 .unwrap();
2839
2840 let outcome = run(&cache, remote(&blob, &eps), Duration::ZERO)
2841 .await
2842 .expect("stale cache stands in");
2843
2844 assert!(outcome.stale);
2845 assert_eq!(outcome.snapshot, earlier);
2846 assert!(
2847 matches!(outcome.last_error, Some((500, _))),
2848 "{:?}",
2849 outcome.last_error
2850 );
2851 }
2852
2853 #[test]
2854 fn source_round_trips_through_the_cache_and_defaults_to_local() {
2855 let mut snap = parsed();
2856 snap.source = AntigravitySource::Remote;
2857 let bytes = serde_json::to_vec(&snap_to_json(&snap)).unwrap();
2858 assert_eq!(
2859 parse_cache_at(&bytes, None, now()).unwrap().source,
2860 AntigravitySource::Remote
2861 );
2862
2863 let mut legacy = snap_to_json(&snap);
2865 legacy.as_object_mut().unwrap().remove("source");
2866 let legacy = serde_json::to_vec(&legacy).unwrap();
2867 assert_eq!(
2868 parse_cache_at(&legacy, None, now()).unwrap().source,
2869 AntigravitySource::Local
2870 );
2871 }
2872}