From 9d0fe80467fc23ca7ab154b86ff1e82666b6cbf4 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:17:57 +0700 Subject: [PATCH 01/15] Add prepaid balance to CostSnapshot --- rust/src/core/usage_snapshot.rs | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 116ffc02f6..dbb4f7f174 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -227,6 +227,10 @@ pub struct CostSnapshot { /// When this snapshot was captured pub updated_at: DateTime, + + /// Remaining prepaid balance (currency units), separate from used/limit spend. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub balance: Option, } impl CostSnapshot { @@ -239,6 +243,7 @@ impl CostSnapshot { period: period.into(), resets_at: None, updated_at: Utc::now(), + balance: None, } } @@ -248,6 +253,12 @@ impl CostSnapshot { self } + /// Builder pattern: set remaining prepaid balance (finite, ≥ 0 only) + pub fn with_balance(mut self, balance: f64) -> Self { + self.balance = finite_amount(balance); + self + } + /// Builder pattern: set reset time pub fn with_resets_at(mut self, resets_at: DateTime) -> Self { self.resets_at = Some(resets_at); @@ -279,6 +290,12 @@ impl CostSnapshot { pub fn format_limit(&self) -> Option { self.limit.map(|l| format_currency(l, &self.currency_code)) } + + /// Format the prepaid balance as a currency string + pub fn format_balance(&self) -> Option { + self.balance + .map(|b| format_currency(b, &self.currency_code)) + } } /// Format a value as currency From 647f76451b9e4b89ee482c300e980402d980e6d7 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:06 +0700 Subject: [PATCH 02/15] Add Qwen Cloud, ZoomMate, and Alibaba Token Plan regions --- .../providers/icons/ProviderIcon-zoommate.svg | 4 + .../src/components/providers/providerIcons.ts | 5 + apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 4 +- .../surfaces/settings/tabs/ProvidersTab.tsx | 2 + .../desktop-tauri/src/test/providerCatalog.ts | 1 + rust/src/core/curl_capture.rs | 271 ++++ rust/src/core/provider.rs | 54 +- rust/src/core/provider_factory.rs | 8 +- rust/src/core/token_accounts.rs | 11 +- .../mod.rs} | 341 ++-- .../providers/alibabatokenplan/personal.rs | 371 +++++ rust/src/providers/alibabatokenplan/region.rs | 220 +++ rust/src/providers/mod.rs | 6 +- rust/src/providers/qwencloud/mod.rs | 1391 +++++++++++++++++ rust/src/providers/zoommate/mod.rs | 886 +++++++++++ 15 files changed, 3455 insertions(+), 120 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg create mode 100644 rust/src/core/curl_capture.rs rename rust/src/providers/{alibabatokenplan.rs => alibabatokenplan/mod.rs} (74%) create mode 100644 rust/src/providers/alibabatokenplan/personal.rs create mode 100644 rust/src/providers/alibabatokenplan/region.rs create mode 100644 rust/src/providers/qwencloud/mod.rs create mode 100644 rust/src/providers/zoommate/mod.rs diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg new file mode 100644 index 0000000000..f8806a5391 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg @@ -0,0 +1,4 @@ + + + + diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 3fe4974de8..11e3da6941 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -22,6 +22,7 @@ import aiand from "./icons/ProviderIcon-aiand.svg?raw"; import clinepass from "./icons/ProviderIcon-clinepass.svg?raw"; import longcat from "./icons/ProviderIcon-longcat.svg?raw"; import neuralwatt from "./icons/ProviderIcon-neuralwatt.svg?raw"; +import zoommate from "./icons/ProviderIcon-zoommate.svg?raw"; import zenmux from "./icons/ProviderIcon-zenmux.svg?raw"; import deepseek from "./icons/ProviderIcon-deepseek.svg?raw"; import doubao from "./icons/ProviderIcon-doubao.svg?raw"; @@ -100,6 +101,7 @@ const RAW: Record = { clinepass: tint(clinepass), longcat: tint(longcat), neuralwatt: tint(neuralwatt), + zoommate: tint(zoommate), zenmux: tint(zenmux), deepseek: tint(deepseek), doubao: tint(doubao), @@ -156,6 +158,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { clinepass: { id: "clinepass", brandColor: "#61a3fa", fallbackLetter: "C", svgPath: RAW.clinepass }, longcat: { id: "longcat", brandColor: "#ffd100", fallbackLetter: "L", svgPath: RAW.longcat }, neuralwatt: { id: "neuralwatt", brandColor: "#38d98c", fallbackLetter: "N", svgPath: RAW.neuralwatt }, + zoommate: { id: "zoommate", brandColor: "#0B5CFF", fallbackLetter: "Z", svgPath: RAW.zoommate }, zenmux: { id: "zenmux", brandColor: "#6c5ce7", fallbackLetter: "Z", svgPath: RAW.zenmux }, deepseek: { id: "deepseek", brandColor: "#527df0", fallbackLetter: "D", svgPath: RAW.deepseek }, elevenlabs: { id: "elevenlabs", brandColor: "#111827", fallbackLetter: "E", svgPath: RAW.elevenlabs }, @@ -241,6 +244,8 @@ const ALIASES: Record = { lc: "longcat", "neural-watt": "neuralwatt", nw: "neuralwatt", + "zoom-mate": "zoommate", + "zoom mate": "zoommate", codeium: "windsurf", "xiaomi mimo": "mimo", xiaomimimo: "mimo", diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index 7d367e7174..4875a1e5fa 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -38,7 +38,7 @@ import AgentSessions from "../components/AgentSessions"; const HAS_DASHBOARD = new Set([ "abacus", "alibaba", "alibabatokenplan", "amp", "augment", "azureopenai", "bedrock", "claude", "codex", "codebuff", - "aiand", "commandcode", "copilot", "crof", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", + "aiand", "commandcode", "copilot", "crof", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", "doubao", "elevenlabs", "factory", "gemini", "grok", "groq", "infini", "jetbrains", "kilo", "kimi", "kimik2", "kiro", "manus", "mimo", "minimax", "mistral", "nanogpt", "ollama", "openaiapi", @@ -49,7 +49,7 @@ const HAS_DASHBOARD = new Set([ /** Provider IDs that have a status page URL in the backend */ const HAS_STATUS_PAGE = new Set([ "alibabatokenplan", "amp", "augment", "azureopenai", "bedrock", - "claude", "codex", "copilot", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "elevenlabs", + "claude", "codex", "copilot", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", "elevenlabs", "gemini", "grok", "groq", "kiro", "mistral", "openaiapi", "openrouter", "vertexai", "windsurf", ]); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx index 9131b60ac5..8ccefe3bfb 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx @@ -215,6 +215,8 @@ function providerSourceHintShort( case "infini": case "manus": case "mimo": + case "zoommate": + case "t3chat": case "commandcode": return t("ProviderSourceWebShort"); case "gemini": diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index 1042ca7d37..bda7cbc154 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -39,6 +39,7 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["clinepass", "ClinePass"], ["longcat", "LongCat"], ["neuralwatt", "Neuralwatt"], + ["zoommate", "ZoomMate"], ["windsurf", "Windsurf"], ["manus", "Manus"], ["mimo", "Xiaomi MiMo"], diff --git a/rust/src/core/curl_capture.rs b/rust/src/core/curl_capture.rs new file mode 100644 index 0000000000..14ed579dd3 --- /dev/null +++ b/rust/src/core/curl_capture.rs @@ -0,0 +1,271 @@ +//! Shared parsing for DevTools "Copy as cURL" captures pasted into manual-auth +//! provider settings. Ported from upstream `CurlCaptureParser` (v0.46.0). + +use regex_lite::Regex; +use reqwest::Url; +use std::collections::HashMap; +use std::sync::OnceLock; + +/// Extracts the request URL from the standard DevTools "Copy as cURL" shape, +/// where the URL is the first argument after `curl`. +/// +/// Returns `None` for malformed captures or option-first forms +/// (e.g. `curl --location 'https://…'`). +pub fn request_url(raw: &str) -> Option { + // (?s)(?:^|\s)curl\s+(?:\$'…'|'…'|"…"|bare) + static RE: OnceLock = OnceLock::new(); + let re = RE.get_or_init(|| { + Regex::new( + r#"(?s)(?:^|\s)curl(?:\.exe)?\s+(?:\$'((?:\\.|[^'])*)'|'([^']*)'|"((?:\\.|[^"])*)"|([^\s\\]+))"#, + ) + .expect("curl url regex") + }); + let caps = re.captures(raw)?; + let value = if let Some(m) = caps.get(1) { + unescape_shell_segment(m.as_str(), true) + } else if let Some(m) = caps.get(2) { + m.as_str().to_string() + } else if let Some(m) = caps.get(3) { + unescape_shell_segment(m.as_str(), false) + } else if let Some(m) = caps.get(4) { + unescape_shell_segment(m.as_str(), false) + } else { + return None; + }; + let value = value.trim(); + if value.is_empty() { + return None; + } + // Minimal URL sanity: scheme + host present. + let parsed = Url::parse(value).ok()?; + if parsed.scheme().is_empty() || parsed.host_str().is_none() { + return None; + } + Some(value.to_string()) +} + +/// Splits a raw cURL capture into the raw `-H "Name: Value"` / `--header '…'` +/// header field strings it contains, unescaping shell quoting (including +/// ANSI-C `$'...'` strings). +pub fn header_fields(raw: &str) -> Vec { + static RE: OnceLock = OnceLock::new(); + let re = RE.get_or_init(|| { + // regex_lite has no lookaround; allow space, `=`, or glued quote (`-H'…'`). + Regex::new( + r#"(?s)(?:^|\s)(?:-H|--header)(?:\s*=\s*|\s+)?(?:\$'((?:\\.|[^'])*)'|'([^']*)'|"((?:\\.|[^"])*)"|(\S+))"#, + ) + .expect("curl header regex") + }); + let mut fields = Vec::new(); + for caps in re.captures_iter(raw) { + if let Some(m) = caps.get(1) { + fields.push(unescape_shell_segment(m.as_str(), true)); + } else if let Some(m) = caps.get(2) { + fields.push(m.as_str().to_string()); + } else if let Some(m) = caps.get(3) { + fields.push(unescape_shell_segment(m.as_str(), false)); + } else if let Some(m) = caps.get(4) { + fields.push(unescape_shell_segment(m.as_str(), false)); + } + } + fields +} + +/// Returns the value of the first header field whose name matches `name` +/// (case-insensitive). +pub fn header_value(name: &str, fields: &[String]) -> Option { + for field in fields { + let (raw_name, value) = split_header_field(field)?; + if raw_name.eq_ignore_ascii_case(name) && !value.is_empty() { + return Some(value.to_string()); + } + } + None +} + +/// Builds a `[canonicalHeaderName: value]` map from `fields`, keeping only +/// headers whose lowercased name appears in `allowlist` (mapping lowercase +/// name → canonical HTTP header name). +pub fn forwarded_headers( + fields: &[String], + allowlist: &HashMap<&str, &str>, +) -> HashMap { + let mut headers = HashMap::new(); + for field in fields { + let Some((raw_name, value)) = split_header_field(field) else { + continue; + }; + if raw_name.is_empty() || value.is_empty() { + continue; + } + let key = raw_name.to_ascii_lowercase(); + if let Some(&canonical) = allowlist.get(key.as_str()) { + headers.insert(canonical.to_string(), value.to_string()); + } + } + headers +} + +/// Convenience: allowlist as a slice of `(lowercase, Canonical)` pairs. +pub fn forwarded_headers_from_pairs( + fields: &[String], + allowlist: &[(&str, &str)], +) -> HashMap { + let map: HashMap<&str, &str> = allowlist.iter().copied().collect(); + forwarded_headers(fields, &map) +} + +fn split_header_field(field: &str) -> Option<(&str, &str)> { + let colon = field.find(':')?; + let name = field[..colon].trim(); + let value = field[colon + 1..].trim(); + Some((name, value)) +} + +/// Unescape a shell-quoted segment. When `ansi` is true (ANSI-C `$'…'`), +/// `\n`/`\r`/`\t` expand; otherwise only backslash-escapes of the next char +/// (and line-continuation `\\\n`) are handled. +pub fn unescape_shell_segment(raw: &str, ansi: bool) -> String { + let mut output = String::with_capacity(raw.len()); + let mut chars = raw.chars().peekable(); + while let Some(ch) = chars.next() { + if ch != '\\' { + output.push(ch); + continue; + } + let Some(next) = chars.next() else { + break; + }; + match next { + 'n' if ansi => output.push('\n'), + 'r' if ansi => output.push('\r'), + 't' if ansi => output.push('\t'), + '\n' => {} + other => output.push(other), + } + } + output +} + +/// True when `raw` looks like a DevTools cURL capture rather than a bare cookie header. +pub fn looks_like_curl_capture(raw: &str) -> bool { + let lower = raw.trim_start().to_ascii_lowercase(); + lower.starts_with("curl ") || lower.starts_with("curl.exe ") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extracts_request_url_from_standard_devtools_capture() { + let curl = "curl 'https://example.com/api/usage' -H 'Authorization: Bearer fake-token'"; + assert_eq!( + request_url(curl).as_deref(), + Some("https://example.com/api/usage") + ); + } + + #[test] + fn extracts_request_url_from_double_quoted_and_bare_captures() { + assert_eq!( + request_url("curl \"https://example.com/double\"") + .as_deref() + .and_then(|u| Url::parse(u).ok()) + .map(|u| u.path().to_string()) + .as_deref(), + Some("/double") + ); + assert_eq!( + request_url("curl https://example.com/bare") + .as_deref() + .and_then(|u| Url::parse(u).ok()) + .map(|u| u.path().to_string()) + .as_deref(), + Some("/bare") + ); + } + + #[test] + fn request_url_returns_none_for_option_first_or_malformed() { + assert_eq!(request_url("curl --location 'https://example.com'"), None); + assert_eq!(request_url("not-curl 'https://example.com'"), None); + } + + #[test] + fn extracts_header_fields_from_double_quoted_headers() { + let curl = r#"curl 'https://example.com' --header "Authorization: Bearer fake-token" --header "Cookie: session=abc""#; + let fields = header_fields(curl); + assert!( + fields + .iter() + .any(|f| f == "Authorization: Bearer fake-token") + ); + assert!(fields.iter().any(|f| f == "Cookie: session=abc")); + } + + #[test] + fn extracts_header_fields_from_single_quoted_h_flags() { + let curl = "curl 'https://example.com' -H 'Authorization: Bearer fake-token'"; + assert_eq!( + header_fields(curl), + vec!["Authorization: Bearer fake-token"] + ); + } + + #[test] + fn header_value_lookup_is_case_insensitive() { + let fields = vec!["AUTHORIZATION: Bearer fake-token".to_string()]; + assert_eq!( + header_value("authorization", &fields).as_deref(), + Some("Bearer fake-token") + ); + } + + #[test] + fn header_value_lookup_returns_none_for_missing() { + let fields = vec!["Cookie: session=abc".to_string()]; + assert_eq!(header_value("Authorization", &fields), None); + } + + #[test] + fn forwarded_headers_respects_allowlist() { + let fields = vec![ + "Authorization: Bearer fake-token".to_string(), + "Cookie: session=abc".to_string(), + "X-Not-Allowed: nope".to_string(), + ]; + let headers = forwarded_headers_from_pairs( + &fields, + &[("authorization", "Authorization"), ("cookie", "Cookie")], + ); + assert_eq!( + headers.get("Authorization").map(String::as_str), + Some("Bearer fake-token") + ); + assert_eq!( + headers.get("Cookie").map(String::as_str), + Some("session=abc") + ); + assert!(!headers.contains_key("X-Not-Allowed")); + } + + #[test] + fn forwarded_headers_can_include_authorization() { + let fields = vec!["authorization: Bearer fake-token".to_string()]; + let headers = forwarded_headers_from_pairs(&fields, &[("authorization", "Authorization")]); + assert_eq!( + headers.get("Authorization").map(String::as_str), + Some("Bearer fake-token") + ); + } + + #[test] + fn unescapes_ansi_c_quoted_segments() { + assert_eq!(unescape_shell_segment(r"a\nb\tc", true), "a\nb\tc"); + assert_eq!( + unescape_shell_segment(r#"say \"hi\""#, false), + r#"say "hi""# + ); + } +} diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 2633b751f4..2f17de6d23 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -75,6 +75,8 @@ pub enum ProviderId { ClinePass, LongCat, Neuralwatt, + ZoomMate, + QwenCloud, } impl ProviderId { @@ -145,6 +147,8 @@ impl ProviderId { ProviderId::ClinePass, ProviderId::LongCat, ProviderId::Neuralwatt, + ProviderId::ZoomMate, + ProviderId::QwenCloud, ] } @@ -215,6 +219,8 @@ impl ProviderId { ProviderId::ClinePass => "clinepass", ProviderId::LongCat => "longcat", ProviderId::Neuralwatt => "neuralwatt", + ProviderId::ZoomMate => "zoommate", + ProviderId::QwenCloud => "qwen-cloud", } } @@ -287,6 +293,8 @@ impl ProviderId { ProviderId::ClinePass => "ClinePass", ProviderId::LongCat => "LongCat", ProviderId::Neuralwatt => "Neuralwatt", + ProviderId::ZoomMate => "ZoomMate", + ProviderId::QwenCloud => "Qwen Cloud", } } @@ -362,6 +370,8 @@ impl ProviderId { ProviderId::ZenMux => None, ProviderId::ClinePass => None, ProviderId::Neuralwatt => None, + ProviderId::ZoomMate => Some("zoommate.zoom.us"), + ProviderId::QwenCloud => Some("qwencloud.com"), } } @@ -394,7 +404,7 @@ impl ProviderId { "jetbrains" | "jetbrains-ai" | "jetbrains ai" | "intellij" => { Some(ProviderId::JetBrains) } - "alibaba" | "tongyi" | "qianwen" | "qwen" => Some(ProviderId::Alibaba), + "alibaba" | "tongyi" | "qianwen" => Some(ProviderId::Alibaba), "alibabatokenplan" | "alibaba-token-plan" | "alibaba token plan" | "alibaba-token" | "bailian-token-plan" => Some(ProviderId::AlibabaTokenPlan), "nanogpt" | "nano-gpt" => Some(ProviderId::NanoGPT), @@ -443,6 +453,10 @@ impl ProviderId { "clinepass" | "cline-pass" | "cline" => Some(ProviderId::ClinePass), "longcat" | "long-cat" | "lc" => Some(ProviderId::LongCat), "neuralwatt" | "neural-watt" | "nw" | "neural" => Some(ProviderId::Neuralwatt), + "qwen-cloud" | "qwencloud" | "qwen" | "qwen-token-plan" | "qwen cloud" => { + Some(ProviderId::QwenCloud) + } + "zoommate" | "zoom-mate" | "zoom mate" => Some(ProviderId::ZoomMate), _ => None, } } @@ -655,7 +669,9 @@ pub fn cli_name_map() -> HashMap<&'static str, ProviderId> { map.insert("aws bedrock", ProviderId::Bedrock); map.insert("tongyi", ProviderId::Alibaba); map.insert("qianwen", ProviderId::Alibaba); - map.insert("qwen", ProviderId::Alibaba); + map.insert("qwen", ProviderId::QwenCloud); + map.insert("qwencloud", ProviderId::QwenCloud); + map.insert("qwen-token-plan", ProviderId::QwenCloud); map.insert("infini-ai", ProviderId::Infini); map.insert("pplx", ProviderId::Perplexity); map.insert("abacus-ai", ProviderId::Abacus); @@ -693,7 +709,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 64); + assert_eq!(all.len(), 66); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Kimi)); @@ -738,6 +754,8 @@ mod tests { assert!(all.contains(&ProviderId::ClinePass)); assert!(all.contains(&ProviderId::LongCat)); assert!(all.contains(&ProviderId::Neuralwatt)); + assert!(all.contains(&ProviderId::ZoomMate)); + assert!(all.contains(&ProviderId::QwenCloud)); } #[test] @@ -901,6 +919,34 @@ mod tests { ProviderId::from_cli_name("qianwen"), Some(ProviderId::Alibaba) ); - assert_eq!(ProviderId::from_cli_name("qwen"), Some(ProviderId::Alibaba)); + } + + #[test] + fn test_provider_id_qwen_cloud() { + assert_eq!(ProviderId::QwenCloud.cli_name(), "qwen-cloud"); + assert_eq!(ProviderId::QwenCloud.display_name(), "Qwen Cloud"); + assert_eq!(ProviderId::QwenCloud.cookie_domain(), Some("qwencloud.com")); + assert_eq!( + ProviderId::from_cli_name("qwen-cloud"), + Some(ProviderId::QwenCloud) + ); + assert_eq!( + ProviderId::from_cli_name("qwencloud"), + Some(ProviderId::QwenCloud) + ); + assert_eq!( + ProviderId::from_cli_name("qwen"), + Some(ProviderId::QwenCloud) + ); + assert_eq!( + ProviderId::from_cli_name("qwen-token-plan"), + Some(ProviderId::QwenCloud) + ); + assert_eq!( + ProviderId::from_cli_name("qwen cloud"), + Some(ProviderId::QwenCloud) + ); + // Bare "qwen" must not resolve to Alibaba Coding Plan. + assert_ne!(ProviderId::from_cli_name("qwen"), Some(ProviderId::Alibaba)); } } diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 3d4e2a8cdb..2d0fd93ad1 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -17,9 +17,9 @@ use crate::providers::{ LongCatProvider, ManusProvider, MiMoProvider, MiniMaxProvider, MistralProvider, NanoGPTProvider, NeuralwattProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, PerplexityProvider, PoeProvider, QoderProvider, - SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, VeniceProvider, - VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, ZaiProvider, ZedProvider, - ZenMuxProvider, + QwenCloudProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, + VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, + ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -92,6 +92,8 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::ClinePass => Box::new(ClinePassProvider::new()), ProviderId::LongCat => Box::new(LongCatProvider::new()), ProviderId::Neuralwatt => Box::new(NeuralwattProvider::new()), + ProviderId::ZoomMate => Box::new(ZoomMateProvider::new()), + ProviderId::QwenCloud => Box::new(QwenCloudProvider::new()), } } diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index d157b0993e..652b6af926 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -141,6 +141,14 @@ impl TokenAccountSupport { requires_manual_cookie_source: true, cookie_name: None, }), + ProviderId::ZoomMate => Some(TokenAccountSupport { + title: "Session tokens", + subtitle: "Store multiple ZoomMate Cookie headers or credits/status cURL captures.", + placeholder: "Cookie: ... or curl 'https://ai.zoom.us/.../credits/status' -H 'Authorization: Bearer ...'", + injection: TokenInjection::CookieHeader, + requires_manual_cookie_source: true, + cookie_name: None, + }), ProviderId::Mistral => Some(TokenAccountSupport { title: "Session tokens", subtitle: "Store multiple Mistral Cookie headers.", @@ -308,7 +316,8 @@ impl TokenAccountSupport { | ProviderId::Zed | ProviderId::CrossModel | ProviderId::LongCat - | ProviderId::Wayfinder => None, + | ProviderId::Wayfinder + | ProviderId::QwenCloud => None, } } diff --git a/rust/src/providers/alibabatokenplan.rs b/rust/src/providers/alibabatokenplan/mod.rs similarity index 74% rename from rust/src/providers/alibabatokenplan.rs rename to rust/src/providers/alibabatokenplan/mod.rs index f6585adacd..a069eb300f 100644 --- a/rust/src/providers/alibabatokenplan.rs +++ b/rust/src/providers/alibabatokenplan/mod.rs @@ -1,8 +1,16 @@ //! Alibaba Token Plan provider implementation. //! -//! Fetches Bailian token-plan credits from the same commerce endpoint used by -//! the upstream macOS provider. Authentication uses Bailian browser cookies or -//! a manually pasted cookie header. +//! Fetches Bailian / Model Studio token-plan credits from the same commerce +//! endpoints used by the upstream macOS provider. Authentication uses browser +//! cookies or a manually pasted cookie header. +//! +//! Team path: `GetSubscriptionSummary` / BssOpenAPI-V3 (+ optional sec_token). +//! Personal/Solo path: OneConsole personal token-plan APIs (no sec_token). + +mod personal; +mod region; + +pub use region::AlibabaTokenPlanRegion; use async_trait::async_trait; use chrono::{DateTime, NaiveDate, NaiveDateTime, TimeZone, Utc}; @@ -15,30 +23,40 @@ use crate::core::{ }; use crate::providers::browser_cookie_header; -const GATEWAY_BASE_URL: &str = "https://bailian.console.aliyun.com"; -const DASHBOARD_URL: &str = +use region::AlibabaTokenPlanRegion as Region; + +const DEFAULT_DASHBOARD_URL: &str = "https://bailian.console.aliyun.com/cn-beijing?tab=plan#/efm/subscription/token-plan"; -const TOKEN_PLAN_COMMODITY_CODE: &str = "sfm_tokenplanteams_dp_cn"; -const CURRENT_REGION_ID: &str = "cn-beijing"; -const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"; - -const COOKIE_DOMAINS: &[&str] = &[ - "bailian-cs.console.aliyun.com", - "bailian.console.aliyun.com", - "aliyun.com", -]; +pub(super) const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"; +pub(super) const LANGUAGE: &str = "en-US"; +pub(super) const PERSONAL_CONSOLE_PRODUCT: &str = "sfm_bailian"; +pub(super) const PERSONAL_USAGE_API: &str = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/usage"; +pub(super) const PERSONAL_SUBSCRIPTION_API: &str = + "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/subscription"; +pub(super) const PERSONAL_QUOTA_CONFIG_API: &str = + "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/quota-config"; + +const FIVE_HOUR_MINUTES: u32 = 5 * 60; +const WEEKLY_MINUTES: u32 = 7 * 24 * 60; +const LEGACY_MINUTES: u32 = 30 * 24 * 60; pub struct AlibabaTokenPlanProvider { metadata: ProviderMetadata, } #[derive(Debug, Clone, PartialEq)] -struct TokenPlanSnapshot { - plan_name: Option, - used_quota: Option, - total_quota: Option, - remaining_quota: Option, - resets_at: Option>, +pub(super) struct TokenPlanSnapshot { + pub(super) plan_name: Option, + pub(super) used_quota: Option, + pub(super) total_quota: Option, + pub(super) remaining_quota: Option, + pub(super) resets_at: Option>, + pub(super) five_hour_used_percent: Option, + pub(super) five_hour_total_quota: Option, + pub(super) five_hour_resets_at: Option>, + pub(super) weekly_used_percent: Option, + pub(super) weekly_total_quota: Option, + pub(super) weekly_resets_at: Option>, } impl AlibabaTokenPlanProvider { @@ -53,25 +71,47 @@ impl AlibabaTokenPlanProvider { supports_credits: false, default_enabled: false, is_primary: false, - dashboard_url: Some(DASHBOARD_URL), + dashboard_url: Some(DEFAULT_DASHBOARD_URL), status_page_url: Some("https://status.aliyun.com"), }, } } + fn resolve_region(ctx: &FetchContext) -> Region { + Region::from_settings_value(ctx.api_region.as_deref()) + } + async fn fetch_via_web(&self, ctx: &FetchContext) -> Result { - let cookie_header = Self::resolve_cookie_header(ctx)?; + let region = Self::resolve_region(ctx); + let cookie_header = Self::resolve_cookie_header(ctx, region)?; let client = crate::core::credentialed_http_client_builder() .timeout(std::time::Duration::from_secs(ctx.web_timeout.max(1))) .redirect(reqwest::redirect::Policy::limited(5)) .build() .map_err(|e| ProviderError::Other(e.to_string()))?; - let sec_token = Self::resolve_sec_token(&client, &cookie_header, ctx).await; + + let snapshot = if region.uses_personal_api() { + personal::fetch_personal_usage(&client, &cookie_header, region, ctx).await? + } else { + self.fetch_team_usage(&client, &cookie_header, region, ctx) + .await? + }; + Self::snapshot_to_usage(snapshot) + } + + async fn fetch_team_usage( + &self, + client: &reqwest::Client, + cookie_header: &str, + region: Region, + ctx: &FetchContext, + ) -> Result { + let sec_token = Self::resolve_sec_token(client, cookie_header, region, ctx).await; let mut form = vec![ ("product", "BssOpenAPI-V3".to_string()), ("action", "GetSubscriptionSummary".to_string()), - ("params", Self::request_params()), - ("region", CURRENT_REGION_ID.to_string()), + ("params", Self::team_request_params(region)), + ("region", region.current_region_id().to_string()), ]; if let Some(token) = sec_token .as_deref() @@ -80,18 +120,18 @@ impl AlibabaTokenPlanProvider { form.push(("sec_token", token.to_string())); } let mut request = client - .post(Self::quota_url()) - .header("Cookie", &cookie_header) + .post(Self::team_quota_url(region)) + .header("Cookie", cookie_header) .header("Accept", "*/*") .header("Content-Type", "application/x-www-form-urlencoded") - .header("Origin", "https://bailian.console.aliyun.com") - .header("Referer", DASHBOARD_URL) + .header("Origin", region.gateway_base_url()) + .header("Referer", region.dashboard_url()) .header("User-Agent", USER_AGENT) .header("X-Requested-With", "XMLHttpRequest") .form(&form); - if let Some(csrf) = cookie_value("login_aliyunid_csrf", &cookie_header) - .or_else(|| cookie_value("csrf", &cookie_header)) + if let Some(csrf) = cookie_value("login_aliyunid_csrf", cookie_header) + .or_else(|| cookie_value("csrf", cookie_header)) { request = request .header("x-xsrf-token", csrf.clone()) @@ -112,11 +152,10 @@ impl AlibabaTokenPlanProvider { ))); } - let snapshot = Self::parse_usage_snapshot(&body)?; - Self::snapshot_to_usage(snapshot) + Self::parse_usage_snapshot(&body) } - fn resolve_cookie_header(ctx: &FetchContext) -> Result { + fn resolve_cookie_header(ctx: &FetchContext, region: Region) -> Result { if let Some(raw) = ctx .manual_cookie_header .as_deref() @@ -135,17 +174,18 @@ impl AlibabaTokenPlanProvider { return Ok(header); } } - browser_cookie_header(COOKIE_DOMAINS) + browser_cookie_header(region.cookie_domains()) .and_then(|header| normalize_cookie_header(&header).ok_or(ProviderError::NoCookies)) } async fn resolve_sec_token( client: &reqwest::Client, cookie_header: &str, + region: Region, ctx: &FetchContext, ) -> Option { let response = client - .get(DASHBOARD_URL) + .get(region.dashboard_url()) .timeout(std::time::Duration::from_secs(ctx.web_timeout.clamp(1, 10))) .header("Cookie", cookie_header) .header( @@ -163,15 +203,16 @@ impl AlibabaTokenPlanProvider { extract_sec_token(&text).or_else(|| cookie_value("sec_token", cookie_header)) } - fn quota_url() -> String { + fn team_quota_url(region: Region) -> String { format!( - "{GATEWAY_BASE_URL}/data/api.json?action=GetSubscriptionSummary&product=BssOpenAPI-V3&_tag=" + "{}/data/api.json?action=GetSubscriptionSummary&product=BssOpenAPI-V3&_tag=", + region.gateway_base_url() ) } - fn request_params() -> String { + fn team_request_params(region: Region) -> String { serde_json::json!({ - "ProductCode": TOKEN_PLAN_COMMODITY_CODE, + "ProductCode": region.product_code(), }) .to_string() } @@ -190,7 +231,7 @@ impl AlibabaTokenPlanProvider { } })?; let expanded = expand_json_strings(value); - Self::throw_if_error_payload(&expanded)?; + throw_if_error_payload(&expanded)?; let instance = find_token_plan_instance(&expanded); let plan_name = instance @@ -228,77 +269,57 @@ impl AlibabaTokenPlanProvider { total_quota: total, remaining_quota: remaining, resets_at, + five_hour_used_percent: None, + five_hour_total_quota: None, + five_hour_resets_at: None, + weekly_used_percent: None, + weekly_total_quota: None, + weekly_resets_at: None, }) } - fn throw_if_error_payload(value: &Value) -> Result<(), ProviderError> { - if let Some(status) = find_first_i64(value, &["statusCode", "status_code", "code"]) - && status != 0 - && status != 200 - { - if status == 401 || status == 403 { - return Err(ProviderError::AuthRequired); - } - let message = - find_first_string(value, &["statusMessage", "status_msg", "message", "msg"]) - .unwrap_or_else(|| format!("status code {status}")); - return Err(ProviderError::Other(format!( - "Alibaba Token Plan API error: {message}" - ))); - } - - if let Some(success) = find_first_bool(value, &["success", "Success"]) - && !success - { - let message = find_first_string(value, &["message", "msg", "Message", "errorMessage"]) - .unwrap_or_else(|| "request failed".to_string()); - let lower = message.to_lowercase(); - if lower.contains("needlogin") - || lower.contains("login") - || lower.contains("log in") - || lower.contains("unauthorized") - { - return Err(ProviderError::AuthRequired); - } - return Err(ProviderError::Other(format!( - "Alibaba Token Plan API error: {message}" - ))); - } - - let code = find_first_string(value, &["code", "status", "statusCode"]) - .unwrap_or_default() - .to_lowercase(); - let message = find_first_string(value, &["message", "msg", "statusMessage"]) - .unwrap_or_default() - .to_lowercase(); - if code.contains("needlogin") - || code.contains("login") - || message.contains("log in") - || message.contains("login") - { - return Err(ProviderError::AuthRequired); - } - Ok(()) - } - - fn snapshot_to_usage(snapshot: TokenPlanSnapshot) -> Result { - let Some(used_percent) = used_percent( - snapshot.used_quota, - snapshot.total_quota, - snapshot.remaining_quota, - ) else { - return Err(ProviderError::Parse( - "Alibaba Token Plan quota totals missing".into(), - )); - }; - let detail = quota_detail( + pub(super) fn snapshot_to_usage( + snapshot: TokenPlanSnapshot, + ) -> Result { + let five_hour = snapshot.five_hour_used_percent.map(|percent| { + RateWindow::with_details( + percent, + Some(FIVE_HOUR_MINUTES), + snapshot.five_hour_resets_at, + quota_detail_percent(percent, snapshot.five_hour_total_quota), + ) + }); + let legacy = used_percent( snapshot.used_quota, snapshot.total_quota, snapshot.remaining_quota, - ); - let primary = - RateWindow::with_details(used_percent, Some(30 * 24 * 60), snapshot.resets_at, detail); + ) + .map(|percent| { + RateWindow::with_details( + percent, + Some(LEGACY_MINUTES), + snapshot.resets_at, + quota_detail( + snapshot.used_quota, + snapshot.total_quota, + snapshot.remaining_quota, + ), + ) + }); + let primary = five_hour.or(legacy).ok_or_else(|| { + ProviderError::Parse("Alibaba Token Plan quota totals missing".into()) + })?; let mut usage = UsageSnapshot::new(primary); + + if let Some(weekly_percent) = snapshot.weekly_used_percent { + usage = usage.with_secondary(RateWindow::with_details( + weekly_percent, + Some(WEEKLY_MINUTES), + snapshot.weekly_resets_at, + quota_detail_percent(weekly_percent, snapshot.weekly_total_quota), + )); + } + if let Some(plan) = snapshot.plan_name.filter(|plan| !plan.trim().is_empty()) { usage = usage.with_login_method(plan); } @@ -343,6 +364,55 @@ impl Provider for AlibabaTokenPlanProvider { } } +pub(super) fn throw_if_error_payload(value: &Value) -> Result<(), ProviderError> { + if let Some(status) = find_first_i64(value, &["statusCode", "status_code", "code"]) + && status != 0 + && status != 200 + { + if status == 401 || status == 403 { + return Err(ProviderError::AuthRequired); + } + let message = find_first_string(value, &["statusMessage", "status_msg", "message", "msg"]) + .unwrap_or_else(|| format!("status code {status}")); + return Err(ProviderError::Other(format!( + "Alibaba Token Plan API error: {message}" + ))); + } + + if let Some(success) = find_first_bool(value, &["success", "Success"]) + && !success + { + let message = find_first_string(value, &["message", "msg", "Message", "errorMessage"]) + .unwrap_or_else(|| "request failed".to_string()); + let lower = message.to_lowercase(); + if lower.contains("needlogin") + || lower.contains("login") + || lower.contains("log in") + || lower.contains("unauthorized") + { + return Err(ProviderError::AuthRequired); + } + return Err(ProviderError::Other(format!( + "Alibaba Token Plan API error: {message}" + ))); + } + + let code = find_first_string(value, &["code", "status", "statusCode"]) + .unwrap_or_default() + .to_lowercase(); + let message = find_first_string(value, &["message", "msg", "statusMessage"]) + .unwrap_or_default() + .to_lowercase(); + if code.contains("needlogin") + || code.contains("login") + || message.contains("log in") + || message.contains("login") + { + return Err(ProviderError::AuthRequired); + } + Ok(()) +} + fn normalize_cookie_header(raw: &str) -> Option { let mut header = raw.trim(); if header @@ -354,7 +424,7 @@ fn normalize_cookie_header(raw: &str) -> Option { (!header.is_empty() && header.contains('=')).then(|| header.to_string()) } -fn cookie_value(name: &str, cookie_header: &str) -> Option { +pub(super) fn cookie_value(name: &str, cookie_header: &str) -> Option { cookie_header.split(';').find_map(|part| { let (key, value) = part.trim().split_once('=')?; (key.trim() == name) @@ -385,13 +455,13 @@ fn extract_sec_token(html: &str) -> Option { None } -fn is_likely_login_html(data: &[u8]) -> bool { +pub(super) fn is_likely_login_html(data: &[u8]) -> bool { let text = String::from_utf8_lossy(data).to_lowercase(); text.contains(" Value { +pub(super) fn expand_json_strings(value: Value) -> Value { match value { Value::Array(values) => Value::Array(values.into_iter().map(expand_json_strings).collect()), Value::Object(map) => Value::Object( @@ -408,6 +478,35 @@ fn expand_json_strings(value: Value) -> Value { } } +pub(super) fn percentage_points(ratio: Option) -> Option { + let ratio = ratio.filter(|v| v.is_finite())?; + Some((ratio.clamp(0.0, 1.0) * 100.0).clamp(0.0, 100.0)) +} + +pub(super) fn number_field(value: &Value, key: &str) -> Option { + value.as_object().and_then(|map| parse_f64(map.get(key))) +} + +pub(super) fn date_field(value: &Value, key: &str) -> Option> { + value.as_object().and_then(|map| parse_date(map.get(key))) +} + +pub(super) fn find_object_containing_any_of(value: &Value, keys: &[&str]) -> Option { + match value { + Value::Object(map) => { + if keys.iter().any(|key| map.contains_key(*key)) { + return Some(Value::Object(map.clone())); + } + map.values() + .find_map(|nested| find_object_containing_any_of(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_object_containing_any_of(nested, keys)), + _ => None, + } +} + const PLAN_NAME_KEYS: &[&str] = &[ "planName", "plan_name", @@ -603,7 +702,7 @@ fn first_string(value: &Value, keys: &[&str]) -> Option { keys.iter().find_map(|key| parse_string(map.get(*key))) } -fn find_first_string(value: &Value, keys: &[&str]) -> Option { +pub(super) fn find_first_string(value: &Value, keys: &[&str]) -> Option { match value { Value::Object(map) => first_string(value, keys).or_else(|| { map.values() @@ -785,6 +884,16 @@ fn quota_detail(used: Option, total: Option, remaining: Option) - remaining.map(|remaining| format!("{} credits left", format_quota(remaining))) } +fn quota_detail_percent(used_percent: f64, total: Option) -> Option { + let total = total.filter(|total| *total > 0.0)?; + let used = total * used_percent / 100.0; + Some(format!( + "{} / {} credits used", + format_quota(used), + format_quota(total) + )) +} + fn format_quota(value: f64) -> String { if (value.round() - value).abs() < f64::EPSILON { format_count(value.round() as i64) @@ -931,4 +1040,18 @@ mod tests { Some("xyz".to_string()) ); } + + #[test] + fn default_region_cn_team_urls_match_legacy() { + let region = Region::Cn; + assert_eq!( + AlibabaTokenPlanProvider::team_quota_url(region), + "https://bailian.console.aliyun.com/data/api.json?action=GetSubscriptionSummary&product=BssOpenAPI-V3&_tag=" + ); + assert_eq!( + AlibabaTokenPlanProvider::team_request_params(region), + serde_json::json!({"ProductCode": "sfm_tokenplanteams_dp_cn"}).to_string() + ); + assert_eq!(region.current_region_id(), "cn-beijing"); + } } diff --git a/rust/src/providers/alibabatokenplan/personal.rs b/rust/src/providers/alibabatokenplan/personal.rs new file mode 100644 index 0000000000..aa9738864b --- /dev/null +++ b/rust/src/providers/alibabatokenplan/personal.rs @@ -0,0 +1,371 @@ +//! Alibaba Token Plan Personal/Solo OneConsole path (upstream 0.46.0). + +use serde_json::{Map, Value, json}; +use uuid::Uuid; + +use super::region::AlibabaTokenPlanRegion; +use super::{ + LANGUAGE, PERSONAL_CONSOLE_PRODUCT, PERSONAL_QUOTA_CONFIG_API, PERSONAL_SUBSCRIPTION_API, + PERSONAL_USAGE_API, TokenPlanSnapshot, USER_AGENT, cookie_value, date_field, + expand_json_strings, find_object_containing_any_of, is_likely_login_html, number_field, + percentage_points, throw_if_error_payload, +}; +use crate::core::{FetchContext, ProviderError}; + +pub(super) async fn fetch_personal_usage( + client: &reqwest::Client, + cookie_header: &str, + region: AlibabaTokenPlanRegion, + ctx: &FetchContext, +) -> Result { + let usage_body = post_personal_api( + client, + PERSONAL_USAGE_API, + Map::new(), + cookie_header, + region, + ctx, + ) + .await?; + + let mut subscription_params = Map::new(); + subscription_params.insert( + "commodityCode".into(), + Value::String(region.product_code().to_string()), + ); + let subscription_body = post_personal_api_optional( + client, + PERSONAL_SUBSCRIPTION_API, + subscription_params, + cookie_header, + region, + ctx, + ) + .await; + + let quota_config_body = post_personal_api_optional( + client, + PERSONAL_QUOTA_CONFIG_API, + Map::new(), + cookie_header, + region, + ctx, + ) + .await; + + parse_personal_usage( + &usage_body, + subscription_body.as_deref(), + quota_config_body.as_deref(), + ) +} + +async fn post_personal_api( + client: &reqwest::Client, + api: &str, + data_parameters: Map, + cookie_header: &str, + region: AlibabaTokenPlanRegion, + ctx: &FetchContext, +) -> Result, ProviderError> { + let url = personal_api_url(api, region); + let params_json = build_personal_params_json(api, data_parameters, cookie_header, region); + let form = [ + ("product", PERSONAL_CONSOLE_PRODUCT.to_string()), + ("action", region.personal_api_action().to_string()), + ("region", region.current_region_id().to_string()), + ("language", LANGUAGE.to_string()), + ("params", params_json), + ]; + + let mut request = client + .post(&url) + .timeout(std::time::Duration::from_secs(ctx.web_timeout.max(1))) + .header("Cookie", cookie_header) + .header("Accept", "application/json, text/plain, */*") + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Origin", region.gateway_base_url()) + .header("Referer", region.dashboard_url()) + .header("User-Agent", USER_AGENT) + .header("X-Requested-With", "XMLHttpRequest") + .form(&form); + + if let Some(csrf) = cookie_value("login_aliyunid_csrf", cookie_header) + .or_else(|| cookie_value("csrf", cookie_header)) + { + request = request + .header("x-xsrf-token", csrf.clone()) + .header("x-csrf-token", csrf); + } + + let response = request.send().await?; + let status = response.status(); + let body = response.bytes().await?; + if !status.is_success() { + if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + return Err(ProviderError::Other(format!( + "Alibaba Token Plan Personal API error: HTTP {status}" + ))); + } + Ok(body.to_vec()) +} + +async fn post_personal_api_optional( + client: &reqwest::Client, + api: &str, + data_parameters: Map, + cookie_header: &str, + region: AlibabaTokenPlanRegion, + ctx: &FetchContext, +) -> Option> { + post_personal_api(client, api, data_parameters, cookie_header, region, ctx) + .await + .ok() +} + +fn personal_api_url(api: &str, region: AlibabaTokenPlanRegion) -> String { + format!( + "{}/data/api.json?action={}&product={}&api={api}&_v=undefined", + region.quota_base_url(), + region.personal_api_action(), + PERSONAL_CONSOLE_PRODUCT, + ) +} + +fn build_personal_params_json( + api: &str, + mut data_parameters: Map, + cookie_header: &str, + region: AlibabaTokenPlanRegion, +) -> String { + let dashboard = region.dashboard_url(); + let domain = dashboard + .trim_start_matches("https://") + .trim_start_matches("http://") + .split('/') + .next() + .unwrap_or_default(); + + let mut cornerstone = Map::new(); + cornerstone.insert( + "feTraceId".into(), + Value::String(Uuid::new_v4().to_string().to_lowercase()), + ); + cornerstone.insert("feURL".into(), Value::String(dashboard.to_string())); + cornerstone.insert("protocol".into(), Value::String("V2".into())); + cornerstone.insert("console".into(), Value::String("ONE_CONSOLE".into())); + cornerstone.insert("productCode".into(), Value::String("p_efm".into())); + cornerstone.insert("switchAgent".into(), json!(1_233_135)); + cornerstone.insert("switchUserType".into(), json!(3)); + cornerstone.insert("domain".into(), Value::String(domain.to_string())); + cornerstone.insert( + "consoleSite".into(), + Value::String(region.personal_console_site().to_string()), + ); + cornerstone.insert("userNickName".into(), Value::String(String::new())); + cornerstone.insert("userPrincipalName".into(), Value::String(String::new())); + cornerstone.insert("xsp_lang".into(), Value::String(LANGUAGE.into())); + if let Some(cna) = cookie_value("cna", cookie_header) { + cornerstone.insert("X-Anonymous-Id".into(), Value::String(cna)); + } + + data_parameters.insert("cornerstoneParam".into(), Value::Object(cornerstone)); + + json!({ + "Api": api, + "V": "1.0", + "Data": Value::Object(data_parameters), + }) + .to_string() +} + +pub(super) fn parse_personal_usage( + usage_data: &[u8], + subscription_data: Option<&[u8]>, + quota_config_data: Option<&[u8]>, +) -> Result { + if usage_data.is_empty() { + return Err(ProviderError::Parse( + "Empty Alibaba Token Plan Personal response".into(), + )); + } + + let value: Value = serde_json::from_slice(usage_data).map_err(|_| { + if is_likely_login_html(usage_data) { + ProviderError::AuthRequired + } else { + ProviderError::Parse("Invalid Alibaba Token Plan Personal JSON response".into()) + } + })?; + let expanded = expand_json_strings(value); + throw_if_error_payload(&expanded)?; + + let usage = + find_object_containing_any_of(&expanded, &["per5HourPercentage", "per1WeekPercentage"]) + .ok_or_else(|| { + ProviderError::Parse("Missing Alibaba Token Plan Personal usage windows".into()) + })?; + + let five_hour = percentage_points(number_field(&usage, "per5HourPercentage")); + let weekly = percentage_points(number_field(&usage, "per1WeekPercentage")); + if five_hour.is_none() && weekly.is_none() { + return Err(ProviderError::Parse( + "Missing Alibaba Token Plan Personal usage windows".into(), + )); + } + + let plan_code = subscription_data.and_then(plan_code_from_bytes); + let plan_name = plan_code + .as_deref() + .map(display_plan_name) + .or_else(|| Some("Personal".to_string())); + let quota = quota_config_data + .zip(plan_code.as_ref()) + .and_then(|(data, code)| quota_totals_from_bytes(data, code)); + + Ok(TokenPlanSnapshot { + plan_name, + used_quota: None, + total_quota: None, + remaining_quota: None, + resets_at: None, + five_hour_used_percent: five_hour, + five_hour_total_quota: quota.map(|q| q.0).unwrap_or(None), + five_hour_resets_at: date_field(&usage, "per5HourResetTime"), + weekly_used_percent: weekly, + weekly_total_quota: quota.map(|q| q.1).unwrap_or(None), + weekly_resets_at: date_field(&usage, "per1WeekResetTime"), + }) +} + +fn plan_code_from_bytes(data: &[u8]) -> Option { + let value: Value = serde_json::from_slice(data).ok()?; + let expanded = expand_json_strings(value); + let plan = find_object_containing_any_of( + &expanded, + &["specCode", "spec_code", "planName", "plan_name"], + )?; + for key in ["specCode", "spec_code", "planName", "plan_name"] { + if let Some(raw) = plan.get(key).and_then(Value::as_str) { + let normalized = raw.trim().to_lowercase(); + if !normalized.is_empty() { + return Some(normalized); + } + } + } + None +} + +fn display_plan_name(plan_code: &str) -> String { + match plan_code { + "lite" => "Lite".into(), + "standard" => "Standard".into(), + "pro" => "Pro".into(), + "max" => "Max".into(), + other => other.to_string(), + } +} + +fn quota_totals_from_bytes(data: &[u8], plan_code: &str) -> Option<(Option, Option)> { + let value: Value = serde_json::from_slice(data).ok()?; + let expanded = expand_json_strings(value); + let quota = find_first_value_for_key(&expanded, plan_code)? + .as_object()? + .clone(); + let five_hour = number_field(&Value::Object(quota.clone()), "five_hour") + .or_else(|| number_field(&Value::Object(quota.clone()), "fiveHour")); + let weekly = number_field(&Value::Object(quota), "weekly"); + if five_hour.is_none() && weekly.is_none() { + return None; + } + Some((five_hour, weekly)) +} + +fn find_first_value_for_key(value: &Value, key: &str) -> Option { + match value { + Value::Object(map) => { + if let Some(nested) = map.get(key) { + return Some(nested.clone()); + } + map.values() + .find_map(|nested| find_first_value_for_key(nested, key)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_value_for_key(nested, key)), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::UsageSnapshot; + use crate::providers::alibabatokenplan::AlibabaTokenPlanProvider; + + #[test] + fn parses_personal_usage_fixture_with_plan_name() { + let usage = serde_json::json!({ + "code": "200", + "data": { + "DataV2": { + "data": { + "success": true, + "data": { + "per5HourPercentage": 0.0009973083333333333, + "per5HourResetTime": 1784813220000_i64, + "per1WeekPercentage": 0.0003014725, + "per1WeekResetTime": 1785234900000_i64 + } + }, + "success": true, + "httpStatus": 200 + } + }, + "successResponse": true + }); + let subscription = serde_json::json!({ + "code": "200", + "data": { + "specCode": "pro", + "planName": "pro" + } + }); + let quota_config = serde_json::json!({ + "code": "200", + "data": { + "pro": { + "five_hour": 1000000, + "weekly": 5000000 + } + } + }); + + let snapshot = parse_personal_usage( + usage.to_string().as_bytes(), + Some(subscription.to_string().as_bytes()), + Some(quota_config.to_string().as_bytes()), + ) + .unwrap(); + + assert_eq!(snapshot.plan_name.as_deref(), Some("Pro")); + assert!((snapshot.five_hour_used_percent.unwrap() - 0.09973083333333333).abs() < 1e-9); + assert!((snapshot.weekly_used_percent.unwrap() - 0.03014725).abs() < 1e-9); + assert_eq!(snapshot.five_hour_total_quota, Some(1_000_000.0)); + assert_eq!(snapshot.weekly_total_quota, Some(5_000_000.0)); + assert!(snapshot.five_hour_resets_at.is_some()); + assert!(snapshot.weekly_resets_at.is_some()); + + let usage_snap: UsageSnapshot = + AlibabaTokenPlanProvider::snapshot_to_usage(snapshot).unwrap(); + assert!((usage_snap.primary.used_percent - 0.09973083333333333).abs() < 1e-9); + assert_eq!(usage_snap.primary.window_minutes, Some(300)); + assert_eq!( + usage_snap.secondary.as_ref().and_then(|w| w.window_minutes), + Some(10080) + ); + assert_eq!(usage_snap.login_method.as_deref(), Some("Pro")); + } +} diff --git a/rust/src/providers/alibabatokenplan/region.rs b/rust/src/providers/alibabatokenplan/region.rs new file mode 100644 index 0000000000..c1b2402739 --- /dev/null +++ b/rust/src/providers/alibabatokenplan/region.rs @@ -0,0 +1,220 @@ +//! Alibaba Token Plan API region (upstream 0.46.0). + +use serde::{Deserialize, Serialize}; + +/// Persisted region id: `"cn" | "intl" | "cn-personal" | "intl-personal"`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum AlibabaTokenPlanRegion { + #[default] + #[serde(rename = "cn")] + Cn, + #[serde(rename = "intl")] + Intl, + #[serde(rename = "cn-personal")] + CnPersonal, + #[serde(rename = "intl-personal")] + IntlPersonal, +} + +impl AlibabaTokenPlanRegion { + /// Serde/settings raw string exactly as persisted. + pub fn as_str(self) -> &'static str { + match self { + Self::Cn => "cn", + Self::Intl => "intl", + Self::CnPersonal => "cn-personal", + Self::IntlPersonal => "intl-personal", + } + } + + /// Parse settings / `FetchContext.api_region` value. Unknown → cn. + pub fn from_settings_value(value: Option<&str>) -> Self { + match value.unwrap_or("").trim().to_ascii_lowercase().as_str() { + "intl" | "international" | "ap-southeast-1" => Self::Intl, + "cn-personal" | "cn_personal" | "personal" | "solo" => Self::CnPersonal, + "intl-personal" + | "intl_personal" + | "international-personal" + | "international_personal" => Self::IntlPersonal, + "cn" | "china" | "cn-beijing" | "beijing" | "" => Self::Cn, + other => { + tracing::warn!( + value = other, + "Unknown Alibaba Token Plan region; falling back to cn" + ); + Self::Cn + } + } + } + + pub fn gateway_base_url(self) -> &'static str { + match self { + Self::Intl | Self::IntlPersonal => "https://modelstudio.console.alibabacloud.com", + Self::Cn | Self::CnPersonal => "https://bailian.console.aliyun.com", + } + } + + pub fn quota_base_url(self) -> &'static str { + match self { + Self::Cn | Self::Intl => self.gateway_base_url(), + Self::CnPersonal => "https://bailian-cs.console.aliyun.com", + Self::IntlPersonal => "https://bailian-singapore-cs.alibabacloud.com", + } + } + + pub fn current_region_id(self) -> &'static str { + match self { + Self::Intl | Self::IntlPersonal => "ap-southeast-1", + Self::Cn | Self::CnPersonal => "cn-beijing", + } + } + + pub fn product_code(self) -> &'static str { + match self { + Self::Cn => "sfm_tokenplanteams_dp_cn", + Self::Intl => "sfm_tokenplanteams_dp_intl", + Self::CnPersonal => "sfm_tokenplansolo_public_cn", + Self::IntlPersonal => "sfm_tokenplansolo_public_intl", + } + } + + pub fn uses_personal_api(self) -> bool { + matches!(self, Self::CnPersonal | Self::IntlPersonal) + } + + pub fn personal_api_action(self) -> &'static str { + match self { + Self::Intl | Self::IntlPersonal => "IntlBroadScopeAspnGateway", + Self::Cn | Self::CnPersonal => "BroadScopeAspnGateway", + } + } + + /// Upstream spelling is intentional (`ALBABACLOUD`). + pub fn personal_console_site(self) -> &'static str { + match self { + Self::Intl | Self::IntlPersonal => "MODELSTUDIO_ALBABACLOUD", + Self::Cn | Self::CnPersonal => "BAILIAN_ALIYUN", + } + } + + pub fn dashboard_url(self) -> &'static str { + match self { + Self::Cn => { + "https://bailian.console.aliyun.com/cn-beijing?tab=plan#/efm/subscription/token-plan" + } + Self::Intl => { + "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=plan#/efm/subscription/token-plan" + } + Self::CnPersonal => { + "https://bailian.console.aliyun.com/cn-beijing?tab=plan#/efm/subscription/token-plan/personal" + } + Self::IntlPersonal => { + "https://modelstudio.console.alibabacloud.com/ap-southeast-1/?tab=plan#/efm/subscription/token-plan/personal" + } + } + } + + pub fn cookie_domains(self) -> &'static [&'static str] { + match self { + Self::Cn | Self::CnPersonal => CN_COOKIE_DOMAINS, + Self::Intl | Self::IntlPersonal => INTL_COOKIE_DOMAINS, + } + } +} + +/// Existing cn superset (byte-stable for default region). +const CN_COOKIE_DOMAINS: &[&str] = &[ + "bailian-cs.console.aliyun.com", + "bailian.console.aliyun.com", + "aliyun.com", +]; + +/// Bailian + Model Studio hosts for intl regions. +const INTL_COOKIE_DOMAINS: &[&str] = &[ + "bailian-cs.console.aliyun.com", + "bailian.console.aliyun.com", + "aliyun.com", + "modelstudio.console.alibabacloud.com", + "bailian-singapore-cs.alibabacloud.com", + "alibabacloud.com", +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn region_serde_roundtrip() { + for region in [ + AlibabaTokenPlanRegion::Cn, + AlibabaTokenPlanRegion::Intl, + AlibabaTokenPlanRegion::CnPersonal, + AlibabaTokenPlanRegion::IntlPersonal, + ] { + let raw = serde_json::to_string(®ion).unwrap(); + assert_eq!(raw, format!("\"{}\"", region.as_str())); + let back: AlibabaTokenPlanRegion = serde_json::from_str(&raw).unwrap(); + assert_eq!(back, region); + } + } + + #[test] + fn region_gateway_and_commodity_mapping() { + let cn = AlibabaTokenPlanRegion::Cn; + assert_eq!(cn.gateway_base_url(), "https://bailian.console.aliyun.com"); + assert_eq!(cn.quota_base_url(), cn.gateway_base_url()); + assert_eq!(cn.product_code(), "sfm_tokenplanteams_dp_cn"); + assert_eq!(cn.current_region_id(), "cn-beijing"); + assert!(!cn.uses_personal_api()); + + let intl = AlibabaTokenPlanRegion::Intl; + assert_eq!( + intl.gateway_base_url(), + "https://modelstudio.console.alibabacloud.com" + ); + assert_eq!(intl.quota_base_url(), intl.gateway_base_url()); + assert_eq!(intl.product_code(), "sfm_tokenplanteams_dp_intl"); + assert_eq!(intl.current_region_id(), "ap-southeast-1"); + assert!(!intl.uses_personal_api()); + + let cn_personal = AlibabaTokenPlanRegion::CnPersonal; + assert_eq!( + cn_personal.gateway_base_url(), + "https://bailian.console.aliyun.com" + ); + assert_eq!( + cn_personal.quota_base_url(), + "https://bailian-cs.console.aliyun.com" + ); + assert_eq!(cn_personal.product_code(), "sfm_tokenplansolo_public_cn"); + assert_eq!(cn_personal.current_region_id(), "cn-beijing"); + assert!(cn_personal.uses_personal_api()); + assert_eq!(cn_personal.personal_api_action(), "BroadScopeAspnGateway"); + assert_eq!(cn_personal.personal_console_site(), "BAILIAN_ALIYUN"); + + let intl_personal = AlibabaTokenPlanRegion::IntlPersonal; + assert_eq!( + intl_personal.gateway_base_url(), + "https://modelstudio.console.alibabacloud.com" + ); + assert_eq!( + intl_personal.quota_base_url(), + "https://bailian-singapore-cs.alibabacloud.com" + ); + assert_eq!( + intl_personal.product_code(), + "sfm_tokenplansolo_public_intl" + ); + assert_eq!(intl_personal.current_region_id(), "ap-southeast-1"); + assert!(intl_personal.uses_personal_api()); + assert_eq!( + intl_personal.personal_api_action(), + "IntlBroadScopeAspnGateway" + ); + assert_eq!( + intl_personal.personal_console_site(), + "MODELSTUDIO_ALBABACLOUD" + ); + } +} diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index f726111d2c..b2b7225f5e 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -55,6 +55,7 @@ pub mod openrouter; pub mod perplexity; pub mod poe; pub mod qoder; +pub mod qwencloud; pub mod sakana; pub mod stepfun; pub mod sub2api; @@ -67,12 +68,13 @@ pub mod windsurf; pub mod zai; pub mod zed; pub mod zenmux; +pub mod zoommate; // Re-export provider implementations pub use abacus::AbacusProvider; pub use aiand::AiAndProvider; pub use alibaba::{AlibabaProvider, AlibabaRegion}; -pub use alibabatokenplan::AlibabaTokenPlanProvider; +pub use alibabatokenplan::{AlibabaTokenPlanProvider, AlibabaTokenPlanRegion}; pub use amp::AmpProvider; pub use antigravity::AntigravityProvider; pub use augment::AugmentProvider; @@ -121,6 +123,7 @@ pub use openrouter::OpenRouterProvider; pub use perplexity::PerplexityProvider; pub use poe::PoeProvider; pub use qoder::QoderProvider; +pub use qwencloud::QwenCloudProvider; pub use sakana::SakanaProvider; pub use stepfun::StepFunProvider; pub use sub2api::Sub2ApiProvider; @@ -133,6 +136,7 @@ pub use windsurf::WindsurfProvider; pub use zai::ZaiProvider; pub use zed::ZedProvider; pub use zenmux::ZenMuxProvider; +pub use zoommate::ZoomMateProvider; pub(crate) fn browser_cookie_header( domains: &[&str], diff --git a/rust/src/providers/qwencloud/mod.rs b/rust/src/providers/qwencloud/mod.rs new file mode 100644 index 0000000000..b30beccbdc --- /dev/null +++ b/rust/src/providers/qwencloud/mod.rs @@ -0,0 +1,1391 @@ +//! Qwen Cloud personal token-plan provider (upstream 0.46.0). +//! +//! Cookie-authenticated OneConsole gateway: +//! `POST https://cs-data.qwencloud.com/data/api.json?...` +//! with `IntlBroadScopeAspnGateway` / `sfm_bailian`. + +use async_trait::async_trait; +use chrono::{DateTime, NaiveDate, NaiveDateTime, TimeZone, Utc}; +use regex_lite::Regex; +use serde_json::{Map, Value, json}; +use uuid::Uuid; + +use crate::core::{ + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, + RateWindow, SourceMode, UsageSnapshot, +}; +use crate::providers::browser_cookie_header; + +const GATEWAY_BASE_URL: &str = "https://home.qwencloud.com"; +const DATA_GATEWAY_BASE_URL: &str = "https://cs-data.qwencloud.com"; +const DASHBOARD_URL: &str = "https://home.qwencloud.com/billing/subscription/token-plan-individual"; +const STATUS_PAGE_URL: &str = "https://status.alibabacloud.com"; +const PRODUCT_CODE: &str = "sfm_tokenplansolo_public_intl"; +const CONSOLE_PRODUCT: &str = "sfm_bailian"; +const CONSOLE_ACTION: &str = "IntlBroadScopeAspnGateway"; +const USAGE_API: &str = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/usage"; +const SUBSCRIPTION_API: &str = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/subscription"; +const QUOTA_CONFIG_API: &str = "zeldaHttp.apikeyMgr./tokenplan/personal/api/v2/quota-config"; +const REGION: &str = "ap-southeast-1"; +const LANGUAGE: &str = "en-US"; +const USER_INFO_PATH: &str = "/tool/user/info.json"; +const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"; + +const COOKIE_DOMAINS: &[&str] = &[ + "qwencloud.com", + "home.qwencloud.com", + "account.qwencloud.com", + "signin.qwencloud.com", + "www.qwencloud.com", + "cs-data.qwencloud.com", + "alibabacloud.com", + "account.alibabacloud.com", + "aliyun.com", + "console.aliyun.com", +]; + +const FIVE_HOUR_MINUTES: u32 = 5 * 60; +const WEEKLY_MINUTES: u32 = 7 * 24 * 60; +const LEGACY_MINUTES: u32 = 30 * 24 * 60; + +pub struct QwenCloudProvider { + metadata: ProviderMetadata, +} + +#[derive(Debug, Clone, PartialEq)] +struct QwenCloudSnapshot { + plan_name: Option, + used_quota: Option, + total_quota: Option, + remaining_quota: Option, + resets_at: Option>, + five_hour_used_percent: Option, + five_hour_total_quota: Option, + five_hour_resets_at: Option>, + weekly_used_percent: Option, + weekly_total_quota: Option, + weekly_resets_at: Option>, +} + +impl QwenCloudProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::QwenCloud, + display_name: "Qwen Cloud", + session_label: "5-hour", + weekly_label: "Weekly", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some(DASHBOARD_URL), + status_page_url: Some(STATUS_PAGE_URL), + }, + } + } + + async fn fetch_via_web(&self, ctx: &FetchContext) -> Result { + let cookie_header = Self::resolve_cookie_header(ctx)?; + let client = crate::core::credentialed_http_client_builder() + .timeout(std::time::Duration::from_secs(ctx.web_timeout.max(1))) + .redirect(reqwest::redirect::Policy::limited(5)) + .build() + .map_err(|e| ProviderError::Other(e.to_string()))?; + + let sec_token = Self::resolve_sec_token(&client, &cookie_header, ctx) + .await + .ok_or(ProviderError::AuthRequired)?; + + let usage_body = Self::post_api( + &client, + USAGE_API, + Map::new(), + &sec_token, + &cookie_header, + ctx, + ) + .await?; + + let subscription_body = Self::post_api_optional( + &client, + SUBSCRIPTION_API, + { + let mut data = Map::new(); + data.insert( + "commodityCode".into(), + Value::String(PRODUCT_CODE.to_string()), + ); + data + }, + &sec_token, + &cookie_header, + ctx, + ) + .await; + + let quota_config_body = Self::post_api_optional( + &client, + QUOTA_CONFIG_API, + Map::new(), + &sec_token, + &cookie_header, + ctx, + ) + .await; + + let snapshot = Self::parse( + &usage_body, + subscription_body.as_deref(), + quota_config_body.as_deref(), + )?; + Self::snapshot_to_usage(snapshot) + } + + fn resolve_cookie_header(ctx: &FetchContext) -> Result { + if let Some(raw) = ctx + .manual_cookie_header + .as_deref() + .and_then(normalize_cookie_header) + { + return Ok(raw); + } + for env_name in ["QWEN_CLOUD_COOKIE", "QWEN_CLOUD_COOKIE_HEADER"] { + if let Ok(raw) = std::env::var(env_name) + && let Some(header) = normalize_cookie_header(&raw) + { + return Ok(header); + } + } + browser_cookie_header(COOKIE_DOMAINS) + .and_then(|header| normalize_cookie_header(&header).ok_or(ProviderError::NoCookies)) + } + + async fn resolve_sec_token( + client: &reqwest::Client, + cookie_header: &str, + ctx: &FetchContext, + ) -> Option { + let timeout = std::time::Duration::from_secs(ctx.web_timeout.clamp(1, 20)); + + // 1. Dashboard HTML (freshest OneConsole inject). + if let Ok(response) = client + .get(DASHBOARD_URL) + .timeout(timeout) + .header("Cookie", cookie_header) + .header( + "Accept", + "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8", + ) + .header("User-Agent", USER_AGENT) + .send() + .await + && response.status().is_success() + && let Ok(text) = response.text().await + { + if looks_like_login_page(&text) { + return None; + } + if let Some(token) = extract_sec_token(&text) { + return Some(token); + } + } + + // 2. Cookie fallback. + if let Some(token) = cookie_value("sec_token", cookie_header) { + return Some(token); + } + + // 3. User-info JSON endpoint. + let user_info_url = format!("{GATEWAY_BASE_URL}{USER_INFO_PATH}"); + if let Ok(response) = client + .get(&user_info_url) + .timeout(timeout) + .header("Cookie", cookie_header) + .header("Accept", "application/json, text/plain, */*") + .header("User-Agent", USER_AGENT) + .send() + .await + && response.status().is_success() + && let Ok(bytes) = response.bytes().await + && let Ok(value) = serde_json::from_slice::(&bytes) + { + let expanded = expand_json_strings(value); + if let Some(token) = + find_first_string(&expanded, &["secToken", "sec_token", "csrfToken", "token"]) + { + return Some(token); + } + } + + None + } + + async fn post_api( + client: &reqwest::Client, + api: &str, + data_parameters: Map, + sec_token: &str, + cookie_header: &str, + ctx: &FetchContext, + ) -> Result, ProviderError> { + let url = api_url(api); + let params_json = build_params_json(api, data_parameters, cookie_header); + let form = [ + ("product", CONSOLE_PRODUCT.to_string()), + ("action", CONSOLE_ACTION.to_string()), + ("sec_token", sec_token.to_string()), + ("region", REGION.to_string()), + ("language", LANGUAGE.to_string()), + ("params", params_json), + ]; + + let mut request = client + .post(&url) + .timeout(std::time::Duration::from_secs(ctx.web_timeout.max(1))) + .header("Cookie", cookie_header) + .header("Accept", "application/json, text/plain, */*") + .header("Content-Type", "application/x-www-form-urlencoded") + .header("Origin", GATEWAY_BASE_URL) + .header("Referer", DASHBOARD_URL) + .header("User-Agent", USER_AGENT) + .header("X-Requested-With", "XMLHttpRequest") + .form(&form); + + if let Some(csrf) = cookie_value("login_aliyunid_csrf", cookie_header) + .or_else(|| cookie_value("csrf", cookie_header)) + { + request = request + .header("x-xsrf-token", csrf.clone()) + .header("x-csrf-token", csrf); + } + + let response = request.send().await?; + let status = response.status(); + let body = response.bytes().await?; + if !status.is_success() { + if status == reqwest::StatusCode::UNAUTHORIZED + || status == reqwest::StatusCode::FORBIDDEN + { + return Err(ProviderError::AuthRequired); + } + return Err(ProviderError::Other(format!( + "Qwen Cloud API error: HTTP {status}" + ))); + } + Ok(body.to_vec()) + } + + async fn post_api_optional( + client: &reqwest::Client, + api: &str, + data_parameters: Map, + sec_token: &str, + cookie_header: &str, + ctx: &FetchContext, + ) -> Option> { + Self::post_api(client, api, data_parameters, sec_token, cookie_header, ctx) + .await + .ok() + } + + fn parse( + usage_data: &[u8], + subscription_data: Option<&[u8]>, + quota_config_data: Option<&[u8]>, + ) -> Result { + if usage_data.is_empty() { + return Err(ProviderError::Parse("Empty Qwen Cloud response".into())); + } + + let value: Value = serde_json::from_slice(usage_data).map_err(|_| { + if is_likely_login_html(usage_data) { + ProviderError::AuthRequired + } else { + ProviderError::Parse("Invalid Qwen Cloud JSON response".into()) + } + })?; + let expanded = expand_json_strings(value); + throw_if_error_payload(&expanded)?; + + if let Some(snapshot) = + parse_current_token_plan(&expanded, subscription_data, quota_config_data) + { + return Ok(snapshot); + } + + parse_legacy_token_plan(&expanded) + } + + fn snapshot_to_usage(snapshot: QwenCloudSnapshot) -> Result { + let five_hour = snapshot.five_hour_used_percent.map(|percent| { + RateWindow::with_details( + percent, + Some(FIVE_HOUR_MINUTES), + snapshot.five_hour_resets_at, + quota_detail_percent(percent, snapshot.five_hour_total_quota), + ) + }); + let legacy = used_percent( + snapshot.used_quota, + snapshot.total_quota, + snapshot.remaining_quota, + ) + .map(|percent| { + RateWindow::with_details( + percent, + Some(LEGACY_MINUTES), + snapshot.resets_at, + quota_detail( + snapshot.used_quota, + snapshot.total_quota, + snapshot.remaining_quota, + ), + ) + }); + let primary = five_hour + .or(legacy) + .ok_or_else(|| ProviderError::Parse("Qwen Cloud usage windows missing".into()))?; + let mut usage = UsageSnapshot::new(primary); + + if let Some(weekly_percent) = snapshot.weekly_used_percent { + usage = usage.with_secondary(RateWindow::with_details( + weekly_percent, + Some(WEEKLY_MINUTES), + snapshot.weekly_resets_at, + quota_detail_percent(weekly_percent, snapshot.weekly_total_quota), + )); + } + + if let Some(plan) = snapshot.plan_name.filter(|plan| !plan.trim().is_empty()) { + usage = usage.with_login_method(plan); + } + Ok(usage) + } +} + +impl Default for QwenCloudProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for QwenCloudProvider { + fn id(&self) -> ProviderId { + ProviderId::QwenCloud + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::Web => { + let usage = self.fetch_via_web(ctx).await?; + Ok(ProviderFetchResult::new(usage, "web")) + } + SourceMode::Cli | SourceMode::OAuth => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::Web] + } + + fn supports_web(&self) -> bool { + true + } +} + +fn api_url(api: &str) -> String { + format!( + "{DATA_GATEWAY_BASE_URL}/data/api.json?action={CONSOLE_ACTION}&product={CONSOLE_PRODUCT}&api={api}&_v=undefined" + ) +} + +fn build_params_json( + api: &str, + mut data_parameters: Map, + cookie_header: &str, +) -> String { + let mut cornerstone = Map::new(); + cornerstone.insert( + "feTraceId".into(), + Value::String(Uuid::new_v4().to_string().to_lowercase()), + ); + cornerstone.insert("feURL".into(), Value::String(DASHBOARD_URL.to_string())); + cornerstone.insert("protocol".into(), Value::String("V2".into())); + cornerstone.insert("console".into(), Value::String("ONE_CONSOLE".into())); + cornerstone.insert("productCode".into(), Value::String("p_efm".into())); + cornerstone.insert("domain".into(), Value::String("home.qwencloud.com".into())); + cornerstone.insert("consoleSite".into(), Value::String("QWENCLOUD".into())); + cornerstone.insert("userNickName".into(), Value::String(String::new())); + cornerstone.insert("userPrincipalName".into(), Value::String(String::new())); + cornerstone.insert("xsp_lang".into(), Value::String(LANGUAGE.into())); + if let Some(cna) = cookie_value("cna", cookie_header) { + cornerstone.insert("X-Anonymous-Id".into(), Value::String(cna)); + } + + data_parameters.insert("cornerstoneParam".into(), Value::Object(cornerstone)); + + json!({ + "Api": api, + "V": "1.0", + "Data": Value::Object(data_parameters), + }) + .to_string() +} + +fn parse_current_token_plan( + expanded: &Value, + subscription_data: Option<&[u8]>, + quota_config_data: Option<&[u8]>, +) -> Option { + let usage = + find_object_containing_any_of(expanded, &["per5HourPercentage", "per1WeekPercentage"])?; + let five_hour = percentage_points(number_field(&usage, "per5HourPercentage")); + let weekly = percentage_points(number_field(&usage, "per1WeekPercentage")); + if five_hour.is_none() && weekly.is_none() { + return None; + } + + let plan_code = subscription_data.and_then(plan_code_from_bytes); + let plan_name = plan_code.as_deref().map(display_plan_name); + let quota = quota_config_data + .zip(plan_code.as_ref()) + .and_then(|(data, code)| quota_totals_from_bytes(data, code)); + + Some(QwenCloudSnapshot { + plan_name, + used_quota: None, + total_quota: None, + remaining_quota: None, + resets_at: None, + five_hour_used_percent: five_hour, + five_hour_total_quota: quota.map(|q| q.0).unwrap_or(None), + five_hour_resets_at: date_field(&usage, "per5HourResetTime"), + weekly_used_percent: weekly, + weekly_total_quota: quota.map(|q| q.1).unwrap_or(None), + weekly_resets_at: date_field(&usage, "per1WeekResetTime"), + }) +} + +fn parse_legacy_token_plan(expanded: &Value) -> Result { + let instance = find_token_plan_instance(expanded); + let plan_name = instance + .as_ref() + .and_then(|v| find_first_string(v, PLAN_NAME_KEYS)) + .or_else(|| find_first_string(expanded, PLAN_NAME_KEYS)); + let quota_source = instance + .as_ref() + .and_then(find_quota_info) + .or_else(|| find_quota_info(expanded)) + .unwrap_or_else(|| expanded.clone()); + let used = first_f64("a_source, USED_QUOTA_KEYS) + .or_else(|| find_first_f64("a_source, USED_QUOTA_KEYS)); + let total = first_f64("a_source, TOTAL_QUOTA_KEYS) + .or_else(|| find_first_f64("a_source, TOTAL_QUOTA_KEYS)); + let remaining = first_f64("a_source, REMAINING_QUOTA_KEYS) + .or_else(|| find_first_f64("a_source, REMAINING_QUOTA_KEYS)); + let resets_at = instance + .as_ref() + .and_then(|v| { + first_date(v, RESET_DATE_KEYS).or_else(|| find_first_date(v, RESET_DATE_KEYS)) + }) + .or_else(|| { + first_date(expanded, RESET_DATE_KEYS) + .or_else(|| find_first_date(expanded, RESET_DATE_KEYS)) + }); + + if used_percent(used, total, remaining).is_none() { + return Err(ProviderError::Parse( + "Qwen Cloud has no active token-plan subscription".into(), + )); + } + + Ok(QwenCloudSnapshot { + plan_name, + used_quota: used, + total_quota: total, + remaining_quota: remaining, + resets_at, + five_hour_used_percent: None, + five_hour_total_quota: None, + five_hour_resets_at: None, + weekly_used_percent: None, + weekly_total_quota: None, + weekly_resets_at: None, + }) +} + +fn plan_code_from_bytes(data: &[u8]) -> Option { + let value: Value = serde_json::from_slice(data).ok()?; + let expanded = expand_json_strings(value); + let plan = find_object_containing_any_of( + &expanded, + &["specCode", "spec_code", "planName", "plan_name"], + )?; + for key in ["specCode", "spec_code", "planName", "plan_name"] { + if let Some(raw) = plan.get(key).and_then(Value::as_str) { + let normalized = raw.trim().to_lowercase(); + if !normalized.is_empty() { + return Some(normalized); + } + } + } + None +} + +fn display_plan_name(plan_code: &str) -> String { + match plan_code { + "lite" => "Lite".into(), + "standard" => "Standard".into(), + "pro" => "Pro".into(), + "max" => "Max".into(), + other => other.to_string(), + } +} + +fn quota_totals_from_bytes(data: &[u8], plan_code: &str) -> Option<(Option, Option)> { + let value: Value = serde_json::from_slice(data).ok()?; + let expanded = expand_json_strings(value); + let quota = find_first_value_for_key(&expanded, plan_code)? + .as_object()? + .clone(); + let five_hour = number_field(&Value::Object(quota.clone()), "five_hour") + .or_else(|| number_field(&Value::Object(quota.clone()), "fiveHour")); + let weekly = number_field(&Value::Object(quota), "weekly"); + if five_hour.is_none() && weekly.is_none() { + return None; + } + Some((five_hour, weekly)) +} + +fn throw_if_error_payload(value: &Value) -> Result<(), ProviderError> { + if let Some(status) = find_first_i64(value, &["statusCode", "status_code", "code"]) + && status != 0 + && status != 200 + { + if status == 401 || status == 403 { + return Err(ProviderError::AuthRequired); + } + let message = find_first_string(value, &["statusMessage", "status_msg", "message", "msg"]) + .unwrap_or_else(|| format!("status code {status}")); + return Err(ProviderError::Other(format!( + "Qwen Cloud API error: {message}" + ))); + } + + if let Some(success) = find_first_bool(value, &["success", "Success", "successResponse"]) + && !success + { + let message = find_first_string(value, &["message", "msg", "Message", "errorMessage"]) + .unwrap_or_else(|| "request failed".to_string()); + let lower = message.to_lowercase(); + if lower.contains("needlogin") + || lower.contains("login") + || lower.contains("log in") + || lower.contains("unauthorized") + { + return Err(ProviderError::AuthRequired); + } + return Err(ProviderError::Other(format!( + "Qwen Cloud API error: {message}" + ))); + } + + let code = find_first_string(value, &["code", "status", "statusCode"]) + .unwrap_or_default() + .to_lowercase(); + let message = find_first_string(value, &["message", "msg", "statusMessage"]) + .unwrap_or_default() + .to_lowercase(); + if code.contains("needlogin") + || code.contains("login") + || message.contains("log in") + || message.contains("login") + { + return Err(ProviderError::AuthRequired); + } + if code.contains("forbidden") || code == "403" || message.contains("forbidden") { + return Err(ProviderError::AuthRequired); + } + Ok(()) +} + +fn normalize_cookie_header(raw: &str) -> Option { + let mut header = raw.trim(); + if header + .get(.."cookie:".len()) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("cookie:")) + { + header = header["cookie:".len()..].trim(); + } + if (header.starts_with('"') && header.ends_with('"')) + || (header.starts_with('\'') && header.ends_with('\'')) + { + header = header[1..header.len().saturating_sub(1)].trim(); + } + (!header.is_empty() && header.contains('=')).then(|| header.to_string()) +} + +fn cookie_value(name: &str, cookie_header: &str) -> Option { + cookie_header.split(';').find_map(|part| { + let (key, value) = part.trim().split_once('=')?; + (key.trim().eq_ignore_ascii_case(name)) + .then(|| value.trim().to_string()) + .filter(|value| !value.is_empty()) + }) +} + +fn extract_sec_token(html: &str) -> Option { + for pattern in [ + r#""secToken"\s*:\s*"([^"]+)""#, + r#""sec_token"\s*:\s*"([^"]+)""#, + r#"secToken['"]?\s*[:=]\s*['"]([^'"]+)['"]"#, + r#"sec_token['"]?\s*[:=]\s*['"]([^'"]+)['"]"#, + r#"csrfToken['"]?\s*[:=]\s*['"]([^'"]+)['"]"#, + ] { + let Ok(regex) = Regex::new(pattern) else { + continue; + }; + if let Some(value) = regex + .captures(html) + .and_then(|captures| captures.get(1)) + .map(|m| m.as_str().trim().to_string()) + .filter(|value| !value.is_empty()) + { + return Some(value); + } + } + None +} + +fn looks_like_login_page(html: &str) -> bool { + let lowered = html.to_lowercase(); + lowered.contains("passport.alibabacloud.com") + || lowered.contains("signin.aliyun.com") + || lowered.contains("account.alibabacloud.com/login") + || lowered.contains("login.qwencloud.com") + || (lowered.contains("login") + && lowered.contains("password") + && lowered.contains("sign in")) +} + +fn is_likely_login_html(data: &[u8]) -> bool { + let text = String::from_utf8_lossy(data).to_lowercase(); + text.contains(" Value { + match value { + Value::Array(values) => Value::Array(values.into_iter().map(expand_json_strings).collect()), + Value::Object(map) => Value::Object( + map.into_iter() + .map(|(key, value)| (key, expand_json_strings(value))) + .collect(), + ), + Value::String(text) => serde_json::from_str::(&text) + .ok() + .filter(|nested| nested.is_object() || nested.is_array()) + .map(expand_json_strings) + .unwrap_or(Value::String(text)), + other => other, + } +} + +fn percentage_points(ratio: Option) -> Option { + let ratio = ratio.filter(|v| v.is_finite())?; + Some((ratio.clamp(0.0, 1.0) * 100.0).clamp(0.0, 100.0)) +} + +fn number_field(value: &Value, key: &str) -> Option { + value.as_object().and_then(|map| parse_f64(map.get(key))) +} + +fn date_field(value: &Value, key: &str) -> Option> { + value.as_object().and_then(|map| parse_date(map.get(key))) +} + +fn find_object_containing_any_of(value: &Value, keys: &[&str]) -> Option { + match value { + Value::Object(map) => { + if keys.iter().any(|key| map.contains_key(*key)) { + return Some(Value::Object(map.clone())); + } + map.values() + .find_map(|nested| find_object_containing_any_of(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_object_containing_any_of(nested, keys)), + _ => None, + } +} + +fn find_first_value_for_key(value: &Value, key: &str) -> Option { + match value { + Value::Object(map) => { + if let Some(nested) = map.get(key) { + return Some(nested.clone()); + } + map.values() + .find_map(|nested| find_first_value_for_key(nested, key)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_value_for_key(nested, key)), + _ => None, + } +} + +const PLAN_NAME_KEYS: &[&str] = &[ + "planName", + "plan_name", + "packageName", + "package_name", + "commodityName", + "commodity_name", + "instanceName", + "instance_name", + "displayName", + "display_name", + "name", + "title", + "planType", + "plan_type", + "ProductName", + "productName", + "InstanceCode", +]; + +const USED_QUOTA_KEYS: &[&str] = &[ + "usedQuota", + "used_quota", + "usedCredits", + "usedCredit", + "consumedCredits", + "usage", + "used", + "usedAmount", + "consumeAmount", + "usedValue", + "UsedValue", + "consumedValue", + "ConsumedValue", +]; + +const TOTAL_QUOTA_KEYS: &[&str] = &[ + "totalQuota", + "total_quota", + "totalCredits", + "totalCredit", + "quota", + "creditLimit", + "creditsTotal", + "monthlyTotalQuota", + "amount", + "totalValue", + "TotalValue", + "CycleTotalValue", + "cycleTotalValue", + "subscriptionTotalNumber", + "SubscriptionTotalNumber", +]; + +const REMAINING_QUOTA_KEYS: &[&str] = &[ + "remainingQuota", + "remainQuota", + "remainingCredits", + "remainingCredit", + "availableCredits", + "balance", + "remaining", + "availableAmount", + "remainAmount", + "totalSurplusValue", + "TotalSurplusValue", + "surplusValue", + "SurplusValue", + "CycleSurplusValue", + "cycleSurplusValue", +]; + +const RESET_DATE_KEYS: &[&str] = &[ + "nextRefreshTime", + "resetTime", + "periodEndTime", + "billingCycleEnd", + "billCycleEndTime", + "expireTime", + "expirationTime", + "endTime", + "EndTime", + "validEndTime", + "instanceEndTime", + "nearestExpireDate", + "NearestExpireDate", +]; + +fn find_token_plan_instance(value: &Value) -> Option { + find_first_object( + value, + &[ + "tokenPlanInstanceInfo", + "token_plan_instance_info", + "instanceInfo", + "instance_info", + ], + ) + .or_else(|| { + find_first_array( + value, + &[ + "tokenPlanInstanceInfos", + "token_plan_instance_infos", + "instanceInfos", + "instances", + "EquityList", + "Data", + "data", + "successResponse", + ], + ) + .and_then(|values| { + values + .into_iter() + .filter(Value::is_object) + .max_by_key(active_signal_score) + }) + }) +} + +fn find_quota_info(value: &Value) -> Option { + find_first_object( + value, + &[ + "quotaInfo", + "quota_info", + "tokenPlanQuotaInfo", + "token_plan_quota_info", + ], + ) + .or_else(|| { + find_first_array(value, &["EquityList", "equityList"]).and_then(|values| { + values.into_iter().find(|item| { + item.as_object().is_some_and(|map| { + map.contains_key("CycleTotalValue") + || map.contains_key("cycleTotalValue") + || map.contains_key("CycleSurplusValue") + }) + }) + }) + }) + .or_else(|| { + let keys: Vec<&str> = USED_QUOTA_KEYS + .iter() + .chain(TOTAL_QUOTA_KEYS.iter()) + .chain(REMAINING_QUOTA_KEYS.iter()) + .copied() + .collect(); + find_first_object_with_any_key(value, &keys) + }) +} + +fn find_first_object(value: &Value, keys: &[&str]) -> Option { + match value { + Value::Object(map) => { + for key in keys { + if let Some(nested) = map.get(*key).filter(|v| v.is_object()) { + return Some(nested.clone()); + } + } + map.values() + .find_map(|nested| find_first_object(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_object(nested, keys)), + _ => None, + } +} + +fn find_first_object_with_any_key(value: &Value, keys: &[&str]) -> Option { + match value { + Value::Object(map) => { + if keys.iter().any(|key| map.contains_key(*key)) { + return Some(Value::Object(map.clone())); + } + map.values() + .find_map(|nested| find_first_object_with_any_key(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_object_with_any_key(nested, keys)), + _ => None, + } +} + +fn find_first_array(value: &Value, keys: &[&str]) -> Option> { + match value { + Value::Object(map) => { + for key in keys { + if let Some(Value::Array(values)) = map.get(*key) { + return Some(values.clone()); + } + } + map.values() + .find_map(|nested| find_first_array(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_array(nested, keys)), + _ => None, + } +} + +fn first_string(value: &Value, keys: &[&str]) -> Option { + value + .as_object() + .and_then(|map| keys.iter().find_map(|key| parse_string(map.get(*key)))) +} + +fn find_first_string(value: &Value, keys: &[&str]) -> Option { + first_string(value, keys).or_else(|| match value { + Value::Object(map) => map + .values() + .find_map(|nested| find_first_string(nested, keys)), + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_string(nested, keys)), + _ => None, + }) +} + +fn first_f64(value: &Value, keys: &[&str]) -> Option { + value + .as_object() + .and_then(|map| keys.iter().find_map(|key| parse_f64(map.get(*key)))) +} + +fn find_first_f64(value: &Value, keys: &[&str]) -> Option { + first_f64(value, keys).or_else(|| match value { + Value::Object(map) => map.values().find_map(|nested| find_first_f64(nested, keys)), + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_f64(nested, keys)), + _ => None, + }) +} + +fn find_first_i64(value: &Value, keys: &[&str]) -> Option { + match value { + Value::Object(map) => { + for key in keys { + if let Some(parsed) = parse_i64(map.get(*key)) { + return Some(parsed); + } + } + map.values().find_map(|nested| find_first_i64(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_i64(nested, keys)), + _ => None, + } +} + +fn find_first_bool(value: &Value, keys: &[&str]) -> Option { + match value { + Value::Object(map) => { + for key in keys { + if let Some(parsed) = parse_bool(map.get(*key)) { + return Some(parsed); + } + } + map.values() + .find_map(|nested| find_first_bool(nested, keys)) + } + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_bool(nested, keys)), + _ => None, + } +} + +fn first_date(value: &Value, keys: &[&str]) -> Option> { + value + .as_object() + .and_then(|map| keys.iter().find_map(|key| parse_date(map.get(*key)))) +} + +fn find_first_date(value: &Value, keys: &[&str]) -> Option> { + first_date(value, keys).or_else(|| match value { + Value::Object(map) => map + .values() + .find_map(|nested| find_first_date(nested, keys)), + Value::Array(values) => values + .iter() + .find_map(|nested| find_first_date(nested, keys)), + _ => None, + }) +} + +fn parse_string(value: Option<&Value>) -> Option { + value? + .as_str() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(ToOwned::to_owned) +} + +fn parse_f64(value: Option<&Value>) -> Option { + match value? { + Value::Number(number) => number.as_f64(), + Value::String(text) => text.trim().replace(',', "").parse().ok(), + _ => None, + } + .filter(|v| v.is_finite()) +} + +fn parse_i64(value: Option<&Value>) -> Option { + match value? { + Value::Number(number) => number + .as_i64() + .or_else(|| number.as_f64().map(|v| v as i64)), + Value::String(text) => text.trim().replace(',', "").parse().ok(), + _ => None, + } +} + +fn parse_bool(value: Option<&Value>) -> Option { + match value? { + Value::Bool(flag) => Some(*flag), + Value::Number(number) => number.as_i64().map(|v| v != 0), + Value::String(text) => match text.trim().to_lowercase().as_str() { + "true" | "1" | "yes" | "active" | "valid" | "normal" => Some(true), + "false" | "0" | "no" | "inactive" | "invalid" | "expired" => Some(false), + _ => None, + }, + _ => None, + } +} + +fn parse_date(value: Option<&Value>) -> Option> { + if let Some(raw) = parse_i64(value) { + if raw > 1_000_000_000_000 { + return Utc.timestamp_opt(raw / 1000, 0).single(); + } + if raw > 1_000_000_000 { + return Utc.timestamp_opt(raw, 0).single(); + } + } + let text = parse_string(value)?; + if let Ok(date) = DateTime::parse_from_rfc3339(&text) { + return Some(date.with_timezone(&Utc)); + } + if let Ok(date) = NaiveDate::parse_from_str(&text, "%Y-%m-%d") + && let Some(date_time) = date.and_hms_opt(0, 0, 0) + { + return Some(date_time.and_utc()); + } + for format in ["%Y-%m-%d %H:%M", "%Y-%m-%d %H:%M:%S"] { + if let Ok(date) = NaiveDateTime::parse_from_str(&text, format) { + return Some(date.and_utc()); + } + } + None +} + +fn active_signal_score(value: &Value) -> i32 { + let status = first_string(value, &["status", "instanceStatus", "state", "Status"]) + .unwrap_or_default() + .to_uppercase(); + if ["VALID", "ACTIVE", "NORMAL"].contains(&status.as_str()) { + return 3; + } + if [ + "EXPIRED", + "INVALID", + "INACTIVE", + "DISABLED", + "TERMINATED", + "STOPPED", + ] + .contains(&status.as_str()) + { + return -1; + } + parse_bool( + value + .as_object() + .and_then(|map| map.get("isActive").or_else(|| map.get("active"))), + ) + .map(|active| if active { 3 } else { -1 }) + .unwrap_or(0) +} + +fn used_percent(used: Option, total: Option, remaining: Option) -> Option { + let total = total.filter(|total| *total > 0.0)?; + let used = used.or_else(|| remaining.map(|remaining| total - remaining))?; + Some((used.clamp(0.0, total) / total * 100.0).clamp(0.0, 100.0)) +} + +fn quota_detail(used: Option, total: Option, remaining: Option) -> Option { + if let (Some(used), Some(total)) = (used, total.filter(|total| *total > 0.0)) { + return Some(format!( + "{} / {} credits used", + format_quota(used), + format_quota(total) + )); + } + if let (Some(remaining), Some(total)) = (remaining, total.filter(|total| *total > 0.0)) { + return Some(format!( + "{} / {} credits left", + format_quota(remaining), + format_quota(total) + )); + } + remaining.map(|remaining| format!("{} credits left", format_quota(remaining))) +} + +fn quota_detail_percent(used_percent: f64, total: Option) -> Option { + let total = total.filter(|total| *total > 0.0)?; + let used = total * used_percent / 100.0; + Some(format!( + "{} / {} credits used", + format_quota(used), + format_quota(total) + )) +} + +fn format_quota(value: f64) -> String { + if (value.round() - value).abs() < f64::EPSILON { + format_count(value.round() as i64) + } else { + let formatted = format!("{value:.2}"); + let trimmed = formatted.trim_end_matches('0').trim_end_matches('.'); + format_count_decimal(trimmed) + } +} + +fn format_count(value: i64) -> String { + let raw = value.to_string(); + let mut output = String::with_capacity(raw.len() + raw.len() / 3); + for (idx, ch) in raw.chars().rev().enumerate() { + if idx > 0 && idx % 3 == 0 { + output.push(','); + } + output.push(ch); + } + output.chars().rev().collect() +} + +fn format_count_decimal(raw: &str) -> String { + let (whole, fraction) = raw.split_once('.').unwrap_or((raw, "")); + if fraction.is_empty() { + format_count(whole.parse().unwrap_or(0)) + } else { + format!("{}.{}", format_count(whole.parse().unwrap_or(0)), fraction) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_current_token_plan_5h_and_weekly() { + let inner = r#"{ + "code": 0, + "data": { + "per5HourPercentage": 0.03, + "per5HourResetTime": 1700003600000, + "per1WeekPercentage": 0.01, + "per1WeekResetTime": 1700086400000 + }, + "success": true + }"#; + let payload = serde_json::json!({ + "data": { + "DataV2": { + "data": inner, + }, + }, + "httpStatusCode": 200, + }); + + let snapshot = + QwenCloudProvider::parse(payload.to_string().as_bytes(), None, None).unwrap(); + assert_eq!(snapshot.five_hour_used_percent, Some(3.0)); + assert_eq!( + snapshot.five_hour_resets_at, + Some(Utc.timestamp_opt(1_700_003_600, 0).single().unwrap()) + ); + assert_eq!(snapshot.weekly_used_percent, Some(1.0)); + assert_eq!( + snapshot.weekly_resets_at, + Some(Utc.timestamp_opt(1_700_086_400, 0).single().unwrap()) + ); + + let usage = QwenCloudProvider::snapshot_to_usage(snapshot).unwrap(); + assert_eq!(usage.primary.used_percent, 3.0); + assert_eq!(usage.primary.window_minutes, Some(FIVE_HOUR_MINUTES)); + assert_eq!(usage.secondary.as_ref().map(|w| w.used_percent), Some(1.0)); + assert_eq!( + usage.secondary.as_ref().and_then(|w| w.window_minutes), + Some(WEEKLY_MINUTES) + ); + } + + #[test] + fn parses_personal_usage_fixture_shape() { + let payload = serde_json::json!({ + "code": "200", + "data": { + "DataV2": { + "data": { + "success": true, + "data": { + "per5HourPercentage": 0.0009973083333333333, + "per5HourResetTime": 1784813220000_i64, + "per1WeekPercentage": 0.0003014725, + "per1WeekResetTime": 1785234900000_i64 + } + }, + "success": true, + "httpStatus": 200 + } + }, + "successResponse": true + }); + let usage = QwenCloudProvider::snapshot_to_usage( + QwenCloudProvider::parse(payload.to_string().as_bytes(), None, None).unwrap(), + ) + .unwrap(); + assert!((usage.primary.used_percent - 0.09973083333333333).abs() < 1e-9); + assert_eq!(usage.primary.window_minutes, Some(300)); + assert!(usage.secondary.is_some()); + assert_eq!( + usage.secondary.as_ref().and_then(|w| w.window_minutes), + Some(10080) + ); + } + + #[test] + fn parses_nested_equity_list_legacy() { + let payload = serde_json::json!({ + "code": "200", + "successResponse": true, + "data": { + "TotalCount": 1, + "Data": [ + { + "InstanceCode": "qwen-token-plan", + "Status": "NORMAL", + "EndTime": 1_701_000_000_000_i64, + "EquityList": [ + { + "Type": "CREDITS", + "CycleTotalValue": "1000", + "CycleSurplusValue": "875" + } + ] + } + ] + } + }); + let snapshot = + QwenCloudProvider::parse(payload.to_string().as_bytes(), None, None).unwrap(); + assert_eq!(snapshot.total_quota, Some(1000.0)); + assert_eq!(snapshot.remaining_quota, Some(875.0)); + let usage = QwenCloudProvider::snapshot_to_usage(snapshot).unwrap(); + assert_eq!(usage.primary.used_percent, 12.5); + assert_eq!(usage.primary.window_minutes, Some(LEGACY_MINUTES)); + } + + #[test] + fn parses_flat_subscription_summary_legacy() { + let payload = serde_json::json!({ + "Success": true, + "Data": { + "TotalCount": 1, + "TotalValue": 2000, + "TotalSurplusValue": 1500 + } + }); + let snapshot = + QwenCloudProvider::parse(payload.to_string().as_bytes(), None, None).unwrap(); + assert_eq!(snapshot.total_quota, Some(2000.0)); + assert_eq!(snapshot.remaining_quota, Some(1500.0)); + assert_eq!( + used_percent( + snapshot.used_quota, + snapshot.total_quota, + snapshot.remaining_quota + ), + Some(25.0) + ); + } + + #[test] + fn login_payload_maps_to_auth_required() { + let err = QwenCloudProvider::parse( + br#"{"code":"ConsoleNeedLogin","message":"You need to log in.","successResponse":false}"#, + None, + None, + ) + .unwrap_err(); + assert!(matches!(err, ProviderError::AuthRequired)); + } + + #[test] + fn attaches_plan_name_from_subscription() { + let usage_payload = serde_json::json!({ + "data": { + "per5HourPercentage": 0.1, + "per5HourResetTime": 1700003600000_i64, + "per1WeekPercentage": 0.2, + "per1WeekResetTime": 1700086400000_i64 + } + }); + let subscription = serde_json::json!({ + "data": { "specCode": "pro" } + }); + let snapshot = QwenCloudProvider::parse( + usage_payload.to_string().as_bytes(), + Some(subscription.to_string().as_bytes()), + None, + ) + .unwrap(); + assert_eq!(snapshot.plan_name.as_deref(), Some("Pro")); + let usage = QwenCloudProvider::snapshot_to_usage(snapshot).unwrap(); + assert_eq!(usage.login_method.as_deref(), Some("Pro")); + } + + #[test] + fn extracts_sec_token_from_html() { + assert_eq!( + extract_sec_token(r#""#).as_deref(), + Some("qwen-html-token") + ); + assert_eq!( + cookie_value("login_aliyunid_csrf", "foo=bar; login_aliyunid_csrf=tok"), + Some("tok".to_string()) + ); + } + + #[test] + fn metadata_labels_match_upstream() { + let provider = QwenCloudProvider::new(); + assert_eq!(provider.metadata().session_label, "5-hour"); + assert_eq!(provider.metadata().weekly_label, "Weekly"); + assert!(!provider.metadata().default_enabled); + assert_eq!(provider.metadata().dashboard_url, Some(DASHBOARD_URL)); + } +} diff --git a/rust/src/providers/zoommate/mod.rs b/rust/src/providers/zoommate/mod.rs new file mode 100644 index 0000000000..1749c47dce --- /dev/null +++ b/rust/src/providers/zoommate/mod.rs @@ -0,0 +1,886 @@ +//! ZoomMate provider (upstream 0.46.0). +//! +//! Auth: manual Cookie header **or** DevTools cURL of credits/status (must include +//! `Authorization: Bearer …`), else browser cookies → mint bearer via login bootstrap. +//! Bearer JWTs are cached in-process by cookie fingerprint until `exp - 60s`. +//! +//! Optional credits/history (Today KPI) is intentionally skipped for the Windows v1 +//! port — primary credits/status is sufficient; history can land later without +//! blocking the snapshot. + +use async_trait::async_trait; +use chrono::{DateTime, TimeZone, Utc}; +use reqwest::Client; +use serde::Deserialize; +use sha2::{Digest, Sha256}; +use std::collections::HashMap; +use std::sync::{Mutex, OnceLock}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use crate::core::curl_capture; +use crate::core::{ + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, + RateWindow, SourceMode, UsageSnapshot, +}; +use crate::providers::browser_cookie_header; + +const API_HOSTS: &[&str] = &["ai.zoom.us", "zoommate.zoom.us"]; +const COOKIE_DOMAINS: &[&str] = &["ai.zoom.us", "zoommate.zoom.us", "zoom.us"]; +const CREDITS_STATUS_PATH: &str = "/ai-computer/api/v1/credits/status"; +const ORIGIN_REFERER: &str = "https://zoommate.zoom.us"; +const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"; +const BEARER_REFRESH_SKEW_SECS: u64 = 60; + +const FORWARDED_MANUAL_HEADERS: &[(&str, &str)] = &[ + ("authorization", "Authorization"), + ("cookie", "Cookie"), + ("origin", "Origin"), + ("referer", "Referer"), + ("user-agent", "User-Agent"), + ("accept", "Accept"), + ("accept-language", "Accept-Language"), +]; + +// ── Response shapes ────────────────────────────────────────────────────────── + +#[derive(Debug, Deserialize, Default, Clone)] +struct CreditsStatusEnvelope { + data: Option, + #[serde(default)] + status_code: Option, + #[serde(default)] + error_message: Option, +} + +#[derive(Debug, Deserialize, Default, Clone)] +struct CreditsDataBox { + credit_status: Option, +} + +#[derive(Debug, Deserialize, Default, Clone)] +pub(crate) struct CreditStatus { + budget_cap: Option, + used_credit: Option, + remaining_credit: Option, + overage_credit: Option, + allow_overage: Option, + cycle_start_date: Option, + cycle_end_date: Option, + is_quota_available: Option, + is_unlimited: Option, +} + +#[derive(Debug, Deserialize, Default)] +struct LoginBootstrapEnvelope { + success: Option, + data: Option, +} + +#[derive(Debug, Deserialize, Default)] +struct LoginDataBox { + nak: Option, + user_profile: Option, +} + +#[derive(Debug, Deserialize, Default)] +struct LoginUserProfile { + email: Option, +} + +// ── Request context ────────────────────────────────────────────────────────── + +#[derive(Debug, Clone)] +struct RequestContext { + authorization: String, + headers: HashMap, + /// Host → Cookie header. Empty map / missing host → no Cookie sent for that host. + cookie_by_host: HashMap, + preferred_host: Option, + account_email: Option, + /// Present only on cookie-mint path so 401 can invalidate the bearer cache entry. + cache_key: Option, +} + +#[derive(Debug, Clone)] +struct MintedToken { + bearer_token: String, + account_email: Option, +} + +// ── In-process bearer cache ────────────────────────────────────────────────── + +struct BearerEntry { + token: String, + account_email: Option, + /// UNIX seconds of JWT `exp`. + exp_unix: u64, +} + +fn bearer_cache() -> &'static Mutex> { + static CACHE: OnceLock>> = OnceLock::new(); + CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn cookie_fingerprint(cookie_by_host: &HashMap) -> String { + let mut pairs: Vec<_> = cookie_by_host.iter().collect(); + pairs.sort_by(|a, b| a.0.cmp(b.0)); + let canonical = pairs + .into_iter() + .map(|(h, c)| format!("{h}={c}")) + .collect::>() + .join("\n"); + let digest = Sha256::digest(canonical.as_bytes()); + digest.iter().map(|b| format!("{b:02x}")).collect() +} + +fn cache_valid_entry(key: &str) -> Option { + let now = SystemTime::now().duration_since(UNIX_EPOCH).ok()?.as_secs(); + let mut guard = bearer_cache().lock().ok()?; + let entry = guard.get(key)?; + if entry.exp_unix.saturating_sub(BEARER_REFRESH_SKEW_SECS) <= now { + guard.remove(key); + return None; + } + Some(MintedToken { + bearer_token: entry.token.clone(), + account_email: entry.account_email.clone(), + }) +} + +fn cache_store(key: String, minted: &MintedToken, exp_unix: u64) { + if let Ok(mut guard) = bearer_cache().lock() { + guard.insert( + key, + BearerEntry { + token: minted.bearer_token.clone(), + account_email: minted.account_email.clone(), + exp_unix, + }, + ); + } +} + +fn cache_invalidate(key: &str) { + if let Ok(mut guard) = bearer_cache().lock() { + guard.remove(key); + } +} + +// ── Provider ───────────────────────────────────────────────────────────────── + +pub struct ZoomMateProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl ZoomMateProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::ZoomMate, + display_name: "ZoomMate", + session_label: "Credits", + weekly_label: "Credits", + supports_opus: false, + supports_credits: true, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://zoommate.zoom.us/#/?settings=credit-usage"), + status_page_url: Some("https://www.zoomstatus.com/"), + }, + client: crate::core::credentialed_http_client_builder() + .timeout(Duration::from_secs(15)) + .build() + .unwrap_or_else(|_| Client::new()), + } + } + + async fn fetch_via_web(&self, ctx: &FetchContext) -> Result { + let request_context = self.resolve_request_context(ctx).await?; + match self + .fetch_credits_status(&request_context, ctx.web_timeout) + .await + { + Ok(snap) => Ok(snap), + Err(ProviderError::AuthRequired) => { + if let Some(key) = request_context.cache_key.as_deref() { + cache_invalidate(key); + } + Err(ProviderError::AuthRequired) + } + Err(e) => Err(e), + } + } + + async fn resolve_request_context( + &self, + ctx: &FetchContext, + ) -> Result { + // 1) Manual paste: cURL capture (bearer) or bare Cookie header. + if let Some(raw) = ctx.manual_cookie_header.as_deref() { + let raw = raw.trim(); + if !raw.is_empty() { + if let Some(ctx) = request_context_from_manual(raw) { + return Ok(ctx); + } + // Bare cookie header → mint bearer. + if !curl_capture::looks_like_curl_capture(raw) { + let mut cookie_by_host = HashMap::new(); + for host in API_HOSTS { + cookie_by_host.insert((*host).to_string(), raw.to_string()); + } + return self + .request_context_from_cookies(cookie_by_host, ctx.web_timeout) + .await; + } + return Err(ProviderError::Other( + "Paste a cURL capture of the HTTPS ZoomMate credits/status request \ + (from ai.zoom.us or zoommate.zoom.us) including Authorization: Bearer …, \ + or paste a Cookie header." + .into(), + )); + } + } + + // 2) Browser cookies → mint bearer. + let cookie_header = browser_cookie_header(COOKIE_DOMAINS)?; + let trimmed = cookie_header.trim(); + if trimmed.is_empty() { + return Err(ProviderError::NoCookies); + } + let mut cookie_by_host = HashMap::new(); + for host in API_HOSTS { + cookie_by_host.insert((*host).to_string(), trimmed.to_string()); + } + self.request_context_from_cookies(cookie_by_host, ctx.web_timeout) + .await + } + + async fn request_context_from_cookies( + &self, + cookie_by_host: HashMap, + timeout_secs: u64, + ) -> Result { + let cache_key = cookie_fingerprint(&cookie_by_host); + let minted = if let Some(entry) = cache_valid_entry(&cache_key) { + entry + } else { + let minted = self + .mint_bearer_token(&cookie_by_host, timeout_secs) + .await?; + if let Some(exp) = jwt_exp_unix(&minted.bearer_token) { + cache_store(cache_key.clone(), &minted, exp); + } + minted + }; + Ok(RequestContext { + authorization: bearer_header_value(&minted.bearer_token), + headers: HashMap::new(), + cookie_by_host, + preferred_host: None, + account_email: minted.account_email, + cache_key: Some(cache_key), + }) + } + + async fn mint_bearer_token( + &self, + cookie_by_host: &HashMap, + timeout_secs: u64, + ) -> Result { + with_api_host_failover(None, |host| async move { + self.mint_bearer_token_on_host(cookie_by_host, host, timeout_secs) + .await + }) + .await + } + + async fn mint_bearer_token_on_host( + &self, + cookie_by_host: &HashMap, + host: &str, + timeout_secs: u64, + ) -> Result { + let url = + format!("https://{host}/ai-computer/api/v1/login/?continue=https://zoommate.zoom.us/"); + let mut req = self + .client + .get(&url) + .timeout(Duration::from_secs(timeout_secs.max(1))) + .header("Accept", "application/json, text/plain, */*") + .header("Accept-Language", "en-US,en;q=0.9") + .header("User-Agent", USER_AGENT) + .header("Origin", ORIGIN_REFERER) + .header("Referer", ORIGIN_REFERER) + .header("Sec-Fetch-Dest", "empty") + .header("Sec-Fetch-Mode", "cors") + .header("Sec-Fetch-Site", "same-site"); + if let Some(cookie) = cookie_by_host + .get(host) + .or_else(|| cookie_by_host.values().next()) + { + req = req.header("Cookie", cookie); + } + let resp = req.send().await?; + let status = resp.status(); + let body = resp.bytes().await?; + if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "ZoomMate login bootstrap error: HTTP {status}" + ))); + } + let envelope: LoginBootstrapEnvelope = serde_json::from_slice(&body).map_err(|e| { + ProviderError::Parse(format!("ZoomMate login bootstrap parse failed: {e}")) + })?; + let nak = envelope + .data + .as_ref() + .and_then(|d| d.nak.as_deref()) + .map(str::trim) + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + ProviderError::Parse("Missing nak in ZoomMate login bootstrap response.".into()) + })?; + let email = envelope + .data + .as_ref() + .and_then(|d| d.user_profile.as_ref()) + .and_then(|p| p.email.as_deref()) + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(str::to_string); + Ok(MintedToken { + bearer_token: nak.to_string(), + account_email: email, + }) + } + + async fn fetch_credits_status( + &self, + ctx: &RequestContext, + timeout_secs: u64, + ) -> Result { + let preferred = ctx.preferred_host.clone(); + let authorization = ctx.authorization.clone(); + let headers = ctx.headers.clone(); + let cookie_by_host = ctx.cookie_by_host.clone(); + let account_email = ctx.account_email.clone(); + + with_api_host_failover(preferred.as_deref(), |host| { + let authorization = authorization.clone(); + let headers = headers.clone(); + let cookie_by_host = cookie_by_host.clone(); + let account_email = account_email.clone(); + async move { + self.fetch_credits_status_on_host( + host, + &authorization, + &headers, + &cookie_by_host, + account_email.as_deref(), + timeout_secs, + ) + .await + } + }) + .await + } + + async fn fetch_credits_status_on_host( + &self, + host: &str, + authorization: &str, + headers: &HashMap, + cookie_by_host: &HashMap, + account_email: Option<&str>, + timeout_secs: u64, + ) -> Result { + let url = format!("https://{host}{CREDITS_STATUS_PATH}"); + let mut req = self + .client + .get(&url) + .timeout(Duration::from_secs(timeout_secs.max(1))) + .header("Accept", "application/json, text/plain, */*") + .header("Accept-Language", "en-US,en;q=0.9") + .header("User-Agent", USER_AGENT) + .header("Sec-Fetch-Dest", "empty") + .header("Sec-Fetch-Mode", "cors") + .header("Sec-Fetch-Site", "same-site"); + + // Captured headers first (except Origin/Referer/Authorization which we pin). + for (name, value) in headers { + if name.eq_ignore_ascii_case("origin") + || name.eq_ignore_ascii_case("referer") + || name.eq_ignore_ascii_case("authorization") + || name.eq_ignore_ascii_case("cookie") + { + continue; + } + req = req.header(name.as_str(), value.as_str()); + } + if let Some(cookie) = cookie_by_host + .get(host) + .or_else(|| cookie_by_host.values().next()) + { + req = req.header("Cookie", cookie); + } + // Fixed Origin/Referer so captured values never widen the first-party boundary. + req = req + .header("Authorization", authorization) + .header("Origin", ORIGIN_REFERER) + .header("Referer", ORIGIN_REFERER); + + let resp = req.send().await?; + let status = resp.status(); + let body = resp.bytes().await?; + if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "ZoomMate API error: HTTP {status}" + ))); + } + + let envelope: CreditsStatusEnvelope = serde_json::from_slice(&body).map_err(|e| { + ProviderError::Parse(format!("ZoomMate credits/status parse failed: {e}")) + })?; + let credit_status = envelope + .data + .and_then(|d| d.credit_status) + .ok_or_else(|| ProviderError::Parse("Missing credit_status object.".into()))?; + + Ok(snapshot_from_credit_status( + &credit_status, + account_email, + Utc::now(), + )) + } +} + +impl Default for ZoomMateProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for ZoomMateProvider { + fn id(&self) -> ProviderId { + ProviderId::ZoomMate + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::Web => { + let usage = self.fetch_via_web(ctx).await?; + Ok(ProviderFetchResult::new(usage, "web")) + } + 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 + } +} + +// ── Pure helpers ───────────────────────────────────────────────────────────── + +fn bearer_header_value(raw: &str) -> String { + let trimmed = raw.trim(); + if trimmed.len() >= 7 && trimmed[..7].eq_ignore_ascii_case("bearer ") { + trimmed.to_string() + } else { + format!("Bearer {trimmed}") + } +} + +fn jwt_exp_unix(token: &str) -> Option { + let raw = bearer_header_value(token); + let token = raw.strip_prefix("Bearer ").unwrap_or(raw.as_str()); + let mut parts = token.split('.'); + let _header = parts.next()?; + let payload_b64 = parts.next()?; + let json = base64url_decode(payload_b64)?; + let value: serde_json::Value = serde_json::from_slice(&json).ok()?; + let exp = value.get("exp")?.as_u64().or_else(|| { + value + .get("exp")? + .as_f64() + .filter(|f| f.is_finite() && *f > 0.0) + .map(|f| f as u64) + })?; + (exp > 0).then_some(exp) +} + +fn base64url_decode(value: &str) -> Option> { + use base64::Engine; + let mut s = value.replace('-', "+").replace('_', "/"); + let pad = (4 - s.len() % 4) % 4; + s.push_str(&"=".repeat(pad)); + base64::engine::general_purpose::STANDARD.decode(s).ok() +} + +fn date_from_millis(raw: Option) -> Option> { + let ms = raw.filter(|&v| v > 0)?; + Utc.timestamp_millis_opt(ms).single() +} + +fn hosts_preferred(preferred: Option<&str>) -> Vec<&'static str> { + let preferred = preferred.map(|h| h.to_ascii_lowercase()); + match preferred.as_deref() { + Some(h) if API_HOSTS.iter().any(|x| x.eq_ignore_ascii_case(h)) => { + let mut out = vec![ + API_HOSTS + .iter() + .copied() + .find(|x| x.eq_ignore_ascii_case(h)) + .unwrap_or(API_HOSTS[0]), + ]; + for host in API_HOSTS { + if !host.eq_ignore_ascii_case(h) { + out.push(*host); + } + } + out + } + _ => API_HOSTS.to_vec(), + } +} + +/// Runs `operation` per host. Auth (401/403 → AuthRequired) and parse errors do +/// **not** fail over; other errors try the next host. +async fn with_api_host_failover( + preferred: Option<&str>, + mut operation: F, +) -> Result +where + F: FnMut(&'static str) -> Fut, + Fut: std::future::Future>, +{ + let hosts = hosts_preferred(preferred); + let mut last_err: Option = None; + for (index, host) in hosts.iter().enumerate() { + match operation(host).await { + Ok(v) => return Ok(v), + Err(ProviderError::AuthRequired) => return Err(ProviderError::AuthRequired), + Err(ProviderError::Parse(msg)) => return Err(ProviderError::Parse(msg)), + Err(e) => { + last_err = Some(e); + if index + 1 < hosts.len() { + tracing::info!("ZoomMate API host unavailable; retrying on the alternate host"); + } + } + } + } + Err(last_err.unwrap_or_else(|| ProviderError::Other("No ZoomMate API host succeeded.".into()))) +} + +/// cURL validation: https, host ∈ {ai.zoom.us, zoommate.zoom.us}, path exactly +/// `/ai-computer/api/v1/credits/status`, no query/fragment, authorization required. +fn is_allowed_capture_url(url: &reqwest::Url) -> bool { + let host = match url.host_str() { + Some(h) => h.to_ascii_lowercase(), + None => return false, + }; + url.scheme().eq_ignore_ascii_case("https") + && API_HOSTS.iter().any(|h| host == *h) + && url.port().is_none() + && url.username().is_empty() + && url.password().is_none() + && url.path() == CREDITS_STATUS_PATH + && url.query().is_none() + && url.fragment().is_none() +} + +fn request_context_from_manual(raw: &str) -> Option { + let raw = raw.trim(); + if raw.is_empty() { + return None; + } + + // Prefer full cURL capture with Authorization. + if curl_capture::looks_like_curl_capture(raw) { + let capture_url = curl_capture::request_url(raw)?; + let parsed = reqwest::Url::parse(&capture_url).ok()?; + if !is_allowed_capture_url(&parsed) { + return None; + } + let capture_host = parsed.host_str()?.to_ascii_lowercase(); + let fields = curl_capture::header_fields(raw); + let authorization = curl_capture::header_value("Authorization", &fields)?; + if authorization.trim().is_empty() { + return None; + } + let mut headers = + curl_capture::forwarded_headers_from_pairs(&fields, FORWARDED_MANUAL_HEADERS); + headers.remove("Authorization"); + let cookie = headers.remove("Cookie"); + let mut cookie_by_host = HashMap::new(); + if let Some(c) = cookie { + cookie_by_host.insert(capture_host.clone(), c); + } + // Pin Origin/Referer at send time; drop any captured values. + headers.remove("Origin"); + headers.remove("Referer"); + return Some(RequestContext { + authorization: bearer_header_value(&authorization), + headers, + cookie_by_host, + preferred_host: Some(capture_host), + account_email: None, + cache_key: None, + }); + } + + None +} + +pub(crate) fn snapshot_from_credit_status( + status: &CreditStatus, + account_email: Option<&str>, + now: DateTime, +) -> UsageSnapshot { + let budget_cap = status.budget_cap.unwrap_or(0.0); + let used_credit = status.used_credit.unwrap_or(0.0); + let is_unlimited = status.is_unlimited.unwrap_or(false); + + let used_percent = + if is_unlimited || budget_cap <= 0.0 || !budget_cap.is_finite() || !used_credit.is_finite() + { + 0.0 + } else { + ((used_credit / budget_cap) * 100.0).clamp(0.0, 100.0) + }; + + let cycle_start = date_from_millis(status.cycle_start_date); + let cycle_end = date_from_millis(status.cycle_end_date); + + let resets_at = if is_unlimited || budget_cap <= 0.0 { + None + } else { + cycle_end + }; + + let window_minutes = match (cycle_start, cycle_end) { + (Some(start), Some(end)) if end > start => { + let mins = (end - start).num_minutes(); + if mins > 0 { + Some(mins.min(u32::MAX as i64) as u32) + } else { + None + } + } + _ => None, + }; + + let primary = RateWindow::with_details( + used_percent, + window_minutes, + resets_at, + Some("Credits".into()), + ); + + let mut snap = UsageSnapshot::new(primary); + snap.updated_at = now; + if let Some(email) = account_email.map(str::trim).filter(|s| !s.is_empty()) { + snap = snap.with_email(email).with_login_method("Cookie"); + } + // Silence unused-field warnings for fields we intentionally decode for forward compat. + let _ = ( + status.remaining_credit, + status.overage_credit, + status.allow_overage, + status.is_quota_available, + ); + snap +} + +// Classify failover errors for pure unit tests. +#[cfg(test)] +pub(crate) fn should_failover(err: &ProviderError) -> bool { + !matches!(err, ProviderError::AuthRequired | ProviderError::Parse(_)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_status_json() -> &'static str { + r#"{ + "data": { + "credit_status": { + "budget_cap": 1000.0, + "used_credit": 250.0, + "remaining_credit": 750.0, + "overage_credit": 0.0, + "allow_overage": false, + "cycle_start_date": 1722470400000, + "cycle_end_date": 1725148800000, + "is_quota_available": true, + "is_unlimited": false + } + }, + "status_code": 200 + }"# + } + + #[test] + fn parses_credits_status_fixture() { + let envelope: CreditsStatusEnvelope = + serde_json::from_str(sample_status_json()).expect("fixture parses"); + let status = envelope.data.unwrap().credit_status.unwrap(); + let snap = snapshot_from_credit_status(&status, Some("user@zoom.us"), Utc::now()); + assert!((snap.primary.used_percent - 25.0).abs() < 0.01); + assert_eq!(snap.primary.reset_description.as_deref(), Some("Credits")); + assert!(snap.primary.resets_at.is_some()); + // 31 days ≈ 44640 minutes (2024-08-01 → 2024-09-01) + assert_eq!(snap.primary.window_minutes, Some(44640)); + assert_eq!(snap.account_email.as_deref(), Some("user@zoom.us")); + assert_eq!(snap.login_method.as_deref(), Some("Cookie")); + } + + #[test] + fn unlimited_or_zero_cap_yields_zero_percent() { + let unlimited = CreditStatus { + budget_cap: Some(100.0), + used_credit: Some(50.0), + is_unlimited: Some(true), + ..Default::default() + }; + let snap = snapshot_from_credit_status(&unlimited, None, Utc::now()); + assert_eq!(snap.primary.used_percent, 0.0); + assert!(snap.primary.resets_at.is_none()); + + let zero_cap = CreditStatus { + budget_cap: Some(0.0), + used_credit: Some(10.0), + is_unlimited: Some(false), + ..Default::default() + }; + let snap = snapshot_from_credit_status(&zero_cap, None, Utc::now()); + assert_eq!(snap.primary.used_percent, 0.0); + } + + #[test] + fn clamps_used_percent_to_100() { + let status = CreditStatus { + budget_cap: Some(100.0), + used_credit: Some(150.0), + is_unlimited: Some(false), + cycle_end_date: Some(1_900_000_000_000), + ..Default::default() + }; + let snap = snapshot_from_credit_status(&status, None, Utc::now()); + assert_eq!(snap.primary.used_percent, 100.0); + } + + #[test] + fn manual_curl_capture_requires_allowed_url_and_authorization() { + let good = "curl 'https://ai.zoom.us/ai-computer/api/v1/credits/status' \ + -H 'Authorization: Bearer tok-abc' -H 'Cookie: session=xyz'"; + let ctx = request_context_from_manual(good).expect("valid capture"); + assert_eq!(ctx.authorization, "Bearer tok-abc"); + assert_eq!(ctx.preferred_host.as_deref(), Some("ai.zoom.us")); + assert_eq!( + ctx.cookie_by_host.get("ai.zoom.us").map(String::as_str), + Some("session=xyz") + ); + + // Wrong path + assert!( + request_context_from_manual( + "curl 'https://ai.zoom.us/ai-computer/api/v1/other' -H 'Authorization: Bearer x'" + ) + .is_none() + ); + // Query rejected + assert!( + request_context_from_manual( + "curl 'https://ai.zoom.us/ai-computer/api/v1/credits/status?x=1' \ + -H 'Authorization: Bearer x'" + ) + .is_none() + ); + // Missing auth + assert!( + request_context_from_manual( + "curl 'https://ai.zoom.us/ai-computer/api/v1/credits/status' -H 'Cookie: a=b'" + ) + .is_none() + ); + // Bad host + assert!( + request_context_from_manual( + "curl 'https://evil.example/ai-computer/api/v1/credits/status' \ + -H 'Authorization: Bearer x'" + ) + .is_none() + ); + } + + #[test] + fn hosts_preferred_promotes_capture_host() { + assert_eq!( + hosts_preferred(Some("zoommate.zoom.us")), + vec!["zoommate.zoom.us", "ai.zoom.us"] + ); + assert_eq!( + hosts_preferred(None), + vec!["ai.zoom.us", "zoommate.zoom.us"] + ); + } + + #[test] + fn failover_skips_auth_and_parse() { + assert!(!should_failover(&ProviderError::AuthRequired)); + assert!(!should_failover(&ProviderError::Parse("x".into()))); + assert!(should_failover(&ProviderError::Other("HTTP 500".into()))); + assert!(should_failover(&ProviderError::NoCookies)); + } + + #[test] + fn bearer_header_normalizes_prefix() { + assert_eq!(bearer_header_value("tok"), "Bearer tok"); + assert_eq!(bearer_header_value("Bearer tok"), "Bearer tok"); + assert_eq!(bearer_header_value("bearer tok"), "bearer tok"); + } + + #[test] + fn jwt_exp_reads_payload() { + // header.payload.sig — payload = {"exp": 2000000000} + let payload = base64::Engine::encode( + &base64::engine::general_purpose::URL_SAFE_NO_PAD, + br#"{"exp":2000000000}"#, + ); + let token = format!("aaa.{payload}.sig"); + assert_eq!(jwt_exp_unix(&token), Some(2_000_000_000)); + assert_eq!(jwt_exp_unix("not-a-jwt"), None); + } + + #[test] + fn cookie_fingerprint_is_stable() { + let mut a = HashMap::new(); + a.insert("ai.zoom.us".into(), "c=1".into()); + a.insert("zoommate.zoom.us".into(), "c=2".into()); + let mut b = HashMap::new(); + b.insert("zoommate.zoom.us".into(), "c=2".into()); + b.insert("ai.zoom.us".into(), "c=1".into()); + assert_eq!(cookie_fingerprint(&a), cookie_fingerprint(&b)); + assert_eq!(cookie_fingerprint(&a).len(), 64); + } +} From e2f13e14070a4a2000ecb2d2249d6135fb90bc75 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:14 +0700 Subject: [PATCH 03/15] Parse Amp subscription windows; fix Chutes/LLMProxy/Grok/Ollama --- rust/src/providers/amp/mod.rs | 153 ++++++++++++++++++++++ rust/src/providers/chutes/mod.rs | 203 +++++++++++++++++++++++++++-- rust/src/providers/grok/mod.rs | 100 ++++++++++++-- rust/src/providers/llmproxy/mod.rs | 55 +++++++- rust/src/providers/ollama/mod.rs | 144 +++++++++++++++++--- 5 files changed, 613 insertions(+), 42 deletions(-) diff --git a/rust/src/providers/amp/mod.rs b/rust/src/providers/amp/mod.rs index 80e3e12561..8d44a989f2 100755 --- a/rust/src/providers/amp/mod.rs +++ b/rust/src/providers/amp/mod.rs @@ -260,6 +260,19 @@ impl Provider for AmpProvider { } } +/// Monthly pace window sentinel used by Amp subscription/pace UI (30 days). +const AMP_MONTHLY_WINDOW_MINUTES: u32 = 30 * 24 * 60; + +/// Parsed Amp subscription (Megawatt-style dual other/orb windows). +#[derive(Debug, Clone, PartialEq)] +pub struct AmpSubscriptionUsage { + pub plan: String, + pub other_used_percent: f64, + pub orb_used_percent: f64, + pub resets_at: chrono::DateTime, + pub reset_description: String, +} + /// Parse Amp Free percentage lines from CLI/display text (upstream 0.42.1+ shape). /// /// Matches lines like: @@ -300,9 +313,101 @@ pub fn parse_amp_free_percent_remaining(text: &str) -> Option { None } +/// Parse Amp subscription display text (Megawatt dual other/orb windows). +/// +/// Matches: +/// `Subscription Megawatt: 42% other usage and 88% orb usage remaining - resets upon renewal in 12 days` +pub fn parse_amp_subscription_usage( + text: &str, + now: chrono::DateTime, +) -> Option { + let re = regex_lite::Regex::new( + r"(?im)^\s*Subscription\s+(.+?):\s*([0-9][0-9,]*(?:\.[0-9]+)?)\s*%\s+other\s+usage\s+and\s+([0-9][0-9,]*(?:\.[0-9]+)?)\s*%\s+orb\s+usage\s+remaining\s*-\s*resets\s+upon\s+renewal\s+in\s+([0-9][0-9,]*)\s+days?(?:\s+-\s+https?://\S+)?\s*$", + ) + .ok()?; + + for line in text.lines() { + let Some(caps) = re.captures(line) else { + continue; + }; + let plan = caps.get(1)?.as_str().trim(); + if plan.is_empty() { + continue; + } + let other_remaining = parse_amp_number(caps.get(2)?.as_str())?; + let orb_remaining = parse_amp_number(caps.get(3)?.as_str())?; + let renewal_days: i64 = caps.get(4)?.as_str().replace(',', "").parse().ok()?; + if renewal_days < 0 { + continue; + } + let reset_description = if renewal_days == 1 { + "renews in 1 day".to_string() + } else { + format!("renews in {renewal_days} days") + }; + return Some(AmpSubscriptionUsage { + plan: plan.to_string(), + other_used_percent: 100.0 - other_remaining.clamp(0.0, 100.0), + orb_used_percent: 100.0 - orb_remaining.clamp(0.0, 100.0), + resets_at: now + chrono::Duration::days(renewal_days), + reset_description, + }); + } + None +} + +/// Build a [`UsageSnapshot`] from Amp Free / subscription display text. +/// +/// Subscription (Megawatt) wins for primary/secondary windows when present: +/// - primary = other usage +/// - secondary = orb usage +/// +/// Free percent path fills primary when there is no subscription match. +pub fn usage_snapshot_from_amp_display_text( + text: &str, + now: chrono::DateTime, +) -> Option { + if let Some(sub) = parse_amp_subscription_usage(text, now) { + // Labels: primary = "Other usage", secondary = "Orb usage" + // (surfaced via provider session/weekly labels + plan login_method). + let other = RateWindow::with_details( + sub.other_used_percent, + Some(AMP_MONTHLY_WINDOW_MINUTES), + Some(sub.resets_at), + Some(sub.reset_description.clone()), + ); + let orb = RateWindow::with_details( + sub.orb_used_percent, + Some(AMP_MONTHLY_WINDOW_MINUTES), + Some(sub.resets_at), + Some(sub.reset_description), + ); + return Some( + UsageSnapshot::new(other) + .with_secondary(orb) + .with_login_method(sub.plan), + ); + } + + let free_used = parse_amp_free_percent_remaining(text)?; + let primary = RateWindow::with_details( + free_used, + Some(24 * 60), + None, + Some("resets daily".to_string()), + ); + Some(UsageSnapshot::new(primary).with_login_method("Amp Free")) +} + +fn parse_amp_number(raw: &str) -> Option { + let value: f64 = raw.replace(',', "").parse().ok()?; + value.is_finite().then_some(value) +} + #[cfg(test)] mod tests { use super::*; + use chrono::{TimeZone, Utc}; #[test] fn dashboard_points_to_current_usage_page() { @@ -337,4 +442,52 @@ mod tests { None ); } + + #[test] + fn parses_megawatt_subscription_dual_windows() { + let now = Utc.with_ymd_and_hms(2026, 7, 1, 12, 0, 0).unwrap(); + let text = "Signed in as user@example.com (Acme)\n\ +Subscription Megawatt: 42% other usage and 88% orb usage remaining - resets upon renewal in 12 days\n"; + let sub = parse_amp_subscription_usage(text, now).expect("subscription"); + assert_eq!(sub.plan, "Megawatt"); + assert!((sub.other_used_percent - 58.0).abs() < f64::EPSILON); + assert!((sub.orb_used_percent - 12.0).abs() < f64::EPSILON); + assert_eq!(sub.reset_description, "renews in 12 days"); + assert_eq!(sub.resets_at, now + chrono::Duration::days(12)); + + let snapshot = usage_snapshot_from_amp_display_text(text, now).expect("snapshot"); + assert!((snapshot.primary.used_percent - 58.0).abs() < f64::EPSILON); + assert_eq!( + snapshot.primary.window_minutes, + Some(AMP_MONTHLY_WINDOW_MINUTES) + ); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("renews in 12 days") + ); + let secondary = snapshot.secondary.expect("orb secondary"); + assert!((secondary.used_percent - 12.0).abs() < f64::EPSILON); + assert_eq!(secondary.window_minutes, Some(AMP_MONTHLY_WINDOW_MINUTES)); + assert_eq!(snapshot.login_method.as_deref(), Some("Megawatt")); + } + + #[test] + fn megawatt_one_day_renewal_wording() { + let now = Utc.with_ymd_and_hms(2026, 7, 1, 0, 0, 0).unwrap(); + let text = "Subscription Megawatt: 0% other usage and 100% orb usage remaining - resets upon renewal in 1 day"; + let sub = parse_amp_subscription_usage(text, now).unwrap(); + assert_eq!(sub.reset_description, "renews in 1 day"); + assert!((sub.other_used_percent - 100.0).abs() < f64::EPSILON); + assert!((sub.orb_used_percent - 0.0).abs() < f64::EPSILON); + } + + #[test] + fn free_path_still_builds_snapshot() { + let now = Utc.with_ymd_and_hms(2026, 7, 1, 0, 0, 0).unwrap(); + let text = "Amp Free: 72% remaining today"; + let snapshot = usage_snapshot_from_amp_display_text(text, now).unwrap(); + assert!((snapshot.primary.used_percent - 28.0).abs() < f64::EPSILON); + assert!(snapshot.secondary.is_none()); + assert_eq!(snapshot.login_method.as_deref(), Some("Amp Free")); + } } diff --git a/rust/src/providers/chutes/mod.rs b/rust/src/providers/chutes/mod.rs index 6561e074a2..9b831a7d03 100644 --- a/rust/src/providers/chutes/mod.rs +++ b/rust/src/providers/chutes/mod.rs @@ -1,4 +1,5 @@ use async_trait::async_trait; +use chrono::{DateTime, TimeZone, Utc}; use reqwest::Client; use serde_json::Value; @@ -116,8 +117,8 @@ fn quota_windows(value: &Value) -> Vec { fn collect_windows(value: &Value, out: &mut Vec) { match value { Value::Object(map) => { - if let Some(percent) = percent_from_object(map) { - out.push(RateWindow::new(percent)); + if let Some(window) = window_from_object(map) { + out.push(window); } for value in map.values() { collect_windows(value, out); @@ -132,6 +133,18 @@ fn collect_windows(value: &Value, out: &mut Vec) { } } +/// Build a rate window from a quota object. +/// +/// Quota counts (`used`/`limit`) go into `reset_description` as detail text +/// (e.g. `"25/100 credits"`). They are never treated as reset schedules. +fn window_from_object(map: &serde_json::Map) -> Option { + let percent = percent_from_object(map)?; + let detail = quota_count_description(map); + // Only parse dedicated reset timestamp keys — never raw quota counts. + let resets_at = first_reset_timestamp(map); + Some(RateWindow::with_details(percent, None, resets_at, detail)) +} + fn percent_from_object(map: &serde_json::Map) -> Option { for key in [ "usage_percent", @@ -143,18 +156,136 @@ fn percent_from_object(map: &serde_json::Map) -> Option { return Some(if v <= 1.0 { v * 100.0 } else { v }); } } - let used = ["used", "usage", "current_usage", "currentUsage"] - .iter() - .find_map(|k| map.get(*k).and_then(Value::as_f64)); - let limit = ["limit", "quota", "quota_limit", "quotaLimit", "total"] - .iter() - .find_map(|k| map.get(*k).and_then(Value::as_f64)); - match (used, limit) { - (Some(used), Some(limit)) if limit > 0.0 => Some(used / limit * 100.0), + let used = first_f64(map, &["used", "usage", "current_usage", "currentUsage"]); + let limit = first_f64( + map, + &["limit", "quota", "quota_limit", "quotaLimit", "total"], + ); + let remaining = first_f64(map, &["remaining", "remaining_quota", "remainingQuota"]); + match (used, limit, remaining) { + (Some(used), Some(limit), _) if limit > 0.0 => Some((used / limit) * 100.0), + (None, Some(limit), Some(remaining)) if limit > 0.0 => { + Some(((limit - remaining).max(0.0) / limit) * 100.0) + } + (Some(used), None, Some(remaining)) => { + let limit = used + remaining; + (limit > 0.0).then_some((used / limit) * 100.0) + } + _ => None, + } +} + +fn quota_count_description(map: &serde_json::Map) -> Option { + let used = first_f64(map, &["used", "usage", "current_usage", "currentUsage"]); + let mut limit = first_f64( + map, + &["limit", "quota", "quota_limit", "quotaLimit", "total"], + ); + let remaining = first_f64(map, &["remaining", "remaining_quota", "remainingQuota"]); + + if limit.is_none() + && let (Some(used), Some(remaining)) = (used, remaining) + { + limit = Some(used + remaining); + } + let limit = limit.filter(|l| *l > 0.0)?; + let used = match used { + Some(u) => u, + None => remaining.map(|r| (limit - r).max(0.0))?, + }; + let unit = first_str(map, &["unit", "units", "quota_unit", "quotaUnit"]).unwrap_or("credits"); + Some(format!( + "{}/{} {}", + format_quota_amount(used), + format_quota_amount(limit), + unit + )) +} + +fn first_reset_timestamp(map: &serde_json::Map) -> Option> { + for key in [ + "resets_at", + "resetsAt", + "reset_at", + "resetAt", + "reset_time", + "resetTime", + "renews_at", + "renewsAt", + ] { + if let Some(dt) = map.get(key).and_then(parse_timestamp_value) { + return Some(dt); + } + } + None +} + +fn parse_timestamp_value(value: &Value) -> Option> { + match value { + Value::String(raw) => { + let text = raw.trim(); + if text.is_empty() { + return None; + } + // ISO-8601 / RFC3339 only for strings that look like dates. + // Bare small numbers as strings are quota counts, not schedules. + if let Ok(dt) = DateTime::parse_from_rfc3339(text) { + return Some(dt.with_timezone(&Utc)); + } + if let Ok(number) = text.parse::() { + return epoch_to_datetime(number); + } + None + } + Value::Number(n) => n.as_f64().and_then(epoch_to_datetime), _ => None, } } +fn epoch_to_datetime(value: f64) -> Option> { + if !value.is_finite() || value <= 0.0 { + return None; + } + // Reject small integers that are clearly quota counts, not unix epochs. + // Unix seconds ~1e9; ms ~1e12. Quota counts are typically << 1e8. + if value < 1_000_000_000.0 { + return None; + } + let seconds = if value > 10_000_000_000.0 { + value / 1000.0 + } else { + value + }; + Utc.timestamp_opt(seconds as i64, 0).single() +} + +fn first_f64(map: &serde_json::Map, keys: &[&str]) -> Option { + keys.iter() + .find_map(|k| map.get(*k).and_then(Value::as_f64)) +} + +fn first_str<'a>(map: &'a serde_json::Map, keys: &[&str]) -> Option<&'a str> { + keys.iter() + .find_map(|k| map.get(*k).and_then(Value::as_str)) + .map(str::trim) + .filter(|s| !s.is_empty()) +} + +fn format_quota_amount(value: f64) -> String { + if (value - value.round()).abs() < 0.0001 { + format!("{}", value.round() as i64) + } else { + let mut text = format!("{value:.2}"); + while text.contains('.') && text.ends_with('0') { + text.pop(); + } + if text.ends_with('.') { + text.pop(); + } + text + } +} + #[cfg(test)] mod tests { use super::*; @@ -165,4 +296,56 @@ mod tests { snapshot_from_usage(&serde_json::json!({"quotas":[{"used":25,"limit":100}]})); assert_eq!(snapshot.primary.used_percent, 25.0); } + + #[test] + fn quota_counts_go_to_reset_description_not_schedule() { + let snapshot = snapshot_from_usage(&serde_json::json!({ + "quotas": [{ + "used": 25, + "limit": 100, + "unit": "credits", + "remaining": 75 + }] + })); + assert_eq!(snapshot.primary.used_percent, 25.0); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("25/100 credits") + ); + assert!(snapshot.primary.resets_at.is_none()); + } + + #[test] + fn small_numeric_fields_are_not_reset_schedules() { + // A mis-keyed payload where "reset_time" is accidentally a remaining count. + // Values below unix-epoch range must not become resets_at. + let snapshot = snapshot_from_usage(&serde_json::json!({ + "quotas": [{ + "used": 10, + "limit": 50, + "reset_time": 40 + }] + })); + assert_eq!(snapshot.primary.used_percent, 20.0); + assert!(snapshot.primary.resets_at.is_none()); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("10/50 credits") + ); + } + + #[test] + fn iso_reset_timestamps_still_parse() { + let snapshot = snapshot_from_usage(&serde_json::json!({ + "quotas": [{ + "used": 1, + "limit": 2, + "resets_at": "2026-08-01T00:00:00Z" + }] + })); + assert_eq!( + snapshot.primary.resets_at, + Some(Utc.with_ymd_and_hms(2026, 8, 1, 0, 0, 0).unwrap()) + ); + } } diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index bf154417b9..27ae9aab45 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -96,6 +96,28 @@ impl GrokProvider { )) } + /// Cookie refresh path (upstream #2458): + /// 1. Try last validated cached cookie header (background reuse) + /// 2. On miss/auth failure: re-import browser cookies, validate, cache + async fn fetch_with_cookie_refresh(&self) -> Result { + use crate::browser::cookie_cache::CookieHeaderCache; + + if let Some(cached) = CookieHeaderCache::load(ProviderId::Grok) { + match self.fetch_with_cookie(&cached.cookie_header).await { + Ok(result) => return Ok(result), + Err(err) if is_cookie_authentication_failure(&err) => { + CookieHeaderCache::clear(ProviderId::Grok); + } + Err(err) => return Err(err), + } + } + + let cookie_header = crate::providers::browser_cookie_header(&["grok.com"])?; + let result = self.fetch_with_cookie(&cookie_header).await?; + let _ = CookieHeaderCache::store(ProviderId::Grok, &cookie_header, "browser"); + Ok(result) + } + async fn fetch_billing( &self, authorization: Option, @@ -181,16 +203,14 @@ impl Provider for GrokProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { SourceMode::Auto | SourceMode::Web => { - if let Some(ref cookie_header) = ctx.manual_cookie_header { + if let Some(cookie_header) = &ctx.manual_cookie_header { return self.fetch_with_cookie(cookie_header).await; } - match crate::providers::browser_cookie_header(&["grok.com"]) { - Ok(cookie_header) => match self.fetch_with_cookie(&cookie_header).await { - Ok(result) => return Ok(result), - Err(ProviderError::AuthRequired) => {} - Err(e) => return Err(e), - }, - Err(ProviderError::NoCookies) => {} + // Explicit cookie path: try validated cache, else re-import browser cookies, + // validate the session, and cache the header for background reuse (#2458). + match self.fetch_with_cookie_refresh().await { + Ok(result) => return Ok(result), + Err(ProviderError::AuthRequired) | Err(ProviderError::NoCookies) => {} Err(e) => return Err(e), } let credentials = Self::load_credentials()?; @@ -332,6 +352,33 @@ fn validate_grpc_headers(headers: &reqwest::header::HeaderMap) -> Result<(), Pro Ok(()) } +/// Whether a cookie-path error should invalidate the cached browser session. +fn is_cookie_authentication_failure(err: &ProviderError) -> bool { + matches!(err, ProviderError::AuthRequired) +} + +/// Decide the next cookie-refresh step given cache presence and last error. +/// Pure helper for unit tests of the #2458 refresh flow. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CookieRefreshAction { + UseCached, + ReimportBrowser, + GiveUp, +} + +fn cookie_refresh_action( + has_cached_header: bool, + last_error: Option<&ProviderError>, +) -> CookieRefreshAction { + match last_error { + None if has_cached_header => CookieRefreshAction::UseCached, + None => CookieRefreshAction::ReimportBrowser, + Some(err) if is_cookie_authentication_failure(err) => CookieRefreshAction::ReimportBrowser, + Some(ProviderError::NoCookies) => CookieRefreshAction::ReimportBrowser, + Some(_) => CookieRefreshAction::GiveUp, + } +} + fn parse_grpc_web_response(data: &[u8]) -> Result { let frames = grpc_web_data_frames(data); if frames.is_empty() { @@ -528,10 +575,45 @@ mod tests { assert_eq!(parsed.access_token, "oidc"); assert_eq!(parsed.login_method().as_deref(), Some("SuperGrok")); } - #[test] fn splits_grpc_web_data_frames() { let data = [0, 0, 0, 0, 2, 1, 2, 0x80, 0, 0, 0, 1, b'x']; assert_eq!(grpc_web_data_frames(&data), vec![vec![1, 2]]); } + + #[test] + fn cookie_refresh_uses_cache_when_present() { + assert_eq!( + cookie_refresh_action(true, None), + CookieRefreshAction::UseCached + ); + } + + #[test] + fn cookie_refresh_reimports_on_auth_failure() { + assert_eq!( + cookie_refresh_action(true, Some(&ProviderError::AuthRequired)), + CookieRefreshAction::ReimportBrowser + ); + assert_eq!( + cookie_refresh_action(false, None), + CookieRefreshAction::ReimportBrowser + ); + } + + #[test] + fn cookie_refresh_gives_up_on_non_auth_errors() { + assert_eq!( + cookie_refresh_action(true, Some(&ProviderError::Other("network down".into()))), + CookieRefreshAction::GiveUp + ); + } + + #[test] + fn is_cookie_auth_failure_only_auth_required() { + assert!(is_cookie_authentication_failure( + &ProviderError::AuthRequired + )); + assert!(!is_cookie_authentication_failure(&ProviderError::NoCookies)); + } } diff --git a/rust/src/providers/llmproxy/mod.rs b/rust/src/providers/llmproxy/mod.rs index b4d581189b..230ee0ce79 100644 --- a/rust/src/providers/llmproxy/mod.rs +++ b/rust/src/providers/llmproxy/mod.rs @@ -300,10 +300,7 @@ fn parse_summary(data: &[u8]) -> Result { .iter() .filter_map(|group| group.remaining_percent) .min_by(|left, right| left.total_cmp(right)), - next_reset_at: quota_groups - .iter() - .filter_map(|group| parse_date(group.reset_time.as_deref())) - .min(), + next_reset_at: next_reset_at_from_groups("a_groups, Utc::now()), top_providers: provider_summaries, }) } @@ -369,6 +366,16 @@ fn token_total(tokens: Option<&TokenStats>) -> u64 { .unwrap_or(0) } +/// Pick the soonest upcoming reset. Already-elapsed times are ignored so a stale +/// past reset cannot win over a real future one (upstream #2335). +fn next_reset_at_from_groups(groups: &[QuotaGroup], now: DateTime) -> Option> { + groups + .iter() + .filter_map(|group| parse_date(group.reset_time.as_deref())) + .filter(|t| *t > now) + .min() +} + fn parse_date(raw: Option<&str>) -> Option> { DateTime::parse_from_rfc3339(raw?) .ok() @@ -461,4 +468,44 @@ mod tests { assert_eq!(snapshot.primary.used_percent, 74.5); assert_eq!(snapshot.extra_rate_windows.len(), 2); } + + #[test] + fn next_reset_at_ignores_elapsed_resets() { + let now = DateTime::parse_from_rfc3339("2026-05-15T12:00:00Z") + .unwrap() + .with_timezone(&Utc); + let groups = vec![ + QuotaGroup { + remaining_percent: Some(10.0), + reset_time: Some("2026-05-10T00:00:00Z".into()), // past + }, + QuotaGroup { + remaining_percent: Some(40.0), + reset_time: Some("2026-05-20T00:00:00Z".into()), // future + }, + QuotaGroup { + remaining_percent: Some(50.0), + reset_time: Some("2026-05-18T00:00:00Z".into()), // sooner future + }, + ]; + let next = next_reset_at_from_groups(&groups, now).expect("future reset"); + assert_eq!( + next, + DateTime::parse_from_rfc3339("2026-05-18T00:00:00Z") + .unwrap() + .with_timezone(&Utc) + ); + } + + #[test] + fn next_reset_at_none_when_all_elapsed() { + let now = DateTime::parse_from_rfc3339("2026-06-01T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + let groups = vec![QuotaGroup { + remaining_percent: Some(10.0), + reset_time: Some("2026-05-20T00:00:00Z".into()), + }]; + assert!(next_reset_at_from_groups(&groups, now).is_none()); + } } diff --git a/rust/src/providers/ollama/mod.rs b/rust/src/providers/ollama/mod.rs index 8ea5047c92..f7a4f82f02 100755 --- a/rust/src/providers/ollama/mod.rs +++ b/rust/src/providers/ollama/mod.rs @@ -88,8 +88,27 @@ impl OllamaProvider { .map_err(|e| ProviderError::Other(e.to_string()))?; let start_url = Url::parse(OLLAMA_SETTINGS_URL).map_err(|e| ProviderError::Other(e.to_string()))?; - let html = fetch_settings_html_at(&client, &cookies, start_url).await?; - self.parse_usage_html(&html) + + match fetch_settings_html_at(&client, &cookies, start_url.clone()).await { + Ok(html) => { + // Only cache non-manual browser/validated sessions for reuse. + if ctx.manual_cookie_header.is_none() { + Self::cache_validated_session_cookie(&cookies); + } + self.parse_usage_html(&html) + } + Err(ProviderError::AuthRequired) if ctx.manual_cookie_header.is_none() => { + // Cached/imported session expired — clear and re-import once. + Self::invalidate_cached_session_cookie(); + let fresh = resolve_browser_cookie_header(true)? + .map(OllamaCookieSource::Manual) + .ok_or(ProviderError::AuthRequired)?; + let html = fetch_settings_html_at(&client, &fresh, start_url).await?; + Self::cache_validated_session_cookie(&fresh); + self.parse_usage_html(&html) + } + Err(err) => Err(err), + } } async fn fetch_usage_api(&self, ctx: &FetchContext) -> Result { @@ -234,13 +253,16 @@ impl OllamaProvider { } } - /// Resolve cookies from manual cookies, browser import, or context. + /// Resolve cookies from manual cookies, validated cache, or browser import. + /// + /// Upstream #2404: reuse the last validated browser session cookie header + /// across refreshes until auth fails, then re-import. fn resolve_cookie_source( &self, ctx: &FetchContext, ) -> Result { // Check manual cookie header first - if let Some(ref cookie) = ctx.manual_cookie_header + if let Some(cookie) = &ctx.manual_cookie_header && let Some(header) = Self::normalize_cookie_header(cookie) { return has_recognized_ollama_session_cookie(&header) @@ -248,24 +270,80 @@ impl OllamaProvider { .ok_or(ProviderError::NoCookies); } - // Try browser cookie extraction - match crate::providers::browser_cookies_for_domain(OLLAMA_COOKIE_DOMAIN) { - Ok(cookies) => { - let source = OllamaCookieSource::Browser(cookies); - source - .header_for_url( - &Url::parse(OLLAMA_SETTINGS_URL) - .map_err(|e| ProviderError::Other(e.to_string()))?, - ) - .is_some() - .then_some(source) - .ok_or(ProviderError::NoCookies) - } - Err(ProviderError::NoCookies) => Err(ProviderError::NoCookies), - Err(err) => Err(err), + match resolve_browser_cookie_header(false)? { + Some(header) => Ok(OllamaCookieSource::Manual(header)), + None => Err(ProviderError::NoCookies), + } + } + + /// After a successful web fetch, cache the validated browser/manual session header. + fn cache_validated_session_cookie(source: &OllamaCookieSource) { + use crate::browser::cookie_cache::CookieHeaderCache; + if let Some(header) = source.header_for_url( + &Url::parse(OLLAMA_SETTINGS_URL) + .unwrap_or_else(|_| Url::parse("https://ollama.com/settings").expect("static url")), + ) { + let label = match source { + OllamaCookieSource::Manual(_) => "validated", + OllamaCookieSource::Browser(_) => "browser", + }; + let _ = CookieHeaderCache::store(ProviderId::Ollama, &header, label); + } + } + + /// Clear cached session after auth failure so the next refresh re-imports. + fn invalidate_cached_session_cookie() { + use crate::browser::cookie_cache::CookieHeaderCache; + CookieHeaderCache::clear(ProviderId::Ollama); + } +} + +/// Resolve a browser/session cookie header for Ollama. +/// +/// When `force_reimport` is false, prefers the last validated cached header +/// (upstream #2404). On force or cache miss, imports from the browser. +fn resolve_browser_cookie_header(force_reimport: bool) -> Result, ProviderError> { + use crate::browser::cookie_cache::CookieHeaderCache; + + if !force_reimport + && let Some(cached) = CookieHeaderCache::load(ProviderId::Ollama) + && has_recognized_ollama_session_cookie(&cached.cookie_header) + { + return Ok(Some(cached.cookie_header)); + } + + match crate::providers::browser_cookies_for_domain(OLLAMA_COOKIE_DOMAIN) { + Ok(cookies) => { + let url = + Url::parse(OLLAMA_SETTINGS_URL).map_err(|e| ProviderError::Other(e.to_string()))?; + Ok(ollama_cookie_header_for_url(&cookies, &url) + .filter(|h| has_recognized_ollama_session_cookie(h))) } + Err(ProviderError::NoCookies) => Ok(None), + Err(err) => Err(err), } +} +/// Pure decision helper for the Ollama session reuse path (unit-tested). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OllamaSessionAction { + UseCached, + ReimportBrowser, +} + +fn ollama_session_action( + has_cached_validated: bool, + auth_failed: bool, + force_reimport: bool, +) -> OllamaSessionAction { + if force_reimport || auth_failed || !has_cached_validated { + OllamaSessionAction::ReimportBrowser + } else { + OllamaSessionAction::UseCached + } +} + +impl OllamaProvider { /// Parse usage data from the Ollama settings HTML page fn parse_usage_html(&self, html: &str) -> Result { // Check if we're signed out @@ -971,4 +1049,32 @@ mod tests { assert_eq!(weekly.used_percent, 84.0); assert!(weekly.resets_at.is_some()); } + + #[test] + fn session_action_reuses_cached_until_auth_fails() { + assert_eq!( + ollama_session_action(true, false, false), + OllamaSessionAction::UseCached + ); + assert_eq!( + ollama_session_action(true, true, false), + OllamaSessionAction::ReimportBrowser + ); + assert_eq!( + ollama_session_action(false, false, false), + OllamaSessionAction::ReimportBrowser + ); + assert_eq!( + ollama_session_action(true, false, true), + OllamaSessionAction::ReimportBrowser + ); + } + + #[test] + fn recognized_session_cookie_required_for_cache_reuse() { + assert!(has_recognized_ollama_session_cookie( + "__Secure-session=abc123; path=/" + )); + assert!(!has_recognized_ollama_session_cookie("foo=bar; baz=qux")); + } } From 07a12ac8465c6750408d6b462675bd530e7bae16 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:17 +0700 Subject: [PATCH 04/15] Add Claude prepaid balance, Daily Routines toggle, settings --- rust/src/providers/claude/oauth/mod.rs | 106 +++++++++- rust/src/providers/claude/web_api.rs | 270 ++++++++++++++++++++++--- rust/src/settings.rs | 25 +++ rust/src/settings/raw.rs | 24 +++ rust/src/settings/tests.rs | 1 + 5 files changed, 383 insertions(+), 43 deletions(-) diff --git a/rust/src/providers/claude/oauth/mod.rs b/rust/src/providers/claude/oauth/mod.rs index 84450e57fb..e52f378495 100644 --- a/rust/src/providers/claude/oauth/mod.rs +++ b/rust/src/providers/claude/oauth/mod.rs @@ -352,6 +352,16 @@ impl ClaudeOAuthFetcher { &self, response: &OAuthUsageResponse, credentials: &ClaudeOAuthCredentials, + ) -> UsageSnapshot { + let show_routines = crate::settings::Settings::load().claude_daily_routines_usage_visible; + self.build_usage_snapshot_with_options(response, credentials, show_routines) + } + + fn build_usage_snapshot_with_options( + &self, + response: &OAuthUsageResponse, + credentials: &ClaudeOAuthCredentials, + show_routines: bool, ) -> UsageSnapshot { // Primary: 5-hour session window let primary = response @@ -390,10 +400,18 @@ impl ClaudeOAuthFetcher { usage = usage.with_model_specific(sonnet); } - if let Some(window) = response - .seven_day_routines - .as_ref() - .and_then(|w| Self::to_rate_window(w, Some(10080))) + // Model-scoped weekly limits first; Daily Routines last (upstream order). + usage + .extra_rate_windows + .extend(super::scoped_weekly::scoped_weekly_windows( + &response.limits, + )); + + if show_routines + && let Some(window) = response + .seven_day_routines + .as_ref() + .and_then(|w| Self::to_rate_window(w, Some(10080))) { usage.extra_rate_windows.push(NamedRateWindow::new( "claude-routines", @@ -401,14 +419,9 @@ impl ClaudeOAuthFetcher { window, )); } - usage - .extra_rate_windows - .extend(super::scoped_weekly::scoped_weekly_windows( - &response.limits, - )); // Login method from rate limit tier or default - if let Some(ref tier) = credentials.rate_limit_tier { + if let Some(tier) = &credentials.rate_limit_tier { usage = usage.with_login_method(super::claude_plan_label(tier)); } else { usage = usage.with_login_method("Claude (OAuth)"); @@ -686,4 +699,77 @@ mod tests { assert!(message.contains("rate limited")); assert!(message.contains("credentials were preserved")); } + + #[test] + fn oauth_extras_put_scoped_weekly_before_routines() { + let response: OAuthUsageResponse = serde_json::from_str( + r#"{ + "five_hour": {"utilization": 10.0}, + "seven_day_routines": {"utilization": 5.0}, + "limits": [{ + "kind": "weekly_scoped", + "group": "weekly", + "percent": 7, + "resets_at": "2026-05-29T10:00:00Z", + "scope": {"model": {"display_name": "Fable"}} + }] + }"#, + ) + .expect("oauth body"); + + let credentials = ClaudeOAuthCredentials { + access_token: "token".to_string(), + refresh_token: None, + expires_at: None, + scopes: vec![], + rate_limit_tier: None, + }; + let usage = ClaudeOAuthFetcher::new().build_usage_snapshot(&response, &credentials); + + let ids: Vec<&str> = usage + .extra_rate_windows + .iter() + .map(|w| w.id.as_str()) + .collect(); + assert_eq!(ids, vec!["claude-weekly-scoped-fable", "claude-routines"]); + } + + #[test] + fn oauth_extras_hide_routines_when_disabled() { + let response: OAuthUsageResponse = serde_json::from_str( + r#"{ + "five_hour": {"utilization": 10.0}, + "seven_day_routines": {"utilization": 5.0}, + "limits": [{ + "kind": "weekly_scoped", + "group": "weekly", + "percent": 7, + "scope": {"model": {"display_name": "Fable"}} + }] + }"#, + ) + .expect("oauth body"); + + let credentials = ClaudeOAuthCredentials { + access_token: "token".to_string(), + refresh_token: None, + expires_at: None, + scopes: vec![], + rate_limit_tier: None, + }; + let usage = ClaudeOAuthFetcher::new().build_usage_snapshot_with_options( + &response, + &credentials, + false, + ); + + assert!( + usage + .extra_rate_windows + .iter() + .all(|w| w.id != "claude-routines") + ); + assert_eq!(usage.extra_rate_windows.len(), 1); + assert_eq!(usage.extra_rate_windows[0].id, "claude-weekly-scoped-fable"); + } } diff --git a/rust/src/providers/claude/web_api.rs b/rust/src/providers/claude/web_api.rs index d4fa736b3e..5de3770525 100755 --- a/rust/src/providers/claude/web_api.rs +++ b/rust/src/providers/claude/web_api.rs @@ -342,39 +342,26 @@ impl ClaudeWebApiFetcher { snapshot = snapshot.with_model_specific(m); } - for (id, title, window) in [ - ( - "claude-oauth-apps", - "OAuth apps", - usage - .seven_day_oauth_apps - .as_ref() - .map(|w| self.to_rate_window(w, Some(10080))), - ), - ( - "claude-routines", - "Daily Routines", - usage - .seven_day_routines - .as_ref() - .map(|w| self.to_rate_window(w, Some(10080))), - ), - ] { - if let Some(window) = window { - snapshot - .extra_rate_windows - .push(NamedRateWindow::new(id, title, window)); - } - } - snapshot - .extra_rate_windows - .extend(super::scoped_weekly::scoped_weekly_windows(&usage.limits)); + let show_routines = crate::settings::Settings::load().claude_daily_routines_usage_visible; + append_web_extra_windows( + &mut snapshot, + usage + .seven_day_oauth_apps + .as_ref() + .map(|w| self.to_rate_window(w, Some(10080))), + super::scoped_weekly::scoped_weekly_windows(&usage.limits), + usage + .seven_day_routines + .as_ref() + .map(|w| self.to_rate_window(w, Some(10080))), + show_routines, + ); - if let Some(ref acc) = account { - if let Some(ref email) = acc.email_address { + if let Some(acc) = &account { + if let Some(email) = &acc.email_address { snapshot = snapshot.with_email(email.clone()); } - if let Some(ref tier) = acc.rate_limit_tier { + if let Some(tier) = &acc.rate_limit_tier { snapshot = snapshot.with_login_method(super::claude_plan_label(tier)); } } @@ -382,9 +369,10 @@ impl ClaudeWebApiFetcher { let mut result = ProviderFetchResult::new(snapshot, "web"); // Add cost info if available - if let Some(extra) = extra_usage - && extra.is_enabled.unwrap_or(false) - { + let mut cost = extra_usage.and_then(|extra| { + if !extra.is_enabled.unwrap_or(false) { + return None; + } let used_cents = extra.used_credits.unwrap_or(0.0); let limit_cents = extra.monthly_credit_limit; let currency = extra.currency.unwrap_or_else(|| "USD".to_string()); @@ -398,7 +386,21 @@ impl ClaudeWebApiFetcher { if let Some(limit) = limit_cents { cost = cost.with_limit(limit / 100.0); } + Some(cost) + }); + + // Best-effort prepaid Extra usage balance (non-fatal). + // Gate: cookie session is already available on this path; skip only when + // cookie source is explicitly off. + let settings = crate::settings::Settings::load(); + let cookie_source = settings.claude_cookie_source(); + if !cookie_source.eq_ignore_ascii_case("off") + && let Some(balance) = self.get_prepaid_credits(&org_id, &headers).await + { + cost = Some(apply_prepaid_balance(balance, cost)); + } + if let Some(cost) = cost { result = result.with_cost(cost); } @@ -554,6 +556,35 @@ impl ClaudeWebApiFetcher { parse_json_with_body(response, "extra usage").await } + /// Best-effort prepaid Extra usage balance. Non-fatal on any failure. + async fn get_prepaid_credits( + &self, + org_id: &str, + headers: &reqwest::header::HeaderMap, + ) -> Option { + let url = format!( + "{}/organizations/{}/prepaid/credits", + Self::BASE_URL, + org_id + ); + + let response = self + .client + .get(&url) + .headers(headers.clone()) + .timeout(std::time::Duration::from_secs(15)) + .send() + .await + .ok()?; + + if !response.status().is_success() { + return None; + } + + let body = response.text().await.ok()?; + parse_prepaid_balance(&body) + } + /// Get account info async fn get_account_info( &self, @@ -658,6 +689,75 @@ fn cookie_value(cookie_header: &str, name: &str) -> Option { }) } +#[derive(Debug, Clone, PartialEq)] +struct PrepaidBalance { + amount_dollars: f64, + currency_code: String, +} + +#[derive(Debug, Deserialize)] +struct PrepaidCreditsResponse { + amount: f64, + currency: String, +} + +/// Parse `{ amount: cents, currency }` → dollars when finite ≥ 0. +fn parse_prepaid_balance(body: &str) -> Option { + let response: PrepaidCreditsResponse = serde_json::from_str(body).ok()?; + if !response.amount.is_finite() || response.amount < 0.0 { + return None; + } + let currency = response.currency.trim().to_ascii_uppercase(); + if currency.is_empty() { + return None; + } + Some(PrepaidBalance { + amount_dollars: response.amount / 100.0, + currency_code: currency, + }) +} + +/// Attach prepaid balance onto an existing same-currency cost, otherwise create +/// an "Extra usage" snapshot carrying only the balance. +fn apply_prepaid_balance(balance: PrepaidBalance, existing: Option) -> CostSnapshot { + match existing { + Some(cost) + if cost + .currency_code + .eq_ignore_ascii_case(&balance.currency_code) => + { + cost.with_balance(balance.amount_dollars) + } + _ => CostSnapshot::new(0.0, balance.currency_code, "Extra usage") + .with_balance(balance.amount_dollars), + } +} + +/// Push extras in upstream order: oauth-apps → scoped weekly → routines (optional). +fn append_web_extra_windows( + snapshot: &mut UsageSnapshot, + oauth_apps: Option, + scoped_weekly: Vec, + routines: Option, + show_routines: bool, +) { + if let Some(window) = oauth_apps { + snapshot.extra_rate_windows.push(NamedRateWindow::new( + "claude-oauth-apps", + "OAuth apps", + window, + )); + } + snapshot.extra_rate_windows.extend(scoped_weekly); + if show_routines && let Some(window) = routines { + snapshot.extra_rate_windows.push(NamedRateWindow::new( + "claude-routines", + "Daily Routines", + window, + )); + } +} + #[cfg(test)] mod tests { use super::{AccountResponse, ClaudeWebApiFetcher, UsageWindow, cookie_value}; @@ -937,4 +1037,108 @@ mod tests { assert_eq!(extra.monthly_credit_limit, Some(2000.0)); assert_eq!(extra.used_credits, Some(550.0)); } + + #[test] + fn parse_prepaid_balance_converts_cents_to_dollars() { + let balance = super::parse_prepaid_balance(r#"{"amount": 2550, "currency": "usd"}"#) + .expect("prepaid balance"); + assert!((balance.amount_dollars - 25.5).abs() < f64::EPSILON); + assert_eq!(balance.currency_code, "USD"); + } + + #[test] + fn parse_prepaid_balance_rejects_negative_or_non_finite() { + assert!(super::parse_prepaid_balance(r#"{"amount": -1, "currency": "USD"}"#).is_none()); + assert!(super::parse_prepaid_balance(r#"{"amount": 10, "currency": " "}"#).is_none()); + } + + #[test] + fn apply_prepaid_balance_attaches_to_same_currency_cost() { + let existing = crate::core::CostSnapshot::new(1.0, "USD", "Monthly").with_limit(20.0); + let balance = super::PrepaidBalance { + amount_dollars: 12.34, + currency_code: "USD".into(), + }; + let cost = super::apply_prepaid_balance(balance, Some(existing)); + assert_eq!(cost.balance, Some(12.34)); + assert!((cost.used - 1.0).abs() < f64::EPSILON); + assert_eq!(cost.limit, Some(20.0)); + assert_eq!(cost.period, "Monthly"); + } + + #[test] + fn apply_prepaid_balance_creates_extra_usage_when_missing_or_mismatch() { + let balance = super::PrepaidBalance { + amount_dollars: 5.0, + currency_code: "USD".into(), + }; + let created = super::apply_prepaid_balance(balance.clone(), None); + assert_eq!(created.balance, Some(5.0)); + assert_eq!(created.period, "Extra usage"); + assert!((created.used - 0.0).abs() < f64::EPSILON); + + let eur = crate::core::CostSnapshot::new(2.0, "EUR", "Monthly"); + let replaced = super::apply_prepaid_balance(balance, Some(eur)); + assert_eq!(replaced.currency_code, "USD"); + assert_eq!(replaced.period, "Extra usage"); + assert_eq!(replaced.balance, Some(5.0)); + } + + #[test] + fn web_extras_order_oauth_scoped_then_routines() { + use crate::core::{NamedRateWindow, RateWindow, UsageSnapshot}; + + let mut snapshot = UsageSnapshot::new(RateWindow::new(10.0)); + super::append_web_extra_windows( + &mut snapshot, + Some(RateWindow::new(1.0)), + vec![NamedRateWindow::new( + "claude-weekly-scoped-fable", + "Fable only", + RateWindow::new(2.0), + )], + Some(RateWindow::new(3.0)), + true, + ); + + let ids: Vec<&str> = snapshot + .extra_rate_windows + .iter() + .map(|w| w.id.as_str()) + .collect(); + assert_eq!( + ids, + vec![ + "claude-oauth-apps", + "claude-weekly-scoped-fable", + "claude-routines" + ] + ); + } + + #[test] + fn web_extras_hide_routines_when_disabled() { + use crate::core::{NamedRateWindow, RateWindow, UsageSnapshot}; + + let mut snapshot = UsageSnapshot::new(RateWindow::new(10.0)); + super::append_web_extra_windows( + &mut snapshot, + Some(RateWindow::new(1.0)), + vec![NamedRateWindow::new( + "claude-weekly-scoped-fable", + "Fable only", + RateWindow::new(2.0), + )], + Some(RateWindow::new(3.0)), + false, + ); + + assert!( + snapshot + .extra_rate_windows + .iter() + .all(|w| w.id != "claude-routines") + ); + assert_eq!(snapshot.extra_rate_windows.len(), 2); + } } diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 262db94e29..8c3826c4aa 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -258,12 +258,30 @@ pub struct Settings { /// Defaults on so upgrades keep the icon pinned to the taskbar notification area. #[serde(default = "default_true")] pub promote_tray_icon: bool, + + /// When true, show Claude Daily Routines usage in extras. + /// Defaults on to match upstream visibility. + #[serde(default = "default_true")] + pub claude_daily_routines_usage_visible: bool, + + /// Optional work-week length [2,6] for session-equivalent weekly forecast. + /// `None` uses wall-clock time until weekly reset. + #[serde(default)] + pub weekly_progress_work_days: Option, + + /// Alibaba Token Plan API region: "cn" | "intl" | "cn-personal" | "intl-personal". + #[serde(default = "default_alibaba_token_plan_region")] + pub alibaba_token_plan_region: String, } fn default_window_scale_percent() -> u16 { 100 } +fn default_alibaba_token_plan_region() -> String { + "cn".to_string() +} + pub fn clamp_window_scale_percent(value: u16) -> u16 { value.clamp(100, 250) } @@ -375,6 +393,7 @@ const DEFAULT_PROVIDER_SOURCE: &str = "auto"; fn default_api_region(id: ProviderId) -> &'static str { match id { ProviderId::Alibaba => crate::providers::AlibabaRegion::Singapore.settings_value(), + ProviderId::AlibabaTokenPlan => "cn", ProviderId::Zai | ProviderId::MiniMax => "global", _ => "", } @@ -449,6 +468,9 @@ impl Default for Settings { float_bar_show_reset_inline: false, float_bar_show_cost: false, promote_tray_icon: true, + claude_daily_routines_usage_visible: true, + weekly_progress_work_days: None, + alibaba_token_plan_region: default_alibaba_token_plan_region(), } } } @@ -765,6 +787,9 @@ impl Settings { /// API region for `id`, or the provider-specific default if unset. pub fn api_region(&self, id: ProviderId) -> &str { + if id == ProviderId::AlibabaTokenPlan && !self.alibaba_token_plan_region.trim().is_empty() { + return self.alibaba_token_plan_region.as_str(); + } self.provider_configs .get(&id) .and_then(|c| c.api_region.as_deref()) diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 43ca496f5e..dce98ea538 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -155,6 +155,12 @@ pub(super) struct RawSettings { float_bar_show_cost: bool, #[serde(default = "default_true")] promote_tray_icon: bool, + #[serde(default = "default_true")] + claude_daily_routines_usage_visible: bool, + #[serde(default)] + weekly_progress_work_days: Option, + #[serde(default = "default_alibaba_token_plan_region")] + alibaba_token_plan_region: String, } impl Default for RawSettings { @@ -246,6 +252,9 @@ impl Default for RawSettings { float_bar_show_reset_inline: s.float_bar_show_reset_inline, float_bar_show_cost: s.float_bar_show_cost, promote_tray_icon: s.promote_tray_icon, + claude_daily_routines_usage_visible: s.claude_daily_routines_usage_visible, + weekly_progress_work_days: s.weekly_progress_work_days, + alibaba_token_plan_region: s.alibaba_token_plan_region, } } } @@ -383,6 +392,11 @@ impl From for Settings { ProviderId::MiniMax, raw.minimax_api_region, ); + set_region( + &mut provider_configs, + ProviderId::AlibabaTokenPlan, + Some(raw.alibaba_token_plan_region.clone()).filter(|v| !v.trim().is_empty()), + ); set_header( &mut provider_configs, @@ -520,6 +534,16 @@ impl From for Settings { float_bar_show_reset_inline: raw.float_bar_show_reset_inline, float_bar_show_cost: raw.float_bar_show_cost, promote_tray_icon: raw.promote_tray_icon, + claude_daily_routines_usage_visible: raw.claude_daily_routines_usage_visible, + weekly_progress_work_days: raw.weekly_progress_work_days, + alibaba_token_plan_region: { + let trimmed = raw.alibaba_token_plan_region.trim(); + if trimmed.is_empty() { + "cn".to_string() + } else { + trimmed.to_string() + } + }, } } } diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index d364b7a3fb..95587c7f78 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -13,6 +13,7 @@ fn test_settings_default() { assert!(!settings.predictive_pace_warning_enabled); assert!(!settings.float_bar_show_cost); assert!(settings.promote_tray_icon); + assert!(settings.claude_daily_routines_usage_visible); } #[test] From 1fe709058ce03194ef8662267f5ea300d055089c Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:20 +0700 Subject: [PATCH 05/15] Redact config dump secrets by default --- rust/src/cli/config.rs | 172 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 167 insertions(+), 5 deletions(-) diff --git a/rust/src/cli/config.rs b/rust/src/cli/config.rs index 497891163f..2ebabe06bc 100755 --- a/rust/src/cli/config.rs +++ b/rust/src/cli/config.rs @@ -25,6 +25,9 @@ pub enum ConfigCommand { /// Output format: json or toml #[arg(short, long, default_value = "json")] format: String, + /// Include raw secrets (default: redact) + #[arg(long = "show-secrets", default_value_t = false)] + show_secrets: bool, }, /// List providers and enabled state Providers, @@ -60,7 +63,10 @@ pub enum ConfigCommand { pub async fn run(args: ConfigArgs) -> anyhow::Result<()> { match args.command { ConfigCommand::Validate => validate_config().await, - ConfigCommand::Dump { format } => dump_config(&format).await, + ConfigCommand::Dump { + format, + show_secrets, + } => dump_config(&format, show_secrets).await, ConfigCommand::Providers => list_providers().await, ConfigCommand::Enable { provider } => set_provider_enabled(&provider, true).await, ConfigCommand::Disable { provider } => set_provider_enabled(&provider, false).await, @@ -213,16 +219,18 @@ fn print_validation_summary(errors: &[String], warnings: &[String]) -> anyhow::R } /// Dump configuration to stdout -async fn dump_config(format: &str) -> anyhow::Result<()> { - let settings = Settings::load(); +async fn dump_config(format: &str, show_secrets: bool) -> anyhow::Result<()> { + let value = build_dump_value()?; + let value = sanitize_settings_for_dump(value, show_secrets); match format.to_lowercase().as_str() { "json" => { - let json = serde_json::to_string_pretty(&settings)?; + let json = serde_json::to_string_pretty(&value)?; println!("{}", json); } "toml" => { - let toml = toml::to_string_pretty(&settings)?; + let toml = toml::to_string_pretty(&value) + .map_err(|e| anyhow::anyhow!("Failed to convert dump to TOML: {e}"))?; println!("{}", toml); } _ => { @@ -233,6 +241,84 @@ async fn dump_config(format: &str) -> anyhow::Result<()> { Ok(()) } +fn build_dump_value() -> anyhow::Result { + let settings = Settings::load(); + let mut root = serde_json::to_value(&settings)?; + + if let Some(obj) = root.as_object_mut() { + obj.insert( + "api_keys".to_string(), + serde_json::to_value(ApiKeys::load())?, + ); + obj.insert( + "manual_cookies".to_string(), + serde_json::to_value(ManualCookies::load())?, + ); + + let token_accounts = TokenAccountStore::new() + .load() + .unwrap_or_default() + .into_iter() + .map(|(id, data)| (id.cli_name().to_string(), data)) + .collect::>(); + obj.insert( + "token_accounts".to_string(), + serde_json::to_value(token_accounts)?, + ); + } + + Ok(root) +} + +/// Recursively redact secret-shaped fields for `config dump`. +/// +/// When `show_secrets` is true the value is returned unchanged. +fn sanitize_settings_for_dump(value: serde_json::Value, show_secrets: bool) -> serde_json::Value { + if show_secrets { + return value; + } + redact_secrets_value(value) +} + +fn is_secret_field_name(key: &str) -> bool { + let normalized: String = key + .chars() + .filter(|c| *c != '_') + .flat_map(|c| c.to_lowercase()) + .collect(); + matches!( + normalized.as_str(), + "apikey" + | "secretkey" + | "cookieheader" + | "manualcookieheader" + | "token" + | "apitoken" + | "httpproxypassword" + | "password" + ) +} + +fn redact_secrets_value(value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(map) => { + let mut out = serde_json::Map::new(); + for (key, child) in map { + if is_secret_field_name(&key) { + out.insert(key, serde_json::Value::String("[REDACTED]".to_string())); + } else { + out.insert(key, redact_secrets_value(child)); + } + } + serde_json::Value::Object(out) + } + serde_json::Value::Array(items) => { + serde_json::Value::Array(items.into_iter().map(redact_secrets_value).collect()) + } + other => other, + } +} + /// List provider enabled state. async fn list_providers() -> anyhow::Result<()> { let settings = Settings::load(); @@ -386,3 +472,79 @@ async fn show_paths() -> anyhow::Result<()> { Ok(()) } + +#[cfg(test)] +mod tests { + use super::sanitize_settings_for_dump; + use serde_json::json; + + #[test] + fn sanitize_settings_for_dump_redacts_secret_fields() { + let raw = json!({ + "http_proxy_password": "proxy-secret", + "provider_configs": { + "claude": { + "manual_cookie_header": "sessionKey=abc", + "api_token": "tok-123" + } + }, + "api_keys": { + "keys": { + "zai": { + "api_key": "cb_test_api_key_456", + "label": "Work", + "saved_at": "2026-01-01" + } + } + }, + "manual_cookies": { + "cookies": { + "claude": { + "cookie_header": "sessionKey=secret-cookie", + "saved_at": "2026-01-01" + } + } + }, + "token_accounts": { + "factory": { + "accounts": [{ + "id": "11111111-1111-1111-1111-111111111111", + "label": "Team", + "token": "raw-token-value", + "added_at": 1 + }] + } + }, + "secret_key": "top-secret" + }); + + let redacted = sanitize_settings_for_dump(raw, false); + let text = serde_json::to_string(&redacted).expect("serialize"); + + assert!(text.contains("[REDACTED]")); + assert!(!text.contains("proxy-secret")); + assert!(!text.contains("sessionKey=abc")); + assert!(!text.contains("tok-123")); + assert!(!text.contains("cb_test_api_key_456")); + assert!(!text.contains("secret-cookie")); + assert!(!text.contains("raw-token-value")); + assert!(!text.contains("top-secret")); + // Non-secret identity fields stay. + assert!(text.contains("Work")); + assert!(text.contains("Team")); + assert!(text.contains("11111111-1111-1111-1111-111111111111")); + } + + #[test] + fn sanitize_settings_for_dump_show_secrets_keeps_raw() { + let raw = json!({ + "api_key": "keep-me", + "nested": { "token": "also-keep", "label": "Team" } + }); + + let out = sanitize_settings_for_dump(raw.clone(), true); + assert_eq!(out, raw); + assert_eq!(out["api_key"], "keep-me"); + assert_eq!(out["nested"]["token"], "also-keep"); + } +} From 09f17884216d7df24b316d0582a5e3bd5c5cfdf6 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:22 +0700 Subject: [PATCH 06/15] Wire Codex cost scan disk cache and resume --- rust/src/cli/cost.rs | 4 +- rust/src/cli/serve.rs | 4 +- rust/src/codex_costs.rs | 60 +++++ rust/src/core/jsonl_scanner.rs | 69 +++++- rust/src/cost_scanner.rs | 439 +++++++++++++++++++++++++++++---- 5 files changed, 512 insertions(+), 64 deletions(-) diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index e5103b2bfc..d85da4eebf 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -5,7 +5,7 @@ use clap::Args; use super::usage::{OutputFormat, ProviderSelection}; -use crate::core::ProviderId; +use crate::core::{CostScanOptions, ProviderId}; use crate::cost_scanner::{CostScanner, CostSummary}; /// Arguments for the cost command @@ -46,7 +46,7 @@ pub async fn run(args: CostArgs) -> anyhow::Result<()> { let providers = ProviderSelection::from_arg(args.provider.as_deref())?; let use_color = !args.no_color && is_terminal(); - let scanner = CostScanner::new(args.days); + let scanner = CostScanner::new(args.days).with_options(CostScanOptions::app_driven()); tracing::debug!( "Running cost command: providers={:?}, format={:?}, days={}", diff --git a/rust/src/cli/serve.rs b/rust/src/cli/serve.rs index 482da729fc..6a0fd6a07c 100644 --- a/rust/src/cli/serve.rs +++ b/rust/src/cli/serve.rs @@ -9,7 +9,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use super::usage::ProviderSelection; -use crate::core::{FetchContext, ProviderId, SourceMode, instantiate_provider}; +use crate::core::{CostScanOptions, FetchContext, ProviderId, SourceMode, instantiate_provider}; use crate::cost_scanner::CostScanner; const DASHBOARD_TOKEN_ENV: &str = "CODEXBAR_DASHBOARD_TOKEN"; @@ -305,7 +305,7 @@ async fn cost_response(provider: Option<&str>) -> String { return json_response(400, serde_json::json!({ "error": error.to_string() })); } }; - let scanner = CostScanner::new(30); + let scanner = CostScanner::new(30).with_options(CostScanOptions::app_driven()); let mut results = Vec::new(); for provider_id in selection.as_list() { let (supported, summary) = match provider_id { diff --git a/rust/src/codex_costs.rs b/rust/src/codex_costs.rs index aa8f697b9c..3141dbc9a6 100644 --- a/rust/src/codex_costs.rs +++ b/rust/src/codex_costs.rs @@ -48,6 +48,65 @@ pub(crate) fn add_codex_records_to_summary( (total_cost, has_tokens) } +/// Merge billable records into a day→model→`[input,cached,output]` map. +pub(crate) fn merge_codex_records_into_days( + days: &mut std::collections::HashMap>>, + records: &[CodexUsageRecord], +) { + for record in records { + let models = days.entry(record.day_key.clone()).or_default(); + let packed = models + .entry(record.model.clone()) + .or_insert_with(|| vec![0, 0, 0]); + if packed.len() < 3 { + packed.resize(3, 0); + } + packed[0] = packed[0].saturating_add(record.input.max(0)); + packed[1] = packed[1].saturating_add(record.cached.max(0)); + packed[2] = packed[2].saturating_add(record.output.max(0)); + } +} + +/// Apply one packed `[input, cached, output]` triple to a summary. +pub(crate) fn add_codex_packed_tokens_to_summary( + summary: &mut CostSummary, + model: &str, + packed: &[i32], +) -> Option { + let input = packed.first().copied().unwrap_or(0); + let cached = packed.get(1).copied().unwrap_or(0); + let output = packed.get(2).copied().unwrap_or(0); + add_codex_tokens_to_summary( + summary, + model, + CodexTokenCounts::from_values(input, cached, output), + ) +} + +/// Fold day→model→packed token maps into a cost summary (range-filtered). +/// Returns `(session_cost, has_tokens)` — caller adds cost to `total_cost_usd`. +pub(crate) fn add_codex_days_map_to_summary( + summary: &mut CostSummary, + days: &std::collections::HashMap>>, + range: &CostUsageDayRange, +) -> (f64, bool) { + let mut total_cost = 0.0; + let mut has_tokens = false; + for (day_key, models) in days { + if !CostUsageDayRange::is_in_range(day_key, &range.since_key, &range.until_key) { + continue; + } + for (model, packed) in models { + if let Some(cost) = add_codex_packed_tokens_to_summary(summary, model, packed) { + total_cost += cost; + has_tokens = true; + } + } + } + (total_cost, has_tokens) +} + +#[cfg_attr(not(test), allow(dead_code))] pub(crate) fn scan_codex_file_cost_for_range(path: &Path, range: &CostUsageDayRange) -> f64 { let parse_result = match JsonlScanner::parse_codex_file(path, range, 0, None, None) { Ok(result) => result, @@ -153,6 +212,7 @@ fn add_codex_tokens_to_summary( Some(cost) } +#[cfg_attr(not(test), allow(dead_code))] fn codex_records_cost(records: &[CodexUsageRecord], range: &CostUsageDayRange) -> f64 { let mut total_cost = 0.0; diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index b5c59a5db7..e8d3b16a79 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -22,11 +22,10 @@ pub const DEFAULT_COST_SCAN_REFRESH_MIN_INTERVAL_SECS: u64 = 60; /// Options for a cost scan pass (disk-cache-backed full inspections). /// -/// **Windows note:** the production [`crate::cost_scanner::CostScanner`] path -/// does not load/save [`CostUsageCache`] today — it always full-walks sessions. -/// These options are ready for the cache path (and unit-tested) so app-driven -/// callers can pass [`CostScanOptions::app_driven`] once that path is wired -/// (upstream issue #2089). Until then they do not change runtime freshness. +/// Default debounce is 60s between full disk inspections when a +/// [`CostUsageCache`] is present. Pass [`CostScanOptions::app_driven`] (interval 0) +/// for explicit/CLI refreshes. Production [`crate::cost_scanner::CostScanner`] +/// honors these options and persists cache under `{cache}/CodexBar/cost-usage/`. #[derive(Debug, Clone, Copy)] pub struct CostScanOptions { /// Minimum seconds between disk-cache-backed full inspections. @@ -68,6 +67,11 @@ pub struct CostUsageCache { pub files: HashMap, /// Aggregated daily data: day_key -> model -> [input, cached, output] pub days: HashMap>>, + /// Inclusive range covered by the last successful full inspection. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scan_since_key: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scan_until_key: Option, } /// Per-file usage tracking @@ -844,26 +848,65 @@ impl JsonlScanner { CostUsageCache::default() } - /// Save cache to disk + /// Save cache to disk (temp sibling + copy into place). pub fn save_cache(provider: ProviderId, cache: &CostUsageCache, cache_root: Option<&Path>) { let cache_path = Self::cache_path(provider, cache_root); - if let Some(parent) = cache_path.parent() { - let _ = fs::create_dir_all(parent); - } + let Some(parent) = cache_path.parent() else { + return; + }; + let _ = fs::create_dir_all(parent); + + let Ok(json) = serde_json::to_string(cache) else { + return; + }; - if let Ok(json) = serde_json::to_string_pretty(cache) { - let _ = fs::write(&cache_path, json); + let tmp_name = format!( + ".{}.{}-{}.tmp", + provider.cli_name(), + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0) + ); + let tmp_path = parent.join(tmp_name); + if fs::write(&tmp_path, json.as_bytes()).is_err() { + return; } + // `copy` replaces an existing target on Windows; prefer it over rename. + if fs::copy(&tmp_path, &cache_path).is_err() { + let _ = fs::write(&cache_path, json.as_bytes()); + } + // Best-effort temp cleanup (ignore errors — unique name avoids clashes). + let _ = fs::File::create(&tmp_path).and_then(|f| f.set_len(0)); + } + + /// Default on-disk cache root: `%LOCALAPPDATA%\CodexBar` (via `dirs::cache_dir`). + pub fn default_cache_root() -> Option { + dirs::cache_dir().map(|d| d.join("CodexBar")) } fn cache_path(provider: ProviderId, cache_root: Option<&Path>) -> PathBuf { let root = cache_root .map(|p| p.to_path_buf()) - .or_else(|| dirs::cache_dir().map(|d| d.join("CodexBar"))) + .or_else(Self::default_cache_root) .unwrap_or_else(|| PathBuf::from(".")); - root.join(format!("{}_cost_cache.json", provider.cli_name())) + // Mirror upstream layout: {cacheRoot}/cost-usage/{provider}-v1.json + root.join("cost-usage") + .join(format!("{}-v1.json", provider.cli_name())) + } + + /// Whether `cache` covers the requested day window (for debounce short-circuit). + pub fn cache_covers_range(cache: &CostUsageCache, range: &CostUsageDayRange) -> bool { + match (&cache.scan_since_key, &cache.scan_until_key) { + (Some(since), Some(until)) => { + since.as_str() <= range.since_key.as_str() + && until.as_str() >= range.until_key.as_str() + } + _ => !cache.days.is_empty() || !cache.files.is_empty(), + } } } diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 2fb203b616..6cac68319e 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -2,9 +2,11 @@ //! //! Scans local JSONL log files to aggregate token usage and calculate costs. //! -//! Note: this path always full-walks session files. The disk-cache / -//! [`crate::core::CostScanOptions`] debounce API in `jsonl_scanner` is not -//! wired here yet (upstream #2089); app-level TTL still owns refresh pacing. +//! Codex production path loads/saves [`crate::core::CostUsageCache`] under +//! `{cache}/CodexBar/cost-usage/`, skips unchanged files by mtime+size, resumes +//! partial files from `parsed_bytes`, honors [`crate::core::CostScanOptions`] +//! debounce (default 60s; `app_driven` forces a fresh inspection), and checks +//! cancel flags between files. use chrono::{DateTime, Duration, Local, NaiveDate, Utc}; use serde::Deserialize; @@ -13,15 +15,19 @@ use std::fs::{self, File}; use std::io::{BufRead, BufReader}; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; #[cfg(test)] use crate::codex_costs::scan_codex_file_cost; use crate::codex_costs::{ - add_codex_records_to_summary, codex_period_start, codex_scan_dates, - scan_codex_file_cost_for_range, + add_codex_days_map_to_summary, add_codex_records_to_summary, codex_period_start, + codex_scan_dates, merge_codex_records_into_days, }; use crate::codex_sessions::{codex_sessions_dir_candidates, default_wsl_roots}; -use crate::core::{CostUsageDayRange, CostUsagePricing, JsonlScanner}; +use crate::core::{ + CostScanOptions, CostUsageCache, CostUsageDayRange, CostUsageFileUsage, CostUsagePricing, + JsonlScanner, ProviderId, +}; use crate::settings::Settings; /// Cost summary from scanning local logs @@ -81,6 +87,40 @@ fn is_cancelled(cancel: Option<&AtomicBool>) -> bool { /// pricing table (unknown or retired IDs). Prices as Sonnet 4.6. const FALLBACK_CLAUDE_MODEL: &str = "claude-sonnet-4-6"; +fn unix_now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis().min(i64::MAX as u128) as i64) + .unwrap_or(0) +} + +fn system_time_to_unix_ms(modified: Option) -> i64 { + modified + .and_then(|t| t.duration_since(UNIX_EPOCH).ok()) + .map(|d| d.as_millis().min(i64::MAX as u128) as i64) + .unwrap_or(0) +} + +fn rebuild_cache_days(cache: &mut CostUsageCache) { + cache.days.clear(); + for usage in cache.files.values() { + for (day, models) in &usage.days { + let day_entry = cache.days.entry(day.clone()).or_default(); + for (model, packed) in models { + let dest = day_entry + .entry(model.clone()) + .or_insert_with(|| vec![0, 0, 0]); + if dest.len() < 3 { + dest.resize(3, 0); + } + for (i, value) in packed.iter().take(3).enumerate() { + dest[i] = dest[i].saturating_add(*value); + } + } + } + } +} + /// Claude cost calculation for the usage scanner. /// /// Per-token rates come from the canonical `CostUsagePricing::claude_cost_usd` @@ -218,15 +258,52 @@ struct ClaudeUsageRecord { cost: f64, } +/// Per-pass counters for cache/resume behavior (tests + diagnostics). +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CostScanStats { + pub files_seen: u32, + pub files_parsed: u32, + pub files_skipped: u32, + pub files_resumed: u32, + pub used_cache_debounce: bool, +} + /// Cost usage scanner pub struct CostScanner { days: u32, + options: CostScanOptions, + cache_root: Option, + /// When set, bypass normal sessions-dir discovery (tests / inject roots). + sessions_dirs_override: Option>, } impl CostScanner { - /// Create a new scanner for the last N days + /// Create a new scanner for the last N days (default 60s cache debounce). pub fn new(days: u32) -> Self { - Self { days } + Self { + days, + options: CostScanOptions::default(), + cache_root: None, + sessions_dirs_override: None, + } + } + + /// Override scan options (e.g. [`CostScanOptions::app_driven`] for force refresh). + pub fn with_options(mut self, options: CostScanOptions) -> Self { + self.options = options; + self + } + + /// Override on-disk cache root (`{root}/cost-usage/…`). + pub fn with_cache_root(mut self, root: impl Into) -> Self { + self.cache_root = Some(root.into()); + self + } + + /// Override Codex sessions roots (primarily for tests). + pub fn with_sessions_dirs(mut self, dirs: Vec) -> Self { + self.sessions_dirs_override = Some(dirs); + self } /// Scan Codex local logs @@ -236,23 +313,78 @@ impl CostScanner { /// Scan Codex local logs, stopping early when the caller cancels the scan. pub fn scan_codex_with_cancel(&self, cancel: Option<&AtomicBool>) -> CostSummary { + self.scan_codex_detailed(cancel).0 + } + + /// Scan Codex and return cache/resume stats alongside the summary. + pub fn scan_codex_detailed(&self, cancel: Option<&AtomicBool>) -> (CostSummary, CostScanStats) { let mut summary = CostSummary::default(); + let mut stats = CostScanStats::default(); let today = Local::now().date_naive(); let start_date = codex_period_start(today, self.days); let range = CostUsageDayRange::new(start_date, today); + let now_ms = unix_now_ms(); summary.period_start = Some(start_date); summary.period_end = Some(today); + let cache_root = self.cache_root.as_deref(); + let mut cache = JsonlScanner::load_cache(ProviderId::Codex, cache_root); + + // Debounce: rebuild from disk cache without re-walking session files. + if JsonlScanner::should_skip_cached_scan(&cache, self.options, now_ms) + && JsonlScanner::cache_covers_range(&cache, &range) + && (!cache.days.is_empty() || !cache.files.is_empty()) + { + stats.used_cache_debounce = true; + let (cost, _) = add_codex_days_map_to_summary(&mut summary, &cache.days, &range); + summary.total_cost_usd += cost; + summary.sessions_count = cache + .files + .values() + .filter(|usage| { + usage.days.keys().any(|day| { + CostUsageDayRange::is_in_range(day, &range.since_key, &range.until_key) + }) + }) + .count() as u32; + + // Pi-compatible sessions are outside the Codex JSONL cache. + let mut seen_pi = HashSet::new(); + crate::pi_session_cost::scan_pi_compatible_into( + &mut summary, + crate::pi_session_cost::PiMappedProvider::Codex, + self.days, + cancel, + &mut seen_pi, + ); + return (summary, stats); + } + for sessions_dir in self.get_codex_sessions_dirs() { if is_cancelled(cancel) { break; } if sessions_dir.exists() { - self.scan_codex_sessions_dir(&sessions_dir, &range, &mut summary, cancel); + self.scan_codex_sessions_dir( + &sessions_dir, + &range, + &mut summary, + &mut cache, + cancel, + &mut stats, + ); } } + if !is_cancelled(cancel) { + rebuild_cache_days(&mut cache); + cache.last_scan_unix_ms = now_ms; + cache.scan_since_key = Some(range.since_key.clone()); + cache.scan_until_key = Some(range.until_key.clone()); + JsonlScanner::save_cache(ProviderId::Codex, &cache, cache_root); + } + // OMP / pi-compatible agent sessions (upstream #2269). Dedup by entry id. let mut seen_pi = HashSet::new(); crate::pi_session_cost::scan_pi_compatible_into( @@ -263,7 +395,7 @@ impl CostScanner { &mut seen_pi, ); - summary + (summary, stats) } /// Scan Claude local logs @@ -312,6 +444,9 @@ impl CostScanner { } fn get_codex_sessions_dirs(&self) -> Vec { + if let Some(dirs) = &self.sessions_dirs_override { + return dirs.clone(); + } let settings = Settings::load(); let codex_home = std::env::var("CODEX_HOME").ok(); codex_sessions_dir_candidates( @@ -327,7 +462,9 @@ impl CostScanner { sessions_dir: &Path, range: &CostUsageDayRange, summary: &mut CostSummary, + cache: &mut CostUsageCache, cancel: Option<&AtomicBool>, + stats: &mut CostScanStats, ) { // Iterate through the date-based directory structure with one day of // padding on each side. Codex JSONL timestamps are UTC, while the tray @@ -352,7 +489,7 @@ impl CostScanner { } let path = entry.path(); if path.extension().is_some_and(|e| e == "jsonl") { - self.parse_codex_file(&path, summary, cancel); + self.parse_codex_file(&path, range, summary, cache, cancel, stats); } } } @@ -381,27 +518,119 @@ impl CostScanner { fn parse_codex_file( &self, path: &Path, + range: &CostUsageDayRange, summary: &mut CostSummary, + cache: &mut CostUsageCache, cancel: Option<&AtomicBool>, + stats: &mut CostScanStats, ) { if is_cancelled(cancel) { return; } - let today = Local::now().date_naive(); - let start_date = codex_period_start(today, self.days); - let range = CostUsageDayRange::new(start_date, today); - let parse_result = match JsonlScanner::parse_codex_file(path, &range, 0, None, None) { + stats.files_seen += 1; + + let metadata = match fs::metadata(path) { + Ok(m) => m, + Err(_) => return, + }; + let size = metadata.len() as i64; + let mtime_ms = system_time_to_unix_ms(metadata.modified().ok()); + let path_key = path.to_string_lossy().to_string(); + let cached = cache.files.get(&path_key).cloned(); + + // Unchanged complete file: reuse packed days, skip re-parse. + if let Some(entry) = &cached + && entry.mtime_unix_ms == mtime_ms + && entry.size == size + && entry.parsed_bytes.unwrap_or(0) >= size + && size > 0 + { + let (session_cost, has_tokens) = + add_codex_days_map_to_summary(summary, &entry.days, range); + if has_tokens { + summary.total_cost_usd += session_cost; + summary.sessions_count += 1; + } + stats.files_skipped += 1; + return; + } + + // Growing file: resume from last parsed offset when safe. + if let Some(entry) = &cached { + let start_offset = entry.parsed_bytes.unwrap_or(0); + if size > entry.size + && start_offset > 0 + && start_offset <= size + && entry.last_totals.is_some() + { + let parse_result = match JsonlScanner::parse_codex_file( + path, + range, + start_offset, + entry.last_model.clone(), + entry.last_totals.clone(), + ) { + Ok(result) => result, + Err(_) => return, + }; + + let mut days = entry.days.clone(); + merge_codex_records_into_days(&mut days, &parse_result.records); + + let (session_cost, has_tokens) = + add_codex_days_map_to_summary(summary, &days, range); + if has_tokens { + summary.total_cost_usd += session_cost; + summary.sessions_count += 1; + } + + cache.files.insert( + path_key, + CostUsageFileUsage { + mtime_unix_ms: mtime_ms, + size, + days, + parsed_bytes: Some(parse_result.parsed_bytes), + last_model: parse_result.last_model.or_else(|| entry.last_model.clone()), + last_totals: parse_result + .last_totals + .or_else(|| entry.last_totals.clone()), + }, + ); + stats.files_resumed += 1; + return; + } + } + + // Full parse from offset 0. + let parse_result = match JsonlScanner::parse_codex_file(path, range, 0, None, None) { Ok(result) => result, Err(_) => return, }; + let mut days = HashMap::new(); + merge_codex_records_into_days(&mut days, &parse_result.records); + let (session_cost, has_tokens) = - add_codex_records_to_summary(summary, &parse_result.records, &range); + add_codex_records_to_summary(summary, &parse_result.records, range); if has_tokens { summary.total_cost_usd += session_cost; summary.sessions_count += 1; } + + cache.files.insert( + path_key, + CostUsageFileUsage { + mtime_unix_ms: mtime_ms, + size, + days, + parsed_bytes: Some(parse_result.parsed_bytes), + last_model: parse_result.last_model, + last_totals: parse_result.last_totals, + }, + ); + stats.files_parsed += 1; } fn walk_claude_files( @@ -649,34 +878,22 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, f64)> { match provider { "codex" => { - // Scan Codex logs by day across Windows and WSL session roots. - let sessions_dirs = scanner.get_codex_sessions_dirs(); - for days_ago in 0..days { - let date = today - Duration::days(days_ago as i64); - let date_str = date.format("%Y-%m-%d").to_string(); - let range = CostUsageDayRange::new(date, date); - let mut day_cost = 0.0; - - for sessions_dir in sessions_dirs.iter().filter(|dir| dir.exists()) { - for scan_date in codex_scan_dates(&range) { - let year = scan_date.format("%Y").to_string(); - let month = scan_date.format("%m").to_string(); - let day = scan_date.format("%d").to_string(); - let day_dir = sessions_dir.join(&year).join(&month).join(&day); - if !day_dir.exists() { - continue; - } - if let Ok(entries) = fs::read_dir(&day_dir) { - for entry in entries.flatten() { - let path = entry.path(); - if path.extension().is_some_and(|e| e == "jsonl") { - day_cost += scan_codex_file_cost_for_range(&path, &range); - } - } - } - } - } - daily_costs.insert(date_str, day_cost); + // Warm/refresh the disk cache (honors debounce), then price from packed days. + let _ = scanner.scan_codex(); + let cache = JsonlScanner::load_cache(ProviderId::Codex, scanner.cache_root.as_deref()); + for (day_key, models) in &cache.days { + let Some(slot) = daily_costs.get_mut(day_key) else { + continue; + }; + let Some(day) = CostUsageDayRange::parse_day_key(day_key) else { + continue; + }; + let day_range = CostUsageDayRange::new(day, day); + let mut one_day = HashMap::new(); + one_day.insert(day_key.clone(), models.clone()); + let mut scratch = CostSummary::default(); + let (cost, _) = add_codex_days_map_to_summary(&mut scratch, &one_day, &day_range); + *slot = cost; } } "claude" => { @@ -835,11 +1052,13 @@ mod tests { ts = recent ) .unwrap(); - drop(file); - let scanner = CostScanner::new(30); let mut summary = CostSummary::default(); - scanner.parse_codex_file(&path, &mut summary, None); + let today = Local::now().date_naive(); + let range = CostUsageDayRange::new(codex_period_start(today, 30), today); + let mut cache = CostUsageCache::default(); + let mut stats = CostScanStats::default(); + scanner.parse_codex_file(&path, &range, &mut summary, &mut cache, None, &mut stats); assert_eq!(summary.sessions_count, 1); assert_eq!(summary.input_tokens, 125); @@ -1027,4 +1246,130 @@ mod tests { assert_eq!(counted, 1, "incomplete final JSONL line must be processed"); let _ = std::fs::remove_file(&path); } + + fn write_codex_session_fixture(sessions_root: &Path, name: &str, input_tokens: u64) -> PathBuf { + let today = Local::now().date_naive(); + let day_dir = sessions_root + .join(today.format("%Y").to_string()) + .join(today.format("%m").to_string()) + .join(today.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let ts = (Utc::now() - Duration::hours(1)) + .format("%Y-%m-%dT%H:%M:%S%.3fZ") + .to_string(); + let body = format!( + r#"{{"timestamp":"{ts}","type":"event_msg","payload":{{"type":"token_count","info":{{"model":"gpt-5","total_token_usage":{{"input_tokens":{input_tokens},"cached_input_tokens":0,"output_tokens":5}}}}}}}} +"# + ); + std::fs::write(&path, body).unwrap(); + path + } + + #[test] + fn cost_scan_second_pass_skips_unchanged_files_via_cache() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + write_codex_session_fixture(&sessions, "a.jsonl", 100); + write_codex_session_fixture(&sessions, "b.jsonl", 200); + + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + + let (summary1, stats1) = scanner.scan_codex_detailed(None); + assert_eq!(stats1.files_parsed, 2, "first pass parses both files"); + assert_eq!(stats1.files_skipped, 0); + assert!(summary1.total_cost_usd > 0.0); + assert_eq!(summary1.sessions_count, 2); + + // Second pass with default debounce still inspects files but skips re-parse. + // Use app_driven so we exercise per-file mtime skip rather than whole-scan debounce. + let (summary2, stats2) = scanner.scan_codex_detailed(None); + assert_eq!(stats2.files_seen, 2); + assert_eq!(stats2.files_skipped, 2, "cache hit skips re-parse"); + assert_eq!(stats2.files_parsed, 0); + assert_eq!(summary2.input_tokens, summary1.input_tokens); + assert!((summary2.total_cost_usd - summary1.total_cost_usd).abs() < 1e-9); + + // Force path already used above; confirm debounce short-circuit with default options. + let debounced = CostScanner::new(7) + .with_options(CostScanOptions::default()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (summary3, stats3) = debounced.scan_codex_detailed(None); + assert!( + stats3.used_cache_debounce, + "default options debounce within 60s" + ); + assert_eq!(stats3.files_seen, 0); + assert_eq!(summary3.input_tokens, summary1.input_tokens); + + // app_driven after debounce still re-reads (skip via mtime, not full re-parse). + let forced = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + let (_, stats4) = forced.scan_codex_detailed(None); + assert!(!stats4.used_cache_debounce); + assert_eq!(stats4.files_skipped, 2); + assert_eq!(stats4.files_parsed, 0); + } + + #[test] + fn cost_scan_cancel_stops_between_files() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + write_codex_session_fixture(&sessions, "a.jsonl", 100); + write_codex_session_fixture(&sessions, "b.jsonl", 200); + write_codex_session_fixture(&sessions, "c.jsonl", 300); + + let cancel = AtomicBool::new(true); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(cache_root) + .with_sessions_dirs(vec![sessions]); + let (summary, stats) = scanner.scan_codex_detailed(Some(&cancel)); + assert_eq!(stats.files_seen, 0, "cancel before first file stops walk"); + assert_eq!(summary.sessions_count, 0); + } + + #[test] + fn cost_scan_resumes_appended_bytes() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let path = write_codex_session_fixture(&sessions, "grow.jsonl", 50); + + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (s1, st1) = scanner.scan_codex_detailed(None); + assert_eq!(st1.files_parsed, 1); + assert_eq!(s1.input_tokens, 50); + + // Append another cumulative token_count event (100 total => +50 delta). + let ts = Utc::now().format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); + let extra = format!( + r#"{{"timestamp":"{ts}","type":"event_msg","payload":{{"type":"token_count","info":{{"model":"gpt-5","total_token_usage":{{"input_tokens":100,"cached_input_tokens":0,"output_tokens":10}}}}}}}} +"# + ); + use std::io::Write as _; + let mut f = std::fs::OpenOptions::new() + .append(true) + .open(&path) + .unwrap(); + f.write_all(extra.as_bytes()).unwrap(); + drop(f); + + // Bump mtime/size visibly on some FS by rewriting metadata via reopen. + let (s2, st2) = scanner.scan_codex_detailed(None); + assert_eq!(st2.files_resumed, 1, "grown file resumes from offset"); + assert_eq!(st2.files_parsed, 0); + assert_eq!(s2.input_tokens, 100); + } } From 837800f6a92eced4cf7024fafcb11396a2bd342d Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:26 +0700 Subject: [PATCH 07/15] Add fractional session quota forecast --- .../desktop-tauri/src/components/MenuCard.tsx | 29 + rust/src/core/mod.rs | 4 + rust/src/core/session_equivalent_forecast.rs | 1137 +++++++++++++++++ 3 files changed, 1170 insertions(+) create mode 100644 rust/src/core/session_equivalent_forecast.rs diff --git a/apps/desktop-tauri/src/components/MenuCard.tsx b/apps/desktop-tauri/src/components/MenuCard.tsx index 875849b01a..6e950937eb 100644 --- a/apps/desktop-tauri/src/components/MenuCard.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.tsx @@ -6,6 +6,7 @@ import type { ProviderLocalUsageSummary, ProviderUsageSnapshot, RateWindowSnapshot, + SessionEquivalentForecastSnapshot, } from "../types/bridge"; import { getProviderChartData } from "../lib/tauri"; import { useLocale } from "../hooks/useLocale"; @@ -97,6 +98,23 @@ function formatReserveDescription( return t("PanelReserveRunsOutInHours").replace("{}", String(h)); } +/** Upstream session-quota estimate: "Estimated: {n} session quota(s) left". */ +function formatSessionEquivalentEstimate( + forecast: SessionEquivalentForecastSnapshot | null | undefined, +): string | null { + if (!forecast) return null; + const raw = forecast.estimatedWindowsToExhaustWeekly; + if (!Number.isFinite(raw)) return null; + const rounded = Math.round(Math.min(Math.max(raw, 0), 1_000_000) * 10) / 10; + const display = + Number.isInteger(rounded) || Math.abs(rounded - Math.round(rounded)) < 1e-9 + ? String(Math.round(rounded)) + : rounded.toFixed(1); + const unit = + rounded > 0 && rounded <= 1 ? "session quota" : "session quotas"; + return `Estimated: ${display} ${unit} left`; +} + function formatCurrency(amount: number, code: string): string { try { return new Intl.NumberFormat("en-US", { @@ -278,6 +296,7 @@ interface MetricEntry { label: string; snap: RateWindowSnapshot; resetFormatMode?: ResetTimeFormatMode; + sessionEquivalentForecast?: SessionEquivalentForecastSnapshot | null; } type MetricPaceView = @@ -316,6 +335,7 @@ function MetricRow({ expanded, onToggleExpanded, resetFormatMode, + sessionEquivalentForecast, }: { title: string; snap: RateWindowSnapshot; @@ -326,6 +346,7 @@ function MetricRow({ expanded: boolean; onToggleExpanded: () => void; resetFormatMode?: ResetTimeFormatMode; + sessionEquivalentForecast?: SessionEquivalentForecastSnapshot | null; }) { const { t } = useLocale(); const isInformational = snap.isInformational === true; @@ -352,6 +373,7 @@ function MetricRow({ resetText !== null; const paceView = getMetricPaceView(snap); const reserveDescription = formatReserveDescription(snap, t); + const forecastText = formatSessionEquivalentEstimate(sessionEquivalentForecast); const formatBudget = (value: number) => value < 10 ? value.toFixed(1).replace(/\.0$/, "") : Math.round(value).toString(); return ( @@ -417,6 +439,11 @@ function MetricRow({ )} )} + {!isInformational && forecastText && ( +
+ {forecastText} +
+ )} ); } @@ -506,6 +533,7 @@ export default function MenuCard({ id: "secondary", label: localizeWindowLabel(provider.secondaryLabel, t) || t("DetailWindowSecondary"), snap: provider.secondary, + sessionEquivalentForecast: provider.sessionEquivalentForecast, }); if (provider.modelSpecific) metrics.push({ @@ -599,6 +627,7 @@ export default function MenuCard({ showAsUsed={showAsUsed} expanded={expandedPaceWindow === m.id} resetFormatMode={m.resetFormatMode} + sessionEquivalentForecast={m.sessionEquivalentForecast} onToggleExpanded={() => { setExpandedPaceWindow((current) => current === m.id ? null : m.id, diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs index d3e9595d4f..d62419dbcb 100755 --- a/rust/src/core/mod.rs +++ b/rust/src/core/mod.rs @@ -3,6 +3,7 @@ mod adaptive_refresh; mod aws_signing; mod cost_pricing; +pub mod curl_capture; mod hooks; mod http; mod http_proxy; @@ -13,6 +14,7 @@ mod provider; mod provider_factory; mod rate_window; mod redactor; +mod session_equivalent_forecast; mod session_quota; mod token_accounts; mod usage_pace; @@ -22,6 +24,7 @@ mod widget_snapshot; pub use adaptive_refresh::*; pub use aws_signing::*; pub use cost_pricing::*; +pub use curl_capture::*; pub use hooks::*; pub use http::*; pub use http_proxy::*; @@ -32,6 +35,7 @@ pub use provider::*; pub use provider_factory::instantiate as instantiate_provider; pub use rate_window::*; pub use redactor::*; +pub use session_equivalent_forecast::*; pub use session_quota::*; pub use token_accounts::*; pub use usage_pace::*; diff --git a/rust/src/core/session_equivalent_forecast.rs b/rust/src/core/session_equivalent_forecast.rs new file mode 100644 index 0000000000..2ecf62ca76 --- /dev/null +++ b/rust/src/core/session_equivalent_forecast.rs @@ -0,0 +1,1137 @@ +//! Session-equivalent weekly forecast. +//! +//! Estimates how many full 5-hour session quotas remain in the weekly window +//! from recent plan-utilization burn history (upstream SessionEquivalentForecast). + +use chrono::{DateTime, Datelike, Duration, Utc, Weekday}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::{LazyLock, Mutex}; + +/// Canonical session window length (5 hours). +pub const SESSION_WINDOW_MINUTES: u32 = 300; +/// Canonical weekly window length (7 days). +pub const WEEKLY_WINDOW_MINUTES: u32 = 10_080; +/// Reset-boundary grouping tolerance. +pub const RESET_TOLERANCE_SECS: i64 = 120; + +/// Median weekly-percent burn observed per completed session window. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct SessionEquivalentBurnEstimate { + pub median_weekly_percent_per_window: f64, + pub sample_count: usize, +} + +/// Forecast of remaining session-equivalent windows inside the weekly quota. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct SessionEquivalentForecast { + pub estimated_windows_to_exhaust_weekly: f64, + pub windows_until_reset: i64, + pub available_windows_until_reset: f64, + pub sample_count: usize, + pub weekly_resets_at: DateTime, + pub weekly_used_percent: f64, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub weekly_window_id: Option, +} + +/// One captured plan-utilization observation. +#[derive(Debug, Clone, PartialEq)] +pub struct PlanUtilizationHistoryEntry { + pub captured_at: DateTime, + pub used_percent: f64, + pub resets_at: Option>, +} + +/// Named plan-utilization series used by the burn estimator. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum PlanUtilizationSeriesName { + Session, + Weekly, +} + +impl PlanUtilizationSeriesName { + pub fn canonical_window_minutes(self) -> u32 { + match self { + Self::Session => SESSION_WINDOW_MINUTES, + Self::Weekly => WEEKLY_WINDOW_MINUTES, + } + } +} + +/// Chronological history for one series. +#[derive(Debug, Clone, PartialEq)] +pub struct PlanUtilizationSeriesHistory { + pub name: PlanUtilizationSeriesName, + pub window_minutes: u32, + pub entries: Vec, +} + +/// Burn-estimator knobs matching upstream SessionEquivalentBurnEstimator. +pub struct SessionEquivalentBurnEstimator; + +impl SessionEquivalentBurnEstimator { + pub const DEFAULT_SAMPLE_LIMIT: usize = 7; + pub const MINIMUM_SAMPLE_COUNT: usize = 3; + + /// Estimate median full-session weekly burn from plan-utilization histories. + pub fn estimate( + histories: &[PlanUtilizationSeriesHistory], + current_session_resets_at: Option>, + now: DateTime, + sample_limit: usize, + ) -> Option { + if sample_limit == 0 { + return None; + } + + let session_history = histories.iter().find(|h| { + h.name == PlanUtilizationSeriesName::Session + && h.window_minutes == SESSION_WINDOW_MINUTES + })?; + let weekly_history = histories.iter().find(|h| { + h.name == PlanUtilizationSeriesName::Weekly && h.window_minutes == WEEKLY_WINDOW_MINUTES + })?; + + let session_duration = Duration::seconds(i64::from(SESSION_WINDOW_MINUTES) * 60); + let weekly_duration = Duration::seconds(i64::from(WEEKLY_WINDOW_MINUTES) * 60); + + if !is_chronologically_ordered(&session_history.entries) + || !is_chronologically_ordered(&weekly_history.entries) + { + return None; + } + + // Expired/invalid-remaining current session = idle: don't gate the burn + // estimate on it. Implausible far-future resets still abort. + let effective_current = match current_session_resets_at { + None => None, + Some(current) => { + let remaining = (current - now).num_seconds() as f64; + let max_remaining = + session_duration.num_seconds() as f64 + RESET_TOLERANCE_SECS as f64; + if !remaining.is_finite() || remaining > max_remaining { + return None; + } + if remaining <= 0.0 { + None + } else { + Some(current) + } + } + }; + + let mut groups: Vec = Vec::new(); + for entry in &session_history.entries { + if !entry.used_percent.is_finite() || !(0.0..=100.0).contains(&entry.used_percent) { + continue; + } + let Some(resets_at) = entry.resets_at else { + continue; + }; + if !is_plausible_reset(resets_at, entry.captured_at, session_duration) { + continue; + } + + if let Some(last) = groups.last_mut() { + let delta = (last.resets_at - resets_at).num_seconds().abs(); + if delta <= RESET_TOLERANCE_SECS { + last.entries.push(entry.clone()); + last.maximum_used_percent = last.maximum_used_percent.max(entry.used_percent); + continue; + } + if last.resets_at > resets_at { + return None; + } + } + + groups.push(SessionGroup { + resets_at, + entries: vec![entry.clone()], + maximum_used_percent: entry.used_percent, + }); + } + + let completed_active_groups: Vec<&SessionGroup> = groups + .iter() + .rev() + .filter(|group| { + let precedes_current = effective_current + .map(|cur| group.resets_at < cur - Duration::seconds(RESET_TOLERANCE_SECS)) + .unwrap_or(true); + precedes_current && group.resets_at <= now && group.maximum_used_percent > 0.0 + }) + .collect(); + + let weekly_entries: Vec<&PlanUtilizationHistoryEntry> = weekly_history + .entries + .iter() + .filter(|entry| { + entry.used_percent.is_finite() + && (0.0..=100.0).contains(&entry.used_percent) + && entry + .resets_at + .is_some_and(|r| is_plausible_reset(r, entry.captured_at, weekly_duration)) + }) + .collect(); + if weekly_entries.is_empty() { + return None; + } + + let mut burns = Vec::new(); + for group in completed_active_groups.into_iter().take(sample_limit) { + if let Some(burn) = normalized_burn(group, &weekly_entries, session_duration) { + burns.push(burn); + } + } + + if burns.len() < Self::MINIMUM_SAMPLE_COUNT { + return None; + } + burns.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let middle = burns.len() / 2; + let median = if burns.len().is_multiple_of(2) { + (burns[middle - 1] + burns[middle]) / 2.0 + } else { + burns[middle] + }; + if !median.is_finite() || median <= 0.0 { + return None; + } + + Some(SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: median, + sample_count: burns.len(), + }) + } +} + +impl SessionEquivalentForecast { + /// Build a forecast from current windows + burn estimate. + pub fn make( + session_window: &crate::core::RateWindow, + weekly_window: &crate::core::RateWindow, + burn_estimate: &SessionEquivalentBurnEstimate, + weekly_window_id: Option, + now: DateTime, + work_days: Option, + ) -> Option { + if session_window.is_informational { + return None; + } + if session_window.window_minutes != Some(SESSION_WINDOW_MINUTES) { + return None; + } + if weekly_window.window_minutes != Some(WEEKLY_WINDOW_MINUTES) { + return None; + } + let weekly_resets_at = weekly_window.resets_at?; + if !weekly_window.used_percent.is_finite() + || !(0.0..=100.0).contains(&weekly_window.used_percent) + { + return None; + } + if !burn_estimate.median_weekly_percent_per_window.is_finite() + || burn_estimate.median_weekly_percent_per_window <= 0.0 + || burn_estimate.sample_count < SessionEquivalentBurnEstimator::MINIMUM_SAMPLE_COUNT + { + return None; + } + + let session_seconds = f64::from(SESSION_WINDOW_MINUTES) * 60.0; + let weekly_seconds = f64::from(WEEKLY_WINDOW_MINUTES) * 60.0; + + // Expired session resets mean the window is idle — still show the learned + // estimate against the weekly lane. Only reject implausible far-future resets. + if let Some(session_resets_at) = session_window.resets_at { + let session_remaining = (session_resets_at - now).num_seconds() as f64; + if !session_remaining.is_finite() + || session_remaining > session_seconds + RESET_TOLERANCE_SECS as f64 + { + return None; + } + } + + let weekly_remaining = (weekly_resets_at - now).num_seconds() as f64; + if !weekly_remaining.is_finite() + || weekly_remaining <= 0.0 + || weekly_remaining > weekly_seconds + RESET_TOLERANCE_SECS as f64 + { + return None; + } + + let remaining_weekly_percent = (100.0 - weekly_window.used_percent).clamp(0.0, 100.0); + if remaining_weekly_percent <= 0.0 { + return None; + } + let estimated_windows = + remaining_weekly_percent / burn_estimate.median_weekly_percent_per_window; + if !estimated_windows.is_finite() || estimated_windows < 0.0 { + return None; + } + + let remaining_seconds = effective_remaining_seconds(now, weekly_resets_at, work_days); + if remaining_seconds < 0.0 { + return None; + } + let available_windows_until_reset = remaining_seconds / session_seconds; + let windows_until_reset = available_windows_until_reset.floor() as i64; + + Some(Self { + estimated_windows_to_exhaust_weekly: estimated_windows, + windows_until_reset, + available_windows_until_reset, + sample_count: burn_estimate.sample_count, + weekly_resets_at, + weekly_used_percent: weekly_window.used_percent, + weekly_window_id, + }) + } +} + +#[derive(Debug, Clone)] +struct SessionGroup { + resets_at: DateTime, + entries: Vec, + maximum_used_percent: f64, +} + +#[derive(Debug, Clone)] +struct BurnObservation { + session_used_percent: f64, + weekly_entry: PlanUtilizationHistoryEntry, +} + +fn normalized_burn( + group: &SessionGroup, + weekly_entries: &[&PlanUtilizationHistoryEntry], + session_duration: Duration, +) -> Option { + let first_session = group.entries.first()?; + let last_session = group.entries.last()?; + + let mut observations: Vec = Vec::new(); + let window_start = group.resets_at - session_duration; + + if let Some(weekly_start) = nearest_entry( + window_start, + weekly_entries, + Duration::seconds(RESET_TOLERANCE_SECS), + true, + ) && weekly_start.captured_at <= window_start + && weekly_start.captured_at < first_session.captured_at + { + observations.push(BurnObservation { + session_used_percent: 0.0, + weekly_entry: weekly_start.clone(), + }); + } + + for session_entry in &group.entries { + // Upstream observationAlignmentTolerance is 0 → exact capture match only. + if let Some(weekly_entry) = nearest_entry( + session_entry.captured_at, + weekly_entries, + Duration::zero(), + false, + ) { + observations.push(BurnObservation { + session_used_percent: session_entry.used_percent, + weekly_entry: weekly_entry.clone(), + }); + } + } + + if group.maximum_used_percent >= 100.0 + && let Some(weekly_end) = nearest_entry( + group.resets_at, + weekly_entries, + Duration::seconds(RESET_TOLERANCE_SECS), + true, + ) + && weekly_end.captured_at <= group.resets_at + && last_session.captured_at < weekly_end.captured_at + { + observations.push(BurnObservation { + session_used_percent: 100.0, + weekly_entry: weekly_end.clone(), + }); + } + + observations.sort_by(|lhs, rhs| { + match lhs + .weekly_entry + .captured_at + .cmp(&rhs.weekly_entry.captured_at) + { + std::cmp::Ordering::Equal => lhs + .session_used_percent + .partial_cmp(&rhs.session_used_percent) + .unwrap_or(std::cmp::Ordering::Equal), + other => other, + } + }); + + let start = observations.first()?; + let end = observations.last()?; + if start.weekly_entry.captured_at >= end.weekly_entry.captured_at { + return None; + } + let start_reset = start.weekly_entry.resets_at?; + let end_reset = end.weekly_entry.resets_at?; + if (start_reset - end_reset).num_seconds().abs() > RESET_TOLERANCE_SECS { + return None; + } + + let session_consumption = end.session_used_percent - start.session_used_percent; + let weekly_burn = end.weekly_entry.used_percent - start.weekly_entry.used_percent; + if !session_consumption.is_finite() + || session_consumption <= 0.0 + || !weekly_burn.is_finite() + || weekly_burn <= 0.0 + { + return None; + } + let full_allowance_burn = 100.0 * weekly_burn / session_consumption; + if !full_allowance_burn.is_finite() || full_allowance_burn <= 0.0 { + return None; + } + Some(full_allowance_burn) +} + +fn nearest_entry<'a>( + target: DateTime, + entries: &[&'a PlanUtilizationHistoryEntry], + tolerance: Duration, + require_not_after_target: bool, +) -> Option<&'a PlanUtilizationHistoryEntry> { + if entries.is_empty() { + return None; + } + let mut lower = 0usize; + let mut upper = entries.len(); + while lower < upper { + let middle = (lower + upper) / 2; + if entries[middle].captured_at < target { + lower = middle + 1; + } else { + upper = middle; + } + } + + let mut candidates = Vec::new(); + if lower < entries.len() { + candidates.push(entries[lower]); + } + if lower > 0 { + candidates.push(entries[lower - 1]); + } + + let tol_secs = tolerance.num_seconds().abs(); + candidates + .into_iter() + .filter(|e| !require_not_after_target || e.captured_at <= target) + .filter(|e| (e.captured_at - target).num_seconds().abs() <= tol_secs) + .min_by_key(|e| (e.captured_at - target).num_seconds().abs()) +} + +fn is_chronologically_ordered(entries: &[PlanUtilizationHistoryEntry]) -> bool { + entries + .windows(2) + .all(|w| w[0].captured_at <= w[1].captured_at) +} + +fn is_plausible_reset( + resets_at: DateTime, + captured_at: DateTime, + duration: Duration, +) -> bool { + let remaining = (resets_at - captured_at).num_seconds(); + remaining >= -RESET_TOLERANCE_SECS && remaining <= duration.num_seconds() + RESET_TOLERANCE_SECS +} + +/// Remaining seconds until weekly reset, optionally counting only work days. +/// +/// `work_days` in `[2, 6]` keeps ISO weekdays `1..=work_days` (Mon..). +/// Anything else falls back to wall-clock remaining time. +pub fn effective_remaining_seconds( + now: DateTime, + resets_at: DateTime, + work_days: Option, +) -> f64 { + let wall = (resets_at - now).num_seconds().max(0) as f64; + let Some(work_days) = work_days else { + return wall; + }; + if !(2..=6).contains(&work_days) { + return wall; + } + + let mut work_seconds = 0.0_f64; + let mut cursor = now; + while cursor < resets_at { + let next_day = match start_of_next_utc_day(cursor) { + Some(d) if d > cursor => d, + _ => return wall, + }; + let slice_end = if next_day < resets_at { + next_day + } else { + resets_at + }; + if is_workday(cursor, work_days) { + work_seconds += (slice_end - cursor).num_seconds() as f64; + } + cursor = slice_end; + } + work_seconds +} + +fn start_of_next_utc_day(dt: DateTime) -> Option> { + let date = dt.date_naive() + Duration::days(1); + date.and_hms_opt(0, 0, 0) + .map(|naive| DateTime::::from_naive_utc_and_offset(naive, Utc)) +} + +fn is_workday(date: DateTime, work_days: u8) -> bool { + let iso = match date.weekday() { + Weekday::Mon => 1, + Weekday::Tue => 2, + Weekday::Wed => 3, + Weekday::Thu => 4, + Weekday::Fri => 5, + Weekday::Sat => 6, + Weekday::Sun => 7, + }; + iso <= work_days +} + +// ── In-process history store ───────────────────────────────────────── + +/// Append-only ring of session/weekly observations for burn estimation. +/// +/// ponytail: history is process-local and lost on restart; upgrade path = disk cache +/// keyed by provider+account once forecast quality justifies persistence. +#[derive(Debug, Default)] +pub struct SessionEquivalentHistoryStore { + by_provider: HashMap, +} + +#[derive(Debug, Default)] +struct ProviderHistory { + session: Vec, + weekly: Vec, +} + +static HISTORY_STORE: LazyLock> = + LazyLock::new(|| Mutex::new(SessionEquivalentHistoryStore::default())); + +impl SessionEquivalentHistoryStore { + /// Record one observation pair for a provider (session + optional weekly). + pub fn record( + &mut self, + provider_id: &str, + session: Option, + weekly: Option, + sample_limit: usize, + ) { + let limit = sample_limit.max(1); + // Keep more raw points than completed groups so grouping still has density. + let ring = limit.saturating_mul(8).max(24); + let hist = self.by_provider.entry(provider_id.to_string()).or_default(); + if let Some(entry) = session { + push_ring(&mut hist.session, entry, ring); + } + if let Some(entry) = weekly { + push_ring(&mut hist.weekly, entry, ring); + } + } + + pub fn histories(&self, provider_id: &str) -> Vec { + let Some(hist) = self.by_provider.get(provider_id) else { + return Vec::new(); + }; + let mut out = Vec::new(); + if !hist.session.is_empty() { + out.push(PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: hist.session.clone(), + }); + } + if !hist.weekly.is_empty() { + out.push(PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: hist.weekly.clone(), + }); + } + out + } +} + +fn push_ring( + buf: &mut Vec, + entry: PlanUtilizationHistoryEntry, + ring: usize, +) { + buf.push(entry); + if buf.len() > ring { + let drop_n = buf.len() - ring; + buf.drain(0..drop_n); + } +} + +/// Global in-process history used by Claude/Codex snapshot refresh. +pub fn global_history_store() -> &'static Mutex { + &HISTORY_STORE +} + +/// Record session/weekly windows from a live provider usage snapshot. +pub fn record_provider_windows( + provider_id: &str, + session: &crate::core::RateWindow, + weekly: Option<&crate::core::RateWindow>, + now: DateTime, +) { + if session.is_informational + || session.window_minutes != Some(SESSION_WINDOW_MINUTES) + || !session.used_percent.is_finite() + { + return; + } + let session_entry = PlanUtilizationHistoryEntry { + captured_at: now, + used_percent: session.used_percent.clamp(0.0, 100.0), + resets_at: session.resets_at, + }; + let weekly_entry = weekly.and_then(|w| { + if w.is_informational + || w.window_minutes != Some(WEEKLY_WINDOW_MINUTES) + || !w.used_percent.is_finite() + { + return None; + } + Some(PlanUtilizationHistoryEntry { + captured_at: now, + used_percent: w.used_percent.clamp(0.0, 100.0), + resets_at: w.resets_at, + }) + }); + + if let Ok(mut guard) = global_history_store().lock() { + guard.record( + provider_id, + Some(session_entry), + weekly_entry, + SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, + ); + } +} + +/// Last learned full-session burn estimate retained across idle refreshes. +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct RetainedFullSessionEstimate { + pub estimate: f64, + pub updated_at: DateTime, +} + +/// Keep the last positive full-session estimate when a refresh yields none. +/// +/// `fresh` wins when present and valid; otherwise `previous` is kept so the UI +/// does not blank while the session window is idle (upstream #2336). +pub fn retain_last_full_session_estimate( + previous: Option, + fresh: Option, + now: DateTime, +) -> Option { + match fresh { + Some(estimate) if estimate.is_finite() && estimate > 0.0 => { + Some(RetainedFullSessionEstimate { + estimate, + updated_at: now, + }) + } + _ => previous, + } +} + +// ponytail: in-memory only, lost on restart; upgrade = persist to cache file. +#[derive(Debug, Default)] +struct LastFullSessionEstimateStore { + by_provider: HashMap, +} + +static LAST_FULL_SESSION_ESTIMATE_STORE: LazyLock> = + LazyLock::new(|| Mutex::new(LastFullSessionEstimateStore::default())); + +fn last_full_session_estimate_store() -> &'static Mutex { + &LAST_FULL_SESSION_ESTIMATE_STORE +} + +/// Remember/recall the last learned full-session burn estimate for a provider. +pub fn remember_full_session_estimate( + provider_id: &str, + fresh: Option, + now: DateTime, +) -> Option { + let Ok(mut guard) = last_full_session_estimate_store().lock() else { + return fresh.filter(|v| v.is_finite() && *v > 0.0); + }; + let previous = guard.by_provider.get(provider_id).copied(); + let retained = retain_last_full_session_estimate(previous, fresh, now); + if let Some(entry) = retained { + guard.by_provider.insert(provider_id.to_string(), entry); + Some(entry.estimate) + } else { + guard.by_provider.remove(provider_id); + None + } +} + +/// Compute forecast for a provider using the in-process history ring. +pub fn forecast_for_provider( + provider_id: &str, + session: &crate::core::RateWindow, + weekly: &crate::core::RateWindow, + now: DateTime, + work_days: Option, +) -> Option { + let histories = global_history_store() + .lock() + .ok() + .map(|g| g.histories(provider_id)) + .unwrap_or_default(); + let fresh_burn = SessionEquivalentBurnEstimator::estimate( + &histories, + session.resets_at, + now, + SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, + ); + let fresh_median = fresh_burn + .as_ref() + .map(|b| b.median_weekly_percent_per_window); + let sample_count = fresh_burn + .as_ref() + .map(|b| b.sample_count) + .unwrap_or(SessionEquivalentBurnEstimator::MINIMUM_SAMPLE_COUNT); + let median = remember_full_session_estimate(provider_id, fresh_median, now)?; + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: median, + sample_count, + }; + SessionEquivalentForecast::make(session, weekly, &burn, None, now, work_days) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::RateWindow; + use chrono::TimeZone; + + fn ts(secs: i64) -> DateTime { + Utc.timestamp_opt(secs, 0).single().unwrap() + } + + fn entry(captured: i64, used: f64, resets: i64) -> PlanUtilizationHistoryEntry { + PlanUtilizationHistoryEntry { + captured_at: ts(captured), + used_percent: used, + resets_at: Some(ts(resets)), + } + } + + fn session_secs() -> i64 { + i64::from(SESSION_WINDOW_MINUTES) * 60 + } + + fn weekly_secs() -> i64 { + i64::from(WEEKLY_WINDOW_MINUTES) * 60 + } + + /// Build three completed full-burn sessions with aligned weekly observations. + fn three_sample_histories( + base: i64, + ) -> ( + Vec, + DateTime, + DateTime, + ) { + let s = session_secs(); + let w = weekly_secs(); + let weekly_reset = base + w; + + let mut session_entries = Vec::new(); + let mut weekly_entries = Vec::new(); + + // Three completed sessions ending at base+s, base+2s, base+3s. + for i in 0..3 { + let start = base + i * s; + let end = start + s; + // Window-start weekly anchor (session used 0). + weekly_entries.push(entry(start - 30, 10.0 + i as f64 * 5.0, weekly_reset)); + // Mid / end session captures with exact weekly alignment. + session_entries.push(entry(start + 60, 40.0, end)); + weekly_entries.push(entry(start + 60, 12.0 + i as f64 * 5.0, weekly_reset)); + session_entries.push(entry(end - 60, 100.0, end)); + weekly_entries.push(entry(end - 60, 15.0 + i as f64 * 5.0, weekly_reset)); + // End-of-window weekly for max>=100 path. + weekly_entries.push(entry(end, 15.5 + i as f64 * 5.0, weekly_reset)); + } + + // Current open session. + let current_start = base + 3 * s; + let current_end = current_start + s; + let now_secs = current_start + 30 * 60; + session_entries.push(entry(now_secs, 20.0, current_end)); + weekly_entries.push(entry(now_secs, 30.0, weekly_reset)); + + session_entries.sort_by_key(|e| e.captured_at); + weekly_entries.sort_by_key(|e| e.captured_at); + + let histories = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: session_entries, + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: weekly_entries, + }, + ]; + (histories, ts(now_secs), ts(current_end)) + } + + #[test] + fn median_math_from_three_full_sessions() { + let base = 1_700_000_000_i64; + let (histories, now, current_end) = three_sample_histories(base); + let estimate = SessionEquivalentBurnEstimator::estimate( + &histories, + Some(current_end), + now, + SessionEquivalentBurnEstimator::DEFAULT_SAMPLE_LIMIT, + ) + .expect("estimate"); + assert_eq!(estimate.sample_count, 3); + // Mid-session captures dominate when start/end anchors are optional: + // weekly 12→15 over session 40→100 ⇒ 100 * 3 / 60 = 5.0 + // (window-start + end anchors yield the same full-allowance burn here). + assert!( + (estimate.median_weekly_percent_per_window - 5.0).abs() < 0.01, + "median={}", + estimate.median_weekly_percent_per_window + ); + } + + #[test] + fn grouping_tolerance_merges_nearby_resets() { + let s = session_secs(); + let base = 1_700_100_000_i64; + let weekly_reset = base + weekly_secs(); + let end = base + s; + + // Nearby resets within ±120s must merge into one group (not out-of-order). + let single_group = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: vec![ + entry(base + 60, 30.0, end), + entry(base + 120, 60.0, end + 100), + entry(base + 180, 90.0, end + 110), + ], + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: vec![ + entry(base - 30, 5.0, weekly_reset), + entry(base + 60, 8.0, weekly_reset), + entry(base + 120, 10.0, weekly_reset), + entry(base + 180, 12.0, weekly_reset), + ], + }, + ]; + // One completed group → below minimum sample count, but not an ordering reject. + assert!( + SessionEquivalentBurnEstimator::estimate( + &single_group, + Some(ts(end + s)), + ts(end + 60), + 7 + ) + .is_none() + ); + + // Strictly decreasing resets_at across groups is rejected. + let out_of_order = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: vec![ + entry(base + 60, 30.0, end + s), + entry(base + 120, 60.0, end), // jumps backward beyond tolerance + ], + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: vec![ + entry(base + 60, 8.0, weekly_reset), + entry(base + 120, 10.0, weekly_reset), + ], + }, + ]; + assert!( + SessionEquivalentBurnEstimator::estimate( + &out_of_order, + Some(ts(end + 2 * s)), + ts(end + s + 60), + 7 + ) + .is_none() + ); + } + + #[test] + fn insufficient_samples_returns_none() { + let base = 1_700_200_000_i64; + let s = session_secs(); + let weekly_reset = base + weekly_secs(); + let end = base + s; + let histories = vec![ + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Session, + window_minutes: SESSION_WINDOW_MINUTES, + entries: vec![entry(base + 60, 50.0, end), entry(base + 120, 80.0, end)], + }, + PlanUtilizationSeriesHistory { + name: PlanUtilizationSeriesName::Weekly, + window_minutes: WEEKLY_WINDOW_MINUTES, + entries: vec![ + entry(base - 30, 10.0, weekly_reset), + entry(base + 60, 12.0, weekly_reset), + entry(base + 120, 14.0, weekly_reset), + ], + }, + ]; + assert!( + SessionEquivalentBurnEstimator::estimate( + &histories, + Some(ts(end + s)), + ts(end + 60), + 7 + ) + .is_none() + ); + } + + #[test] + fn forecast_make_fractional_windows() { + let now = ts(1_700_300_000); + let session = RateWindow::with_details( + 25.0, + Some(SESSION_WINDOW_MINUTES), + Some(now + Duration::hours(3)), + None, + ); + let weekly = RateWindow::with_details( + 40.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(3)), + None, + ); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 12.0, + sample_count: 3, + }; + let forecast = SessionEquivalentForecast::make(&session, &weekly, &burn, None, now, None) + .expect("forecast"); + // remaining weekly 60 / 12 = 5.0 + assert!((forecast.estimated_windows_to_exhaust_weekly - 5.0).abs() < 1e-9); + assert!(forecast.available_windows_until_reset > 0.0); + assert_eq!( + forecast.windows_until_reset, + forecast.available_windows_until_reset.floor() as i64 + ); + assert_eq!(forecast.sample_count, 3); + assert!((forecast.weekly_used_percent - 40.0).abs() < 1e-9); + } + + #[test] + fn work_days_reduces_available_windows_vs_wall_clock() { + // Friday 18:00 UTC → Monday 18:00 UTC (weekend in between). + let friday = Utc + .with_ymd_and_hms(2024, 6, 14, 18, 0, 0) + .single() + .unwrap(); + assert_eq!(friday.weekday(), Weekday::Fri); + let monday = Utc + .with_ymd_and_hms(2024, 6, 17, 18, 0, 0) + .single() + .unwrap(); + + let wall = effective_remaining_seconds(friday, monday, None); + let work5 = effective_remaining_seconds(friday, monday, Some(5)); + assert!((wall - 3.0 * 24.0 * 3600.0).abs() < 1.0); + // Fri 18→24 = 6h, Mon 0→18 = 18h → 24h work seconds (no Sat/Sun). + assert!((work5 - 24.0 * 3600.0).abs() < 1.0); + assert!(work5 < wall); + + let session = RateWindow::with_details( + 10.0, + Some(SESSION_WINDOW_MINUTES), + Some(friday + Duration::hours(4)), + None, + ); + let weekly = + RateWindow::with_details(20.0, Some(WEEKLY_WINDOW_MINUTES), Some(monday), None); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 4, + }; + let wall_f = + SessionEquivalentForecast::make(&session, &weekly, &burn, None, friday, None).unwrap(); + let work_f = + SessionEquivalentForecast::make(&session, &weekly, &burn, None, friday, Some(5)) + .unwrap(); + assert!(work_f.available_windows_until_reset < wall_f.available_windows_until_reset); + assert_eq!( + work_f.windows_until_reset, + work_f.available_windows_until_reset.floor() as i64 + ); + } + + #[test] + fn work_days_out_of_range_uses_wall_clock() { + let now = ts(1_700_400_000); + let reset = now + Duration::days(2); + let wall = effective_remaining_seconds(now, reset, None); + assert!((effective_remaining_seconds(now, reset, Some(1)) - wall).abs() < 1e-9); + assert!((effective_remaining_seconds(now, reset, Some(7)) - wall).abs() < 1e-9); + } + + #[test] + fn history_ring_retains_latest_samples() { + let mut store = SessionEquivalentHistoryStore::default(); + let base = ts(1_700_500_000); + for i in 0..30 { + store.record( + "claude", + Some(PlanUtilizationHistoryEntry { + captured_at: base + Duration::minutes(i), + used_percent: i as f64, + resets_at: Some(base + Duration::hours(5)), + }), + None, + 7, + ); + } + let h = store.histories("claude"); + assert_eq!(h.len(), 1); + assert!(h[0].entries.len() <= 7 * 8); + assert!((h[0].entries.last().unwrap().used_percent - 29.0).abs() < 1e-9); + } + + #[test] + fn make_rejects_bad_windows_and_zero_remaining() { + let now = ts(1_700_600_000); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 3, + }; + let session = RateWindow::with_details( + 10.0, + Some(SESSION_WINDOW_MINUTES), + Some(now + Duration::hours(2)), + None, + ); + let exhausted = RateWindow::with_details( + 100.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(2)), + None, + ); + assert!( + SessionEquivalentForecast::make(&session, &exhausted, &burn, None, now, None).is_none() + ); + + let low_samples = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 2, + }; + let weekly = RateWindow::with_details( + 50.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(2)), + None, + ); + assert!( + SessionEquivalentForecast::make(&session, &weekly, &low_samples, None, now, None) + .is_none() + ); + } + + #[test] + fn retain_last_full_session_estimate_survives_idle_refresh() { + let t0 = ts(1_700_000_000); + let t1 = t0 + Duration::hours(1); + let learned = retain_last_full_session_estimate(None, Some(12.5), t0) + .expect("fresh estimate should stick"); + assert_eq!(learned.estimate, 12.5); + assert_eq!(learned.updated_at, t0); + + // Idle refresh yields no fresh sample — keep the last learned value. + let idle = retain_last_full_session_estimate(Some(learned), None, t1) + .expect("idle refresh must keep last estimate"); + assert_eq!(idle.estimate, 12.5); + assert_eq!(idle.updated_at, t0); + + // Non-positive / non-finite fresh values also keep previous. + let still = retain_last_full_session_estimate(Some(idle), Some(0.0), t1) + .expect("zero fresh must not clear"); + assert_eq!(still.estimate, 12.5); + let still = retain_last_full_session_estimate(Some(still), Some(f64::NAN), t1) + .expect("nan fresh must not clear"); + assert_eq!(still.estimate, 12.5); + } + + #[test] + fn retain_last_full_session_estimate_replaced_by_fresh_sample() { + let t0 = ts(1_700_000_000); + let t1 = t0 + Duration::hours(2); + let previous = retain_last_full_session_estimate(None, Some(8.0), t0).unwrap(); + let next = retain_last_full_session_estimate(Some(previous), Some(15.0), t1) + .expect("fresh sample replaces previous"); + assert_eq!(next.estimate, 15.0); + assert_eq!(next.updated_at, t1); + } + + #[test] + fn make_keeps_forecast_when_session_reset_is_expired_idle() { + let now = ts(1_700_600_000); + let burn = SessionEquivalentBurnEstimate { + median_weekly_percent_per_window: 10.0, + sample_count: 3, + }; + // Session reset already passed (idle) — forecast still derives from weekly. + let session = RateWindow::with_details( + 0.0, + Some(SESSION_WINDOW_MINUTES), + Some(now - Duration::hours(1)), + None, + ); + let weekly = RateWindow::with_details( + 40.0, + Some(WEEKLY_WINDOW_MINUTES), + Some(now + Duration::days(3)), + None, + ); + let forecast = SessionEquivalentForecast::make(&session, &weekly, &burn, None, now, None) + .expect("idle expired session should still forecast"); + assert!((forecast.estimated_windows_to_exhaust_weekly - 6.0).abs() < 1e-9); + } +} From c2df183fd92097b5e4f08f33b352cfa0bbc8aa80 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:32 +0700 Subject: [PATCH 08/15] Add Codex Workspaces local indexing foundation --- .../src-tauri/src/commands/bridge.rs | 48 + .../src/commands/codex_workspaces.rs | 35 + .../src-tauri/src/commands/mod.rs | 2 + apps/desktop-tauri/src-tauri/src/main.rs | 1 + apps/desktop-tauri/src-tauri/src/powertoys.rs | 2 + apps/desktop-tauri/src/lib/tauri.ts | 11 + apps/desktop-tauri/src/types/bridge.ts | 76 ++ rust/src/cli/mod.rs | 4 + rust/src/cli/workspaces.rs | 146 +++ rust/src/codex_workspaces/indexer.rs | 887 ++++++++++++++++++ rust/src/codex_workspaces/mod.rs | 18 + rust/src/codex_workspaces/sidecar.rs | 190 ++++ rust/src/codex_workspaces/types.rs | 182 ++++ rust/src/lib.rs | 1 + rust/src/main.rs | 1 + 15 files changed, 1604 insertions(+) create mode 100644 apps/desktop-tauri/src-tauri/src/commands/codex_workspaces.rs create mode 100644 rust/src/cli/workspaces.rs create mode 100644 rust/src/codex_workspaces/indexer.rs create mode 100644 rust/src/codex_workspaces/mod.rs create mode 100644 rust/src/codex_workspaces/sidecar.rs create mode 100644 rust/src/codex_workspaces/types.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 424c820179..348cbf53e7 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -61,6 +61,8 @@ pub struct CostSnapshotBridge { pub resets_at: Option, pub formatted_used: String, pub formatted_limit: Option, + pub balance: Option, + pub formatted_balance: Option, } #[derive(Debug, Clone, Serialize)] @@ -83,6 +85,18 @@ pub struct PaceSnapshot { pub actual_used_percent: f64, } +/// Session-equivalent weekly forecast for Claude/Codex menu secondary line. +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionEquivalentForecastSnapshot { + pub estimated_windows_to_exhaust_weekly: f64, + pub windows_until_reset: i64, + pub available_windows_until_reset: f64, + pub sample_count: usize, + pub weekly_resets_at: String, + pub weekly_used_percent: f64, +} + /// A frontend-friendly snapshot of one provider's usage data. #[derive(Debug, Clone, Serialize)] #[serde(rename_all = "camelCase")] @@ -107,6 +121,8 @@ pub struct ProviderUsageSnapshot { pub tray_status_label: Option, pub fetch_duration_ms: Option, pub wayfinder_usage: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub session_equivalent_forecast: Option, } pub(crate) fn filter_hidden_codex_spark_rows( @@ -168,6 +184,9 @@ impl ProviderUsageSnapshot { s }); + let session_equivalent_forecast = + session_equivalent_forecast_for(id, &usage.primary, usage.secondary.as_ref()); + Self { provider_id: id.cli_name().to_string(), display_name: id.display_name().to_string(), @@ -204,6 +223,8 @@ impl ProviderUsageSnapshot { resets_at: c.resets_at.map(|dt| dt.to_rfc3339()), formatted_used: c.format_used(), formatted_limit: c.format_limit(), + balance: c.balance, + formatted_balance: c.format_balance(), }), plan_name: usage.login_method.clone(), account_email: usage.account_email.clone(), @@ -215,6 +236,7 @@ impl ProviderUsageSnapshot { tray_status_label: None, fetch_duration_ms: None, wayfinder_usage: result.wayfinder_usage.clone(), + session_equivalent_forecast, } } @@ -253,10 +275,36 @@ impl ProviderUsageSnapshot { tray_status_label: None, fetch_duration_ms: None, wayfinder_usage: None, + session_equivalent_forecast: None, } } } +fn session_equivalent_forecast_for( + id: ProviderId, + session: &RateWindow, + weekly: Option<&RateWindow>, +) -> Option { + if !matches!(id, ProviderId::Claude | ProviderId::Codex) { + return None; + } + let weekly = weekly?; + let now = chrono::Utc::now(); + let provider_id = id.cli_name(); + codexbar::core::record_provider_windows(provider_id, session, Some(weekly), now); + let work_days = Settings::load().weekly_progress_work_days; + let forecast = + codexbar::core::forecast_for_provider(provider_id, session, weekly, now, work_days)?; + Some(SessionEquivalentForecastSnapshot { + estimated_windows_to_exhaust_weekly: forecast.estimated_windows_to_exhaust_weekly, + windows_until_reset: forecast.windows_until_reset, + available_windows_until_reset: forecast.available_windows_until_reset, + sample_count: forecast.sample_count, + weekly_resets_at: forecast.weekly_resets_at.to_rfc3339(), + weekly_used_percent: forecast.weekly_used_percent, + }) +} + /// Build a compact tray status label from a raw snapshot using the current language. /// Localization is done at render time so cached snapshots stay language-neutral. pub(crate) fn compact_tray_status_label( diff --git a/apps/desktop-tauri/src-tauri/src/commands/codex_workspaces.rs b/apps/desktop-tauri/src-tauri/src/commands/codex_workspaces.rs new file mode 100644 index 0000000000..88a6cb7368 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/codex_workspaces.rs @@ -0,0 +1,35 @@ +//! Codex local Workspaces snapshot bridge. + +use codexbar::codex_workspaces::{CodexLocalProjectUsageSnapshot, CodexWorkspacesIndex}; +use codexbar::settings::Settings; +use tauri::State; + +use crate::state::AppState; +use std::sync::Mutex; + +/// Return the local Codex workspaces usage snapshot. +/// +/// Privacy-redacts paths/titles when `settings.hide_personal_info` is set. +#[tauri::command] +pub async fn get_codex_workspaces_snapshot( + _state: State<'_, Mutex>, + force_refresh: Option, + history_days: Option, +) -> Result { + let force = force_refresh.unwrap_or(false); + let days = history_days.unwrap_or(30); + let hide = Settings::load().hide_personal_info; + + tauri::async_runtime::spawn_blocking(move || { + let index = CodexWorkspacesIndex::new(days); + let mut snapshot = index + .load_snapshot(force, |_| {}) + .map_err(|e| e.to_string())?; + if hide { + snapshot.redact_for_privacy(); + } + Ok(snapshot) + }) + .await + .map_err(|e| format!("workspaces worker failed: {e}"))? +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index 04bcdfcdd3..2e32d549ad 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -30,6 +30,7 @@ mod usage_spend; mod agent_sessions; mod bridge; mod browser_import; +mod codex_workspaces; mod credential_detection; mod credentials; mod locale_cmd; @@ -44,6 +45,7 @@ mod system; pub use agent_sessions::*; pub(crate) use bridge::*; pub use browser_import::*; +pub use codex_workspaces::*; pub use credential_detection::*; pub use credentials::*; pub use locale_cmd::*; diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index fbe9070488..aa5d09fd22 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -180,6 +180,7 @@ fn main() { commands::get_provider_chart_data, commands::get_provider_local_usage_summary, commands::get_usage_spend_summary, + commands::get_codex_workspaces_snapshot, commands::reorder_providers, commands::set_provider_cookie_source, commands::get_provider_cookie_source_options, diff --git a/apps/desktop-tauri/src-tauri/src/powertoys.rs b/apps/desktop-tauri/src-tauri/src/powertoys.rs index 8633332d24..a03f2c5977 100644 --- a/apps/desktop-tauri/src-tauri/src/powertoys.rs +++ b/apps/desktop-tauri/src-tauri/src/powertoys.rs @@ -205,6 +205,7 @@ mod tests { tray_status_label: None, fetch_duration_ms: None, wayfinder_usage: None, + session_equivalent_forecast: None, }); let value = serde_json::to_value(snapshot).unwrap(); @@ -248,6 +249,7 @@ mod tests { tray_status_label: None, fetch_duration_ms: None, wayfinder_usage: None, + session_equivalent_forecast: None, }); let value = serde_json::to_value(snapshot).unwrap(); diff --git a/apps/desktop-tauri/src/lib/tauri.ts b/apps/desktop-tauri/src/lib/tauri.ts index 40f659efb2..46866fe8ec 100644 --- a/apps/desktop-tauri/src/lib/tauri.ts +++ b/apps/desktop-tauri/src/lib/tauri.ts @@ -32,6 +32,7 @@ import type { SessionFocusResult, TrayVisibilityStatusDto, UsageSpendSummary, + CodexLocalProjectUsageSnapshot, } from "../types/bridge"; export function getBootstrapState(): Promise { @@ -235,6 +236,16 @@ export function getUsageSpendSummary(): Promise { return invoke("get_usage_spend_summary"); } +export function getCodexWorkspacesSnapshot(options?: { + forceRefresh?: boolean; + historyDays?: number; +}): Promise { + return invoke("get_codex_workspaces_snapshot", { + forceRefresh: options?.forceRefresh ?? null, + historyDays: options?.historyDays ?? null, + }); +} + // ── Token account bridge ───────────────────────────────────────────── export function getTokenAccountProviders(): Promise { diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 12d634205a..f17c1e124a 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -282,6 +282,70 @@ export interface UsageSpendSummary { rows: UsageSpendRow[]; } +/** Codex local Workspaces snapshot (get_codex_workspaces_snapshot). */ +export type CodexWorkspacesSourceStatus = + | "complete" + | "catalogMissing" + | "catalogLocked" + | "catalogCorrupt" + | "catalogIncompatible"; + +export interface CodexWorkspacesUsageTotals { + inputTokens: number; + cachedInputTokens: number; + outputTokens: number; + totalTokens: number; +} + +export interface CodexWorkspacesCostEstimate { + knownUsd: number; + unknownTokens: number; +} + +export interface CodexWorkspacesDailyPoint { + day: string; + totalTokens: number; + cachedInputTokens: number; + estimatedCostUsd: number | null; +} + +export interface CodexWorkspacesSessionUsage { + id: string; + projectId: string; + displayTitle: string; + cwd: string | null; + startedAt: string | null; + latestActivity: string | null; + totals: CodexWorkspacesUsageTotals; + costEstimate: CodexWorkspacesCostEstimate; + topModel: string | null; +} + +export interface CodexWorkspacesProjectUsage { + id: string; + displayName: string; + path: string | null; + totals: CodexWorkspacesUsageTotals; + costEstimate: CodexWorkspacesCostEstimate; + sessionCount: number; + latestActivity: string | null; + topModel: string | null; + topSessions: CodexWorkspacesSessionUsage[]; +} + +export interface CodexLocalProjectUsageSnapshot { + updatedAt: string; + historyDays: number; + scopeSignature: string; + indexedFileCount: number; + skippedFileCount: number; + total: CodexWorkspacesUsageTotals; + projects: CodexWorkspacesProjectUsage[]; + daily: CodexWorkspacesDailyPoint[]; + sourceStatus: CodexWorkspacesSourceStatus; +} + + export interface BootstrapState { contractVersion: string; providers: ProviderCatalogEntry[]; @@ -313,6 +377,8 @@ export interface CostSnapshotBridge { resetsAt: string | null; formattedUsed: string; formattedLimit: string | null; + balance?: number | null; + formattedBalance?: string | null; } export interface PaceSnapshot { @@ -324,6 +390,15 @@ export interface PaceSnapshot { actualUsedPercent: number; } +export interface SessionEquivalentForecastSnapshot { + estimatedWindowsToExhaustWeekly: number; + windowsUntilReset: number; + availableWindowsUntilReset: number; + sampleCount: number; + weeklyResetsAt: string; + weeklyUsedPercent: number; +} + export interface ProviderUsageSnapshot { providerId: string; displayName: string; @@ -349,6 +424,7 @@ export interface ProviderUsageSnapshot { trayStatusLabel: string | null; fetchDurationMs?: number | null; wayfinderUsage?: WayfinderUsageSnapshot | null; + sessionEquivalentForecast?: SessionEquivalentForecastSnapshot | null; } export interface WayfinderRouteSummary { diff --git a/rust/src/cli/mod.rs b/rust/src/cli/mod.rs index 2f7253060b..c8f5498e88 100755 --- a/rust/src/cli/mod.rs +++ b/rust/src/cli/mod.rs @@ -19,6 +19,7 @@ pub mod serve; pub mod sessions; pub mod tty_runner; pub mod usage; +pub mod workspaces; use clap::{Parser, Subcommand}; @@ -93,6 +94,9 @@ pub enum Commands { /// List, enable, disable, or test external hooks Hooks(hooks::HooksArgs), + + /// List local Codex project/workspace usage + Workspaces(workspaces::WorkspacesArgs), } #[cfg(test)] diff --git a/rust/src/cli/workspaces.rs b/rust/src/cli/workspaces.rs new file mode 100644 index 0000000000..e270487f3c --- /dev/null +++ b/rust/src/cli/workspaces.rs @@ -0,0 +1,146 @@ +//! `codexbar workspaces` — list local Codex project usage. + +use clap::{Args, Subcommand}; +use serde::Serialize; + +use crate::codex_workspaces::{CodexWorkspacesIndex, ProgressPhase}; + +#[derive(Args, Debug, Default)] +pub struct WorkspacesArgs { + #[command(subcommand)] + pub command: Option, + + /// Shorthand for list --json + #[arg(long, global = true)] + pub json: bool, + + /// Pretty-print JSON + #[arg(long, global = true)] + pub pretty: bool, + + /// History window in days (default 30) + #[arg(short, long, default_value = "30", global = true)] + pub days: u32, + + /// Force a full re-index even when a cached snapshot exists + #[arg(long, global = true)] + pub force: bool, +} + +#[derive(Subcommand, Debug)] +pub enum WorkspacesCommand { + /// List projects (default) + List(WorkspacesListArgs), +} + +#[derive(Args, Debug, Default)] +pub struct WorkspacesListArgs {} + +#[derive(Debug, Serialize)] +#[serde(rename_all = "camelCase")] +struct ProjectRow { + name: String, + sessions: u32, + cost_usd: f64, + unknown_tokens: u64, + latest_activity: Option, + top_model: Option, + path: Option, +} + +pub async fn run(args: WorkspacesArgs) -> anyhow::Result<()> { + match args.command { + None | Some(WorkspacesCommand::List(_)) => {} + } + + let index = CodexWorkspacesIndex::new(args.days); + let snapshot = index.load_snapshot(args.force, |p| { + if matches!( + p.phase, + ProgressPhase::ScanningLogs | ProgressPhase::IndexingProjects + ) { + tracing::debug!( + phase = ?p.phase, + processed = ?p.processed_file_count, + total = ?p.total_file_count, + "workspaces index progress" + ); + } + })?; + + if args.json { + if args.pretty { + println!("{}", serde_json::to_string_pretty(&snapshot)?); + } else { + println!("{}", serde_json::to_string(&snapshot)?); + } + return Ok(()); + } + + println!( + "Codex workspaces ({} day{}, {} projects, status={:?})", + snapshot.history_days, + if snapshot.history_days == 1 { "" } else { "s" }, + snapshot.projects.len(), + snapshot.source_status + ); + println!( + "Indexed {} file(s), skipped {}.", + snapshot.indexed_file_count, snapshot.skipped_file_count + ); + println!(); + println!( + "{:<28} {:>8} {:>10} {:>20} MODEL", + "PROJECT", "SESSIONS", "COST", "LATEST" + ); + + if snapshot.projects.is_empty() { + println!("(no projects)"); + return Ok(()); + } + + for project in &snapshot.projects { + let cost = format!("${:.4}", project.cost_estimate.known_usd); + let latest = project + .latest_activity + .map(|t| t.format("%Y-%m-%d %H:%M").to_string()) + .unwrap_or_else(|| "-".into()); + let model = project.top_model.as_deref().unwrap_or("-"); + println!( + "{:<28} {:>8} {:>10} {:>20} {}", + truncate(&project.display_name, 28), + project.session_count, + cost, + latest, + model + ); + let _ = ProjectRow { + name: project.display_name.clone(), + sessions: project.session_count, + cost_usd: project.cost_estimate.known_usd, + unknown_tokens: project.cost_estimate.unknown_tokens, + latest_activity: project.latest_activity.map(|t| t.to_rfc3339()), + top_model: project.top_model.clone(), + path: project.path.clone(), + }; + } + + println!(); + println!( + "Total tokens: {} (in {} / out {} / cached {})", + snapshot.total.total_tokens, + snapshot.total.input_tokens, + snapshot.total.output_tokens, + snapshot.total.cached_input_tokens + ); + Ok(()) +} + +fn truncate(s: &str, max: usize) -> String { + if s.chars().count() <= max { + return s.to_string(); + } + let mut out: String = s.chars().take(max.saturating_sub(1)).collect(); + out.push('…'); + out +} diff --git a/rust/src/codex_workspaces/indexer.rs b/rust/src/codex_workspaces/indexer.rs new file mode 100644 index 0000000000..fe2d1baea4 --- /dev/null +++ b/rust/src/codex_workspaces/indexer.rs @@ -0,0 +1,887 @@ +//! Scan Codex rollout JSONL + read-only catalog → project usage snapshot. + +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Local, NaiveDate, TimeZone, Utc}; +use rusqlite::{Connection, OpenFlags}; + +use crate::agent_sessions::CodexRolloutFirstLineParser; +use crate::codex_costs::codex_period_start; +use crate::core::{CostUsageDayRange, CostUsagePricing, JsonlScanner, sha256_hex}; + +use super::sidecar::{SidecarError, WorkspaceUsageSidecar}; +use super::types::{ + CodexLocalProjectUsageSnapshot, CostEstimate, DailyPoint, Progress, ProgressPhase, + ProjectUsage, SessionUsage, SourceStatus, UsageTotals, +}; +use super::{CHATS_DISPLAY_NAME, CHATS_PROJECT_ID}; + +#[derive(Debug, thiserror::Error)] +pub enum IndexError { + #[error("codex home could not be resolved")] + HomeUnresolved, + #[error(transparent)] + Sidecar(#[from] SidecarError), + #[error("workspaces index io failed: {0}")] + Io(#[from] std::io::Error), +} + +/// Resolved local Codex sources for one home scope. +#[derive(Debug, Clone)] +pub struct CodexLocalDataScope { + pub identifier: String, + pub codex_home: PathBuf, + pub sessions_root: PathBuf, + pub archived_sessions_root: PathBuf, + pub state_database: PathBuf, +} + +impl CodexLocalDataScope { + /// `CODEX_HOME` → `CODEX_SQLITE_HOME` → `~/.codex`. + pub fn resolve() -> Option { + let home = non_empty_env("CODEX_HOME") + .or_else(|| non_empty_env("CODEX_SQLITE_HOME")) + .map(PathBuf::from) + .or_else(|| dirs::home_dir().map(|h| h.join(".codex")))?; + Some(Self::from_home(home)) + } + + pub fn from_home(codex_home: PathBuf) -> Self { + let codex_home = normalize_path(&codex_home); + let path_str = codex_home.to_string_lossy(); + Self { + identifier: format!("codex-workspaces:{}", scope_fingerprint(&path_str)), + sessions_root: codex_home.join("sessions"), + archived_sessions_root: codex_home.join("archived_sessions"), + state_database: codex_home.join("state_5.sqlite"), + codex_home, + } + } + + pub fn scope_signature(&self) -> &str { + &self.identifier + } +} + +/// Public workspaces index API. +#[derive(Debug, Clone)] +pub struct CodexWorkspacesIndex { + history_days: u32, + codex_home_override: Option, + sidecar_path_override: Option, +} + +impl Default for CodexWorkspacesIndex { + fn default() -> Self { + Self::new(30) + } +} + +impl CodexWorkspacesIndex { + pub fn new(history_days: u32) -> Self { + Self { + history_days: history_days.clamp(1, 365), + codex_home_override: None, + sidecar_path_override: None, + } + } + + pub fn with_codex_home(mut self, home: impl Into) -> Self { + self.codex_home_override = Some(home.into()); + self + } + + pub fn with_sidecar_path(mut self, path: impl Into) -> Self { + self.sidecar_path_override = Some(path.into()); + self + } + + pub fn history_days(&self) -> u32 { + self.history_days + } + + fn scope(&self) -> Result { + if let Some(home) = &self.codex_home_override { + return Ok(CodexLocalDataScope::from_home(home.clone())); + } + CodexLocalDataScope::resolve().ok_or(IndexError::HomeUnresolved) + } + + fn sidecar(&self) -> Result { + let path = if let Some(path) = &self.sidecar_path_override { + path.clone() + } else { + WorkspaceUsageSidecar::default_path().ok_or(SidecarError::PathUnavailable)? + }; + Ok(WorkspaceUsageSidecar::new(path)) + } + + pub fn load_cached_snapshot( + &self, + ) -> Result, IndexError> { + let scope = self.scope()?; + let sidecar = self.sidecar()?; + Ok(sidecar.load_latest_snapshot(scope.scope_signature(), self.history_days)?) + } + + pub fn clear_cached_snapshot(&self) -> Result<(), IndexError> { + let sidecar = self.sidecar()?; + sidecar.clear_snapshots()?; + Ok(()) + } + + pub fn load_snapshot( + &self, + force_refresh: bool, + mut progress: F, + ) -> Result + where + F: FnMut(Progress), + { + let scope = self.scope()?; + let sidecar = self.sidecar()?; + let source_status = read_catalog_status(&scope.state_database); + + if !force_refresh + && let Ok(Some(cached)) = + sidecar.load_latest_snapshot(scope.scope_signature(), self.history_days) + { + let mut snap = cached; + snap.source_status = source_status; + return Ok(snap); + } + + progress(Progress::phase(ProgressPhase::ScanningLogs)); + let today = Local::now().date_naive(); + let since = codex_period_start(today, self.history_days); + let range = CostUsageDayRange::new(since, today); + + let mut files = list_session_files(&scope.sessions_root, &range); + files.extend(list_session_files(&scope.archived_sessions_root, &range)); + files.sort(); + files.dedup(); + + let total_file_count = files.len() as u32; + progress(Progress { + phase: ProgressPhase::ScanningLogs, + processed_file_count: Some(0), + total_file_count: Some(total_file_count), + indexed_file_count: 0, + skipped_file_count: 0, + }); + + progress(Progress::phase(ProgressPhase::IndexingProjects)); + let mut session_buckets: HashMap = HashMap::new(); + let mut daily_acc: HashMap = HashMap::new(); + let mut indexed = 0u32; + let mut skipped = 0u32; + + for (idx, path) in files.iter().enumerate() { + match index_one_file(path, &range) { + Some(parsed) => { + merge_daily(&mut daily_acc, &parsed.day_models); + let entry = session_buckets + .entry(parsed.session_id.clone()) + .or_insert_with(|| SessionBucket::new(&parsed)); + entry.merge(parsed); + indexed = indexed.saturating_add(1); + } + None => { + skipped = skipped.saturating_add(1); + } + } + if idx == 0 || (idx + 1) % 25 == 0 || idx + 1 == files.len() { + progress(Progress { + phase: ProgressPhase::IndexingProjects, + processed_file_count: Some((idx + 1) as u32), + total_file_count: Some(total_file_count), + indexed_file_count: indexed, + skipped_file_count: skipped, + }); + } + } + + let mut projects = build_projects(&session_buckets); + projects.sort_by(|a, b| { + b.latest_activity + .cmp(&a.latest_activity) + .then_with(|| a.display_name.cmp(&b.display_name)) + }); + + let mut total = UsageTotals::default(); + for p in &projects { + total.add_assign(&p.totals); + } + + let mut daily: Vec = daily_acc + .into_iter() + .map(|(day, acc)| DailyPoint { + day, + total_tokens: acc.total_tokens, + cached_input_tokens: acc.cached_input_tokens, + estimated_cost_usd: if acc.unknown_tokens == 0 || acc.known_usd > 0.0 { + Some(acc.known_usd) + } else { + None + }, + }) + .collect(); + daily.sort_by(|a, b| a.day.cmp(&b.day)); + + let snapshot = CodexLocalProjectUsageSnapshot { + updated_at: Utc::now(), + history_days: self.history_days, + scope_signature: scope.scope_signature().to_string(), + indexed_file_count: indexed, + skipped_file_count: skipped, + total, + projects, + daily, + source_status, + }; + + progress(Progress::phase(ProgressPhase::Saving)); + sidecar.publish_snapshot(&snapshot)?; + Ok(snapshot) + } +} + +#[derive(Debug, Clone)] +struct ParsedFile { + session_id: String, + cwd: Option, + title: Option, + project_id: String, + project_display_name: String, + project_path: Option, + totals: UsageTotals, + cost: CostEstimate, + model_tokens: HashMap, + day_models: HashMap>, + started_at: Option>, + latest_activity: Option>, +} + +#[derive(Debug, Clone)] +struct SessionBucket { + id: String, + project_id: String, + project_display_name: String, + project_path: Option, + cwd: Option, + title: Option, + totals: UsageTotals, + cost: CostEstimate, + model_tokens: HashMap, + started_at: Option>, + latest_activity: Option>, +} + +impl SessionBucket { + fn new(parsed: &ParsedFile) -> Self { + Self { + id: parsed.session_id.clone(), + project_id: parsed.project_id.clone(), + project_display_name: parsed.project_display_name.clone(), + project_path: parsed.project_path.clone(), + cwd: parsed.cwd.clone(), + title: parsed.title.clone(), + totals: UsageTotals::default(), + cost: CostEstimate::default(), + model_tokens: HashMap::new(), + started_at: None, + latest_activity: None, + } + } + + fn merge(&mut self, parsed: ParsedFile) { + self.totals.add_assign(&parsed.totals); + self.cost.add_assign(&parsed.cost); + if self.cwd.is_none() { + self.cwd = parsed.cwd; + } + if self.title.is_none() { + self.title = parsed.title; + } + if self.project_path.is_none() { + self.project_path = parsed.project_path; + self.project_display_name = parsed.project_display_name; + self.project_id = parsed.project_id; + } + self.started_at = min_time(self.started_at, parsed.started_at); + self.latest_activity = max_time(self.latest_activity, parsed.latest_activity); + for (model, tokens) in parsed.model_tokens { + *self.model_tokens.entry(model).or_default() = self + .model_tokens + .get(&model) + .copied() + .unwrap_or(0) + .saturating_add(tokens); + } + } + + fn top_model(&self) -> Option { + self.model_tokens + .iter() + .max_by(|a, b| a.1.cmp(b.1).then_with(|| a.0.cmp(b.0))) + .map(|(m, _)| m.clone()) + } + + fn to_session_usage(&self) -> SessionUsage { + SessionUsage { + id: self.id.clone(), + project_id: self.project_id.clone(), + display_title: self + .title + .clone() + .filter(|t| !t.trim().is_empty()) + .unwrap_or_else(|| "Local Codex chat".to_string()), + cwd: self.cwd.clone(), + started_at: self.started_at, + latest_activity: self.latest_activity, + totals: self.totals, + cost_estimate: self.cost, + top_model: self.top_model(), + } + } +} + +#[derive(Default)] +struct DailyAcc { + total_tokens: u64, + cached_input_tokens: u64, + known_usd: f64, + unknown_tokens: u64, +} + +fn build_projects(sessions: &HashMap) -> Vec { + let mut by_project: HashMap> = HashMap::new(); + for session in sessions.values() { + by_project + .entry(session.project_id.clone()) + .or_default() + .push(session); + } + + by_project + .into_values() + .map(|mut buckets| { + buckets.sort_by(|a, b| { + b.latest_activity + .cmp(&a.latest_activity) + .then_with(|| a.id.cmp(&b.id)) + }); + let first = buckets[0]; + let mut totals = UsageTotals::default(); + let mut cost = CostEstimate::default(); + let mut model_tokens: HashMap = HashMap::new(); + let mut latest = None; + for b in &buckets { + totals.add_assign(&b.totals); + cost.add_assign(&b.cost); + latest = max_time(latest, b.latest_activity); + for (model, tokens) in &b.model_tokens { + *model_tokens.entry(model.clone()).or_default() += *tokens; + } + } + let top_model = model_tokens + .into_iter() + .max_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0))) + .map(|(m, _)| m); + let top_sessions = buckets + .iter() + .take(5) + .map(|b| b.to_session_usage()) + .collect(); + ProjectUsage { + id: first.project_id.clone(), + display_name: first.project_display_name.clone(), + path: first.project_path.clone(), + totals, + cost_estimate: cost, + session_count: buckets.len() as u32, + latest_activity: latest, + top_model, + top_sessions, + } + }) + .collect() +} + +fn index_one_file(path: &Path, range: &CostUsageDayRange) -> Option { + let first_line = CodexRolloutFirstLineParser::read_first_line(path)?; + let meta = CodexRolloutFirstLineParser::parse(&first_line); + let (session_id, cwd) = if let Some(meta) = meta { + (meta.session_id, meta.cwd) + } else { + let stem = path.file_stem()?.to_string_lossy().into_owned(); + let id = stem + .rsplit('-') + .next() + .filter(|s| s.len() >= 8) + .map(|s| s.to_string()) + .unwrap_or(stem); + (id, None) + }; + + let parsed = JsonlScanner::parse_codex_file(path, range, 0, None, None).ok()?; + if parsed.records.is_empty() && parsed.last_model.is_none() { + // Still index empty sessions with metadata when they fall in range by mtime/name. + // Skip only when there is truly nothing usable. + if cwd.is_none() && meta_day_out_of_range(path, range) { + return None; + } + } + + let identity = project_identity(cwd.as_deref()); + let mut totals = UsageTotals::default(); + let mut cost = CostEstimate::default(); + let mut model_tokens: HashMap = HashMap::new(); + let mut day_models: HashMap> = HashMap::new(); + + for record in &parsed.records { + let input = record.input.max(0) as u64; + let cached = (record.cached.max(0) as u64).min(input); + let output = record.output.max(0) as u64; + let part = UsageTotals::from_parts(input, cached, output); + totals.add_assign(&part); + + let model = CostUsagePricing::normalize_codex_model(&record.model); + let tokens = input.saturating_add(output); + *model_tokens.entry(model.clone()).or_default() += tokens; + let day_entry = day_models + .entry(record.day_key.clone()) + .or_default() + .entry(model.clone()) + .or_default(); + day_entry.0 = day_entry.0.saturating_add(input); + day_entry.1 = day_entry.1.saturating_add(cached); + day_entry.2 = day_entry.2.saturating_add(output); + + match CostUsagePricing::codex_cost_usd(&model, input, cached, output) { + Some(usd) => cost.known_usd += usd, + None => cost.unknown_tokens = cost.unknown_tokens.saturating_add(tokens), + } + } + + if model_tokens.is_empty() + && let Some(model) = parsed + .last_model + .as_deref() + .map(CostUsagePricing::normalize_codex_model) + { + model_tokens.insert(model, 0); + } + + let (started_at, latest_activity) = file_activity_bounds(path, &day_models); + + Some(ParsedFile { + session_id, + cwd, + title: None, + project_id: identity.0, + project_display_name: identity.1, + project_path: identity.2, + totals, + cost, + model_tokens, + day_models, + started_at, + latest_activity, + }) +} + +fn merge_daily( + daily: &mut HashMap, + day_models: &HashMap>, +) { + for (day, models) in day_models { + let acc = daily.entry(day.clone()).or_default(); + for (model, (input, cached, output)) in models { + let tokens = input.saturating_add(*output); + acc.total_tokens = acc.total_tokens.saturating_add(tokens); + acc.cached_input_tokens = acc.cached_input_tokens.saturating_add(*cached); + match CostUsagePricing::codex_cost_usd(model, *input, *cached, *output) { + Some(usd) => acc.known_usd += usd, + None => acc.unknown_tokens = acc.unknown_tokens.saturating_add(tokens), + } + } + } +} + +fn list_session_files(root: &Path, range: &CostUsageDayRange) -> Vec { + if !root.exists() { + return Vec::new(); + } + JsonlScanner::list_codex_session_files(root, &range.scan_since_key, &range.scan_until_key) +} + +fn read_catalog_status(db_path: &Path) -> SourceStatus { + if !db_path.exists() { + return SourceStatus::CatalogMissing; + } + let open = Connection::open_with_flags( + db_path, + OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ); + let conn = match open { + Ok(c) => c, + Err(e) => { + let msg = e.to_string().to_ascii_lowercase(); + return if msg.contains("locked") || msg.contains("busy") { + SourceStatus::CatalogLocked + } else if msg.contains("not a database") || msg.contains("corrupt") { + SourceStatus::CatalogCorrupt + } else { + SourceStatus::CatalogIncompatible + }; + } + }; + let _ = conn.busy_timeout(std::time::Duration::from_millis(250)); + let _ = conn.execute_batch("PRAGMA query_only = ON;"); + + let has_threads: Result = conn.query_row( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'threads' LIMIT 1", + [], + |_| Ok(true), + ); + match has_threads { + Ok(true) => {} + Ok(false) => return SourceStatus::CatalogIncompatible, + Err(e) => { + let msg = e.to_string().to_ascii_lowercase(); + return if msg.contains("locked") || msg.contains("busy") { + SourceStatus::CatalogLocked + } else if msg.contains("corrupt") || msg.contains("not a database") { + SourceStatus::CatalogCorrupt + } else { + SourceStatus::CatalogIncompatible + }; + } + } + + // Probe a cheap read so locked/corrupt states surface. + match conn.query_row("SELECT COUNT(*) FROM threads", [], |r| r.get::<_, i64>(0)) { + Ok(_) => SourceStatus::Complete, + Err(e) => { + let msg = e.to_string().to_ascii_lowercase(); + if msg.contains("locked") || msg.contains("busy") { + SourceStatus::CatalogLocked + } else if msg.contains("corrupt") { + SourceStatus::CatalogCorrupt + } else { + SourceStatus::CatalogIncompatible + } + } + } +} + +/// Project identity from cwd: `(id, display_name, path)`. +pub fn project_identity(cwd: Option<&str>) -> (String, String, Option) { + let Some(cwd) = cwd.map(str::trim).filter(|s| !s.is_empty()) else { + return ( + CHATS_PROJECT_ID.to_string(), + CHATS_DISPLAY_NAME.to_string(), + None, + ); + }; + let root = resolve_project_root(Path::new(cwd)); + let path = normalize_path(&root); + let path_str = path.to_string_lossy().replace('\\', "/"); + let id = format!("project-{}", &sha256_hex(path_str.as_bytes())[..16]); + let display = path + .file_name() + .map(|s| s.to_string_lossy().into_owned()) + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| path_str.clone()); + (id, display, Some(path_str)) +} + +fn resolve_project_root(cwd: &Path) -> PathBuf { + if !cwd.exists() { + return cwd.to_path_buf(); + } + let mut current = if cwd.is_file() { + cwd.parent().unwrap_or(cwd).to_path_buf() + } else { + cwd.to_path_buf() + }; + loop { + if current.join(".git").exists() { + return current; + } + match current.parent() { + Some(parent) if parent != current => current = parent.to_path_buf(), + _ => return cwd.to_path_buf(), + } + } +} + +fn scope_fingerprint(value: &str) -> String { + // FNV-1a 64-bit, matching upstream CodexLocalDataScope. + let mut hash: u64 = 14_695_981_039_346_656_037; + for byte in value.as_bytes() { + hash ^= u64::from(*byte); + hash = hash.wrapping_mul(1_099_511_628_211); + } + format!("{hash:x}") +} + +fn non_empty_env(key: &str) -> Option { + std::env::var(key).ok().and_then(|v| { + let t = v.trim(); + if t.is_empty() { + None + } else { + Some(t.to_string()) + } + }) +} + +fn normalize_path(path: &Path) -> PathBuf { + fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf()) +} + +fn min_time(a: Option>, b: Option>) -> Option> { + match (a, b) { + (Some(x), Some(y)) => Some(x.min(y)), + (Some(x), None) => Some(x), + (None, Some(y)) => Some(y), + (None, None) => None, + } +} + +fn max_time(a: Option>, b: Option>) -> Option> { + match (a, b) { + (Some(x), Some(y)) => Some(x.max(y)), + (Some(x), None) => Some(x), + (None, Some(y)) => Some(y), + (None, None) => None, + } +} + +fn file_activity_bounds( + path: &Path, + day_models: &HashMap>, +) -> (Option>, Option>) { + let mut earliest: Option = None; + let mut latest: Option = None; + for day in day_models.keys() { + if let Ok(d) = NaiveDate::parse_from_str(day, "%Y-%m-%d") { + earliest = Some(earliest.map_or(d, |e| e.min(d))); + latest = Some(latest.map_or(d, |e| e.max(d))); + } + } + let started = earliest + .and_then(|d| d.and_hms_opt(0, 0, 0)) + .map(|n| Utc.from_utc_datetime(&n)); + let mut end = latest + .and_then(|d| d.and_hms_opt(23, 59, 59)) + .map(|n| Utc.from_utc_datetime(&n)); + if end.is_none() + && let Ok(meta) = fs::metadata(path) + && let Ok(modified) = meta.modified() + { + end = Some(DateTime::::from(modified)); + } + (started.or(end), end) +} + +fn meta_day_out_of_range(path: &Path, range: &CostUsageDayRange) -> bool { + // Best-effort: parent folder name YYYY/MM/DD or mtime. + if let Some(day) = path + .parent() + .and_then(|p| p.file_name()) + .and_then(|d| d.to_str()) + .filter(|d| d.len() == 2) + { + // day folder — reconstruct from parents + let parts: Vec<_> = path + .components() + .filter_map(|c| c.as_os_str().to_str()) + .collect(); + if parts.len() >= 3 { + let y = parts[parts.len() - 3]; + let m = parts[parts.len() - 2]; + let d = day; + if let Ok(date) = NaiveDate::parse_from_str(&format!("{y}-{m}-{d}"), "%Y-%m-%d") { + let key = CostUsageDayRange::day_key(date); + return !CostUsageDayRange::is_in_range(&key, &range.since_key, &range.until_key); + } + } + } + false +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::codex_workspaces::types::SourceStatus; + use rusqlite::Connection; + use std::fs::File; + use std::io::Write; + use tempfile::TempDir; + + fn write_session( + dir: &Path, + day: &str, + name: &str, + cwd: &str, + model: &str, + input: i32, + output: i32, + ) { + // day = YYYY-MM-DD + let parts: Vec<_> = day.split('-').collect(); + let folder = dir.join(parts[0]).join(parts[1]).join(parts[2]); + fs::create_dir_all(&folder).unwrap(); + let path = folder.join(name); + let mut f = File::create(&path).unwrap(); + writeln!( + f, + r#"{{"timestamp":"{day}T10:00:00.000Z","type":"session_meta","payload":{{"session_id":"{id}","cwd":"{cwd}","originator":"codex_exec","source":"cli"}}}}"#, + id = name.trim_end_matches(".jsonl"), + cwd = cwd.replace('\\', "\\\\"), + ) + .unwrap(); + writeln!( + f, + r#"{{"timestamp":"{day}T10:00:01.000Z","type":"turn_context","payload":{{"model":"{model}"}}}}"# + ) + .unwrap(); + writeln!( + f, + r#"{{"timestamp":"{day}T10:00:02.000Z","type":"event_msg","payload":{{"type":"token_count","info":{{"last_token_usage":{{"input_tokens":{input},"cached_input_tokens":0,"output_tokens":{output}}}}}}}}}"# + ) + .unwrap(); + } + + #[test] + fn indexes_two_sessions_into_projects_and_daily() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("codex"); + let sessions = home.join("sessions"); + fs::create_dir_all(&sessions).unwrap(); + let day = Local::now().date_naive().format("%Y-%m-%d").to_string(); + + write_session( + &sessions, + &day, + "sess-aaa.jsonl", + &tmp.path().join("proj-a").to_string_lossy(), + "gpt-5", + 1000, + 500, + ); + write_session( + &sessions, + &day, + "sess-bbb.jsonl", + &tmp.path().join("proj-b").to_string_lossy(), + "gpt-5", + 2000, + 100, + ); + + // Ensure project roots exist so resolve doesn't collapse oddly. + fs::create_dir_all(tmp.path().join("proj-a")).unwrap(); + fs::create_dir_all(tmp.path().join("proj-b")).unwrap(); + + let sidecar = tmp.path().join("sidecar.sqlite"); + let index = CodexWorkspacesIndex::new(30) + .with_codex_home(&home) + .with_sidecar_path(&sidecar); + + let snap = index.load_snapshot(true, |_| {}).expect("snapshot"); + assert_eq!(snap.indexed_file_count, 2); + assert_eq!(snap.projects.len(), 2); + assert!(!snap.daily.is_empty()); + assert_eq!(snap.source_status, SourceStatus::CatalogMissing); + let total_sessions: u32 = snap.projects.iter().map(|p| p.session_count).sum(); + assert_eq!(total_sessions, 2); + assert!(snap.total.input_tokens >= 3000); + assert!(snap.total.output_tokens >= 600); + + // Cached load works. + let cached = index.load_cached_snapshot().unwrap().expect("cached"); + assert_eq!(cached.indexed_file_count, 2); + } + + #[test] + fn foreign_user_version_is_rejected() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("bad.sqlite"); + { + let conn = Connection::open(&path).unwrap(); + conn.pragma_update(None, "user_version", 99).unwrap(); + } + let home = tmp.path().join("codex"); + fs::create_dir_all(home.join("sessions")).unwrap(); + let index = CodexWorkspacesIndex::new(7) + .with_codex_home(home) + .with_sidecar_path(path); + let err = index.load_snapshot(true, |_| {}).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("incompatible") || msg.contains("user_version"), + "unexpected error: {msg}" + ); + } + + #[test] + fn missing_catalog_reports_catalog_missing() { + let tmp = TempDir::new().unwrap(); + let home = tmp.path().join("codex"); + fs::create_dir_all(home.join("sessions")).unwrap(); + let index = CodexWorkspacesIndex::new(7) + .with_codex_home(&home) + .with_sidecar_path(tmp.path().join("side.sqlite")); + let snap = index.load_snapshot(true, |_| {}).unwrap(); + assert_eq!(snap.source_status, SourceStatus::CatalogMissing); + } + + #[test] + fn privacy_redaction_strips_paths_and_titles() { + let mut snap = CodexLocalProjectUsageSnapshot { + updated_at: Utc::now(), + history_days: 30, + scope_signature: "x".into(), + indexed_file_count: 1, + skipped_file_count: 0, + total: UsageTotals::from_parts(10, 0, 5), + projects: vec![ProjectUsage { + id: "project-abc".into(), + display_name: "secret-repo".into(), + path: Some("/Users/me/secret-repo".into()), + totals: UsageTotals::from_parts(10, 0, 5), + cost_estimate: CostEstimate::default(), + session_count: 1, + latest_activity: None, + top_model: Some("gpt-5".into()), + top_sessions: vec![SessionUsage { + id: "s1".into(), + project_id: "project-abc".into(), + display_title: "do not leak".into(), + cwd: Some("/Users/me/secret-repo".into()), + started_at: None, + latest_activity: None, + totals: UsageTotals::from_parts(10, 0, 5), + cost_estimate: CostEstimate::default(), + top_model: Some("gpt-5".into()), + }], + }], + daily: vec![], + source_status: SourceStatus::Complete, + }; + snap.redact_for_privacy(); + assert_eq!(snap.projects[0].display_name, "Workspace"); + assert!(snap.projects[0].path.is_none()); + assert_eq!( + snap.projects[0].top_sessions[0].display_title, + "Local Codex chat" + ); + assert!(snap.projects[0].top_sessions[0].cwd.is_none()); + } +} diff --git a/rust/src/codex_workspaces/mod.rs b/rust/src/codex_workspaces/mod.rs new file mode 100644 index 0000000000..1f5115da8d --- /dev/null +++ b/rust/src/codex_workspaces/mod.rs @@ -0,0 +1,18 @@ +//! Codex local Workspaces indexing foundation. +//! +//! Attributes local Codex rollout usage to projects, sessions, models, and days. +//! Codex-only; never writes the Codex catalog database. + +mod indexer; +mod sidecar; +mod types; + +pub use indexer::{CodexLocalDataScope, CodexWorkspacesIndex, IndexError, project_identity}; +pub use sidecar::{PAYLOAD_FORMAT_VERSION, SCHEMA_VERSION, SidecarError, WorkspaceUsageSidecar}; +pub use types::{ + CodexLocalProjectUsageSnapshot, CostEstimate, DailyPoint, Progress, ProgressPhase, + ProjectUsage, SessionUsage, SourceStatus, UsageTotals, +}; + +pub const CHATS_PROJECT_ID: &str = "chats"; +pub const CHATS_DISPLAY_NAME: &str = "Chats"; diff --git a/rust/src/codex_workspaces/sidecar.rs b/rust/src/codex_workspaces/sidecar.rs new file mode 100644 index 0000000000..732e7e435a --- /dev/null +++ b/rust/src/codex_workspaces/sidecar.rs @@ -0,0 +1,190 @@ +//! CodexBar-owned SQLite sidecar for Workspaces snapshots. +//! +//! Never attaches to or writes Codex's `state_5.sqlite`. Schema version 5 and +//! payload format 3 match upstream; foreign `user_version` values are refused. + +use std::fs; +use std::path::{Path, PathBuf}; + +use rusqlite::{Connection, OpenFlags, OptionalExtension, params}; + +use super::types::CodexLocalProjectUsageSnapshot; + +pub const SCHEMA_VERSION: i32 = 5; +pub const PAYLOAD_FORMAT_VERSION: i32 = 3; + +#[derive(Debug, thiserror::Error)] +pub enum SidecarError { + #[error("workspaces sidecar path unavailable")] + PathUnavailable, + #[error("workspaces sidecar open failed: {0}")] + Open(String), + #[error( + "workspaces sidecar schema incompatible (user_version={found}, expected {SCHEMA_VERSION})" + )] + Incompatible { found: i32 }, + #[error("workspaces sidecar error: {0}")] + Sqlite(#[from] rusqlite::Error), + #[error("workspaces sidecar encode/decode failed: {0}")] + Serde(#[from] serde_json::Error), + #[error("workspaces sidecar io failed: {0}")] + Io(#[from] std::io::Error), +} + +/// Persistence for complete workspaces snapshots. +#[derive(Debug, Clone)] +pub struct WorkspaceUsageSidecar { + path: PathBuf, +} + +impl WorkspaceUsageSidecar { + pub fn new(path: PathBuf) -> Self { + Self { path } + } + + /// Default: `%LOCALAPPDATA%\CodexBar\local-usage\codex-workspaces-v1.sqlite`. + pub fn default_path() -> Option { + dirs::data_local_dir().map(|root| { + root.join("CodexBar") + .join("local-usage") + .join("codex-workspaces-v1.sqlite") + }) + } + + pub fn path(&self) -> &Path { + &self.path + } + + pub fn load_latest_snapshot( + &self, + scope_signature: &str, + history_days: u32, + ) -> Result, SidecarError> { + if !self.path.exists() { + return Ok(None); + } + let conn = self.open(true)?; + let row: Option<(Vec, i32)> = conn + .query_row( + "SELECT payload, payload_format_version + FROM snapshot_payloads + WHERE scope_signature = ?1 + AND history_days = ?2 + AND is_complete = 1 + ORDER BY updated_at_ms DESC + LIMIT 1", + params![scope_signature, history_days as i64], + |r| Ok((r.get(0)?, r.get(1)?)), + ) + .optional()?; + let Some((payload, format_version)) = row else { + return Ok(None); + }; + if format_version != PAYLOAD_FORMAT_VERSION { + return Ok(None); + } + let snapshot = serde_json::from_slice(&payload)?; + Ok(Some(snapshot)) + } + + /// Atomically replace the complete snapshot for `(scope, history_days)`. + pub fn publish_snapshot( + &self, + snapshot: &CodexLocalProjectUsageSnapshot, + ) -> Result<(), SidecarError> { + if let Some(parent) = self.path.parent() { + fs::create_dir_all(parent)?; + } + let conn = self.open(false)?; + let payload = serde_json::to_vec(snapshot)?; + let updated_at_ms = snapshot.updated_at.timestamp_millis(); + let tx = conn.unchecked_transaction()?; + tx.execute( + "INSERT INTO snapshot_payloads ( + scope_signature, history_days, updated_at_ms, is_complete, + payload_format_version, payload + ) VALUES (?1, ?2, ?3, 1, ?4, ?5) + ON CONFLICT(scope_signature, history_days) DO UPDATE SET + updated_at_ms = excluded.updated_at_ms, + is_complete = 1, + payload_format_version = excluded.payload_format_version, + payload = excluded.payload", + params![ + snapshot.scope_signature, + snapshot.history_days as i64, + updated_at_ms, + PAYLOAD_FORMAT_VERSION, + payload, + ], + )?; + tx.commit()?; + Ok(()) + } + + pub fn clear_snapshots(&self) -> Result<(), SidecarError> { + if !self.path.exists() { + return Ok(()); + } + let conn = self.open(false)?; + conn.execute("DELETE FROM snapshot_payloads", [])?; + Ok(()) + } + + fn open(&self, read_only: bool) -> Result { + if read_only { + let conn = Connection::open_with_flags( + &self.path, + OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) + .map_err(|e| SidecarError::Open(e.to_string()))?; + conn.busy_timeout(std::time::Duration::from_millis(250))?; + Self::assert_compatible(&conn)?; + return Ok(conn); + } + + if let Some(parent) = self.path.parent() { + fs::create_dir_all(parent)?; + } + let conn = Connection::open_with_flags( + &self.path, + OpenFlags::SQLITE_OPEN_READ_WRITE + | OpenFlags::SQLITE_OPEN_CREATE + | OpenFlags::SQLITE_OPEN_NO_MUTEX, + ) + .map_err(|e| SidecarError::Open(e.to_string()))?; + conn.busy_timeout(std::time::Duration::from_millis(250))?; + Self::ensure_schema(&conn)?; + Ok(conn) + } + + fn assert_compatible(conn: &Connection) -> Result<(), SidecarError> { + let version: i32 = conn.pragma_query_value(None, "user_version", |r| r.get(0))?; + if version == 0 || version == SCHEMA_VERSION { + Ok(()) + } else { + Err(SidecarError::Incompatible { found: version }) + } + } + + fn ensure_schema(conn: &Connection) -> Result<(), SidecarError> { + let version: i32 = conn.pragma_query_value(None, "user_version", |r| r.get(0))?; + if version != 0 && version != SCHEMA_VERSION { + return Err(SidecarError::Incompatible { found: version }); + } + if version == 0 { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS snapshot_payloads ( + scope_signature TEXT NOT NULL, + history_days INTEGER NOT NULL, + updated_at_ms INTEGER NOT NULL, + is_complete INTEGER NOT NULL, + payload_format_version INTEGER NOT NULL, + payload BLOB NOT NULL, + PRIMARY KEY (scope_signature, history_days) + ); + PRAGMA user_version = 5;", + )?; + } + Ok(()) + } +} diff --git a/rust/src/codex_workspaces/types.rs b/rust/src/codex_workspaces/types.rs new file mode 100644 index 0000000000..f2944d7133 --- /dev/null +++ b/rust/src/codex_workspaces/types.rs @@ -0,0 +1,182 @@ +//! Snapshot DTOs for Codex local Workspaces indexing. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +/// Progress phases emitted while building a workspaces snapshot. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum ProgressPhase { + ScanningLogs, + IndexingProjects, + Saving, +} + +/// Progress callback payload. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct Progress { + pub phase: ProgressPhase, + pub processed_file_count: Option, + pub total_file_count: Option, + pub indexed_file_count: u32, + pub skipped_file_count: u32, +} + +impl Progress { + pub fn phase(phase: ProgressPhase) -> Self { + Self { + phase, + processed_file_count: None, + total_file_count: None, + indexed_file_count: 0, + skipped_file_count: 0, + } + } +} + +/// Completeness of the read-only Codex thread catalog (`state_5.sqlite`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum SourceStatus { + Complete, + CatalogMissing, + CatalogLocked, + CatalogCorrupt, + CatalogIncompatible, +} + +impl SourceStatus { + pub fn is_partial(self) -> bool { + self != Self::Complete + } +} + +/// Known vs unknown cost split. Unknown pricing never becomes a synthetic zero. +#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CostEstimate { + pub known_usd: f64, + pub unknown_tokens: u64, +} + +impl Default for CostEstimate { + fn default() -> Self { + Self { + known_usd: 0.0, + unknown_tokens: 0, + } + } +} + +impl CostEstimate { + pub fn add_assign(&mut self, other: &Self) { + self.known_usd += other.known_usd; + self.unknown_tokens = self.unknown_tokens.saturating_add(other.unknown_tokens); + } +} + +/// Aggregated token totals. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "camelCase")] +pub struct UsageTotals { + pub input_tokens: u64, + pub cached_input_tokens: u64, + pub output_tokens: u64, + pub total_tokens: u64, +} + +impl UsageTotals { + pub fn add_assign(&mut self, other: &Self) { + self.input_tokens = self.input_tokens.saturating_add(other.input_tokens); + self.cached_input_tokens = self + .cached_input_tokens + .saturating_add(other.cached_input_tokens); + self.output_tokens = self.output_tokens.saturating_add(other.output_tokens); + self.total_tokens = self.total_tokens.saturating_add(other.total_tokens); + } + + pub fn from_parts(input: u64, cached: u64, output: u64) -> Self { + let cached = cached.min(input); + Self { + input_tokens: input, + cached_input_tokens: cached, + output_tokens: output, + total_tokens: input.saturating_add(output), + } + } +} + +/// One calendar-day aggregate point. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DailyPoint { + pub day: String, + pub total_tokens: u64, + pub cached_input_tokens: u64, + pub estimated_cost_usd: Option, +} + +/// Per-session usage row (top sessions only on projects). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionUsage { + pub id: String, + pub project_id: String, + pub display_title: String, + pub cwd: Option, + pub started_at: Option>, + pub latest_activity: Option>, + pub totals: UsageTotals, + pub cost_estimate: CostEstimate, + pub top_model: Option, +} + +/// Per-project usage aggregate. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProjectUsage { + pub id: String, + pub display_name: String, + pub path: Option, + pub totals: UsageTotals, + pub cost_estimate: CostEstimate, + pub session_count: u32, + pub latest_activity: Option>, + pub top_model: Option, + pub top_sessions: Vec, +} + +/// Complete workspaces snapshot published to the sidecar and presentation layer. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CodexLocalProjectUsageSnapshot { + pub updated_at: DateTime, + pub history_days: u32, + pub scope_signature: String, + pub indexed_file_count: u32, + pub skipped_file_count: u32, + pub total: UsageTotals, + pub projects: Vec, + pub daily: Vec, + pub source_status: SourceStatus, +} + +impl CodexLocalProjectUsageSnapshot { + /// Strip paths/titles for presentation when hide-personal-info is on. + /// Does not rewrite the sidecar. + pub fn redact_for_privacy(&mut self) { + for project in &mut self.projects { + if project.id == crate::codex_workspaces::CHATS_PROJECT_ID { + project.display_name = crate::codex_workspaces::CHATS_DISPLAY_NAME.to_string(); + } else { + project.display_name = "Workspace".to_string(); + } + project.path = None; + for session in &mut project.top_sessions { + session.display_title = "Local Codex chat".to_string(); + session.cwd = None; + } + } + } +} diff --git a/rust/src/lib.rs b/rust/src/lib.rs index f3ec8157bd..35e06c8475 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -6,6 +6,7 @@ pub mod agent_sessions; pub mod browser; pub mod cli; +pub mod codex_workspaces; pub mod core; pub mod cost_scanner; pub mod host; diff --git a/rust/src/main.rs b/rust/src/main.rs index 83ce1b1fb9..c028ea5850 100644 --- a/rust/src/main.rs +++ b/rust/src/main.rs @@ -103,6 +103,7 @@ fn dispatch_command(rt: &Runtime, command: Option) -> i32 { Some(Commands::Account(args)) => run_unexpected(rt, cli::account::run(args)), Some(Commands::Config(args)) => run_unexpected(rt, cli::config::run(args)), Some(Commands::Hooks(args)) => run_unexpected(rt, cli::hooks::run(args)), + Some(Commands::Workspaces(args)) => run_categorized(rt, cli::workspaces::run(args)), None => missing_subcommand(), } } From 175c45e7da24ae0b9eb1ec1138ac3fa7266fbbdb Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:35 +0700 Subject: [PATCH 09/15] Prioritize exhausted windows and keep idle session estimates --- .../src-tauri/src/tray_bridge.rs | 170 +++++++++++++----- 1 file changed, 127 insertions(+), 43 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index db79fbc505..41957c508f 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -680,43 +680,64 @@ fn automatic_metric_percent( provider: Option, ) -> Option { match provider { - Some(ProviderId::Cursor) => max_metric_percent([ - Some(snapshot.primary.used_percent), - snapshot.secondary.as_ref().map(|w| w.used_percent), - snapshot.tertiary.as_ref().map(|w| w.used_percent), - ]), - Some(ProviderId::Zai) => max_metric_percent([ - Some(snapshot.primary.used_percent), - snapshot.tertiary.as_ref().map(|w| w.used_percent), - None, - ]) - .or_else(|| snapshot.secondary.as_ref().map(|w| w.used_percent)), - Some(ProviderId::Factory) | Some(ProviderId::Kimi) => snapshot - .secondary - .as_ref() - .filter(|w| !w.is_informational) - .map(|w| w.used_percent) - .or_else(|| { - (!snapshot.primary.is_informational).then_some(snapshot.primary.used_percent) - }), - Some(ProviderId::Copilot) => max_metric_percent([ - (!snapshot.primary.is_informational).then_some(snapshot.primary.used_percent), - snapshot - .secondary - .as_ref() - .filter(|w| !w.is_informational) - .map(|w| w.used_percent), - extra_rate_window_percent(snapshot), - ]), - // Informational primary (e.g. Claude null five_hour) is not a real quota — - // prefer secondary / non-informational signals for every provider. - _ if snapshot.primary.is_informational => snapshot - .secondary - .as_ref() - .filter(|w| !w.is_informational) - .map(|w| w.used_percent), - _ => Some(snapshot.primary.used_percent), + // When a model carve-out is exhausted but account weekly still has + // remaining, prefer weekly for Automatic display (upstream 0.46). + Some(ProviderId::Claude) => claude_automatic_metric_percent(snapshot), + // Highest used_percent across windows so exhausted (≥100%) surfaces first + // (upstream #2352). Provider-specific overrides above stay intact. + _ => highest_window_metric_percent(snapshot), + } +} + +fn claude_automatic_metric_percent( + snapshot: &crate::commands::ProviderUsageSnapshot, +) -> Option { + let weekly = snapshot.secondary.as_ref().filter(|w| !w.is_informational); + let model = snapshot.model_specific.as_ref(); + + if let (Some(model), Some(weekly)) = (model, weekly) { + let model_exhausted = model.is_exhausted || model.used_percent >= 100.0; + let weekly_has_remaining = !weekly.is_exhausted && weekly.used_percent < 100.0; + if model_exhausted && weekly_has_remaining { + return Some(weekly.used_percent); + } + } + + if snapshot.primary.is_informational { + return weekly.map(|w| w.used_percent); } + + // Fall through to highest-window among Claude lanes when no carve-out rule hits. + highest_window_metric_percent(snapshot) +} + +/// Automatic tray metric: highest used_percent across non-informational windows. +fn highest_window_metric_percent(snapshot: &crate::commands::ProviderUsageSnapshot) -> Option { + let mut values = Vec::with_capacity(4 + snapshot.extra_rate_windows.len()); + if !snapshot.primary.is_informational { + values.push(snapshot.primary.used_percent); + } + if let Some(w) = snapshot.secondary.as_ref().filter(|w| !w.is_informational) { + values.push(w.used_percent); + } + if let Some(w) = snapshot + .model_specific + .as_ref() + .filter(|w| !w.is_informational) + { + values.push(w.used_percent); + } + if let Some(w) = snapshot.tertiary.as_ref().filter(|w| !w.is_informational) { + values.push(w.used_percent); + } + for extra in &snapshot.extra_rate_windows { + if !extra.window.is_informational { + values.push(extra.window.used_percent); + } + } + values + .into_iter() + .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) } fn average_metric_percent(snapshot: &crate::commands::ProviderUsageSnapshot) -> Option { @@ -744,13 +765,6 @@ fn extra_rate_window_percent(snapshot: &crate::commands::ProviderUsageSnapshot) .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) } -fn max_metric_percent(values: [Option; N]) -> Option { - values - .into_iter() - .flatten() - .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) -} - /// Build a compact multi-line tooltip string from provider snapshots. fn build_tooltip( snapshots: &[crate::commands::ProviderUsageSnapshot], @@ -1136,6 +1150,8 @@ mod tests { resets_at: None, formatted_used: format!("${used:.2}"), formatted_limit: Some(format!("${limit:.2}")), + balance: None, + formatted_balance: None, }), plan_name: None, account_email: None, @@ -1147,6 +1163,7 @@ mod tests { tray_status_label: None, fetch_duration_ms: None, wayfinder_usage: None, + session_equivalent_forecast: None, } } @@ -1455,4 +1472,71 @@ mod tests { None ); } + + #[test] + fn claude_automatic_prefers_weekly_when_model_exhausted() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 40.0, Some(22.0), None, None); + snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { + used_percent: 100.0, + remaining_percent: 0.0, + window_minutes: Some(10080), + resets_at: None, + reset_description: None, + is_exhausted: true, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 22.0); + + // Explicit model override is untouched. + let mut overridden = settings.clone(); + overridden.set_provider_metric(ProviderId::Claude, MetricPreference::Model); + let (primary, _) = selected_tray_percents(&snapshot, &overridden); + assert_eq!(primary, 100.0); + } + + #[test] + fn automatic_prefers_exhausted_weekly_over_low_session() { + let settings = Settings::default(); + let snapshot = fake_snapshot_with("codex", "Codex", 20.0, Some(100.0), None, None); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 100.0); + + // Explicit session override still wins. + let mut overridden = settings.clone(); + overridden.set_provider_metric(ProviderId::Codex, MetricPreference::Session); + let (primary, _) = selected_tray_percents(&snapshot, &overridden); + assert_eq!(primary, 20.0); + } + + #[test] + fn automatic_picks_highest_among_model_and_extra_windows() { + let settings = Settings::default(); + let mut snapshot = + fake_snapshot_with("gemini", "Gemini", 10.0, Some(30.0), Some(40.0), None); + snapshot.model_specific = Some(crate::commands::RateWindowSnapshot { + used_percent: 55.0, + remaining_percent: 45.0, + window_minutes: None, + resets_at: None, + reset_description: None, + is_exhausted: false, + is_informational: false, + reserve_percent: None, + reserve_description: None, + reserve_will_last_to_reset: false, + reserve_eta_seconds: None, + }); + snapshot.extra_rate_windows.push(fake_extra_window(90.0)); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 90.0); + } } From b9deb10c0dfb9d8a33d0a75a4d5625859658d54d Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:18:38 +0700 Subject: [PATCH 10/15] Prepare Win-CodexBar 0.46.0 --- CHANGELOG.md | 28 ++++++++++++++++++++ Cargo.lock | 4 +-- apps/desktop-tauri/package.json | 2 +- apps/desktop-tauri/src-tauri/Cargo.toml | 2 +- apps/desktop-tauri/src-tauri/tauri.conf.json | 2 +- rust/Cargo.toml | 2 +- version.env | 2 +- 7 files changed, 35 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2a7bd96ea7..66a1fcc83d 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,33 @@ # Changelog +## [Windows] 0.46.0 - 2026-07-30 + +Windows port of upstream CodexBar **0.45.2 → 0.46.0**. macOS-only shell items (WidgetKit, Sparkle, AppKit menu layout, Homebrew, Safari cookie APIs) remain deferred. + +### Added +- Providers: Qwen Cloud (Individual Token Plans with 5-hour and weekly rolling windows) and ZoomMate (credit status with host failover, cURL/cookie auth, bearer mint). +- Alibaba Token Plan: Personal/Solo variants for mainland (Bailian) and international (Model Studio) accounts via `alibaba_token_plan_region` setting. +- Claude: prepaid credit balance in cost surfaces from cached/manual web sessions. +- Claude: setting to hide the Daily Routines row (`claude_daily_routines_usage_visible`, default on). +- Menu: fractional session-quota estimates on the weekly row ("Estimated: N session quotas left") when session+weekly history qualifies; optional `weekly_progress_work_days`. +- Codex: local Workspaces indexing foundation — per-project/session/model usage attribution with a local sidecar index; new `codexbar workspaces` CLI and desktop snapshot bridge. +- CLI: `config dump` redacts stored credentials by default; `--show-secrets` restores raw output. + +### Changed +- Codex local cost scans use the disk cache: unchanged files skipped by mtime/size, grown logs resumed mid-file, 256 KiB line bound — repeat scans are incremental. +- Automatic tray metric surfaces the highest-used (exhausted) window across providers, preserving per-provider overrides. +- CLI alias `qwen` now resolves to Qwen Cloud (use `alibaba` for the Coding Plan). + +### Fixed +- Claude: model-scoped weekly rows above Daily Routines; automatic metric prefers account Weekly over exhausted model carve-outs; learned full-session estimate stays visible while the session window is idle. +- Amp: subscription plans (e.g. Megawatt) parse into Other/Orb percentage windows instead of a misleading cookie error. +- Grok: explicit cookie refresh with validated session caching for background reuse. +- Chutes: quota counts render as detail text instead of being misread as reset schedules. +- LLMProxy: skip already-elapsed reset times when picking the next reset. +- Ollama: reuse validated browser sessions across refreshes. + +--- + ## [Windows] 0.45.3 - 2026-07-29 ### Added diff --git a/Cargo.lock b/Cargo.lock index 50748c42e0..a6d9823729 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -655,7 +655,7 @@ checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" [[package]] name = "codexbar" -version = "0.45.3" +version = "0.46.0" dependencies = [ "aes-gcm", "anyhow", @@ -695,7 +695,7 @@ dependencies = [ [[package]] name = "codexbar-desktop-tauri" -version = "0.45.3" +version = "0.46.0" dependencies = [ "chrono", "codexbar", diff --git a/apps/desktop-tauri/package.json b/apps/desktop-tauri/package.json index 0a9fcf4b98..02be55aff8 100644 --- a/apps/desktop-tauri/package.json +++ b/apps/desktop-tauri/package.json @@ -1,7 +1,7 @@ { "name": "desktop-tauri", "private": true, - "version": "0.45.3", + "version": "0.46.0", "packageManager": "pnpm@10.18.1", "type": "module", "scripts": { diff --git a/apps/desktop-tauri/src-tauri/Cargo.toml b/apps/desktop-tauri/src-tauri/Cargo.toml index 9a16d92133..02030cfc8c 100644 --- a/apps/desktop-tauri/src-tauri/Cargo.toml +++ b/apps/desktop-tauri/src-tauri/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "codexbar-desktop-tauri" -version = "0.45.3" +version = "0.46.0" edition = "2024" publish = false diff --git a/apps/desktop-tauri/src-tauri/tauri.conf.json b/apps/desktop-tauri/src-tauri/tauri.conf.json index 13be29665e..68ccd3af6f 100644 --- a/apps/desktop-tauri/src-tauri/tauri.conf.json +++ b/apps/desktop-tauri/src-tauri/tauri.conf.json @@ -1,7 +1,7 @@ { "$schema": "https://schema.tauri.app/config/2", "productName": "CodexBar Desktop", - "version": "0.45.3", + "version": "0.46.0", "identifier": "com.codexbar.desktop", "build": { "beforeDevCommand": "pnpm run dev", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 73191a01aa..1b5299c47a 100755 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "codexbar" -version = "0.45.3" +version = "0.46.0" edition = "2024" authors = ["CodexBar Contributors"] description = "Windows and WSL system tray app for monitoring AI provider usage limits" diff --git a/version.env b/version.env index 42ce0dbfe3..b2b58d53d9 100755 --- a/version.env +++ b/version.env @@ -1,2 +1,2 @@ -MARKETING_VERSION=0.45.3 +MARKETING_VERSION=0.46.0 BUILD_NUMBER=88 From cbd1c00670a88bc593deb04bde356d5ac5dc4590 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 10:28:24 +0700 Subject: [PATCH 11/15] Wire 0.46 settings and balance through the desktop bridge --- .../src-tauri/src/commands/bridge.rs | 6 ++ .../src/commands/provider_settings.rs | 9 +++ .../src-tauri/src/commands/settings.rs | 19 +++++ .../src-tauri/src/commands/tests.rs | 17 ++++ apps/desktop-tauri/src/App.test.tsx | 3 + .../src/components/MenuCard.test.tsx | 60 +++++++++++++++ .../desktop-tauri/src/components/MenuCard.tsx | 77 +++++++++++++------ .../src/floatbar/FloatBar.test.tsx | 3 + .../src/floatbar/SettingsSection.test.tsx | 3 + apps/desktop-tauri/src/i18n/keys.ts | 3 + .../src/surfaces/PopOutPanel.test.tsx | 3 + .../src/surfaces/TrayPanel.test.tsx | 3 + .../sections/credentials/ClaudeCreds.tsx | 60 ++++++++++++--- .../surfaces/settings/tabs/AboutTab.test.tsx | 3 + .../settings/tabs/GeneralTab.test.tsx | 3 + apps/desktop-tauri/src/types/bridge.test.ts | 3 + apps/desktop-tauri/src/types/bridge.ts | 9 +++ rust/src/locale.rs | 3 + rust/src/locale/en-US.ftl | 3 + rust/src/locale/es-MX.ftl | 3 + rust/src/locale/ja-JP.ftl | 3 + rust/src/locale/ko-KR.ftl | 3 + rust/src/locale/tests.rs | 1 + rust/src/locale/zh-CN.ftl | 3 + rust/src/locale/zh-TW.ftl | 3 + rust/src/providers/alibabatokenplan/region.rs | 11 +++ rust/src/settings.rs | 6 +- 27 files changed, 288 insertions(+), 35 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 348cbf53e7..c1734c38e8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -512,6 +512,9 @@ pub struct SettingsSnapshot { float_bar_show_reset_inline: bool, float_bar_show_cost: bool, promote_tray_icon: bool, + claude_daily_routines_usage_visible: bool, + alibaba_token_plan_region: String, + weekly_progress_work_days: Option, } #[tauri::command] @@ -612,6 +615,9 @@ impl From for SettingsSnapshot { float_bar_show_reset_inline: settings.float_bar_show_reset_inline, float_bar_show_cost: settings.float_bar_show_cost, promote_tray_icon: settings.promote_tray_icon, + claude_daily_routines_usage_visible: settings.claude_daily_routines_usage_visible, + alibaba_token_plan_region: settings.alibaba_token_plan_region, + weekly_progress_work_days: settings.weekly_progress_work_days, } } } 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 3bc04a8de9..1e44f3346c 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -127,6 +127,7 @@ fn region_provider(provider_id: &str) -> Option { use codexbar::core::ProviderId; Some(match provider_id { "alibaba" => ProviderId::Alibaba, + "alibabatokenplan" => ProviderId::AlibabaTokenPlan, "zai" => ProviderId::Zai, "minimax" => ProviderId::MiniMax, _ => return None, @@ -586,6 +587,14 @@ pub fn region_options_for(provider_id: &str) -> Vec { .to_string(), }, ], + "alibabatokenplan" => codexbar::providers::AlibabaTokenPlanRegion::ALL + .iter() + .copied() + .map(|region| RegionOption { + value: region.as_str().to_string(), + label: region.display_name().to_string(), + }) + .collect(), _ => Vec::new(), } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index 0835364f62..070059c67b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -65,11 +65,17 @@ pub struct SettingsUpdate { pub float_bar_show_reset_inline: Option, pub float_bar_show_cost: Option, pub promote_tray_icon: Option, + pub claude_daily_routines_usage_visible: Option, + pub alibaba_token_plan_region: Option, + pub weekly_progress_work_days: Option, } impl SettingsUpdate { fn refreshes_provider_data(&self) -> bool { self.enabled_providers.is_some() + || self.claude_daily_routines_usage_visible.is_some() + || self.alibaba_token_plan_region.is_some() + || self.weekly_progress_work_days.is_some() } fn notifies_float_bar(&self) -> bool { @@ -292,6 +298,19 @@ impl SettingsUpdate { settings.set_claude_avoid_keychain_prompts(true); } } + if let Some(v) = self.claude_daily_routines_usage_visible { + settings.claude_daily_routines_usage_visible = v; + } + if let Some(v) = self.alibaba_token_plan_region.as_deref() { + let region = codexbar::providers::AlibabaTokenPlanRegion::from_settings_value(Some(v)); + settings.set_api_region( + codexbar::core::ProviderId::AlibabaTokenPlan, + region.as_str(), + ); + } + if let Some(v) = self.weekly_progress_work_days { + settings.weekly_progress_work_days = if (2..=6).contains(&v) { Some(v) } else { None }; + } self } diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 18f3b3b167..2323a7e9ed 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -1149,6 +1149,23 @@ fn region_options_for_regional_provider() { assert_eq!(values, vec!["singapore", "us", "germany", "hongkong", "cn"]); } +#[test] +fn alibaba_token_plan_region_options() { + let opts = super::region_options_for("alibabatokenplan"); + let values: Vec<_> = opts.iter().map(|o| o.value.as_str()).collect(); + let labels: Vec<_> = opts.iter().map(|o| o.label.as_str()).collect(); + assert_eq!(values, vec!["cn", "intl", "cn-personal", "intl-personal"]); + assert_eq!( + labels, + vec![ + "China Team", + "International Team", + "China Personal/Solo", + "International Personal/Solo" + ] + ); +} + #[test] fn minimax_region_options_match_upstream_hosts() { let opts = super::region_options_for("minimax"); diff --git a/apps/desktop-tauri/src/App.test.tsx b/apps/desktop-tauri/src/App.test.tsx index cd073ed89d..2c8e3e8eed 100644 --- a/apps/desktop-tauri/src/App.test.tsx +++ b/apps/desktop-tauri/src/App.test.tsx @@ -110,6 +110,9 @@ function settings(overrides: Partial = {}): SettingsSnapshot { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, ...overrides, }; } diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index 1fbd5005ad..14309a68d7 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -459,6 +459,66 @@ describe("MenuCard", () => { expect(screen.queryByText("On-pace budget")).not.toBeInTheDocument(); }); + + it("shows spend used/limit plus balance secondary line", async () => { + tauriMocks.getLocaleStrings.mockResolvedValue( + buildBundle({ + DetailCostTitle: "Cost", + DetailCostUsed: "Used", + DetailCostBalance: "Balance", + DetailCostRemaining: "Remaining", + }), + ); + const snapshot = provider(null, 20); + snapshot.cost = { + used: 12.5, + limit: 100, + remaining: 87.5, + currencyCode: "USD", + period: "Extra usage", + resetsAt: null, + formattedUsed: "$12.50", + formattedLimit: "$100.00", + balance: 25.5, + formattedBalance: "$25.50", + }; + + renderCard(snapshot); + + expect(await screen.findByText(/Cost — Extra usage/)).toBeInTheDocument(); + expect(screen.getByText(/Used:\s*\$12\.50\s*\/\s*\$100\.00/)).toBeInTheDocument(); + expect(screen.getByText(/Balance:\s*\$25\.50/)).toBeInTheDocument(); + }); + + it("renders balance-only cost as credits-style value", async () => { + tauriMocks.getLocaleStrings.mockResolvedValue( + buildBundle({ + CreditsLabel: "Credits", + DetailCostTitle: "Cost", + DetailCostUsed: "Used", + }), + ); + const snapshot = provider(null, 20); + snapshot.cost = { + used: 0, + limit: null, + remaining: null, + currencyCode: "USD", + period: "Extra usage", + resetsAt: null, + formattedUsed: "$0.00", + formattedLimit: null, + balance: 25.5, + formattedBalance: "$25.50", + }; + + renderCard(snapshot); + + expect(await screen.findByText("Extra usage")).toBeInTheDocument(); + expect(screen.getByText("$25.50")).toBeInTheDocument(); + expect(screen.queryByText(/Used:/)).not.toBeInTheDocument(); + }); + it("localizes the relative updated-at time in Japanese without duplicated prefix", async () => { tauriMocks.getLocaleStrings.mockResolvedValue( buildBundle({ diff --git a/apps/desktop-tauri/src/components/MenuCard.tsx b/apps/desktop-tauri/src/components/MenuCard.tsx index 6e950937eb..68218e23c9 100644 --- a/apps/desktop-tauri/src/components/MenuCard.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.tsx @@ -654,30 +654,63 @@ export default function MenuCard({ {provider.cost && (
- {t("DetailCostTitle")} — {provider.cost.period} + {provider.cost.balance != null && provider.cost.limit == null + ? provider.cost.period || t("CreditsLabel") + : `${t("DetailCostTitle")} — ${provider.cost.period}`}
-
- {t("DetailCostUsed")}:{" "} - {provider.cost.formattedUsed || - formatCurrency(provider.cost.used, provider.cost.currencyCode)} - {provider.cost.limit != null && ( - <> - {" / "} - {provider.cost.formattedLimit || - formatCurrency(provider.cost.limit, provider.cost.currencyCode)} - - )} -
- {provider.cost.remaining != null && ( -
- {t("DetailCostRemaining")}:{" "} - {formatCurrency(provider.cost.remaining, provider.cost.currencyCode)} -
- )} - {formattedCostReset && ( -
- {t("DetailCostResets")}: {formattedCostReset} + {provider.cost.balance != null && provider.cost.limit == null ? ( +
+ {provider.cost.formattedBalance || + formatCurrency( + provider.cost.balance, + provider.cost.currencyCode, + )}
+ ) : ( + <> +
+ {t("DetailCostUsed")}:{" "} + {provider.cost.formattedUsed || + formatCurrency( + provider.cost.used, + provider.cost.currencyCode, + )} + {provider.cost.limit != null && ( + <> + {" / "} + {provider.cost.formattedLimit || + formatCurrency( + provider.cost.limit, + provider.cost.currencyCode, + )} + + )} +
+ {provider.cost.balance != null && ( +
+ {t("DetailCostBalance")}:{" "} + {provider.cost.formattedBalance || + formatCurrency( + provider.cost.balance, + provider.cost.currencyCode, + )} +
+ )} + {provider.cost.remaining != null && ( +
+ {t("DetailCostRemaining")}:{" "} + {formatCurrency( + provider.cost.remaining, + provider.cost.currencyCode, + )} +
+ )} + {formattedCostReset && ( +
+ {t("DetailCostResets")}: {formattedCostReset} +
+ )} + )}
)} diff --git a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx index 5722be9a19..c592f87451 100644 --- a/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx +++ b/apps/desktop-tauri/src/floatbar/FloatBar.test.tsx @@ -137,6 +137,9 @@ function settings(overrides: Partial = {}): SettingsSnapshot { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, ...overrides, }; } diff --git a/apps/desktop-tauri/src/floatbar/SettingsSection.test.tsx b/apps/desktop-tauri/src/floatbar/SettingsSection.test.tsx index 73b5bb842c..b481e12f8b 100644 --- a/apps/desktop-tauri/src/floatbar/SettingsSection.test.tsx +++ b/apps/desktop-tauri/src/floatbar/SettingsSection.test.tsx @@ -14,6 +14,9 @@ const settings = { floatBarOrientation: "horizontal", floatBarStyle: "floating", floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, floatBarShowResetInline: false, floatBarDarkText: false, floatBarClickThrough: false, diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 2e79d3980a..c44b8634c2 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -268,6 +268,8 @@ export const ALL_LOCALE_KEYS = [ "ProviderClaudeCookiesHelp", "ProviderClaudeAvoidKeychainPrompts", "ProviderClaudeAvoidKeychainPromptsHelp", + "ProviderClaudeDailyRoutinesUsage", + "ProviderClaudeDailyRoutinesUsageHelp", "ProviderCodexSparkUsage", "ProviderCodexSparkUsageHelp", "ProviderCursorCookieSourceHelp", @@ -487,6 +489,7 @@ export const ALL_LOCALE_KEYS = [ "DetailCostUsed", "DetailCostLimit", "DetailCostRemaining", + "DetailCostBalance", "DetailCostResets", "DetailChartCost", "DetailChartCredits", diff --git a/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx b/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx index dc96b99c34..74dbd82f1e 100644 --- a/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/PopOutPanel.test.tsx @@ -163,6 +163,9 @@ function settings(): SettingsSnapshot { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, }; } diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index e1f9318d20..31cbcee301 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -149,6 +149,9 @@ function settings(overrides: Partial = {}): SettingsSnapshot { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, ...overrides, }; } diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/ClaudeCreds.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/ClaudeCreds.tsx index e0e08f7f44..84d91d15fe 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/ClaudeCreds.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/ClaudeCreds.tsx @@ -7,36 +7,56 @@ interface Props { } /** - * Claude-specific credential options. + * Claude-specific credential/options. * * Port of the `ProviderId::Claude` branch of the "Options" block in - * `rust/src/native_ui/preferences.rs::render_provider_detail_panel` (~6662). - * Currently exposes the "Avoid keychain prompts" toggle. The broader - * `disable_keychain_access` master switch lives in the Advanced tab and - * is intentionally not duplicated here. + * `rust/src/native_ui/preferences.rs::render_provider_detail_panel`. + * Exposes "Avoid keychain prompts" and "Show Daily Routines usage". + * The broader `disable_keychain_access` master switch lives in Advanced. */ export function ClaudeCreds({ t }: Props) { - const [value, setValue] = useState(null); + const [avoidKeychain, setAvoidKeychain] = useState(null); + const [showDailyRoutines, setShowDailyRoutines] = useState( + null, + ); const [error, setError] = useState(null); const [saving, setSaving] = useState(false); useEffect(() => { let cancelled = false; getSettingsSnapshot() - .then((s) => !cancelled && setValue(s.claudeAvoidKeychainPrompts)) + .then((s) => { + if (cancelled) return; + setAvoidKeychain(s.claudeAvoidKeychainPrompts); + setShowDailyRoutines(s.claudeDailyRoutinesUsageVisible ?? true); + }) .catch((e) => !cancelled && setError(String(e))); return () => { cancelled = true; }; }, []); - const toggle = async (next: boolean) => { + const toggleAvoidKeychain = async (next: boolean) => { setSaving(true); try { const updated = await updateSettings({ claudeAvoidKeychainPrompts: next, }); - setValue(updated.claudeAvoidKeychainPrompts); + setAvoidKeychain(updated.claudeAvoidKeychainPrompts); + } catch (e) { + setError(String(e)); + } finally { + setSaving(false); + } + }; + + const toggleDailyRoutines = async (next: boolean) => { + setSaving(true); + try { + const updated = await updateSettings({ + claudeDailyRoutinesUsageVisible: next, + }); + setShowDailyRoutines(updated.claudeDailyRoutinesUsageVisible ?? next); } catch (e) { setError(String(e)); } finally { @@ -44,7 +64,7 @@ export function ClaudeCreds({ t }: Props) { } }; - if (value === null) return null; + if (avoidKeychain === null || showDailyRoutines === null) return null; return (
@@ -52,9 +72,9 @@ export function ClaudeCreds({ t }: Props) { + {error &&
{error}
}
); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx index 592fc6cf08..ca22a017cb 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AboutTab.test.tsx @@ -85,6 +85,9 @@ const settings: SettingsSnapshot = { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, }; describe("AboutTab", () => { diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx index fdd684fffa..e7f9e415e6 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/GeneralTab.test.tsx @@ -67,6 +67,9 @@ const settings: SettingsSnapshot = { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, showResetWhenExhausted: false, }; diff --git a/apps/desktop-tauri/src/types/bridge.test.ts b/apps/desktop-tauri/src/types/bridge.test.ts index 549351dbc2..a689ab8df0 100644 --- a/apps/desktop-tauri/src/types/bridge.test.ts +++ b/apps/desktop-tauri/src/types/bridge.test.ts @@ -86,6 +86,9 @@ describe("Language type", () => { floatBarDarkText: false, floatBarShowResetInline: false, floatBarShowCost: false, + claudeDailyRoutinesUsageVisible: true, + alibabaTokenPlanRegion: "cn", + weeklyProgressWorkDays: null, }; expect(snap.uiLanguage).toBe("spanish"); diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index f17c1e124a..474ab42765 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -200,6 +200,12 @@ export interface SettingsSnapshot { floatBarShowCost: boolean; /** Promote the tray icon out of the Windows hidden-icons overflow (Win11 only). */ promoteTrayIcon?: boolean; + /** When true, show Claude Daily Routines quota row (default true). */ + claudeDailyRoutinesUsageVisible: boolean; + /** Alibaba Token Plan region: cn | intl | cn-personal | intl-personal. */ + alibabaTokenPlanRegion: string; + /** Optional work-week length [2,6] for session-equivalent weekly forecast. */ + weeklyProgressWorkDays?: number | null; } /** Partial settings object — only include fields you want to change. */ @@ -261,6 +267,9 @@ export interface SettingsUpdate { floatBarShowResetInline?: boolean; floatBarShowCost?: boolean; promoteTrayIcon?: boolean; + claudeDailyRoutinesUsageVisible?: boolean; + alibabaTokenPlanRegion?: string; + weeklyProgressWorkDays?: number | null; } export interface UsageThresholdOverride { diff --git a/rust/src/locale.rs b/rust/src/locale.rs index 5b27061714..819cbc8309 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -482,6 +482,8 @@ locale_keys! { ProviderClaudeCookiesHelp, ProviderClaudeAvoidKeychainPrompts, ProviderClaudeAvoidKeychainPromptsHelp, + ProviderClaudeDailyRoutinesUsage, + ProviderClaudeDailyRoutinesUsageHelp, ProviderCodexSparkUsage, ProviderCodexSparkUsageHelp, ProviderCursorCookieSourceHelp, @@ -733,6 +735,7 @@ locale_keys! { DetailCostUsed, DetailCostLimit, DetailCostRemaining, + DetailCostBalance, DetailCostResets, DetailChartCost, DetailChartCredits, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index 51321892b9..cea035ab8b 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -263,6 +263,8 @@ ProviderClaudeCookies = Claude cookies ProviderClaudeCookiesHelp = Browser cookies/sessionKey are preferred because they match Claude's settings usage page. ProviderClaudeAvoidKeychainPrompts = Avoid Keychain prompts ProviderClaudeAvoidKeychainPromptsHelp = Use /usr/bin/security to read Claude credentials and avoid CodexBar keychain prompts. +ProviderClaudeDailyRoutinesUsage = Show Daily Routines usage +ProviderClaudeDailyRoutinesUsageHelp = Show the Daily Routines quota row for Claude web and OAuth usage. ProviderCodexSparkUsage = Show Codex Spark usage ProviderCodexSparkUsageHelp = Show Codex Spark quota rows without hiding credits or other extra usage. ProviderCursorCookieSourceHelp = Automatic imports browser cookies or stored sessions. @@ -469,6 +471,7 @@ DetailCostTitle = Cost DetailCostUsed = Used DetailCostLimit = Limit DetailCostRemaining = Remaining +DetailCostBalance = Balance DetailCostResets = Resets DetailChartCost = Cost (30 days) DetailChartCredits = Credits used (30 days) diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index bce913fae5..fc6e81096e 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -257,6 +257,8 @@ ProviderClaudeCookies = Cookies de Claude ProviderClaudeCookiesHelp = Se prefieren cookies/sessionKey del navegador porque coinciden con la página de uso de configuración de Claude. ProviderClaudeAvoidKeychainPrompts = Evitar avisos de llavero ProviderClaudeAvoidKeychainPromptsHelp = Usar /usr/bin/security para leer credenciales de Claude y evitar avisos de llavero de CodexBar. +ProviderClaudeDailyRoutinesUsage = Mostrar uso de Daily Routines +ProviderClaudeDailyRoutinesUsageHelp = Muestra la fila de cuota Daily Routines en el uso web/OAuth de Claude. ProviderCodexSparkUsage = Show Codex Spark usage ProviderCodexSparkUsageHelp = Show Codex Spark quota rows without hiding credits or other extra usage. ProviderCursorCookieSourceHelp = Importa automáticamente cookies del navegador o sesiones almacenadas. @@ -468,6 +470,7 @@ DetailCostTitle = Costo DetailCostUsed = Usado DetailCostLimit = Límite DetailCostRemaining = Restante +DetailCostBalance = Saldo DetailCostResets = Reinicia DetailChartCost = Costo (30 días) DetailChartCredits = Créditos usados (30 días) diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index e8670760f3..76dc028e9d 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -257,6 +257,8 @@ ProviderClaudeCookies = Claude Cookie ProviderClaudeCookiesHelp = Claude の設定使用量ページに合致するため、ブラウザ Cookie / sessionKey が優先されます。 ProviderClaudeAvoidKeychainPrompts = キーチェーン要求を出さない ProviderClaudeAvoidKeychainPromptsHelp = /usr/bin/security を使って Claude 認証情報を読み込み、CodexBar のキーチェーン要求を回避します。 +ProviderClaudeDailyRoutinesUsage = Daily Routines の使用量を表示 +ProviderClaudeDailyRoutinesUsageHelp = Claude の Web/OAuth 使用量に Daily Routines 行を表示します。 ProviderCodexSparkUsage = Codex Spark の使用量を表示 ProviderCodexSparkUsageHelp = クレジットやその他の追加使用量を非表示にせず、Codex Spark のクレジット行を表示します。 ProviderCursorCookieSourceHelp = ブラウザ Cookie または保存済みセッションを自動インポートします。 @@ -463,6 +465,7 @@ DetailCostTitle = コスト DetailCostUsed = 使用済み DetailCostLimit = 上限 DetailCostRemaining = 残り +DetailCostBalance = 残高 DetailCostResets = リセット DetailChartCost = コスト（30日間） DetailChartCredits = 使用クレジット（30日間） diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index b537aa53c0..7ea587819e 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -257,6 +257,8 @@ ProviderClaudeCookies = Claude 쿠키 ProviderClaudeCookiesHelp = 브라우저 쿠키/`sessionKey`가 Claude의 설정 사용량 페이지와 일치하므로 권장됩니다. ProviderClaudeAvoidKeychainPrompts = 키체인 메시지 표시 방지 ProviderClaudeAvoidKeychainPromptsHelp = Claude 자격 증명을 읽을 때 `/usr/bin/security`를 사용하여 CodexBar 키체인 권한 창이 뜨지 않도록 합니다. +ProviderClaudeDailyRoutinesUsage = Daily Routines 사용량 표시 +ProviderClaudeDailyRoutinesUsageHelp = Claude 웹/OAuth 사용량에 Daily Routines 할당량 행을 표시합니다. ProviderCodexSparkUsage = Show Codex Spark usage ProviderCodexSparkUsageHelp = Show Codex Spark quota rows without hiding credits or other extra usage. ProviderCursorCookieSourceHelp = 브라우저 쿠키 또는 저장된 세션을 자동으로 가져옵니다. @@ -468,6 +470,7 @@ DetailCostTitle = 비용 DetailCostUsed = 사용량 DetailCostLimit = 한도 DetailCostRemaining = 남음 +DetailCostBalance = 잔액 DetailCostResets = 초기화 DetailChartCost = 비용 (30일) DetailChartCredits = 사용 크레딧 (30일) diff --git a/rust/src/locale/tests.rs b/rust/src/locale/tests.rs index b197f3fb97..aefdc4dd42 100644 --- a/rust/src/locale/tests.rs +++ b/rust/src/locale/tests.rs @@ -79,6 +79,7 @@ fn test_japanese_menu_card_locale_values_are_translated() { (LocaleKey::DetailCostUsed, "使用済み"), (LocaleKey::DetailCostLimit, "上限"), (LocaleKey::DetailCostRemaining, "残り"), + (LocaleKey::DetailCostBalance, "残高"), (LocaleKey::DetailCostResets, "リセット"), (LocaleKey::DetailChartCost, "コスト（30日間）"), (LocaleKey::DetailChartCredits, "使用クレジット（30日間）"), diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index 7b5d9061c0..a5c94dc82e 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -257,6 +257,8 @@ ProviderClaudeCookies = Claude Cookie ProviderClaudeCookiesHelp = 优先使用浏览器 Cookie/sessionKey，因为它与 Claude 设置页的用量一致。 ProviderClaudeAvoidKeychainPrompts = 避免钥匙串提示 ProviderClaudeAvoidKeychainPromptsHelp = 使用 /usr/bin/security 读取 Claude 凭据，避免 CodexBar 的钥匙串提示。 +ProviderClaudeDailyRoutinesUsage = 显示 Daily Routines 用量 +ProviderClaudeDailyRoutinesUsageHelp = 显示 Claude 网页/OAuth 用量中的 Daily Routines 配额行。 ProviderCodexSparkUsage = 显示 Codex Spark 用量 ProviderCodexSparkUsageHelp = 显示 Codex Spark 配额行,不隐藏额度或其他额外用量。 ProviderCursorCookieSourceHelp = 自动导入浏览器 Cookie 或已保存会话。 @@ -463,6 +465,7 @@ DetailCostTitle = 费用 DetailCostUsed = 已用 DetailCostLimit = 限额 DetailCostRemaining = 剩余 +DetailCostBalance = 余额 DetailCostResets = 重置 DetailChartCost = 费用（30 天） DetailChartCredits = 已用额度（30 天） diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index 19138220a8..ecd54db806 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -257,6 +257,8 @@ ProviderClaudeCookies = Claude Cookie ProviderClaudeCookiesHelp = 優先使用瀏覽器 Cookie/sessionKey，因為它與 Claude 設定頁的用量一致。 ProviderClaudeAvoidKeychainPrompts = 避免鑰匙串提示 ProviderClaudeAvoidKeychainPromptsHelp = 使用 /usr/bin/security 讀取 Claude 憑據，避免 CodexBar 的鑰匙串提示。 +ProviderClaudeDailyRoutinesUsage = 顯示 Daily Routines 用量 +ProviderClaudeDailyRoutinesUsageHelp = 顯示 Claude 網頁/OAuth 用量中的 Daily Routines 配額行。 ProviderCodexSparkUsage = 顯示 Codex Spark 用量 ProviderCodexSparkUsageHelp = 顯示 Codex Spark 配額行,不隱藏額度或其他額外用量。 ProviderCursorCookieSourceHelp = 自動匯入瀏覽器 Cookie 或已儲存會話。 @@ -463,6 +465,7 @@ DetailCostTitle = 費用 DetailCostUsed = 已用 DetailCostLimit = 限額 DetailCostRemaining = 剩餘 +DetailCostBalance = 餘額 DetailCostResets = 重置 DetailChartCost = 費用（30 天） DetailChartCredits = 已用額度（30 天） diff --git a/rust/src/providers/alibabatokenplan/region.rs b/rust/src/providers/alibabatokenplan/region.rs index c1b2402739..9019fd8850 100644 --- a/rust/src/providers/alibabatokenplan/region.rs +++ b/rust/src/providers/alibabatokenplan/region.rs @@ -18,6 +18,17 @@ pub enum AlibabaTokenPlanRegion { } impl AlibabaTokenPlanRegion { + pub const ALL: [Self; 4] = [Self::Cn, Self::Intl, Self::CnPersonal, Self::IntlPersonal]; + + pub fn display_name(self) -> &'static str { + match self { + Self::Cn => "China Team", + Self::Intl => "International Team", + Self::CnPersonal => "China Personal/Solo", + Self::IntlPersonal => "International Personal/Solo", + } + } + /// Serde/settings raw string exactly as persisted. pub fn as_str(self) -> &'static str { match self { diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 8c3826c4aa..b80d1ba4fd 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -797,7 +797,11 @@ impl Settings { } pub fn set_api_region(&mut self, id: ProviderId, region: impl Into) { - self.provider_config_mut(id).api_region = Some(region.into()); + let region = region.into(); + if id == ProviderId::AlibabaTokenPlan { + self.alibaba_token_plan_region = region.clone(); + } + self.provider_config_mut(id).api_region = Some(region); } /// Manual cookie header for `id`, or `""` if unset. From b365874c7b5a17d91366a526e005d4a8676024d7 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:12:40 +0700 Subject: [PATCH 12/15] Fix ZoomMate tray/settings icon letter --- .../src/components/providers/icons/ProviderIcon-zoommate.svg | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg index f8806a5391..bed342df05 100644 --- a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-zoommate.svg @@ -1,4 +1,3 @@ - - + From 74daa0eadc3204cef8293b3c760d12f7ae7aa183 Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:14:36 +0700 Subject: [PATCH 13/15] Register Qwen Cloud provider icon and fix qwen alias --- apps/desktop-tauri/src/components/providers/providerIcons.ts | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 11e3da6941..c442890fc2 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -208,6 +208,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { poe: { id: "poe", brandColor: "#5d5fef", fallbackLetter: "P" }, devin: { id: "devin", brandColor: "#111827", fallbackLetter: "D" }, zed: { id: "zed", brandColor: "#084ccf", fallbackLetter: "Z" }, + qwencloud: { id: "qwencloud", brandColor: "#615CED", fallbackLetter: "Q" }, }; const ALIASES: Record = { @@ -217,7 +218,9 @@ const ALIASES: Record = { "jetbrains ai": "jetbrains", "kimi k2": "kimik2", tongyi: "alibaba", - qwen: "alibaba", + qwen: "qwencloud", + "qwen cloud": "qwencloud", + "qwen-cloud": "qwencloud", qianwen: "alibaba", "alibaba token plan": "alibabatokenplan", "alibaba-token-plan": "alibabatokenplan", From 16858f178d699a9717beb5d8e522fc61a69dadee Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:20:48 +0700 Subject: [PATCH 14/15] fix(clippy): use ? in curl_capture URL parse fallback --- rust/src/core/curl_capture.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/rust/src/core/curl_capture.rs b/rust/src/core/curl_capture.rs index 14ed579dd3..9d2f793b37 100644 --- a/rust/src/core/curl_capture.rs +++ b/rust/src/core/curl_capture.rs @@ -27,10 +27,9 @@ pub fn request_url(raw: &str) -> Option { m.as_str().to_string() } else if let Some(m) = caps.get(3) { unescape_shell_segment(m.as_str(), false) - } else if let Some(m) = caps.get(4) { - unescape_shell_segment(m.as_str(), false) } else { - return None; + let m = caps.get(4)?; + unescape_shell_segment(m.as_str(), false) }; let value = value.trim(); if value.is_empty() { From f127afa2377dbc2164728382d6fe313adb2927bf Mon Sep 17 00:00:00 2001 From: Finesssee <90105158+Finesssee@users.noreply.github.com> Date: Thu, 30 Jul 2026 12:25:37 +0700 Subject: [PATCH 15/15] fix(test): harden ProvidersSidebar reorder and add qwencloud catalog --- .../providers/ProvidersSidebar.test.tsx | 52 ++++++++++++------- .../desktop-tauri/src/test/providerCatalog.ts | 1 + 2 files changed, 35 insertions(+), 18 deletions(-) diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx index 2a51ad5887..8b947a7dc0 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/ProvidersSidebar.test.tsx @@ -1,3 +1,4 @@ +import { useState } from "react"; import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { LocaleProvider } from "../../../i18n/LocaleProvider"; @@ -91,22 +92,42 @@ describe("ProvidersSidebar", () => { it("reorders providers through explicit move buttons", async () => { const onReorder = vi.fn(); - const { container } = render( - - - , - ); + const initial = rows(); + function Harness() { + const [providers, setProviders] = useState(initial); + return ( + + { + onReorder(orderedIds); + setProviders((prev) => { + const byId = new Map(prev.map((p) => [p.id, p])); + return orderedIds + .map((id) => byId.get(id)) + .filter((p): p is ProviderSidebarRow => Boolean(p)); + }); + }} + onToggleEnabled={vi.fn()} + /> + + ); + } + const { container } = render(); fireEvent.click(await screen.findByRole("button", { name: "Move down Codex" })); + await waitFor(() => { + expect(onReorder).toHaveBeenCalledWith([ + "claude", + "codex", + ...TEST_PROVIDER_CATALOG.slice(2).map(([id]) => id), + ]); + }); await waitFor(() => { const names = Array.from( container.querySelectorAll(".providers-sidebar__name"), @@ -114,11 +135,6 @@ describe("ProvidersSidebar", () => { ); expect(names.slice(0, 3)).toEqual(["Claude", "Codex", "Cursor"]); }); - expect(onReorder).toHaveBeenCalledWith([ - "claude", - "codex", - ...TEST_PROVIDER_CATALOG.slice(2).map(([id]) => id), - ]); }); it("does not let the first provider move up", async () => { diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index bda7cbc154..20e04840c8 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -63,4 +63,5 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["qoder", "Qoder"], ["sakana", "Sakana AI"], ["sub2api", "sub2api"], + ["qwencloud", "Qwen Cloud"], ];