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    #[tokio::test]
452    async fn payg_cache_without_a_balance_is_refetched_unless_backed_off() {
453        for backed_off in [false, true] {
454            let mut server = mockito::Server::new_async().await;
455            let expected_requests = usize::from(!backed_off);
456            let subscription = server
457                .mock("GET", "/billing/subscription")
458                .with_status(200)
459                .with_body(PAYG_SUBSCRIPTION)
460                .expect(expected_requests)
461                .create_async()
462                .await;
463            let balance = server
464                .mock("GET", "/billing/balance")
465                .with_status(200)
466                .with_body(BALANCE)
467                .expect(expected_requests)
468                .create_async()
469                .await;
470            let (_dir, cache) = cache();
471            let incomplete = format!(r#"{{"subscription":{PAYG_SUBSCRIPTION}}}"#);
472            cache.write_payload(incomplete.as_bytes()).unwrap();
473            let deadline = backed_off.then(|| {
474                cache.write_last_error(429, "rate limited");
475                std::fs::read(cache.retry_after_path()).unwrap()
476            });
477
478            let result = fetch_snapshot(
479                &reqwest::Client::new(),
480                "drl-test",
481                &cache,
482                &endpoints(&server),
483                Duration::from_secs(60),
484            )
485            .await;
486
487            if let Some(deadline) = deadline {
488                assert!(matches!(result, Err(AppError::Http { status: 429, .. })));
489                assert_eq!(std::fs::read(cache.retry_after_path()).unwrap(), deadline);
490                assert_eq!(
491                    std::fs::read(cache.payload_path()).unwrap(),
492                    incomplete.as_bytes()
493                );
494            } else {
495                let outcome = result.unwrap();
496                assert_eq!(outcome.snapshot.balance_usd, Some(12.34));
497                assert!(outcome.off_the_wire());
498                assert!(!outcome.stale);
499            }
500            subscription.assert_async().await;
501            balance.assert_async().await;
502        }
503    }
504
505    #[tokio::test]
506    async fn a_warm_cache_serves_the_previous_plan_without_a_request() {
507        let mut server = mockito::Server::new_async().await;
508        let request = server
509            .mock("GET", "/billing/subscription")
510            .expect(0)
511            .create_async()
512            .await;
513        let (_dir, cache) = cache();
514        cache.write_payload(&cached_body()).unwrap();
515
516        let outcome = fetch_snapshot(
517            &reqwest::Client::new(),
518            "drl-test",
519            &cache,
520            &endpoints(&server),
521            Duration::from_secs(60),
522        )
523        .await
524        .unwrap();
525        assert_eq!(outcome.snapshot.plan, "Plus");
526        assert!(!outcome.off_the_wire());
527        assert!(!outcome.stale);
528        request.assert_async().await;
529    }
530
531    #[tokio::test]
532    async fn an_insufficient_scope_key_says_which_scope_to_add() {
533        let mut server = mockito::Server::new_async().await;
534        server
535            .mock("GET", "/billing/subscription")
536            .with_status(403)
537            .with_body(
538                r#"{"type":"https://api.deeprelay.ai/errors/insufficient_scope",
539                    "title":"Forbidden","status":403,
540                    "detail":"This key is not allowed to read billing",
541                    "code":"insufficient_scope","request_id":"abc"}"#,
542            )
543            .create_async()
544            .await;
545        let (_dir, cache) = cache();
546
547        let error = fetch_snapshot(
548            &reqwest::Client::new(),
549            "drl-test",
550            &cache,
551            &endpoints(&server),
552            Duration::ZERO,
553        )
554        .await
555        .expect_err("scope failures have no cache to fall back on");
556        assert!(matches!(
557            &error,
558            AppError::Http { status: 403, body } if body.contains("billing:read")
559        ));
560    }
561
562    #[tokio::test]
563    async fn malformed_wallet_bodies_are_a_schema_mismatch() {
564        let mut server = mockito::Server::new_async().await;
565        server
566            .mock("GET", "/billing/subscription")
567            .with_status(200)
568            .with_body(r#"{"error":"unauthorized"}"#)
569            .create_async()
570            .await;
571        let (_dir, cache) = cache();
572
573        let error = fetch_snapshot(
574            &reqwest::Client::new(),
575            "drl-test",
576            &cache,
577            &endpoints(&server),
578            Duration::ZERO,
579        )
580        .await
581        .expect_err("a body that is not a subscription must not project one");
582        assert!(matches!(error, AppError::Schema(_)));
583    }
584
585    #[tokio::test]
586    async fn auth_failure_keeps_the_previous_plan_stale_and_redacted() {
587        let mut server = mockito::Server::new_async().await;
588        server
589            .mock("GET", "/billing/subscription")
590            .with_status(401)
591            .with_body(r#"{"detail":"bad key canary","code":"unauthenticated"}"#)
592            .create_async()
593            .await;
594        let (_dir, cache) = cache();
595        cache.write_payload(&cached_body()).unwrap();
596
597        let outcome = fetch_snapshot(
598            &reqwest::Client::new(),
599            "drl-test",
600            &cache,
601            &endpoints(&server),
602            Duration::ZERO,
603        )
604        .await
605        .unwrap();
606        assert!(outcome.stale);
607        let (_, message) = outcome.last_error.expect("auth diagnostic");
608        assert_eq!(message, crate::error::AUTH_FAILURE_MESSAGE);
609        assert!(!message.contains("canary"));
610    }
611
612    #[tokio::test]
613    async fn an_unparseable_cache_does_not_bypass_the_429_backoff() {
614        let mut server = mockito::Server::new_async().await;
615        let request = server
616            .mock("GET", "/billing/subscription")
617            .expect(0)
618            .create_async()
619            .await;
620        let (_dir, cache) = cache();
621        cache.write_payload(br#"{"unexpected":true}"#).unwrap();
622        cache.write_last_error(429, "rate limited");
623
624        let error = fetch_snapshot(
625            &reqwest::Client::new(),
626            "drl-test",
627            &cache,
628            &endpoints(&server),
629            Duration::ZERO,
630        )
631        .await
632        .expect_err("an invalid cache must not skip the backoff");
633        assert!(matches!(
634            error,
635            AppError::Http { status: 429, body } if body.contains("next attempt in")
636        ));
637        request.assert_async().await;
638    }
639}
ai_usagebar/deeprelay/fetch.rs

ai_usagebar/deeprelay/
fetch.rs
