Skip to main content

ai_usagebar/antigravity/
cloud.rs

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