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),
66        Err(AppError::Http { status, body }) => {
67            cache.mark_stale();
68            cache.write_last_error(status, &body);
69            fallback_with_error(cache, Some((status, body)))
70        }
71        Err(e) => {
72            cache.mark_stale();
73            cache.write_last_error(0, &e.to_string());
74            fallback_with_error(cache, Some((0, e.to_string())))
75        }
76    }
77}
78
79fn fallback_silent(cache: &Cache) -> Result<FetchOutcome> {
80    let Some(bytes) = cache.fallback_payload(MAX_STALE)? else {
81        return Err(AppError::Transport(
82            "deepseek: no cache and network unreachable".into(),
83        ));
84    };
85    reuse_cache(bytes, cache, true)
86}
87
88fn fallback_with_error(cache: &Cache, last_error: Option<(u16, String)>) -> Result<FetchOutcome> {
89    let Some(bytes) = cache.fallback_payload(MAX_STALE)? else {
90        return Err(AppError::Other("deepseek: no usable cache".into()));
91    };
92    let mut outcome = reuse_cache(bytes, cache, true)?;
93    outcome.last_error = last_error;
94    Ok(outcome)
95}
96
97fn reuse_cache(bytes: Vec<u8>, cache: &Cache, stale: bool) -> Result<FetchOutcome> {
98    let snap = parse_cache(&bytes)?;
99    Ok(FetchOutcome {
100        snapshot: snap,
101        stale,
102        last_error: cache.read_last_error(),
103        cache_age: cache.payload_age(),
104    })
105}
106
107/// Cached money is required, not optional: a truncated or half-written payload
108/// must be refetched rather than rendered as a $0.00 balance.
109fn parse_cache(bytes: &[u8]) -> Result<DeepseekSnapshot> {
110    let v: serde_json::Value = serde_json::from_slice(bytes)?;
111    let money = |name: &str| -> Result<f64> {
112        let n = v[name]
113            .as_f64()
114            .ok_or_else(|| AppError::Schema(format!("deepseek cache missing '{name}'")))?;
115        if n.is_finite() {
116            Ok(n)
117        } else {
118            Err(AppError::Schema(format!(
119                "deepseek cache '{name}' is not finite"
120            )))
121        }
122    };
123    let currency = v["currency"]
124        .as_str()
125        .ok_or_else(|| AppError::Schema("deepseek cache missing 'currency'".into()))?;
126    if !matches!(currency, "USD" | "CNY") {
127        return Err(AppError::Schema(format!(
128            "deepseek cache has unsupported currency {currency:?}"
129        )));
130    }
131    Ok(DeepseekSnapshot {
132        is_available: v["is_available"]
133            .as_bool()
134            .ok_or_else(|| AppError::Schema("deepseek cache missing 'is_available'".into()))?,
135        balance: money("balance")?,
136        granted: money("granted")?,
137        topped_up: money("topped_up")?,
138        currency: currency.to_string(),
139    })
140}
141
142fn snap_to_json(snap: &DeepseekSnapshot) -> serde_json::Value {
143    serde_json::json!({
144        "is_available": snap.is_available,
145        "balance": snap.balance,
146        "granted": snap.granted,
147        "topped_up": snap.topped_up,
148        "currency": snap.currency,
149    })
150}
151
152async fn fetch_live(
153    client: &reqwest::Client,
154    url: &str,
155    api_key: &str,
156) -> Result<DeepseekSnapshot> {
157    let resp = tokio::time::timeout(
158        HTTP_TIMEOUT,
159        client
160            .get(url)
161            .header("Authorization", format!("Bearer {api_key}"))
162            .header("Accept", "application/json")
163            .send(),
164    )
165    .await
166    .map_err(|_| AppError::Transport(format!("deepseek timeout: {url}")))??;
167
168    let status = resp.status();
169    let bytes = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
170
171    if !status.is_success() {
172        let body = String::from_utf8_lossy(&bytes).chars().take(200).collect();
173        return Err(AppError::Http {
174            status: status.as_u16(),
175            body,
176        });
177    }
178
179    let r: BalanceResponse = serde_json::from_slice(&bytes)
180        .map_err(|e| AppError::Schema(format!("deepseek balance response: {e}")))?;
181    r.into_snapshot()
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use tempfile::TempDir;
188
189    fn cache_fixture() -> (TempDir, Cache) {
190        let td = TempDir::new().unwrap();
191        let cache = Cache::at(td.path().join("deepseek"));
192        cache.ensure_dir().unwrap();
193        (td, cache)
194    }
195
196    #[test]
197    fn cached_unknown_currency_is_rejected_like_a_live_response() {
198        let cache = serde_json::json!({
199            "is_available": true,
200            "balance": 10.0,
201            "granted": 10.0,
202            "topped_up": 0.0,
203            "currency": "EUR"
204        });
205        let error = parse_cache(cache.to_string().as_bytes()).unwrap_err();
206        assert!(
207            error.to_string().contains("unsupported currency"),
208            "{error}"
209        );
210    }
211
212    #[tokio::test]
213    async fn live_200_returns_snapshot() {
214        let mut server = mockito::Server::new_async().await;
215        server
216            .mock("GET", "/user/balance")
217            .with_status(200)
218            .with_body(r#"{
219                "is_available": true,
220                "balance_infos": [
221                    {"currency": "USD", "total_balance": "5.00", "granted_balance": "5.00", "topped_up_balance": "0.00"}
222                ]
223            }"#)
224            .create_async()
225            .await;
226
227        let (_td, cache) = cache_fixture();
228        let client = reqwest::Client::new();
229        let endpoints = Endpoints {
230            balance: format!("{}/user/balance", server.url()),
231        };
232        let out = fetch_snapshot(
233            &client,
234            "sk-test",
235            &cache,
236            &endpoints,
237            Duration::from_secs(0),
238        )
239        .await
240        .unwrap();
241        assert!(out.snapshot.is_available);
242        assert!((out.snapshot.balance - 5.0).abs() < 1e-9);
243        assert_eq!(out.snapshot.currency, "USD");
244        assert!(!out.stale);
245    }
246
247    #[tokio::test]
248    async fn http_401_falls_back_to_cache() {
249        let mut server = mockito::Server::new_async().await;
250        server
251            .mock("GET", "/user/balance")
252            .with_status(401)
253            .with_body(r#"{"error": "invalid api key"}"#)
254            .create_async()
255            .await;
256
257        let (_td, cache) = cache_fixture();
258        let seed = serde_json::json!({
259            "is_available": true,
260            "balance": 3.0,
261            "granted": 3.0,
262            "topped_up": 0.0,
263            "currency": "USD"
264        });
265        cache.write_payload(seed.to_string().as_bytes()).unwrap();
266
267        let client = reqwest::Client::new();
268        let endpoints = Endpoints {
269            balance: format!("{}/user/balance", server.url()),
270        };
271        let out = fetch_snapshot(
272            &client,
273            "bad-key",
274            &cache,
275            &endpoints,
276            Duration::from_secs(0),
277        )
278        .await
279        .unwrap();
280        assert!(out.stale);
281        assert!((out.snapshot.balance - 3.0).abs() < 1e-9);
282        assert_eq!(out.last_error.as_ref().map(|(c, _)| *c), Some(401));
283    }
284}
ai_usagebar/deepseek/fetch.rs

ai_usagebar/deepseek/
fetch.rs
