Skip to main content

ai_usagebar/openai/
fetch.rs

1//! Orchestrate: read ~/.codex/auth.json → maybe refresh → fetch usage → cache.
2//!
3//! Mirrors `anthropic::fetch::fetch_snapshot` but for the Codex OAuth flow.
4
5use std::path::{Path, PathBuf};
6use std::time::Duration;
7
8use chrono::Utc;
9
10use crate::cache::{Cache, LockGuard, acquire_lock_async};
11use crate::error::{AppError, Result};
12use crate::usage::OpenAiSnapshot;
13
14use super::creds::{self, Tokens};
15use super::oauth;
16use super::types::UsageResponse;
17
18pub const USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";
19const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
20const REFRESH_TIMEOUT: Duration = Duration::from_secs(25);
21const LOCK_TIMEOUT: Duration = Duration::from_secs(45);
22/// A switch holds the credentials lock for a few file writes only.
23const CREDENTIALS_LOCK_TIMEOUT: Duration = Duration::from_secs(30);
24/// How often a route that moved under the lock is followed. Each move means a
25/// switch landed between two looks, so a second one is already rare.
26const ROUTE_ATTEMPTS: usize = 3;
27
28#[derive(Debug, Clone)]
29pub struct Endpoints {
30    pub usage: String,
31    pub token: String,
32}
33
34impl Default for Endpoints {
35    fn default() -> Self {
36        Self {
37            usage: USAGE_URL.into(),
38            token: oauth::TOKEN_URL.into(),
39        }
40    }
41}
42
43/// This vendor's [`Outcome`](crate::outcome::Outcome) — the shared shape,
44/// specialised to its snapshot.
45pub type FetchOutcome = crate::outcome::Outcome<OpenAiSnapshot>;
46
47pub async fn fetch_snapshot(
48    client: &reqwest::Client,
49    creds_path: &Path,
50    cache: &Cache,
51    endpoints: &Endpoints,
52    cache_ttl: Duration,
53) -> Result<FetchOutcome> {
54    let route = || Ok(creds_path.to_path_buf());
55    fetch_snapshot_routed(client, route, cache, endpoints, cache_ttl).await
56}
57
58/// [`fetch_snapshot`] for a login `account switch --codex` can move: `route`
59/// names the file to read, and is asked again once that file's lock is held.
60/// A path chosen before the lock can name the default slot after a switch
61/// installed another account there, and that account's usage would then land
62/// in this one's cache.
63pub async fn fetch_snapshot_routed(
64    client: &reqwest::Client,
65    route: impl Fn() -> Result<PathBuf>,
66    cache: &Cache,
67    endpoints: &Endpoints,
68    cache_ttl: Duration,
69) -> Result<FetchOutcome> {
70    cache.ensure_dir()?;
71    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
72    let (creds_path, _credentials_lock) = lock_route(&route).await?;
73    let creds_path = creds_path.as_path();
74
75    let mut auth = creds::read_from(creds_path)?;
76    let plan_hint = auth.tokens.plan_type_from_id_token();
77
78    // Corrupt fresh cache falls through to a live fetch rather than returning
79    // an all-zero snapshot.
80    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
81        && let Ok(outcome) = reuse(bytes, cache, false, plan_hint.as_deref())
82    {
83        return Ok(outcome);
84    }
85
86    // Maybe refresh — Codex CLI doesn't always populate expires_at, so we use
87    // the id_token's exp claim.
88    let now = Utc::now().timestamp();
89    if oauth::needs_refresh(auth.tokens.expires_at_secs(), now) {
90        match tokio::time::timeout(
91            REFRESH_TIMEOUT,
92            oauth::refresh(client, &endpoints.token, &auth.tokens.refresh_token),
93        )
94        .await
95        {
96            Ok(Ok(rr)) => {
97                auth.tokens.access_token = rr.access_token;
98                // A rotated refresh token exists only in memory until it is
99                // persisted; losing it silently logs the user out on the next
100                // run. See the matching comment in `anthropic::fetch`.
101                let rotated = rr.refresh_token.is_some();
102                if let Some(rt) = rr.refresh_token {
103                    auth.tokens.refresh_token = rt;
104                }
105                if let Some(id) = rr.id_token {
106                    auth.tokens.id_token = id;
107                }
108                // Expiry is normally read from the id_token's exp claim, so a
109                // refresh that returns no new id_token would leave the old
110                // (expired) claim in place and make every later run refresh
111                // again. Record the response's own `expires_in` as the explicit
112                // `expires_at` so the expiry is known either way.
113                if let Some(secs) = rr.expires_in
114                    && let Some(dt) = chrono::DateTime::from_timestamp(now + secs as i64, 0)
115                {
116                    auth.tokens.expires_at = Some(dt.to_rfc3339());
117                }
118                if let Err(e) = creds::write_back(creds_path, &auth)
119                    && rotated
120                {
121                    let msg = format!(
122                        "refreshed token could not be saved ({e}); the rotated \
123                         refresh token is lost — re-run `codex login`"
124                    );
125                    cache.write_last_error(0, &msg);
126                    return handle_auth_failure(cache, plan_hint.as_deref(), false);
127                }
128            }
129            Ok(Err(AppError::Http { status, body })) => {
130                cache.write_last_error(status, &body);
131                return handle_auth_failure(cache, plan_hint.as_deref(), false);
132            }
133            Ok(Err(e)) if e.is_transient() => {
134                return handle_auth_failure(cache, plan_hint.as_deref(), true);
135            }
136            Ok(Err(e)) => {
137                cache.write_last_error(0, &e.to_string());
138                return handle_auth_failure(cache, plan_hint.as_deref(), false);
139            }
140            Err(_) => return handle_auth_failure(cache, plan_hint.as_deref(), true),
141        }
142    }
143
144    match tokio::time::timeout(
145        HTTP_TIMEOUT,
146        fetch_usage(client, &endpoints.usage, &auth.tokens),
147    )
148    .await
149    {
150        Ok(Ok(response)) => {
151            // Cache what we parse, not what arrived. The raw body carries the
152            // account's `user_id`, `account_id` and `email`, none of which any
153            // renderer reads — writing the parsed response is an allowlist by
154            // construction, so a field OpenAI adds later cannot quietly start
155            // living on disk. Same rule Command Code follows.
156            cache.write_payload(&serde_json::to_vec(&response)?)?;
157            let snap = response.into_snapshot(plan_hint.as_deref())?;
158            Ok(crate::outcome::Outcome::fresh(snap))
159        }
160        Ok(Err(AppError::Http { status, body })) => {
161            cache.mark_stale();
162            let last_error = Some(cache.write_last_error(status, &body));
163            fallback(
164                cache,
165                plan_hint.as_deref(),
166                last_error,
167                AppError::Http { status, body },
168            )
169        }
170        Ok(Err(e)) if e.is_transient() => fallback_silent(cache, plan_hint.as_deref(), e),
171        Ok(Err(e)) => {
172            cache.mark_stale();
173            let last_error = Some(cache.write_last_error(0, &e.to_string()));
174            fallback(cache, plan_hint.as_deref(), last_error, e)
175        }
176        Err(_) => fallback_silent(
177            cache,
178            plan_hint.as_deref(),
179            AppError::Transport("openai: usage request timed out".into()),
180        ),
181    }
182}
183
184/// Lock the credential file `route` names and confirm, under that lock, that
185/// it still names the same one. A switch holds every directory it touches, the
186/// default slot's included, so an answer given under the lock stays true until
187/// the read-refresh-write below is done.
188async fn lock_route(route: &impl Fn() -> Result<PathBuf>) -> Result<(PathBuf, Option<LockGuard>)> {
189    let mut path = route()?;
190    for _ in 0..ROUTE_ATTEMPTS {
191        let lock = lock_credentials(&path).await?;
192        let settled = route()?;
193        if settled == path {
194            return Ok((path, lock));
195        }
196        path = settled;
197    }
198    Err(AppError::Transport(
199        "openai: the Codex login kept moving between accounts".into(),
200    ))
201}
202
203/// `account switch --codex` moves these files under the same lock, so a
204/// refresh can never write one account's tokens over another's. Only an
205/// existing directory is locked: creating one is not a fetch's business.
206async fn lock_credentials(creds_path: &Path) -> Result<Option<LockGuard>> {
207    let lock = super::account::lock_path(creds_path);
208    Ok(match lock.parent() {
209        Some(dir) if dir.is_dir() => {
210            Some(acquire_lock_async(&lock, CREDENTIALS_LOCK_TIMEOUT).await?)
211        }
212        _ => None,
213    })
214}
215
216fn reuse(
217    bytes: Vec<u8>,
218    cache: &Cache,
219    stale: bool,
220    plan_hint: Option<&str>,
221) -> Result<FetchOutcome> {
222    let snap = parse_payload(&bytes, plan_hint)?;
223    Ok(crate::outcome::Outcome::cached(snap, cache, stale))
224}
225
226fn fallback(
227    cache: &Cache,
228    plan_hint: Option<&str>,
229    last_error: Option<(u16, String)>,
230    original: AppError,
231) -> Result<FetchOutcome> {
232    crate::outcome::fallback(cache, last_error, original, |bytes| {
233        parse_payload(bytes, plan_hint)
234    })
235}
236
237fn fallback_silent(
238    cache: &Cache,
239    plan_hint: Option<&str>,
240    original: AppError,
241) -> Result<FetchOutcome> {
242    crate::outcome::fallback(cache, None, original, |bytes| {
243        parse_payload(bytes, plan_hint)
244    })
245}
246
247/// The one place a *synthesized* error beats the original: the refresh failed,
248/// and "run `codex login` to re-auth" tells the user what to do about it,
249/// which the underlying OAuth error does not.
250fn handle_auth_failure(
251    cache: &Cache,
252    plan_hint: Option<&str>,
253    transient: bool,
254) -> Result<FetchOutcome> {
255    let original = if transient {
256        AppError::Transport("openai: no cache and refresh failed transiently".into())
257    } else {
258        AppError::Credentials("openai: token refresh failed; run `codex login` to re-auth".into())
259    };
260    crate::outcome::fallback(cache, None, original, |bytes| {
261        parse_payload(bytes, plan_hint)
262    })
263}
264
265fn parse_payload(bytes: &[u8], plan_hint: Option<&str>) -> Result<OpenAiSnapshot> {
266    parse_response(bytes)?.into_snapshot(plan_hint)
267}
268
269/// The wire response, before it becomes a snapshot. Split out so the live path
270/// can cache the parsed form rather than the raw body.
271fn parse_response(bytes: &[u8]) -> Result<UsageResponse> {
272    Ok(serde_json::from_slice(bytes)?)
273}
274
275fn authorized(client: &reqwest::Client, url: String, t: &Tokens) -> reqwest::RequestBuilder {
276    let mut req = client
277        .get(url)
278        .header("Authorization", format!("Bearer {}", t.access_token))
279        .header("User-Agent", "codex-cli");
280    if let Some(aid) = t.account_id.as_deref() {
281        req = req.header("ChatGPT-Account-Id", aid);
282    }
283    req
284}
285
286/// Cache the usage payload, grafting on the reset-credit expiries from the
287/// second endpoint when they exist. The graft is a field insert into the
288/// original JSON: re-serializing our typed `UsageResponse` would drop every
289/// unknown key the API still sends (`user_id`, tomorrow's new window, …),
290/// and a cache holding only the usage bytes would keep the count while the
291/// deadline beside it vanished for the rest of the TTL. The inserted
292/// `credits` array is our typed projection — status + expiry, never the
293/// redemption `id`.
294async fn fetch_usage(client: &reqwest::Client, url: &str, t: &Tokens) -> Result<UsageResponse> {
295    let resp = authorized(client, url.to_string(), t).send().await?;
296    let status = resp.status();
297    let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
298
299    if !status.is_success() {
300        let body: String = String::from_utf8_lossy(&bytes).chars().take(200).collect();
301        return Err(AppError::Http {
302            status: status.as_u16(),
303            body,
304        });
305    }
306    let mut parsed: UsageResponse = serde_json::from_slice(&bytes)
307        .map_err(|e| AppError::Schema(format!("openai usage response: {e}")))?;
308    parsed.rate_limit_reset_credits = enrich_reset_credits(client, url, t, &parsed).await;
309    // Reject drift here, while the caller can still fall back to a good cache.
310    parsed.clone().into_snapshot(None)?;
311    Ok(parsed)
312}
313
314/// The usage endpoint reports how many banked resets exist but not when they
315/// expire; a second call carries the per-credit detail. That call is strictly
316/// additive: its failure leaves the count exactly as the usage endpoint
317/// reported it, because a count with no deadline is still true, and it is not
318/// worth failing a whole refresh over the deadline alone. The count itself
319/// always stays the usage endpoint's — the two responses can disagree across a
320/// redemption, and the one that also carries the quota figures is the one the
321/// rest of the snapshot is consistent with.
322async fn enrich_reset_credits(
323    client: &reqwest::Client,
324    usage_url: &str,
325    t: &Tokens,
326    parsed: &UsageResponse,
327) -> Option<super::types::ResetCreditsBlock> {
328    let mut block = parsed.rate_limit_reset_credits.clone()?;
329    if block.available_count == 0 {
330        return Some(block);
331    }
332    if let Ok(details) = fetch_reset_credits(client, usage_url, t).await {
333        block.credits = details.credits;
334    }
335    Some(block)
336}
337
338async fn fetch_reset_credits(
339    client: &reqwest::Client,
340    usage_url: &str,
341    t: &Tokens,
342) -> Result<super::types::ResetCreditsBlock> {
343    let base = usage_url.strip_suffix("/usage").unwrap_or(usage_url);
344    let resp = authorized(client, format!("{base}/rate-limit-reset-credits"), t)
345        .send()
346        .await?;
347    let status = resp.status();
348    let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
349    if !status.is_success() {
350        // The body is discarded rather than reported: this call is optional,
351        // its failure never reaches the user, and it would only carry an
352        // account-identifying error into a log.
353        return Err(AppError::Http {
354            status: status.as_u16(),
355            body: String::new(),
356        });
357    }
358    serde_json::from_slice(&bytes)
359        .map_err(|_| AppError::Schema("openai reset credits response is invalid".into()))
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365    use base64::Engine;
366    use std::io::Write;
367    use std::sync::atomic::{AtomicUsize, Ordering};
368    use tempfile::{NamedTempFile, TempDir};
369
370    fn fake_jwt(claims: serde_json::Value) -> String {
371        let h = base64::engine::general_purpose::URL_SAFE_NO_PAD
372            .encode(br#"{"alg":"none","typ":"JWT"}"#);
373        let p =
374            base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(claims.to_string().as_bytes());
375        format!("{h}.{p}.sig")
376    }
377
378    fn future_creds() -> NamedTempFile {
379        // exp 1h in the future.
380        let exp = Utc::now().timestamp() + 3600;
381        let jwt = fake_jwt(serde_json::json!({
382            "exp": exp,
383            "https://api.openai.com/auth": {"chatgpt_plan_type": "plus"}
384        }));
385        let body = format!(
386            r#"{{"tokens":{{"access_token":"AT","refresh_token":"RT","id_token":"{jwt}",
387                "account_id":"acc"}}}}"#
388        );
389        let mut f = NamedTempFile::new().unwrap();
390        f.write_all(body.as_bytes()).unwrap();
391        f.flush().unwrap();
392        f
393    }
394
395    fn cache_fixture() -> (TempDir, Cache) {
396        let td = TempDir::new().unwrap();
397        let c = Cache::at(td.path().join("openai"));
398        c.ensure_dir().unwrap();
399        (td, c)
400    }
401
402    #[tokio::test]
403    async fn live_200_returns_snapshot_with_plan_from_id_token() {
404        let mut server = mockito::Server::new_async().await;
405        server
406            .mock("GET", "/backend-api/wham/usage")
407            .with_status(200)
408            .with_body(
409                r#"{"plan_type":"plus","rate_limit":{
410                "primary_window":{"used_percent":1,"limit_window_seconds":18000,"reset_at":1779597324},
411                "secondary_window":{"used_percent":0,"limit_window_seconds":604800,"reset_at":1780184124}
412            }}"#,
413            )
414            .create_async()
415            .await;
416        let (_td, cache) = cache_fixture();
417        let creds = future_creds();
418        let client = reqwest::Client::new();
419        let endpoints = Endpoints {
420            usage: format!("{}/backend-api/wham/usage", server.url()),
421            token: format!("{}/oauth/token", server.url()),
422        };
423        let out = fetch_snapshot(
424            &client,
425            creds.path(),
426            &cache,
427            &endpoints,
428            Duration::from_secs(0),
429        )
430        .await
431        .unwrap();
432        assert_eq!(out.snapshot.plan, "ChatGPT Plus");
433        assert_eq!(out.snapshot.session.as_ref().unwrap().utilization_pct, 1);
434        assert!(!out.stale);
435    }
436
437    #[tokio::test]
438    async fn weekly_only_primary_returns_weekly_snapshot() {
439        let mut server = mockito::Server::new_async().await;
440        server
441            .mock("GET", "/backend-api/wham/usage")
442            .with_status(200)
443            .with_body(
444                r#"{"plan_type":"prolite","rate_limit":{
445                "primary_window":{"used_percent":66,"limit_window_seconds":604800,"reset_at":1785261834},
446                "secondary_window":null
447            }}"#,
448            )
449            .create_async()
450            .await;
451        let (_td, cache) = cache_fixture();
452        let creds = future_creds();
453        let endpoints = Endpoints {
454            usage: format!("{}/backend-api/wham/usage", server.url()),
455            token: format!("{}/oauth/token", server.url()),
456        };
457        let out = fetch_snapshot(
458            &reqwest::Client::new(),
459            creds.path(),
460            &cache,
461            &endpoints,
462            Duration::from_secs(0),
463        )
464        .await
465        .unwrap();
466        assert!(out.snapshot.session.is_none());
467        assert_eq!(out.snapshot.weekly.unwrap().utilization_pct, 66);
468    }
469
470    #[tokio::test]
471    async fn corrupt_fresh_cache_refetches_instead_of_showing_an_empty_snapshot() {
472        // `reuse` used to swallow a parse failure into `empty(plan_hint)` and
473        // serve that all-zero snapshot for the rest of the TTL.
474        let mut server = mockito::Server::new_async().await;
475        server
476            .mock("GET", "/backend-api/wham/usage")
477            .with_status(200)
478            .with_body(
479                r#"{"plan_type":"pro","rate_limit":{"primary_window":{"used_percent":37,"limit_window_seconds":18000}}}"#,
480            )
481            .create_async()
482            .await;
483
484        let (_td, cache) = cache_fixture();
485        cache.write_payload(b"{ truncated").unwrap();
486
487        let creds = future_creds();
488        let client = reqwest::Client::new();
489        let endpoints = Endpoints {
490            usage: format!("{}/backend-api/wham/usage", server.url()),
491            token: format!("{}/oauth/token", server.url()),
492        };
493        // A long TTL: the payload IS fresh, it is simply unusable.
494        let out = fetch_snapshot(
495            &client,
496            creds.path(),
497            &cache,
498            &endpoints,
499            Duration::from_secs(3600),
500        )
501        .await
502        .unwrap();
503        assert_eq!(out.snapshot.session.as_ref().unwrap().utilization_pct, 37);
504        assert!(!out.stale);
505    }
506
507    /// The expiry lives behind a second endpoint. It has to reach the cache
508    /// with the usage figures, or the deadline disappears for the rest of the
509    /// TTL while the count beside it stays on screen.
510    #[tokio::test]
511    async fn banked_reset_expiries_are_fetched_and_cached_with_the_usage_figures() {
512        let mut server = mockito::Server::new_async().await;
513        let usage = server
514            .mock("GET", "/backend-api/wham/usage")
515            .with_body(
516                r#"{"plan_type":"plus","future_field":true,"rate_limit":{
517                    "primary_window":{"used_percent":81,"limit_window_seconds":18000,"reset_at":1786536977}},
518                    "rate_limit_reset_credits":{"available_count":2}}"#,
519            )
520            .create_async()
521            .await;
522        let details = server
523            .mock("GET", "/backend-api/wham/rate-limit-reset-credits")
524            .with_body(
525                r#"{"available_count":2,"credits":[
526                    {"id":"c1","status":"available","title":"Full reset (Weekly + 5 hr)","expires_at":"2026-07-17T00:00:00Z"},
527                    {"id":"c2","status":"available","title":"Full reset (Weekly + 5 hr)","expires_at":"2026-08-01T00:00:00Z"}]}"#,
528            )
529            .create_async()
530            .await;
531
532        let (_td, cache) = cache_fixture();
533        let creds = future_creds();
534        let endpoints = Endpoints {
535            usage: format!("{}/backend-api/wham/usage", server.url()),
536            token: format!("{}/oauth/token", server.url()),
537        };
538        let out = fetch_snapshot(
539            &reqwest::Client::new(),
540            creds.path(),
541            &cache,
542            &endpoints,
543            Duration::from_secs(0),
544        )
545        .await
546        .unwrap();
547        usage.assert_async().await;
548        details.assert_async().await;
549        assert_eq!(out.snapshot.reset_credits.available, 2);
550        assert_eq!(
551            out.snapshot.reset_credits.next_expiry(),
552            Some("2026-07-17T00:00:00Z".parse().unwrap())
553        );
554
555        // The redemption id is what spends a credit; it must not be written to
556        // disk just because it shared a response with the expiry.
557        let cached = std::fs::read_to_string(cache.payload_path()).unwrap();
558        assert!(!cached.contains("\"c1\""), "{cached}");
559        // The cache holds the parsed response, so it holds only what a
560        // renderer reads. An unknown field is dropped rather than kept — the
561        // cache is a short-lived copy of something refetchable, and keeping
562        // the whole body is how the account's identity ended up on disk.
563        assert!(
564            !cached.contains("future_field"),
565            "the cache must not carry fields nothing parses: {cached}"
566        );
567        let reused = parse_payload(cached.as_bytes(), None).unwrap();
568        assert_eq!(reused.reset_credits, out.snapshot.reset_credits);
569    }
570
571    /// The response carries the account's identity — `user_id`, `account_id`
572    /// and `email` — and no renderer reads any of it. Caching the raw body put
573    /// all three on disk for the life of the TTL. Caching the *parsed*
574    /// response is an allowlist by construction: a field OpenAI adds later
575    /// cannot start living there without someone adding it to the type first.
576    #[tokio::test]
577    async fn the_cache_holds_no_account_identity() {
578        let mut server = mockito::Server::new_async().await;
579        let usage = server
580            .mock("GET", "/backend-api/wham/usage")
581            .with_body(
582                r#"{"plan_type":"pro",
583                    "user_id":"user_abc123",
584                    "account_id":"acct_abc123",
585                    "email":"person@example.test",
586                    "rate_limit":{"primary_window":{"used_percent":5,
587                                  "limit_window_seconds":604800}}}"#,
588            )
589            .create_async()
590            .await;
591
592        let (_td, cache) = cache_fixture();
593        let creds = future_creds();
594        let endpoints = Endpoints {
595            usage: format!("{}/backend-api/wham/usage", server.url()),
596            token: format!("{}/oauth/token", server.url()),
597        };
598        let out = fetch_snapshot(
599            &reqwest::Client::new(),
600            creds.path(),
601            &cache,
602            &endpoints,
603            Duration::from_secs(0),
604        )
605        .await
606        .unwrap();
607        usage.assert_async().await;
608
609        // The figures still arrive.
610        assert_eq!(out.snapshot.weekly.as_ref().unwrap().utilization_pct, 5);
611
612        let cached = std::fs::read_to_string(cache.payload_path()).unwrap();
613        for identity in ["user_abc123", "acct_abc123", "person@example.test"] {
614            assert!(
615                !cached.contains(identity),
616                "{identity} reached the cache: {cached}"
617            );
618        }
619        for key in ["user_id", "account_id", "email"] {
620            assert!(!cached.contains(key), "{key} reached the cache: {cached}");
621        }
622    }
623
624    /// The detail call is an extra. When it fails, the count the usage
625    /// endpoint reported is still true and still worth showing — refusing the
626    /// whole refresh over a missing deadline would cost the quota figures too.
627    #[tokio::test]
628    async fn a_failed_detail_call_keeps_the_count_from_the_usage_response() {
629        let mut server = mockito::Server::new_async().await;
630        server
631            .mock("GET", "/backend-api/wham/usage")
632            .with_body(
633                r#"{"plan_type":"plus","rate_limit":{
634                    "primary_window":{"used_percent":10,"limit_window_seconds":18000}},
635                    "rate_limit_reset_credits":{"available_count":1}}"#,
636            )
637            .create_async()
638            .await;
639        let details = server
640            .mock("GET", "/backend-api/wham/rate-limit-reset-credits")
641            .with_status(404)
642            .create_async()
643            .await;
644
645        let (_td, cache) = cache_fixture();
646        let creds = future_creds();
647        let endpoints = Endpoints {
648            usage: format!("{}/backend-api/wham/usage", server.url()),
649            token: format!("{}/oauth/token", server.url()),
650        };
651        let out = fetch_snapshot(
652            &reqwest::Client::new(),
653            creds.path(),
654            &cache,
655            &endpoints,
656            Duration::from_secs(0),
657        )
658        .await
659        .unwrap();
660        details.assert_async().await;
661        assert!(!out.stale);
662        assert_eq!(out.snapshot.session.as_ref().unwrap().utilization_pct, 10);
663        assert_eq!(out.snapshot.reset_credits.available, 1);
664        assert!(out.snapshot.reset_credits.credits.is_empty());
665    }
666
667    /// With nothing banked there is nothing to detail. Asking anyway would
668    /// double every refresh's request count for every account that has none.
669    #[tokio::test]
670    async fn no_banked_resets_means_no_second_request() {
671        let mut server = mockito::Server::new_async().await;
672        server
673            .mock("GET", "/backend-api/wham/usage")
674            .with_body(
675                r#"{"plan_type":"plus","rate_limit":{
676                    "primary_window":{"used_percent":10,"limit_window_seconds":18000}},
677                    "rate_limit_reset_credits":{"available_count":0}}"#,
678            )
679            .create_async()
680            .await;
681        let details = server
682            .mock("GET", "/backend-api/wham/rate-limit-reset-credits")
683            .expect(0)
684            .create_async()
685            .await;
686
687        let (_td, cache) = cache_fixture();
688        let creds = future_creds();
689        let endpoints = Endpoints {
690            usage: format!("{}/backend-api/wham/usage", server.url()),
691            token: format!("{}/oauth/token", server.url()),
692        };
693        let out = fetch_snapshot(
694            &reqwest::Client::new(),
695            creds.path(),
696            &cache,
697            &endpoints,
698            Duration::from_secs(0),
699        )
700        .await
701        .unwrap();
702        details.assert_async().await;
703        assert!(out.snapshot.reset_credits.is_empty());
704    }
705
706    #[tokio::test]
707    async fn http_500_falls_back_to_cache_when_present() {
708        let mut server = mockito::Server::new_async().await;
709        server
710            .mock("GET", "/backend-api/wham/usage")
711            .with_status(500)
712            .with_body(r#"{"error":{"message":"upstream"}}"#)
713            .create_async()
714            .await;
715        let (_td, cache) = cache_fixture();
716        cache
717            .write_payload(
718                br#"{"plan_type":"pro","rate_limit":{"primary_window":{"used_percent":50,"limit_window_seconds":18000}}}"#,
719            )
720            .unwrap();
721        let creds = future_creds();
722        let client = reqwest::Client::new();
723        let endpoints = Endpoints {
724            usage: format!("{}/backend-api/wham/usage", server.url()),
725            token: format!("{}/oauth/token", server.url()),
726        };
727        let out = fetch_snapshot(
728            &client,
729            creds.path(),
730            &cache,
731            &endpoints,
732            Duration::from_secs(0),
733        )
734        .await
735        .unwrap();
736        assert!(out.stale);
737        assert_eq!(out.snapshot.session.as_ref().unwrap().utilization_pct, 50);
738        assert_eq!(out.last_error.as_ref().map(|(c, _)| *c), Some(500));
739    }
740
741    /// Future-dated credentials at `path` whose access token is `token`.
742    fn write_creds(path: &Path, token: &str) {
743        let exp = Utc::now().timestamp() + 3600;
744        let jwt = fake_jwt(serde_json::json!({
745            "exp": exp,
746            "https://api.openai.com/auth": {"chatgpt_plan_type": "plus"}
747        }));
748        std::fs::create_dir_all(path.parent().unwrap()).unwrap();
749        std::fs::write(
750            path,
751            format!(
752                r#"{{"tokens":{{"access_token":"{token}","refresh_token":"RT",
753                    "id_token":"{jwt}","account_id":"acc"}}}}"#
754            ),
755        )
756        .unwrap();
757    }
758
759    #[tokio::test]
760    async fn a_route_that_moves_under_the_lock_is_followed() {
761        let mut server = mockito::Server::new_async().await;
762        let usage = server
763            .mock("GET", "/backend-api/wham/usage")
764            .match_header("authorization", "Bearer AT-B")
765            .with_status(200)
766            .with_body(
767                r#"{"plan_type":"plus","rate_limit":{
768                "primary_window":{"used_percent":7,"limit_window_seconds":18000,"reset_at":1779597324},
769                "secondary_window":null
770            }}"#,
771            )
772            .expect(1)
773            .create_async()
774            .await;
775        let (td, cache) = cache_fixture();
776        let a = td.path().join("a").join("auth.json");
777        let b = td.path().join("b").join("auth.json");
778        write_creds(&a, "AT-A");
779        write_creds(&b, "AT-B");
780        let endpoints = Endpoints {
781            usage: format!("{}/backend-api/wham/usage", server.url()),
782            token: format!("{}/oauth/token", server.url()),
783        };
784        // The first look names A; by the time A's lock is held, a switch has
785        // moved this account's login to B.
786        let looks = AtomicUsize::new(0);
787        let route = || {
788            Ok(if looks.fetch_add(1, Ordering::SeqCst) == 0 {
789                a.clone()
790            } else {
791                b.clone()
792            })
793        };
794        let out = fetch_snapshot_routed(
795            &reqwest::Client::new(),
796            route,
797            &cache,
798            &endpoints,
799            Duration::ZERO,
800        )
801        .await
802        .unwrap();
803        usage.assert_async().await;
804        assert_eq!(out.snapshot.session.as_ref().unwrap().utilization_pct, 7);
805        assert_eq!(looks.load(Ordering::SeqCst), 3);
806    }
807
808    #[tokio::test]
809    async fn a_route_that_never_settles_is_transient() {
810        let mut server = mockito::Server::new_async().await;
811        let usage = server
812            .mock("GET", "/backend-api/wham/usage")
813            .expect(0)
814            .create_async()
815            .await;
816        let (td, cache) = cache_fixture();
817        let endpoints = Endpoints {
818            usage: format!("{}/backend-api/wham/usage", server.url()),
819            token: format!("{}/oauth/token", server.url()),
820        };
821        let looks = AtomicUsize::new(0);
822        let route = || {
823            let look = looks.fetch_add(1, Ordering::SeqCst);
824            Ok(td.path().join(format!("auth-{look}.json")))
825        };
826        let error = fetch_snapshot_routed(
827            &reqwest::Client::new(),
828            route,
829            &cache,
830            &endpoints,
831            Duration::ZERO,
832        )
833        .await
834        .unwrap_err();
835        usage.assert_async().await;
836        assert!(error.is_transient(), "{error}");
837        assert_eq!(looks.load(Ordering::SeqCst), ROUTE_ATTEMPTS + 1);
838    }
839}