ai_usagebar/openrouter/
fetch.rs

1//! OpenRouter fetch — combines `/api/v1/credits` and `/api/v1/key` under
2//! the shared cache + flock primitives.
3
4use std::time::Duration;
5
6use crate::cache::{Cache, MAX_STALE, acquire_lock_async};
7use crate::error::{AppError, Result};
8use crate::usage::OpenRouterSnapshot;
9
10use super::types::{CreditsData, KeyData, OrEnvelope, combine};
11
12pub const BASE_URL: &str = "https://openrouter.ai/api/v1";
13const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
14const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
15
16#[derive(Debug, Clone)]
17pub struct Endpoints {
18    pub credits: String,
19    pub key: String,
20}
21
22impl Default for Endpoints {
23    fn default() -> Self {
24        Self {
25            credits: format!("{BASE_URL}/credits"),
26            key: format!("{BASE_URL}/key"),
27        }
28    }
29}
30
31#[derive(Debug, Clone)]
32pub struct FetchOutcome {
33    pub snapshot: OpenRouterSnapshot,
34    pub stale: bool,
35    pub last_error: Option<(u16, String)>,
36    pub cache_age: Option<Duration>,
37}
38
39/// Cache-aware fetch. Mirrors `anthropic::fetch::fetch_snapshot` semantics:
40/// fresh cache short-circuits; on failure, fall back to cache + mark stale.
41pub async fn fetch_snapshot(
42    client: &reqwest::Client,
43    api_key: &str,
44    cache: &Cache,
45    endpoints: &Endpoints,
46    cache_ttl: Duration,
47) -> Result<FetchOutcome> {
48    cache.ensure_dir()?;
49    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
50
51    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
52        && let Ok(outcome) = reuse_cache(bytes, cache, false)
53    {
54        return Ok(outcome);
55    }
56    // Corrupt fresh cache: fall through to live fetch rather than return a
57    // fabricated zero-credit snapshot.
58
59    match fetch_live(client, endpoints, api_key).await {
60        Ok((credits, key)) => {
61            let snap = combine(credits, key);
62            // Serialize back to JSON for the cache.
63            let cache_repr = serde_json::json!({
64                "snapshot": serde_repr(&snap),
65            });
66            let bytes = serde_json::to_vec(&cache_repr)?;
67            cache.write_payload(&bytes)?;
68            Ok(FetchOutcome {
69                snapshot: snap,
70                stale: false,
71                last_error: None,
72                cache_age: Some(Duration::ZERO),
73            })
74        }
75        Err(e) if e.is_transient() => fallback_silent(cache),
76        Err(AppError::Http { status, body }) => {
77            cache.mark_stale();
78            cache.write_last_error(status, &body);
79            fallback_with_error(cache, Some((status, body)))
80        }
81        Err(e) => {
82            cache.mark_stale();
83            cache.write_last_error(0, &e.to_string());
84            fallback_with_error(cache, Some((0, e.to_string())))
85        }
86    }
87}
88
89fn fallback_silent(cache: &Cache) -> Result<FetchOutcome> {
90    let Some(bytes) = cache.fallback_payload(MAX_STALE)? else {
91        return Err(AppError::Transport(
92            "openrouter: no cache and network unreachable".into(),
93        ));
94    };
95    reuse_cache(bytes, cache, true)
96}
97
98fn fallback_with_error(cache: &Cache, last_error: Option<(u16, String)>) -> Result<FetchOutcome> {
99    let Some(bytes) = cache.fallback_payload(MAX_STALE)? else {
100        return Err(AppError::Other("openrouter: no usable cache".into()));
101    };
102    let mut outcome = reuse_cache(bytes, cache, true)?;
103    outcome.last_error = last_error;
104    Ok(outcome)
105}
106
107fn reuse_cache(bytes: Vec<u8>, cache: &Cache, stale: bool) -> Result<FetchOutcome> {
108    let snap = parse_cache(&bytes)?;
109    Ok(FetchOutcome {
110        snapshot: snap,
111        stale,
112        last_error: cache.read_last_error(),
113        cache_age: cache.payload_age(),
114    })
115}
116
117/// Cached money is required, not optional: a truncated or half-written payload
118/// must be refetched rather than rendered as $0.00 with a free-tier badge.
119/// `limit`/`limit_remaining` stay optional — the API itself returns them null.
120fn parse_cache(bytes: &[u8]) -> Result<OpenRouterSnapshot> {
121    let v: serde_json::Value = serde_json::from_slice(bytes)?;
122    let s = v
123        .get("snapshot")
124        .ok_or_else(|| AppError::Schema("openrouter cache missing 'snapshot' field".into()))?;
125    let money = |name: &str| -> Result<f64> {
126        let n = s
127            .get(name)
128            .and_then(serde_json::Value::as_f64)
129            .ok_or_else(|| AppError::Schema(format!("openrouter cache missing '{name}'")))?;
130        if n.is_finite() && n >= 0.0 {
131            Ok(n)
132        } else {
133            Err(AppError::Schema(format!(
134                "openrouter cache '{name}' is not finite and non-negative"
135            )))
136        }
137    };
138    let optional_money = |name: &str, nonnegative: bool| -> Result<Option<f64>> {
139        match s.get(name) {
140            None | Some(serde_json::Value::Null) => Ok(None),
141            Some(value) => {
142                let number = value.as_f64().ok_or_else(|| {
143                    AppError::Schema(format!("openrouter cache '{name}' is not numeric or null"))
144                })?;
145                if number.is_finite() && (!nonnegative || number >= 0.0) {
146                    Ok(Some(number))
147                } else {
148                    Err(AppError::Schema(format!(
149                        "openrouter cache '{name}' is outside its valid range"
150                    )))
151                }
152            }
153        }
154    };
155    Ok(OpenRouterSnapshot {
156        label: s
157            .get("label")
158            .and_then(serde_json::Value::as_str)
159            .ok_or_else(|| AppError::Schema("openrouter cache missing 'label'".into()))?
160            .to_string(),
161        total_credits: money("total_credits")?,
162        total_usage: money("total_usage")?,
163        usage_daily: money("usage_daily")?,
164        usage_weekly: money("usage_weekly")?,
165        usage_monthly: money("usage_monthly")?,
166        is_free_tier: s["is_free_tier"]
167            .as_bool()
168            .ok_or_else(|| AppError::Schema("openrouter cache missing 'is_free_tier'".into()))?,
169        limit: optional_money("limit", true)?,
170        limit_remaining: optional_money("limit_remaining", false)?,
171    })
172}
173
174fn serde_repr(snap: &OpenRouterSnapshot) -> serde_json::Value {
175    serde_json::json!({
176        "label": snap.label,
177        "total_credits": snap.total_credits,
178        "total_usage": snap.total_usage,
179        "usage_daily": snap.usage_daily,
180        "usage_weekly": snap.usage_weekly,
181        "usage_monthly": snap.usage_monthly,
182        "is_free_tier": snap.is_free_tier,
183        "limit": snap.limit,
184        "limit_remaining": snap.limit_remaining,
185    })
186}
187
188async fn fetch_live(
189    client: &reqwest::Client,
190    endpoints: &Endpoints,
191    api_key: &str,
192) -> Result<(CreditsData, KeyData)> {
193    // Fetch in parallel.
194    let credits_fut = fetch_one::<CreditsData>(client, &endpoints.credits, api_key);
195    let key_fut = fetch_one::<KeyData>(client, &endpoints.key, api_key);
196    let (credits, key) = tokio::join!(credits_fut, key_fut);
197    Ok((credits?, key?))
198}
199
200async fn fetch_one<T: for<'de> serde::Deserialize<'de>>(
201    client: &reqwest::Client,
202    url: &str,
203    api_key: &str,
204) -> Result<T> {
205    let resp = tokio::time::timeout(
206        HTTP_TIMEOUT,
207        client
208            .get(url)
209            .header("Authorization", format!("Bearer {api_key}"))
210            .send(),
211    )
212    .await
213    .map_err(|_| AppError::Transport(format!("openrouter timeout: {url}")))??;
214
215    let status = resp.status();
216    let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
217
218    if !status.is_success() {
219        let body = String::from_utf8_lossy(&bytes).chars().take(200).collect();
220        return Err(AppError::Http {
221            status: status.as_u16(),
222            body,
223        });
224    }
225    let env: OrEnvelope<T> = serde_json::from_slice(&bytes)
226        .map_err(|e| AppError::Schema(format!("openrouter {url}: {e}")))?;
227    Ok(env.data)
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233    use tempfile::TempDir;
234
235    fn cache_fixture() -> (TempDir, Cache) {
236        let td = TempDir::new().unwrap();
237        let cache = Cache::at(td.path().join("openrouter"));
238        cache.ensure_dir().unwrap();
239        (td, cache)
240    }
241
242    #[tokio::test]
243    async fn live_fetch_combines_both_endpoints() {
244        let mut server = mockito::Server::new_async().await;
245        server
246            .mock("GET", "/api/v1/credits")
247            .with_status(200)
248            .with_body(r#"{"data":{"total_credits":100.0,"total_usage":25.5}}"#)
249            .create_async()
250            .await;
251        server
252            .mock("GET", "/api/v1/key")
253            .with_status(200)
254            .with_body(
255                r#"{"data":{"label":"prod","limit":50.0,"limit_remaining":24.5,
256                "usage":25.5,"usage_daily":1.0,"usage_weekly":7.0,"usage_monthly":25.5,
257                "is_free_tier":false}}"#,
258            )
259            .create_async()
260            .await;
261
262        let (_td, cache) = cache_fixture();
263        let client = reqwest::Client::new();
264        let endpoints = Endpoints {
265            credits: format!("{}/api/v1/credits", server.url()),
266            key: format!("{}/api/v1/key", server.url()),
267        };
268        let out = fetch_snapshot(
269            &client,
270            "sk-or-test",
271            &cache,
272            &endpoints,
273            Duration::from_secs(0),
274        )
275        .await
276        .unwrap();
277        assert_eq!(out.snapshot.total_credits, 100.0);
278        assert_eq!(out.snapshot.total_usage, 25.5);
279        assert!((out.snapshot.balance() - 74.5).abs() < 1e-9);
280        assert_eq!(out.snapshot.label, "OpenRouter — prod");
281        assert!(!out.stale);
282    }
283
284    #[tokio::test]
285    async fn http_error_falls_back_to_cache_when_present() {
286        let mut server = mockito::Server::new_async().await;
287        server
288            .mock("GET", "/api/v1/credits")
289            .with_status(401)
290            .with_body(r#"{"error":"unauthorized"}"#)
291            .create_async()
292            .await;
293        server
294            .mock("GET", "/api/v1/key")
295            .with_status(401)
296            .with_body(r#"{"error":"unauthorized"}"#)
297            .create_async()
298            .await;
299
300        let (_td, cache) = cache_fixture();
301        // Seed cache with a "snapshot" repr.
302        let seed = serde_json::json!({
303            "snapshot": {
304                "label":"OpenRouter — seed","total_credits": 50.0,
305                "total_usage": 10.0,"usage_daily":1.0,"usage_weekly":3.0,
306                "usage_monthly":10.0,"is_free_tier":false,
307                "limit":null,"limit_remaining":null
308            }
309        });
310        cache.write_payload(seed.to_string().as_bytes()).unwrap();
311
312        let client = reqwest::Client::new();
313        let endpoints = Endpoints {
314            credits: format!("{}/api/v1/credits", server.url()),
315            key: format!("{}/api/v1/key", server.url()),
316        };
317        let out = fetch_snapshot(&client, "k", &cache, &endpoints, Duration::from_secs(0))
318            .await
319            .unwrap();
320        assert!(out.stale);
321        assert_eq!(out.snapshot.label, "OpenRouter — seed");
322        assert_eq!(out.last_error.as_ref().map(|(c, _)| *c), Some(401));
323    }
324}
ai_usagebar/openrouter/fetch.rs

ai_usagebar/openrouter/
fetch.rs
