ai_usagebar/deepinfra/
fetch.rs

1//! DeepInfra fetch for current prepaid balance and monthly usage.
2
3use std::time::Duration;
4
5use crate::cache::{Cache, acquire_lock_async};
6use crate::error::{AppError, Result};
7use crate::usage::DeepInfraSnapshot;
8
9use super::types::{Checklist, UsageResponse, combine};
10
11pub const BASE_URL: &str = "https://api.deepinfra.com";
12const HTTP_TIMEOUT: Duration = Duration::from_secs(20);
13const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
14
15#[derive(Debug, Clone)]
16pub struct Endpoints {
17    pub checklist: String,
18    pub usage: String,
19}
20
21impl Default for Endpoints {
22    fn default() -> Self {
23        Self {
24            checklist: format!("{BASE_URL}/payment/checklist"),
25            usage: format!("{BASE_URL}/payment/usage?from=current"),
26        }
27    }
28}
29
30pub type FetchOutcome = crate::outcome::Outcome<DeepInfraSnapshot>;
31
32pub async fn fetch_snapshot(
33    client: &reqwest::Client,
34    api_key: &str,
35    cache: &Cache,
36    endpoints: &Endpoints,
37    cache_ttl: Duration,
38) -> Result<FetchOutcome> {
39    cache.ensure_dir()?;
40    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
41
42    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
43        && let Ok(outcome) = reuse_cache(bytes, cache, false)
44    {
45        return Ok(outcome);
46    }
47
48    match fetch_live(client, endpoints, api_key)
49        .await
50        .and_then(|(checklist, usage)| combine(checklist, usage))
51    {
52        Ok(snapshot) => {
53            let bytes = serde_json::to_vec(&serde_json::json!({
54                "snapshot": serde_repr(&snapshot),
55            }))?;
56            cache.write_payload(&bytes)?;
57            Ok(crate::outcome::Outcome::fresh(snapshot))
58        }
59        Err(error) if error.is_transient() => fallback_silent(cache, error),
60        Err(AppError::Http { status, body }) => {
61            cache.mark_stale();
62            let last_error = Some(cache.write_last_error(status, &body));
63            fallback_with_error(cache, last_error, AppError::Http { status, body })
64        }
65        Err(error) => {
66            cache.mark_stale();
67            let last_error = Some(cache.write_last_error(0, &error.to_string()));
68            fallback_with_error(cache, last_error, error)
69        }
70    }
71}
72
73fn fallback_silent(cache: &Cache, original: AppError) -> Result<FetchOutcome> {
74    crate::outcome::fallback(cache, None, original, parse_cache)
75}
76
77fn fallback_with_error(
78    cache: &Cache,
79    last_error: Option<(u16, String)>,
80    original: AppError,
81) -> Result<FetchOutcome> {
82    crate::outcome::fallback(cache, last_error, original, parse_cache)
83}
84
85fn reuse_cache(bytes: Vec<u8>, cache: &Cache, stale: bool) -> Result<FetchOutcome> {
86    let snapshot = parse_cache(&bytes)?;
87    Ok(crate::outcome::Outcome::cached(snapshot, cache, stale))
88}
89
90fn parse_cache(bytes: &[u8]) -> Result<DeepInfraSnapshot> {
91    let value: serde_json::Value = serde_json::from_slice(bytes)?;
92    let snapshot = value
93        .get("snapshot")
94        .ok_or_else(|| AppError::Schema("deepinfra cache missing 'snapshot' field".into()))?;
95    let money = |name: &str, nonnegative: bool| -> Result<f64> {
96        let number = snapshot
97            .get(name)
98            .and_then(serde_json::Value::as_f64)
99            .ok_or_else(|| AppError::Schema(format!("deepinfra cache missing '{name}'")))?;
100        if number.is_finite() && (!nonnegative || number >= 0.0) {
101            Ok(number)
102        } else {
103            Err(AppError::Schema(format!(
104                "deepinfra cache '{name}' is outside its valid range"
105            )))
106        }
107    };
108    let monthly_limit = match snapshot.get("monthly_limit") {
109        None | Some(serde_json::Value::Null) => None,
110        Some(_) => Some(money("monthly_limit", true)?),
111    };
112    let period = snapshot
113        .get("period")
114        .and_then(serde_json::Value::as_str)
115        .filter(|period| !period.trim().is_empty())
116        .ok_or_else(|| AppError::Schema("deepinfra cache missing 'period'".into()))?;
117
118    Ok(DeepInfraSnapshot {
119        balance: money("balance", false)?,
120        monthly_spend: money("monthly_spend", true)?,
121        monthly_limit,
122        period: period.to_string(),
123    })
124}
125
126fn serde_repr(snapshot: &DeepInfraSnapshot) -> serde_json::Value {
127    serde_json::json!({
128        "balance": snapshot.balance,
129        "monthly_spend": snapshot.monthly_spend,
130        "monthly_limit": snapshot.monthly_limit,
131        "period": snapshot.period,
132    })
133}
134
135async fn fetch_live(
136    client: &reqwest::Client,
137    endpoints: &Endpoints,
138    api_key: &str,
139) -> Result<(Checklist, UsageResponse)> {
140    let checklist = fetch_one::<Checklist>(client, &endpoints.checklist, api_key);
141    let usage = fetch_one::<UsageResponse>(client, &endpoints.usage, api_key);
142    let (checklist, usage) = tokio::join!(checklist, usage);
143    Ok((checklist?, usage?))
144}
145
146async fn fetch_one<T: for<'de> serde::Deserialize<'de>>(
147    client: &reqwest::Client,
148    url: &str,
149    api_key: &str,
150) -> Result<T> {
151    let response = tokio::time::timeout(HTTP_TIMEOUT, client.get(url).bearer_auth(api_key).send())
152        .await
153        .map_err(|_| AppError::Transport(format!("deepinfra timeout: {url}")))??;
154    let status = response.status();
155    let bytes = crate::vendor::read_body_capped(response, crate::vendor::MAX_BODY_BYTES).await?;
156
157    if !status.is_success() {
158        let body = String::from_utf8_lossy(&bytes).chars().take(200).collect();
159        return Err(AppError::Http {
160            status: status.as_u16(),
161            body,
162        });
163    }
164
165    serde_json::from_slice(&bytes)
166        .map_err(|error| AppError::Schema(format!("deepinfra response: {error}")))
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172    use tempfile::TempDir;
173
174    fn cache_fixture() -> (TempDir, Cache) {
175        let directory = TempDir::new().unwrap();
176        let cache = Cache::at(directory.path().join("deepinfra"));
177        cache.ensure_dir().unwrap();
178        (directory, cache)
179    }
180
181    #[tokio::test]
182    async fn live_fetch_combines_balance_and_usage() {
183        let mut server = mockito::Server::new_async().await;
184        let checklist = server
185            .mock("GET", "/payment/checklist")
186            .match_header("authorization", "Bearer test-key")
187            .with_status(200)
188            .with_body(r#"{"stripe_balance":-5.0,"recent":2.08,"limit":null}"#)
189            .create_async()
190            .await;
191        let usage = server
192            .mock("GET", "/payment/usage?from=current")
193            .match_header("authorization", "Bearer test-key")
194            .with_status(200)
195            .with_body(r#"{"months":[{"period":"2026.09","total_cost":208}]}"#)
196            .create_async()
197            .await;
198
199        let (_directory, cache) = cache_fixture();
200        let endpoints = Endpoints {
201            checklist: format!("{}/payment/checklist", server.url()),
202            usage: format!("{}/payment/usage?from=current", server.url()),
203        };
204        let outcome = fetch_snapshot(
205            &reqwest::Client::new(),
206            "test-key",
207            &cache,
208            &endpoints,
209            Duration::ZERO,
210        )
211        .await
212        .unwrap();
213
214        checklist.assert_async().await;
215        usage.assert_async().await;
216        assert!((outcome.snapshot.balance - 2.92).abs() < 1e-9);
217        assert!((outcome.snapshot.monthly_spend - 2.08).abs() < 1e-9);
218        assert!(!outcome.stale);
219    }
220
221    #[tokio::test]
222    async fn an_http_error_falls_back_to_cache() {
223        let mut server = mockito::Server::new_async().await;
224        server
225            .mock("GET", "/payment/checklist")
226            .with_status(401)
227            .with_body(r#"{"detail":"unauthorized"}"#)
228            .create_async()
229            .await;
230
231        let (_directory, cache) = cache_fixture();
232        cache
233            .write_payload(
234                &serde_json::to_vec(&serde_json::json!({
235                    "snapshot": {
236                        "balance": 2.92,
237                        "monthly_spend": 2.08,
238                        "monthly_limit": null,
239                        "period": "2026.09"
240                    }
241                }))
242                .unwrap(),
243            )
244            .unwrap();
245        let endpoints = Endpoints {
246            checklist: format!("{}/payment/checklist", server.url()),
247            usage: format!("{}/payment/usage?from=current", server.url()),
248        };
249        let outcome = fetch_snapshot(
250            &reqwest::Client::new(),
251            "revoked-key",
252            &cache,
253            &endpoints,
254            Duration::ZERO,
255        )
256        .await
257        .unwrap();
258
259        assert!(outcome.stale);
260        assert_eq!(outcome.last_error.as_ref().map(|error| error.0), Some(401));
261        assert!((outcome.snapshot.balance - 2.92).abs() < 1e-9);
262    }
263
264    #[tokio::test]
265    async fn fresh_cache_skips_both_endpoints() {
266        let (_directory, cache) = cache_fixture();
267        cache
268            .write_payload(
269                &serde_json::to_vec(&serde_json::json!({
270                    "snapshot": {
271                        "balance": 2.92,
272                        "monthly_spend": 2.08,
273                        "monthly_limit": null,
274                        "period": "2026.09"
275                    }
276                }))
277                .unwrap(),
278            )
279            .unwrap();
280        let endpoints = Endpoints {
281            checklist: "http://127.0.0.1:1/payment/checklist".into(),
282            usage: "http://127.0.0.1:1/payment/usage?from=current".into(),
283        };
284
285        let outcome = fetch_snapshot(
286            &reqwest::Client::new(),
287            "test-key",
288            &cache,
289            &endpoints,
290            Duration::from_secs(60),
291        )
292        .await
293        .unwrap();
294
295        assert!(!outcome.stale);
296        assert!((outcome.snapshot.balance - 2.92).abs() < 1e-9);
297    }
298}
ai_usagebar/deepinfra/fetch.rs

ai_usagebar/deepinfra/
fetch.rs
