ai_usagebar/antigravity/
cloud.rs

1//! Google Cloud Code client for the Antigravity "app closed" fallback.
2//!
3//! When no Antigravity product is running there is no loopback RPC to ask, but
4//! the Google session Antigravity saved in the OS keyring (`credential.rs`)
5//! is still valid. The same Cloud Code endpoints the product itself calls —
6//! `retrieveUserQuotaSummary` for the quota windows and `loadCodeAssist` for
7//! the plan tier — accept that bearer token directly. Access tokens live about
8//! an hour; the refresh token is exchanged at Google's standard token endpoint
9//! and the refreshed access token is persisted in ai-usagebar's own cache,
10//! never written back to the keyring.
11//!
12//! Upstream error bodies are never surfaced: a rejected bearer token or an
13//! `invalid_grant` response can carry account detail, so every non-2xx maps to
14//! a fixed message — same rule as `kiro::oauth::refresh`.
15
16use std::path::{Path, PathBuf};
17use std::time::Duration;
18
19use chrono::{DateTime, Utc};
20use serde::{Deserialize, Serialize};
21
22use crate::cache::{Cache, atomic_write};
23use crate::error::{AppError, Result};
24use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
25
26const DAILY_BASE: &str = "https://daily-cloudcode-pa.googleapis.com";
27const PROD_BASE: &str = "https://cloudcode-pa.googleapis.com";
28const QUOTA_PATH: &str = "/v1internal:retrieveUserQuotaSummary";
29const LOAD_CODE_ASSIST_PATH: &str = "/v1internal:loadCodeAssist";
30const TOKEN_URL: &str = "https://oauth2.googleapis.com/token";
31
32const HTTP_TIMEOUT: Duration = Duration::from_secs(15);
33const QUOTA_USER_AGENT: &str = "antigravity";
34const LOAD_CODE_ASSIST_USER_AGENT: &str = "agy";
35const OAUTH_CACHE_FILE: &str = "oauth.json";
36
37/// Refresh this far ahead of the recorded expiry so a slow round-trip never
38/// races the token's death. Mirrors `kiro::oauth::REFRESH_BUFFER_SECS`.
39pub const REFRESH_BUFFER_SECS: i64 = 300;
40
41/// Longest plan label kept from an unrecognised tier name.
42const MAX_PLAN_CHARS: usize = 32;
43
44const SESSION_REJECTED: &str = "Antigravity's Google session was rejected";
45const REFRESH_FAILED: &str = "Antigravity token refresh failed";
46const SESSION_EXPIRED: &str =
47    "Antigravity's saved Google session expired; open Antigravity to sign in again";
48
49/// Where the fallback talks to. Quota and plan each try every base in order
50/// (the daily channel first, as the product does); tests point them at mockito.
51#[derive(Debug, Clone)]
52pub struct Endpoints {
53    pub quota: Vec<String>,
54    pub load_code_assist: Vec<String>,
55    pub token: String,
56}
57
58impl Default for Endpoints {
59    fn default() -> Self {
60        Self {
61            quota: vec![
62                format!("{DAILY_BASE}{QUOTA_PATH}"),
63                format!("{PROD_BASE}{QUOTA_PATH}"),
64            ],
65            load_code_assist: vec![
66                format!("{DAILY_BASE}{LOAD_CODE_ASSIST_PATH}"),
67                format!("{PROD_BASE}{LOAD_CODE_ASSIST_PATH}"),
68            ],
69            token: TOKEN_URL.to_string(),
70        }
71    }
72}
73
74/// The OAuth client used for refresh: config overrides or the defaults.
75#[derive(Debug, Clone)]
76pub struct OauthClient {
77    pub id: String,
78    pub secret: String,
79}
80
81impl OauthClient {
82    /// `[antigravity] oauth_client_id` + `oauth_client_secret`. Both are
83    /// required for a refresh; nothing is embedded in source, so without them
84    /// the saved access token is used only while it lasts. The pair is
85    /// Antigravity's own public installed-app client (RFC 8252 §8.5: such a
86    /// secret grants nothing on its own — the refresh token in the keyring is
87    /// the credential), but a secret-shaped literal in a tracked file trips
88    /// every secret scanner, so it is the user's line to write.
89    pub fn from_config(id: Option<&str>, secret: Option<&str>) -> Option<Self> {
90        let pick = |value: Option<&str>| {
91            value
92                .map(str::trim)
93                .filter(|value| !value.is_empty())
94                .map(str::to_string)
95        };
96        Some(Self {
97            id: pick(id)?,
98            secret: pick(secret)?,
99        })
100    }
101}
102
103/// A freshly exchanged access token.
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct Refreshed {
106    pub access_token: String,
107    pub expires_at: DateTime<Utc>,
108}
109
110/// ai-usagebar's own record of a refreshed access token, scoped to the
111/// keyring session it was minted from (`credential::StoredToken::fingerprint`).
112#[derive(Debug, Clone, PartialEq, Eq, Deserialize, Serialize)]
113pub struct PersistedOAuth {
114    pub fingerprint: String,
115    pub access_token: String,
116    pub expires_at: DateTime<Utc>,
117}
118
119async fn post_json(
120    client: &reqwest::Client,
121    url: &str,
122    access_token: &str,
123    user_agent: &str,
124) -> std::result::Result<reqwest::Response, reqwest::Error> {
125    client
126        .post(url)
127        .timeout(HTTP_TIMEOUT)
128        .header("Authorization", format!("Bearer {access_token}"))
129        .header("Content-Type", "application/json")
130        .header("Accept", "application/json")
131        .header("User-Agent", user_agent)
132        .body("{}")
133        .send()
134        .await
135}
136
137/// Ask each quota base in turn for the user's quota summary. The raw JSON is
138/// returned as-is — either the bare summary or the `{"response": …}` wrapper,
139/// depending on the channel — for the caller to project.
140///
141/// A 401/403 stops immediately: the token is the problem, not the base. Any
142/// other status or a transport failure moves on to the next base; when every
143/// base failed the last error is returned.
144pub async fn fetch_quota(
145    client: &reqwest::Client,
146    endpoints: &Endpoints,
147    access_token: &str,
148) -> Result<serde_json::Value> {
149    let mut last_error = AppError::Transport("no Antigravity quota endpoint configured".into());
150    for url in &endpoints.quota {
151        let resp = match post_json(client, url, access_token, QUOTA_USER_AGENT).await {
152            Ok(resp) => resp,
153            Err(e) => {
154                last_error = e.into();
155                continue;
156            }
157        };
158        let status = resp.status();
159        if matches!(status.as_u16(), 401 | 403) {
160            return Err(AppError::Http {
161                status: status.as_u16(),
162                body: SESSION_REJECTED.into(),
163            });
164        }
165        if !status.is_success() {
166            last_error = AppError::Http {
167                status: status.as_u16(),
168                body: format!(
169                    "Antigravity quota endpoint returned HTTP {}",
170                    status.as_u16()
171                ),
172            };
173            continue;
174        }
175        match read_body_capped(resp, MAX_BODY_BYTES).await {
176            Ok(bytes) => match serde_json::from_slice::<serde_json::Value>(&bytes) {
177                Ok(value) => return Ok(value),
178                Err(e) => {
179                    last_error = AppError::Schema(format!("antigravity quota response: {e}"));
180                }
181            },
182            Err(e) => last_error = e,
183        }
184    }
185    Err(last_error)
186}
187
188/// Best-effort plan label from `loadCodeAssist`: the first base that answers
189/// 2xx wins, `paidTier.name` beats `currentTier.name`, and any failure is
190/// simply `None` — the quota figures are the product, the tier is a garnish.
191pub async fn fetch_plan(
192    client: &reqwest::Client,
193    endpoints: &Endpoints,
194    access_token: &str,
195) -> Option<String> {
196    for url in &endpoints.load_code_assist {
197        let Ok(resp) = post_json(client, url, access_token, LOAD_CODE_ASSIST_USER_AGENT).await
198        else {
199            continue;
200        };
201        let status = resp.status();
202        if matches!(status.as_u16(), 401 | 403) {
203            return None;
204        }
205        if !status.is_success() {
206            continue;
207        }
208        let Ok(bytes) = read_body_capped(resp, MAX_BODY_BYTES).await else {
209            continue;
210        };
211        let Ok(value) = serde_json::from_slice::<serde_json::Value>(&bytes) else {
212            continue;
213        };
214        return plan_from_value(&value);
215    }
216    None
217}
218
219fn plan_from_value(value: &serde_json::Value) -> Option<String> {
220    let body = value.get("response").unwrap_or(value);
221    ["paidTier", "currentTier"]
222        .iter()
223        .filter_map(|tier| body.get(tier)?.get("name")?.as_str())
224        .map(format_plan)
225        .find(|plan| !plan.is_empty())
226}
227
228/// Google's tier ids and display names both reach here (`google_ai_ultra`,
229/// `Google AI Pro`, `free-tier`); the widget shows one short word.
230pub fn format_plan(raw: &str) -> String {
231    let words: Vec<String> = raw
232        .split(|c: char| c.is_whitespace() || c == '_' || c == '-')
233        .filter(|word| !word.is_empty())
234        .map(str::to_lowercase)
235        .collect();
236    for (needle, label) in [("ultra", "Ultra"), ("pro", "Pro"), ("free", "Free")] {
237        if words.iter().any(|word| word == needle) {
238            return label.to_string();
239        }
240    }
241    let title: Vec<String> = words
242        .iter()
243        .map(|word| {
244            let mut chars = word.chars();
245            match chars.next() {
246                Some(first) => first.to_uppercase().chain(chars).collect(),
247                None => String::new(),
248            }
249        })
250        .collect();
251    title
252        .join(" ")
253        .chars()
254        .take(MAX_PLAN_CHARS)
255        .collect::<String>()
256        .trim_end()
257        .to_string()
258}
259
260#[derive(Debug, Deserialize)]
261struct TokenResponse {
262    access_token: String,
263    expires_in: Option<serde_json::Value>,
264}
265
266/// Exchange the keyring's refresh token for a new access token at `token_url`
267/// ([`Endpoints::token`] in production; mockito in tests).
268///
269/// A definitive 4xx (anything but 408/429) means the session is gone —
270/// revoked, or the refresh token rotated by a newer sign-in — and is reported
271/// as a `Credentials` error that tells the user what to do. Everything else is
272/// an `Http` error with a fixed body.
273pub async fn refresh(
274    client: &reqwest::Client,
275    token_url: &str,
276    oauth: &OauthClient,
277    refresh_token: &str,
278) -> Result<Refreshed> {
279    let form = [
280        ("grant_type", "refresh_token"),
281        ("refresh_token", refresh_token),
282        ("client_id", oauth.id.as_str()),
283        ("client_secret", oauth.secret.as_str()),
284    ];
285    let resp = client
286        .post(token_url)
287        .timeout(HTTP_TIMEOUT)
288        .header("Accept", "application/json")
289        .form(&form)
290        .send()
291        .await?;
292    let status = resp.status();
293    let body = read_body_capped(resp, MAX_BODY_BYTES).await?;
294    if !status.is_success() {
295        let code = status.as_u16();
296        if (400..500).contains(&code) && !matches!(code, 408 | 429) {
297            return Err(AppError::Credentials(SESSION_EXPIRED.into()));
298        }
299        return Err(AppError::Http {
300            status: code,
301            body: REFRESH_FAILED.into(),
302        });
303    }
304    let parsed: TokenResponse = serde_json::from_slice(&body)
305        .map_err(|e| AppError::Schema(format!("antigravity token refresh response: {e}")))?;
306    if parsed.access_token.trim().is_empty() {
307        return Err(AppError::Schema(
308            "antigravity token refresh response: access_token is empty".into(),
309        ));
310    }
311    let expires_in = parsed
312        .expires_in
313        .as_ref()
314        .and_then(expires_in_secs)
315        .ok_or_else(|| {
316            AppError::Schema("antigravity token refresh response: invalid expires_in".into())
317        })?;
318    let expires_at_secs = Utc::now()
319        .timestamp()
320        .checked_add(expires_in)
321        .ok_or_else(|| AppError::Schema("antigravity token refresh expiry overflowed".into()))?;
322    let expires_at = DateTime::from_timestamp(expires_at_secs, 0).ok_or_else(|| {
323        AppError::Schema("antigravity token refresh expiry is out of range".into())
324    })?;
325    Ok(Refreshed {
326        access_token: parsed.access_token,
327        expires_at,
328    })
329}
330
331/// Google returns `expires_in` as a JSON number; some proxies stringify it.
332/// Only a positive integer within a sane range is accepted.
333fn expires_in_secs(value: &serde_json::Value) -> Option<i64> {
334    const MAX_SECS: i64 = 366 * 24 * 60 * 60;
335    let secs = match value {
336        serde_json::Value::Number(number) => number.as_i64()?,
337        serde_json::Value::String(text) => text.trim().parse::<i64>().ok()?,
338        _ => return None,
339    };
340    (1..=MAX_SECS).contains(&secs).then_some(secs)
341}
342
343/// `None` (no recorded expiry) always refreshes; otherwise refresh once the
344/// token is within [`REFRESH_BUFFER_SECS`] of dying.
345pub fn needs_refresh(expires_at: Option<DateTime<Utc>>, now: DateTime<Utc>) -> bool {
346    match expires_at {
347        None => true,
348        Some(expires_at) => expires_at.timestamp() < now.timestamp() + REFRESH_BUFFER_SECS,
349    }
350}
351
352/// Where the refreshed access token is persisted, inside the vendor cache.
353pub fn oauth_cache_path(cache: &Cache) -> PathBuf {
354    cache.dir().join(OAUTH_CACHE_FILE)
355}
356
357/// Whether a usable Google session has already been persisted for Antigravity,
358/// without touching the keyring.
359///
360/// `detect` needs to know that the remote fallback would work, but it promises
361/// to be cheap and silent, and reading the keyring is not unconditionally
362/// silent — on macOS `security find-generic-password` can raise a Keychain
363/// prompt, and a background probe that pops a system dialog is worse than the
364/// provider it would have enabled (see #148 for how badly that goes). This
365/// reads only our own cache file, so it cannot prompt, cannot block and cannot
366/// reach the network. It misses the very first run, before any fetch has
367/// persisted a token; that is the deliberate trade.
368pub fn has_persisted_session(path: &Path) -> bool {
369    let Ok(bytes) = std::fs::read(path) else {
370        return false;
371    };
372    serde_json::from_slice::<PersistedOAuth>(&bytes)
373        .is_ok_and(|saved| !saved.access_token.is_empty() && !saved.fingerprint.is_empty())
374}
375
376/// The persisted token for exactly this keyring session. Absent, unreadable,
377/// malformed, empty, or minted from a different session all read as `None`:
378/// the worst case is one extra refresh round-trip.
379pub fn read_persisted(path: &Path, fingerprint: &str) -> Option<PersistedOAuth> {
380    let bytes = std::fs::read(path).ok()?;
381    let persisted: PersistedOAuth = serde_json::from_slice(&bytes).ok()?;
382    if persisted.fingerprint != fingerprint || persisted.access_token.trim().is_empty() {
383        return None;
384    }
385    Some(persisted)
386}
387
388/// Atomically write the persisted token, owner-only on unix (kiro pattern).
389pub fn write_persisted(path: &Path, value: &PersistedOAuth) -> Result<()> {
390    let bytes = serde_json::to_vec_pretty(value)?;
391    atomic_write(path, &bytes)?;
392    #[cfg(unix)]
393    {
394        use std::os::unix::fs::PermissionsExt;
395        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
396            .map_err(|e| AppError::io_at(path, e))?;
397    }
398    Ok(())
399}
400
401#[cfg(test)]
402mod tests {
403
404    /// The detect probe must treat anything it cannot positively read as
405    /// "no session", so a corrupt cache never advertises a provider that
406    /// would then fail its first fetch.
407    #[test]
408    fn has_persisted_session_is_true_only_for_a_readable_token() {
409        let td = tempfile::TempDir::new().unwrap();
410        let path = td.path().join("oauth.json");
411
412        assert!(!has_persisted_session(&path), "missing file");
413
414        std::fs::write(&path, b"").unwrap();
415        assert!(!has_persisted_session(&path), "empty file");
416
417        std::fs::write(&path, b"{ not json").unwrap();
418        assert!(!has_persisted_session(&path), "malformed");
419
420        let no_token = serde_json::json!({
421            "fingerprint": "abc", "access_token": "",
422            "expires_at": "2026-01-01T00:00:00Z"
423        });
424        std::fs::write(&path, no_token.to_string()).unwrap();
425        assert!(!has_persisted_session(&path), "empty access token");
426
427        let good = serde_json::json!({
428            "fingerprint": "abc", "access_token": "ya29.test",
429            "expires_at": "2026-01-01T00:00:00Z"
430        });
431        std::fs::write(&path, good.to_string()).unwrap();
432        assert!(
433            has_persisted_session(&path),
434            "a readable token is a session"
435        );
436    }
437
438    /// An expired token still counts: `detect` is asking "would the remote
439    /// path work", and an expired access token refreshes rather than failing.
440    #[test]
441    fn an_expired_persisted_token_still_counts_as_a_session() {
442        let td = tempfile::TempDir::new().unwrap();
443        let path = td.path().join("oauth.json");
444        let stale = serde_json::json!({
445            "fingerprint": "abc", "access_token": "ya29.stale",
446            "expires_at": "2000-01-01T00:00:00Z"
447        });
448        std::fs::write(&path, stale.to_string()).unwrap();
449        assert!(has_persisted_session(&path));
450    }
451    use super::*;
452    use mockito::Matcher;
453    use tempfile::TempDir;
454
455    fn endpoints(server: &mockito::Server) -> Endpoints {
456        let base = server.url();
457        Endpoints {
458            quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
459            load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
460            token: format!("{base}/token"),
461        }
462    }
463
464    #[test]
465    fn default_endpoints_try_daily_then_prod() {
466        let e = Endpoints::default();
467        assert_eq!(
468            e.quota,
469            vec![
470                "https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary",
471                "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary",
472            ]
473        );
474        assert_eq!(
475            e.load_code_assist,
476            vec![
477                "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist",
478                "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist",
479            ]
480        );
481        assert_eq!(e.token, "https://oauth2.googleapis.com/token");
482    }
483
484    fn test_client() -> OauthClient {
485        OauthClient {
486            id: "test-client".into(),
487            secret: "test-client-secret".into(),
488        }
489    }
490
491    #[test]
492    fn oauth_client_needs_both_halves_and_trims_them() {
493        assert!(OauthClient::from_config(None, None).is_none());
494        assert!(OauthClient::from_config(Some("id"), Some("")).is_none());
495        assert!(OauthClient::from_config(Some(" "), Some("s")).is_none());
496        let both = OauthClient::from_config(Some(" my-id "), Some(" my-secret ")).unwrap();
497        assert_eq!(both.id, "my-id");
498        assert_eq!(both.secret, "my-secret");
499    }
500
501    #[tokio::test]
502    async fn quota_falls_through_from_daily_to_prod() {
503        let mut server = mockito::Server::new_async().await;
504        let daily = server
505            .mock("POST", "/daily/quota")
506            .match_header("authorization", "Bearer AT")
507            .match_header("content-type", "application/json")
508            .match_header("accept", "application/json")
509            .match_header("user-agent", "antigravity")
510            .match_body("{}")
511            .with_status(500)
512            .with_body("boom")
513            .expect(1)
514            .create_async()
515            .await;
516        let prod = server
517            .mock("POST", "/prod/quota")
518            .match_header("authorization", "Bearer AT")
519            .match_header("user-agent", "antigravity")
520            .with_status(200)
521            .with_body(r#"{"response":{"buckets":[{"modelId":"gemini","remainingFraction":0.5}]}}"#)
522            .expect(1)
523            .create_async()
524            .await;
525
526        let value = fetch_quota(&reqwest::Client::new(), &endpoints(&server), "AT")
527            .await
528            .unwrap();
529        assert_eq!(value["response"]["buckets"][0]["remainingFraction"], 0.5);
530        daily.assert_async().await;
531        prod.assert_async().await;
532    }
533
534    #[tokio::test]
535    async fn quota_401_stops_without_trying_the_next_base() {
536        let mut server = mockito::Server::new_async().await;
537        let daily = server
538            .mock("POST", "/daily/quota")
539            .with_status(401)
540            .with_body(r#"{"error":{"message":"sensitive detail"}}"#)
541            .expect(1)
542            .create_async()
543            .await;
544        let prod = server
545            .mock("POST", "/prod/quota")
546            .with_status(200)
547            .with_body("{}")
548            .expect(0)
549            .create_async()
550            .await;
551
552        let err = fetch_quota(&reqwest::Client::new(), &endpoints(&server), "AT")
553            .await
554            .unwrap_err();
555        match err {
556            AppError::Http { status, body } => {
557                assert_eq!(status, 401);
558                assert_eq!(body, SESSION_REJECTED);
559            }
560            other => panic!("expected Http, got {other:?}"),
561        }
562        daily.assert_async().await;
563        prod.assert_async().await;
564    }
565
566    #[tokio::test]
567    async fn quota_reports_the_last_error_when_every_base_fails() {
568        let mut server = mockito::Server::new_async().await;
569        server
570            .mock("POST", "/daily/quota")
571            .with_status(500)
572            .create_async()
573            .await;
574        server
575            .mock("POST", "/prod/quota")
576            .with_status(503)
577            .with_body("private upstream text")
578            .create_async()
579            .await;
580
581        let err = fetch_quota(&reqwest::Client::new(), &endpoints(&server), "AT")
582            .await
583            .unwrap_err();
584        match err {
585            AppError::Http { status, body } => {
586                assert_eq!(status, 503);
587                assert!(!body.contains("private upstream text"));
588            }
589            other => panic!("expected Http, got {other:?}"),
590        }
591    }
592
593    #[tokio::test]
594    async fn plan_prefers_paid_tier_and_sends_the_agy_user_agent() {
595        let mut server = mockito::Server::new_async().await;
596        let m = server
597            .mock("POST", "/daily/plan")
598            .match_header("authorization", "Bearer AT")
599            .match_header("user-agent", "agy")
600            .match_body("{}")
601            .with_status(200)
602            .with_body(
603                r#"{"currentTier":{"id":"free-tier","name":"Free"},"paidTier":{"id":"google_ai_pro","name":"Google AI Pro"}}"#,
604            )
605            .create_async()
606            .await;
607
608        let plan = fetch_plan(&reqwest::Client::new(), &endpoints(&server), "AT").await;
609        assert_eq!(plan.as_deref(), Some("Pro"));
610        m.assert_async().await;
611    }
612
613    #[tokio::test]
614    async fn plan_falls_back_to_current_tier_and_to_the_next_base() {
615        let mut server = mockito::Server::new_async().await;
616        server
617            .mock("POST", "/daily/plan")
618            .with_status(500)
619            .create_async()
620            .await;
621        server
622            .mock("POST", "/prod/plan")
623            .with_status(200)
624            .with_body(r#"{"response":{"currentTier":{"name":"GOOGLE_AI_ULTRA"}}}"#)
625            .create_async()
626            .await;
627
628        let plan = fetch_plan(&reqwest::Client::new(), &endpoints(&server), "AT").await;
629        assert_eq!(plan.as_deref(), Some("Ultra"));
630    }
631
632    #[tokio::test]
633    async fn plan_is_none_when_nothing_answers_usefully() {
634        let mut server = mockito::Server::new_async().await;
635        server
636            .mock("POST", "/daily/plan")
637            .with_status(200)
638            .with_body("not json")
639            .create_async()
640            .await;
641        server
642            .mock("POST", "/prod/plan")
643            .with_status(200)
644            .with_body(r#"{"currentTier":{"id":"x"}}"#)
645            .create_async()
646            .await;
647        assert_eq!(
648            fetch_plan(&reqwest::Client::new(), &endpoints(&server), "AT").await,
649            None
650        );
651    }
652
653    #[test]
654    fn format_plan_normalises_known_tiers() {
655        assert_eq!(format_plan("Google AI Pro"), "Pro");
656        assert_eq!(format_plan("google_ai_ultra"), "Ultra");
657        assert_eq!(format_plan("GOOGLE_AI_ULTRA"), "Ultra");
658        assert_eq!(format_plan("free-tier"), "Free");
659        assert_eq!(format_plan("Free"), "Free");
660        assert_eq!(format_plan("  legacy_team plan "), "Legacy Team Plan");
661        assert_eq!(format_plan(""), "");
662        let long = format_plan(&"word ".repeat(20));
663        assert!(long.chars().count() <= MAX_PLAN_CHARS);
664        assert!(!long.ends_with(' '));
665    }
666
667    #[tokio::test]
668    async fn refresh_sends_the_form_body_and_parses_the_token() {
669        let mut server = mockito::Server::new_async().await;
670        let m = server
671            .mock("POST", "/token")
672            .match_header("content-type", "application/x-www-form-urlencoded")
673            .match_body(Matcher::AllOf(vec![
674                Matcher::UrlEncoded("grant_type".into(), "refresh_token".into()),
675                Matcher::UrlEncoded("refresh_token".into(), "old-rt".into()),
676                Matcher::UrlEncoded("client_id".into(), "cid".into()),
677                Matcher::UrlEncoded("client_secret".into(), "csecret".into()),
678            ]))
679            .with_status(200)
680            .with_body(r#"{"access_token":"new-at","expires_in":3599,"token_type":"Bearer"}"#)
681            .create_async()
682            .await;
683        let oauth = OauthClient {
684            id: "cid".into(),
685            secret: "csecret".into(),
686        };
687        let before = Utc::now();
688        let refreshed = refresh(
689            &reqwest::Client::new(),
690            &format!("{}/token", server.url()),
691            &oauth,
692            "old-rt",
693        )
694        .await
695        .unwrap();
696        assert_eq!(refreshed.access_token, "new-at");
697        let delta = refreshed.expires_at.timestamp() - before.timestamp();
698        assert!((3590..=3610).contains(&delta), "{delta}");
699        m.assert_async().await;
700    }
701
702    #[tokio::test]
703    async fn refresh_400_is_a_credentials_error_that_does_not_echo_the_body() {
704        let mut server = mockito::Server::new_async().await;
705        server
706            .mock("POST", "/token")
707            .with_status(400)
708            .with_body(r#"{"error":"invalid_grant","error_description":"sensitive detail"}"#)
709            .create_async()
710            .await;
711        let err = refresh(
712            &reqwest::Client::new(),
713            &format!("{}/token", server.url()),
714            &test_client(),
715            "old-rt",
716        )
717        .await
718        .unwrap_err();
719        assert!(matches!(err, AppError::Credentials(_)), "{err:?}");
720        let text = err.to_string();
721        assert!(!text.contains("sensitive detail"));
722        assert!(!text.contains("invalid_grant"));
723        assert!(text.contains("sign in again"));
724    }
725
726    #[tokio::test]
727    async fn refresh_5xx_and_429_are_http_errors_with_a_fixed_body() {
728        for status in [429u16, 503] {
729            let mut server = mockito::Server::new_async().await;
730            server
731                .mock("POST", "/token")
732                .with_status(status.into())
733                .with_body("private upstream text")
734                .create_async()
735                .await;
736            let err = refresh(
737                &reqwest::Client::new(),
738                &format!("{}/token", server.url()),
739                &test_client(),
740                "old-rt",
741            )
742            .await
743            .unwrap_err();
744            match err {
745                AppError::Http { status: got, body } => {
746                    assert_eq!(got, status);
747                    assert_eq!(body, REFRESH_FAILED);
748                }
749                other => panic!("expected Http, got {other:?}"),
750            }
751        }
752    }
753
754    #[tokio::test]
755    async fn refresh_rejects_malformed_success_bodies() {
756        for body in [
757            r#"{"access_token":"","expires_in":3600}"#,
758            r#"{"access_token":"new","expires_in":0}"#,
759            r#"{"access_token":"new","expires_in":"soon"}"#,
760            r#"{"access_token":"new"}"#,
761            "not json",
762        ] {
763            let mut server = mockito::Server::new_async().await;
764            server
765                .mock("POST", "/token")
766                .with_status(200)
767                .with_body(body)
768                .create_async()
769                .await;
770            let err = refresh(
771                &reqwest::Client::new(),
772                &format!("{}/token", server.url()),
773                &test_client(),
774                "old-rt",
775            )
776            .await
777            .unwrap_err();
778            assert!(matches!(err, AppError::Schema(_)), "{body}: {err:?}");
779        }
780    }
781
782    #[test]
783    fn needs_refresh_threshold() {
784        let now = DateTime::parse_from_rfc3339("2026-08-03T12:00:00Z")
785            .unwrap()
786            .with_timezone(&Utc);
787        assert!(needs_refresh(None, now));
788        assert!(needs_refresh(
789            Some(now + chrono::Duration::seconds(REFRESH_BUFFER_SECS - 1)),
790            now
791        ));
792        assert!(needs_refresh(Some(now - chrono::Duration::hours(1)), now));
793        assert!(!needs_refresh(
794            Some(now + chrono::Duration::seconds(REFRESH_BUFFER_SECS + 60)),
795            now
796        ));
797    }
798
799    #[test]
800    fn persisted_token_round_trips_and_is_scoped_to_its_fingerprint() {
801        let td = TempDir::new().unwrap();
802        let cache = Cache::at(td.path().join("antigravity"));
803        cache.ensure_dir().unwrap();
804        let path = oauth_cache_path(&cache);
805        assert_eq!(read_persisted(&path, "abcd"), None);
806
807        let value = PersistedOAuth {
808            fingerprint: "abcd".into(),
809            access_token: "AT".into(),
810            expires_at: DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
811                .unwrap()
812                .with_timezone(&Utc),
813        };
814        write_persisted(&path, &value).unwrap();
815
816        assert_eq!(read_persisted(&path, "abcd"), Some(value.clone()));
817        assert_eq!(read_persisted(&path, "other"), None);
818
819        #[cfg(unix)]
820        {
821            use std::os::unix::fs::PermissionsExt;
822            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
823            assert_eq!(mode & 0o077, 0);
824        }
825    }
826
827    #[test]
828    fn malformed_or_empty_persisted_token_reads_as_none() {
829        let td = TempDir::new().unwrap();
830        let path = td.path().join("oauth.json");
831        std::fs::write(&path, b"{not json").unwrap();
832        assert_eq!(read_persisted(&path, "abcd"), None);
833        std::fs::write(
834            &path,
835            br#"{"fingerprint":"abcd","access_token":"  ","expires_at":"2030-01-01T00:00:00Z"}"#,
836        )
837        .unwrap();
838        assert_eq!(read_persisted(&path, "abcd"), None);
839    }
840}
ai_usagebar/antigravity/cloud.rs

ai_usagebar/antigravity/
cloud.rs
