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