ai_usagebar/deepseek/
fetch.rs

1//! Fetch DeepSeek usage from `/user/balance`.
2
3use std::time::Duration;
4
5use crate::cache::{Cache, MAX_STALE, acquire_lock_async};
6use crate::error::{AppError, Result};
7use crate::usage::DeepseekSnapshot;
8
9use super::types::BalanceResponse;
10
11pub const BASE_URL: &str = "https://api.deepseek.com";
12const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
13const LOCK_TIMEOUT: Duration = Duration::from_secs(15);
14
15#[derive(Debug, Clone)]
16pub struct Endpoints {
17    pub balance: String,
18}
19
20impl Default for Endpoints {
21    fn default() -> Self {
22        Self {
23            balance: format!("{BASE_URL}/user/balance"),
24        }
25    }
26}
27
28#[derive(Debug, Clone)]
29pub struct FetchOutcome {
30    pub snapshot: DeepseekSnapshot,
31    pub stale: bool,
32    pub last_error: Option<(u16, String)>,
33    pub cache_age: Option<Duration>,
34}
35
36pub async fn fetch_snapshot(
37    client: &reqwest::Client,
38    api_key: &str,
39    cache: &Cache,
40    endpoints: &Endpoints,
41    cache_ttl: Duration,
42) -> Result<FetchOutcome> {
43    cache.ensure_dir()?;
44    let _lock = acquire_lock_async(&cache.lock_path(), LOCK_TIMEOUT).await?;
45
46    if let Some(bytes) = cache.fresh_payload(cache_ttl)?
47        && let Ok(outcome) = reuse_cache(bytes, cache, false)
48    {
49        return Ok(outcome);
50    }
51    // Corrupt fresh cache: fall through to live fetch rather than return a
52    // fabricated zero balance.
53
54    match fetch_live(client, &endpoints.balance, api_key).await {
55        Ok(snap) => {
56            let bytes = serde_json::to_vec(&snap_to_json(&snap))?;
57            cache.write_payload(&bytes)?;
58            Ok(FetchOutcome {
59                snapshot: snap,
60                stale: false,
61                last_error: None,
62                cache_age: Some(Duration::ZERO),
63            })
64        }
65        Err(e) if e.is_transient() => fallback_silent(cache, e),
66        Err(AppError::Http { status, body }) => {
67            cache.mark_stale();
68            let last_error = Some(cache.write_last_error(status, &body));
69            fallback_with_error(cache, last_error, AppError::Http { status, body })
70        }
71        Err(e) => {
72            cache.mark_stale();
73            let last_error = Some(cache.write_last_error(0, &e.to_string()));
74            fallback_with_error(cache, last_error, e)
75        }
76    }
77}
78
79fn fallback_silent(cache: &Cache, original: AppError) -> Result<FetchOutcome> {
80    let Some(bytes) = cache.fallback_payload(MAX_STALE)? else {
81        return Err(original);
82    };
83    reuse_cache(bytes, cache, true)
84}
85
86/// On failure we show the last good figure with the error alongside it. With
87/// nothing usable cached there is nothing to show, so the **original** error is
88/// returned rather than a generic "no usable cache" that hides what went wrong
89/// — a cold cache and an expired key would otherwise look identical.
90fn fallback_with_error(
91    cache: &Cache,
92    last_error: Option<(u16, String)>,
93    original: AppError,
94) -> Result<FetchOutcome> {
95    let Some(bytes) = cache.fallback_payload(MAX_STALE)? else {
96        return Err(original);
97    };
98    let mut outcome = reuse_cache(bytes, cache, true)?;
99    outcome.last_error = last_error;
100    Ok(outcome)
101}
102
103fn reuse_cache(bytes: Vec<u8>, cache: &Cache, stale: bool) -> Result<FetchOutcome> {
104    let snap = parse_cache(&bytes)?;
105    Ok(FetchOutcome {
106        snapshot: snap,
107        stale,
108        last_error: cache.read_last_error(),
109        cache_age: cache.payload_age(),
110    })
111}
112
113/// Cached money is required, not optional: a truncated or half-written payload
114/// must be refetched rather than rendered as a $0.00 balance.
115fn parse_cache(bytes: &[u8]) -> Result<DeepseekSnapshot> {
116    let v: serde_json::Value = serde_json::from_slice(bytes)?;
117    let money = |name: &str| -> Result<f64> {
118        let n = v[name]
119            .as_f64()
120            .ok_or_else(|| AppError::Schema(format!("deepseek cache missing '{name}'")))?;
121        if n.is_finite() {
122            Ok(n)
123        } else {
124            Err(AppError::Schema(format!(
125                "deepseek cache '{name}' is not finite"
126            )))
127        }
128    };
129    let currency = v["currency"]
130        .as_str()
131        .ok_or_else(|| AppError::Schema("deepseek cache missing 'currency'".into()))?;
132    if !matches!(currency, "USD" | "CNY") {
133        return Err(AppError::Schema(format!(
134            "deepseek cache has unsupported currency {currency:?}"
135        )));
136    }
137    Ok(DeepseekSnapshot {
138        is_available: v["is_available"]
139            .as_bool()
140            .ok_or_else(|| AppError::Schema("deepseek cache missing 'is_available'".into()))?,
141        balance: money("balance")?,
142        granted: money("granted")?,
143        topped_up: money("topped_up")?,
144        currency: currency.to_string(),
145    })
146}
147
148fn snap_to_json(snap: &DeepseekSnapshot) -> serde_json::Value {
149    serde_json::json!({
150        "is_available": snap.is_available,
151        "balance": snap.balance,
152        "granted": snap.granted,
153        "topped_up": snap.topped_up,
154        "currency": snap.currency,
155    })
156}
157
158async fn fetch_live(
159    client: &reqwest::Client,
160    url: &str,
161    api_key: &str,
162) -> Result<DeepseekSnapshot> {
163    let resp = tokio::time::timeout(
164        HTTP_TIMEOUT,
165        client
166            .get(url)
167            .header("Authorization", format!("Bearer {api_key}"))
168            .header("Accept", "application/json")
169            .send(),
170    )
171    .await
172    .map_err(|_| AppError::Transport(format!("deepseek timeout: {url}")))??;
173
174    let status = resp.status();
175    let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
176
177    if !status.is_success() {
178        let body = String::from_utf8_lossy(&bytes).chars().take(200).collect();
179        return Err(AppError::Http {
180            status: status.as_u16(),
181            body,
182        });
183    }
184
185    let r: BalanceResponse = serde_json::from_slice(&bytes)
186        .map_err(|e| AppError::Schema(format!("deepseek balance response: {e}")))?;
187    r.into_snapshot()
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use tempfile::TempDir;
194
195    fn cache_fixture() -> (TempDir, Cache) {
196        let td = TempDir::new().unwrap();
197        let cache = Cache::at(td.path().join("deepseek"));
198        cache.ensure_dir().unwrap();
199        (td, cache)
200    }
201
202    #[test]
203    fn cached_unknown_currency_is_rejected_like_a_live_response() {
204        let cache = serde_json::json!({
205            "is_available": true,
206            "balance": 10.0,
207            "granted": 10.0,
208            "topped_up": 0.0,
209            "currency": "EUR"
210        });
211        let error = parse_cache(cache.to_string().as_bytes()).unwrap_err();
212        assert!(
213            error.to_string().contains("unsupported currency"),
214            "{error}"
215        );
216    }
217
218    #[tokio::test]
219    async fn live_200_returns_snapshot() {
220        let mut server = mockito::Server::new_async().await;
221        server
222            .mock("GET", "/user/balance")
223            .with_status(200)
224            .with_body(r#"{
225                "is_available": true,
226                "balance_infos": [
227                    {"currency": "USD", "total_balance": "5.00", "granted_balance": "5.00", "topped_up_balance": "0.00"}
228                ]
229            }"#)
230            .create_async()
231            .await;
232
233        let (_td, cache) = cache_fixture();
234        let client = reqwest::Client::new();
235        let endpoints = Endpoints {
236            balance: format!("{}/user/balance", server.url()),
237        };
238        let out = fetch_snapshot(
239            &client,
240            "sk-test",
241            &cache,
242            &endpoints,
243            Duration::from_secs(0),
244        )
245        .await
246        .unwrap();
247        assert!(out.snapshot.is_available);
248        assert!((out.snapshot.balance - 5.0).abs() < 1e-9);
249        assert_eq!(out.snapshot.currency, "USD");
250        assert!(!out.stale);
251    }
252
253    /// The whole bug in one path. `reuse_cache` already puts the *redacted*
254    /// message in the outcome — it reads it back from disk — and
255    /// `fallback_with_error` then overwrote it with a pair built from the raw
256    /// body. So the run that hit the `401` displayed the body and every run
257    /// after it displayed the neutral message.
258    ///
259    /// Written against Deepseek because it is the smallest harness that reaches
260    /// the shared path; the same two lines were repeated in five other vendors.
261    #[tokio::test]
262    async fn a_401_body_does_not_reach_the_outcome_when_a_cache_is_warm() {
263        let mut server = mockito::Server::new_async().await;
264        server
265            .mock("GET", "/user/balance")
266            .with_status(401)
267            .with_body("PANCEA user@example.test <credential>&token")
268            .create_async()
269            .await;
270
271        let (_td, cache) = cache_fixture();
272        let warm = serde_json::json!({
273            "is_available": true,
274            "balance": 5.0,
275            "granted": 5.0,
276            "topped_up": 0.0,
277            "currency": "USD"
278        });
279        cache.write_payload(warm.to_string().as_bytes()).unwrap();
280
281        let endpoints = Endpoints {
282            balance: format!("{}/user/balance", server.url()),
283        };
284        let out = fetch_snapshot(
285            &reqwest::Client::new(),
286            "sk-test",
287            &cache,
288            &endpoints,
289            Duration::from_secs(0),
290        )
291        .await
292        .unwrap();
293
294        let (code, msg) = out.last_error.expect("the 401 must still be reported");
295        assert_eq!(code, 401);
296        assert_eq!(msg, crate::error::AUTH_FAILURE_MESSAGE);
297        assert!(!msg.contains("PANCEA"), "{msg}");
298        assert!(!msg.contains("<credential>"), "{msg}");
299        assert!(out.stale);
300    }
301
302    #[tokio::test]
303    async fn http_401_falls_back_to_cache() {
304        let mut server = mockito::Server::new_async().await;
305        server
306            .mock("GET", "/user/balance")
307            .with_status(401)
308            .with_body(r#"{"error": "invalid api key"}"#)
309            .create_async()
310            .await;
311
312        let (_td, cache) = cache_fixture();
313        let seed = serde_json::json!({
314            "is_available": true,
315            "balance": 3.0,
316            "granted": 3.0,
317            "topped_up": 0.0,
318            "currency": "USD"
319        });
320        cache.write_payload(seed.to_string().as_bytes()).unwrap();
321
322        let client = reqwest::Client::new();
323        let endpoints = Endpoints {
324            balance: format!("{}/user/balance", server.url()),
325        };
326        let out = fetch_snapshot(
327            &client,
328            "bad-key",
329            &cache,
330            &endpoints,
331            Duration::from_secs(0),
332        )
333        .await
334        .unwrap();
335        assert!(out.stale);
336        assert!((out.snapshot.balance - 3.0).abs() < 1e-9);
337        assert_eq!(out.last_error.as_ref().map(|(c, _)| *c), Some(401));
338    }
339}
ai_usagebar/deepseek/fetch.rs

ai_usagebar/deepseek/
fetch.rs
