ai_usagebar/droid/
fetch.rs

1//! Read the Factory subscription limits and keep only normalized metrics in
2//! the local cache.
3//!
4//! The JWT travels only as the `Authorization: Bearer` header of a GET to the
5//! fixed Factory host, over a client that rejects every redirect. Neither the
6//! raw response nor the credential is persisted — the cache holds the
7//! projected windows and the login's truncated fingerprint.
8
9use std::sync::LazyLock;
10use std::time::Duration;
11
12use crate::cache::{Cache, acquire_lock_async};
13use crate::error::{AppError, Result};
14use crate::usage::{DroidSnapshot, UsageWindow};
15use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
16
17use super::creds::{self, Credentials};
18use super::{renew, types};
19
20const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
21const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
22const LIMITS_PATH: &str = "/api/billing/limits";
23pub const API_BASE: &str = "https://api.factory.ai";
24
25#[derive(Debug, Clone)]
26pub struct Endpoints {
27    pub base: String,
28}
29
30impl Endpoints {
31    pub fn limits_url(&self) -> String {
32        format!("{}{}", self.base.trim_end_matches('/'), LIMITS_PATH)
33    }
34}
35
36impl Default for Endpoints {
37    fn default() -> Self {
38        Self {
39            base: API_BASE.to_string(),
40        }
41    }
42}
43
44pub type FetchOutcome = crate::outcome::Outcome<DroidSnapshot>;
45
46/// The client refuses every redirect, so the CLI token is never forwarded or
47/// re-sent through one, including a same-origin redirect.
48fn build_client(timeout: Duration) -> reqwest::Result<reqwest::Client> {
49    reqwest::Client::builder()
50        .timeout(timeout)
51        .redirect(reqwest::redirect::Policy::none())
52        .build()
53}
54
55/// Droid keeps its own client instead of the process-wide one the widget and
56/// TUI share: that client follows redirects, and this provider must never
57/// send the CLI token through one.
58static CLIENT: LazyLock<Option<reqwest::Client>> =
59    LazyLock::new(|| build_client(HTTP_TIMEOUT).ok());
60
61fn http_client() -> Result<&'static reqwest::Client> {
62    CLIENT
63        .as_ref()
64        .ok_or_else(|| AppError::Other("Droid HTTP client initialization failed".into()))
65}
66
67/// Production entry point; the CLI login is resolved on every call — before
68/// any client or cache directory is touched — so a changed login is reported
69/// before a stale cache can answer for it. A JWT whose `exp` has passed is
70/// never sent to the API: it gets the same cache fallback a 401 does.
71pub async fn fetch_snapshot(
72    cfg: &crate::config::DroidConfig,
73    cache: &Cache,
74    cache_ttl: Duration,
75) -> Result<FetchOutcome> {
76    let mut credentials = super::resolve_credentials(cfg)?;
77    if cfg.cli_renewal && credentials.expires_within(renew::RENEW_MARGIN, chrono::Utc::now()) {
78        credentials = renew_login(cfg, cache, credentials).await?;
79    }
80    fetch_snapshot_with(
81        http_client()?,
82        &credentials,
83        cache,
84        &Endpoints::default(),
85        cache_ttl,
86    )
87    .await
88}
89
90/// Let the `droid` CLI renew a login that is expiring, then read it again.
91/// Under the vendor lock, so two bars refreshing together spawn the CLI once;
92/// the login is re-read first because the other one may already have done it.
93async fn renew_login(
94    cfg: &crate::config::DroidConfig,
95    cache: &Cache,
96    credentials: Credentials,
97) -> Result<Credentials> {
98    cache.ensure_dir()?;
99    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
100    let current = super::resolve_credentials(cfg)?;
101    if !current.expires_within(renew::RENEW_MARGIN, chrono::Utc::now()) {
102        return Ok(current);
103    }
104    if renew::renew(cfg, cache).await {
105        // A renewal the CLI could not finish leaves the old login in place,
106        // which the expired path below answers like before.
107        return super::resolve_credentials(cfg).or(Ok(credentials));
108    }
109    Ok(current)
110}
111
112async fn fetch_snapshot_with(
113    client: &reqwest::Client,
114    credentials: &Credentials,
115    cache: &Cache,
116    endpoints: &Endpoints,
117    cache_ttl: Duration,
118) -> Result<FetchOutcome> {
119    cache.ensure_dir()?;
120    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
121
122    // A JWT whose `exp` already passed would 401: skip the wire call and run
123    // the same cache fallback a rejected token gets — the last good payload
124    // plus the re-login note when a cache holds one.
125    if credentials.expired {
126        let error = AppError::Credentials(creds::EXPIRED_LOGIN_MESSAGE.into());
127        cache.mark_stale();
128        let (code, message) = error_to_pair(&error);
129        let diagnostic = cache.write_last_error(code, &message);
130        return crate::outcome::fallback(cache, Some(diagnostic), error, |bytes| {
131            parse_cache_at(bytes, &credentials.fingerprint)
132        });
133    }
134
135    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
136        && let Ok(outcome) = reuse_cache(&bytes, cache, &credentials.fingerprint)
137    {
138        return Ok(outcome);
139    }
140
141    match limits_call(client, credentials, endpoints).await {
142        Ok(snapshot) => {
143            cache.write_payload(&serde_json::to_vec(&snapshot_to_cache(
144                &snapshot,
145                &credentials.fingerprint,
146            ))?)?;
147            Ok(crate::outcome::Outcome::fresh(snapshot))
148        }
149        Err(error) if error.is_transient() => {
150            crate::outcome::fallback(cache, None, error, |bytes| {
151                parse_cache_at(bytes, &credentials.fingerprint)
152            })
153        }
154        Err(error) => {
155            cache.mark_stale();
156            let (code, message) = error_to_pair(&error);
157            let diagnostic = cache.write_last_error(code, &message);
158            crate::outcome::fallback(cache, Some(diagnostic), error, |bytes| {
159                parse_cache_at(bytes, &credentials.fingerprint)
160            })
161        }
162    }
163}
164
165async fn limits_call(
166    client: &reqwest::Client,
167    credentials: &Credentials,
168    endpoints: &Endpoints,
169) -> Result<DroidSnapshot> {
170    let response = client
171        .get(endpoints.limits_url())
172        .header(reqwest::header::ACCEPT, "application/json")
173        .bearer_auth(&credentials.access_token)
174        .send()
175        .await
176        .map_err(|error| {
177            AppError::Transport(if error.is_timeout() {
178                "Droid limits request timed out".into()
179            } else {
180                "Droid limits request failed".into()
181            })
182        })?;
183
184    let status = response.status();
185    let bytes = read_body_capped(response, MAX_BODY_BYTES).await?;
186    if status.as_u16() == 401 {
187        // A rejected token means the CLI's JWT lapsed between the exp check
188        // and the request. The `droid` CLI is what renews it — say so. Only
189        // 401 gets this: a 403 is an ordinary HTTP error, like elsewhere.
190        return Err(AppError::Credentials(creds::EXPIRED_LOGIN_MESSAGE.into()));
191    }
192    if !status.is_success() {
193        // The response body is never copied into the error: an error from a
194        // token-bearing endpoint may echo account detail.
195        return Err(AppError::Http {
196            status: status.as_u16(),
197            body: format!("Droid limits endpoint returned HTTP {}", status.as_u16()),
198        });
199    }
200    types::parse_response(&bytes)
201}
202
203/// Only an HTTP status or a schema message survives into `.last_error` and
204/// the tooltip. Everything else collapses to fixed text on purpose: those
205/// errors can carry a path or a request detail, and this is a token-bearing
206/// request. A credentials error keeps its own message — it is a fixed string
207/// that tells the user the one useful thing (run `droid` once).
208fn error_to_pair(error: &AppError) -> (u16, String) {
209    match error {
210        AppError::Http { status, body } => (*status, body.clone()),
211        AppError::Schema(message) => (0, message.clone()),
212        AppError::Credentials(message) => (0, message.clone()),
213        _ => (0, "Droid limits request failed; details suppressed".into()),
214    }
215}
216
217fn reuse_cache(bytes: &[u8], cache: &Cache, fingerprint: &str) -> Result<FetchOutcome> {
218    let snapshot = parse_cache_at(bytes, fingerprint)?;
219    Ok(crate::outcome::Outcome::cached(snapshot, cache, false))
220}
221
222fn snapshot_to_cache(snapshot: &DroidSnapshot, fingerprint: &str) -> serde_json::Value {
223    let window = |window: &Option<UsageWindow>| match window {
224        Some(window) => serde_json::json!({
225            "used_pct": window.utilization_pct,
226            "reset_unix": window.resets_at.map(|at| at.timestamp()),
227        }),
228        None => serde_json::Value::Null,
229    };
230    serde_json::json!({
231        "account": fingerprint,
232        "session": window(&snapshot.session),
233        "weekly": window(&snapshot.weekly),
234        "monthly": window(&snapshot.monthly),
235        "core_session": window(&snapshot.core_session),
236        "core_weekly": window(&snapshot.core_weekly),
237        "core_monthly": window(&snapshot.core_monthly),
238        "extra_balance_cents": snapshot.extra_balance_cents,
239    })
240}
241
242fn parse_cache_at(bytes: &[u8], fingerprint: &str) -> Result<DroidSnapshot> {
243    let value: serde_json::Value = serde_json::from_slice(bytes)
244        .map_err(|_| AppError::Schema("Droid cache was invalid JSON".into()))?;
245    if value.get("account").and_then(serde_json::Value::as_str) != Some(fingerprint) {
246        return Err(AppError::Schema(
247            "Droid cache belongs to a different login; refetching".into(),
248        ));
249    }
250    let invalid = |field: &str| AppError::Schema(format!("Droid cache: invalid {field}"));
251    let window = |field: &str, duration| -> Result<Option<UsageWindow>> {
252        let Some(value) = value.get(field) else {
253            return Err(invalid(field));
254        };
255        if value.is_null() {
256            return Ok(None);
257        }
258        let pct = value
259            .get("used_pct")
260            .and_then(serde_json::Value::as_i64)
261            .filter(|pct| (0..=100).contains(pct))
262            .ok_or_else(|| invalid(field))? as i32;
263        let reset = match value.get("reset_unix") {
264            None | Some(serde_json::Value::Null) => None,
265            Some(value) => {
266                let seconds = value
267                    .as_i64()
268                    .filter(|seconds| *seconds >= 0)
269                    .ok_or_else(|| invalid(field))?;
270                Some(chrono::DateTime::from_timestamp(seconds, 0).ok_or_else(|| invalid(field))?)
271            }
272        };
273        Ok(Some(UsageWindow {
274            utilization_pct: pct,
275            resets_at: reset,
276            window_duration: duration,
277        }))
278    };
279    let extra_balance_cents = match value.get("extra_balance_cents") {
280        Some(serde_json::Value::Null) => None,
281        Some(value) => Some(value.as_i64().ok_or_else(|| invalid("extra balance"))?),
282        None => return Err(invalid("extra balance")),
283    };
284    Ok(DroidSnapshot {
285        session: window("session", types::SESSION_WINDOW)?,
286        weekly: window("weekly", types::WEEKLY_WINDOW)?,
287        monthly: window("monthly", types::MONTHLY_WINDOW)?,
288        core_session: window("core_session", types::SESSION_WINDOW)?,
289        core_weekly: window("core_weekly", types::WEEKLY_WINDOW)?,
290        core_monthly: window("core_monthly", types::MONTHLY_WINDOW)?,
291        extra_balance_cents,
292    })
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298    use tempfile::TempDir;
299
300    const CANARY: &str = "DroidTokenCanary-748201";
301
302    fn credentials() -> Credentials {
303        Credentials {
304            access_token: CANARY.to_string(),
305            fingerprint: crate::cache::fingerprint_of(CANARY),
306            expires_at: None,
307            expired: false,
308        }
309    }
310
311    fn cache_fixture() -> (TempDir, Cache) {
312        let dir = TempDir::new().unwrap();
313        let cache = Cache::at(dir.path().join("droid"));
314        cache.ensure_dir().unwrap();
315        (dir, cache)
316    }
317
318    fn safe_client() -> reqwest::Client {
319        http_client().unwrap().clone()
320    }
321
322    fn response_body() -> String {
323        serde_json::json!({
324            "limits": {
325                "standard": {
326                    "fiveHour": {"usedPercent": 42, "windowEnd": "2026-10-22T01:38:45Z"},
327                    "weekly": {"usedPercent": 31, "windowEnd": "2026-10-28T20:38:45Z"},
328                    "monthly": {"usedPercent": 11, "windowEnd": null}
329                },
330                "core": {
331                    "fiveHour": {"usedPercent": 0, "windowEnd": null},
332                    "weekly": {"usedPercent": 0, "windowEnd": null},
333                    "monthly": {"usedPercent": 0, "windowEnd": null}
334                }
335            },
336            "extraUsageBalanceCents": 250
337        })
338        .to_string()
339    }
340
341    fn request_mock(server: &mut mockito::ServerGuard) -> mockito::Mock {
342        server
343            .mock("GET", LIMITS_PATH)
344            .match_header("authorization", format!("Bearer {CANARY}").as_str())
345            .match_header("accept", "application/json")
346    }
347
348    fn snapshot() -> DroidSnapshot {
349        let window = |pct, reset, duration| {
350            Some(UsageWindow {
351                utilization_pct: pct,
352                resets_at: reset,
353                window_duration: duration,
354            })
355        };
356        DroidSnapshot {
357            session: window(
358                42,
359                chrono::DateTime::from_timestamp(1_791_014_400, 0),
360                types::SESSION_WINDOW,
361            ),
362            weekly: window(31, None, types::WEEKLY_WINDOW),
363            monthly: None,
364            core_session: window(0, None, types::SESSION_WINDOW),
365            core_weekly: None,
366            core_monthly: None,
367            extra_balance_cents: Some(250),
368        }
369    }
370
371    #[tokio::test]
372    async fn request_sends_the_bearer_and_cache_keeps_only_normalized_values() {
373        let credentials = credentials();
374        let mut server = mockito::Server::new_async().await;
375        let request = request_mock(&mut server)
376            .with_status(200)
377            .with_body(response_body())
378            .create_async()
379            .await;
380        let (_cache_dir, cache) = cache_fixture();
381        let outcome = fetch_snapshot_with(
382            &safe_client(),
383            &credentials,
384            &cache,
385            &Endpoints { base: server.url() },
386            Duration::ZERO,
387        )
388        .await
389        .unwrap();
390        request.assert_async().await;
391        assert_eq!(
392            outcome.snapshot.session.as_ref().unwrap().utilization_pct,
393            42
394        );
395        assert_eq!(
396            outcome.snapshot.monthly.as_ref().unwrap().utilization_pct,
397            11
398        );
399        assert_eq!(outcome.snapshot.monthly.as_ref().unwrap().resets_at, None);
400        assert_eq!(outcome.snapshot.extra_balance_cents, Some(250));
401        let stored = std::fs::read_to_string(cache.payload_path()).unwrap();
402        assert!(!stored.contains(CANARY));
403        assert!(!stored.contains("usedPercent"));
404        assert_eq!(
405            parse_cache_at(stored.as_bytes(), &credentials.fingerprint).unwrap(),
406            outcome.snapshot
407        );
408    }
409
410    #[tokio::test]
411    async fn a_rejected_token_reports_the_droid_refresh_fix() {
412        let credentials = credentials();
413        let mut server = mockito::Server::new_async().await;
414        let request = request_mock(&mut server)
415            .with_status(401)
416            .with_body(format!("credential echo: {CANARY}"))
417            .create_async()
418            .await;
419        let (_cache_dir, cache) = cache_fixture();
420
421        let error = fetch_snapshot_with(
422            &safe_client(),
423            &credentials,
424            &cache,
425            &Endpoints { base: server.url() },
426            Duration::ZERO,
427        )
428        .await
429        .unwrap_err();
430
431        request.assert_async().await;
432        assert!(matches!(error, AppError::Credentials(_)), "{error:?}");
433        assert!(error.to_string().contains(creds::EXPIRED_LOGIN_MESSAGE));
434        assert!(!error.to_string().contains(CANARY));
435    }
436
437    #[tokio::test]
438    async fn a_forbidden_response_is_a_plain_http_error() {
439        let credentials = credentials();
440        let mut server = mockito::Server::new_async().await;
441        let request = request_mock(&mut server)
442            .with_status(403)
443            .with_body(format!("credential echo: {CANARY}"))
444            .create_async()
445            .await;
446        let (_cache_dir, cache) = cache_fixture();
447
448        let error = fetch_snapshot_with(
449            &safe_client(),
450            &credentials,
451            &cache,
452            &Endpoints { base: server.url() },
453            Duration::ZERO,
454        )
455        .await
456        .unwrap_err();
457
458        request.assert_async().await;
459        assert!(
460            matches!(error, AppError::Http { status: 403, .. }),
461            "{error:?}"
462        );
463        assert!(!error.to_string().contains(creds::EXPIRED_LOGIN_MESSAGE));
464        assert!(!error.to_string().contains(CANARY));
465    }
466
467    #[tokio::test]
468    async fn an_expired_jwt_with_a_cache_serves_the_last_good_payload() {
469        let mut credentials = credentials();
470        credentials.expired = true;
471        let mut server = mockito::Server::new_async().await;
472        // The wire is never touched: an expired token is not even sent.
473        let request = request_mock(&mut server).expect(0).create_async().await;
474        let (_cache_dir, cache) = cache_fixture();
475        cache
476            .write_payload(
477                &serde_json::to_vec(&snapshot_to_cache(&snapshot(), &credentials.fingerprint))
478                    .unwrap(),
479            )
480            .unwrap();
481
482        let outcome = fetch_snapshot_with(
483            &safe_client(),
484            &credentials,
485            &cache,
486            &Endpoints { base: server.url() },
487            Duration::ZERO,
488        )
489        .await
490        .unwrap();
491
492        request.assert_async().await;
493        assert!(outcome.stale);
494        assert_eq!(outcome.snapshot, snapshot());
495        let (_, message) = outcome.last_error.unwrap();
496        assert!(message.contains(creds::EXPIRED_LOGIN_MESSAGE));
497        assert!(!message.contains(CANARY));
498    }
499
500    #[tokio::test]
501    async fn an_expired_jwt_without_a_cache_reports_the_relogin_fix() {
502        let mut credentials = credentials();
503        credentials.expired = true;
504        let mut server = mockito::Server::new_async().await;
505        let request = request_mock(&mut server).expect(0).create_async().await;
506        let (_cache_dir, cache) = cache_fixture();
507
508        let error = fetch_snapshot_with(
509            &safe_client(),
510            &credentials,
511            &cache,
512            &Endpoints { base: server.url() },
513            Duration::ZERO,
514        )
515        .await
516        .unwrap_err();
517
518        request.assert_async().await;
519        assert!(matches!(error, AppError::Credentials(_)), "{error:?}");
520        assert!(error.to_string().contains(creds::EXPIRED_LOGIN_MESSAGE));
521    }
522
523    #[tokio::test]
524    async fn a_rejected_token_with_a_cache_serves_the_last_good_payload() {
525        let credentials = credentials();
526        let mut server = mockito::Server::new_async().await;
527        let request = request_mock(&mut server)
528            .with_status(401)
529            .with_body(format!("credential echo: {CANARY}"))
530            .create_async()
531            .await;
532        let (_cache_dir, cache) = cache_fixture();
533        cache
534            .write_payload(
535                &serde_json::to_vec(&snapshot_to_cache(&snapshot(), &credentials.fingerprint))
536                    .unwrap(),
537            )
538            .unwrap();
539
540        let outcome = fetch_snapshot_with(
541            &safe_client(),
542            &credentials,
543            &cache,
544            &Endpoints { base: server.url() },
545            Duration::ZERO,
546        )
547        .await
548        .unwrap();
549
550        request.assert_async().await;
551        assert!(outcome.stale);
552        assert_eq!(outcome.snapshot, snapshot());
553        let (_, message) = outcome.last_error.unwrap();
554        assert!(message.contains(creds::EXPIRED_LOGIN_MESSAGE));
555        assert!(!message.contains(CANARY));
556        for item in std::fs::read_dir(cache.dir()).unwrap() {
557            let item = item.unwrap();
558            if item.file_type().unwrap().is_file() {
559                let bytes = std::fs::read(item.path()).unwrap();
560                assert!(!String::from_utf8_lossy(&bytes).contains(CANARY));
561            }
562        }
563    }
564
565    #[tokio::test]
566    async fn redirects_are_not_followed_with_the_cli_token() {
567        let credentials = credentials();
568        let mut server = mockito::Server::new_async().await;
569        let request = request_mock(&mut server)
570            .with_status(307)
571            .with_header("location", "https://attacker.invalid/collect")
572            .create_async()
573            .await;
574        let redirected = server
575            .mock("GET", "/collect")
576            .expect(0)
577            .create_async()
578            .await;
579        let (_cache_dir, cache) = cache_fixture();
580        let error = fetch_snapshot_with(
581            &safe_client(),
582            &credentials,
583            &cache,
584            &Endpoints { base: server.url() },
585            Duration::ZERO,
586        )
587        .await
588        .unwrap_err();
589        request.assert_async().await;
590        redirected.assert_async().await;
591        assert!(matches!(error, AppError::Http { status: 307, .. }));
592        assert!(!error.to_string().contains(CANARY));
593    }
594
595    #[tokio::test]
596    async fn a_cache_for_another_login_is_refetched_not_reused() {
597        let mut server = mockito::Server::new_async().await;
598        let request = request_mock(&mut server)
599            .with_status(200)
600            .with_body(response_body())
601            .create_async()
602            .await;
603        let (_cache_dir, cache) = cache_fixture();
604        let old = Credentials {
605            access_token: "old-account-token".to_string(),
606            fingerprint: crate::cache::fingerprint_of("old-account-token"),
607            expires_at: None,
608            expired: false,
609        };
610        let mut old_snapshot = snapshot();
611        old_snapshot.session.as_mut().unwrap().utilization_pct = 99;
612        cache
613            .write_payload(
614                &serde_json::to_vec(&snapshot_to_cache(&old_snapshot, &old.fingerprint)).unwrap(),
615            )
616            .unwrap();
617
618        let outcome = fetch_snapshot_with(
619            &safe_client(),
620            &credentials(),
621            &cache,
622            &Endpoints { base: server.url() },
623            Duration::from_secs(3600),
624        )
625        .await
626        .unwrap();
627
628        request.assert_async().await;
629        assert_eq!(
630            outcome.snapshot.session.as_ref().unwrap().utilization_pct,
631            42
632        );
633    }
634
635    #[tokio::test]
636    async fn a_429_arms_the_backoff_and_the_next_refresh_stays_off_the_wire() {
637        let credentials = credentials();
638        let mut server = mockito::Server::new_async().await;
639        let request = request_mock(&mut server)
640            .with_status(429)
641            .with_body(format!("credential echo: {CANARY}"))
642            .expect(1)
643            .create_async()
644            .await;
645        let (_cache_dir, cache) = cache_fixture();
646        let endpoints = Endpoints { base: server.url() };
647
648        let error = fetch_snapshot_with(
649            &safe_client(),
650            &credentials,
651            &cache,
652            &endpoints,
653            Duration::ZERO,
654        )
655        .await
656        .unwrap_err();
657        assert!(
658            matches!(error, AppError::Http { status: 429, .. }),
659            "{error:?}"
660        );
661        assert!(!error.to_string().contains(CANARY));
662        assert!(
663            cache
664                .backoff_remaining_at(std::time::SystemTime::now())
665                .is_some()
666        );
667
668        // Armed: the second refresh is refused before any request.
669        let error = fetch_snapshot_with(
670            &safe_client(),
671            &credentials,
672            &cache,
673            &endpoints,
674            Duration::ZERO,
675        )
676        .await
677        .unwrap_err();
678        assert!(
679            matches!(error, AppError::Http { status: 429, .. }),
680            "{error:?}"
681        );
682        request.assert_async().await;
683    }
684
685    #[tokio::test]
686    async fn a_429_with_a_cache_keeps_serving_the_last_good_payload() {
687        let credentials = credentials();
688        let mut server = mockito::Server::new_async().await;
689        let request = request_mock(&mut server)
690            .with_status(429)
691            .expect(1)
692            .create_async()
693            .await;
694        let (_cache_dir, cache) = cache_fixture();
695        cache
696            .write_payload(
697                &serde_json::to_vec(&snapshot_to_cache(&snapshot(), &credentials.fingerprint))
698                    .unwrap(),
699            )
700            .unwrap();
701        let endpoints = Endpoints { base: server.url() };
702
703        for _ in 0..2 {
704            let outcome = fetch_snapshot_with(
705                &safe_client(),
706                &credentials,
707                &cache,
708                &endpoints,
709                Duration::ZERO,
710            )
711            .await
712            .unwrap();
713            assert_eq!(outcome.snapshot, snapshot());
714        }
715        request.assert_async().await;
716    }
717
718    #[tokio::test]
719    async fn an_oversized_body_is_refused_without_echoing_it() {
720        let credentials = credentials();
721        let mut server = mockito::Server::new_async().await;
722        let request = request_mock(&mut server)
723            .with_status(200)
724            .with_body(format!("{CANARY}{}", "x".repeat(MAX_BODY_BYTES + 1)))
725            .create_async()
726            .await;
727        let (_cache_dir, cache) = cache_fixture();
728
729        let error = fetch_snapshot_with(
730            &safe_client(),
731            &credentials,
732            &cache,
733            &Endpoints { base: server.url() },
734            Duration::ZERO,
735        )
736        .await
737        .unwrap_err();
738
739        request.assert_async().await;
740        assert!(matches!(error, AppError::Schema(_)), "{error:?}");
741        assert!(!error.to_string().contains(CANARY));
742        assert!(!cache.payload_path().exists());
743    }
744
745    #[tokio::test]
746    async fn a_silent_server_times_out_and_falls_back_to_the_cache() {
747        // Accepts the connection and never answers.
748        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
749        let address = listener.local_addr().unwrap();
750        let _server = tokio::spawn(async move {
751            let mut held = Vec::new();
752            while let Ok((socket, _)) = listener.accept().await {
753                held.push(socket);
754            }
755        });
756        let client = build_client(Duration::from_millis(300)).unwrap();
757        let credentials = credentials();
758        let (_cache_dir, cache) = cache_fixture();
759        let endpoints = Endpoints {
760            base: format!("http://{address}"),
761        };
762
763        let error = fetch_snapshot_with(&client, &credentials, &cache, &endpoints, Duration::ZERO)
764            .await
765            .unwrap_err();
766        assert!(
767            matches!(&error, AppError::Transport(message) if message.contains("timed out")),
768            "{error:?}"
769        );
770
771        cache
772            .write_payload(
773                &serde_json::to_vec(&snapshot_to_cache(&snapshot(), &credentials.fingerprint))
774                    .unwrap(),
775            )
776            .unwrap();
777        let outcome =
778            fetch_snapshot_with(&client, &credentials, &cache, &endpoints, Duration::ZERO)
779                .await
780                .unwrap();
781        assert!(outcome.stale);
782        assert_eq!(outcome.snapshot, snapshot());
783    }
784
785    #[tokio::test]
786    async fn a_rotated_token_for_the_same_organization_reuses_the_cache() {
787        let mut server = mockito::Server::new_async().await;
788        let request = server
789            .mock("GET", LIMITS_PATH)
790            .expect(0)
791            .create_async()
792            .await;
793        let (_cache_dir, cache) = cache_fixture();
794        let scope = crate::cache::fingerprint_of("droid-org:org-1");
795        cache
796            .write_payload(&serde_json::to_vec(&snapshot_to_cache(&snapshot(), &scope)).unwrap())
797            .unwrap();
798        let rotated = Credentials {
799            access_token: "rotated-token".to_string(),
800            fingerprint: scope,
801            expires_at: None,
802            expired: false,
803        };
804
805        let outcome = fetch_snapshot_with(
806            &safe_client(),
807            &rotated,
808            &cache,
809            &Endpoints { base: server.url() },
810            Duration::from_secs(3600),
811        )
812        .await
813        .unwrap();
814
815        request.assert_async().await;
816        assert_eq!(outcome.snapshot, snapshot());
817    }
818}
ai_usagebar/droid/fetch.rs

ai_usagebar/droid/
fetch.rs
