diff --git a/rust/src/providers/mistral/mod.rs b/rust/src/providers/mistral/mod.rs index 5fb8bdf2a3..9a70005f8f 100644 --- a/rust/src/providers/mistral/mod.rs +++ b/rust/src/providers/mistral/mod.rs @@ -9,16 +9,24 @@ use reqwest::Client; use serde::Deserialize; use std::collections::HashMap; +mod subscription; mod token_math; +use subscription::{SubscriptionBudget, SubscriptionBudgets}; + use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + CostSnapshot, FetchContext, NamedRateWindow, Provider, ProviderError, ProviderFetchResult, + ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; const BASE_URL: &str = "https://admin.mistral.ai"; const COOKIE_DOMAINS: [&str; 3] = ["admin.mistral.ai", "mistral.ai", "auth.mistral.ai"]; const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36"; +const CLIENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +/// Optional subscription-page enrichment joins on a fast deadline so a slow +/// `/subscription` render can never stall the refresh; degraded enrichment is +/// logged and skipped, never fatal. +const SUBSCRIPTION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(4); #[derive(Debug, Deserialize)] struct BillingResponse { @@ -168,7 +176,7 @@ impl MistralProvider { status_page_url: Some("https://status.mistral.ai"), }, client: crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(30)) + .timeout(CLIENT_TIMEOUT) .build() .unwrap_or_else(|_| Client::new()), } @@ -224,7 +232,48 @@ impl MistralProvider { .map_err(|e| ProviderError::Parse(format!("Failed to parse Mistral usage: {e}")))?; let summary = Self::summarize_billing(billing)?; - Ok(Self::build_result(summary)) + let budgets = match self.fetch_subscription_budgets(cookie_header).await { + Ok(budgets) => Some(budgets), + Err(error) => { + tracing::debug!(error = %error, "Mistral subscription allowance enrichment unavailable"); + None + } + }; + Ok(Self::build_result(summary, budgets)) + } + + async fn fetch_subscription_budgets( + &self, + cookie_header: &str, + ) -> Result { + let response = self + .client + .get(format!("{BASE_URL}/subscription")) + .timeout(SUBSCRIPTION_TIMEOUT) + .header("Accept", "text/html") + .header("Accept-Language", "en-US,en;q=0.9") + .header("Cookie", cookie_header) + .header("Referer", format!("{BASE_URL}/subscription")) + .header("User-Agent", USER_AGENT) + .send() + .await?; + let status = response.status(); + if status.as_u16() == 401 || status.as_u16() == 403 { + return Err(ProviderError::AuthRequired); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "Mistral subscription API returned {status}" + ))); + } + let final_url = response.url(); + if final_url.scheme() != "https" || final_url.host_str() != Some("admin.mistral.ai") { + return Err(ProviderError::Parse( + "Mistral subscription response came from an unexpected host".into(), + )); + } + let body = response.text().await?; + subscription::parse(&body).map_err(ProviderError::Parse) } fn summarize_billing(billing: BillingResponse) -> Result { @@ -305,7 +354,10 @@ impl MistralProvider { }) } - fn build_result(summary: MistralUsageSummary) -> ProviderFetchResult { + fn build_result( + summary: MistralUsageSummary, + budgets: Option, + ) -> ProviderFetchResult { let reset_date = summary.end_date.map(|dt| dt + chrono::Duration::seconds(1)); let cost_description = if summary.total_cost > 0.0 { format!( @@ -338,9 +390,39 @@ impl MistralProvider { token_detail )); + if let Some(budgets) = budgets { + if let Some(api) = budgets.api { + usage.primary = Self::budget_window(&api); + usage.primary_label = Some("Included API".to_string()); + } + if let Some(vibe) = budgets.vibe { + usage.extra_rate_windows.push(NamedRateWindow::new( + "mistral-monthly-plan", + "Monthly Plan", + Self::budget_window(&vibe), + )); + } + } + ProviderFetchResult::new(usage, "web").with_cost(cost) } + fn budget_window(budget: &SubscriptionBudget) -> RateWindow { + let used = budget.used_amount(); + let remaining = budget.remaining_amount(); + let description = format!( + "{used:.2} {currency} / {limit:.2} {currency} · {remaining:.2} {currency} remaining", + currency = budget.currency, + limit = budget.limit, + ); + RateWindow::with_details( + budget.used_percent, + None, + budget.resets_at, + Some(description), + ) + } + fn build_price_index(prices: Vec) -> HashMap { prices .into_iter() @@ -502,7 +584,7 @@ mod tests { assert!((summary.total_cost - 0.005).abs() < 0.000001); assert_eq!(summary.model_count, 1); - let result = MistralProvider::build_result(summary); + let result = MistralProvider::build_result(summary, None); assert_eq!( result.cost.as_ref().map(|c| c.currency_code.as_str()), Some("EUR") @@ -518,6 +600,46 @@ mod tests { ); } + #[test] + fn attaches_subscription_allowances_without_replacing_billing_cost() { + let summary = MistralUsageSummary { + total_cost: 12.5, + currency: "EUR".to_string(), + currency_symbol: "€".to_string(), + total_input_tokens: 100, + total_output_tokens: 50, + total_cached_tokens: 0, + model_count: 1, + end_date: None, + }; + let result = MistralProvider::build_result( + summary, + Some(SubscriptionBudgets { + api: Some(SubscriptionBudget { + used_percent: 25.0, + limit: 100.0, + currency: "USD".to_string(), + resets_at: None, + }), + vibe: Some(SubscriptionBudget { + used_percent: 50.0, + limit: 20.0, + currency: "EUR".to_string(), + resets_at: None, + }), + }), + ); + + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!(result.usage.primary_label.as_deref(), Some("Included API")); + assert_eq!(result.usage.extra_rate_windows.len(), 1); + assert_eq!( + result.usage.extra_rate_windows[0].id, + "mistral-monthly-plan" + ); + assert_eq!(result.cost.as_ref().map(|cost| cost.used), Some(12.5)); + } + #[test] fn extracts_csrf_token_from_cookie_header() { assert_eq!( diff --git a/rust/src/providers/mistral/subscription.rs b/rust/src/providers/mistral/subscription.rs new file mode 100644 index 0000000000..0d09b3e84e --- /dev/null +++ b/rust/src/providers/mistral/subscription.rs @@ -0,0 +1,272 @@ +//! Parser for the authenticated Mistral subscription page. +//! +//! The page embeds a small React Flight stream: `self.__next_f.push` calls +//! whose second element is a chunk of the Flight text. This parser scans the +//! markers, joins the chunks, and reads the `0::` lines. Keep it +//! local to the Mistral provider so a change in that page format cannot +//! affect other providers or the shared snapshot contract. + +use chrono::{DateTime, Utc}; +use serde_json::Value; + +#[derive(Debug, Clone, PartialEq)] +pub(super) struct SubscriptionBudgets { + pub(super) api: Option, + pub(super) vibe: Option, +} + +#[derive(Debug, Clone, PartialEq)] +pub(super) struct SubscriptionBudget { + pub(super) used_percent: f64, + pub(super) limit: f64, + pub(super) currency: String, + pub(super) resets_at: Option>, +} + +impl SubscriptionBudget { + pub(super) fn used_amount(&self) -> f64 { + self.limit * (self.used_percent / 100.0) + } + + pub(super) fn remaining_amount(&self) -> f64 { + (self.limit - self.used_amount()).max(0.0) + } +} + +pub(super) fn parse(html: &str) -> Result { + let stream = flight_chunks(html).join(""); + let mut matches = Vec::new(); + collect_models(stream.as_bytes(), &mut |value| { + collect_budgets(value, &mut matches); + })?; + + match matches.as_slice() { + [] => Err("Mistral subscription budgets were not found".into()), + [budgets] => Ok(budgets.clone()), + _ => Err("Mistral subscription budgets were ambiguous".into()), + } +} + +fn flight_chunks(html: &str) -> Vec { + let bytes = html.as_bytes(); + let marker = b"self.__next_f.push("; + let mut cursor = 0; + let mut chunks = Vec::new(); + + while cursor < bytes.len() { + let Some(offset) = find_bytes(&bytes[cursor..], marker) else { + break; + }; + let marker_start = cursor + offset; + let mut start = marker_start + marker.len(); + while start < bytes.len() && bytes[start].is_ascii_whitespace() { + start += 1; + } + if start >= bytes.len() || bytes[start] != b'[' { + cursor = marker_start + marker.len(); + continue; + } + let Some(end) = json_container_end(bytes, start) else { + cursor = start + 1; + continue; + }; + if let Ok(Value::Array(values)) = serde_json::from_slice::(&bytes[start..end]) + && values.first().and_then(Value::as_u64) == Some(1) + && let Some(chunk) = values.get(1).and_then(Value::as_str) + { + chunks.push(chunk.to_string()); + } + cursor = end; + } + + chunks +} + +fn collect_models(data: &[u8], body: &mut impl FnMut(Value)) -> Result<(), String> { + // Advance past a scanned line, stopping at the final newlineless tail. + let advance = |line_end: usize| line_end.min(data.len().saturating_sub(1) + 1) + 1; + let mut cursor = 0; + while cursor < data.len() { + let line_end = data[cursor..] + .iter() + .position(|byte| *byte == b'\n') + .map_or(data.len(), |offset| cursor + offset); + let line = &data[cursor..line_end]; + let colon = line + .iter() + .position(|byte| *byte == b':') + .filter(|colon| *colon != 0 && line[..*colon].iter().all(u8::is_ascii_hexdigit)); + let Some(colon) = colon else { + cursor = advance(line_end); + continue; + }; + + let start = cursor + colon + 1; + let payload = data[start..line_end].trim_ascii(); + if matches!(payload.first(), Some(b'[' | b'{')) + && let Ok(value) = serde_json::from_slice::(payload) + { + body(value); + } + cursor = advance(line_end); + } + Ok(()) +} + +fn collect_budgets(value: Value, matches: &mut Vec) { + match value { + Value::Object(object) => { + if let Some(Value::Object(budget)) = object.get("budget") { + let candidate = SubscriptionBudgets { + api: budget.get("api_budget").and_then(parse_budget), + vibe: budget.get("vibe_budget").and_then(parse_budget), + }; + if candidate.api.is_some() || candidate.vibe.is_some() { + matches.push(candidate); + } + } + for child in object.into_values() { + collect_budgets(child, matches); + } + } + Value::Array(values) => { + for child in values { + collect_budgets(child, matches); + } + } + _ => {} + } +} + +fn parse_budget(value: &Value) -> Option { + let object = value.as_object()?; + let used_percent = number(object.get("usage_percentage")?)?; + let limit = number(object.get("initial_budget")?)?; + let currency = object + .get("currency") + .and_then(Value::as_str)? + .trim() + .to_uppercase(); + if !used_percent.is_finite() || used_percent < 0.0 || !limit.is_finite() || limit <= 0.0 { + return None; + } + if currency.is_empty() { + return None; + } + let used_amount = limit * (used_percent / 100.0); + if !used_amount.is_finite() { + return None; + } + let resets_at = object + .get("reset_at") + .and_then(Value::as_str) + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|value| value.with_timezone(&Utc)); + Some(SubscriptionBudget { + used_percent, + limit, + currency, + resets_at, + }) +} + +fn number(value: &Value) -> Option { + value + .as_f64() + .or_else(|| value.as_str()?.trim().parse::().ok()) +} + +fn json_container_end(data: &[u8], start: usize) -> Option { + let mut closers = Vec::new(); + let mut in_string = false; + let mut escaped = false; + let mut index = start; + while index < data.len() { + let byte = data[index]; + if in_string { + if escaped { + escaped = false; + } else if byte == b'\\' { + escaped = true; + } else if byte == b'"' { + in_string = false; + } + } else { + match byte { + b'"' => in_string = true, + b'[' => closers.push(b']'), + b'{' => closers.push(b'}'), + b']' | b'}' => { + if closers.last().copied() != Some(byte) { + return None; + } + closers.pop(); + if closers.is_empty() { + return Some(index + 1); + } + } + _ => {} + } + } + index += 1; + } + None +} + +fn find_bytes(haystack: &[u8], needle: &[u8]) -> Option { + haystack + .windows(needle.len()) + .position(|window| window == needle) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn flight_html(model: &str) -> String { + let row = format!("0:{model}\n"); + format!( + "", + serde_json::to_string(&row).unwrap() + ) + } + + #[test] + fn parses_api_and_vibe_budgets_from_flight_data() { + let html = flight_html( + r#"{"budget":{"api_budget":{"usage_percentage":25,"initial_budget":100,"currency":"usd","reset_at":"2026-09-30T00:00:00Z"},"vibe_budget":{"usage_percentage":50,"initial_budget":20,"currency":"EUR","reset_at":null}}}"#, + ); + let budgets = parse(&html).unwrap(); + assert_eq!( + budgets.api.as_ref().map(|budget| budget.used_percent), + Some(25.0) + ); + assert_eq!( + budgets.api.as_ref().map(|budget| budget.currency.as_str()), + Some("USD") + ); + assert_eq!( + budgets + .vibe + .as_ref() + .map(|budget| budget.remaining_amount()), + Some(10.0) + ); + } + + #[test] + fn ignores_invalid_budget_values() { + let html = flight_html( + r#"{"budget":{"api_budget":{"usage_percentage":-1,"initial_budget":100,"currency":"USD"}}}"#, + ); + assert!(parse(&html).is_err()); + } + + #[test] + fn handles_brackets_inside_flight_strings() { + let html = flight_html( + r#"{"note":"fake [brackets]", "budget":{"api_budget":{"usage_percentage":1,"initial_budget":10,"currency":"USD"}}}"#, + ); + assert!(parse(&html).unwrap().api.is_some()); + } +}