usage_monitor_cli/provider/
deepseek.rs

1use async_trait::async_trait;
2
3use crate::error::SpendPanelError;
4use crate::model::{CreditsSnapshot, UsageSnapshot};
5use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
6
7#[derive(Debug, serde::Deserialize)]
8struct DeepSeekBalanceResponse {
9    is_available: bool,
10    balance_infos: Vec<DeepSeekBalanceInfo>,
11}
12
13#[derive(Debug, serde::Deserialize)]
14struct DeepSeekBalanceInfo {
15    currency: String,
16    total_balance: String,
17    granted_balance: String,
18    topped_up_balance: String,
19}
20
21#[derive(Debug, Clone, PartialEq)]
22struct ParsedBalance {
23    is_available: bool,
24    currency: String,
25    total_balance: f64,
26    granted_balance: f64,
27    topped_up_balance: f64,
28}
29
30/// DeepSeek API balance provider.
31pub struct DeepSeekProvider {
32    metadata: ProviderMetadata,
33    /// Base URL override for tests.
34    base_url: Option<String>,
35}
36
37impl DeepSeekProvider {
38    pub fn new() -> Self {
39        Self {
40            metadata: ProviderMetadata {
41                id: "deepseek",
42                name: "DeepSeek",
43                description: "DeepSeek API balance monitor",
44                auth_methods: &["api_key", "env"],
45                website: Some("https://platform.deepseek.com"),
46            },
47            base_url: None,
48        }
49    }
50
51    /// Creates a provider with a custom base URL (for tests).
52    pub fn with_base_url(url: &str) -> Self {
53        let mut p = Self::new();
54        p.base_url = Some(url.to_string());
55        p
56    }
57
58    fn api_base<'a>(&'a self, ctx: &'a ProviderContext) -> &'a str {
59        ctx.config
60            .get("base_url")
61            .map(String::as_str)
62            .filter(|value| !value.is_empty())
63            .or(self.base_url.as_deref())
64            .unwrap_or("https://api.deepseek.com")
65    }
66
67    fn clean_key(raw: &str) -> String {
68        let mut value = raw.trim();
69        if value.len() >= 2
70            && ((value.starts_with('"') && value.ends_with('"'))
71                || (value.starts_with('\'') && value.ends_with('\'')))
72        {
73            value = &value[1..value.len() - 1];
74        }
75        value.trim().to_string()
76    }
77
78    fn detect_credentials_from(primary: Option<&str>, fallback: Option<&str>) -> bool {
79        primary
80            .or(fallback)
81            .map(Self::clean_key)
82            .is_some_and(|key| !key.is_empty())
83    }
84
85    fn resolve_api_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
86        for key in ["api_key", "token"] {
87            if let Some(value) = ctx.config.get(key) {
88                let cleaned = Self::clean_key(value);
89                if !cleaned.is_empty() {
90                    return Ok(cleaned);
91                }
92            }
93        }
94
95        for env in ["DEEPSEEK_API_KEY", "DEEPSEEK_KEY"] {
96            if let Ok(value) = std::env::var(env) {
97                let cleaned = Self::clean_key(&value);
98                if !cleaned.is_empty() {
99                    return Ok(cleaned);
100                }
101            }
102        }
103
104        Err(SpendPanelError::AuthFailed(
105            "deepseek".into(),
106            "no API key found in config, token, DEEPSEEK_API_KEY, or DEEPSEEK_KEY".into(),
107        ))
108    }
109
110    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
111        reqwest::Client::builder()
112            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
113            .build()
114            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
115    }
116
117    async fn fetch_balance(
118        base_url: &str,
119        client: &reqwest::Client,
120        api_key: &str,
121    ) -> Result<DeepSeekBalanceResponse, SpendPanelError> {
122        let resp = client
123            .get(format!("{}/user/balance", base_url.trim_end_matches('/')))
124            .header("Authorization", format!("Bearer {}", api_key))
125            .header("Accept", "application/json")
126            .send()
127            .await
128            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
129
130        let status = resp.status();
131        let body = resp
132            .text()
133            .await
134            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
135
136        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
137            return Err(SpendPanelError::AuthFailed(
138                "deepseek".into(),
139                format!("invalid API key (HTTP {})", status.as_u16()),
140            ));
141        }
142        if !status.is_success() {
143            return Err(SpendPanelError::ProviderError(
144                "deepseek".into(),
145                format!("HTTP {}: {}", status, body),
146            ));
147        }
148
149        serde_json::from_str(&body)
150            .map_err(|e| SpendPanelError::ParseError("deepseek".into(), e.to_string()))
151    }
152
153    fn parse_balance_info(info: &DeepSeekBalanceInfo) -> Result<ParsedBalance, SpendPanelError> {
154        let parse = |label: &str, value: &str| {
155            value.parse::<f64>().map_err(|_| {
156                SpendPanelError::ParseError(
157                    "deepseek".into(),
158                    format!("non-numeric {} balance value: {}", label, value),
159                )
160            })
161        };
162
163        Ok(ParsedBalance {
164            is_available: true,
165            currency: info.currency.clone(),
166            total_balance: parse("total", &info.total_balance)?,
167            granted_balance: parse("granted", &info.granted_balance)?,
168            topped_up_balance: parse("topped_up", &info.topped_up_balance)?,
169        })
170    }
171
172    fn select_balance(resp: DeepSeekBalanceResponse) -> Result<ParsedBalance, SpendPanelError> {
173        let mut balances = resp
174            .balance_infos
175            .iter()
176            .map(Self::parse_balance_info)
177            .collect::<Result<Vec<_>, _>>()?;
178
179        if balances.is_empty() {
180            return Ok(ParsedBalance {
181                is_available: false,
182                currency: "USD".into(),
183                total_balance: 0.0,
184                granted_balance: 0.0,
185                topped_up_balance: 0.0,
186            });
187        }
188
189        for balance in &mut balances {
190            balance.is_available = resp.is_available;
191        }
192
193        let selected = balances
194            .iter()
195            .find(|b| b.currency == "USD" && b.total_balance > 0.0)
196            .or_else(|| balances.iter().find(|b| b.total_balance > 0.0))
197            .or_else(|| balances.iter().find(|b| b.currency == "USD"))
198            .unwrap_or(&balances[0]);
199
200        Ok(selected.clone())
201    }
202
203    fn snapshot_from_balance(balance: ParsedBalance) -> UsageSnapshot {
204        let mut credits = CreditsSnapshot::new(balance.total_balance, balance.currency.clone());
205        credits.bonus = Some(balance.granted_balance);
206        credits.purchased = Some(balance.topped_up_balance);
207
208        let mut snapshot = UsageSnapshot::new("deepseek");
209        snapshot.credits = Some(credits);
210        if !balance.is_available || balance.total_balance <= 0.0 {
211            snapshot.primary_rate_window = Some(crate::model::RateWindow::new(1, 1, "Balance", 0));
212        } else {
213            snapshot.primary_rate_window = Some(crate::model::RateWindow::new(0, 1, "Balance", 0));
214        }
215        snapshot
216    }
217}
218
219impl Default for DeepSeekProvider {
220    fn default() -> Self {
221        Self::new()
222    }
223}
224
225#[async_trait]
226impl UsageProvider for DeepSeekProvider {
227    fn metadata(&self) -> &ProviderMetadata {
228        &self.metadata
229    }
230
231    fn detect_credentials(&self) -> bool {
232        Self::detect_credentials_from(
233            std::env::var("DEEPSEEK_API_KEY").ok().as_deref(),
234            std::env::var("DEEPSEEK_KEY").ok().as_deref(),
235        )
236    }
237
238    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
239        let api_key = Self::resolve_api_key(ctx)?;
240        let client = Self::build_client(ctx)?;
241        let response = Self::fetch_balance(self.api_base(ctx), &client, &api_key).await?;
242        let balance = Self::select_balance(response)?;
243        Ok(Self::snapshot_from_balance(balance))
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use pretty_assertions::assert_eq;
251    use wiremock::matchers::{header, method, path};
252    use wiremock::{Mock, MockServer, ResponseTemplate};
253
254    fn response(body: &str) -> DeepSeekBalanceResponse {
255        serde_json::from_str(body).unwrap()
256    }
257
258    #[test]
259    fn test_provider_metadata() {
260        let provider = DeepSeekProvider::new();
261        let meta = provider.metadata();
262        assert_eq!(meta.id, "deepseek");
263        assert_eq!(meta.name, "DeepSeek");
264        assert!(meta.auth_methods.contains(&"api_key"));
265    }
266
267    #[test]
268    fn test_clean_key_trims_and_unquotes() {
269        assert_eq!(DeepSeekProvider::clean_key("  sk-test  "), "sk-test");
270        assert_eq!(DeepSeekProvider::clean_key("\"sk-test\""), "sk-test");
271        assert_eq!(DeepSeekProvider::clean_key("'sk-test'"), "sk-test");
272    }
273
274    #[test]
275    fn test_resolve_api_key_from_context_api_key() {
276        let ctx = ProviderContext::with_api_key(" sk-test ");
277        assert_eq!(DeepSeekProvider::resolve_api_key(&ctx).unwrap(), "sk-test");
278    }
279
280    #[test]
281    fn test_resolve_api_key_from_context_token() {
282        let mut ctx = ProviderContext::new();
283        ctx.config.insert("token".into(), "sk-token".into());
284        assert_eq!(DeepSeekProvider::resolve_api_key(&ctx).unwrap(), "sk-token");
285    }
286
287    #[test]
288    fn test_resolve_api_key_missing_is_error() {
289        let err = DeepSeekProvider::resolve_api_key(&ProviderContext::new()).unwrap_err();
290        assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
291    }
292
293    #[test]
294    fn test_select_balance_prefers_funded_usd() {
295        let json = r#"{
296          "is_available": true,
297          "balance_infos": [
298            {"currency":"CNY","total_balance":"100.00","granted_balance":"0.00","topped_up_balance":"100.00"},
299            {"currency":"USD","total_balance":"20.00","granted_balance":"5.00","topped_up_balance":"15.00"}
300          ]
301        }"#;
302        let balance = DeepSeekProvider::select_balance(response(json)).unwrap();
303        assert_eq!(balance.currency, "USD");
304        assert_eq!(balance.total_balance, 20.0);
305        assert!(balance.is_available);
306    }
307
308    #[test]
309    fn test_select_balance_prefers_positive_cny_over_empty_usd() {
310        let json = r#"{
311          "is_available": true,
312          "balance_infos": [
313            {"currency":"USD","total_balance":"0.00","granted_balance":"0.00","topped_up_balance":"0.00"},
314            {"currency":"CNY","total_balance":"100.00","granted_balance":"0.00","topped_up_balance":"100.00"}
315          ]
316        }"#;
317        let balance = DeepSeekProvider::select_balance(response(json)).unwrap();
318        assert_eq!(balance.currency, "CNY");
319        assert_eq!(balance.total_balance, 100.0);
320    }
321
322    #[test]
323    fn test_select_balance_empty_returns_unavailable_usd_zero() {
324        let balance = DeepSeekProvider::select_balance(response(
325            r#"{"is_available":true,"balance_infos":[]}"#,
326        ))
327        .unwrap();
328        assert_eq!(balance.currency, "USD");
329        assert_eq!(balance.total_balance, 0.0);
330        assert!(!balance.is_available);
331    }
332
333    #[test]
334    fn test_select_balance_malformed_number_fails() {
335        let err = DeepSeekProvider::select_balance(response(
336            r#"{"is_available":true,"balance_infos":[{"currency":"USD","total_balance":"NaN?","granted_balance":"0.00","topped_up_balance":"0.00"}]}"#,
337        ))
338        .unwrap_err();
339        assert!(matches!(err, SpendPanelError::ParseError(_, _)));
340    }
341
342    #[test]
343    fn test_snapshot_contains_credit_breakdown() {
344        let snapshot = DeepSeekProvider::snapshot_from_balance(ParsedBalance {
345            is_available: true,
346            currency: "USD".into(),
347            total_balance: 50.0,
348            granted_balance: 10.0,
349            topped_up_balance: 40.0,
350        });
351        let credits = snapshot.credits.unwrap();
352        assert_eq!(credits.balance, 50.0);
353        assert_eq!(credits.currency, "USD");
354        assert_eq!(credits.bonus, Some(10.0));
355        assert_eq!(credits.purchased, Some(40.0));
356        assert_eq!(snapshot.primary_rate_window.unwrap().used, Some(0));
357    }
358
359    #[tokio::test]
360    async fn test_fetch_usage_success() {
361        let server = MockServer::start().await;
362        Mock::given(method("GET"))
363            .and(path("/user/balance"))
364            .and(header("authorization", "Bearer sk-test"))
365            .and(header("accept", "application/json"))
366            .respond_with(ResponseTemplate::new(200).set_body_raw(
367                r#"{"is_available":true,"balance_infos":[{"currency":"USD","total_balance":"50.00","granted_balance":"10.00","topped_up_balance":"40.00"}]}"#,
368                "application/json",
369            ))
370            .mount(&server)
371            .await;
372
373        let provider = DeepSeekProvider::with_base_url(&server.uri());
374        let snapshot = provider
375            .fetch_usage(&ProviderContext::with_api_key("sk-test"))
376            .await
377            .unwrap();
378        let credits = snapshot.credits.unwrap();
379        assert_eq!(credits.balance, 50.0);
380        assert_eq!(credits.bonus, Some(10.0));
381        assert_eq!(credits.purchased, Some(40.0));
382    }
383
384    #[tokio::test]
385    async fn test_fetch_usage_401_is_auth_failed() {
386        let server = MockServer::start().await;
387        Mock::given(method("GET"))
388            .and(path("/user/balance"))
389            .respond_with(ResponseTemplate::new(401))
390            .mount(&server)
391            .await;
392
393        let provider = DeepSeekProvider::with_base_url(&server.uri());
394        let err = provider
395            .fetch_usage(&ProviderContext::with_api_key("bad"))
396            .await
397            .unwrap_err();
398        assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
399    }
400}
usage_monitor_cli/provider/deepseek.rs

usage_monitor_cli/provider/
deepseek.rs
