From 540fa80ebe7a68f27a33731dfefff3af115a7cf9 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:48:50 +0700 Subject: [PATCH 01/11] Add v0.64 provider adapters --- .../src/commands/provider_settings.rs | 36 ++ .../src/components/providers/providerIcons.ts | 3 + .../desktop-tauri/src/test/providerCatalog.ts | 3 + rust/src/core/provider.rs | 29 +- rust/src/core/provider_factory.rs | 22 +- rust/src/core/token_accounts.rs | 3 + rust/src/providers/helmcode.rs | 396 ++++++++++++++++ rust/src/providers/huggingface/mod.rs | 213 ++++++++- rust/src/providers/mod.rs | 6 + rust/src/providers/muse/mod.rs | 60 ++- rust/src/providers/typesafe.rs | 443 ++++++++++++++++++ rust/src/providers/v0.rs | 367 +++++++++++++++ rust/src/settings/api_keys.rs | 11 + rust/src/settings/provider_workspace.rs | 15 + 14 files changed, 1582 insertions(+), 25 deletions(-) create mode 100644 rust/src/providers/helmcode.rs create mode 100644 rust/src/providers/typesafe.rs create mode 100644 rust/src/providers/v0.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs index 7d26ba8fdc..f9344fcfba 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -224,6 +224,8 @@ fn cookie_source_provider(provider_id: &str) -> Option ProviderId::Notion, "grok" => ProviderId::Grok, "replicate" => ProviderId::Replicate, + "helmcode" => ProviderId::Helmcode, + "typesafe" => ProviderId::TypeSafe, _ => return None, }) } @@ -325,6 +327,8 @@ fn workspace_provider(provider_id: &str) -> Option { "opencodego" => ProviderId::OpenCodeGo, "zed" => ProviderId::Zed, "xai" => ProviderId::Xai, + "v0" => ProviderId::V0, + "helmcode" => ProviderId::Helmcode, _ => return None, }) } @@ -745,6 +749,38 @@ pub fn cookie_source_options_for(provider_id: &str, lang: Language) -> Vec vec![ + cookie_option( + lang, + "auto", + "Automatically imports the signed-in Helmcode or NaN Builders browser session.", + "", + None, + ), + cookie_option( + lang, + "manual", + "", + "Paste a Cookie header and select the tenant in the workspace field.", + None, + ), + ], + "typesafe" => vec![ + cookie_option( + lang, + "auto", + "Automatically imports the signed-in TypeSafe console session.", + "", + None, + ), + cookie_option( + lang, + "manual", + "", + "Paste a Cookie header from the TypeSafe billing page.", + None, + ), + ], _ => Vec::new(), } } diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index ccfa304078..2fdc4380d7 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -188,6 +188,9 @@ export const PROVIDER_ICON_REGISTRY: Record = { grok: { id: "grok", brandColor: "#111827", fallbackLetter: "G", svgPath: RAW.grok }, groq: { id: "groq", brandColor: "#f55036", fallbackLetter: "G", svgPath: RAW.groq }, huggingface: { id: "huggingface", brandColor: "#ffd21e", fallbackLetter: "H", svgPath: RAW.huggingface }, + helmcode: { id: "helmcode", brandColor: "#4f46e5", fallbackLetter: "H" }, + v0: { id: "v0", brandColor: "#111827", fallbackLetter: "V" }, + typesafe: { id: "typesafe", brandColor: "#2563eb", fallbackLetter: "T" }, jetbrains: { id: "jetbrains", brandColor: "#ff3399", fallbackLetter: "J", svgPath: RAW.jetbrains }, kilo: { id: "kilo", brandColor: "#5d87ff", fallbackLetter: "K", svgPath: RAW.kilo }, bedrock: { id: "bedrock", brandColor: "#ff9900", fallbackLetter: "B", svgPath: RAW.bedrock }, diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index c8ab2daa5d..dda4c24441 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -57,6 +57,9 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["deepgram", "Deepgram"], ["groq", "Groq"], ["huggingface", "Hugging Face"], + ["helmcode", "Helmcode"], + ["v0", "v0"], + ["typesafe", "TypeSafe"], ["llmproxy", "LLM Proxy"], ["chutes", "Chutes"], ["litellm", "LiteLLM"], diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index d6a38df551..cc954839f7 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -64,6 +64,9 @@ pub enum ProviderId { Deepgram, Groq, HuggingFace, + Helmcode, + V0, + TypeSafe, LLMProxy, Chutes, LiteLLM, @@ -148,6 +151,9 @@ impl ProviderId { ProviderId::Deepgram, ProviderId::Groq, ProviderId::HuggingFace, + ProviderId::Helmcode, + ProviderId::V0, + ProviderId::TypeSafe, ProviderId::LLMProxy, ProviderId::Chutes, ProviderId::LiteLLM, @@ -235,6 +241,9 @@ impl ProviderId { ProviderId::Deepgram => "deepgram", ProviderId::Groq => "groq", ProviderId::HuggingFace => "huggingface", + ProviderId::Helmcode => "helmcode", + ProviderId::V0 => "v0", + ProviderId::TypeSafe => "typesafe", ProviderId::LLMProxy => "llmproxy", ProviderId::Chutes => "chutes", ProviderId::LiteLLM => "litellm", @@ -319,6 +328,9 @@ impl ProviderId { ProviderId::Deepgram => "Deepgram", ProviderId::Groq => "Groq", ProviderId::HuggingFace => "Hugging Face", + ProviderId::Helmcode => "Helmcode", + ProviderId::V0 => "v0", + ProviderId::TypeSafe => "TypeSafe", ProviderId::LLMProxy => "LLM Proxy", ProviderId::Chutes => "Chutes", ProviderId::LiteLLM => "LiteLLM", @@ -413,6 +425,9 @@ impl ProviderId { ProviderId::Deepgram => None, ProviderId::Groq => None, ProviderId::HuggingFace => None, + ProviderId::Helmcode => Some("helmcode.com"), + ProviderId::TypeSafe => Some("typesafe.ai"), + ProviderId::V0 => None, ProviderId::LLMProxy => None, ProviderId::Chutes => None, ProviderId::LiteLLM => None, @@ -502,6 +517,9 @@ impl ProviderId { "deepgram" | "dg" => Some(ProviderId::Deepgram), "groq" | "groqcloud" | "groq-cloud" | "groq cloud" => Some(ProviderId::Groq), "huggingface" | "hugging-face" | "hugging face" | "hf" => Some(ProviderId::HuggingFace), + "helmcode" | "nan-builders" | "nan builders" => Some(ProviderId::Helmcode), + "v0" | "v0-dev" | "v0.dev" => Some(ProviderId::V0), + "typesafe" | "type-safe" | "type safe" => Some(ProviderId::TypeSafe), "llmproxy" | "llm-proxy" | "llm proxy" => Some(ProviderId::LLMProxy), "chutes" | "chutes-ai" | "chutes ai" => Some(ProviderId::Chutes), "litellm" | "lite-llm" | "lite llm" => Some(ProviderId::LiteLLM), @@ -988,6 +1006,9 @@ pub fn cli_name_map() -> HashMap<&'static str, ProviderId> { map.insert("groq-cloud", ProviderId::Groq); map.insert("hugging-face", ProviderId::HuggingFace); map.insert("hf", ProviderId::HuggingFace); + map.insert("nan-builders", ProviderId::Helmcode); + map.insert("v0-dev", ProviderId::V0); + map.insert("type-safe", ProviderId::TypeSafe); map.insert("chutes-ai", ProviderId::Chutes); map.insert("lite-llm", ProviderId::LiteLLM); map.insert("zed-ai", ProviderId::Zed); @@ -1057,6 +1078,9 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Deepgram => "#13EF93", ProviderId::Groq => "#F55036", ProviderId::HuggingFace => "#FFD21E", + ProviderId::Helmcode => "#4F46E5", + ProviderId::V0 => "#111827", + ProviderId::TypeSafe => "#2563EB", ProviderId::LLMProxy => "#4F46E5", ProviderId::Chutes => "#FF5C35", ProviderId::LiteLLM => "#0EA5E9", @@ -1096,7 +1120,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 77); + assert_eq!(all.len(), 80); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); @@ -1130,6 +1154,9 @@ mod tests { assert!(all.contains(&ProviderId::Deepgram)); assert!(all.contains(&ProviderId::Groq)); assert!(all.contains(&ProviderId::HuggingFace)); + assert!(all.contains(&ProviderId::Helmcode)); + assert!(all.contains(&ProviderId::V0)); + assert!(all.contains(&ProviderId::TypeSafe)); assert!(all.contains(&ProviderId::LLMProxy)); assert!(all.contains(&ProviderId::Chutes)); assert!(all.contains(&ProviderId::LiteLLM)); diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index aff4b6f31d..712c1d4fb9 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -13,15 +13,16 @@ use crate::providers::{ CodexProvider, CommandCodeProvider, CopilotProvider, CrofProvider, CrossModelProvider, CursorProvider, DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevinProvider, DoubaoProvider, ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, - GrokProvider, GroqProvider, HuggingFaceProvider, InfiniProvider, JetBrainsProvider, - KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, LiteLLMProvider, - LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, - MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, NousProvider, - OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, - PerplexityProvider, PiProvider, PoeProvider, QoderProvider, QwenCloudProvider, - ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, - VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, - XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, + GrokProvider, GroqProvider, HelmcodeProvider, HuggingFaceProvider, InfiniProvider, + JetBrainsProvider, KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, + LiteLLMProvider, LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, + MistralProvider, MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, + NousProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, + OpenRouterProvider, PerplexityProvider, PiProvider, PoeProvider, QoderProvider, + QwenCloudProvider, ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, + T3ChatProvider, TypeSafeProvider, V0Provider, VeniceProvider, VertexAIProvider, WarpProvider, + WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, + ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -82,6 +83,9 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::Deepgram => Box::new(DeepgramProvider::new()), ProviderId::Groq => Box::new(GroqProvider::new()), ProviderId::HuggingFace => Box::new(HuggingFaceProvider::new()), + ProviderId::Helmcode => Box::new(HelmcodeProvider::new()), + ProviderId::V0 => Box::new(V0Provider::new()), + ProviderId::TypeSafe => Box::new(TypeSafeProvider::new()), ProviderId::LLMProxy => Box::new(LLMProxyProvider::new()), ProviderId::Chutes => Box::new(ChutesProvider::new()), ProviderId::LiteLLM => Box::new(LiteLLMProvider::new()), diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index a8f3cc71ae..6dedc0fbde 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -361,6 +361,9 @@ impl TokenAccountSupport { | ProviderId::ElevenLabs | ProviderId::Deepgram | ProviderId::Groq + | ProviderId::Helmcode + | ProviderId::V0 + | ProviderId::TypeSafe | ProviderId::LLMProxy | ProviderId::Chutes | ProviderId::LiteLLM diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs new file mode 100644 index 0000000000..02ac002f2a --- /dev/null +++ b/rust/src/providers/helmcode.rs @@ -0,0 +1,396 @@ +//! Helmcode Cloud and NaN Builders dashboard quota provider. + +use async_trait::async_trait; +use chrono::{DateTime, Datelike, TimeZone, Utc}; +use reqwest::{Client, StatusCode}; +use serde_json::Value; +use std::time::Duration; + +use crate::core::{ + CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, + ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, +}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Tenant { + Helmcode, + NanBuilders, +} + +impl Tenant { + fn domain(self) -> &'static str { + match self { + Self::Helmcode => "helmcode.com", + Self::NanBuilders => "nan.builders", + } + } + fn name(self) -> &'static str { + match self { + Self::Helmcode => "Helmcode Cloud", + Self::NanBuilders => "NaN Builders", + } + } + fn api(self) -> String { + format!("https://cloud-api.{}", self.domain()) + } + fn origin(self) -> String { + format!("https://cloud.{}", self.domain()) + } +} + +#[derive(Debug, Clone, PartialEq)] +struct ModelQuota { + name: String, + cap: f64, + used: f64, + credit: f64, + window_hours: Option, + resets_at: Option>, +} + +pub struct HelmcodeProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl HelmcodeProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::Helmcode, + display_name: "Helmcode", + session_label: "Quota", + weekly_label: "Quota", + supports_opus: false, + supports_credits: true, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://cloud.helmcode.com/dashboard"), + status_page_url: None, + tertiary_label_key: None, + }, + client: crate::core::credentialed_http_client_builder() + .timeout(Duration::from_secs(10)) + .build() + .unwrap_or_else(|_| Client::new()), + } + } + + async fn fetch_web(&self, ctx: &FetchContext) -> Result { + let manual_tenant = ctx + .workspace_id + .as_deref() + .is_some_and(|value| value.eq_ignore_ascii_case("nanBuilders")); + let candidates = if ctx.manual_cookie_header.is_some() { + vec![if manual_tenant { + Tenant::NanBuilders + } else { + Tenant::Helmcode + }] + } else { + vec![Tenant::Helmcode, Tenant::NanBuilders] + }; + let mut rejected = false; + for tenant in candidates { + let cookie = match ctx.manual_cookie_header.as_deref() { + Some(raw) => normalize_cookie(raw).ok_or(ProviderError::NoCookies)?, + None => match crate::providers::browser_cookie_header(&[tenant.domain()]) { + Ok(header) => header, + Err(_) => continue, + }, + }; + match self.fetch_tenant(tenant, &cookie).await { + Ok(result) => return Ok(result), + Err(ProviderError::AuthRequired) => rejected = true, + Err(error) => return Err(error), + } + } + if rejected { + Err(ProviderError::AuthRequired) + } else { + Err(ProviderError::NotInstalled( + "Sign in to cloud.helmcode.com or cloud.nan.builders, or paste a Cookie header." + .into(), + )) + } + } + + async fn fetch_tenant( + &self, + tenant: Tenant, + cookie: &str, + ) -> Result { + let quota = self + .get(tenant, cookie, "/api/usage/quota", false) + .await? + .ok_or_else(|| parse_failure("quota"))?; + let billing = self.get(tenant, cookie, "/api/billing", true).await?; + let premium = billing + .as_ref() + .and_then(|value| value.get("subscription")) + .and_then(|value| value.get("premium")) + .and_then(Value::as_bool) + == Some(true); + let mut models = parse_models("a, premium)?; + models.sort_by(|a, b| { + (b.used / b.cap) + .partial_cmp(&(a.used / a.cap)) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.name.cmp(&b.name)) + }); + let primary = models + .first() + .map(model_window) + .unwrap_or_else(|| RateWindow::informational("No active model quota")); + let mut usage = UsageSnapshot::new(primary) + .with_organization(tenant.name()) + .with_login_method("Dashboard session"); + for model in models.iter().skip(1) { + usage = usage.with_extra_rate_window( + format!("helmcode-{}", model.name), + model.name.clone(), + model_window(model), + ); + } + let mut result = ProviderFetchResult::new(usage, "web"); + if tenant == Tenant::Helmcode + && let Some(credits) = self + .get(tenant, cookie, "/api/billing/credits", true) + .await? + && let Some(balance_micros) = credits.get("balanceMicros").and_then(Value::as_i64) + { + let currency = credits + .get("currency") + .and_then(Value::as_str) + .unwrap_or("EUR") + .to_ascii_uppercase(); + if currency.len() == 3 && currency.chars().all(|ch| ch.is_ascii_uppercase()) { + result = result.with_cost( + CostSnapshot::new(0.0, currency, "Prepaid balance") + .with_balance((balance_micros.max(0) as f64) / 1_000_000.0), + ); + } + } + Ok(result) + } + + async fn get( + &self, + tenant: Tenant, + cookie: &str, + path: &str, + optional: bool, + ) -> Result, ProviderError> { + let response = self + .client + .get(format!("{}{path}", tenant.api())) + .header("Cookie", cookie) + .header("Origin", tenant.origin()) + .header("Referer", format!("{}/dashboard", tenant.origin())) + .send() + .await?; + let status = response.status(); + if status == StatusCode::UNAUTHORIZED + || status == StatusCode::FORBIDDEN + || status.is_redirection() + { + return Err(ProviderError::AuthRequired); + } + if optional && !status.is_success() { + return Ok(None); + } + if status == StatusCode::TOO_MANY_REQUESTS { + return Err(ProviderError::Other("Helmcode rate limit reached.".into())); + } + if status == StatusCode::REQUEST_TIMEOUT || status.is_server_error() { + return Err(ProviderError::Other( + "Helmcode dashboard is unavailable.".into(), + )); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "Helmcode dashboard returned HTTP {status}." + ))); + } + let value = response + .json() + .await + .map_err(|_| parse_failure("invalid JSON"))?; + Ok(Some(value)) + } +} + +impl Default for HelmcodeProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for HelmcodeProvider { + fn id(&self) -> ProviderId { + ProviderId::Helmcode + } + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::Web => self.fetch_web(ctx).await, + SourceMode::OAuth | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::Web] + } + fn supports_web(&self) -> bool { + true + } + fn owns_browser_cookie_resolution(&self) -> bool { + true + } +} + +fn parse_models(quota: &Value, premium: bool) -> Result, ProviderError> { + let object = quota + .as_object() + .ok_or_else(|| parse_failure("quota object"))?; + let period_start = object + .get("periodStart") + .and_then(Value::as_str) + .ok_or_else(|| parse_failure("periodStart"))?; + let fallback = DateTime::parse_from_rfc3339(period_start) + .ok() + .map(|date| { + let date = date.with_timezone(&Utc); + Utc.with_ymd_and_hms(date.year(), date.month(), 1, 0, 0, 0) + .single() + }) + .flatten(); + let models = object + .get("models") + .and_then(Value::as_array) + .ok_or_else(|| parse_failure("models"))?; + models + .iter() + .map(|value| { + let row = value.as_object().ok_or_else(|| parse_failure("model"))?; + let name = row + .get("model") + .and_then(Value::as_str) + .map(str::trim) + .filter(|name| { + !name.is_empty() + && name.chars().count() <= 120 + && !name.chars().any(char::is_control) + }) + .ok_or_else(|| parse_failure("model name"))? + .to_string(); + let cap = nonnegative(row.get("cap"), "cap")?; + let used = nonnegative(row.get("tokensUsed"), "tokensUsed")?; + let credit = + optional_nonnegative(row.get("creditTokens"), "creditTokens")?.unwrap_or(0.0); + let window_hours = optional_nonnegative(row.get("windowHours"), "windowHours")? + .map(|value| value as u32); + if window_hours.is_some_and(|hours| hours == 0 || hours > 8_760) { + return Err(parse_failure("windowHours")); + } + let resets_at = row + .get("periodEnd") + .and_then(Value::as_str) + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|date| date.with_timezone(&Utc)) + .or(fallback); + Ok(ModelQuota { + name, + cap, + used, + credit, + window_hours, + resets_at, + }) + }) + .filter_map(|result| match result { + Ok(model) if model.cap > 0.0 && (model.window_hours.is_none() || premium) => { + Some(Ok(model)) + } + Ok(_) => None, + Err(error) => Some(Err(error)), + }) + .collect() +} + +fn model_window(model: &ModelQuota) -> RateWindow { + let mut window = RateWindow::with_details( + (model.used / model.cap * 100.0).clamp(0.0, 100.0), + model.window_hours.map(|hours| hours * 60), + model.resets_at, + None, + ); + window.reset_description = Some(format!( + "{} · {:.0} / {:.0} tokens{}", + model.name, + model.used, + model.cap, + if model.credit > 0.0 { + format!(" · {:.0} credit-funded", model.credit) + } else { + String::new() + } + )); + window +} + +fn nonnegative(value: Option<&Value>, field: &str) -> Result { + value + .and_then(Value::as_f64) + .filter(|value| value.is_finite() && *value >= 0.0 && value.fract() == 0.0) + .ok_or_else(|| parse_failure(field)) +} +fn optional_nonnegative(value: Option<&Value>, field: &str) -> Result, ProviderError> { + match value { + None | Some(Value::Null) => Ok(None), + Some(value) => nonnegative(Some(value), field).map(Some), + } +} +fn normalize_cookie(raw: &str) -> Option { + let value = raw + .trim() + .strip_prefix("Cookie:") + .unwrap_or(raw.trim()) + .trim(); + (!value.is_empty() && !value.chars().any(char::is_control)).then(|| value.to_string()) +} +fn parse_failure(field: impl AsRef) -> ProviderError { + ProviderError::Parse(format!( + "Helmcode quota response format changed: {}", + field.as_ref() + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn parses_models_filters_nonpremium_rolling_windows_and_orders_are_external() { + let value = json!({"periodStart":"2030-02-12T00:00:00Z","models":[ + {"model":"monthly","cap":1000,"tokensUsed":250,"creditTokens":50,"periodEnd":null}, + {"model":"rolling","cap":100,"tokensUsed":90,"windowHours":5} + ]}); + let free = parse_models(&value, false).unwrap(); + assert_eq!(free.len(), 1); + assert_eq!(free[0].name, "monthly"); + assert_eq!(free[0].resets_at.unwrap().day(), 1); + assert_eq!(parse_models(&value, true).unwrap().len(), 2); + } + + #[test] + fn rejects_fractional_or_negative_quota_counts() { + assert!(parse_models(&json!({"periodStart":"2030-01-01T00:00:00Z","models":[{"model":"x","cap":1.5,"tokensUsed":0}]}), true).is_err()); + assert!(parse_models(&json!({"periodStart":"2030-01-01T00:00:00Z","models":[{"model":"x","cap":1,"tokensUsed":-1}]}), true).is_err()); + } +} diff --git a/rust/src/providers/huggingface/mod.rs b/rust/src/providers/huggingface/mod.rs index 3b3f5b3d16..857bbc0c26 100644 --- a/rust/src/providers/huggingface/mod.rs +++ b/rust/src/providers/huggingface/mod.rs @@ -47,6 +47,7 @@ struct ZeroGpuSnapshot { #[derive(Debug, Clone, PartialEq, Eq)] struct IdentitySnapshot { + user_id: Option, name: Option, email: Option, plan: Option, @@ -173,8 +174,72 @@ impl HuggingFaceProvider { let billing = parse_billing(billing?)?; let identity = identity.and_then(|value| parse_identity(&value)); let zerogpu = zerogpu.and_then(|value| parse_zerogpu(&value)); + let balance = match identity + .as_ref() + .and_then(|identity| identity.user_id.as_deref()) + { + Some(user_id) => self.fetch_optional_wallet(user_id).await, + None => None, + }; - Ok(build_result(billing, identity, zerogpu)) + Ok(build_result(billing, identity, zerogpu, balance)) + } + + async fn fetch_optional_wallet(&self, expected_user_id: &str) -> Option { + let cookie = crate::providers::browser_cookie_header(&["huggingface.co"]).ok()?; + let billing = self + .fetch_cookie_text( + "https://huggingface.co/settings/billing", + &cookie, + "text/html", + ) + .await + .ok()?; + let candidate = parse_wallet_balance(&billing).ok()?; + let whoami = self + .fetch_cookie_text(WHOAMI_URL, &cookie, "application/json") + .await + .ok()?; + let profile: Value = serde_json::from_str(&whoami).ok()?; + let observed_user_id = profile + .get("type") + .and_then(Value::as_str) + .filter(|kind| *kind == "user") + .and_then(|_| profile.get("id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty())?; + (observed_user_id == expected_user_id).then_some(candidate) + } + + async fn fetch_cookie_text( + &self, + url: &str, + cookie: &str, + accept: &str, + ) -> Result { + let response = tokio::time::timeout( + OPTIONAL_TIMEOUT, + self.client + .get(url) + .header(reqwest::header::COOKIE, cookie) + .header(reqwest::header::ACCEPT, accept) + .send(), + ) + .await + .map_err(|_| ProviderError::Timeout)??; + if !response.status().is_success() { + return Err(classify_status(response.status())); + } + let bytes = response.bytes().await?; + if bytes.len() > MAX_RESPONSE_BYTES { + return Err(ProviderError::Parse( + "Hugging Face returned an oversized wallet response.".to_string(), + )); + } + String::from_utf8(bytes.to_vec()).map_err(|_| { + ProviderError::Parse("Hugging Face returned invalid wallet text.".to_string()) + }) } async fn fetch_optional_json(&self, url: Url, token: &str) -> Option { @@ -400,17 +465,23 @@ fn parse_timestamp(value: &Value) -> Option> { } fn parse_identity(value: &Value) -> Option { + let user_id = (value.get("type").and_then(Value::as_str) == Some("user")) + .then(|| safe_text(value.get("id").and_then(Value::as_str))) + .flatten(); let name = safe_text(value.get("name").and_then(Value::as_str)); let email = safe_text(value.get("email").and_then(Value::as_str)); let plan = value .get("isPro") .and_then(Value::as_bool) .map(|is_pro| if is_pro { "Pro" } else { "Free" }.to_string()); - (name.is_some() || email.is_some() || plan.is_some()).then_some(IdentitySnapshot { - name, - email, - plan, - }) + (user_id.is_some() || name.is_some() || email.is_some() || plan.is_some()).then_some( + IdentitySnapshot { + user_id, + name, + email, + plan, + }, + ) } fn safe_text(value: Option<&str>) -> Option { @@ -421,10 +492,116 @@ fn safe_text(value: Option<&str>) -> Option { Some(value.to_string()) } +fn parse_wallet_balance(html: &str) -> Result { + let mut current = Vec::new(); + let mut legacy = Vec::new(); + let mut rest = html; + while let Some(index) = rest.find("data-props") { + rest = &rest[index + "data-props".len()..]; + let trimmed = rest.trim_start(); + let Some(after_equals) = trimmed.strip_prefix('=') else { + continue; + }; + let after_equals = after_equals.trim_start(); + let Some(quote) = after_equals + .chars() + .next() + .filter(|value| matches!(value, '\'' | '"')) + else { + continue; + }; + let payload = &after_equals[quote.len_utf8()..]; + let Some(end) = payload.find(quote) else { + break; + }; + let decoded = decode_html_entities(&payload[..end])?; + rest = &payload[end + quote.len_utf8()..]; + let Ok(value) = serde_json::from_str::(&decoded) else { + continue; + }; + let Some(object) = value.as_object() else { + continue; + }; + if let Some(entity) = object.get("entity").and_then(Value::as_object) + && entity.contains_key("currentBalanceUsd") + { + if entity.get("type").and_then(Value::as_str) != Some("user") { + return Err(invalid_wallet("wallet entity type")); + } + current.push(wallet_number( + entity.get("currentBalanceUsd"), + "currentBalanceUsd", + )?); + } + if object.contains_key("invoiceCreditsCents") { + let cents = wallet_number(object.get("invoiceCreditsCents"), "invoiceCreditsCents")?; + if cents.fract() != 0.0 { + return Err(invalid_wallet("invoiceCreditsCents")); + } + legacy.push(cents / 100.0); + } + } + match (current.as_slice(), legacy.as_slice()) { + ([balance], _) => Ok(*balance), + ([], [balance]) => Ok(*balance), + ([], _) => Err(invalid_wallet("missing or ambiguous legacy wallet")), + _ => Err(invalid_wallet("ambiguous current wallet")), + } +} + +fn wallet_number(value: Option<&Value>, field: &str) -> Result { + value + .and_then(Value::as_f64) + .filter(|value| value.is_finite() && *value >= 0.0) + .ok_or_else(|| invalid_wallet(field)) +} + +fn invalid_wallet(field: &str) -> ProviderError { + ProviderError::Parse(format!("Hugging Face wallet field '{field}' was invalid.")) +} + +fn decode_html_entities(raw: &str) -> Result { + let mut output = String::with_capacity(raw.len()); + let mut rest = raw; + while let Some(index) = rest.find('&') { + output.push_str(&rest[..index]); + rest = &rest[index + 1..]; + let Some(end) = rest.find(';') else { + return Err(invalid_wallet("HTML entity")); + }; + let entity = &rest[..end]; + let decoded = match entity { + "amp" => '&', + "apos" => '\'', + "gt" => '>', + "lt" => '<', + "nbsp" => '\u{00a0}', + "quot" => '"', + value if value.starts_with("#x") || value.starts_with("#X") => { + u32::from_str_radix(&value[2..], 16) + .ok() + .and_then(char::from_u32) + .ok_or_else(|| invalid_wallet("HTML entity"))? + } + value if value.starts_with('#') => value[1..] + .parse::() + .ok() + .and_then(char::from_u32) + .ok_or_else(|| invalid_wallet("HTML entity"))?, + _ => return Err(invalid_wallet("HTML entity")), + }; + output.push(decoded); + rest = &rest[end + 1..]; + } + output.push_str(rest); + Ok(output) +} + fn build_result( billing: BillingSnapshot, identity: Option, zerogpu: Option, + balance: Option, ) -> ProviderFetchResult { let mut result = ProviderFetchResult::new( UsageSnapshot::new(RateWindow::informational("Hugging Face billing")) @@ -437,6 +614,9 @@ fn build_result( if let Some(limit) = billing.limit_usd { cost = cost.with_limit(limit); } + if let Some(balance) = balance { + cost = cost.with_balance(balance); + } result = result.with_cost(cost); let mut details: Vec<(&str, &str, String)> = vec![ @@ -462,6 +642,9 @@ fn build_result( if let Some(requests) = billing.requests { details.push(("inference-requests", "Requests", requests.to_string())); } + if let Some(balance) = balance { + details.push(("prepaid-balance", "Prepaid balance", format_usd(balance))); + } let mut rows: Vec> = details .into_iter() @@ -751,6 +934,7 @@ mod tests { total_minutes: 1500.0, resets_at: None, }), + None, ); assert_eq!(result.source_label, "api"); assert_eq!(result.cost.as_ref().and_then(|cost| cost.limit), Some(10.0)); @@ -759,4 +943,21 @@ mod tests { assert!(result.usage.secondary.is_none()); assert!(!result.pace_authoritative); } + + #[test] + fn wallet_parser_prefers_unique_current_balance_and_supports_legacy_cents() { + let current = r#"
"#; + assert_eq!(parse_wallet_balance(current).unwrap(), 12.5); + + let legacy = r#"
"#; + assert_eq!(parse_wallet_balance(legacy).unwrap(), 7.25); + } + + #[test] + fn wallet_parser_rejects_ambiguous_or_non_user_balances() { + let ambiguous = r#"
"#; + assert!(parse_wallet_balance(ambiguous).is_err()); + let organization = r#"
"#; + assert!(parse_wallet_balance(organization).is_err()); + } } diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index ad697ad321..2acbfcf3bd 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -38,6 +38,7 @@ pub mod fireworks; pub mod gemini; pub mod grok; pub mod groq; +pub mod helmcode; pub mod huggingface; pub mod infini; pub mod jetbrains; @@ -74,6 +75,8 @@ pub mod sakana; pub mod stepfun; pub mod sub2api; pub mod t3chat; +pub mod typesafe; +pub mod v0; pub mod venice; pub mod vertexai; pub mod warp; @@ -118,6 +121,7 @@ pub use fireworks::FireworksProvider; pub use gemini::GeminiProvider; pub use grok::GrokProvider; pub use groq::GroqProvider; +pub use helmcode::HelmcodeProvider; pub use huggingface::HuggingFaceProvider; pub use infini::InfiniProvider; pub use jetbrains::JetBrainsProvider; @@ -153,6 +157,8 @@ pub use sakana::SakanaProvider; pub use stepfun::StepFunProvider; pub use sub2api::Sub2ApiProvider; pub use t3chat::T3ChatProvider; +pub use typesafe::TypeSafeProvider; +pub use v0::V0Provider; pub use venice::VeniceProvider; pub use vertexai::VertexAIProvider; pub use warp::WarpProvider; diff --git a/rust/src/providers/muse/mod.rs b/rust/src/providers/muse/mod.rs index 887f3c4c32..940981888f 100644 --- a/rust/src/providers/muse/mod.rs +++ b/rust/src/providers/muse/mod.rs @@ -267,11 +267,29 @@ fn parse_response(body: &[u8]) -> Result { )); } - let usage = object( - root.get("subs_usage") - .ok_or_else(|| parse_failure("missing subs_usage"))?, - "subs_usage", - )?; + let plan = optional_text(root.get("subs_tier_name"), "subs_tier_name")?; + let email = optional_text(root.get("user_email"), "user_email")?; + let Some(raw_usage) = root.get("subs_usage").filter(|value| !value.is_null()) else { + let mut usage = UsageSnapshot::new(RateWindow::informational( + "Subscription active; quota was not included in this login response", + )) + .with_login_method("Muse login"); + if let Some(email) = email { + usage = usage.with_email(email); + } + let mut result = ProviderFetchResult::new(usage, "oauth") + .with_non_authoritative_pace() + .with_display_detail(ProviderDisplayDetail::new( + "quota", + "Quota", + "Not included in this login response", + )); + if let Some(plan) = plan { + result = result.with_display_detail(ProviderDisplayDetail::new("plan", "Plan", plan)); + } + return Ok(result); + }; + let usage = object(raw_usage, "subs_usage")?; let window = object( usage .get("window") @@ -305,9 +323,6 @@ fn parse_response(body: &[u8]) -> Result { )?; let primary_reset = parse_reset(window.get("resets_at"))?; let weekly_reset = parse_reset(weekly.get("resets_at"))?; - let plan = optional_text(root.get("subs_tier_name"), "subs_tier_name")?; - let email = optional_text(root.get("user_email"), "user_email")?; - let primary = RateWindow::with_details( primary_percent.clamp(0.0, 100.0), Some(duration), @@ -603,10 +618,37 @@ mod tests { assert!(result.usage.secondary.as_ref().unwrap().resets_at.is_some()); } + #[test] + fn active_windowless_subscription_preserves_identity_without_inventing_quota() { + let payload = serde_json::json!({ + "require_payment": false, + "is_subs_active": true, + "subs_tier_name": "Pro", + "user_email": "muse@example.com" + }); + let result = parse_response(&serde_json::to_vec(&payload).unwrap()).unwrap(); + + assert_eq!(result.usage.primary.used_percent, 0.0); + assert!(result.usage.primary.window_minutes.is_none()); + assert_eq!( + result.usage.account_email.as_deref(), + Some("muse@example.com") + ); + assert!(!result.pace_authoritative); + assert!(result.display_details().iter().any(|row| { + row.title() == "Quota" && row.value() == "Not included in this login response" + })); + assert!( + result + .display_details() + .iter() + .any(|row| row.title() == "Plan" && row.value() == "Pro") + ); + } + #[test] fn malformed_required_fields_fail_closed_without_echoing_payload() { for payload in [ - serde_json::json!({"is_subs_active": true}), serde_json::json!({"is_subs_active": true, "subs_usage": {"window": {}, "weekly": {}}}), serde_json::json!({"require_payment": "yes", "is_subs_active": true}), serde_json::json!({"require_payment": false, "is_subs_active": true, "subs_usage": {"window": {"window_duration_mins": "300"}, "weekly": {"used_percent": 1}}}), diff --git a/rust/src/providers/typesafe.rs b/rust/src/providers/typesafe.rs new file mode 100644 index 0000000000..cf532218e6 --- /dev/null +++ b/rust/src/providers/typesafe.rs @@ -0,0 +1,443 @@ +//! TypeSafe console billing provider. + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use reqwest::{Client, StatusCode, redirect::Policy}; +use serde_json::Value; +use std::time::Duration; + +use crate::core::{ + CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, + ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, +}; + +const BILLING_URL: &str = "https://console.typesafe.ai/settings/billing"; +const ORIGIN: &str = "https://console.typesafe.ai"; +const MAX_CHUNKS: usize = 60; +const MAX_BODY_BYTES: usize = 2 * 1024 * 1024; + +#[derive(Debug, Clone, PartialEq)] +struct Credit { + amount: f64, + remaining: f64, + expires_at: DateTime, +} + +#[derive(Debug, Clone, PartialEq)] +struct Billing { + spent: f64, + balance: f64, + cycle_label: Option, + plan: Option, + credits: Vec, +} + +pub struct TypeSafeProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl TypeSafeProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::TypeSafe, + display_name: "TypeSafe", + session_label: "Balance", + weekly_label: "Spend", + supports_opus: false, + supports_credits: true, + default_enabled: false, + is_primary: false, + dashboard_url: Some(BILLING_URL), + status_page_url: None, + tertiary_label_key: None, + }, + client: crate::core::credentialed_http_client_builder() + .redirect(Policy::none()) + .timeout(Duration::from_secs(8)) + .build() + .unwrap_or_else(|_| Client::new()), + } + } + + async fn fetch_web(&self, ctx: &FetchContext) -> Result { + let cookie = match ctx.manual_cookie_header.as_deref() { + Some(raw) => normalize_cookie(raw).ok_or_else(|| { + ProviderError::Other( + "TypeSafe needs a nonempty Cookie header from the billing page.".into(), + ) + })?, + None => { + crate::providers::browser_cookie_header(&["console.typesafe.ai", "typesafe.ai"])? + } + }; + let page = self.get(BILLING_URL, &cookie, "text/html").await?; + let action_id = self.discover_action(&cookie, &page).await?; + let response = match self.post_action(&cookie, &action_id).await? { + Some(response) => response, + None => { + let refreshed_page = self.get(BILLING_URL, &cookie, "text/html").await?; + let refreshed_action = self.discover_action(&cookie, &refreshed_page).await?; + self.post_action(&cookie, &refreshed_action) + .await? + .ok_or_else(|| parse_failure("server action stayed stale after rediscovery"))? + } + }; + let billing = parse_rsc_billing(&response)?; + Ok(build_result(billing)) + } + + async fn discover_action(&self, cookie: &str, page: &str) -> Result { + for url in extract_chunk_urls(page).into_iter().take(MAX_CHUNKS) { + let chunk = self.get(&url, &cookie, "application/javascript").await?; + if let Some(found) = find_action_id(&chunk) { + return Ok(found); + } + } + Err(parse_failure("action id not found")) + } + + async fn get(&self, url: &str, cookie: &str, accept: &str) -> Result { + let response = self + .client + .get(url) + .header("Cookie", cookie) + .header("Accept", accept) + .send() + .await?; + read_response(response).await + } + + async fn post_action( + &self, + cookie: &str, + action_id: &str, + ) -> Result, ProviderError> { + let response = self + .client + .post(BILLING_URL) + .header("Cookie", cookie) + .header("Origin", ORIGIN) + .header("Next-Action", action_id) + .header("Accept", "text/x-component") + .header("Content-Type", "application/json") + .body("[]") + .send() + .await?; + if response.status() == StatusCode::NOT_FOUND + && response + .headers() + .get("x-nextjs-action-not-found") + .and_then(|value| value.to_str().ok()) + == Some("1") + { + return Ok(None); + } + read_response(response).await.map(Some) + } +} + +impl Default for TypeSafeProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for TypeSafeProvider { + fn id(&self) -> ProviderId { + ProviderId::TypeSafe + } + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::Web => self.fetch_web(ctx).await, + SourceMode::OAuth | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::Web] + } + fn supports_web(&self) -> bool { + true + } + fn owns_browser_cookie_resolution(&self) -> bool { + true + } +} + +async fn read_response(response: reqwest::Response) -> Result { + let status = response.status(); + if status == StatusCode::UNAUTHORIZED + || status == StatusCode::FORBIDDEN + || status.is_redirection() + { + return Err(ProviderError::AuthRequired); + } + if status == StatusCode::TOO_MANY_REQUESTS { + return Err(ProviderError::Other("TypeSafe rate limit reached.".into())); + } + if status == StatusCode::REQUEST_TIMEOUT || status.is_server_error() { + return Err(ProviderError::Other( + "TypeSafe billing is unavailable.".into(), + )); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "TypeSafe returned HTTP {status}." + ))); + } + let bytes = response.bytes().await?; + if bytes.len() > MAX_BODY_BYTES { + return Err(parse_failure("response too large")); + } + let body = + String::from_utf8(bytes.to_vec()).map_err(|_| parse_failure("response was not UTF-8"))?; + if body.contains("\\\"(auth)\\\",{\\\"children\\\":[\\\"login\\\"") { + return Err(ProviderError::AuthRequired); + } + Ok(body) +} + +fn extract_chunk_urls(html: &str) -> Vec { + let mut urls = Vec::new(); + let mut rest = html; + while let Some(index) = rest.find("src=") { + rest = &rest[index + 4..]; + let Some(quote) = rest + .chars() + .next() + .filter(|value| matches!(value, '\'' | '"')) + else { + continue; + }; + rest = &rest[quote.len_utf8()..]; + let Some(end) = rest.find(quote) else { + break; + }; + let source = &rest[..end]; + rest = &rest[end + quote.len_utf8()..]; + let normalized = if source.starts_with('/') { + format!("{ORIGIN}{source}") + } else { + source.to_string() + }; + if normalized.starts_with(&format!("{ORIGIN}/")) + && normalized + .split('?') + .next() + .is_some_and(|path| path.ends_with(".js")) + && !urls.contains(&normalized) + { + urls.push(normalized); + } + if urls.len() >= MAX_CHUNKS { + break; + } + } + urls +} + +fn find_action_id(chunk: &str) -> Option { + let marker = chunk.find("getBillingOverviewResult")?; + let prefix = &chunk[marker.saturating_sub(200)..marker]; + prefix.split('"').rev().find_map(|candidate| { + (candidate.len() >= 40 && candidate.chars().all(|ch| ch.is_ascii_hexdigit())) + .then(|| candidate.to_string()) + }) +} + +fn parse_rsc_billing(body: &str) -> Result { + let result = body + .lines() + .find_map(|line| { + let (_, json) = line.split_once(':')?; + let value: Value = serde_json::from_str(json).ok()?; + value.as_object()?.contains_key("ok").then_some(value) + }) + .ok_or_else(|| parse_failure("missing result"))?; + if result.get("ok").and_then(Value::as_bool) != Some(true) { + return Err(ProviderError::Other( + "TypeSafe billing request failed.".into(), + )); + } + let billing = result + .get("data") + .and_then(|value| value.get("billing")) + .and_then(Value::as_object) + .ok_or_else(|| parse_failure("missing billing"))?; + let spent = finite(billing.get("spent"), "spent")?; + let balance = finite(billing.get("balance"), "balance")?; + let cycle_label = clean_text(billing.get("cycleLabel")); + let plan = clean_text(billing.get("plan")); + let credits = billing + .get("credits") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|item| { + let object = item.as_object()?; + let amount = finite(object.get("amount"), "credit amount").ok()?; + let remaining = finite(object.get("remaining"), "credit remaining").ok()?; + if remaining <= 0.0 { + return None; + } + let expires_at = object + .get("expiresAt") + .and_then(Value::as_str) + .and_then(|value| DateTime::parse_from_rfc3339(value).ok())? + .with_timezone(&Utc); + Some(Credit { + amount, + remaining, + expires_at, + }) + }) + .collect(); + Ok(Billing { + spent, + balance, + cycle_label, + plan, + credits, + }) +} + +fn build_result(billing: Billing) -> ProviderFetchResult { + let mut usage = UsageSnapshot::new(RateWindow::informational(format!( + "Balance ${:.2}", + billing.balance + ))) + .with_login_method(format!("Balance ${:.2}", billing.balance)); + usage.updated_at = Utc::now(); + let mut cost = CostSnapshot::new( + billing.spent, + "USD", + billing + .cycle_label + .clone() + .unwrap_or_else(|| "Billing cycle".into()), + ) + .with_balance(billing.balance) + .always_visible(); + cost.updated_at = usage.updated_at; + let spent_title = billing + .cycle_label + .as_deref() + .map_or_else(|| "Spent".into(), |cycle| format!("Spent ({cycle})")); + let mut result = ProviderFetchResult::new(usage, "web") + .with_non_authoritative_pace() + .with_cost(cost) + .with_display_detail(ProviderDisplayDetail::new( + "spent", + spent_title, + format!("${:.2}", billing.spent), + )); + if let Some(plan) = billing.plan.as_deref() { + let label = if plan == "free_plan" { + "Free".into() + } else { + title_case(plan) + }; + result = result.with_display_detail(ProviderDisplayDetail::new("plan", "Plan", label)); + } + let available = 24usize.saturating_sub(result.display_details().len()); + let visible = billing + .credits + .iter() + .take(available.saturating_sub(1)) + .collect::>(); + for (index, credit) in visible.iter().enumerate() { + result = result.with_display_detail(ProviderDisplayDetail::new( + format!("credit-{index}"), + "Credit", + format!( + "{:.2} of {:.2}, expires {}", + credit.remaining, + credit.amount, + credit.expires_at.format("%b %d") + ), + )); + } + if billing.credits.len() > visible.len() { + result = result.with_display_detail(ProviderDisplayDetail::new( + "additional-credits", + "Additional credits", + (billing.credits.len() - visible.len()).to_string(), + )); + } + result +} + +fn normalize_cookie(raw: &str) -> Option { + let value = raw + .trim() + .strip_prefix("Cookie:") + .unwrap_or(raw.trim()) + .trim(); + (!value.is_empty() && !value.chars().any(char::is_control)).then(|| value.to_string()) +} +fn finite(value: Option<&Value>, field: &str) -> Result { + value + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) + .ok_or_else(|| parse_failure(field)) +} +fn clean_text(value: Option<&Value>) -> Option { + value + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) +} +fn title_case(value: &str) -> String { + value + .split(['_', '-']) + .filter(|part| !part.is_empty()) + .map(|part| { + let mut chars = part.chars(); + chars.next().map_or_else(String::new, |first| { + first.to_uppercase().collect::() + &chars.as_str().to_ascii_lowercase() + }) + }) + .collect::>() + .join(" ") +} +fn parse_failure(field: impl AsRef) -> ProviderError { + ProviderError::Parse(format!( + "TypeSafe billing response format changed: {}.", + field.as_ref() + )) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn discovers_only_same_origin_javascript_and_action_id() { + let urls = extract_chunk_urls( + r#""#, + ); + assert_eq!(urls, vec!["https://console.typesafe.ai/_next/a.js"]); + let id = "a".repeat(40); + assert_eq!( + find_action_id(&format!(r#"x("{id}")xxx"getBillingOverviewResult""#)).as_deref(), + Some(id.as_str()) + ); + } + + #[test] + fn parses_billing_result_and_skips_expired_or_empty_credits() { + let body = r#"1:{"ok":true,"data":{"billing":{"spent":4.5,"balance":10,"cycleLabel":"September","plan":"free_plan","credits":[{"amount":8,"remaining":3,"expiresAt":"2030-01-02T00:00:00Z"},{"amount":1,"remaining":0,"expiresAt":"2030-01-02T00:00:00Z"}]}}}"#; + let parsed = parse_rsc_billing(body).unwrap(); + assert_eq!(parsed.spent, 4.5); + assert_eq!(parsed.balance, 10.0); + assert_eq!(parsed.credits.len(), 1); + } +} diff --git a/rust/src/providers/v0.rs b/rust/src/providers/v0.rs new file mode 100644 index 0000000000..7685230f60 --- /dev/null +++ b/rust/src/providers/v0.rs @@ -0,0 +1,367 @@ +//! v0 billing and API rate-limit provider. + +use async_trait::async_trait; +use chrono::{TimeZone, Utc}; +use reqwest::{Client, StatusCode, Url}; +use serde_json::Value; +use std::time::Duration; + +use crate::core::{ + FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, + ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, +}; + +const API_BASE: &str = "https://api.v0.dev/v1"; +const CREDENTIAL_TARGET: &str = "codexbar-v0"; +const ENV_KEYS: &[&str] = &["V0_API_KEY"]; + +#[derive(Debug, Clone, PartialEq)] +struct Quota { + used_percent: Option, + resets_at: Option>, + remaining: Option, + limit: f64, +} + +#[derive(Debug, Clone, PartialEq)] +struct Billing { + quota: Quota, + billing_type: String, + on_demand_balance: Option, +} + +pub struct V0Provider { + metadata: ProviderMetadata, + client: Client, +} + +impl V0Provider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::V0, + display_name: "v0", + session_label: "Billing", + weekly_label: "Rate limit", + supports_opus: false, + supports_credits: true, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://v0.app/chat/settings/billing"), + status_page_url: Some("https://www.vercel-status.com/"), + tertiary_label_key: None, + }, + client: crate::core::credentialed_http_client_builder() + .timeout(Duration::from_secs(15)) + .build() + .unwrap_or_else(|_| Client::new()), + } + } + + async fn fetch_api(&self, ctx: &FetchContext) -> Result { + let api_key = + crate::providers::resolve_api_key(ctx.api_key.as_deref(), CREDENTIAL_TARGET, ENV_KEYS)?; + let scope = ctx + .workspace_id + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string) + .or_else(|| { + std::env::var("V0_SCOPE") + .ok() + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + }); + let billing = parse_billing( + &self + .get_json("/user/billing", scope.as_deref(), &api_key) + .await?, + )?; + let rate_limit = parse_quota( + &self + .get_json("/rate-limits", scope.as_deref(), &api_key) + .await?, + "rate limit response", + )?; + Ok(build_result(billing, rate_limit, scope.as_deref())) + } + + async fn get_json( + &self, + path: &str, + scope: Option<&str>, + api_key: &str, + ) -> Result { + let mut url = Url::parse(&format!("{API_BASE}{path}")) + .map_err(|_| ProviderError::Other("Invalid v0 API URL.".into()))?; + if let Some(scope) = scope { + url.query_pairs_mut().append_pair("scope", scope); + } + let response = self.client.get(url).bearer_auth(api_key).send().await?; + classify_status(response.status())?; + response.json().await.map_err(|_| { + ProviderError::Parse(format!( + "Could not parse v0 usage: {path} returned invalid JSON" + )) + }) + } +} + +impl Default for V0Provider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for V0Provider { + fn id(&self) -> ProviderId { + ProviderId::V0 + } + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::OAuth => self.fetch_api(ctx).await, + SourceMode::Web | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::OAuth] + } +} + +fn classify_status(status: StatusCode) -> Result<(), ProviderError> { + match status { + status if status.is_success() => Ok(()), + StatusCode::UNAUTHORIZED => Err(ProviderError::AuthRequired), + StatusCode::FORBIDDEN => Err(ProviderError::Other( + "v0 denied access to this scope.".into(), + )), + StatusCode::TOO_MANY_REQUESTS => { + Err(ProviderError::Other("v0 API rate limit reached.".into())) + } + status if status.is_server_error() => Err(ProviderError::Other(format!( + "v0 API returned HTTP {status}." + ))), + status => Err(ProviderError::Other(format!( + "v0 API returned HTTP {status}." + ))), + } +} + +fn parse_billing(value: &Value) -> Result { + let object = value + .as_object() + .ok_or_else(|| parse_failure("billing response"))?; + let billing_type = text(object.get("billingType"), "billingType")? + .ok_or_else(|| parse_failure("billingType"))?; + let data = object + .get("data") + .ok_or_else(|| parse_failure("billing.data"))?; + match billing_type.as_str() { + "legacy" => Ok(Billing { + quota: parse_quota(data, "billing.data")?, + billing_type, + on_demand_balance: None, + }), + "token" => { + let data = data + .as_object() + .ok_or_else(|| parse_failure("billing.data"))?; + let balance = data + .get("balance") + .and_then(Value::as_object) + .ok_or_else(|| parse_failure("billing.data.balance"))?; + let total = finite_number(balance.get("total"), "billing.data.balance.total")?; + let remaining = + finite_number(balance.get("remaining"), "billing.data.balance.remaining")?; + if total < 0.0 { + return Err(parse_failure( + "billing.data.balance.total must not be negative", + )); + } + let resets_at = match data.get("billingCycle").and_then(Value::as_object) { + Some(cycle) => parse_reset(cycle.get("end"))?, + None => None, + }; + let on_demand_balance = data + .get("onDemand") + .filter(|value| !value.is_null()) + .and_then(Value::as_object) + .map(|on_demand| { + finite_number(on_demand.get("balance"), "billing.data.onDemand.balance") + }) + .transpose()?; + Ok(Billing { + quota: Quota { + used_percent: percent(total - remaining, total), + resets_at, + remaining: Some(remaining), + limit: total, + }, + billing_type, + on_demand_balance, + }) + } + _ => Err(parse_failure("billingType")), + } +} + +fn parse_quota(value: &Value, field: &str) -> Result { + let object = value.as_object().ok_or_else(|| parse_failure(field))?; + let limit = finite_number(object.get("limit"), &format!("{field}.limit"))?; + if limit < 0.0 { + return Err(parse_failure(format!("{field}.limit must not be negative"))); + } + let remaining = optional_number(object.get("remaining"), &format!("{field}.remaining"))?; + Ok(Quota { + used_percent: remaining.and_then(|remaining| percent(limit - remaining, limit)), + resets_at: parse_reset(object.get("reset"))?, + remaining, + limit, + }) +} + +fn build_result(billing: Billing, rate_limit: Quota, scope: Option<&str>) -> ProviderFetchResult { + let primary = billing.quota.used_percent.map_or_else( + || RateWindow::informational("Billing usage unavailable"), + |used| RateWindow::with_details(used, None, billing.quota.resets_at, None), + ); + let mut usage = UsageSnapshot::new(primary).with_login_method("API key"); + if let Some(used) = rate_limit.used_percent { + usage = usage.with_secondary(RateWindow::with_details( + used, + None, + rate_limit.resets_at, + None, + )); + } + let mut result = ProviderFetchResult::new(usage, "api"); + let remaining = billing + .quota + .remaining + .map_or_else(|| "Unavailable".into(), format_number); + result = result.with_display_detail(ProviderDisplayDetail::new( + "billing-remaining", + "Billing remaining", + format!("{remaining} of {}", format_number(billing.quota.limit)), + )); + if let Some(balance) = billing.on_demand_balance { + result = result.with_display_detail(ProviderDisplayDetail::new( + "on-demand", + "On-demand balance", + format_number(balance), + )); + } + let rate_remaining = rate_limit + .remaining + .map_or_else(|| "Unavailable".into(), format_number); + result = result.with_display_detail(ProviderDisplayDetail::new( + "rate-limit-remaining", + "Rate-limit remaining", + format!("{rate_remaining} of {}", format_number(rate_limit.limit)), + )); + result = result.with_display_detail(ProviderDisplayDetail::new( + "billing-type", + "Billing type", + billing.billing_type, + )); + if let Some(scope) = scope { + result = result.with_display_detail(ProviderDisplayDetail::new( + "scope", + "Scope", + scope.chars().take(120).collect::(), + )); + } + result +} + +fn finite_number(value: Option<&Value>, field: &str) -> Result { + value + .and_then(Value::as_f64) + .filter(|value| value.is_finite()) + .ok_or_else(|| parse_failure(field)) +} + +fn optional_number(value: Option<&Value>, field: &str) -> Result, ProviderError> { + match value { + None | Some(Value::Null) => Ok(None), + Some(value) => finite_number(Some(value), field).map(Some), + } +} + +fn text(value: Option<&Value>, field: &str) -> Result, ProviderError> { + match value { + None | Some(Value::Null) => Ok(None), + Some(Value::String(value)) => { + Ok((!value.trim().is_empty()).then(|| value.trim().to_string())) + } + Some(_) => Err(parse_failure(field)), + } +} + +fn parse_reset(value: Option<&Value>) -> Result>, ProviderError> { + let Some(value) = value.filter(|value| !value.is_null()) else { + return Ok(None); + }; + let raw = finite_number(Some(value), "reset")?; + if raw <= 0.0 { + return Ok(None); + } + let seconds = if raw >= 1_000_000_000_000.0 { + raw / 1000.0 + } else { + raw + }; + Ok(Utc.timestamp_opt(seconds.trunc() as i64, 0).single()) +} + +fn percent(used: f64, limit: f64) -> Option { + (limit > 0.0 && used.is_finite()).then(|| (used.max(0.0) / limit * 100.0).clamp(0.0, 100.0)) +} + +fn format_number(value: f64) -> String { + format!("{value:.2}") + .trim_end_matches('0') + .trim_end_matches('.') + .to_string() +} +fn parse_failure(field: impl AsRef) -> ProviderError { + ProviderError::Parse(format!("Could not parse v0 usage: {}", field.as_ref())) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn parses_token_billing_and_rate_limit_without_fabricating_windows() { + let billing = parse_billing(&json!({ + "billingType": "token", + "data": {"balance": {"total": 100, "remaining": 75}, "billingCycle": {"end": 1_800_000_000}, "onDemand": {"balance": 12.5}} + })).unwrap(); + assert_eq!(billing.quota.used_percent, Some(25.0)); + assert_eq!(billing.on_demand_balance, Some(12.5)); + let unknown = parse_quota( + &json!({"limit": 50, "remaining": null, "reset": null}), + "rate", + ) + .unwrap(); + assert_eq!(unknown.used_percent, None); + assert_eq!(unknown.remaining, None); + } + + #[test] + fn rejects_unknown_billing_type_and_negative_limits() { + assert!(parse_billing(&json!({"billingType": "future", "data": {}})).is_err()); + assert!(parse_quota(&json!({"limit": -1}), "rate").is_err()); + } +} diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index 3a7b166386..c2adf3565c 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -363,6 +363,17 @@ pub fn get_api_key_providers() -> Vec { config_file_path: None, dashboard_url: Some("https://huggingface.co/settings/billing"), }, + ProviderConfigInfo { + id: ProviderId::V0, + name: "v0", + requires_api_key: true, + api_key_env_var: Some("V0_API_KEY"), + api_key_help: Some( + "Add a v0 Platform API key. An optional scope can use the provider workspace field or V0_SCOPE.", + ), + config_file_path: None, + dashboard_url: Some("https://v0.app/chat/settings/billing"), + }, ProviderConfigInfo { id: ProviderId::Fireworks, name: "Fireworks", diff --git a/rust/src/settings/provider_workspace.rs b/rust/src/settings/provider_workspace.rs index 913670510e..338747c8ef 100644 --- a/rust/src/settings/provider_workspace.rs +++ b/rust/src/settings/provider_workspace.rs @@ -40,6 +40,21 @@ pub fn validate_provider_workspace_value( } Ok(trimmed.to_string()) } + ProviderId::V0 => validate_id(trimmed, "v0 scope", |value| { + value.len() <= 128 + && value + .chars() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, '_' | '-')) + }), + ProviderId::Helmcode => { + if trimmed.eq_ignore_ascii_case("helmcode") { + Ok("helmcode".to_string()) + } else if trimmed.eq_ignore_ascii_case("nanBuilders") { + Ok("nanBuilders".to_string()) + } else { + Err("Helmcode tenant must be 'helmcode' or 'nanBuilders'".to_string()) + } + } ProviderId::LiteLLM => validate_token_endpoint(trimmed, "LiteLLM base URL", |_| true), ProviderId::Sub2Api => validate_sub2api_base_url(trimmed), _ => Ok(trimmed.to_string()), From ec068d212c7e3c92c909804f840bd009d89ba3de Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:16:35 +0700 Subject: [PATCH 02/11] Fix provider conversion lint errors --- rust/src/providers/helmcode.rs | 16 ++++++++++++---- rust/src/providers/typesafe.rs | 2 +- rust/src/providers/v0.rs | 6 +++++- 3 files changed, 18 insertions(+), 6 deletions(-) diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs index 02ac002f2a..04dc5073d7 100644 --- a/rust/src/providers/helmcode.rs +++ b/rust/src/providers/helmcode.rs @@ -263,12 +263,11 @@ fn parse_models(quota: &Value, premium: bool) -> Result, Provide .ok_or_else(|| parse_failure("periodStart"))?; let fallback = DateTime::parse_from_rfc3339(period_start) .ok() - .map(|date| { + .and_then(|date| { let date = date.with_timezone(&Utc); Utc.with_ymd_and_hms(date.year(), date.month(), 1, 0, 0, 0) .single() - }) - .flatten(); + }); let models = object .get("models") .and_then(Value::as_array) @@ -293,7 +292,9 @@ fn parse_models(quota: &Value, premium: bool) -> Result, Provide let credit = optional_nonnegative(row.get("creditTokens"), "creditTokens")?.unwrap_or(0.0); let window_hours = optional_nonnegative(row.get("windowHours"), "windowHours")? - .map(|value| value as u32); + .map(|value| format!("{value:.0}").parse::()) + .transpose() + .map_err(|_| parse_failure("windowHours"))?; if window_hours.is_some_and(|hours| hours == 0 || hours > 8_760) { return Err(parse_failure("windowHours")); } @@ -392,5 +393,12 @@ mod tests { fn rejects_fractional_or_negative_quota_counts() { assert!(parse_models(&json!({"periodStart":"2030-01-01T00:00:00Z","models":[{"model":"x","cap":1.5,"tokensUsed":0}]}), true).is_err()); assert!(parse_models(&json!({"periodStart":"2030-01-01T00:00:00Z","models":[{"model":"x","cap":1,"tokensUsed":-1}]}), true).is_err()); + assert!( + parse_models( + &json!({"periodStart":"2030-01-01T00:00:00Z","models":[{"model":"x","cap":1,"tokensUsed":0,"windowHours":4_294_967_296_u64}]}), + true, + ) + .is_err() + ); } } diff --git a/rust/src/providers/typesafe.rs b/rust/src/providers/typesafe.rs index cf532218e6..bcddac39f3 100644 --- a/rust/src/providers/typesafe.rs +++ b/rust/src/providers/typesafe.rs @@ -90,7 +90,7 @@ impl TypeSafeProvider { async fn discover_action(&self, cookie: &str, page: &str) -> Result { for url in extract_chunk_urls(page).into_iter().take(MAX_CHUNKS) { - let chunk = self.get(&url, &cookie, "application/javascript").await?; + let chunk = self.get(&url, cookie, "application/javascript").await?; if let Some(found) = find_action_id(&chunk) { return Ok(found); } diff --git a/rust/src/providers/v0.rs b/rust/src/providers/v0.rs index 7685230f60..9ce848fdfa 100644 --- a/rust/src/providers/v0.rs +++ b/rust/src/providers/v0.rs @@ -320,7 +320,10 @@ fn parse_reset(value: Option<&Value>) -> Result>, P } else { raw }; - Ok(Utc.timestamp_opt(seconds.trunc() as i64, 0).single()) + let seconds = format!("{:.0}", seconds.trunc()) + .parse::() + .map_err(|_| parse_failure("reset"))?; + Ok(Utc.timestamp_opt(seconds, 0).single()) } fn percent(used: f64, limit: f64) -> Option { @@ -363,5 +366,6 @@ mod tests { fn rejects_unknown_billing_type_and_negative_limits() { assert!(parse_billing(&json!({"billingType": "future", "data": {}})).is_err()); assert!(parse_quota(&json!({"limit": -1}), "rate").is_err()); + assert!(parse_reset(Some(&json!(1e30))).is_err()); } } From 4a8daf5ecb9b906ae9f5bab6ff1c60e31819a717 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:56:00 +0700 Subject: [PATCH 03/11] Fix provider response boundaries --- rust/src/providers/huggingface/mod.rs | 86 ++++++++--------- rust/src/providers/typesafe.rs | 128 +++++++++++++++++--------- rust/src/providers/v0.rs | 80 +++++++++++++--- 3 files changed, 199 insertions(+), 95 deletions(-) diff --git a/rust/src/providers/huggingface/mod.rs b/rust/src/providers/huggingface/mod.rs index 857bbc0c26..6bef977198 100644 --- a/rust/src/providers/huggingface/mod.rs +++ b/rust/src/providers/huggingface/mod.rs @@ -187,19 +187,17 @@ impl HuggingFaceProvider { async fn fetch_optional_wallet(&self, expected_user_id: &str) -> Option { let cookie = crate::providers::browser_cookie_header(&["huggingface.co"]).ok()?; - let billing = self - .fetch_cookie_text( + let (billing, whoami) = tokio::join!( + self.fetch_cookie_text( "https://huggingface.co/settings/billing", &cookie, "text/html", - ) - .await - .ok()?; + ), + self.fetch_cookie_text(WHOAMI_URL, &cookie, "application/json"), + ); + let billing = billing.ok()?; + let whoami = whoami.ok()?; let candidate = parse_wallet_balance(&billing).ok()?; - let whoami = self - .fetch_cookie_text(WHOAMI_URL, &cookie, "application/json") - .await - .ok()?; let profile: Value = serde_json::from_str(&whoami).ok()?; let observed_user_id = profile .get("type") @@ -218,28 +216,24 @@ impl HuggingFaceProvider { cookie: &str, accept: &str, ) -> Result { - let response = tokio::time::timeout( - OPTIONAL_TIMEOUT, - self.client + tokio::time::timeout(OPTIONAL_TIMEOUT, async { + let response = self + .client .get(url) .header(reqwest::header::COOKIE, cookie) .header(reqwest::header::ACCEPT, accept) - .send(), - ) - .await - .map_err(|_| ProviderError::Timeout)??; - if !response.status().is_success() { - return Err(classify_status(response.status())); - } - let bytes = response.bytes().await?; - if bytes.len() > MAX_RESPONSE_BYTES { - return Err(ProviderError::Parse( - "Hugging Face returned an oversized wallet response.".to_string(), - )); - } - String::from_utf8(bytes.to_vec()).map_err(|_| { - ProviderError::Parse("Hugging Face returned invalid wallet text.".to_string()) + .send() + .await?; + if !response.status().is_success() { + return Err(classify_status(response.status())); + } + let body = read_bounded_body(response, "wallet response").await?; + String::from_utf8(body).map_err(|_| { + ProviderError::Parse("Hugging Face returned invalid wallet text.".to_string()) + }) }) + .await + .map_err(|_| ProviderError::Timeout)? } async fn fetch_optional_json(&self, url: Url, token: &str) -> Option { @@ -269,21 +263,7 @@ impl HuggingFaceProvider { return Err(classify_status(status)); } - let mut body = Vec::new(); - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(|_| { - ProviderError::Parse( - "Hugging Face returned an unreadable JSON body.".to_string(), - ) - })?; - body.extend_from_slice(&chunk); - if body.len() > MAX_RESPONSE_BYTES { - return Err(ProviderError::Parse( - "Hugging Face returned an oversized JSON body.".to_string(), - )); - } - } + let body = read_bounded_body(response, "JSON body").await?; serde_json::from_slice(&body).map_err(|_| { ProviderError::Parse("Hugging Face returned invalid JSON.".to_string()) }) @@ -293,6 +273,28 @@ impl HuggingFaceProvider { } } +async fn read_bounded_body( + response: reqwest::Response, + response_kind: &str, +) -> Result, ProviderError> { + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(|_| { + ProviderError::Parse(format!( + "Hugging Face returned an unreadable {response_kind}." + )) + })?; + if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(body.len()) { + return Err(ProviderError::Parse(format!( + "Hugging Face returned an oversized {response_kind}." + ))); + } + body.extend_from_slice(&chunk); + } + Ok(body) +} + impl Default for HuggingFaceProvider { fn default() -> Self { Self::new() diff --git a/rust/src/providers/typesafe.rs b/rust/src/providers/typesafe.rs index bcddac39f3..59dd26ecd1 100644 --- a/rust/src/providers/typesafe.rs +++ b/rust/src/providers/typesafe.rs @@ -2,6 +2,7 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; +use futures::{StreamExt, stream}; use reqwest::{Client, StatusCode, redirect::Policy}; use serde_json::Value; use std::time::Duration; @@ -14,7 +15,9 @@ use crate::core::{ const BILLING_URL: &str = "https://console.typesafe.ai/settings/billing"; const ORIGIN: &str = "https://console.typesafe.ai"; const MAX_CHUNKS: usize = 60; +const CHUNK_SCAN_CONCURRENCY: usize = 6; const MAX_BODY_BYTES: usize = 2 * 1024 * 1024; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(8); #[derive(Debug, Clone, PartialEq)] struct Credit { @@ -55,7 +58,7 @@ impl TypeSafeProvider { }, client: crate::core::credentialed_http_client_builder() .redirect(Policy::none()) - .timeout(Duration::from_secs(8)) + .timeout(REQUEST_TIMEOUT) .build() .unwrap_or_else(|_| Client::new()), } @@ -89,24 +92,40 @@ impl TypeSafeProvider { } async fn discover_action(&self, cookie: &str, page: &str) -> Result { - for url in extract_chunk_urls(page).into_iter().take(MAX_CHUNKS) { - let chunk = self.get(&url, cookie, "application/javascript").await?; - if let Some(found) = find_action_id(&chunk) { - return Ok(found); + let mut chunks = stream::iter(extract_chunk_urls(page)) + .map(|url| async move { self.get(&url, cookie, "application/javascript").await }) + .buffer_unordered(CHUNK_SCAN_CONCURRENCY); + let mut first_error = None; + while let Some(result) = chunks.next().await { + match result { + Ok(chunk) => { + if let Some(found) = find_action_id(&chunk) { + return Ok(found); + } + } + Err(error) if first_error.is_none() => first_error = Some(error), + Err(_) => {} } } + if let Some(error) = first_error { + return Err(error); + } Err(parse_failure("action id not found")) } async fn get(&self, url: &str, cookie: &str, accept: &str) -> Result { - let response = self - .client - .get(url) - .header("Cookie", cookie) - .header("Accept", accept) - .send() - .await?; - read_response(response).await + tokio::time::timeout(REQUEST_TIMEOUT, async { + let response = self + .client + .get(url) + .header("Cookie", cookie) + .header("Accept", accept) + .send() + .await?; + read_response(response).await + }) + .await + .map_err(|_| ProviderError::Timeout)? } async fn post_action( @@ -114,27 +133,31 @@ impl TypeSafeProvider { cookie: &str, action_id: &str, ) -> Result, ProviderError> { - let response = self - .client - .post(BILLING_URL) - .header("Cookie", cookie) - .header("Origin", ORIGIN) - .header("Next-Action", action_id) - .header("Accept", "text/x-component") - .header("Content-Type", "application/json") - .body("[]") - .send() - .await?; - if response.status() == StatusCode::NOT_FOUND - && response - .headers() - .get("x-nextjs-action-not-found") - .and_then(|value| value.to_str().ok()) - == Some("1") - { - return Ok(None); - } - read_response(response).await.map(Some) + tokio::time::timeout(REQUEST_TIMEOUT, async { + let response = self + .client + .post(BILLING_URL) + .header("Cookie", cookie) + .header("Origin", ORIGIN) + .header("Next-Action", action_id) + .header("Accept", "text/x-component") + .header("Content-Type", "application/json") + .body("[]") + .send() + .await?; + if response.status() == StatusCode::NOT_FOUND + && response + .headers() + .get("x-nextjs-action-not-found") + .and_then(|value| value.to_str().ok()) + == Some("1") + { + return Ok(None); + } + read_response(response).await.map(Some) + }) + .await + .map_err(|_| ProviderError::Timeout)? } } @@ -192,12 +215,16 @@ async fn read_response(response: reqwest::Response) -> Result MAX_BODY_BYTES { - return Err(parse_failure("response too large")); + let mut body = Vec::new(); + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk?; + if chunk.len() > MAX_BODY_BYTES.saturating_sub(body.len()) { + return Err(parse_failure("response too large")); + } + body.extend_from_slice(&chunk); } - let body = - String::from_utf8(bytes.to_vec()).map_err(|_| parse_failure("response was not UTF-8"))?; + let body = String::from_utf8(body).map_err(|_| parse_failure("response was not UTF-8"))?; if body.contains("\\\"(auth)\\\",{\\\"children\\\":[\\\"login\\\"") { return Err(ProviderError::AuthRequired); } @@ -245,7 +272,10 @@ fn extract_chunk_urls(html: &str) -> Vec { fn find_action_id(chunk: &str) -> Option { let marker = chunk.find("getBillingOverviewResult")?; - let prefix = &chunk[marker.saturating_sub(200)..marker]; + let start = (marker.saturating_sub(200)..=marker) + .find(|index| chunk.is_char_boundary(*index)) + .unwrap_or(marker); + let prefix = &chunk[start..marker]; prefix.split('"').rev().find_map(|candidate| { (candidate.len() >= 40 && candidate.chars().all(|ch| ch.is_ascii_hexdigit())) .then(|| candidate.to_string()) @@ -292,6 +322,9 @@ fn parse_rsc_billing(body: &str) -> Result { .and_then(Value::as_str) .and_then(|value| DateTime::parse_from_rfc3339(value).ok())? .with_timezone(&Utc); + if expires_at <= Utc::now() { + return None; + } Some(Credit { amount, remaining, @@ -432,9 +465,22 @@ mod tests { ); } + #[test] + fn action_discovery_handles_multibyte_text_at_scan_boundary() { + let id = "b".repeat(40); + let chunk = format!( + "{}\u{00e9}{}x(\"{id}\")getBillingOverviewResult", + "x".repeat(10), + "x".repeat(154) + ); + let marker = chunk.find("getBillingOverviewResult").unwrap(); + assert!(!chunk.is_char_boundary(marker - 200)); + assert_eq!(find_action_id(&chunk).as_deref(), Some(id.as_str())); + } + #[test] fn parses_billing_result_and_skips_expired_or_empty_credits() { - let body = r#"1:{"ok":true,"data":{"billing":{"spent":4.5,"balance":10,"cycleLabel":"September","plan":"free_plan","credits":[{"amount":8,"remaining":3,"expiresAt":"2030-01-02T00:00:00Z"},{"amount":1,"remaining":0,"expiresAt":"2030-01-02T00:00:00Z"}]}}}"#; + let body = r#"1:{"ok":true,"data":{"billing":{"spent":4.5,"balance":10,"cycleLabel":"September","plan":"free_plan","credits":[{"amount":8,"remaining":3,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":1,"remaining":0,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":5,"remaining":2,"expiresAt":"2000-01-02T00:00:00Z"}]}}}"#; let parsed = parse_rsc_billing(body).unwrap(); assert_eq!(parsed.spent, 4.5); assert_eq!(parsed.balance, 10.0); diff --git a/rust/src/providers/v0.rs b/rust/src/providers/v0.rs index 9ce848fdfa..7b18b76926 100644 --- a/rust/src/providers/v0.rs +++ b/rust/src/providers/v0.rs @@ -181,7 +181,7 @@ fn parse_billing(value: &Value) -> Result { .ok_or_else(|| parse_failure("billing.data.balance"))?; let total = finite_number(balance.get("total"), "billing.data.balance.total")?; let remaining = - finite_number(balance.get("remaining"), "billing.data.balance.remaining")?; + nonnegative_number(balance.get("remaining"), "billing.data.balance.remaining")?; if total < 0.0 { return Err(parse_failure( "billing.data.balance.total must not be negative", @@ -196,7 +196,7 @@ fn parse_billing(value: &Value) -> Result { .filter(|value| !value.is_null()) .and_then(Value::as_object) .map(|on_demand| { - finite_number(on_demand.get("balance"), "billing.data.onDemand.balance") + nonnegative_number(on_demand.get("balance"), "billing.data.onDemand.balance") }) .transpose()?; Ok(Billing { @@ -220,7 +220,8 @@ fn parse_quota(value: &Value, field: &str) -> Result { if limit < 0.0 { return Err(parse_failure(format!("{field}.limit must not be negative"))); } - let remaining = optional_number(object.get("remaining"), &format!("{field}.remaining"))?; + let remaining = + optional_nonnegative_number(object.get("remaining"), &format!("{field}.remaining"))?; Ok(Quota { used_percent: remaining.and_then(|remaining| percent(limit - remaining, limit)), resets_at: parse_reset(object.get("reset"))?, @@ -290,10 +291,21 @@ fn finite_number(value: Option<&Value>, field: &str) -> Result, field: &str) -> Result, ProviderError> { +fn nonnegative_number(value: Option<&Value>, field: &str) -> Result { + let value = finite_number(value, field)?; + if value < 0.0 { + return Err(parse_failure(format!("{field} must not be negative"))); + } + Ok(value) +} + +fn optional_nonnegative_number( + value: Option<&Value>, + field: &str, +) -> Result, ProviderError> { match value { None | Some(Value::Null) => Ok(None), - Some(value) => finite_number(Some(value), field).map(Some), + Some(value) => nonnegative_number(Some(value), field).map(Some), } } @@ -315,15 +327,20 @@ fn parse_reset(value: Option<&Value>) -> Result>, P if raw <= 0.0 { return Ok(None); } - let seconds = if raw >= 1_000_000_000_000.0 { - raw / 1000.0 - } else { - raw - }; - let seconds = format!("{:.0}", seconds.trunc()) + if raw.fract() != 0.0 { + return Err(parse_failure( + "reset must use integral seconds or milliseconds", + )); + } + let raw = format!("{raw:.0}") .parse::() .map_err(|_| parse_failure("reset"))?; - Ok(Utc.timestamp_opt(seconds, 0).single()) + let (seconds, nanos) = if raw >= 1_000_000_000_000 { + (raw / 1000, ((raw % 1000) as u32) * 1_000_000) + } else { + (raw, 0) + }; + Ok(Utc.timestamp_opt(seconds, nanos).single()) } fn percent(used: f64, limit: f64) -> Option { @@ -368,4 +385,43 @@ mod tests { assert!(parse_quota(&json!({"limit": -1}), "rate").is_err()); assert!(parse_reset(Some(&json!(1e30))).is_err()); } + + #[test] + fn rejects_negative_balances_in_every_response_shape() { + assert!( + parse_billing(&json!({ + "billingType": "token", + "data": {"balance": {"total": 100, "remaining": -1}} + })) + .is_err() + ); + assert!( + parse_billing(&json!({ + "billingType": "token", + "data": { + "balance": {"total": 100, "remaining": 75}, + "onDemand": {"balance": -1} + } + })) + .is_err() + ); + assert!( + parse_billing(&json!({ + "billingType": "legacy", + "data": {"limit": 100, "remaining": -1} + })) + .is_err() + ); + assert!(parse_quota(&json!({"limit": 100, "remaining": -1}), "rate").is_err()); + } + + #[test] + fn rejects_fractional_resets_and_preserves_integral_milliseconds() { + assert!(parse_reset(Some(&json!(1_800_000_000.5))).is_err()); + let reset = parse_reset(Some(&json!(1_800_000_000_500_i64))) + .unwrap() + .unwrap(); + assert_eq!(reset.timestamp(), 1_800_000_000); + assert_eq!(reset.timestamp_subsec_millis(), 500); + } } From 581c9a6f237841e46f9295d03115a643af5780ed Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:59:15 +0700 Subject: [PATCH 04/11] Fix PR 604 thermo findings --- rust/src/providers/helmcode.rs | 11 +---- rust/src/providers/huggingface/mod.rs | 69 +++++++++++++++++++-------- rust/src/providers/mod.rs | 47 ++++++++++++++++++ rust/src/providers/typesafe.rs | 10 +--- 4 files changed, 100 insertions(+), 37 deletions(-) diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs index 04dc5073d7..b844c8b2da 100644 --- a/rust/src/providers/helmcode.rs +++ b/rust/src/providers/helmcode.rs @@ -93,7 +93,8 @@ impl HelmcodeProvider { let mut rejected = false; for tenant in candidates { let cookie = match ctx.manual_cookie_header.as_deref() { - Some(raw) => normalize_cookie(raw).ok_or(ProviderError::NoCookies)?, + Some(raw) => crate::providers::normalize_cookie_header(raw) + .ok_or(ProviderError::NoCookies)?, None => match crate::providers::browser_cookie_header(&[tenant.domain()]) { Ok(header) => header, Err(_) => continue, @@ -356,14 +357,6 @@ fn optional_nonnegative(value: Option<&Value>, field: &str) -> Result nonnegative(Some(value), field).map(Some), } } -fn normalize_cookie(raw: &str) -> Option { - let value = raw - .trim() - .strip_prefix("Cookie:") - .unwrap_or(raw.trim()) - .trim(); - (!value.is_empty() && !value.chars().any(char::is_control)).then(|| value.to_string()) -} fn parse_failure(field: impl AsRef) -> ProviderError { ProviderError::Parse(format!( "Helmcode quota response format changed: {}", diff --git a/rust/src/providers/huggingface/mod.rs b/rust/src/providers/huggingface/mod.rs index 6bef977198..eabf2dd0a8 100644 --- a/rust/src/providers/huggingface/mod.rs +++ b/rust/src/providers/huggingface/mod.rs @@ -53,6 +53,12 @@ struct IdentitySnapshot { plan: Option, } +#[derive(Debug, Clone, PartialEq)] +struct WalletCandidate { + user_id: String, + balance: f64, +} + #[derive(Debug, Clone, Default)] struct TokenEnvironment { config_api_key: Option, @@ -166,26 +172,21 @@ impl HuggingFaceProvider { let zerogpu_url = Url::parse(ZEROGPU_URL) .map_err(|_| ProviderError::Other("Invalid Hugging Face ZeroGPU URL.".to_string()))?; - let (billing, identity, zerogpu) = tokio::join!( + let (billing, identity, zerogpu, wallet_candidate) = tokio::join!( self.fetch_json(billing_url, &token, PRIMARY_TIMEOUT), self.fetch_optional_json(whoami_url, &token), self.fetch_optional_json(zerogpu_url, &token), + self.fetch_optional_wallet_candidate(), ); let billing = parse_billing(billing?)?; let identity = identity.and_then(|value| parse_identity(&value)); let zerogpu = zerogpu.and_then(|value| parse_zerogpu(&value)); - let balance = match identity - .as_ref() - .and_then(|identity| identity.user_id.as_deref()) - { - Some(user_id) => self.fetch_optional_wallet(user_id).await, - None => None, - }; + let balance = matching_wallet_balance(identity.as_ref(), wallet_candidate); Ok(build_result(billing, identity, zerogpu, balance)) } - async fn fetch_optional_wallet(&self, expected_user_id: &str) -> Option { + async fn fetch_optional_wallet_candidate(&self) -> Option { let cookie = crate::providers::browser_cookie_header(&["huggingface.co"]).ok()?; let (billing, whoami) = tokio::join!( self.fetch_cookie_text( @@ -197,17 +198,10 @@ impl HuggingFaceProvider { ); let billing = billing.ok()?; let whoami = whoami.ok()?; - let candidate = parse_wallet_balance(&billing).ok()?; + let balance = parse_wallet_balance(&billing).ok()?; let profile: Value = serde_json::from_str(&whoami).ok()?; - let observed_user_id = profile - .get("type") - .and_then(Value::as_str) - .filter(|kind| *kind == "user") - .and_then(|_| profile.get("id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty())?; - (observed_user_id == expected_user_id).then_some(candidate) + let user_id = parse_identity(&profile)?.user_id?; + Some(WalletCandidate { user_id, balance }) } async fn fetch_cookie_text( @@ -486,6 +480,15 @@ fn parse_identity(value: &Value) -> Option { ) } +fn matching_wallet_balance( + identity: Option<&IdentitySnapshot>, + candidate: Option, +) -> Option { + let expected_user_id = identity?.user_id.as_deref()?; + let candidate = candidate?; + (candidate.user_id == expected_user_id).then_some(candidate.balance) +} + fn safe_text(value: Option<&str>) -> Option { let value = value?.trim(); if value.is_empty() || value.chars().count() > 256 || value.chars().any(char::is_control) { @@ -896,6 +899,34 @@ mod tests { assert!(parse_identity(&json!({"email": "bad\nemail"})).is_none()); } + #[test] + fn wallet_candidate_is_attached_only_to_the_matching_token_identity() { + let identity = IdentitySnapshot { + user_id: Some("user-a".to_string()), + name: None, + email: None, + plan: None, + }; + let candidate = WalletCandidate { + user_id: "user-a".to_string(), + balance: 12.5, + }; + assert_eq!( + matching_wallet_balance(Some(&identity), Some(candidate.clone())), + Some(12.5) + ); + + let other_identity = IdentitySnapshot { + user_id: Some("user-b".to_string()), + ..identity.clone() + }; + assert_eq!( + matching_wallet_balance(Some(&other_identity), Some(candidate.clone())), + None + ); + assert_eq!(matching_wallet_balance(None, Some(candidate)), None); + } + #[test] fn status_errors_are_classified_without_response_body_or_token() { let token = "hf_secret_fixture"; diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index 2acbfcf3bd..6fc0df23c9 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -205,6 +205,22 @@ pub(crate) fn cookie_values<'a>(cookie_header: &'a str, name: &str) -> Vec<&'a s .collect() } +/// Normalize a user-supplied `Cookie` header value at the shared provider boundary. +/// +/// Accepts either the raw header value or a full, case-insensitive `Cookie:` line. +/// Empty values and control characters are rejected before the value reaches an +/// HTTP client. +pub(crate) fn normalize_cookie_header(raw: &str) -> Option { + let trimmed = raw.trim(); + let value = trimmed + .get(.."cookie:".len()) + .filter(|prefix| prefix.eq_ignore_ascii_case("cookie:")) + .map_or(trimmed, |_| &trimmed["cookie:".len()..]) + .trim(); + + (!value.is_empty() && !value.chars().any(char::is_control)).then(|| value.to_string()) +} + pub(crate) fn browser_cookies_for_domain( domain: &str, ) -> Result, crate::core::ProviderError> { @@ -340,3 +356,34 @@ pub(crate) fn extract_renewal(text: &str) -> Option Result { let cookie = match ctx.manual_cookie_header.as_deref() { - Some(raw) => normalize_cookie(raw).ok_or_else(|| { + Some(raw) => crate::providers::normalize_cookie_header(raw).ok_or_else(|| { ProviderError::Other( "TypeSafe needs a nonempty Cookie header from the billing page.".into(), ) @@ -407,14 +407,6 @@ fn build_result(billing: Billing) -> ProviderFetchResult { result } -fn normalize_cookie(raw: &str) -> Option { - let value = raw - .trim() - .strip_prefix("Cookie:") - .unwrap_or(raw.trim()) - .trim(); - (!value.is_empty() && !value.chars().any(char::is_control)).then(|| value.to_string()) -} fn finite(value: Option<&Value>, field: &str) -> Result { value .and_then(Value::as_f64) From 7249ea4b8960226b9e74f4602ea53bdfda130794 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 08:50:47 +0700 Subject: [PATCH 05/11] Bound provider response bodies --- rust/src/providers/helmcode.rs | 12 +- rust/src/providers/huggingface/mod.rs | 195 ++--------------------- rust/src/providers/huggingface/wallet.rs | 174 ++++++++++++++++++++ rust/src/providers/mod.rs | 86 +++++++++- rust/src/providers/typesafe.rs | 18 +-- rust/src/providers/v0.rs | 15 +- 6 files changed, 300 insertions(+), 200 deletions(-) create mode 100644 rust/src/providers/huggingface/wallet.rs diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs index b844c8b2da..6da2257b11 100644 --- a/rust/src/providers/helmcode.rs +++ b/rust/src/providers/helmcode.rs @@ -10,6 +10,9 @@ use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum Tenant { @@ -213,10 +216,13 @@ impl HelmcodeProvider { "Helmcode dashboard returned HTTP {status}." ))); } - let value = response - .json() + let body = read_bounded_response(response, MAX_RESPONSE_BYTES) .await - .map_err(|_| parse_failure("invalid JSON"))?; + .map_err(|error| match error { + BoundedBodyError::TooLarge => parse_failure("response too large"), + BoundedBodyError::Read(_) => parse_failure("invalid JSON"), + })?; + let value = serde_json::from_slice(&body).map_err(|_| parse_failure("invalid JSON"))?; Ok(Some(value)) } } diff --git a/rust/src/providers/huggingface/mod.rs b/rust/src/providers/huggingface/mod.rs index eabf2dd0a8..102cbd3d35 100644 --- a/rust/src/providers/huggingface/mod.rs +++ b/rust/src/providers/huggingface/mod.rs @@ -7,17 +7,21 @@ use async_trait::async_trait; use chrono::{DateTime, Datelike, TimeZone, Utc}; -use futures::StreamExt; use reqwest::{Client, StatusCode, Url}; use serde_json::Value; use std::path::{Path, PathBuf}; use std::time::Duration; +use super::{BoundedBodyError, read_bounded_response}; use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; +mod wallet; + +use wallet::{WalletCandidate, matching_wallet_balance, parse_wallet_balance}; + const BILLING_URL: &str = "https://huggingface.co/api/settings/billing/usage-v2"; const WHOAMI_URL: &str = "https://huggingface.co/api/whoami-v2"; const ZEROGPU_URL: &str = "https://huggingface.co/api/spaces/zero-gpu/quota"; @@ -53,12 +57,6 @@ struct IdentitySnapshot { plan: Option, } -#[derive(Debug, Clone, PartialEq)] -struct WalletCandidate { - user_id: String, - balance: f64, -} - #[derive(Debug, Clone, Default)] struct TokenEnvironment { config_api_key: Option, @@ -271,22 +269,16 @@ async fn read_bounded_body( response: reqwest::Response, response_kind: &str, ) -> Result, ProviderError> { - let mut body = Vec::new(); - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.map_err(|_| { - ProviderError::Parse(format!( - "Hugging Face returned an unreadable {response_kind}." - )) - })?; - if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(body.len()) { - return Err(ProviderError::Parse(format!( + read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => ProviderError::Parse(format!( "Hugging Face returned an oversized {response_kind}." - ))); - } - body.extend_from_slice(&chunk); - } - Ok(body) + )), + BoundedBodyError::Read(_) => ProviderError::Parse(format!( + "Hugging Face returned an unreadable {response_kind}." + )), + }) } impl Default for HuggingFaceProvider { @@ -480,15 +472,6 @@ fn parse_identity(value: &Value) -> Option { ) } -fn matching_wallet_balance( - identity: Option<&IdentitySnapshot>, - candidate: Option, -) -> Option { - let expected_user_id = identity?.user_id.as_deref()?; - let candidate = candidate?; - (candidate.user_id == expected_user_id).then_some(candidate.balance) -} - fn safe_text(value: Option<&str>) -> Option { let value = value?.trim(); if value.is_empty() || value.chars().count() > 256 || value.chars().any(char::is_control) { @@ -497,111 +480,6 @@ fn safe_text(value: Option<&str>) -> Option { Some(value.to_string()) } -fn parse_wallet_balance(html: &str) -> Result { - let mut current = Vec::new(); - let mut legacy = Vec::new(); - let mut rest = html; - while let Some(index) = rest.find("data-props") { - rest = &rest[index + "data-props".len()..]; - let trimmed = rest.trim_start(); - let Some(after_equals) = trimmed.strip_prefix('=') else { - continue; - }; - let after_equals = after_equals.trim_start(); - let Some(quote) = after_equals - .chars() - .next() - .filter(|value| matches!(value, '\'' | '"')) - else { - continue; - }; - let payload = &after_equals[quote.len_utf8()..]; - let Some(end) = payload.find(quote) else { - break; - }; - let decoded = decode_html_entities(&payload[..end])?; - rest = &payload[end + quote.len_utf8()..]; - let Ok(value) = serde_json::from_str::(&decoded) else { - continue; - }; - let Some(object) = value.as_object() else { - continue; - }; - if let Some(entity) = object.get("entity").and_then(Value::as_object) - && entity.contains_key("currentBalanceUsd") - { - if entity.get("type").and_then(Value::as_str) != Some("user") { - return Err(invalid_wallet("wallet entity type")); - } - current.push(wallet_number( - entity.get("currentBalanceUsd"), - "currentBalanceUsd", - )?); - } - if object.contains_key("invoiceCreditsCents") { - let cents = wallet_number(object.get("invoiceCreditsCents"), "invoiceCreditsCents")?; - if cents.fract() != 0.0 { - return Err(invalid_wallet("invoiceCreditsCents")); - } - legacy.push(cents / 100.0); - } - } - match (current.as_slice(), legacy.as_slice()) { - ([balance], _) => Ok(*balance), - ([], [balance]) => Ok(*balance), - ([], _) => Err(invalid_wallet("missing or ambiguous legacy wallet")), - _ => Err(invalid_wallet("ambiguous current wallet")), - } -} - -fn wallet_number(value: Option<&Value>, field: &str) -> Result { - value - .and_then(Value::as_f64) - .filter(|value| value.is_finite() && *value >= 0.0) - .ok_or_else(|| invalid_wallet(field)) -} - -fn invalid_wallet(field: &str) -> ProviderError { - ProviderError::Parse(format!("Hugging Face wallet field '{field}' was invalid.")) -} - -fn decode_html_entities(raw: &str) -> Result { - let mut output = String::with_capacity(raw.len()); - let mut rest = raw; - while let Some(index) = rest.find('&') { - output.push_str(&rest[..index]); - rest = &rest[index + 1..]; - let Some(end) = rest.find(';') else { - return Err(invalid_wallet("HTML entity")); - }; - let entity = &rest[..end]; - let decoded = match entity { - "amp" => '&', - "apos" => '\'', - "gt" => '>', - "lt" => '<', - "nbsp" => '\u{00a0}', - "quot" => '"', - value if value.starts_with("#x") || value.starts_with("#X") => { - u32::from_str_radix(&value[2..], 16) - .ok() - .and_then(char::from_u32) - .ok_or_else(|| invalid_wallet("HTML entity"))? - } - value if value.starts_with('#') => value[1..] - .parse::() - .ok() - .and_then(char::from_u32) - .ok_or_else(|| invalid_wallet("HTML entity"))?, - _ => return Err(invalid_wallet("HTML entity")), - }; - output.push(decoded); - rest = &rest[end + 1..]; - } - output.push_str(rest); - Ok(output) -} - fn build_result( billing: BillingSnapshot, identity: Option, @@ -899,34 +777,6 @@ mod tests { assert!(parse_identity(&json!({"email": "bad\nemail"})).is_none()); } - #[test] - fn wallet_candidate_is_attached_only_to_the_matching_token_identity() { - let identity = IdentitySnapshot { - user_id: Some("user-a".to_string()), - name: None, - email: None, - plan: None, - }; - let candidate = WalletCandidate { - user_id: "user-a".to_string(), - balance: 12.5, - }; - assert_eq!( - matching_wallet_balance(Some(&identity), Some(candidate.clone())), - Some(12.5) - ); - - let other_identity = IdentitySnapshot { - user_id: Some("user-b".to_string()), - ..identity.clone() - }; - assert_eq!( - matching_wallet_balance(Some(&other_identity), Some(candidate.clone())), - None - ); - assert_eq!(matching_wallet_balance(None, Some(candidate)), None); - } - #[test] fn status_errors_are_classified_without_response_body_or_token() { let token = "hf_secret_fixture"; @@ -976,21 +826,4 @@ mod tests { assert!(result.usage.secondary.is_none()); assert!(!result.pace_authoritative); } - - #[test] - fn wallet_parser_prefers_unique_current_balance_and_supports_legacy_cents() { - let current = r#"
"#; - assert_eq!(parse_wallet_balance(current).unwrap(), 12.5); - - let legacy = r#"
"#; - assert_eq!(parse_wallet_balance(legacy).unwrap(), 7.25); - } - - #[test] - fn wallet_parser_rejects_ambiguous_or_non_user_balances() { - let ambiguous = r#"
"#; - assert!(parse_wallet_balance(ambiguous).is_err()); - let organization = r#"
"#; - assert!(parse_wallet_balance(organization).is_err()); - } } diff --git a/rust/src/providers/huggingface/wallet.rs b/rust/src/providers/huggingface/wallet.rs new file mode 100644 index 0000000000..9958c6891f --- /dev/null +++ b/rust/src/providers/huggingface/wallet.rs @@ -0,0 +1,174 @@ +use serde_json::Value; + +use super::IdentitySnapshot; +use crate::core::ProviderError; + +#[derive(Debug, Clone, PartialEq)] +pub(super) struct WalletCandidate { + pub(super) user_id: String, + pub(super) balance: f64, +} + +pub(super) fn matching_wallet_balance( + identity: Option<&IdentitySnapshot>, + candidate: Option, +) -> Option { + let expected_user_id = identity?.user_id.as_deref()?; + let candidate = candidate?; + (candidate.user_id == expected_user_id).then_some(candidate.balance) +} + +pub(super) fn parse_wallet_balance(html: &str) -> Result { + let mut current = Vec::new(); + let mut legacy = Vec::new(); + let mut rest = html; + while let Some(index) = rest.find("data-props") { + rest = &rest[index + "data-props".len()..]; + let trimmed = rest.trim_start(); + let Some(after_equals) = trimmed.strip_prefix('=') else { + continue; + }; + let after_equals = after_equals.trim_start(); + let Some(quote) = after_equals + .chars() + .next() + .filter(|value| matches!(value, '\'' | '"')) + else { + continue; + }; + let payload = &after_equals[quote.len_utf8()..]; + let Some(end) = payload.find(quote) else { + break; + }; + let decoded = decode_html_entities(&payload[..end])?; + rest = &payload[end + quote.len_utf8()..]; + let Ok(value) = serde_json::from_str::(&decoded) else { + continue; + }; + let Some(object) = value.as_object() else { + continue; + }; + if let Some(entity) = object.get("entity").and_then(Value::as_object) + && entity.contains_key("currentBalanceUsd") + { + if entity.get("type").and_then(Value::as_str) != Some("user") { + return Err(invalid_wallet("wallet entity type")); + } + current.push(wallet_number( + entity.get("currentBalanceUsd"), + "currentBalanceUsd", + )?); + } + if object.contains_key("invoiceCreditsCents") { + let cents = wallet_number(object.get("invoiceCreditsCents"), "invoiceCreditsCents")?; + if cents.fract() != 0.0 { + return Err(invalid_wallet("invoiceCreditsCents")); + } + legacy.push(cents / 100.0); + } + } + match (current.as_slice(), legacy.as_slice()) { + ([balance], _) => Ok(*balance), + ([], [balance]) => Ok(*balance), + ([], _) => Err(invalid_wallet("missing or ambiguous legacy wallet")), + _ => Err(invalid_wallet("ambiguous current wallet")), + } +} + +fn wallet_number(value: Option<&Value>, field: &str) -> Result { + value + .and_then(Value::as_f64) + .filter(|value| value.is_finite() && *value >= 0.0) + .ok_or_else(|| invalid_wallet(field)) +} + +fn invalid_wallet(field: &str) -> ProviderError { + ProviderError::Parse(format!("Hugging Face wallet field '{field}' was invalid.")) +} + +fn decode_html_entities(raw: &str) -> Result { + let mut output = String::with_capacity(raw.len()); + let mut rest = raw; + while let Some(index) = rest.find('&') { + output.push_str(&rest[..index]); + rest = &rest[index + 1..]; + let Some(end) = rest.find(';') else { + return Err(invalid_wallet("HTML entity")); + }; + let entity = &rest[..end]; + let decoded = match entity { + "amp" => '&', + "apos" => '\'', + "gt" => '>', + "lt" => '<', + "nbsp" => '\u{00a0}', + "quot" => '"', + value if value.starts_with("#x") || value.starts_with("#X") => { + u32::from_str_radix(&value[2..], 16) + .ok() + .and_then(char::from_u32) + .ok_or_else(|| invalid_wallet("HTML entity"))? + } + value if value.starts_with('#') => value[1..] + .parse::() + .ok() + .and_then(char::from_u32) + .ok_or_else(|| invalid_wallet("HTML entity"))?, + _ => return Err(invalid_wallet("HTML entity")), + }; + output.push(decoded); + rest = &rest[end + 1..]; + } + output.push_str(rest); + Ok(output) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn candidate_is_attached_only_to_the_matching_token_identity() { + let identity = IdentitySnapshot { + user_id: Some("user-a".to_string()), + name: None, + email: None, + plan: None, + }; + let candidate = WalletCandidate { + user_id: "user-a".to_string(), + balance: 12.5, + }; + assert_eq!( + matching_wallet_balance(Some(&identity), Some(candidate.clone())), + Some(12.5) + ); + + let other_identity = IdentitySnapshot { + user_id: Some("user-b".to_string()), + ..identity.clone() + }; + assert_eq!( + matching_wallet_balance(Some(&other_identity), Some(candidate.clone())), + None + ); + assert_eq!(matching_wallet_balance(None, Some(candidate)), None); + } + + #[test] + fn parser_prefers_unique_current_balance_and_supports_legacy_cents() { + let current = r#"
"#; + assert_eq!(parse_wallet_balance(current).unwrap(), 12.5); + + let legacy = r#"
"#; + assert_eq!(parse_wallet_balance(legacy).unwrap(), 7.25); + } + + #[test] + fn parser_rejects_ambiguous_or_non_user_balances() { + let ambiguous = r#"
"#; + assert!(parse_wallet_balance(ambiguous).is_err()); + let organization = r#"
"#; + assert!(parse_wallet_balance(organization).is_err()); + } +} diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index 6fc0df23c9..de9d265f6e 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -5,6 +5,8 @@ reason = "provider traits and helpers are shared across modules; not all are consumed in every build" )] +use futures::{Stream, StreamExt}; + pub mod abacus; pub mod aiand; pub mod alibaba; @@ -170,6 +172,46 @@ pub use zed::ZedProvider; pub use zenmux::ZenMuxProvider; pub use zoommate::ZoomMateProvider; +#[derive(Debug, PartialEq)] +pub(crate) enum BoundedBodyError { + TooLarge, + Read(E), +} + +pub(crate) async fn read_bounded_response( + response: reqwest::Response, + max_bytes: usize, +) -> Result, BoundedBodyError> { + if response + .content_length() + .is_some_and(|length| length > max_bytes as u64) + { + return Err(BoundedBodyError::TooLarge); + } + read_bounded_stream(response.bytes_stream(), max_bytes).await +} + +async fn read_bounded_stream( + stream: S, + max_bytes: usize, +) -> Result, BoundedBodyError> +where + S: Stream>, + B: AsRef<[u8]>, +{ + let mut stream = Box::pin(stream); + let mut body = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(BoundedBodyError::Read)?; + let chunk = chunk.as_ref(); + if chunk.len() > max_bytes.saturating_sub(body.len()) { + return Err(BoundedBodyError::TooLarge); + } + body.extend_from_slice(chunk); + } + Ok(body) +} + pub(crate) fn browser_cookie_header( domains: &[&str], ) -> Result { @@ -359,7 +401,49 @@ pub(crate) fn extract_renewal(text: &str) -> Option(vec![0_u8; MAX_BYTES - 1]), Ok(vec![1_u8; 2])]); + + assert_eq!( + read_bounded_stream(body, MAX_BYTES).await, + Err(BoundedBodyError::TooLarge) + ); + } + + #[tokio::test] + async fn bounded_stream_accepts_body_below_limit() { + let body = stream::iter([Ok::<_, ()>(vec![1_u8, 2, 3]), Ok(vec![4_u8])]); + + assert_eq!(read_bounded_stream(body, 8).await, Ok(vec![1, 2, 3, 4])); + } + + #[tokio::test] + async fn bounded_stream_accepts_body_at_limit() { + let body = stream::iter([Ok::<_, ()>(vec![1_u8; 4]), Ok(vec![2_u8; 4])]); + + assert_eq!( + read_bounded_stream(body, 8).await, + Ok(vec![1, 1, 1, 1, 2, 2, 2, 2]) + ); + } + + #[tokio::test] + async fn bounded_stream_reports_read_failure() { + let body = stream::iter([ + Ok::<_, &'static str>(vec![1_u8]), + Err("network read failed"), + ]); + + assert_eq!( + read_bounded_stream(body, 8).await, + Err(BoundedBodyError::Read("network read failed")) + ); + } #[test] fn normalizes_raw_and_prefixed_cookie_headers() { diff --git a/rust/src/providers/typesafe.rs b/rust/src/providers/typesafe.rs index c6562ba35a..091e1f9ca3 100644 --- a/rust/src/providers/typesafe.rs +++ b/rust/src/providers/typesafe.rs @@ -2,11 +2,12 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; -use futures::{StreamExt, stream}; +use futures::stream; use reqwest::{Client, StatusCode, redirect::Policy}; use serde_json::Value; use std::time::Duration; +use super::{BoundedBodyError, read_bounded_response}; use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, @@ -215,15 +216,12 @@ async fn read_response(response: reqwest::Response) -> Result MAX_BODY_BYTES.saturating_sub(body.len()) { - return Err(parse_failure("response too large")); - } - body.extend_from_slice(&chunk); - } + let body = read_bounded_response(response, MAX_BODY_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => parse_failure("response too large"), + BoundedBodyError::Read(error) => ProviderError::Network(error), + })?; let body = String::from_utf8(body).map_err(|_| parse_failure("response was not UTF-8"))?; if body.contains("\\\"(auth)\\\",{\\\"children\\\":[\\\"login\\\"") { return Err(ProviderError::AuthRequired); diff --git a/rust/src/providers/v0.rs b/rust/src/providers/v0.rs index 7b18b76926..c0bddadb06 100644 --- a/rust/src/providers/v0.rs +++ b/rust/src/providers/v0.rs @@ -10,10 +10,12 @@ use crate::core::{ FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; +use crate::providers::{BoundedBodyError, read_bounded_response}; const API_BASE: &str = "https://api.v0.dev/v1"; const CREDENTIAL_TARGET: &str = "codexbar-v0"; const ENV_KEYS: &[&str] = &["V0_API_KEY"]; +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; #[derive(Debug, Clone, PartialEq)] struct Quota { @@ -100,11 +102,14 @@ impl V0Provider { } let response = self.client.get(url).bearer_auth(api_key).send().await?; classify_status(response.status())?; - response.json().await.map_err(|_| { - ProviderError::Parse(format!( - "Could not parse v0 usage: {path} returned invalid JSON" - )) - }) + let body = read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => parse_failure(format!("{path} response too large")), + BoundedBodyError::Read(_) => parse_failure(format!("{path} returned invalid JSON")), + })?; + serde_json::from_slice(&body) + .map_err(|_| parse_failure(format!("{path} returned invalid JSON"))) } } From 34bb26100936cee35a4238a7625f6e46b8873ab3 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 09:06:50 +0700 Subject: [PATCH 06/11] Import stream extension for Typesafe provider --- rust/src/providers/typesafe.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rust/src/providers/typesafe.rs b/rust/src/providers/typesafe.rs index 091e1f9ca3..a4139eda41 100644 --- a/rust/src/providers/typesafe.rs +++ b/rust/src/providers/typesafe.rs @@ -2,7 +2,7 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; -use futures::stream; +use futures::{StreamExt, stream}; use reqwest::{Client, StatusCode, redirect::Policy}; use serde_json::Value; use std::time::Duration; From d73ffeb4a61e2b68755eff26c3cc41892f2afd38 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 09:44:52 +0700 Subject: [PATCH 07/11] Reject negative provider monetary values --- rust/src/providers/helmcode.rs | 20 +++++++++++++++++-- rust/src/providers/typesafe.rs | 36 +++++++++++++++++++++++++++------- 2 files changed, 47 insertions(+), 9 deletions(-) diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs index 6da2257b11..4800c1a71b 100644 --- a/rust/src/providers/helmcode.rs +++ b/rust/src/providers/helmcode.rs @@ -161,7 +161,9 @@ impl HelmcodeProvider { && let Some(credits) = self .get(tenant, cookie, "/api/billing/credits", true) .await? - && let Some(balance_micros) = credits.get("balanceMicros").and_then(Value::as_i64) + && let Some(balance_micros) = credits + .get("balanceMicros") + .and_then(nonnegative_balance_micros) { let currency = credits .get("currency") @@ -171,7 +173,7 @@ impl HelmcodeProvider { if currency.len() == 3 && currency.chars().all(|ch| ch.is_ascii_uppercase()) { result = result.with_cost( CostSnapshot::new(0.0, currency, "Prepaid balance") - .with_balance((balance_micros.max(0) as f64) / 1_000_000.0), + .with_balance((balance_micros as f64) / 1_000_000.0), ); } } @@ -363,6 +365,9 @@ fn optional_nonnegative(value: Option<&Value>, field: &str) -> Result nonnegative(Some(value), field).map(Some), } } +fn nonnegative_balance_micros(value: &Value) -> Option { + value.as_i64().filter(|amount| *amount >= 0) +} fn parse_failure(field: impl AsRef) -> ProviderError { ProviderError::Parse(format!( "Helmcode quota response format changed: {}", @@ -400,4 +405,15 @@ mod tests { .is_err() ); } + + #[test] + fn accepts_nonnegative_balance_micros_and_rejects_negative_values() { + assert_eq!(nonnegative_balance_micros(&json!(0)), Some(0)); + assert_eq!( + nonnegative_balance_micros(&json!(1_250_000)), + Some(1_250_000) + ); + assert_eq!(nonnegative_balance_micros(&json!(-1)), None); + assert_eq!(nonnegative_balance_micros(&json!(1.5)), None); + } } diff --git a/rust/src/providers/typesafe.rs b/rust/src/providers/typesafe.rs index a4139eda41..67136d6013 100644 --- a/rust/src/providers/typesafe.rs +++ b/rust/src/providers/typesafe.rs @@ -299,8 +299,8 @@ fn parse_rsc_billing(body: &str) -> Result { .and_then(|value| value.get("billing")) .and_then(Value::as_object) .ok_or_else(|| parse_failure("missing billing"))?; - let spent = finite(billing.get("spent"), "spent")?; - let balance = finite(billing.get("balance"), "balance")?; + let spent = finite_nonnegative(billing.get("spent"), "spent")?; + let balance = finite_nonnegative(billing.get("balance"), "balance")?; let cycle_label = clean_text(billing.get("cycleLabel")); let plan = clean_text(billing.get("plan")); let credits = billing @@ -310,8 +310,8 @@ fn parse_rsc_billing(body: &str) -> Result { .flatten() .filter_map(|item| { let object = item.as_object()?; - let amount = finite(object.get("amount"), "credit amount").ok()?; - let remaining = finite(object.get("remaining"), "credit remaining").ok()?; + let amount = finite_nonnegative(object.get("amount"), "credit amount").ok()?; + let remaining = finite_nonnegative(object.get("remaining"), "credit remaining").ok()?; if remaining <= 0.0 { return None; } @@ -405,10 +405,10 @@ fn build_result(billing: Billing) -> ProviderFetchResult { result } -fn finite(value: Option<&Value>, field: &str) -> Result { +fn finite_nonnegative(value: Option<&Value>, field: &str) -> Result { value .and_then(Value::as_f64) - .filter(|value| value.is_finite()) + .filter(|value| value.is_finite() && *value >= 0.0) .ok_or_else(|| parse_failure(field)) } fn clean_text(value: Option<&Value>) -> Option { @@ -470,10 +470,32 @@ mod tests { #[test] fn parses_billing_result_and_skips_expired_or_empty_credits() { - let body = r#"1:{"ok":true,"data":{"billing":{"spent":4.5,"balance":10,"cycleLabel":"September","plan":"free_plan","credits":[{"amount":8,"remaining":3,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":1,"remaining":0,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":5,"remaining":2,"expiresAt":"2000-01-02T00:00:00Z"}]}}}"#; + let body = r#"1:{"ok":true,"data":{"billing":{"spent":4.5,"balance":10,"cycleLabel":"September","plan":"free_plan","credits":[{"amount":8,"remaining":3,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":1,"remaining":0,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":5,"remaining":2,"expiresAt":"2000-01-02T00:00:00Z"},{"amount":-2,"remaining":1,"expiresAt":"2100-01-02T00:00:00Z"},{"amount":2,"remaining":-1,"expiresAt":"2100-01-02T00:00:00Z"}]}}}"#; let parsed = parse_rsc_billing(body).unwrap(); assert_eq!(parsed.spent, 4.5); assert_eq!(parsed.balance, 10.0); assert_eq!(parsed.credits.len(), 1); } + + #[test] + fn rejects_negative_spent_and_balance() { + let negative_spent = r#"1:{"ok":true,"data":{"billing":{"spent":-0.01,"balance":10}}}"#; + let negative_balance = r#"1:{"ok":true,"data":{"billing":{"spent":0,"balance":-0.01}}}"#; + + assert!(parse_rsc_billing(negative_spent).is_err()); + assert!(parse_rsc_billing(negative_balance).is_err()); + } + + #[test] + fn accepts_zero_and_positive_monetary_values() { + assert_eq!( + finite_nonnegative(Some(&serde_json::json!(0)), "amount").unwrap(), + 0.0 + ); + assert_eq!( + finite_nonnegative(Some(&serde_json::json!(1.25)), "amount").unwrap(), + 1.25 + ); + assert!(finite_nonnegative(Some(&serde_json::json!(-0.01)), "amount").is_err()); + } } From a0eea7f584b7b13f7737b4d5167def93b20e26f3 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:10:34 +0700 Subject: [PATCH 08/11] Improve v0 request and balance validation --- rust/src/providers/v0.rs | 52 ++++++++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 21 deletions(-) diff --git a/rust/src/providers/v0.rs b/rust/src/providers/v0.rs index c0bddadb06..35264564ee 100644 --- a/rust/src/providers/v0.rs +++ b/rust/src/providers/v0.rs @@ -75,17 +75,12 @@ impl V0Provider { .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()) }); - let billing = parse_billing( - &self - .get_json("/user/billing", scope.as_deref(), &api_key) - .await?, - )?; - let rate_limit = parse_quota( - &self - .get_json("/rate-limits", scope.as_deref(), &api_key) - .await?, - "rate limit response", - )?; + let (billing_result, rate_limit_result) = tokio::join!( + self.get_json("/user/billing", scope.as_deref(), &api_key), + self.get_json("/rate-limits", scope.as_deref(), &api_key), + ); + let billing = parse_billing(&billing_result?)?; + let rate_limit = parse_quota(&rate_limit_result?, "rate limit response")?; Ok(build_result(billing, rate_limit, scope.as_deref())) } @@ -184,14 +179,9 @@ fn parse_billing(value: &Value) -> Result { .get("balance") .and_then(Value::as_object) .ok_or_else(|| parse_failure("billing.data.balance"))?; - let total = finite_number(balance.get("total"), "billing.data.balance.total")?; + let total = nonnegative_number(balance.get("total"), "billing.data.balance.total")?; let remaining = nonnegative_number(balance.get("remaining"), "billing.data.balance.remaining")?; - if total < 0.0 { - return Err(parse_failure( - "billing.data.balance.total must not be negative", - )); - } let resets_at = match data.get("billingCycle").and_then(Value::as_object) { Some(cycle) => parse_reset(cycle.get("end"))?, None => None, @@ -221,10 +211,7 @@ fn parse_billing(value: &Value) -> Result { fn parse_quota(value: &Value, field: &str) -> Result { let object = value.as_object().ok_or_else(|| parse_failure(field))?; - let limit = finite_number(object.get("limit"), &format!("{field}.limit"))?; - if limit < 0.0 { - return Err(parse_failure(format!("{field}.limit must not be negative"))); - } + let limit = nonnegative_number(object.get("limit"), &format!("{field}.limit"))?; let remaining = optional_nonnegative_number(object.get("remaining"), &format!("{field}.remaining"))?; Ok(Quota { @@ -374,16 +361,39 @@ mod tests { "data": {"balance": {"total": 100, "remaining": 75}, "billingCycle": {"end": 1_800_000_000}, "onDemand": {"balance": 12.5}} })).unwrap(); assert_eq!(billing.quota.used_percent, Some(25.0)); + assert_eq!(billing.quota.limit, 100.0); assert_eq!(billing.on_demand_balance, Some(12.5)); let unknown = parse_quota( &json!({"limit": 50, "remaining": null, "reset": null}), "rate", ) .unwrap(); + assert_eq!(unknown.limit, 50.0); assert_eq!(unknown.used_percent, None); assert_eq!(unknown.remaining, None); } + #[test] + fn rejects_negative_token_total_and_accepts_zero_limits() { + assert!( + parse_billing(&json!({ + "billingType": "token", + "data": {"balance": {"total": -1, "remaining": 0}} + })) + .is_err() + ); + + let zero_billing = parse_billing(&json!({ + "billingType": "token", + "data": {"balance": {"total": 0, "remaining": 0}} + })) + .unwrap(); + assert_eq!(zero_billing.quota.limit, 0.0); + + let zero_quota = parse_quota(&json!({"limit": 0, "remaining": 0}), "rate").unwrap(); + assert_eq!(zero_quota.limit, 0.0); + } + #[test] fn rejects_unknown_billing_type_and_negative_limits() { assert!(parse_billing(&json!({"billingType": "future", "data": {}})).is_err()); From 2b55915ffd808ee10037f54adaaac62fcacf096e Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:28:39 +0700 Subject: [PATCH 09/11] Parallelize Helmcode provider requests --- rust/src/providers/helmcode.rs | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs index 4800c1a71b..7ab79659a5 100644 --- a/rust/src/providers/helmcode.rs +++ b/rust/src/providers/helmcode.rs @@ -124,11 +124,20 @@ impl HelmcodeProvider { tenant: Tenant, cookie: &str, ) -> Result { - let quota = self - .get(tenant, cookie, "/api/usage/quota", false) - .await? - .ok_or_else(|| parse_failure("quota"))?; - let billing = self.get(tenant, cookie, "/api/billing", true).await?; + let quota_request = self.get(tenant, cookie, "/api/usage/quota", false); + let billing_request = self.get(tenant, cookie, "/api/billing", true); + let credits_request = async { + if tenant == Tenant::Helmcode { + self.get(tenant, cookie, "/api/billing/credits", true).await + } else { + Ok(None) + } + }; + let (quota_result, billing_result, credits_result) = + tokio::join!(quota_request, billing_request, credits_request); + + let quota = quota_result?.ok_or_else(|| parse_failure("quota"))?; + let billing = billing_result?; let premium = billing .as_ref() .and_then(|value| value.get("subscription")) @@ -157,10 +166,7 @@ impl HelmcodeProvider { ); } let mut result = ProviderFetchResult::new(usage, "web"); - if tenant == Tenant::Helmcode - && let Some(credits) = self - .get(tenant, cookie, "/api/billing/credits", true) - .await? + if let Some(credits) = credits_result? && let Some(balance_micros) = credits .get("balanceMicros") .and_then(nonnegative_balance_micros) From ce6ee941f9f837d075f222610d9022097e480fb0 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:43:23 +0700 Subject: [PATCH 10/11] Validate quotas before fetching credits --- rust/src/providers/helmcode.rs | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/rust/src/providers/helmcode.rs b/rust/src/providers/helmcode.rs index 7ab79659a5..34e550c41c 100644 --- a/rust/src/providers/helmcode.rs +++ b/rust/src/providers/helmcode.rs @@ -126,15 +126,7 @@ impl HelmcodeProvider { ) -> Result { let quota_request = self.get(tenant, cookie, "/api/usage/quota", false); let billing_request = self.get(tenant, cookie, "/api/billing", true); - let credits_request = async { - if tenant == Tenant::Helmcode { - self.get(tenant, cookie, "/api/billing/credits", true).await - } else { - Ok(None) - } - }; - let (quota_result, billing_result, credits_result) = - tokio::join!(quota_request, billing_request, credits_request); + let (quota_result, billing_result) = tokio::join!(quota_request, billing_request); let quota = quota_result?.ok_or_else(|| parse_failure("quota"))?; let billing = billing_result?; @@ -151,6 +143,12 @@ impl HelmcodeProvider { .unwrap_or(std::cmp::Ordering::Equal) .then_with(|| a.name.cmp(&b.name)) }); + let credits = if tenant == Tenant::Helmcode { + self.get(tenant, cookie, "/api/billing/credits", true) + .await? + } else { + None + }; let primary = models .first() .map(model_window) @@ -166,7 +164,7 @@ impl HelmcodeProvider { ); } let mut result = ProviderFetchResult::new(usage, "web"); - if let Some(credits) = credits_result? + if let Some(credits) = credits && let Some(balance_micros) = credits .get("balanceMicros") .and_then(nonnegative_balance_micros) From 98da7a476acb4f2963b49379e1fc73c4befdd0bd Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:11:17 +0700 Subject: [PATCH 11/11] Document v0.64 provider adapters --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index c610d8f17d..2e052ad2f9 100755 --- a/README.md +++ b/README.md @@ -80,6 +80,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | MiniMax | API / Cookies | Usage, Billing Summary | | Kiro | Cookies / CLI | Monthly Credits, Overage | | Vertex AI | gcloud OAuth | Cost | +| v0 | API Key | Billing quota, API rate limits, on-demand balance | | Augment | Cookies | Credits | | OpenCode | Local Config | Usage | | Kimi | Cookies | 5h Rate, Weekly | @@ -89,6 +90,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Ollama | Cookies / API Key | Usage, Cloud Models, Pace windows | | Azure OpenAI | API Key | Deployment | | T3 Chat | Cookies / cURL | Base, Overage | +| TypeSafe | Browser cookies / manual Cookie header | Billing-cycle spend, balance, expiring credits | | OpenRouter | API Key | Credits | | JetBrains AI | Local Config | Usage | | Alibaba | Cookies | Usage | @@ -112,6 +114,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Venice | API Key | USD / DIEM Balance | | OpenAI | Admin API / API Key | Usage, Requests, Project-scoped cost, Credit Balance | | Grok | Cookies / auth.json | Billing | +| Helmcode (also NaN Builders) | Browser cookies / manual Cookie header | Per-model token quotas, reset windows, Helmcode prepaid balance | | Replicate | Cookies / token accounts | Monthly spend, credit balance | | ElevenLabs | API Key | Subscription Credits, Voice Slots | | Deepgram | API Key | Project Usage |