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/// The persisted token for exactly this keyring session. Absent, unreadable,
358/// malformed, empty, or minted from a different session all read as `None`:
359/// the worst case is one extra refresh round-trip.
360pub fn read_persisted(path: &Path, fingerprint: &str) -> Option<PersistedOAuth> {
361    let bytes = std::fs::read(path).ok()?;
362    let persisted: PersistedOAuth = serde_json::from_slice(&bytes).ok()?;
363    if persisted.fingerprint != fingerprint || persisted.access_token.trim().is_empty() {
364        return None;
365    }
366    Some(persisted)
367}
368
369/// Atomically write the persisted token, owner-only on unix (kiro pattern).
370pub fn write_persisted(path: &Path, value: &PersistedOAuth) -> Result<()> {
371    let bytes = serde_json::to_vec_pretty(value)?;
372    atomic_write(path, &bytes)?;
373    #[cfg(unix)]
374    {
375        use std::os::unix::fs::PermissionsExt;
376        std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))
377            .map_err(|e| AppError::io_at(path, e))?;
378    }
379    Ok(())
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385    use mockito::Matcher;
386    use tempfile::TempDir;
387
388    fn endpoints(server: &mockito::Server) -> Endpoints {
389        let base = server.url();
390        Endpoints {
391            quota: vec![format!("{base}/daily/quota"), format!("{base}/prod/quota")],
392            load_code_assist: vec![format!("{base}/daily/plan"), format!("{base}/prod/plan")],
393            token: format!("{base}/token"),
394        }
395    }
396
397    #[test]
398    fn default_endpoints_try_daily_then_prod() {
399        let e = Endpoints::default();
400        assert_eq!(
401            e.quota,
402            vec![
403                "https://daily-cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary",
404                "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuotaSummary",
405            ]
406        );
407        assert_eq!(
408            e.load_code_assist,
409            vec![
410                "https://daily-cloudcode-pa.googleapis.com/v1internal:loadCodeAssist",
411                "https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist",
412            ]
413        );
414        assert_eq!(e.token, "https://oauth2.googleapis.com/token");
415    }
416
417    fn test_client() -> OauthClient {
418        OauthClient {
419            id: "test-client".into(),
420            secret: "test-client-secret".into(),
421        }
422    }
423
424    #[test]
425    fn oauth_client_needs_both_halves_and_trims_them() {
426        assert!(OauthClient::from_config(None, None).is_none());
427        assert!(OauthClient::from_config(Some("id"), Some("")).is_none());
428        assert!(OauthClient::from_config(Some(" "), Some("s")).is_none());
429        let both = OauthClient::from_config(Some(" my-id "), Some(" my-secret ")).unwrap();
430        assert_eq!(both.id, "my-id");
431        assert_eq!(both.secret, "my-secret");
432    }
433
434    #[tokio::test]
435    async fn quota_falls_through_from_daily_to_prod() {
436        let mut server = mockito::Server::new_async().await;
437        let daily = server
438            .mock("POST", "/daily/quota")
439            .match_header("authorization", "Bearer AT")
440            .match_header("content-type", "application/json")
441            .match_header("accept", "application/json")
442            .match_header("user-agent", "antigravity")
443            .match_body("{}")
444            .with_status(500)
445            .with_body("boom")
446            .expect(1)
447            .create_async()
448            .await;
449        let prod = server
450            .mock("POST", "/prod/quota")
451            .match_header("authorization", "Bearer AT")
452            .match_header("user-agent", "antigravity")
453            .with_status(200)
454            .with_body(r#"{"response":{"buckets":[{"modelId":"gemini","remainingFraction":0.5}]}}"#)
455            .expect(1)
456            .create_async()
457            .await;
458
459        let value = fetch_quota(&reqwest::Client::new(), &endpoints(&server), "AT")
460            .await
461            .unwrap();
462        assert_eq!(value["response"]["buckets"][0]["remainingFraction"], 0.5);
463        daily.assert_async().await;
464        prod.assert_async().await;
465    }
466
467    #[tokio::test]
468    async fn quota_401_stops_without_trying_the_next_base() {
469        let mut server = mockito::Server::new_async().await;
470        let daily = server
471            .mock("POST", "/daily/quota")
472            .with_status(401)
473            .with_body(r#"{"error":{"message":"sensitive detail"}}"#)
474            .expect(1)
475            .create_async()
476            .await;
477        let prod = server
478            .mock("POST", "/prod/quota")
479            .with_status(200)
480            .with_body("{}")
481            .expect(0)
482            .create_async()
483            .await;
484
485        let err = fetch_quota(&reqwest::Client::new(), &endpoints(&server), "AT")
486            .await
487            .unwrap_err();
488        match err {
489            AppError::Http { status, body } => {
490                assert_eq!(status, 401);
491                assert_eq!(body, SESSION_REJECTED);
492            }
493            other => panic!("expected Http, got {other:?}"),
494        }
495        daily.assert_async().await;
496        prod.assert_async().await;
497    }
498
499    #[tokio::test]
500    async fn quota_reports_the_last_error_when_every_base_fails() {
501        let mut server = mockito::Server::new_async().await;
502        server
503            .mock("POST", "/daily/quota")
504            .with_status(500)
505            .create_async()
506            .await;
507        server
508            .mock("POST", "/prod/quota")
509            .with_status(503)
510            .with_body("private upstream text")
511            .create_async()
512            .await;
513
514        let err = fetch_quota(&reqwest::Client::new(), &endpoints(&server), "AT")
515            .await
516            .unwrap_err();
517        match err {
518            AppError::Http { status, body } => {
519                assert_eq!(status, 503);
520                assert!(!body.contains("private upstream text"));
521            }
522            other => panic!("expected Http, got {other:?}"),
523        }
524    }
525
526    #[tokio::test]
527    async fn plan_prefers_paid_tier_and_sends_the_agy_user_agent() {
528        let mut server = mockito::Server::new_async().await;
529        let m = server
530            .mock("POST", "/daily/plan")
531            .match_header("authorization", "Bearer AT")
532            .match_header("user-agent", "agy")
533            .match_body("{}")
534            .with_status(200)
535            .with_body(
536                r#"{"currentTier":{"id":"free-tier","name":"Free"},"paidTier":{"id":"google_ai_pro","name":"Google AI Pro"}}"#,
537            )
538            .create_async()
539            .await;
540
541        let plan = fetch_plan(&reqwest::Client::new(), &endpoints(&server), "AT").await;
542        assert_eq!(plan.as_deref(), Some("Pro"));
543        m.assert_async().await;
544    }
545
546    #[tokio::test]
547    async fn plan_falls_back_to_current_tier_and_to_the_next_base() {
548        let mut server = mockito::Server::new_async().await;
549        server
550            .mock("POST", "/daily/plan")
551            .with_status(500)
552            .create_async()
553            .await;
554        server
555            .mock("POST", "/prod/plan")
556            .with_status(200)
557            .with_body(r#"{"response":{"currentTier":{"name":"GOOGLE_AI_ULTRA"}}}"#)
558            .create_async()
559            .await;
560
561        let plan = fetch_plan(&reqwest::Client::new(), &endpoints(&server), "AT").await;
562        assert_eq!(plan.as_deref(), Some("Ultra"));
563    }
564
565    #[tokio::test]
566    async fn plan_is_none_when_nothing_answers_usefully() {
567        let mut server = mockito::Server::new_async().await;
568        server
569            .mock("POST", "/daily/plan")
570            .with_status(200)
571            .with_body("not json")
572            .create_async()
573            .await;
574        server
575            .mock("POST", "/prod/plan")
576            .with_status(200)
577            .with_body(r#"{"currentTier":{"id":"x"}}"#)
578            .create_async()
579            .await;
580        assert_eq!(
581            fetch_plan(&reqwest::Client::new(), &endpoints(&server), "AT").await,
582            None
583        );
584    }
585
586    #[test]
587    fn format_plan_normalises_known_tiers() {
588        assert_eq!(format_plan("Google AI Pro"), "Pro");
589        assert_eq!(format_plan("google_ai_ultra"), "Ultra");
590        assert_eq!(format_plan("GOOGLE_AI_ULTRA"), "Ultra");
591        assert_eq!(format_plan("free-tier"), "Free");
592        assert_eq!(format_plan("Free"), "Free");
593        assert_eq!(format_plan("  legacy_team plan "), "Legacy Team Plan");
594        assert_eq!(format_plan(""), "");
595        let long = format_plan(&"word ".repeat(20));
596        assert!(long.chars().count() <= MAX_PLAN_CHARS);
597        assert!(!long.ends_with(' '));
598    }
599
600    #[tokio::test]
601    async fn refresh_sends_the_form_body_and_parses_the_token() {
602        let mut server = mockito::Server::new_async().await;
603        let m = server
604            .mock("POST", "/token")
605            .match_header("content-type", "application/x-www-form-urlencoded")
606            .match_body(Matcher::AllOf(vec![
607                Matcher::UrlEncoded("grant_type".into(), "refresh_token".into()),
608                Matcher::UrlEncoded("refresh_token".into(), "old-rt".into()),
609                Matcher::UrlEncoded("client_id".into(), "cid".into()),
610                Matcher::UrlEncoded("client_secret".into(), "csecret".into()),
611            ]))
612            .with_status(200)
613            .with_body(r#"{"access_token":"new-at","expires_in":3599,"token_type":"Bearer"}"#)
614            .create_async()
615            .await;
616        let oauth = OauthClient {
617            id: "cid".into(),
618            secret: "csecret".into(),
619        };
620        let before = Utc::now();
621        let refreshed = refresh(
622            &reqwest::Client::new(),
623            &format!("{}/token", server.url()),
624            &oauth,
625            "old-rt",
626        )
627        .await
628        .unwrap();
629        assert_eq!(refreshed.access_token, "new-at");
630        let delta = refreshed.expires_at.timestamp() - before.timestamp();
631        assert!((3590..=3610).contains(&delta), "{delta}");
632        m.assert_async().await;
633    }
634
635    #[tokio::test]
636    async fn refresh_400_is_a_credentials_error_that_does_not_echo_the_body() {
637        let mut server = mockito::Server::new_async().await;
638        server
639            .mock("POST", "/token")
640            .with_status(400)
641            .with_body(r#"{"error":"invalid_grant","error_description":"sensitive detail"}"#)
642            .create_async()
643            .await;
644        let err = refresh(
645            &reqwest::Client::new(),
646            &format!("{}/token", server.url()),
647            &test_client(),
648            "old-rt",
649        )
650        .await
651        .unwrap_err();
652        assert!(matches!(err, AppError::Credentials(_)), "{err:?}");
653        let text = err.to_string();
654        assert!(!text.contains("sensitive detail"));
655        assert!(!text.contains("invalid_grant"));
656        assert!(text.contains("sign in again"));
657    }
658
659    #[tokio::test]
660    async fn refresh_5xx_and_429_are_http_errors_with_a_fixed_body() {
661        for status in [429u16, 503] {
662            let mut server = mockito::Server::new_async().await;
663            server
664                .mock("POST", "/token")
665                .with_status(status.into())
666                .with_body("private upstream text")
667                .create_async()
668                .await;
669            let err = refresh(
670                &reqwest::Client::new(),
671                &format!("{}/token", server.url()),
672                &test_client(),
673                "old-rt",
674            )
675            .await
676            .unwrap_err();
677            match err {
678                AppError::Http { status: got, body } => {
679                    assert_eq!(got, status);
680                    assert_eq!(body, REFRESH_FAILED);
681                }
682                other => panic!("expected Http, got {other:?}"),
683            }
684        }
685    }
686
687    #[tokio::test]
688    async fn refresh_rejects_malformed_success_bodies() {
689        for body in [
690            r#"{"access_token":"","expires_in":3600}"#,
691            r#"{"access_token":"new","expires_in":0}"#,
692            r#"{"access_token":"new","expires_in":"soon"}"#,
693            r#"{"access_token":"new"}"#,
694            "not json",
695        ] {
696            let mut server = mockito::Server::new_async().await;
697            server
698                .mock("POST", "/token")
699                .with_status(200)
700                .with_body(body)
701                .create_async()
702                .await;
703            let err = refresh(
704                &reqwest::Client::new(),
705                &format!("{}/token", server.url()),
706                &test_client(),
707                "old-rt",
708            )
709            .await
710            .unwrap_err();
711            assert!(matches!(err, AppError::Schema(_)), "{body}: {err:?}");
712        }
713    }
714
715    #[test]
716    fn needs_refresh_threshold() {
717        let now = DateTime::parse_from_rfc3339("2026-08-03T12:00:00Z")
718            .unwrap()
719            .with_timezone(&Utc);
720        assert!(needs_refresh(None, now));
721        assert!(needs_refresh(
722            Some(now + chrono::Duration::seconds(REFRESH_BUFFER_SECS - 1)),
723            now
724        ));
725        assert!(needs_refresh(Some(now - chrono::Duration::hours(1)), now));
726        assert!(!needs_refresh(
727            Some(now + chrono::Duration::seconds(REFRESH_BUFFER_SECS + 60)),
728            now
729        ));
730    }
731
732    #[test]
733    fn persisted_token_round_trips_and_is_scoped_to_its_fingerprint() {
734        let td = TempDir::new().unwrap();
735        let cache = Cache::at(td.path().join("antigravity"));
736        cache.ensure_dir().unwrap();
737        let path = oauth_cache_path(&cache);
738        assert_eq!(read_persisted(&path, "abcd"), None);
739
740        let value = PersistedOAuth {
741            fingerprint: "abcd".into(),
742            access_token: "AT".into(),
743            expires_at: DateTime::parse_from_rfc3339("2030-01-01T00:00:00Z")
744                .unwrap()
745                .with_timezone(&Utc),
746        };
747        write_persisted(&path, &value).unwrap();
748
749        assert_eq!(read_persisted(&path, "abcd"), Some(value.clone()));
750        assert_eq!(read_persisted(&path, "other"), None);
751
752        #[cfg(unix)]
753        {
754            use std::os::unix::fs::PermissionsExt;
755            let mode = std::fs::metadata(&path).unwrap().permissions().mode();
756            assert_eq!(mode & 0o077, 0);
757        }
758    }
759
760    #[test]
761    fn malformed_or_empty_persisted_token_reads_as_none() {
762        let td = TempDir::new().unwrap();
763        let path = td.path().join("oauth.json");
764        std::fs::write(&path, b"{not json").unwrap();
765        assert_eq!(read_persisted(&path, "abcd"), None);
766        std::fs::write(
767            &path,
768            br#"{"fingerprint":"abcd","access_token":"  ","expires_at":"2030-01-01T00:00:00Z"}"#,
769        )
770        .unwrap();
771        assert_eq!(read_persisted(&path, "abcd"), None);
772    }
773}
ai_usagebar/antigravity/cloud.rs

ai_usagebar/antigravity/
cloud.rs
