ai_usagebar/commandcode/
fetch.rs

1//! Fetch Command Code usage, with the same cache/lock discipline as every
2//! other vendor.
3//!
4//! The official CLI's `/usage` view is assembled from three calls: `whoami`
5//! names the org that scopes the rest, the credit ledger carries the balance
6//! and the rolling spend windows, and the subscription names the plan. Only
7//! the ledger is load-bearing — a failure anywhere else costs its own detail
8//! and nothing more.
9
10use std::fmt::Write as _;
11use std::time::Duration;
12
13use chrono::DateTime;
14use serde_json::Value;
15use sha2::{Digest, Sha256};
16
17use crate::cache::{Cache, acquire_lock_async};
18use crate::error::{AppError, Result};
19use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
20
21use super::types::{Snapshot, apply_subscription, parse_credits};
22
23pub const BASE_URL: &str = "https://api.commandcode.ai";
24
25pub const WHOAMI_PATH: &str = "/alpha/whoami";
26pub const CREDITS_PATH: &str = "/alpha/billing/credits";
27pub const SUBSCRIPTIONS_PATH: &str = "/alpha/billing/subscriptions";
28
29const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
30const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
31const SCHEMA_ERROR: &str = "Command Code usage response schema mismatch";
32
33#[derive(Debug, Clone)]
34pub struct Endpoints {
35    pub base: String,
36}
37
38impl Default for Endpoints {
39    fn default() -> Self {
40        Self {
41            base: BASE_URL.to_string(),
42        }
43    }
44}
45
46impl Endpoints {
47    fn url(&self, path: &str) -> String {
48        format!("{}{path}", self.base.trim_end_matches('/'))
49    }
50}
51
52/// This vendor's [`Outcome`](crate::outcome::Outcome) — the shared shape,
53/// specialised to its snapshot.
54pub type FetchOutcome = crate::outcome::Outcome<Snapshot>;
55
56pub async fn fetch_snapshot(
57    client: &reqwest::Client,
58    token: &str,
59    cache: &Cache,
60    endpoints: &Endpoints,
61    ttl: Duration,
62) -> Result<FetchOutcome> {
63    cache.ensure_dir()?;
64    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
65    let target = target_key(endpoints, token);
66
67    if let Some(bytes) = cache.fresh_payload(ttl)?
68        && let Ok(snapshot) = parse_cache(&bytes, &target)
69    {
70        return Ok(crate::outcome::Outcome::cached(snapshot, cache, false));
71    }
72
73    match fetch_live(client, token, endpoints).await {
74        Ok(snapshot) => {
75            let body = serde_json::to_vec(&serde_json::json!({
76                "target": target,
77                "response": snapshot_repr(&snapshot),
78            }))?;
79            cache.write_payload(&body)?;
80            Ok(crate::outcome::Outcome::fresh(snapshot))
81        }
82        Err(error @ AppError::Transport(_)) => fallback_or_error(cache, None, &target, error),
83        Err(AppError::Http { status, .. }) => {
84            let message = status_message(status).to_string();
85            cache.mark_stale();
86            cache.write_last_error(status, &message);
87            fallback_or_error(
88                cache,
89                Some((status, message.clone())),
90                &target,
91                AppError::Http {
92                    status,
93                    body: message,
94                },
95            )
96        }
97        Err(AppError::Schema(_)) => {
98            let message = SCHEMA_ERROR.to_string();
99            cache.mark_stale();
100            cache.write_last_error(0, &message);
101            fallback_or_error(
102                cache,
103                Some((0, message.clone())),
104                &target,
105                AppError::Schema(message),
106            )
107        }
108        Err(error) => fallback_or_error(cache, None, &target, error),
109    }
110}
111
112async fn fetch_live(
113    client: &reqwest::Client,
114    token: &str,
115    endpoints: &Endpoints,
116) -> Result<Snapshot> {
117    // whoami scopes the ledger to an org. A personal account has none, and a
118    // failure here only costs that scoping, so the error is not propagated.
119    let org_id = get_json(client, token, &endpoints.url(WHOAMI_PATH), &[])
120        .await
121        .ok()
122        .and_then(|value| value.get("org")?.get("id")?.as_str().map(str::to_string));
123    let scope: Vec<(&str, String)> = org_id.iter().map(|id| ("orgId", id.clone())).collect();
124
125    let credits = get_json(client, token, &endpoints.url(CREDITS_PATH), &scope).await?;
126    let mut snapshot =
127        parse_credits(&credits).map_err(|error| AppError::Schema(error.to_string()))?;
128
129    // The plan is presentation detail; without it the windows still render.
130    if let Ok(subscription) =
131        get_json(client, token, &endpoints.url(SUBSCRIPTIONS_PATH), &scope).await
132    {
133        apply_subscription(&mut snapshot, &subscription);
134    }
135
136    Ok(snapshot)
137}
138
139async fn get_json(
140    client: &reqwest::Client,
141    token: &str,
142    url: &str,
143    query: &[(&str, String)],
144) -> Result<Value> {
145    let response = tokio::time::timeout(
146        HTTP_TIMEOUT,
147        client
148            .get(url)
149            .bearer_auth(token)
150            .query(query)
151            .header(reqwest::header::ACCEPT, "application/json")
152            .send(),
153    )
154    .await
155    .map_err(|_| AppError::Transport("Command Code request timed out".to_string()))??;
156
157    let status = response.status();
158    let body = read_body_capped(response, MAX_BODY_BYTES).await?;
159    if !status.is_success() {
160        return Err(AppError::Http {
161            status: status.as_u16(),
162            body: status_message(status.as_u16()).to_string(),
163        });
164    }
165    serde_json::from_slice(&body).map_err(|error| AppError::Schema(error.to_string()))
166}
167
168/// Stable, non-secret identity for the endpoint and the account the token
169/// resolves to. Cache reuse must fail closed when either input changes.
170fn target_key(endpoints: &Endpoints, token: &str) -> String {
171    let digest = Sha256::digest(token.as_bytes());
172    let mut fingerprint = String::with_capacity(digest.len() * 2);
173    for byte in digest {
174        let _ = write!(fingerprint, "{byte:02x}");
175    }
176    format!("{}|key:{fingerprint}", endpoints.base)
177}
178
179fn status_message(status: u16) -> &'static str {
180    match status {
181        401 | 403 => crate::error::AUTH_FAILURE_MESSAGE,
182        429 => "Command Code rate limited the usage request",
183        500..=599 => "Command Code usage endpoint is unavailable",
184        _ => "Command Code usage request failed",
185    }
186}
187
188/// Cache the parsed snapshot rather than the raw bodies: the raw ledger is
189/// account data with no reason to sit on disk longer than it must.
190fn snapshot_repr(snapshot: &Snapshot) -> Value {
191    let window = |window: &Option<super::types::SpendWindow>| {
192        window.as_ref().map(|w| {
193            serde_json::json!({
194                "used": w.used,
195                "cap": w.cap,
196                "resetAt": w.resets_at.map(|at| at.timestamp_millis()),
197            })
198        })
199    };
200    serde_json::json!({
201        "plan": snapshot.plan,
202        "creditPool": snapshot.credit_pool,
203        "periodEnd": snapshot.period_end.map(|at| at.timestamp_millis()),
204        "credits": snapshot.credits.as_ref().map(|c| serde_json::json!({
205            "monthlyCredits": c.monthly,
206            "purchasedCredits": c.purchased,
207            "freeCredits": c.free,
208        })),
209        "windowLimits": {
210            "fiveHour": window(&snapshot.five_hour),
211            "weekly": window(&snapshot.weekly),
212        },
213    })
214}
215
216fn parse_cache(bytes: &[u8], target: &str) -> Result<Snapshot> {
217    let value: Value = serde_json::from_slice(bytes)
218        .map_err(|_| AppError::Schema("Command Code cache is invalid".into()))?;
219    if value.get("target").and_then(Value::as_str) != Some(target) {
220        return Err(AppError::Schema(
221            "Command Code cache belongs to a different account".into(),
222        ));
223    }
224    let response = value
225        .get("response")
226        .ok_or_else(|| AppError::Schema("Command Code cache is missing its response".into()))?;
227    let mut snapshot =
228        parse_credits(response).map_err(|error| AppError::Schema(error.to_string()))?;
229    snapshot.plan = response
230        .get("plan")
231        .and_then(Value::as_str)
232        .map(str::to_string);
233    snapshot.credit_pool = response.get("creditPool").and_then(Value::as_f64);
234    // Older caches predate the field; a missing entry simply clears it and
235    // the next live refresh restores it.
236    snapshot.period_end = response
237        .get("periodEnd")
238        .and_then(Value::as_i64)
239        .and_then(DateTime::from_timestamp_millis);
240    Ok(snapshot)
241}
242
243fn fallback_or_error(
244    cache: &Cache,
245    last_error: Option<(u16, String)>,
246    target: &str,
247    error: AppError,
248) -> Result<FetchOutcome> {
249    crate::outcome::fallback(cache, last_error, error, |body| parse_cache(body, target))
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use crate::cache::Cache;
256
257    const CREDITS_BODY: &str = include_str!("../../tests/fixtures/commandcode/credits.json");
258    const SUBSCRIPTION_BODY: &str =
259        include_str!("../../tests/fixtures/commandcode/subscriptions.json");
260
261    fn cache_in(dir: &std::path::Path) -> Cache {
262        Cache::at(dir.join("commandcode"))
263    }
264
265    fn endpoints(server: &mockito::Server) -> Endpoints {
266        Endpoints { base: server.url() }
267    }
268
269    #[tokio::test]
270    async fn walks_the_chain_and_scopes_calls_to_the_org() {
271        let mut server = mockito::Server::new_async().await;
272        let whoami = server
273            .mock("GET", WHOAMI_PATH)
274            .with_body(r#"{"org":{"id":"org-42"}}"#)
275            .create_async()
276            .await;
277        let credits = server
278            .mock("GET", CREDITS_PATH)
279            .match_query(mockito::Matcher::UrlEncoded(
280                "orgId".into(),
281                "org-42".into(),
282            ))
283            .with_body(CREDITS_BODY)
284            .create_async()
285            .await;
286        let subscription = server
287            .mock("GET", SUBSCRIPTIONS_PATH)
288            .match_query(mockito::Matcher::UrlEncoded(
289                "orgId".into(),
290                "org-42".into(),
291            ))
292            .with_body(SUBSCRIPTION_BODY)
293            .create_async()
294            .await;
295        let dir = tempfile::tempdir().unwrap();
296
297        let outcome = fetch_snapshot(
298            &reqwest::Client::new(),
299            "tok",
300            &cache_in(dir.path()),
301            &endpoints(&server),
302            Duration::from_secs(60),
303        )
304        .await
305        .expect("fetch must succeed");
306
307        whoami.assert_async().await;
308        credits.assert_async().await;
309        subscription.assert_async().await;
310        assert_eq!(outcome.snapshot.plan.as_deref(), Some("GOAT"));
311        assert_eq!(outcome.snapshot.five_hour.unwrap().pct(), 25);
312        assert!(!outcome.stale);
313    }
314
315    #[tokio::test]
316    async fn a_personal_account_without_an_org_still_gets_its_ledger() {
317        let mut server = mockito::Server::new_async().await;
318        server
319            .mock("GET", WHOAMI_PATH)
320            .with_body(r#"{"org":null}"#)
321            .create_async()
322            .await;
323        let credits = server
324            .mock("GET", CREDITS_PATH)
325            .match_query(mockito::Matcher::Missing)
326            .with_body(CREDITS_BODY)
327            .create_async()
328            .await;
329        server
330            .mock("GET", SUBSCRIPTIONS_PATH)
331            .with_body(SUBSCRIPTION_BODY)
332            .create_async()
333            .await;
334        let dir = tempfile::tempdir().unwrap();
335
336        let outcome = fetch_snapshot(
337            &reqwest::Client::new(),
338            "tok",
339            &cache_in(dir.path()),
340            &endpoints(&server),
341            Duration::from_secs(60),
342        )
343        .await
344        .expect("fetch must succeed without an org");
345
346        credits.assert_async().await;
347        assert!(outcome.snapshot.weekly.is_some());
348    }
349
350    #[tokio::test]
351    async fn a_failing_subscription_costs_the_plan_not_the_windows() {
352        let mut server = mockito::Server::new_async().await;
353        server
354            .mock("GET", WHOAMI_PATH)
355            .with_body("{}")
356            .create_async()
357            .await;
358        server
359            .mock("GET", CREDITS_PATH)
360            .with_body(CREDITS_BODY)
361            .create_async()
362            .await;
363        server
364            .mock("GET", SUBSCRIPTIONS_PATH)
365            .with_status(500)
366            .create_async()
367            .await;
368        let dir = tempfile::tempdir().unwrap();
369
370        let outcome = fetch_snapshot(
371            &reqwest::Client::new(),
372            "tok",
373            &cache_in(dir.path()),
374            &endpoints(&server),
375            Duration::from_secs(60),
376        )
377        .await
378        .expect("windows must survive a subscription failure");
379
380        assert!(outcome.snapshot.plan.is_none());
381        assert_eq!(outcome.snapshot.weekly.unwrap().pct(), 30);
382    }
383
384    #[tokio::test]
385    async fn a_failing_ledger_is_fatal_and_maps_401_to_the_auth_message() {
386        let mut server = mockito::Server::new_async().await;
387        server
388            .mock("GET", WHOAMI_PATH)
389            .with_body("{}")
390            .create_async()
391            .await;
392        server
393            .mock("GET", CREDITS_PATH)
394            .with_status(401)
395            .create_async()
396            .await;
397        let dir = tempfile::tempdir().unwrap();
398
399        let error = fetch_snapshot(
400            &reqwest::Client::new(),
401            "tok",
402            &cache_in(dir.path()),
403            &endpoints(&server),
404            Duration::from_secs(60),
405        )
406        .await
407        .expect_err("the ledger is load-bearing");
408
409        assert!(
410            error
411                .to_string()
412                .contains(crate::error::AUTH_FAILURE_MESSAGE),
413            "{error}"
414        );
415    }
416
417    #[tokio::test]
418    async fn a_fresh_cache_is_served_without_touching_the_network() {
419        let mut server = mockito::Server::new_async().await;
420        let whoami = server
421            .mock("GET", WHOAMI_PATH)
422            .expect(1)
423            .with_body("{}")
424            .create_async()
425            .await;
426        server
427            .mock("GET", CREDITS_PATH)
428            .expect(1)
429            .with_body(CREDITS_BODY)
430            .create_async()
431            .await;
432        server
433            .mock("GET", SUBSCRIPTIONS_PATH)
434            .expect(1)
435            .with_body(SUBSCRIPTION_BODY)
436            .create_async()
437            .await;
438        let dir = tempfile::tempdir().unwrap();
439        let cache = cache_in(dir.path());
440        let endpoints = endpoints(&server);
441        let client = reqwest::Client::new();
442
443        let first = fetch_snapshot(&client, "tok", &cache, &endpoints, Duration::from_secs(600))
444            .await
445            .unwrap();
446        let second = fetch_snapshot(&client, "tok", &cache, &endpoints, Duration::from_secs(600))
447            .await
448            .unwrap();
449
450        // Each endpoint was called exactly once across both fetches.
451        whoami.assert_async().await;
452        assert_eq!(first.snapshot, second.snapshot);
453        assert_eq!(second.snapshot.plan.as_deref(), Some("GOAT"));
454    }
455
456    #[tokio::test]
457    async fn a_cache_written_for_another_token_is_not_reused() {
458        let mut server = mockito::Server::new_async().await;
459        server
460            .mock("GET", WHOAMI_PATH)
461            .with_body("{}")
462            .expect_at_least(2)
463            .create_async()
464            .await;
465        server
466            .mock("GET", CREDITS_PATH)
467            .with_body(CREDITS_BODY)
468            .expect_at_least(2)
469            .create_async()
470            .await;
471        server
472            .mock("GET", SUBSCRIPTIONS_PATH)
473            .with_body(SUBSCRIPTION_BODY)
474            .expect_at_least(2)
475            .create_async()
476            .await;
477        let dir = tempfile::tempdir().unwrap();
478        let cache = cache_in(dir.path());
479        let endpoints = endpoints(&server);
480        let client = reqwest::Client::new();
481
482        fetch_snapshot(
483            &client,
484            "token-a",
485            &cache,
486            &endpoints,
487            Duration::from_secs(600),
488        )
489        .await
490        .unwrap();
491        // A different account must re-fetch rather than read the first's cache.
492        fetch_snapshot(
493            &client,
494            "token-b",
495            &cache,
496            &endpoints,
497            Duration::from_secs(600),
498        )
499        .await
500        .unwrap();
501    }
502
503    #[tokio::test]
504    async fn a_stale_cache_carries_the_widget_through_an_outage() {
505        let mut server = mockito::Server::new_async().await;
506        let whoami = server
507            .mock("GET", WHOAMI_PATH)
508            .with_body("{}")
509            .create_async()
510            .await;
511        let credits = server
512            .mock("GET", CREDITS_PATH)
513            .with_body(CREDITS_BODY)
514            .create_async()
515            .await;
516        let subscription = server
517            .mock("GET", SUBSCRIPTIONS_PATH)
518            .with_body(SUBSCRIPTION_BODY)
519            .create_async()
520            .await;
521        let dir = tempfile::tempdir().unwrap();
522        let cache = cache_in(dir.path());
523        let endpoints = endpoints(&server);
524        let client = reqwest::Client::new();
525
526        fetch_snapshot(&client, "tok", &cache, &endpoints, Duration::ZERO)
527            .await
528            .unwrap();
529
530        // The endpoint starts failing; the cached snapshot still renders.
531        whoami.remove_async().await;
532        credits.remove_async().await;
533        subscription.remove_async().await;
534        server
535            .mock("GET", CREDITS_PATH)
536            .with_status(503)
537            .create_async()
538            .await;
539        server
540            .mock("GET", WHOAMI_PATH)
541            .with_status(503)
542            .create_async()
543            .await;
544
545        let outcome = fetch_snapshot(&client, "tok", &cache, &endpoints, Duration::ZERO)
546            .await
547            .expect("a stale snapshot beats no snapshot");
548
549        assert!(outcome.stale);
550        assert_eq!(outcome.snapshot.plan.as_deref(), Some("GOAT"));
551        assert_eq!(outcome.last_error.unwrap().0, 503);
552    }
553
554    #[tokio::test]
555    async fn a_schema_mismatch_is_reported_as_drift_not_as_success() {
556        let mut server = mockito::Server::new_async().await;
557        server
558            .mock("GET", WHOAMI_PATH)
559            .with_body("{}")
560            .create_async()
561            .await;
562        server
563            .mock("GET", CREDITS_PATH)
564            .with_body(r#"{"unexpected":true}"#)
565            .create_async()
566            .await;
567        let dir = tempfile::tempdir().unwrap();
568
569        let error = fetch_snapshot(
570            &reqwest::Client::new(),
571            "tok",
572            &cache_in(dir.path()),
573            &endpoints(&server),
574            Duration::from_secs(60),
575        )
576        .await
577        .expect_err("an unrecognised ledger must not pass as usage");
578
579        assert!(error.to_string().contains("schema"), "{error}");
580    }
581}
ai_usagebar/commandcode/fetch.rs

ai_usagebar/commandcode/
fetch.rs
