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

ai_usagebar/commandcode/
fetch.rs
