usage_monitor_cli/provider/
mistral.rs

1//! Mistral API spend provider (admin.mistral.ai, browser cookie auth).
2//!
3//! Ports CodexBar's read of `GET /api/billing/v2/usage`, aggregating the
4//! current month's metered cost across every usage category into a
5//! `CostSnapshot`.
6
7use async_trait::async_trait;
8use std::collections::HashMap;
9
10use crate::error::SpendPanelError;
11use crate::model::{CostSnapshot, UsageSnapshot};
12use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
13
14/// Mistral API spend provider.
15pub struct MistralProvider {
16    metadata: ProviderMetadata,
17    base_url: Option<String>,
18}
19
20impl MistralProvider {
21    pub fn new() -> Self {
22        Self {
23            metadata: ProviderMetadata {
24                id: "mistral",
25                name: "Mistral",
26                description: "Mistral API monthly spend monitor (browser cookie)",
27                auth_methods: &["cookie", "env"],
28                website: Some("https://mistral.ai"),
29            },
30            base_url: None,
31        }
32    }
33
34    pub fn with_base_url(url: &str) -> Self {
35        let mut p = Self::new();
36        p.base_url = Some(url.to_string());
37        p
38    }
39
40    fn api_base(&self) -> &str {
41        self.base_url
42            .as_deref()
43            .unwrap_or("https://admin.mistral.ai")
44    }
45
46    fn clean(raw: &str) -> String {
47        let mut v = raw.trim();
48        if v.len() >= 2
49            && ((v.starts_with('"') && v.ends_with('"'))
50                || (v.starts_with('\'') && v.ends_with('\'')))
51        {
52            v = &v[1..v.len() - 1];
53        }
54        v.trim().to_string()
55    }
56
57    fn resolve_cookie(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
58        for key in ["cookie", "token"] {
59            if let Some(v) = ctx.config.get(key) {
60                let c = Self::clean(v);
61                if !c.is_empty() {
62                    return Ok(c);
63                }
64            }
65        }
66        if let Ok(v) = std::env::var("MISTRAL_COOKIE") {
67            let c = Self::clean(&v);
68            if !c.is_empty() {
69                return Ok(c);
70            }
71        }
72        Err(SpendPanelError::AuthFailed(
73            "mistral".into(),
74            "no session cookie in cookie config or MISTRAL_COOKIE".into(),
75        ))
76    }
77
78    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
79        reqwest::Client::builder()
80            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
81            .build()
82            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
83    }
84
85    fn price_index(json: &serde_json::Value) -> HashMap<String, f64> {
86        let mut index = HashMap::new();
87        if let Some(prices) = json.get("prices").and_then(|p| p.as_array()) {
88            for price in prices {
89                let metric = price.get("billing_metric").and_then(|v| v.as_str());
90                let group = price.get("billing_group").and_then(|v| v.as_str());
91                let value = price
92                    .get("price")
93                    .and_then(|v| v.as_str())
94                    .and_then(|s| s.parse::<f64>().ok());
95                if let (Some(metric), Some(group), Some(value)) = (metric, group, value) {
96                    index.insert(format!("{}::{}", metric, group), value);
97                }
98            }
99        }
100        index
101    }
102
103    /// Recursively sums `input`/`output`/`cached` usage-entry arrays into a cost,
104    /// using the price index keyed by `metric::group`.
105    fn accumulate_cost(node: &serde_json::Value, prices: &HashMap<String, f64>, total: &mut f64) {
106        match node {
107            serde_json::Value::Object(map) => {
108                for (key, value) in map {
109                    let entries = value
110                        .as_array()
111                        .filter(|_| matches!(key.as_str(), "input" | "output" | "cached"));
112                    if let Some(entries) = entries {
113                        for entry in entries {
114                            let units = entry
115                                .get("value_paid")
116                                .and_then(|v| v.as_i64())
117                                .or_else(|| entry.get("value").and_then(|v| v.as_i64()))
118                                .unwrap_or(0);
119                            let metric = entry.get("billing_metric").and_then(|v| v.as_str());
120                            let group = entry.get("billing_group").and_then(|v| v.as_str());
121                            let price = match (metric, group) {
122                                (Some(m), Some(g)) => prices.get(&format!("{}::{}", m, g)),
123                                _ => None,
124                            };
125                            if let Some(price) = price {
126                                *total += units as f64 * price;
127                            }
128                        }
129                        continue;
130                    }
131                    Self::accumulate_cost(value, prices, total);
132                }
133            }
134            serde_json::Value::Array(items) => {
135                for item in items {
136                    Self::accumulate_cost(item, prices, total);
137                }
138            }
139            _ => {}
140        }
141    }
142
143    fn parse(body: &str) -> Result<UsageSnapshot, SpendPanelError> {
144        let json: serde_json::Value = serde_json::from_str(body)
145            .map_err(|e| SpendPanelError::ParseError("mistral".into(), e.to_string()))?;
146        let prices = Self::price_index(&json);
147        let mut total = 0.0;
148        // Sum each top-level usage category (skip the prices array itself).
149        if let Some(obj) = json.as_object() {
150            for (key, value) in obj {
151                if key == "prices" {
152                    continue;
153                }
154                Self::accumulate_cost(value, &prices, &mut total);
155            }
156        }
157        let currency = json
158            .get("currency")
159            .and_then(|v| v.as_str())
160            .unwrap_or("EUR")
161            .to_string();
162
163        let mut snapshot = UsageSnapshot::new("mistral");
164        snapshot.cost = Some(CostSnapshot {
165            total_cost: Some(total.max(0.0)),
166            currency,
167            daily_costs: Vec::new(),
168            spend_limit: None,
169        });
170        Ok(snapshot)
171    }
172}
173
174impl Default for MistralProvider {
175    fn default() -> Self {
176        Self::new()
177    }
178}
179
180#[async_trait]
181impl UsageProvider for MistralProvider {
182    fn metadata(&self) -> &ProviderMetadata {
183        &self.metadata
184    }
185
186    fn detect_credentials(&self) -> bool {
187        std::env::var("MISTRAL_COOKIE")
188            .map(|v| !v.trim().is_empty())
189            .unwrap_or(false)
190    }
191
192    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
193        let cookie = Self::resolve_cookie(ctx)?;
194        let client = Self::build_client(ctx)?;
195        let now = chrono::Utc::now();
196        let (month, year) = (chrono::Datelike::month(&now), chrono::Datelike::year(&now));
197        let url = format!(
198            "{}/api/billing/v2/usage?month={}&year={}",
199            self.api_base().trim_end_matches('/'),
200            month,
201            year
202        );
203        let mut req = client
204            .get(url)
205            .header("Accept", "*/*")
206            .header("Cookie", &cookie)
207            .header("Referer", "https://admin.mistral.ai/organization/usage")
208            .header("Origin", "https://admin.mistral.ai");
209        if let Some(csrf) = ctx.config.get("csrf_token").filter(|v| !v.is_empty()) {
210            req = req.header("X-CSRFTOKEN", csrf);
211        }
212        let resp = req
213            .send()
214            .await
215            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
216        let status = resp.status();
217        let body = resp
218            .text()
219            .await
220            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
221        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
222            return Err(SpendPanelError::AuthFailed(
223                "mistral".into(),
224                format!("session cookie rejected (HTTP {})", status.as_u16()),
225            ));
226        }
227        if !status.is_success() {
228            return Err(SpendPanelError::ProviderError(
229                "mistral".into(),
230                format!("HTTP {}: {}", status, body),
231            ));
232        }
233        Self::parse(&body)
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use pretty_assertions::assert_eq;
241    use wiremock::matchers::{method, path};
242    use wiremock::{Mock, MockServer, ResponseTemplate};
243
244    const SAMPLE: &str = r#"{
245      "currency": "USD",
246      "prices": [
247        {"billing_metric": "tokens_in", "billing_group": "mistral-large", "price": "0.001"},
248        {"billing_metric": "tokens_out", "billing_group": "mistral-large", "price": "0.003"}
249      ],
250      "completion": {
251        "models": {
252          "mistral-large": {
253            "input": [{"value": 1000, "billing_metric": "tokens_in", "billing_group": "mistral-large"}],
254            "output": [{"value_paid": 500, "billing_metric": "tokens_out", "billing_group": "mistral-large"}]
255          }
256        }
257      }
258    }"#;
259
260    #[test]
261    fn test_metadata() {
262        assert_eq!(MistralProvider::new().metadata().id, "mistral");
263    }
264
265    #[test]
266    fn test_parse_aggregates_cost() {
267        let snap = MistralProvider::parse(SAMPLE).unwrap();
268        let cost = snap.cost.unwrap();
269        // 1000*0.001 + 500*0.003 = 1.0 + 1.5 = 2.5
270        assert_eq!(cost.total_cost, Some(2.5));
271        assert_eq!(cost.currency, "USD");
272    }
273
274    #[test]
275    fn test_aggregates_across_categories() {
276        // Cost from both completion and ocr categories is summed.
277        let body = r#"{
278          "currency": "USD",
279          "prices": [
280            {"billing_metric": "t_in", "billing_group": "g", "price": "0.01"},
281            {"billing_metric": "pages", "billing_group": "ocr", "price": "0.5"}
282          ],
283          "completion": {"models": {"m": {
284            "input": [{"value": 100, "billing_metric": "t_in", "billing_group": "g"}]
285          }}},
286          "ocr": {"models": {"o": {
287            "input": [{"value_paid": 4, "billing_metric": "pages", "billing_group": "ocr"}]
288          }}}
289        }"#;
290        // 100*0.01 + 4*0.5 = 1.0 + 2.0 = 3.0
291        assert_eq!(
292            MistralProvider::parse(body)
293                .unwrap()
294                .cost
295                .unwrap()
296                .total_cost,
297            Some(3.0)
298        );
299    }
300
301    #[test]
302    fn test_default_currency_eur() {
303        let snap = MistralProvider::parse(r#"{"prices":[],"completion":{"models":{}}}"#).unwrap();
304        let cost = snap.cost.unwrap();
305        assert_eq!(cost.currency, "EUR");
306        assert_eq!(cost.total_cost, Some(0.0));
307    }
308
309    #[tokio::test]
310    async fn test_fetch_success() {
311        let server = MockServer::start().await;
312        Mock::given(method("GET"))
313            .and(path("/api/billing/v2/usage"))
314            .respond_with(ResponseTemplate::new(200).set_body_raw(SAMPLE, "application/json"))
315            .mount(&server)
316            .await;
317        let provider = MistralProvider::with_base_url(&server.uri());
318        let mut ctx = ProviderContext::new();
319        ctx.config.insert("cookie".into(), "sid=abc".into());
320        let snap = provider.fetch_usage(&ctx).await.unwrap();
321        assert_eq!(snap.cost.unwrap().total_cost, Some(2.5));
322    }
323
324    #[tokio::test]
325    async fn test_fetch_401() {
326        let server = MockServer::start().await;
327        Mock::given(method("GET"))
328            .and(path("/api/billing/v2/usage"))
329            .respond_with(ResponseTemplate::new(403))
330            .mount(&server)
331            .await;
332        let provider = MistralProvider::with_base_url(&server.uri());
333        let mut ctx = ProviderContext::new();
334        ctx.config.insert("cookie".into(), "bad".into());
335        assert!(matches!(
336            provider.fetch_usage(&ctx).await.unwrap_err(),
337            SpendPanelError::AuthFailed(_, _)
338        ));
339    }
340}
usage_monitor_cli/provider/mistral.rs

usage_monitor_cli/provider/
mistral.rs
