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

ai_usagebar/commandcode/
fetch.rs
