ai_usagebar/deeprelay/
fetch.rs

1use std::time::Duration;
2
3use serde::Deserialize;
4
5use crate::cache::{Cache, acquire_lock_async};
6use crate::error::{AppError, Result};
7use crate::usage::DeepRelaySnapshot;
8use crate::vendor::{MAX_BODY_BYTES, read_body_capped};
9
10use super::types::{Balance, CachedPayload, Subscription};
11
12pub const BASE_URL: &str = "https://api.deeprelay.ai/v1";
13const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
14const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
15
16pub type FetchOutcome = crate::outcome::Outcome<DeepRelaySnapshot>;
17
18#[derive(Debug, Clone)]
19pub struct Endpoints {
20    pub subscription: String,
21    pub balance: String,
22}
23
24impl Default for Endpoints {
25    fn default() -> Self {
26        Self {
27            subscription: format!("{BASE_URL}/billing/subscription"),
28            balance: format!("{BASE_URL}/billing/balance"),
29        }
30    }
31}
32
33pub async fn fetch_snapshot(
34    client: &reqwest::Client,
35    api_key: &str,
36    cache: &Cache,
37    endpoints: &Endpoints,
38    cache_ttl: Duration,
39) -> Result<FetchOutcome> {
40    cache.ensure_dir()?;
41    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
42    let parse_cache = |bytes: &[u8]| -> Result<DeepRelaySnapshot> {
43        let payload: CachedPayload = serde_json::from_slice(bytes)
44            .map_err(|_| AppError::Schema("DeepRelay billing cache is invalid".into()))?;
45        payload.snapshot()
46    };
47
48    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
49        && let Ok(snapshot) = parse_cache(&bytes)
50    {
51        return Ok(crate::outcome::Outcome::cached(snapshot, cache, false));
52    }
53    // An unparseable cache still must not bypass an active 429 backoff.
54    if let Some(remaining) = cache.backoff_remaining() {
55        return Err(AppError::Http {
56            status: 429,
57            body: format!(
58                "rate limited; next attempt in {}",
59                crate::cache::human_backoff(remaining)
60            ),
61        });
62    }
63
64    match request(client, api_key, endpoints).await {
65        Ok((payload, snapshot)) => {
66            cache.write_payload(&serde_json::to_vec(&payload)?)?;
67            Ok(crate::outcome::Outcome::fresh(snapshot))
68        }
69        Err(error) if error.is_transient() => {
70            crate::outcome::fallback(cache, None, error, parse_cache)
71        }
72        Err(error) => {
73            let code = match &error {
74                AppError::Http { status, .. } => *status,
75                _ => 0,
76            };
77            cache.mark_stale();
78            let diagnostic = cache.write_last_error(code, &error.to_string());
79            crate::outcome::fallback(cache, Some(diagnostic), error, parse_cache)
80        }
81    }
82}
83
84/// The wallet is only worth a second request when it can change what the user
85/// does: a pay-as-you-go account *is* its balance, and a subscriber only needs
86/// it once something has started billing past the plan.
87fn wants_balance(subscription: &Subscription) -> bool {
88    if !subscription.subscribed {
89        return true;
90    }
91    subscription
92        .usage
93        .as_ref()
94        .and_then(|usage| usage.recent_payg_calls_24h)
95        .is_some_and(|calls| calls > 0)
96}
97
98async fn request(
99    client: &reqwest::Client,
100    api_key: &str,
101    endpoints: &Endpoints,
102) -> Result<(CachedPayload, DeepRelaySnapshot)> {
103    let subscription: Subscription =
104        get_json(client, api_key, &endpoints.subscription, "subscription").await?;
105    // PAYG has no plan figures to fall back on: the wallet is its primary
106    // reading. Only a subscriber's supplementary wallet may fail independently.
107    let balance = if !subscription.subscribed {
108        Some(get_json::<Balance>(client, api_key, &endpoints.balance, "balance").await?)
109    } else if wants_balance(&subscription) {
110        get_json::<Balance>(client, api_key, &endpoints.balance, "balance")
111            .await
112            .ok()
113    } else {
114        None
115    };
116    let payload = CachedPayload {
117        subscription,
118        balance,
119    };
120    let snapshot = payload.snapshot()?;
121    Ok((payload, snapshot))
122}
123
124async fn get_json<T: serde::de::DeserializeOwned>(
125    client: &reqwest::Client,
126    api_key: &str,
127    url: &str,
128    what: &str,
129) -> Result<T> {
130    let response = tokio::time::timeout(
131        HTTP_TIMEOUT,
132        client
133            .get(url)
134            .bearer_auth(api_key)
135            .header(reqwest::header::ACCEPT, "application/json")
136            .send(),
137    )
138    .await
139    .map_err(|_| AppError::Transport(format!("DeepRelay {what} request timed out")))??;
140    let status = response.status();
141    let bytes = read_body_capped(response, MAX_BODY_BYTES).await?;
142    if !status.is_success() {
143        return Err(http_error(status.as_u16(), &bytes));
144    }
145    serde_json::from_slice(&bytes)
146        .map_err(|_| AppError::Schema(format!("DeepRelay {what} response schema mismatch")))
147}
148
149/// RFC 9457 problem body. `code` is the stable field — `detail` is prose and
150/// must be sanitized before it can reach a tooltip.
151#[derive(Deserialize)]
152struct Problem {
153    #[serde(default)]
154    detail: Option<String>,
155    #[serde(default)]
156    code: Option<String>,
157}
158
159fn http_error(status: u16, bytes: &[u8]) -> AppError {
160    let problem: Option<Problem> = serde_json::from_slice(bytes).ok();
161    let (code, detail) = match problem {
162        Some(Problem { code, detail }) => (code, detail),
163        None => (None, None),
164    };
165    let code = code.as_deref();
166    let detail =
167        detail.unwrap_or_else(|| String::from_utf8_lossy(bytes).chars().take(200).collect());
168    let detail = crate::display::sanitize_untrusted_line(&detail);
169    let body = match code {
170        // The only billing failure a user can act on: the key is valid but not
171        // allowed to read the plan.
172        Some("insufficient_scope") => format!(
173            "{detail} — this key needs the `billing:read` scope; create one with that scope at deeprelay.ai/cloud/api-keys"
174        ),
175        Some("unauthenticated" | "invalid_api_key") => {
176            format!("{detail} — check the key, or create a new one at deeprelay.ai/cloud/api-keys")
177        }
178        _ => detail,
179    };
180    AppError::Http { status, body }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186    use tempfile::TempDir;
187
188    fn cache() -> (TempDir, Cache) {
189        let dir = TempDir::new().unwrap();
190        let cache = Cache::at(dir.path().join("deeprelay"));
191        cache.ensure_dir().unwrap();
192        (dir, cache)
193    }
194
195    const SUBSCRIPTION: &str = include_str!("../../tests/fixtures/deeprelay/subscription.json");
196
197    /// The cache holds the subscription *and* the wallet, so a seeded cache is
198    /// the wrapped shape, not the bare response.
199    fn cached_body() -> Vec<u8> {
200        format!(r#"{{"subscription":{SUBSCRIPTION}}}"#).into_bytes()
201    }
202    const BALANCE: &str = r#"{"balance_cents":1234,"auto_pay_enabled":false,"auto_pay_threshold_cents":0,"auto_pay_amount_cents":0,"burn_cents_per_hour":0}"#;
203    const PAYG_SUBSCRIPTION: &str = r#"{"subscribed":false,"plan":null,"quota":null,"usage":null}"#;
204
205    fn cached_payg_body() -> Vec<u8> {
206        format!(r#"{{"subscription":{PAYG_SUBSCRIPTION},"balance":{BALANCE}}}"#).into_bytes()
207    }
208
209    fn endpoints(server: &mockito::Server) -> Endpoints {
210        Endpoints {
211            subscription: format!("{}/billing/subscription", server.url()),
212            balance: format!("{}/billing/balance", server.url()),
213        }
214    }
215
216    #[tokio::test]
217    async fn sends_bearer_key_and_parses_the_plan() {
218        let mut server = mockito::Server::new_async().await;
219        let subscription = server
220            .mock("GET", "/billing/subscription")
221            .match_header("authorization", "Bearer drl-test")
222            .match_header("accept", "application/json")
223            .with_status(200)
224            .with_body(SUBSCRIPTION)
225            .create_async()
226            .await;
227        // A subscriber with no wallet activity must not cost a second request.
228        let balance = server
229            .mock("GET", "/billing/balance")
230            .expect(0)
231            .create_async()
232            .await;
233        let (_dir, cache) = cache();
234
235        let outcome = fetch_snapshot(
236            &reqwest::Client::new(),
237            "drl-test",
238            &cache,
239            &endpoints(&server),
240            Duration::ZERO,
241        )
242        .await
243        .unwrap();
244
245        assert_eq!(outcome.snapshot.plan, "Plus");
246        assert_eq!(
247            outcome.snapshot.monthly_input.and_then(|m| m.pct()),
248            Some(32)
249        );
250        assert_eq!(outcome.snapshot.weekly_window().unwrap().resets_at, None);
251        assert_eq!(outcome.snapshot.balance_usd, None);
252        assert!(outcome.off_the_wire());
253        subscription.assert_async().await;
254        balance.assert_async().await;
255    }
256
257    #[tokio::test]
258    async fn reads_the_wallet_for_pay_as_you_go_accounts() {
259        let mut server = mockito::Server::new_async().await;
260        server
261            .mock("GET", "/billing/subscription")
262            .with_status(200)
263            .with_body(
264                r#"{"subscribed":false,"plan":null,"quota":null,
265                    "usage":null}"#,
266            )
267            .create_async()
268            .await;
269        server
270            .mock("GET", "/billing/balance")
271            .with_status(200)
272            .with_body(BALANCE)
273            .create_async()
274            .await;
275        let (_dir, cache) = cache();
276
277        let outcome = fetch_snapshot(
278            &reqwest::Client::new(),
279            "drl-test",
280            &cache,
281            &endpoints(&server),
282            Duration::ZERO,
283        )
284        .await
285        .unwrap();
286        assert_eq!(outcome.snapshot.plan, "Pay as you go");
287        assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
288    }
289
290    #[tokio::test]
291    async fn a_failed_balance_call_keeps_the_plan_figures() {
292        let mut server = mockito::Server::new_async().await;
293        let mut body: serde_json::Value = serde_json::from_str(SUBSCRIPTION).unwrap();
294        body["usage"]["recent_payg_calls_24h"] = 1.into();
295        let subscription = server
296            .mock("GET", "/billing/subscription")
297            .with_status(200)
298            .with_body(body.to_string())
299            .create_async()
300            .await;
301        let balance = server
302            .mock("GET", "/billing/balance")
303            .with_status(500)
304            .with_body("{}")
305            .create_async()
306            .await;
307        let (_dir, cache) = cache();
308
309        let outcome = fetch_snapshot(
310            &reqwest::Client::new(),
311            "drl-test",
312            &cache,
313            &endpoints(&server),
314            Duration::ZERO,
315        )
316        .await
317        .unwrap();
318        assert_eq!(outcome.snapshot.plan, "Plus");
319        assert_eq!(outcome.snapshot.balance_usd, None);
320        assert!(!outcome.stale);
321        subscription.assert_async().await;
322        balance.assert_async().await;
323    }
324
325    #[tokio::test]
326    async fn payg_balance_failures_preserve_cache_or_report_the_original_error() {
327        for warm_cache in [false, true] {
328            for (status, body, error_code) in [
329                (500, r#"{"detail":"wallet unavailable"}"#, 500),
330                (429, r#"{"detail":"rate limited"}"#, 429),
331                (200, r#"{"unexpected":true}"#, 0),
332            ] {
333                let mut server = mockito::Server::new_async().await;
334                let subscription = server
335                    .mock("GET", "/billing/subscription")
336                    .with_status(200)
337                    .with_body(PAYG_SUBSCRIPTION)
338                    .expect(1)
339                    .create_async()
340                    .await;
341                let balance = server
342                    .mock("GET", "/billing/balance")
343                    .with_status(status)
344                    .with_body(body)
345                    .expect(1)
346                    .create_async()
347                    .await;
348                let (_dir, cache) = cache();
349                let cached = cached_payg_body();
350                if warm_cache {
351                    cache.write_payload(&cached).unwrap();
352                }
353                let client = reqwest::Client::new();
354                let endpoints = endpoints(&server);
355
356                let result =
357                    fetch_snapshot(&client, "drl-test", &cache, &endpoints, Duration::ZERO).await;
358                if warm_cache {
359                    let outcome = result.unwrap();
360                    assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
361                    assert!(outcome.stale);
362                    assert!(!outcome.off_the_wire());
363                    assert_eq!(
364                        outcome.last_error.as_ref().map(|(code, _)| *code),
365                        Some(error_code)
366                    );
367                    assert_eq!(std::fs::read(cache.payload_path()).unwrap(), cached);
368                } else {
369                    let error = result.expect_err("PAYG requires a usable wallet balance");
370                    assert!(
371                        matches!(
372                            error,
373                            AppError::Http { status, .. } if status == error_code
374                        ) || matches!(error, AppError::Schema(_)) && error_code == 0
375                    );
376                    assert!(!cache.payload_path().exists());
377                }
378                assert!(cache.is_stale());
379                assert_eq!(
380                    cache.read_last_error().map(|(code, _)| code),
381                    Some(error_code)
382                );
383                if status == 429 {
384                    let deadline = std::fs::read(cache.retry_after_path()).unwrap();
385                    let retry = fetch_snapshot(
386                        &client,
387                        "drl-test",
388                        &cache,
389                        &endpoints,
390                        Duration::from_secs(60),
391                    )
392                    .await;
393                    if warm_cache {
394                        let outcome = retry.unwrap();
395                        assert!(outcome.stale);
396                        assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
397                    } else {
398                        assert!(matches!(retry, Err(AppError::Http { status: 429, .. })));
399                    }
400                    let wire_only =
401                        fetch_snapshot(&client, "drl-test", &cache, &endpoints, Duration::ZERO)
402                            .await;
403                    assert!(matches!(wire_only, Err(AppError::Http { status: 429, .. })));
404                    assert_eq!(std::fs::read(cache.retry_after_path()).unwrap(), deadline);
405                }
406                subscription.assert_async().await;
407                balance.assert_async().await;
408            }
409        }
410    }
411
412    #[tokio::test]
413    async fn payg_balance_transport_failure_keeps_the_previous_wallet_stale() {
414        let mut server = mockito::Server::new_async().await;
415        let subscription = server
416            .mock("GET", "/billing/subscription")
417            .with_status(200)
418            .with_body(PAYG_SUBSCRIPTION)
419            .create_async()
420            .await;
421        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
422        let unavailable = listener.local_addr().unwrap();
423        let reset_connection = tokio::spawn(async move {
424            let (socket, _) = listener.accept().await.unwrap();
425            drop(socket);
426        });
427        let mut endpoints = endpoints(&server);
428        endpoints.balance = format!("http://{unavailable}/billing/balance");
429        let (_dir, cache) = cache();
430        let cached = cached_payg_body();
431        cache.write_payload(&cached).unwrap();
432
433        let outcome = fetch_snapshot(
434            &reqwest::Client::new(),
435            "drl-test",
436            &cache,
437            &endpoints,
438            Duration::ZERO,
439        )
440        .await
441        .unwrap();
442
443        assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
444        assert!(outcome.stale);
445        assert!(outcome.last_error.is_none());
446        assert_eq!(std::fs::read(cache.payload_path()).unwrap(), cached);
447        subscription.assert_async().await;
448        reset_connection.await.unwrap();
449    }
450
451    // #589: key vendors serve the last good reading with an explicit auth
452    // message rather than a bare error — the cached figure is still what this
453    // account last used, and it is marked stale. Antigravity is the deliberate
454    // exception (a signed-out app's quota is not a reading at all).
455    #[tokio::test]
456    async fn a_payg_wallet_401_keeps_the_balance_stale_and_redacted() {
457        let mut server = mockito::Server::new_async().await;
458        server
459            .mock("GET", "/billing/subscription")
460            .with_status(200)
461            .with_body(PAYG_SUBSCRIPTION)
462            .create_async()
463            .await;
464        server
465            .mock("GET", "/billing/balance")
466            .with_status(401)
467            .with_body(r#"{"detail":"wallet key canary","code":"unauthenticated"}"#)
468            .create_async()
469            .await;
470        let (_dir, cache) = cache();
471        let cached = cached_payg_body();
472        cache.write_payload(&cached).unwrap();
473
474        let outcome = fetch_snapshot(
475            &reqwest::Client::new(),
476            "drl-test",
477            &cache,
478            &endpoints(&server),
479            Duration::ZERO,
480        )
481        .await
482        .unwrap();
483
484        assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
485        assert!(outcome.stale);
486        let (_, message) = outcome.last_error.expect("auth diagnostic");
487        assert_eq!(message, crate::error::AUTH_FAILURE_MESSAGE);
488        assert!(!message.contains("canary"));
489        assert_eq!(std::fs::read(cache.payload_path()).unwrap(), cached);
490    }
491
492    #[tokio::test]
493    async fn payg_cache_without_a_balance_is_refetched_unless_backed_off() {
494        for backed_off in [false, true] {
495            let mut server = mockito::Server::new_async().await;
496            let expected_requests = usize::from(!backed_off);
497            let subscription = server
498                .mock("GET", "/billing/subscription")
499                .with_status(200)
500                .with_body(PAYG_SUBSCRIPTION)
501                .expect(expected_requests)
502                .create_async()
503                .await;
504            let balance = server
505                .mock("GET", "/billing/balance")
506                .with_status(200)
507                .with_body(BALANCE)
508                .expect(expected_requests)
509                .create_async()
510                .await;
511            let (_dir, cache) = cache();
512            let incomplete = format!(r#"{{"subscription":{PAYG_SUBSCRIPTION}}}"#);
513            cache.write_payload(incomplete.as_bytes()).unwrap();
514            let deadline = backed_off.then(|| {
515                cache.write_last_error(429, "rate limited");
516                std::fs::read(cache.retry_after_path()).unwrap()
517            });
518
519            let result = fetch_snapshot(
520                &reqwest::Client::new(),
521                "drl-test",
522                &cache,
523                &endpoints(&server),
524                Duration::from_secs(60),
525            )
526            .await;
527
528            if let Some(deadline) = deadline {
529                assert!(matches!(result, Err(AppError::Http { status: 429, .. })));
530                assert_eq!(std::fs::read(cache.retry_after_path()).unwrap(), deadline);
531                assert_eq!(
532                    std::fs::read(cache.payload_path()).unwrap(),
533                    incomplete.as_bytes()
534                );
535            } else {
536                let outcome = result.unwrap();
537                assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
538                assert!(outcome.off_the_wire());
539                assert!(!outcome.stale);
540            }
541            subscription.assert_async().await;
542            balance.assert_async().await;
543        }
544    }
545
546    #[tokio::test]
547    async fn a_warm_cache_serves_the_previous_plan_without_a_request() {
548        let mut server = mockito::Server::new_async().await;
549        let request = server
550            .mock("GET", "/billing/subscription")
551            .expect(0)
552            .create_async()
553            .await;
554        let (_dir, cache) = cache();
555        cache.write_payload(&cached_body()).unwrap();
556
557        let outcome = fetch_snapshot(
558            &reqwest::Client::new(),
559            "drl-test",
560            &cache,
561            &endpoints(&server),
562            Duration::from_secs(60),
563        )
564        .await
565        .unwrap();
566        assert_eq!(outcome.snapshot.plan, "Plus");
567        assert!(!outcome.off_the_wire());
568        assert!(!outcome.stale);
569        request.assert_async().await;
570    }
571
572    #[tokio::test]
573    async fn an_insufficient_scope_key_says_which_scope_to_add() {
574        let mut server = mockito::Server::new_async().await;
575        server
576            .mock("GET", "/billing/subscription")
577            .with_status(403)
578            .with_body(
579                r#"{"type":"https://api.deeprelay.ai/errors/insufficient_scope",
580                    "title":"Forbidden","status":403,
581                    "detail":"This key is not allowed to read billing",
582                    "code":"insufficient_scope","request_id":"abc"}"#,
583            )
584            .create_async()
585            .await;
586        let (_dir, cache) = cache();
587
588        let error = fetch_snapshot(
589            &reqwest::Client::new(),
590            "drl-test",
591            &cache,
592            &endpoints(&server),
593            Duration::ZERO,
594        )
595        .await
596        .expect_err("scope failures have no cache to fall back on");
597        assert!(matches!(
598            &error,
599            AppError::Http { status: 403, body } if body.contains("billing:read")
600        ));
601    }
602
603    #[tokio::test]
604    async fn malformed_wallet_bodies_are_a_schema_mismatch() {
605        let mut server = mockito::Server::new_async().await;
606        server
607            .mock("GET", "/billing/subscription")
608            .with_status(200)
609            .with_body(r#"{"error":"unauthorized"}"#)
610            .create_async()
611            .await;
612        let (_dir, cache) = cache();
613
614        let error = fetch_snapshot(
615            &reqwest::Client::new(),
616            "drl-test",
617            &cache,
618            &endpoints(&server),
619            Duration::ZERO,
620        )
621        .await
622        .expect_err("a body that is not a subscription must not project one");
623        assert!(matches!(error, AppError::Schema(_)));
624    }
625
626    #[tokio::test]
627    async fn auth_failure_keeps_the_previous_plan_stale_and_redacted() {
628        let mut server = mockito::Server::new_async().await;
629        server
630            .mock("GET", "/billing/subscription")
631            .with_status(401)
632            .with_body(r#"{"detail":"bad key canary","code":"unauthenticated"}"#)
633            .create_async()
634            .await;
635        let (_dir, cache) = cache();
636        cache.write_payload(&cached_body()).unwrap();
637
638        let outcome = fetch_snapshot(
639            &reqwest::Client::new(),
640            "drl-test",
641            &cache,
642            &endpoints(&server),
643            Duration::ZERO,
644        )
645        .await
646        .unwrap();
647        assert!(outcome.stale);
648        let (_, message) = outcome.last_error.expect("auth diagnostic");
649        assert_eq!(message, crate::error::AUTH_FAILURE_MESSAGE);
650        assert!(!message.contains("canary"));
651    }
652
653    #[tokio::test]
654    async fn an_unparseable_cache_does_not_bypass_the_429_backoff() {
655        let mut server = mockito::Server::new_async().await;
656        let request = server
657            .mock("GET", "/billing/subscription")
658            .expect(0)
659            .create_async()
660            .await;
661        let (_dir, cache) = cache();
662        cache.write_payload(br#"{"unexpected":true}"#).unwrap();
663        cache.write_last_error(429, "rate limited");
664
665        let error = fetch_snapshot(
666            &reqwest::Client::new(),
667            "drl-test",
668            &cache,
669            &endpoints(&server),
670            Duration::ZERO,
671        )
672        .await
673        .expect_err("an invalid cache must not skip the backoff");
674        assert!(matches!(
675            error,
676            AppError::Http { status: 429, body } if body.contains("next attempt in")
677        ));
678        request.assert_async().await;
679    }
680}
ai_usagebar/deeprelay/fetch.rs

ai_usagebar/deeprelay/
fetch.rs
