ai_usagebar/deepseek/
fetch.rs

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

ai_usagebar/deepseek/
fetch.rs
