diff --git a/README.es-MX.md b/README.es-MX.md index dbd389b54f..389d3186b8 100644 --- a/README.es-MX.md +++ b/README.es-MX.md @@ -102,7 +102,6 @@ Consulta el historial completo en [CHANGELOG.md](CHANGELOG.md). | Xiaomi MiMo | Cookies | Saldo, Plan de tokens | | Doubao | Clave API | Límites de solicitudes | | Command Code | Cookies | Créditos mensuales, Créditos comprados | -| Crof | Clave API | Créditos, Cuota de solicitudes | | StepFun | Token Oasis | 5h, Semanal, Refresco de token | | Venice | Clave API | Saldo USD / DIEM | | OpenAI | Admin API / Clave API | Uso, Solicitudes, Costo con ámbito de proyecto, Saldo de créditos | diff --git a/README.md b/README.md index 2e052ad2f9..ec974c6265 100755 --- a/README.md +++ b/README.md @@ -109,7 +109,6 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Xiaomi MiMo | Cookies | Balance, Token Plan | | Doubao | API Key | Request Limits | | Command Code | Cookies | Monthly Credits, Purchased Credits | -| Crof | API Key | Credits, Request Quota | | StepFun | Oasis Token | 5h, Weekly, Token refresh | | Venice | API Key | USD / DIEM Balance | | OpenAI | Admin API / API Key | Usage, Requests, Project-scoped cost, Credit Balance | diff --git a/README.tr-TR.md b/README.tr-TR.md index 47ab7bb13b..0ae7410775 100644 --- a/README.tr-TR.md +++ b/README.tr-TR.md @@ -107,7 +107,6 @@ Tüm geçmişi [CHANGELOG.md](CHANGELOG.md) dosyasında görebilirsiniz. | Xiaomi MiMo | Çerezler | Bakiye, Token Planı | | Doubao | API Anahtarı | İstek Sınırları | | Command Code | Çerezler | Aylık Krediler, Satın Alınan Krediler | -| Crof | API Anahtarı | Krediler, İstek Kotası | | StepFun | Oasis Tokenı | 5 saatlik, Haftalık, Token yenileme | | Venice | API Anahtarı | USD / DIEM Bakiyesi | | OpenAI | Yönetici API'si / API Anahtarı | Kullanım, İstekler, Proje kapsamlı maliyet, Kredi Bakiyesi | diff --git a/README.zh-CN.md b/README.zh-CN.md index 2776c0fa95..55550f15e8 100644 --- a/README.zh-CN.md +++ b/README.zh-CN.md @@ -14,7 +14,7 @@ ## 功能特性 -- **56 个 AI 服务商** — Codex、Claude、Cursor、Factory、Gemini、Copilot、Antigravity、z.ai、MiniMax、Kiro、Vertex AI、Augment、OpenCode、Kimi、Kimi K2、Amp、Warp、Ollama、Azure OpenAI、T3 Chat、OpenRouter、JetBrains AI、Alibaba、Alibaba Token Plan、NanoGPT、Infini、Perplexity、Abacus AI、Mistral、OpenCode Go、Kilo、AWS Bedrock、Codebuff、DeepSeek、Windsurf、Manus、小米 MiMo、Doubao、Command Code、Crof、StepFun、Venice、OpenAI、Grok、ElevenLabs、Deepgram、Groq、LLM Proxy、Chutes、LiteLLM、Poe、Devin、Zed、CrossModel、Qoder、Sakana AI +- **56 个 AI 服务商** — Codex、Claude、Cursor、Factory、Gemini、Copilot、Antigravity、z.ai、MiniMax、Kiro、Vertex AI、Augment、OpenCode、Kimi、Kimi K2、Amp、Warp、Ollama、Azure OpenAI、T3 Chat、OpenRouter、JetBrains AI、Alibaba、Alibaba Token Plan、NanoGPT、Infini、Perplexity、Abacus AI、Mistral、OpenCode Go、Kilo、AWS Bedrock、Codebuff、DeepSeek、Windsurf、Manus、小米 MiMo、Doubao、Command Code、StepFun、Venice、OpenAI、Grok、ElevenLabs、Deepgram、Groq、LLM Proxy、Chutes、LiteLLM、Poe、Devin、Zed、CrossModel、Qoder、Sakana AI - **系统托盘图标** — 动态双条进度显示会话与周用量 - **Floating Bar** — 可选的置顶透明用量条，支持方向、透明度和点击穿透控制 - **浏览器 Cookie 导入** — Chrome、Edge、Brave、Firefox（Windows DPAPI 解密） @@ -217,7 +217,6 @@ codexbar cost -p codex # 本地成本（JSONL 日志） | 小米 MiMo | Cookies | 余额、Token 套餐 | | Doubao | API Key | 请求限制 | | Command Code | Cookies | 月度 Credits、已购 Credits | -| Crof | API Key | Credits、请求配额 | | StepFun | Oasis Token | 5h、周用量 | | Venice | API Key | USD / DIEM 余额 | | OpenAI | Admin API / API Key | 用量、请求数、余额 | diff --git a/README.zh-TW.md b/README.zh-TW.md index cbda22a135..b2f3dfbe65 100644 --- a/README.zh-TW.md +++ b/README.zh-TW.md @@ -14,7 +14,7 @@ ## 功能特性 -- **56 個 AI 提供者** — Codex、Claude、Cursor、Factory、Gemini、Copilot、Antigravity、z.ai、MiniMax、Kiro、Vertex AI、Augment、OpenCode、Kimi、Kimi K2、Amp、Warp、Ollama、Azure OpenAI、T3 Chat、OpenRouter、JetBrains AI、Alibaba、Alibaba Token Plan、NanoGPT、Infini、Perplexity、Abacus AI、Mistral、OpenCode Go、Kilo、AWS Bedrock、Codebuff、DeepSeek、Windsurf、Manus、小米 MiMo、Doubao、Command Code、Crof、StepFun、Venice、OpenAI、Grok、ElevenLabs、Deepgram、Groq、LLM Proxy、Chutes、LiteLLM、Poe、Devin、Zed、CrossModel、Qoder、Sakana AI +- **56 個 AI 提供者** — Codex、Claude、Cursor、Factory、Gemini、Copilot、Antigravity、z.ai、MiniMax、Kiro、Vertex AI、Augment、OpenCode、Kimi、Kimi K2、Amp、Warp、Ollama、Azure OpenAI、T3 Chat、OpenRouter、JetBrains AI、Alibaba、Alibaba Token Plan、NanoGPT、Infini、Perplexity、Abacus AI、Mistral、OpenCode Go、Kilo、AWS Bedrock、Codebuff、DeepSeek、Windsurf、Manus、小米 MiMo、Doubao、Command Code、StepFun、Venice、OpenAI、Grok、ElevenLabs、Deepgram、Groq、LLM Proxy、Chutes、LiteLLM、Poe、Devin、Zed、CrossModel、Qoder、Sakana AI - **系統系統匣圖示** — 動態雙條進度顯示會話與周用量 - **Floating Bar** — 可選的置頂透明用量條，支援方向、透明度和點選穿透控制 - **瀏覽器 Cookie 匯入** — Chrome、Edge、Brave、Firefox（Windows DPAPI 解密） @@ -217,7 +217,6 @@ codexbar cost -p codex # 本機成本（JSONL 日誌） | 小米 MiMo | Cookies | 餘額、Token 套餐 | | Doubao | API Key | 請求限制 | | Command Code | Cookies | 月度 Credits、已購 Credits | -| Crof | API Key | Credits、請求配額 | | StepFun | Oasis Token | 5h、周用量 | | Venice | API Key | USD / DIEM 餘額 | | OpenAI | Admin API / API Key | 用量、請求數、餘額 | diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index ab5efbef8a..96955cf740 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -661,6 +661,8 @@ pub struct SettingsSnapshot { predictive_pace_warning_enabled: bool, show_pace: bool, tray_icon_mode: &'static str, + stacked_tray_top_provider: Option, + stacked_tray_bottom_provider: Option, switcher_shows_icons: bool, menu_bar_shows_highest_usage: bool, menu_bar_shows_percent: bool, @@ -783,6 +785,8 @@ impl From for SettingsSnapshot { predictive_pace_warning_enabled: settings.predictive_pace_warning_enabled, show_pace: settings.show_pace, tray_icon_mode: tray_icon_mode_label(settings.tray_icon_mode), + stacked_tray_top_provider: settings.stacked_tray_top_provider, + stacked_tray_bottom_provider: settings.stacked_tray_bottom_provider, switcher_shows_icons: settings.switcher_shows_icons, menu_bar_shows_highest_usage: settings.menu_bar_shows_highest_usage, menu_bar_shows_percent: settings.menu_bar_shows_percent, @@ -876,6 +880,7 @@ fn tray_icon_mode_label(mode: TrayIconMode) -> &'static str { match mode { TrayIconMode::Single => "single", TrayIconMode::PerProvider => "perProvider", + TrayIconMode::Stacked => "stacked", } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index fdfaeda967..06721a6e8e 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -104,6 +104,25 @@ fn canonical_provider_arg(provider_id: &str) -> Result { Ok(parse_provider_arg(provider_id)?.cli_name().to_string()) } +fn provider_dashboard_url(id: ProviderId, settings: &Settings) -> Option { + match id { + ProviderId::MiniMax => Some( + codexbar::providers::MiniMaxProvider::dashboard_url_for_region(Some( + settings.api_region(id), + )), + ), + ProviderId::Kimi => Some( + codexbar::providers::KimiRegion::from_settings(Some(settings.api_region(id))) + .console_url() + .to_string(), + ), + _ => instantiate_provider(id) + .metadata() + .dashboard_url + .map(str::to_string), + } +} + fn validate_single_line_secret(value: &str, field: &str, max_len: usize) -> Result<(), String> { let trimmed = value.trim(); if trimmed.is_empty() { diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs index 418bd5c146..6feeea4868 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs @@ -73,15 +73,7 @@ pub(crate) fn build_provider_detail( let provider = instantiate_provider(id); let metadata = provider.metadata(); let resume_supported = auto_resume_supported(id); - let dashboard_url = if id == codexbar::core::ProviderId::MiniMax { - Some( - codexbar::providers::MiniMaxProvider::dashboard_url_for_region(Some( - settings.api_region(id), - )), - ) - } else { - metadata.dashboard_url.map(|s| s.to_string()) - }; + let dashboard_url = provider_dashboard_url(id, &settings); let detail = ProviderDetail { id: id.cli_name().to_string(), diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs index 80675e1eca..a92aa1bdc4 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_refresh.rs @@ -87,7 +87,7 @@ fn provider_cache_can_skip_refresh( .iter() .any(|snapshot| snapshot.provider_id == id.cli_name()) }); - if !force && crate::proof_harness::seed_usage_json_active() && cache_has_all { + if !force && state.provider_cache_seeded && cache_has_all { return true; } !force @@ -125,6 +125,38 @@ pub(super) fn complete_provider_refresh( mod tests { use super::*; + #[test] + fn validated_seed_pins_only_complete_nonforced_cache() { + let mut state = AppState::new(); + state.provider_cache.push( + crate::proof_harness::parse_seed_usage_snapshot( + r#"{"providerId":"codex","primary":{"usedPercent":25.0}}"#, + ) + .unwrap(), + ); + assert!(!provider_cache_can_skip_refresh( + &state, + false, + &[ProviderId::Codex] + )); + state.provider_cache_seeded = true; + assert!(provider_cache_can_skip_refresh( + &state, + false, + &[ProviderId::Codex] + )); + assert!(!provider_cache_can_skip_refresh( + &state, + true, + &[ProviderId::Codex] + )); + assert!(!provider_cache_can_skip_refresh( + &state, + false, + &[ProviderId::Codex, ProviderId::Claude] + )); + } + #[test] fn stale_inputs_cannot_reserve_a_new_generation_after_invalidation() { let mut state = AppState::new(); 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 f9344fcfba..4028a688e9 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -272,6 +272,7 @@ fn region_provider(provider_id: &str) -> Option { "alibabatokenplan" => ProviderId::AlibabaTokenPlan, "zai" => ProviderId::Zai, "minimax" => ProviderId::MiniMax, + "kimi" => ProviderId::Kimi, _ => return None, }) } @@ -284,6 +285,10 @@ pub(crate) fn provider_region_lookup(settings: &Settings, provider_id: &str) -> )) .settings_value() .to_string() + } else if id == codexbar::core::ProviderId::Kimi { + codexbar::providers::KimiRegion::from_settings(Some(settings.api_region(id))) + .settings_value() + .to_string() } else { settings.api_region(id).to_string() } @@ -820,6 +825,14 @@ pub fn region_options_for(provider_id: &str) -> Vec { .to_string(), }, ], + "kimi" => codexbar::providers::KimiRegion::ALL + .iter() + .copied() + .map(|region| RegionOption { + value: region.settings_value().to_string(), + label: region.display_name().to_string(), + }) + .collect(), "alibabatokenplan" => codexbar::providers::AlibabaTokenPlanRegion::ALL .iter() .copied() diff --git a/apps/desktop-tauri/src-tauri/src/commands/settings.rs b/apps/desktop-tauri/src-tauri/src/commands/settings.rs index df0e6621a9..b687c9fd48 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/settings.rs @@ -26,6 +26,8 @@ pub struct SettingsUpdate { pub predictive_pace_warning_enabled: Option, pub show_pace: Option, pub tray_icon_mode: Option, + pub stacked_tray_top_provider: Option, + pub stacked_tray_bottom_provider: Option, pub switcher_shows_icons: Option, pub menu_bar_shows_highest_usage: Option, pub menu_bar_shows_percent: Option, @@ -121,6 +123,8 @@ impl SettingsUpdate { fn refreshes_tray_presentation(&self) -> bool { self.tray_icon_mode.is_some() + || self.stacked_tray_top_provider.is_some() + || self.stacked_tray_bottom_provider.is_some() || self.switcher_shows_icons.is_some() || self.menu_bar_shows_highest_usage.is_some() || self.menu_bar_shows_percent.is_some() @@ -190,6 +194,12 @@ impl SettingsUpdate { { settings.tray_icon_mode = mode; } + if let Some(provider) = self.stacked_tray_top_provider.clone() { + settings.stacked_tray_top_provider = normalize_optional_provider_id(provider); + } + if let Some(provider) = self.stacked_tray_bottom_provider.clone() { + settings.stacked_tray_bottom_provider = normalize_optional_provider_id(provider); + } if let Some(v) = self.provider_metrics.clone() { apply_provider_metrics(settings, v); } @@ -482,10 +492,16 @@ fn parse_tray_icon_mode(s: &str) -> Option { match s { "single" => Some(TrayIconMode::Single), "perProvider" => Some(TrayIconMode::PerProvider), + "stacked" => Some(TrayIconMode::Stacked), _ => None, } } +fn normalize_optional_provider_id(value: String) -> Option { + let value = value.trim(); + (!value.is_empty()).then(|| value.to_string()) +} + fn parse_update_channel(s: &str) -> Option { match s { "stable" => Some(UpdateChannel::Stable), @@ -700,6 +716,36 @@ mod tests { } .refreshes_tray_presentation() ); + assert!( + SettingsUpdate { + stacked_tray_top_provider: Some("claude".to_string()), + ..Default::default() + } + .refreshes_tray_presentation() + ); + } + + #[test] + fn stacked_tray_update_accepts_mode_and_clears_automatic_provider() { + let mut settings = Settings { + stacked_tray_top_provider: Some("codex".to_string()), + ..Settings::default() + }; + + SettingsUpdate { + tray_icon_mode: Some("stacked".to_string()), + stacked_tray_top_provider: Some(String::new()), + stacked_tray_bottom_provider: Some("claude".to_string()), + ..Default::default() + } + .apply_provider_settings(&mut settings); + + assert_eq!(settings.tray_icon_mode, TrayIconMode::Stacked); + assert_eq!(settings.stacked_tray_top_provider, None); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("claude") + ); } #[test] diff --git a/apps/desktop-tauri/src-tauri/src/commands/system.rs b/apps/desktop-tauri/src-tauri/src/commands/system.rs index 9ae17f57ce..d7f56b124d 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/system.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/system.rs @@ -205,36 +205,9 @@ pub fn quit_app(app: tauri::AppHandle) { } fn dashboard_url_for_provider(provider_id: &str) -> Option { - if provider_id == ProviderId::MiniMax.cli_name() { - let settings = Settings::load(); - return Some( - codexbar::providers::MiniMaxProvider::dashboard_url_for_region(Some( - settings.api_region(ProviderId::MiniMax), - )), - ); - } - - // OpenRouter's Usage Dashboard is the Activity page. Resolve it from the - // provider metadata before the legacy API-key catalog entry, which still - // points at the credits settings page. - if provider_id == ProviderId::OpenRouter.cli_name() { - return instantiate_provider(ProviderId::OpenRouter) - .metadata() - .dashboard_url - .map(|s| s.to_string()); - } - - if let Some(url) = codexbar::settings::get_api_key_providers() - .into_iter() - .find(|p| p.id.cli_name() == provider_id) - .and_then(|p| p.dashboard_url.map(|s| s.to_string())) - { - return Some(url); - } - let id = ProviderId::from_cli_name(provider_id)?; - let provider = instantiate_provider(id); - provider.metadata().dashboard_url.map(|s| s.to_string()) + let settings = Settings::load(); + provider_dashboard_url(id, &settings) } fn status_page_url_for_provider(provider_id: &str) -> Option { diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 444e93da0a..b18b7c949a 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -276,6 +276,20 @@ fn minimax_region_lookup_normalizes_legacy_china_value() { assert_eq!(provider_region_lookup(&s, "minimax").as_deref(), Some("cn")); } +#[test] +fn kimi_region_lookup_defaults_to_china_and_roundtrips_international() { + let mut settings = Settings::default(); + assert_eq!( + provider_region_lookup(&settings, "kimi").as_deref(), + Some("china") + ); + super::provider_region_set(&mut settings, "kimi", "international".to_string()).unwrap(); + assert_eq!( + provider_region_lookup(&settings, "kimi").as_deref(), + Some("international") + ); +} + #[test] fn minimax_cookie_domain_follows_selected_region() { let mut s = Settings::default(); @@ -703,6 +717,40 @@ fn fetch_context_includes_minimax_region() { assert_eq!(ctx.api_region.as_deref(), Some("cn")); } +#[test] +fn provider_dashboard_url_uses_selected_regional_console() { + let mut settings = Settings::default(); + settings.set_api_region(ProviderId::MiniMax, "cn"); + settings.set_api_region(ProviderId::Kimi, "international"); + + assert_eq!( + super::provider_dashboard_url(ProviderId::MiniMax, &settings).as_deref(), + Some("https://platform.minimaxi.com/user-center/payment/coding-plan?cycle_type=3") + ); + assert_eq!( + super::provider_dashboard_url(ProviderId::Kimi, &settings).as_deref(), + Some("https://www.kimi.ai/code/console") + ); +} + +/// Provider metadata is the only source of the provider dashboard link; the +/// API-key catalog URL is the key-management link shown next to the key field. +/// A provider that only has a catalog URL must get a metadata URL instead of +/// silently borrowing the key page. +#[test] +fn api_key_catalog_providers_have_metadata_dashboard_urls() { + let settings = Settings::default(); + for provider in codexbar::settings::get_api_key_providers() { + if provider.dashboard_url.is_some() { + assert!( + super::provider_dashboard_url(provider.id, &settings).is_some(), + "{:?} has an API-key page but no metadata dashboard URL", + provider.id + ); + } + } +} + #[test] fn fetch_context_token_account_uses_web_cookie_header() { let settings = Settings::default(); @@ -1837,6 +1885,13 @@ fn minimax_region_options_match_upstream_hosts() { ); } +#[test] +fn kimi_region_options_match_regional_hosts() { + let opts = super::region_options_for("kimi"); + let values: Vec<_> = opts.iter().map(|option| option.value.as_str()).collect(); + assert_eq!(values, vec!["china", "international"]); +} + #[test] fn region_options_empty_for_non_regional_provider() { assert!(super::region_options_for("claude").is_empty()); diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index 01f562e111..ee7ae528c8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -26,6 +26,10 @@ pub struct UsageSpendRow { pub display_name: String, pub seven_day: Option, pub thirty_day: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub seven_day_estimate: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub thirty_day_estimate: Option, pub seven_day_tokens: Option, pub thirty_day_tokens: Option, pub currency: String, @@ -496,6 +500,7 @@ fn build_usage_spend_summary( }) .unwrap_or_else(|| provider_id.clone()); + let mut local_cost_estimates = None; let spend = match provider_id.as_str() { "codex" => SpendValues { seven_day: codex_7_contract.known_cost_usd, @@ -586,17 +591,11 @@ fn build_usage_spend_summary( spend } "antigravity" => { - use codexbar::providers::antigravity::local_sessions::LocalHistoryCoverage; let seven = codexbar::providers::antigravity::local_sessions::summarize(7); let thirty = codexbar::providers::antigravity::local_sessions::summarize(30); - let mut spend = cached_spend(cached_snapshot); - spend.seven_day_tokens = matches!(seven.coverage, LocalHistoryCoverage::Complete) - .then_some(seven.total_tokens); - spend.thirty_day_tokens = matches!(thirty.coverage, LocalHistoryCoverage::Complete) - .then_some(thirty.total_tokens); - if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { - spend.source = "local Antigravity history".to_string(); - } + let spend = + antigravity_spend_values(cached_spend(cached_snapshot), &seven, &thirty); + local_cost_estimates = Some((seven.cost_estimate, thirty.cost_estimate)); spend } _ => cached_spend(cached_snapshot), @@ -618,11 +617,16 @@ fn build_usage_spend_summary( .collect() }) .unwrap_or_default(); + let (seven_day_estimate, thirty_day_estimate) = local_cost_estimates + .map(|(seven, thirty)| (Some(seven), Some(thirty))) + .unwrap_or((None, None)); rows.push(UsageSpendRow { provider_id: provider_id.clone(), display_name, seven_day: spend.seven_day, thirty_day: spend.thirty_day, + seven_day_estimate, + thirty_day_estimate, seven_day_tokens: spend.seven_day_tokens, thirty_day_tokens: spend.thirty_day_tokens, currency, @@ -701,6 +705,29 @@ fn total_token_mix(mix: &codexbar::spend_contract::SpendTokenMix) -> Option saw.then_some(total) } +fn antigravity_spend_values( + mut spend: SpendValues, + seven: &codexbar::spend_contract::LocalTokenHistorySummary, + thirty: &codexbar::spend_contract::LocalTokenHistorySummary, +) -> SpendValues { + use codexbar::spend_contract::LocalHistoryCoverage; + + spend.seven_day = seven.total_usd(); + spend.thirty_day = thirty.total_usd(); + spend.seven_day_tokens = + (seven.coverage == LocalHistoryCoverage::Complete).then_some(seven.total_tokens); + spend.thirty_day_tokens = + (thirty.coverage == LocalHistoryCoverage::Complete).then_some(thirty.total_tokens); + if spend.thirty_day.is_some() { + spend.source = "local Antigravity history · API list-price estimate".to_string(); + } else if thirty.cost_estimate.known_subtotal_usd.is_some() { + spend.source = "local Antigravity history · known API list-price subtotal".to_string(); + } else if thirty.coverage == LocalHistoryCoverage::Complete { + spend.source = "local Antigravity history · unpriced".to_string(); + } + spend +} + fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { let Some(snapshot) = snapshot else { return SpendValues { @@ -778,6 +805,27 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { mod cache_key_tests { use super::*; + fn local_history( + total_tokens: u64, + coverage: codexbar::spend_contract::LocalHistoryCoverage, + known_subtotal_usd: Option, + unpriced: u32, + ) -> codexbar::spend_contract::LocalTokenHistorySummary { + codexbar::spend_contract::LocalTokenHistorySummary { + total_tokens, + session_count: if total_tokens > 0 { 1 } else { 0 }, + coverage, + cost_estimate: codexbar::spend_contract::LocalCostEstimate { + known_subtotal_usd, + coverage: codexbar::spend_contract::CostCoverageCounts { + estimated: if known_subtotal_usd.is_some() { 1 } else { 0 }, + unpriced, + ..Default::default() + }, + }, + } + } + #[test] fn invalidated_owner_clears_orphaned_indexing_activity() { let mut coordinator = UsageSpendCoordinator::default(); @@ -834,4 +882,34 @@ mod cache_key_tests { assert!(include_in_shared_overview("claude", false, true)); assert!(!include_in_shared_overview("codex", false, false)); } + + #[test] + fn antigravity_partial_history_exposes_only_the_known_subtotal() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let seven = local_history(100, LocalHistoryCoverage::Partial, Some(1.25), 0); + let thirty = local_history(200, LocalHistoryCoverage::Partial, Some(2.50), 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &thirty); + + assert_eq!(spend.seven_day, None); + assert_eq!(spend.thirty_day, None); + assert_eq!(spend.seven_day_tokens, None); + assert_eq!(spend.thirty_day_tokens, None); + assert!(spend.source.contains("known API list-price subtotal")); + } + + #[test] + fn antigravity_complete_empty_history_is_a_known_zero() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let seven = local_history(0, LocalHistoryCoverage::Complete, None, 0); + let thirty = local_history(0, LocalHistoryCoverage::Complete, None, 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &thirty); + + assert_eq!(spend.seven_day, Some(0.0)); + assert_eq!(spend.thirty_day, Some(0.0)); + assert_eq!(spend.seven_day_tokens, Some(0)); + assert_eq!(spend.thirty_day_tokens, Some(0)); + assert!(spend.source.contains("API list-price estimate")); + } } diff --git a/apps/desktop-tauri/src-tauri/src/main.rs b/apps/desktop-tauri/src-tauri/src/main.rs index df9a86e148..04f9388e80 100644 --- a/apps/desktop-tauri/src-tauri/src/main.rs +++ b/apps/desktop-tauri/src-tauri/src/main.rs @@ -19,6 +19,7 @@ mod surface_target; mod tray_accounts; mod tray_bridge; mod tray_menu; +mod tray_presentation; mod tray_visibility; mod usage_metric; mod window_positioner; @@ -140,17 +141,27 @@ fn main() { let mut initial_state = AppState::new(); initial_state.proof_config = proof_config; - // Proof-harness seed: CODEXBAR_SEED_USAGE_JSON plants one synthetic Codex - // ProviderUsageSnapshot before the event loop and any WebView read. The - // cache timestamp makes the seeded cache count as fresh so the first - // frontend refresh-if-stale call does not evict the synthetic data. - if let Some(snapshot) = proof_harness::seed_usage_snapshot_from_env() { - tracing::info!( - "proof-harness: seeded provider snapshot for '{}'", - snapshot.provider_id - ); - initial_state.provider_cache.push(snapshot); - initial_state.provider_cache_updated_at = Some(std::time::Instant::now()); + // Validate the complete proof seed before installing any snapshots, so an + // invalid multi-provider fixture cannot leave a partial cache behind. + if let Some(snapshots) = + proof_harness::seed_usage_snapshots_from_env(initial_state.proof_config.as_ref()) + { + let seeded_at = std::time::Instant::now(); + for snapshot in &snapshots { + tracing::info!( + "proof-harness: seeded provider snapshot for '{}'", + snapshot.provider_id + ); + if let Some(provider) = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id) + { + initial_state + .provider_cache_updated_at_by_provider + .insert(provider, seeded_at); + } + } + initial_state.provider_cache.extend(snapshots); + initial_state.provider_cache_seeded = true; + initial_state.provider_cache_updated_at = Some(seeded_at); } tauri::Builder::default() diff --git a/apps/desktop-tauri/src-tauri/src/proof_harness.rs b/apps/desktop-tauri/src-tauri/src/proof_harness.rs index 4a08484f39..8cc2979a87 100644 --- a/apps/desktop-tauri/src-tauri/src/proof_harness.rs +++ b/apps/desktop-tauri/src-tauri/src/proof_harness.rs @@ -17,14 +17,17 @@ //! and suppresses blur-dismiss so the window stays visible for automated //! screenshot capture. //! -//! `CODEXBAR_SEED_USAGE_JSON=` additionally seeds one synthetic, -//! bridge-shaped Codex [`ProviderUsageSnapshot`] into the provider cache at -//! launch (before the first event/WebView read) and pins it against refresh -//! eviction for the run. Malformed files log a warning and the shell -//! continues without seeding — proof runs must never crash on the seed. - +//! `CODEXBAR_SEED_USAGE_JSON=` additionally seeds either the legacy +//! synthetic Codex [`ProviderUsageSnapshot`] object or, in valid proof mode, +//! a nonempty array of unique supported provider snapshots. The validated set +//! is installed before the first event/WebView read and pinned against refresh +//! eviction. Malformed files log a warning and the shell continues without +//! seeding — proof runs must never crash on the seed. + +use std::collections::HashSet; use std::sync::Mutex; +use codexbar::core::ProviderId; use serde::Serialize; use tauri::{AppHandle, Manager}; @@ -140,6 +143,13 @@ pub fn activate(app: &AppHandle) { } } +/// Seeded snapshots must use an exact bridge provider id, not a CLI alias. +fn is_seedable_provider_id(provider_id: &str) -> bool { + ProviderId::all() + .iter() + .any(|provider| provider.cli_name() == provider_id) +} + /// Bottom inset (physical px) kept between the proof panel's bottom edge and /// the monitor work-area bottom (#265). const PROOF_BOTTOM_INSET_PX: i32 = 8; @@ -260,22 +270,17 @@ pub fn is_proof_mode(app: &AppHandle) -> bool { // ── Provider-usage seed (CODEXBAR_SEED_USAGE_JSON) ─────────────────── /// Environment variable pointing at a JSON file with one synthetic, -/// bridge-shaped `ProviderUsageSnapshot` for the codex provider. +/// bridge-shaped Codex snapshot or a proof-only array of provider snapshots. pub const SEED_USAGE_ENV_VAR: &str = "CODEXBAR_SEED_USAGE_JSON"; -/// Whether a seed path was configured at launch. While set, the provider -/// cache is pinned fresh so the synthetic snapshot is never evicted by an -/// automatic refresh during a proof/capture run. -pub fn seed_usage_json_active() -> bool { - std::env::var_os(SEED_USAGE_ENV_VAR).is_some() -} - /// Read and validate the seed file referenced by `CODEXBAR_SEED_USAGE_JSON`. /// /// Returns `None` (with a warn, never a crash) when the variable is unset, -/// the file is unreadable, the JSON is malformed, or the snapshot is not -/// for the `codex` provider. -pub fn seed_usage_snapshot_from_env() -> Option { +/// the file is unreadable, or the seed does not satisfy the selected legacy +/// object or proof-only array contract. +pub fn seed_usage_snapshots_from_env( + proof_config: Option<&ProofConfig>, +) -> Option> { let path = std::env::var_os(SEED_USAGE_ENV_VAR)?; let path = std::path::PathBuf::from(path); let raw = match std::fs::read_to_string(&path) { @@ -288,8 +293,8 @@ pub fn seed_usage_snapshot_from_env() -> Option { return None; } }; - match parse_seed_usage_snapshot(&raw) { - Ok(snapshot) => Some(snapshot), + match parse_seed_usage_snapshots(&raw, proof_config) { + Ok(snapshots) => Some(snapshots), Err(msg) => { tracing::warn!("{SEED_USAGE_ENV_VAR}: {msg} in {}", path.display()); None @@ -319,6 +324,51 @@ pub fn parse_seed_usage_snapshot(json: &str) -> Result, +) -> Result, String> { + if !json.trim_start().starts_with('[') { + return parse_seed_usage_snapshot(json).map(|snapshot| vec![snapshot]); + } + if !proof_config.is_some_and(is_valid_proof_config) { + return Err("provider snapshot arrays require valid proof mode".into()); + } + + let mut snapshots: Vec = + serde_json::from_str(json).map_err(|e| format!("malformed JSON: {e}"))?; + if snapshots.is_empty() { + return Err("provider snapshot array must not be empty".into()); + } + + let mut providers = HashSet::with_capacity(snapshots.len()); + for snapshot in &mut snapshots { + if !is_seedable_provider_id(&snapshot.provider_id) { + return Err(format!( + "unsupported snapshot providerId '{}', ignoring", + snapshot.provider_id + )); + } + if !providers.insert(snapshot.provider_id.clone()) { + return Err(format!( + "duplicate snapshot providerId '{}', ignoring", + snapshot.provider_id + )); + } + normalize_seed_snapshot(snapshot); + } + Ok(snapshots) +} + +fn normalize_seed_snapshot(snapshot: &mut ProviderUsageSnapshot) { normalize_rate_window(&mut snapshot.primary); snapshot.secondary.as_mut().map(normalize_rate_window); snapshot.model_specific.as_mut().map(normalize_rate_window); @@ -331,8 +381,11 @@ pub fn parse_seed_usage_snapshot(json: &str) -> Result bool { + SurfaceMode::parse(&config.target_surface) + .is_some_and(|mode| proof_payload_is_supported(mode, config.target_payload.as_deref())) } /// Recompute `remaining_percent` from `used_percent` (matching the canonical @@ -584,4 +637,72 @@ mod tests { assert_eq!(cost.period, "month"); assert_eq!(cost.formatted_used, "$12.50"); } + + fn valid_proof_config() -> ProofConfig { + ProofConfig { + target_surface: "trayPanel".into(), + settings_tab: None, + target_payload: None, + } + } + + fn seed_snapshot(provider_id: &str, used_percent: f64) -> serde_json::Value { + serde_json::json!({ + "providerId": provider_id, + "primary": { "usedPercent": used_percent, "windowMinutes": 300 } + }) + } + + #[test] + fn seed_snapshot_arrays_preserve_legacy_codex_object_behavior() { + let json = seed_snapshot("codex", 61.0).to_string(); + let snapshots = parse_seed_usage_snapshots(&json, None).expect("legacy seed parses"); + assert_eq!(snapshots.len(), 1); + assert_eq!(snapshots[0].provider_id, "codex"); + assert_eq!(snapshots[0].primary.remaining_percent, 39.0); + } + + #[test] + fn seed_snapshot_array_normalizes_multiple_supported_providers() { + let json = serde_json::json!([seed_snapshot("codex", 61.0), seed_snapshot("claude", 24.0)]) + .to_string(); + let snapshots = parse_seed_usage_snapshots(&json, Some(&valid_proof_config())) + .expect("supported snapshots parse in proof mode"); + + assert_eq!(snapshots.len(), 2); + assert_eq!(snapshots[0].provider_id, "codex"); + assert_eq!(snapshots[0].primary.remaining_percent, 39.0); + assert_eq!(snapshots[1].provider_id, "claude"); + assert_eq!(snapshots[1].primary.remaining_percent, 76.0); + assert!( + snapshots + .iter() + .all(|snapshot| !snapshot.updated_at.is_empty()) + ); + } + + #[test] + fn seed_snapshot_arrays_reject_empty_unknown_duplicate_and_non_finite_values() { + let proof = valid_proof_config(); + for json in [ + "[]".to_string(), + serde_json::json!([seed_snapshot("unknown-provider", 1.0)]).to_string(), + // CLI aliases are not bridge provider ids. + serde_json::json!([seed_snapshot("openai", 1.0)]).to_string(), + serde_json::json!([seed_snapshot("codex", 1.0), seed_snapshot("codex", 2.0)]) + .to_string(), + r#"[{"providerId":"codex","primary":{"usedPercent":1e400}}]"#.to_string(), + ] { + assert!( + parse_seed_usage_snapshots(&json, Some(&proof)).is_err(), + "{json}" + ); + } + } + + #[test] + fn seed_snapshot_arrays_are_rejected_without_valid_proof_config() { + let json = serde_json::json!([seed_snapshot("codex", 10.0)]).to_string(); + assert!(parse_seed_usage_snapshots(&json, None).is_err()); + } } diff --git a/apps/desktop-tauri/src-tauri/src/state.rs b/apps/desktop-tauri/src-tauri/src/state.rs index 10d626ec43..174c13d113 100644 --- a/apps/desktop-tauri/src-tauri/src/state.rs +++ b/apps/desktop-tauri/src-tauri/src/state.rs @@ -143,6 +143,8 @@ pub struct AppState { pub installer_path: Option, /// Proof-harness configuration (set when `CODEXBAR_PROOF_MODE` is active). pub proof_config: Option, + /// True only after a validated proof seed was installed at startup. + pub provider_cache_seeded: bool, /// Persistent notification manager — tracks which alerts have fired to prevent spam. pub notification_manager: codexbar::notifications::NotificationManager, /// Instant when the tray panel was last shown — used to suppress @@ -204,6 +206,7 @@ impl AppState { last_update_check_ms: None, installer_path: None, proof_config: None, + provider_cache_seeded: false, notification_manager: codexbar::notifications::NotificationManager::new(), last_shown_at: None, last_blur_dismissed_at: None, diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index cb11675767..40b038547a 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -3,22 +3,18 @@ use std::sync::Mutex; use crate::commands::ProviderCatalogEntry; -#[cfg(test)] -use codexbar::core::ProviderId; -use codexbar::settings::MetricPreference; -use codexbar::settings::{Settings, TrayIconMode}; +use codexbar::settings::Settings; use tauri::image::Image; use tauri::menu::{CheckMenuItemBuilder, IsMenuItem, Menu, MenuItem, PredefinedMenuItem, Submenu}; use tauri::tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent}; use tauri::{AppHandle, Manager}; -use codexbar::tray::{render_bar_icon_rgba, render_percent_icon_rgba}; - use crate::shell; use crate::state::{AppState, TrayAnchor}; #[cfg(test)] use crate::tray_menu::build_tray_menu; use crate::tray_menu::{TrayMenuEntry, build_tray_menu_with}; +use crate::tray_presentation::{TrayPresentationPlan, headline_window}; #[derive(Debug, Clone, Copy)] struct MonitorScaleInfo { @@ -337,7 +333,8 @@ pub(crate) fn rebuild_tray_menu(app: &AppHandle) { let settings = Settings::load(); let status_labels = if let Some(st) = app.try_state::>() { let guard = st.lock().unwrap(); - status_labels_for_settings(&settings, &guard.provider_cache, settings.ui_language) + TrayPresentationPlan::resolve(&settings, &guard.provider_cache) + .status_labels(settings.ui_language) } else { vec![] }; @@ -355,7 +352,8 @@ pub fn update_tray_status_items( ) { let catalog = crate::commands::get_provider_catalog(); let settings = Settings::load(); - let status_labels = status_labels_for_settings(&settings, snapshots, settings.ui_language); + let status_labels = + TrayPresentationPlan::resolve(&settings, snapshots).status_labels(settings.ui_language); if let Ok(menu) = build_native_tray_menu(app, &catalog, &status_labels) && let Some(tray) = app.tray_by_id("codexbar-main") @@ -393,228 +391,16 @@ pub fn update_tray_icon_and_tooltip( return; }; - // ── Icon ───────────────────────────────────────────────────────────── let settings = Settings::load(); - let snapshots = snapshots.to_vec(); - let ordered_snapshots = ordered_snapshot_refs(&settings, &snapshots); - let ok_snapshots: Vec<_> = ordered_snapshots - .iter() - .copied() - .filter(|s| s.error.is_none()) - .collect(); - let all_error = ok_snapshots.is_empty() && !snapshots.is_empty(); - - let prefer_highest = settings.menu_bar_shows_highest_usage - || settings.menu_bar_display_mode.as_str() == "minimal"; - - let picked = pick_tray_provider(&ok_snapshots, prefer_highest); - - let (session_pct, weekly_pct) = match picked { - Some(s) => selected_tray_percents(s, &settings), - None => ( - ok_snapshots - .iter() - .map(|s| selected_tray_percents(s, &settings).0) - .fold(0.0_f64, f64::max), - None, - ), - }; - - let (rgba, w, h) = render_tray_icon_for_settings(&settings, session_pct, weekly_pct, all_error); + let plan = TrayPresentationPlan::resolve(&settings, snapshots); + let (rgba, w, h) = plan.render_icon(); let icon = Image::new_owned(rgba, w, h); let _ = tray.set_icon(Some(icon)); - // ── Tooltip ─────────────────────────────────────────────────────────── - let tooltip = build_tooltip(&snapshots, settings.ui_language); + let tooltip = build_tooltip(snapshots, settings.ui_language); let _ = tray.set_tooltip(Some(tooltip)); } -fn status_labels_for_settings( - settings: &Settings, - snapshots: &[crate::commands::ProviderUsageSnapshot], - lang: codexbar::settings::Language, -) -> Vec<(String, String)> { - let ordered_snapshots = ordered_snapshot_refs(settings, snapshots); - let healthy: Vec<_> = ordered_snapshots - .into_iter() - .filter(|s| s.error.is_none()) - .collect(); - if settings.tray_icon_mode == TrayIconMode::PerProvider { - return healthy - .into_iter() - .map(|s| provider_status_label(s, lang)) - .collect::>(); - } - - let Some(selected) = pick_tray_provider( - &healthy, - settings.menu_bar_shows_highest_usage || settings.menu_bar_display_mode == "minimal", - ) else { - return vec![]; - }; - - let (_, label) = provider_status_label(selected, lang); - vec![("status_summary".to_string(), label)] -} - -fn ordered_snapshot_refs<'a>( - settings: &Settings, - snapshots: &'a [crate::commands::ProviderUsageSnapshot], -) -> Vec<&'a crate::commands::ProviderUsageSnapshot> { - let order = settings - .provider_display_order_names() - .into_iter() - .enumerate() - .map(|(index, provider_id)| (provider_id, index)) - .collect::>(); - let mut ordered = snapshots.iter().collect::>(); - ordered.sort_by(|a, b| { - let a_order = order.get(&a.provider_id); - let b_order = order.get(&b.provider_id); - match (a_order, b_order) { - (Some(a_order), Some(b_order)) if a_order != b_order => a_order.cmp(b_order), - (Some(_), None) => std::cmp::Ordering::Less, - (None, Some(_)) => std::cmp::Ordering::Greater, - _ => a.display_name.cmp(&b.display_name), - } - }); - ordered -} - -fn provider_status_label( - snapshot: &crate::commands::ProviderUsageSnapshot, - lang: codexbar::settings::Language, -) -> (String, String) { - // MonthlyPlan metric (PAYG spend, e.g. Mistral): show formatted cost. - let provider = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id); - let preference = provider - .map(|id| Settings::load().get_provider_metric(id)) - .unwrap_or_default(); - if preference == MetricPreference::MonthlyPlan - && let Some(cost) = snapshot.cost.as_ref() - { - let amount = if !cost.formatted_used.is_empty() { - cost.formatted_used.clone() - } else { - crate::commands::format_cost_amount(cost) - }; - return ( - snapshot.provider_id.clone(), - format!("{} {}", snapshot.display_name, amount), - ); - } - - let label = crate::commands::compact_tray_status_label(headline_window(snapshot), lang); - ( - snapshot.provider_id.clone(), - format!("{} {}", snapshot.display_name, label), - ) -} - -/// Window that headline tray surfaces should label for a provider. -/// -/// F5 (upstream 0.48.0): for Codex, prefer the first non-informational lane so -/// a monthly-only plan shows the monthly window with its reset countdown -/// instead of the informational "No active 5h session" placeholder. -/// -/// Shared by the tray menu rows (`provider_status_label`) and the tray tooltip -/// (`build_tooltip`) so the two cannot drift apart. -fn headline_window( - snapshot: &crate::commands::ProviderUsageSnapshot, -) -> &crate::commands::RateWindowSnapshot { - if snapshot.provider_id == "codex" { - codex_lane_headline_window(snapshot) - } else { - &snapshot.primary - } -} - -/// F5 (upstream 0.48.0): pick the first non-informational Codex lane in -/// session → weekly → monthly order. When all lanes are informational -/// (no active session at all), fall back to the primary for the -/// "No active 5h session" placeholder. -pub(crate) fn codex_lane_headline_window( - snapshot: &crate::commands::ProviderUsageSnapshot, -) -> &crate::commands::RateWindowSnapshot { - if !snapshot.primary.is_informational { - return &snapshot.primary; - } - if let Some(ref secondary) = snapshot.secondary - && !secondary.is_informational - { - return secondary; - } - if let Some(ref tertiary) = snapshot.tertiary - && !tertiary.is_informational - { - return tertiary; - } - &snapshot.primary -} - -fn render_tray_icon_for_settings( - settings: &Settings, - session_pct: f64, - weekly_pct: Option, - all_error: bool, -) -> (Vec, u32, u32) { - if settings.menu_bar_shows_percent { - render_percent_icon_rgba(session_pct, all_error) - } else { - render_bar_icon_rgba(session_pct, weekly_pct, all_error) - } -} - -/// Pick the provider whose usage the tray icon should render. -/// -/// Exposed so that the unit tests can exercise both `highest` and `first` -/// paths without needing a live Tauri app handle. -fn pick_tray_provider<'a>( - ok_snapshots: &'a [&'a crate::commands::ProviderUsageSnapshot], - prefer_highest: bool, -) -> Option<&'a crate::commands::ProviderUsageSnapshot> { - if ok_snapshots.is_empty() { - return None; - } - if prefer_highest { - ok_snapshots.iter().copied().max_by(|a, b| { - a.primary - .used_percent - .partial_cmp(&b.primary.used_percent) - .unwrap_or(std::cmp::Ordering::Equal) - }) - } else { - Some(ok_snapshots[0]) - } -} - -fn selected_tray_percents( - snapshot: &crate::commands::ProviderUsageSnapshot, - settings: &Settings, -) -> (f64, Option) { - let (selected, companion) = - crate::usage_metric::selected_usage_icon_windows(snapshot, settings); - ( - display_metric_percent(&selected, settings.show_as_used), - companion - .as_ref() - .map(|window| display_metric_percent(window, settings.show_as_used)), - ) -} - -fn display_metric_percent(window: &crate::commands::RateWindowSnapshot, show_as_used: bool) -> f64 { - if window.is_informational { - return 0.0; - } - if window.is_exhausted || window.used_percent >= 100.0 { - return if show_as_used { 100.0 } else { 0.0 }; - } - - let used_percent = window.used_percent; - let used = used_percent.clamp(0.0, 100.0); - if show_as_used { used } else { 100.0 - used } -} - /// Build a compact multi-line tooltip string from provider snapshots. fn build_tooltip( snapshots: &[crate::commands::ProviderUsageSnapshot], @@ -964,131 +750,6 @@ mod tests { fake_snapshot_with(id, display, used_percent, None, None, None) } - fn fake_extra_window(percent: f64) -> crate::commands::NamedRateWindowSnapshot { - crate::commands::NamedRateWindowSnapshot { - id: "additional_budget".to_string(), - title: "Additional Budget".to_string(), - fallback_lane: false, - window: crate::commands::RateWindowSnapshot { - used_percent: percent, - remaining_percent: 100.0 - percent, - 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, - }, - } - } - - #[test] - fn pick_tray_provider_highest_picks_max_primary() { - let a = fake_snapshot("codex", "Codex", 30.0); - let b = fake_snapshot("claude", "Claude", 72.5); - let c = fake_snapshot("gemini", "Gemini", 50.0); - let refs: Vec<&crate::commands::ProviderUsageSnapshot> = vec![&a, &b, &c]; - - let picked = pick_tray_provider(&refs, /* prefer_highest = */ true) - .expect("highest mode should pick a provider"); - assert_eq!(picked.provider_id, "claude"); - } - - #[test] - fn pick_tray_provider_first_preserves_catalog_order() { - let a = fake_snapshot("codex", "Codex", 30.0); - let b = fake_snapshot("claude", "Claude", 72.5); - let refs: Vec<&crate::commands::ProviderUsageSnapshot> = vec![&a, &b]; - - let picked = pick_tray_provider(&refs, /* prefer_highest = */ false) - .expect("non-highest mode should still pick the first entry"); - assert_eq!(picked.provider_id, "codex"); - } - - #[test] - fn pick_tray_provider_none_when_empty() { - let refs: Vec<&crate::commands::ProviderUsageSnapshot> = vec![]; - assert!(pick_tray_provider(&refs, true).is_none()); - assert!(pick_tray_provider(&refs, false).is_none()); - } - - #[test] - fn status_labels_per_provider_mode_lists_each_healthy_provider() { - let settings = Settings { - tray_icon_mode: TrayIconMode::PerProvider, - provider_order: codexbar::settings::normalize_provider_order(&[ - "claude".to_string(), - "codex".to_string(), - ]), - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let labels = status_labels_for_settings( - &settings, - &snapshots, - codexbar::settings::Language::English, - ); - - assert_eq!( - labels, - vec![ - ("claude".to_string(), "Claude 72%".to_string()), - ("codex".to_string(), "Codex 30%".to_string()), - ] - ); - } - - #[test] - fn status_labels_single_mode_collapses_to_selected_provider() { - let settings = Settings { - tray_icon_mode: TrayIconMode::Single, - menu_bar_shows_highest_usage: true, - ..Settings::default() - }; - let snapshots = vec![ - fake_snapshot("codex", "Codex", 30.0), - fake_snapshot("claude", "Claude", 72.0), - ]; - - let labels = status_labels_for_settings( - &settings, - &snapshots, - codexbar::settings::Language::English, - ); - - assert_eq!( - labels, - vec![("status_summary".to_string(), "Claude 72%".to_string())] - ); - } - - #[test] - fn tray_icon_renderer_uses_percent_mode_when_enabled() { - let bar_settings = Settings { - menu_bar_shows_percent: false, - ..Settings::default() - }; - let percent_settings = Settings { - menu_bar_shows_percent: true, - ..Settings::default() - }; - - let (bar, bar_w, bar_h) = - render_tray_icon_for_settings(&bar_settings, 72.0, Some(40.0), false); - let (percent, pct_w, pct_h) = - render_tray_icon_for_settings(&percent_settings, 72.0, Some(40.0), false); - - assert_eq!((bar_w, bar_h), (pct_w, pct_h)); - assert_ne!(bar, percent); - } - #[test] fn tooltip_uses_compact_status_labels() { let mut claude = fake_snapshot("claude", "Claude", 13.0); @@ -1165,319 +826,16 @@ mod tests { "{japanese_tooltip}" ); - let (_, english_label) = - provider_status_label(&claude, codexbar::settings::Language::English); - let (_, japanese_label) = - provider_status_label(&claude, codexbar::settings::Language::Japanese); + let settings = Settings::default(); + let snapshots = vec![claude]; + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + let english_label = plan.status_labels(codexbar::settings::Language::English)[0] + .1 + .clone(); + let japanese_label = plan.status_labels(codexbar::settings::Language::Japanese)[0] + .1 + .clone(); assert!(english_label.contains("Resets in"), "{english_label}"); assert!(japanese_label.contains("リセットまで"), "{japanese_label}"); } - - #[test] - fn selected_tray_percent_uses_cursor_extra_usage_cost() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - ); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 15.0); - assert_eq!(secondary, Some(20.0)); - } - - #[test] - fn selected_tray_percent_tracks_extra_rate_window() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Copilot, MetricPreference::ExtraUsage); - let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); - snapshot.extra_rate_windows.push(fake_extra_window(42.0)); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, None); - } - - #[test] - fn copilot_automatic_tracks_highest_extra_rate_window() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); - snapshot.extra_rate_windows.push(fake_extra_window(42.0)); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - } - - #[test] - fn selected_tray_percent_respects_remaining_display_mode() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with( - "cursor", - "Cursor", - 10.0, - Some(20.0), - Some(72.0), - Some((15.0, 100.0)), - ); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 85.0); - assert_eq!(secondary, Some(80.0)); - } - - #[test] - fn exhausted_automatic_window_never_renders_as_remaining_progress() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with( - "opencodego", - "OpenCode Go", - 20.0, - Some(60.0), - Some(40.0), - None, - ); - snapshot - .tertiary - .as_mut() - .expect("monthly quota") - .is_exhausted = true; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(remaining, 0.0); - - settings.show_as_used = true; - let (used, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(used, 100.0); - } - - #[test] - fn full_automatic_window_without_exhausted_flag_has_zero_remaining_progress() { - let mut settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with( - "opencodego", - "OpenCode Go", - 20.0, - Some(60.0), - Some(100.0), - None, - ); - snapshot - .tertiary - .as_mut() - .expect("monthly quota") - .is_exhausted = false; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(remaining, 0.0); - - settings.show_as_used = true; - let (used, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(used, 100.0); - } - - #[test] - fn missing_automatic_window_does_not_look_like_available_remaining_progress() { - let settings = Settings { - show_as_used: false, - ..Settings::default() - }; - let mut snapshot = fake_snapshot_with("opencodego", "OpenCode Go", 0.0, None, None, None); - snapshot.primary.is_informational = true; - - let (remaining, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(remaining, 0.0); - } - - #[test] - fn selected_tray_percent_falls_back_when_extra_usage_missing() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); - let snapshot = fake_snapshot_with("cursor", "Cursor", 10.0, Some(72.0), None, None); - - let (primary, _) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 72.0); - } - - #[test] - fn single_meaningful_secondary_quota_uses_full_single_meter() { - let settings = Settings::default(); - let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); - snapshot.primary.is_informational = true; - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, None); - } - - #[test] - fn selected_secondary_quota_is_not_duplicated_when_tertiary_is_meaningful() { - let settings = Settings::default(); - let mut snapshot = - fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), Some(30.0), None); - snapshot.primary.is_informational = true; - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 42.0); - assert_eq!(secondary, Some(30.0)); - } - - #[test] - fn two_meaningful_quotas_keep_two_meter_layout() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Cursor, MetricPreference::Session); - let snapshot = fake_snapshot_with("cursor", "Cursor", 15.0, Some(40.0), None, None); - - let (primary, secondary) = selected_tray_percents(&snapshot, &settings); - - assert_eq!(primary, 15.0); - assert_eq!(secondary, Some(40.0)); - } - - #[test] - fn informational_primary_skips_session_and_automatic_phantom_zero() { - let mut settings = Settings::default(); - settings.set_provider_metric(ProviderId::Claude, MetricPreference::Session); - let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); - snapshot.primary.is_informational = true; - - // Session preference must not paint the synthetic 0% primary; - // it falls through to Automatic which prefers weekly (42%). - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 42.0); - assert_ne!(primary, 0.0); - - // Automatic also prefers weekly over informational primary. - settings.set_provider_metric(ProviderId::Claude, MetricPreference::Automatic); - let (primary, _) = selected_tray_percents(&snapshot, &settings); - assert_eq!(primary, 42.0); - } - - #[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); - } - - #[test] - fn f5_headline_prefers_non_informational_primary() { - let snapshot = fake_snapshot_with("codex", "Codex", 50.0, Some(20.0), Some(30.0), None); - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 50.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_falls_back_to_secondary_when_primary_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); - snapshot.primary.is_informational = true; - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 25.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_falls_back_to_tertiary_when_primary_and_secondary_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(35.0), None); - snapshot.primary.is_informational = true; - snapshot.secondary.as_mut().unwrap().is_informational = true; - let headline = codex_lane_headline_window(&snapshot); - assert!((headline.used_percent - 35.0).abs() < f64::EPSILON); - } - - #[test] - fn f5_headline_returns_primary_when_all_informational() { - let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(0.0), None); - snapshot.primary.is_informational = true; - if let Some(sec) = &mut snapshot.secondary { - sec.is_informational = true; - } - if let Some(ter) = &mut snapshot.tertiary { - ter.is_informational = true; - } - let headline = codex_lane_headline_window(&snapshot); - // Falls back to primary (the placeholder) when all are informational. - assert!(headline.is_informational); - } } diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs new file mode 100644 index 0000000000..6292a2c1fa --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation.rs @@ -0,0 +1,370 @@ +//! Pure tray presentation policy shared by the native tray surfaces. + +use crate::commands::{ProviderUsageSnapshot, RateWindowSnapshot}; +use codexbar::settings::{Language, MetricPreference, Settings, TrayIconMode}; +use codexbar::tray::{ + render_bar_icon_rgba, render_percent_icon_rgba, render_stacked_bar_icon_rgba, +}; + +#[derive(Debug, Clone, Copy, PartialEq)] +enum TrayIconPlan { + Bars { + primary_percent: f64, + secondary_percent: Option, + has_error: bool, + }, + Percent { + percent: f64, + has_error: bool, + }, + Stacked { + top_percent: f64, + bottom_percent: f64, + has_error: bool, + }, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TrayStatusKey { + Summary, + Provider, +} + +#[derive(Debug, Clone, Copy)] +struct TrayStatusRow<'a> { + key: TrayStatusKey, + snapshot: &'a ProviderUsageSnapshot, +} + +/// Fully resolved tray presentation, independent of Tauri and operating-system state. +/// +/// The plan is the single policy boundary for provider ordering, mode-specific +/// selection, metric selection, status rows, and icon renderer choice. +pub(crate) struct TrayPresentationPlan<'a> { + settings: &'a Settings, + icon: TrayIconPlan, + status_rows: Vec>, +} + +impl<'a> TrayPresentationPlan<'a> { + pub(crate) fn resolve(settings: &'a Settings, snapshots: &'a [ProviderUsageSnapshot]) -> Self { + let ordered = ordered_snapshot_refs(settings, snapshots); + let healthy = ordered + .into_iter() + .filter(|snapshot| snapshot.error.is_none()) + .collect::>(); + let has_error = healthy.is_empty() && !snapshots.is_empty(); + let prefer_highest = + settings.menu_bar_shows_highest_usage || settings.menu_bar_display_mode == "minimal"; + let selected = pick_tray_provider(&healthy, prefer_highest); + + let (icon, status_rows) = match settings.tray_icon_mode { + TrayIconMode::Stacked => { + if let Some((top, bottom)) = pick_stacked_tray_providers(&healthy, settings) { + ( + TrayIconPlan::Stacked { + top_percent: selected_tray_percents(top, settings).0, + bottom_percent: selected_tray_percents(bottom, settings).0, + has_error, + }, + vec![ + TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot: top, + }, + TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot: bottom, + }, + ], + ) + } else { + let percents = selected + .map(|snapshot| selected_tray_percents(snapshot, settings)) + .unwrap_or((0.0, None)); + let rows = healthy + .first() + .map(|snapshot| TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot, + }) + .into_iter() + .collect(); + ( + resolve_single_provider_icon_plan( + settings, percents.0, percents.1, has_error, + ), + rows, + ) + } + } + TrayIconMode::PerProvider => { + let percents = selected + .map(|snapshot| selected_tray_percents(snapshot, settings)) + .unwrap_or_else(|| fallback_percents(&healthy, settings)); + let rows = healthy + .iter() + .copied() + .map(|snapshot| TrayStatusRow { + key: TrayStatusKey::Provider, + snapshot, + }) + .collect(); + ( + resolve_single_provider_icon_plan(settings, percents.0, percents.1, has_error), + rows, + ) + } + TrayIconMode::Single => { + let percents = selected + .map(|snapshot| selected_tray_percents(snapshot, settings)) + .unwrap_or_else(|| fallback_percents(&healthy, settings)); + let rows = selected + .map(|snapshot| TrayStatusRow { + key: TrayStatusKey::Summary, + snapshot, + }) + .into_iter() + .collect(); + ( + resolve_single_provider_icon_plan(settings, percents.0, percents.1, has_error), + rows, + ) + } + }; + + Self { + settings, + icon, + status_rows, + } + } + + pub(crate) fn render_icon(&self) -> (Vec, u32, u32) { + match self.icon { + TrayIconPlan::Bars { + primary_percent, + secondary_percent, + has_error, + } => render_bar_icon_rgba(primary_percent, secondary_percent, has_error), + TrayIconPlan::Percent { percent, has_error } => { + render_percent_icon_rgba(percent, has_error) + } + TrayIconPlan::Stacked { + top_percent, + bottom_percent, + has_error, + } => render_stacked_bar_icon_rgba(top_percent, bottom_percent, has_error), + } + } + + pub(crate) fn status_labels(&self, language: Language) -> Vec<(String, String)> { + self.status_rows + .iter() + .map(|row| { + let (_, label) = provider_status_label(row.snapshot, self.settings, language); + let key = match row.key { + TrayStatusKey::Summary => "status_summary".to_string(), + TrayStatusKey::Provider => row.snapshot.provider_id.clone(), + }; + (key, label) + }) + .collect() + } +} + +fn resolve_single_provider_icon_plan( + settings: &Settings, + primary_percent: f64, + secondary_percent: Option, + has_error: bool, +) -> TrayIconPlan { + if settings.menu_bar_shows_percent { + TrayIconPlan::Percent { + percent: primary_percent, + has_error, + } + } else { + TrayIconPlan::Bars { + primary_percent, + secondary_percent, + has_error, + } + } +} + +fn fallback_percents( + healthy: &[&ProviderUsageSnapshot], + settings: &Settings, +) -> (f64, Option) { + ( + healthy + .iter() + .map(|snapshot| selected_tray_percents(snapshot, settings).0) + .fold(0.0_f64, f64::max), + None, + ) +} + +fn ordered_snapshot_refs<'a>( + settings: &Settings, + snapshots: &'a [ProviderUsageSnapshot], +) -> Vec<&'a ProviderUsageSnapshot> { + let order = settings + .provider_display_order_names() + .into_iter() + .enumerate() + .map(|(index, provider_id)| (provider_id, index)) + .collect::>(); + let mut ordered = snapshots.iter().collect::>(); + ordered.sort_by(|a, b| { + let a_order = order.get(&a.provider_id); + let b_order = order.get(&b.provider_id); + match (a_order, b_order) { + (Some(a_order), Some(b_order)) if a_order != b_order => a_order.cmp(b_order), + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + _ => a.display_name.cmp(&b.display_name), + } + }); + ordered +} + +fn provider_status_label( + snapshot: &ProviderUsageSnapshot, + settings: &Settings, + language: Language, +) -> (String, String) { + let provider = codexbar::core::ProviderId::from_cli_name(&snapshot.provider_id); + let preference = provider + .map(|id| settings.get_provider_metric(id)) + .unwrap_or_default(); + if preference == MetricPreference::MonthlyPlan + && let Some(cost) = snapshot.cost.as_ref() + { + let amount = if !cost.formatted_used.is_empty() { + cost.formatted_used.clone() + } else { + crate::commands::format_cost_amount(cost) + }; + return ( + snapshot.provider_id.clone(), + format!("{} {}", snapshot.display_name, amount), + ); + } + + let label = crate::commands::compact_tray_status_label(headline_window(snapshot), language); + ( + snapshot.provider_id.clone(), + format!("{} {}", snapshot.display_name, label), + ) +} + +/// Window that headline tray surfaces should label for a provider. +pub(crate) fn headline_window(snapshot: &ProviderUsageSnapshot) -> &RateWindowSnapshot { + if snapshot.provider_id == "codex" { + codex_lane_headline_window(snapshot) + } else { + &snapshot.primary + } +} + +/// Pick the first non-informational Codex lane in session, weekly, monthly order. +pub(crate) fn codex_lane_headline_window(snapshot: &ProviderUsageSnapshot) -> &RateWindowSnapshot { + if !snapshot.primary.is_informational { + return &snapshot.primary; + } + if let Some(ref secondary) = snapshot.secondary + && !secondary.is_informational + { + return secondary; + } + if let Some(ref tertiary) = snapshot.tertiary + && !tertiary.is_informational + { + return tertiary; + } + &snapshot.primary +} + +/// Resolve a stable top/bottom pair while retaining stale saved preferences. +fn pick_stacked_tray_providers<'a>( + healthy: &[&'a ProviderUsageSnapshot], + settings: &Settings, +) -> Option<(&'a ProviderUsageSnapshot, &'a ProviderUsageSnapshot)> { + if healthy.len() < 2 { + return None; + } + + let preferred = |provider_id: Option<&str>| { + provider_id.and_then(|id| { + healthy + .iter() + .copied() + .find(|snapshot| snapshot.provider_id == id) + }) + }; + let preferred_bottom = preferred(settings.stacked_tray_bottom_provider.as_deref()); + let top = preferred(settings.stacked_tray_top_provider.as_deref()).or_else(|| { + healthy.iter().copied().find(|snapshot| { + preferred_bottom.map(|bottom| bottom.provider_id.as_str()) + != Some(snapshot.provider_id.as_str()) + }) + })?; + let bottom = preferred_bottom + .filter(|snapshot| snapshot.provider_id != top.provider_id) + .or_else(|| { + healthy + .iter() + .copied() + .find(|snapshot| snapshot.provider_id != top.provider_id) + })?; + + Some((top, bottom)) +} + +fn pick_tray_provider<'a>( + healthy: &[&'a ProviderUsageSnapshot], + prefer_highest: bool, +) -> Option<&'a ProviderUsageSnapshot> { + if prefer_highest { + healthy.iter().copied().max_by(|a, b| { + a.primary + .used_percent + .partial_cmp(&b.primary.used_percent) + .unwrap_or(std::cmp::Ordering::Equal) + }) + } else { + healthy.first().copied() + } +} + +fn selected_tray_percents( + snapshot: &ProviderUsageSnapshot, + settings: &Settings, +) -> (f64, Option) { + let (selected, companion) = + crate::usage_metric::selected_usage_icon_windows(snapshot, settings); + ( + display_metric_percent(&selected, settings.show_as_used), + companion + .as_ref() + .map(|window| display_metric_percent(window, settings.show_as_used)), + ) +} + +fn display_metric_percent(window: &RateWindowSnapshot, show_as_used: bool) -> f64 { + if window.is_informational { + return 0.0; + } + if window.is_exhausted || window.used_percent >= 100.0 { + return if show_as_used { 100.0 } else { 0.0 }; + } + + let used = window.used_percent.clamp(0.0, 100.0); + if show_as_used { used } else { 100.0 - used } +} + +#[cfg(test)] +#[path = "tray_presentation_tests.rs"] +mod tests; diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs new file mode 100644 index 0000000000..3fce3f03d7 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs @@ -0,0 +1,686 @@ +use super::*; +use codexbar::core::{ProviderId, ProviderStateKind}; + +fn fake_snapshot(id: &str, display_name: &str, used_percent: f64) -> ProviderUsageSnapshot { + fake_snapshot_with(id, display_name, used_percent, None, None, None) +} + +fn fake_snapshot_with( + id: &str, + display_name: &str, + used_percent: f64, + secondary_percent: Option, + tertiary_percent: Option, + cost: Option<(f64, f64)>, +) -> ProviderUsageSnapshot { + let window = |percent: f64| RateWindowSnapshot { + used_percent: percent, + remaining_percent: 100.0 - percent, + 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, + }; + + ProviderUsageSnapshot { + provider_id: id.into(), + display_name: display_name.into(), + primary: window(used_percent), + primary_label: None, + secondary: secondary_percent.map(window), + secondary_label: None, + model_specific: None, + tertiary: tertiary_percent.map(window), + tertiary_label: None, + extra_rate_windows: Vec::new(), + inventory: Vec::new(), + display_details: Vec::new(), + cost: cost.map(|(used, limit)| crate::commands::CostSnapshotBridge { + used, + limit: Some(limit), + remaining: Some((limit - used).max(0.0)), + currency_code: "USD".to_string(), + currency_symbol: None, + period: "monthly".to_string(), + resets_at: None, + formatted_used: format!("${used:.2}"), + formatted_limit: Some(format!("${limit:.2}")), + balance: None, + formatted_balance: None, + balance_updated_at: None, + account_id: None, + daily: Vec::new(), + always_visible: false, + }), + plan_name: None, + account_email: None, + subscription: None, + source_label: String::new(), + has_successful_claude_cli_quota: false, + updated_at: "2025-01-01T00:00:00Z".into(), + error: None, + error_state: ProviderStateKind::Ready, + pace: None, + account_organization: None, + tray_status_label: None, + fetch_duration_ms: None, + wayfinder_usage: None, + session_equivalent_forecast: None, + } +} + +#[test] +fn single_plan_uses_highest_provider_for_icon_and_summary() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Single, + menu_bar_shows_highest_usage: true, + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 72.0, + secondary_percent: None, + has_error: false, + } + ); + assert_eq!( + plan.status_labels(Language::English), + vec![("status_summary".to_string(), "Claude 72%".to_string())] + ); +} + +#[test] +fn single_plan_borrows_selected_snapshot_from_stable_input() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Single, + menu_bar_shows_highest_usage: true, + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + // `resolve` drops its temporary ordered/healthy vectors before returning. + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); +} + +#[test] +fn per_provider_plan_preserves_configured_order_for_status_rows() { + let settings = Settings { + tray_icon_mode: TrayIconMode::PerProvider, + provider_order: codexbar::settings::normalize_provider_order(&[ + "claude".to_string(), + "codex".to_string(), + ]), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let labels = + TrayPresentationPlan::resolve(&settings, &snapshots).status_labels(Language::English); + + assert_eq!( + labels, + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); +} + +#[test] +fn stacked_plan_resolves_distinct_preferences_once() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + fake_snapshot("gemini", "Gemini", 44.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 72.0, + bottom_percent: 30.0, + has_error: false, + } + ); + assert_eq!( + plan.status_labels(Language::English), + vec![ + ("claude".to_string(), "Claude 72%".to_string()), + ("codex".to_string(), "Codex 30%".to_string()), + ] + ); +} + +#[test] +fn stacked_plan_borrows_both_snapshots_from_stable_input() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + // The plan retains references to the caller-owned snapshots, not the + // temporary vector of references used during selection. + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert!(std::ptr::eq(plan.status_rows[0].snapshot, &snapshots[1])); + assert!(std::ptr::eq(plan.status_rows[1].snapshot, &snapshots[0])); +} + +#[test] +fn stacked_plan_falls_back_around_stale_and_duplicate_preferences() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("missing".to_string()), + stacked_tray_bottom_provider: Some("claude".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 30.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Stacked { + top_percent: 30.0, + bottom_percent: 72.0, + has_error: false, + } + ); + assert_eq!(plan.status_rows[0].snapshot.provider_id, "codex"); + assert_eq!(plan.status_rows[1].snapshot.provider_id, "claude"); +} + +#[test] +fn one_provider_stacked_mode_falls_back_to_single_provider_bars() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + ..Settings::default() + }; + let snapshots = vec![fake_snapshot_with( + "codex", + "Codex", + 30.0, + Some(65.0), + None, + None, + )]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 65.0, + secondary_percent: Some(30.0), + has_error: false, + } + ); + assert_eq!(plan.status_rows.len(), 1); +} + +#[test] +fn one_healthy_provider_never_uses_stacked_renderer() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + menu_bar_shows_percent: true, + ..Settings::default() + }; + let healthy = fake_snapshot("codex", "Codex", 30.0); + let mut failed = fake_snapshot("claude", "Claude", 72.0); + failed.error = Some("offline".to_string()); + let snapshots = vec![healthy, failed]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Percent { + percent: 30.0, + has_error: false, + } + ); + assert_eq!(plan.status_rows.len(), 1); + assert_eq!(plan.status_rows[0].snapshot.provider_id, "codex"); +} + +#[test] +fn all_errors_produce_error_styled_zero_percent_plan() { + let settings = Settings { + menu_bar_shows_percent: true, + ..Settings::default() + }; + let mut snapshot = fake_snapshot("codex", "Codex", 30.0); + snapshot.error = Some("offline".to_string()); + let snapshots = vec![snapshot]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Percent { + percent: 0.0, + has_error: true, + } + ); + assert!(plan.status_rows.is_empty()); +} + +#[test] +fn plan_uses_selected_metric_and_remaining_display_mode() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshots = vec![fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + )]; + + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.icon, + TrayIconPlan::Bars { + primary_percent: 85.0, + secondary_percent: Some(80.0), + has_error: false, + } + ); +} + +#[test] +fn render_icon_delegates_to_resolved_stacked_renderer() { + let settings = Settings { + tray_icon_mode: TrayIconMode::Stacked, + stacked_tray_top_provider: Some("claude".to_string()), + stacked_tray_bottom_provider: Some("codex".to_string()), + ..Settings::default() + }; + let snapshots = vec![ + fake_snapshot("codex", "Codex", 40.0), + fake_snapshot("claude", "Claude", 72.0), + ]; + let plan = TrayPresentationPlan::resolve(&settings, &snapshots); + + assert_eq!( + plan.render_icon(), + render_stacked_bar_icon_rgba(72.0, 40.0, false) + ); +} + +#[test] +fn codex_headline_skips_informational_primary() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); + snapshot.primary.is_informational = true; + + assert_eq!(codex_lane_headline_window(&snapshot).used_percent, 25.0); +} +fn fake_extra_window(percent: f64) -> crate::commands::NamedRateWindowSnapshot { + crate::commands::NamedRateWindowSnapshot { + id: "additional_budget".to_string(), + title: "Additional Budget".to_string(), + fallback_lane: false, + window: crate::commands::RateWindowSnapshot { + used_percent: percent, + remaining_percent: 100.0 - percent, + 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, + }, + } +} + +#[test] +fn selected_tray_percent_uses_cursor_extra_usage_cost() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + ); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 15.0); + assert_eq!(secondary, Some(20.0)); +} + +#[test] +fn selected_tray_percent_tracks_extra_rate_window() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Copilot, MetricPreference::ExtraUsage); + let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); + snapshot.extra_rate_windows.push(fake_extra_window(42.0)); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, None); +} + +#[test] +fn copilot_automatic_tracks_highest_extra_rate_window() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot("copilot", "Copilot", 20.0); + snapshot.extra_rate_windows.push(fake_extra_window(42.0)); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); +} + +#[test] +fn selected_tray_percent_respects_remaining_display_mode() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with( + "cursor", + "Cursor", + 10.0, + Some(20.0), + Some(72.0), + Some((15.0, 100.0)), + ); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 85.0); + assert_eq!(secondary, Some(80.0)); +} + +#[test] +fn exhausted_automatic_window_never_renders_as_remaining_progress() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with( + "opencodego", + "OpenCode Go", + 20.0, + Some(60.0), + Some(40.0), + None, + ); + snapshot + .tertiary + .as_mut() + .expect("monthly quota") + .is_exhausted = true; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(remaining, 0.0); + + settings.show_as_used = true; + let (used, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(used, 100.0); +} + +#[test] +fn full_automatic_window_without_exhausted_flag_has_zero_remaining_progress() { + let mut settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with( + "opencodego", + "OpenCode Go", + 20.0, + Some(60.0), + Some(100.0), + None, + ); + snapshot + .tertiary + .as_mut() + .expect("monthly quota") + .is_exhausted = false; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(remaining, 0.0); + + settings.show_as_used = true; + let (used, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(used, 100.0); +} + +#[test] +fn missing_automatic_window_does_not_look_like_available_remaining_progress() { + let settings = Settings { + show_as_used: false, + ..Settings::default() + }; + let mut snapshot = fake_snapshot_with("opencodego", "OpenCode Go", 0.0, None, None, None); + snapshot.primary.is_informational = true; + + let (remaining, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(remaining, 0.0); +} + +#[test] +fn selected_tray_percent_falls_back_when_extra_usage_missing() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::ExtraUsage); + let snapshot = fake_snapshot_with("cursor", "Cursor", 10.0, Some(72.0), None, None); + + let (primary, _) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 72.0); +} + +#[test] +fn single_meaningful_secondary_quota_uses_full_single_meter() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); + snapshot.primary.is_informational = true; + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, None); +} + +#[test] +fn selected_secondary_quota_is_not_duplicated_when_tertiary_is_meaningful() { + let settings = Settings::default(); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), Some(30.0), None); + snapshot.primary.is_informational = true; + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 42.0); + assert_eq!(secondary, Some(30.0)); +} + +#[test] +fn two_meaningful_quotas_keep_two_meter_layout() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Cursor, MetricPreference::Session); + let snapshot = fake_snapshot_with("cursor", "Cursor", 15.0, Some(40.0), None, None); + + let (primary, secondary) = selected_tray_percents(&snapshot, &settings); + + assert_eq!(primary, 15.0); + assert_eq!(secondary, Some(40.0)); +} + +#[test] +fn informational_primary_skips_session_and_automatic_phantom_zero() { + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::Claude, MetricPreference::Session); + let mut snapshot = fake_snapshot_with("claude", "Claude", 0.0, Some(42.0), None, None); + snapshot.primary.is_informational = true; + + // Session preference must not paint the synthetic 0% primary; + // it falls through to Automatic which prefers weekly (42%). + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 42.0); + assert_ne!(primary, 0.0); + + // Automatic also prefers weekly over informational primary. + settings.set_provider_metric(ProviderId::Claude, MetricPreference::Automatic); + let (primary, _) = selected_tray_percents(&snapshot, &settings); + assert_eq!(primary, 42.0); +} + +#[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); +} + +#[test] +fn f5_headline_prefers_non_informational_primary() { + let snapshot = fake_snapshot_with("codex", "Codex", 50.0, Some(20.0), Some(30.0), None); + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 50.0).abs() < f64::EPSILON); +} + +#[test] +fn f5_headline_falls_back_to_secondary_when_primary_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(25.0), Some(30.0), None); + snapshot.primary.is_informational = true; + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 25.0).abs() < f64::EPSILON); +} + +#[test] +fn f5_headline_falls_back_to_tertiary_when_primary_and_secondary_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(35.0), None); + snapshot.primary.is_informational = true; + snapshot.secondary.as_mut().unwrap().is_informational = true; + let headline = codex_lane_headline_window(&snapshot); + assert!((headline.used_percent - 35.0).abs() < f64::EPSILON); +} + +#[test] +fn f5_headline_returns_primary_when_all_informational() { + let mut snapshot = fake_snapshot_with("codex", "Codex", 0.0, Some(0.0), Some(0.0), None); + snapshot.primary.is_informational = true; + if let Some(sec) = &mut snapshot.secondary { + sec.is_informational = true; + } + if let Some(ter) = &mut snapshot.tertiary { + ter.is_informational = true; + } + let headline = codex_lane_headline_window(&snapshot); + // Falls back to primary (the placeholder) when all are informational. + assert!(headline.is_informational); +} diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg deleted file mode 100644 index fdde018b8f..0000000000 --- a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-crof.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.test.ts b/apps/desktop-tauri/src/components/providers/providerIcons.test.ts index 27c677024b..a204e1155a 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.test.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.test.ts @@ -8,4 +8,8 @@ describe("provider icon registry", () => { expect(PROVIDER_ICON_REGISTRY[id], id).toBeDefined(); } }); + + it("does not expose the retired Crof provider", () => { + expect(PROVIDER_ICON_REGISTRY).not.toHaveProperty("crof"); + }); }); diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 2fdc4380d7..aefe0035aa 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -14,7 +14,6 @@ import coderabbit from "./icons/ProviderIcon-coderabbit.svg?raw"; import codex from "./icons/ProviderIcon-codex.svg?raw"; import commandcode from "./icons/ProviderIcon-commandcode.svg?raw"; import copilot from "./icons/ProviderIcon-copilot.svg?raw"; -import crof from "./icons/ProviderIcon-crof.svg?raw"; import crossmodel from "./icons/ProviderIcon-crossmodel.svg?raw"; import cursor from "./icons/ProviderIcon-cursor.svg?raw"; import deepgram from "./icons/ProviderIcon-deepgram.svg?raw"; @@ -102,7 +101,6 @@ const RAW: Record = { codex: tint(codex), commandcode: tint(commandcode), copilot: tint(copilot), - crof: tint(crof), crossmodel: tint(crossmodel), cursor: tint(cursor), deepgram: tint(deepgram), @@ -221,7 +219,6 @@ export const PROVIDER_ICON_REGISTRY: Record = { mimo: { id: "mimo", brandColor: "#ff6900", fallbackLetter: "M", svgPath: RAW.mimo }, doubao: { id: "doubao", brandColor: "#2563eb", fallbackLetter: "D", svgPath: RAW.doubao }, commandcode: { id: "commandcode", brandColor: "#44ff00", fallbackLetter: "C", svgPath: RAW.commandcode }, - crof: { id: "crof", brandColor: "#7c3aed", fallbackLetter: "C", svgPath: RAW.crof }, crossmodel: { id: "crossmodel", brandColor: "#c084fc", fallbackLetter: "X", svgPath: RAW.crossmodel }, qoder: { id: "qoder", brandColor: "#2563eb", fallbackLetter: "Q", svgPath: RAW.qoder }, replicate: { id: "replicate", brandColor: "#000000", fallbackLetter: "R", svgPath: RAW.replicate }, diff --git a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx index 7dbf6d6f7d..1e70c015ff 100644 --- a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx +++ b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx @@ -258,7 +258,7 @@ describe("useTrayPanelLayout sizing", () => { await nudgePass(result, 417, "421px"); // → 421 → 526 phys expect(lastResize()).toEqual({ width: 328, height: 421 }); expect(surface.style.maxHeight).toBe("421px"); - }); + }, 30_000); // 8 bounded 3s settling passes + 3s readiness can exceed Vitest's 5s default. it("reconciles to the applied physical frame after an OS snap (no churn, no cycle)", async () => { // Deliberate 5-physical snap: requesting 539 logical (→674 phys) yields an diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index 90c0984f49..0ad26af95b 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -415,6 +415,7 @@ export const ALL_LOCALE_KEYS = [ "UsageSpendHourlyActivity", "UsageSpendRequests", "UsageSpendTokens", + "UsageSpendKnownSubtotal", "UsageSpendAllTimeHistory", "UsageSpendCustomPricing", "OverviewSpendTitle", @@ -540,6 +541,9 @@ export const ALL_LOCALE_KEYS = [ "TrayIconModeHelper", "TrayIconModeSingle", "TrayIconModePerProvider", + "TrayIconModeStacked", + "StackedTrayTopProvider", + "StackedTrayBottomProvider", "ShowProviderIcons", "ShowProviderIconsHelper", "PreferHighestUsage", diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index bf693ce7b8..13d30f1fdf 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -2,6 +2,8 @@ import { describe, expect, it } from "vitest"; import { formatUsageSpendReportingDay, + formatSpendMetric, + formatUsd, filterUsageSpendSummaryForOverview, renderUsageSpendSharePng, usageSpendShareFooter, @@ -10,6 +12,24 @@ import { import type { SpendContract, UsageSpendRow, UsageSpendSummary } from "../types/bridge"; describe("usage spend sharing", () => { + it("labels a mixed-pricing subtotal without presenting it as a total", () => { + expect(formatSpendMetric(null, 1_500, "USD", "tokens", 0.0125)).toMatch(/^≥.* known/); + }); + + it("uses the caller's localized subtotal label", () => { + const metric = formatSpendMetric(null, 500, "USD", "Token", 6, "{} (teilweise)"); + expect(metric).toBe(`${formatUsd(6, "USD")} (teilweise) · 500 Token`); + expect(metric).not.toContain("known"); + }); + + it("renders a complete known-zero total instead of a subtotal", () => { + const metric = formatSpendMetric(0, 0, "USD", "tokens", 9); + expect(metric).not.toBe("—"); + expect(metric).not.toContain("≥"); + expect(metric).not.toContain("9.00"); + expect(metric).toContain("0 tokens"); + }); + it.each([ [0, "0 subscriptions"], [1, "1 subscription"], diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 2246e98505..353686f540 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -124,15 +124,24 @@ export function formatUsd(value: number | null | undefined, currency: string): s } } -/** Canonical "cost · tokens" cell for spend tables and share renders. */ +/** + * Canonical "cost · tokens" cell for spend tables and share renders. + * `knownSubtotalTemplate` is the localized subtotal label; `{}` is the amount. + */ export function formatSpendMetric( cost: number | null | undefined, tokens: number | null | undefined, currency: string, tokenLabel: string, + knownSubtotal?: number | null, + knownSubtotalTemplate = "≥{} known", ): string { const parts: string[] = []; - if (cost != null && Number.isFinite(cost)) parts.push(formatUsd(cost, currency)); + if (cost != null && Number.isFinite(cost)) { + parts.push(formatUsd(cost, currency)); + } else if (knownSubtotal != null && Number.isFinite(knownSubtotal)) { + parts.push(knownSubtotalTemplate.replace("{}", formatUsd(knownSubtotal, currency))); + } if (tokens != null && Number.isFinite(tokens)) { parts.push(`${Math.max(0, tokens).toLocaleString()} ${tokenLabel}`); } @@ -202,8 +211,20 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri const y = y0 + (index + 1) * rowH; const cells = [ row.displayName, - formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, "tokens"), - formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, "tokens"), + formatSpendMetric( + row.sevenDay, + row.sevenDayTokens, + row.currency, + "tokens", + row.sevenDayEstimate?.knownSubtotalUsd, + ), + formatSpendMetric( + row.thirtyDay, + row.thirtyDayTokens, + row.currency, + "tokens", + row.thirtyDayEstimate?.knownSubtotalUsd, + ), row.currency || "USD", row.source, ]; diff --git a/apps/desktop-tauri/src/styles.css b/apps/desktop-tauri/src/styles.css index c3465d0ec0..36a304a58c 100644 --- a/apps/desktop-tauri/src/styles.css +++ b/apps/desktop-tauri/src/styles.css @@ -5741,31 +5741,6 @@ html:has(.menu-surface--tray) { color: inherit; } -.accent-color-swatch-row { - display: flex; - align-items: center; - gap: 8px; - margin-top: 8px; - font-size: 12px; -} - -.accent-color-swatch { - display: inline-block; - width: 16px; - height: 16px; - border-radius: 4px; - border: 1px solid var(--border-color); -} - -.accent-color-swatch-label { - color: var(--text-secondary); -} - -.accent-color-swatch-value { - font-family: var(--font-mono, monospace); - color: var(--text-secondary); -} - /* ── Mistral monthly spend row (#2821, #2947) ──────────────────── */ .menu-card__monthly-spend { margin-top: 4px; diff --git a/apps/desktop-tauri/src/surfaces/Settings.tsx b/apps/desktop-tauri/src/surfaces/Settings.tsx index 21b5828dda..f3fb995d44 100644 --- a/apps/desktop-tauri/src/surfaces/Settings.tsx +++ b/apps/desktop-tauri/src/surfaces/Settings.tsx @@ -253,7 +253,13 @@ export default function Settings({ state, initialTab: propTab }: { state: Bootst )} {activeTab === "menuBar" && ( - + )} {activeTab === "menu" && ( diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 4655c3b58f..949d275d26 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -322,31 +322,25 @@ describe("TrayPanel provider grid", () => { reportingDay: "2026-09-19", dashboardTimezone: "UTC", rows: [ - { - providerId: "codex", - displayName: "Codex", - sevenDay: 1, - thirtyDay: 2, - currency: "USD", - source: "local", - includedInOverview: true, - }, - { - providerId: "claude", - displayName: "Claude", - sevenDay: 3, - thirtyDay: 4, - currency: "USD", - source: "hidden", - includedInOverview: false, - }, + { providerId: "codex", displayName: "Codex", sevenDay: 1, thirtyDay: 2, + currency: "USD", source: "local", includedInOverview: true }, + { providerId: "claude", displayName: "Claude", sevenDay: 3, thirtyDay: 4, + currency: "USD", source: "hidden", includedInOverview: false }, + // A known subtotal is a partial estimate: counted for coverage, not in the total. + { providerId: "antigravity", displayName: "Antigravity", sevenDay: null, thirtyDay: null, + thirtyDayEstimate: { + knownSubtotalUsd: 9, + coverage: { priced: 0, unpriced: 1, unmetered: 0, estimated: 1 }, + }, + currency: "USD", source: "known subtotal", includedInOverview: true }, ], }); renderTrayPanel([provider("codex", "Codex", 35)]); expect(await screen.findByRole("button", { name: "UsageSpendShare" })).toBeInTheDocument(); - expect(screen.getByText(/1 of 1 OverviewSpendProviderCoverage/)).toBeInTheDocument(); + expect(screen.getByText("~$2.00")).toBeInTheDocument(); + expect(screen.getByText(/1 of 2 OverviewSpendProviderCoverage/)).toBeInTheDocument(); }); it("dismisses the tray panel on unmodified Escape", async () => { diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index 2cfdd609e1..034dfc09a0 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -30,7 +30,7 @@ import { 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", "zoommate", + "aiand", "commandcode", "copilot", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", "doubao", "elevenlabs", "factory", "gemini", "grok", "groq", "infini", "jetbrains", "kilo", "kimi", "kimik2", "kiro", "manus", "replicate", "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx new file mode 100644 index 0000000000..6a5a385660 --- /dev/null +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.test.tsx @@ -0,0 +1,21 @@ +import { render } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { AccentColorSection } from "./AccentColorSection"; + +describe("AccentColorSection", () => { + it("uses the native color input as the only color preview", () => { + const { container } = render( + key} + onChange={vi.fn()} + />, + ); + + expect(container.querySelector('input[type="color"]')).toHaveValue("#123456"); + expect(container.querySelector(".accent-color-swatch-row")).toBeNull(); + expect(container.querySelector(".accent-color-swatch")).toBeNull(); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx index 9e0f8dca41..112e09f856 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/AccentColorSection.tsx @@ -94,16 +94,6 @@ export function AccentColorSection({ {t("ProviderAccentColorReset")} -
- - {t("ProviderAccentColor")} - - - {effective} -
{error &&

{error}

} ); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx index 77ed218f79..e1a95c64aa 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/RegionSection.tsx @@ -12,7 +12,7 @@ interface Props { } /** - * API-region dropdown for Alibaba / Z.ai / MiniMax. + * API-region dropdown for providers with regional endpoints, including Kimi. * * Port of the region ComboBox rows in * `rust/src/native_ui/preferences.rs::render_provider_detail_panel`. diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx index bff42a55b3..d11d431091 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.test.tsx @@ -14,7 +14,10 @@ import DisplayTab from "./DisplayTab"; import type { SettingsSnapshot } from "../../../types/bridge"; const baseSettings = { + enabledProviders: ["codex", "claude"], trayIconMode: "single", + stackedTrayTopProvider: null, + stackedTrayBottomProvider: null, trayPanelAlwaysOnTop: false, switcherShowsIcons: false, menuBarShowsHighestUsage: false, @@ -84,3 +87,28 @@ describe("DisplayTab menu settings", () => { expect(set).toHaveBeenCalledWith({ trayPanelAlwaysOnTop: true }); }); }); + +describe("DisplayTab stacked tray providers", () => { + it("persists explicit top and bottom provider choices", () => { + const set = vi.fn(); + render( + , + ); + const selects = screen.getAllByRole("combobox"); + + fireEvent.change(selects[1], { target: { value: "claude" } }); + fireEvent.change(selects[2], { target: { value: "codex" } }); + + expect(set).toHaveBeenCalledWith({ stackedTrayTopProvider: "claude" }); + expect(set).toHaveBeenCalledWith({ stackedTrayBottomProvider: "codex" }); + }); +}); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx index 52588deb84..1955c5ebf2 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/DisplayTab.tsx @@ -4,6 +4,7 @@ import { Field, Select, Toggle } from "../../../components/FormControls"; import type { MenuBarDisplayMode, OverviewLayout, + ProviderCatalogEntry, TrayIconMode, TrayVisibilityStatusDto, } from "../../../types/bridge"; @@ -16,7 +17,11 @@ export default function DisplayTab({ settings, set, saving, -}: TabProps & { mode?: "menuBar" | "menu" }) { + providers = [], +}: TabProps & { + mode?: "menuBar" | "menu"; + providers?: ProviderCatalogEntry[]; +}) { const { t } = useLocale(); const [trayVisibility, setTrayVisibility] = useState(null); @@ -26,6 +31,13 @@ export default function DisplayTab({ .catch(() => setTrayVisibility(null)); }, []); + const providerName = new Map( + providers.map((provider) => [provider.id, provider.displayName]), + ); + const stackedProviderOptions = settings.enabledProviders.map((providerId) => ({ + value: providerId, + label: providerName.get(providerId) ?? providerId, + })); return ( <> {/* ── Menu bar ─────────────────────────────────────────────── */} @@ -42,10 +54,47 @@ export default function DisplayTab({ options={[ { value: "single", label: t("TrayIconModeSingle") }, { value: "perProvider", label: t("TrayIconModePerProvider") }, + { value: "stacked", label: t("TrayIconModeStacked") }, ]} onChange={(v) => set({ trayIconMode: v as TrayIconMode })} /> + {settings.trayIconMode === "stacked" && ( + <> + + + provider.value !== settings.stackedTrayTopProvider, + ), + ]} + onChange={(provider) => + set({ stackedTrayBottomProvider: provider }) + } + /> + + + )} set({ menuBarShowsHighestUsage: v })} /> @@ -75,7 +124,7 @@ export default function DisplayTab({ > set({ menuBarShowsPercent: v })} /> diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx index 48c200779c..ccf5295187 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/ProvidersTab.tsx @@ -242,7 +242,6 @@ function providerSourceHintShort( case "clinepass": case "neuralwatt": case "doubao": - case "crof": case "stepfun": case "venice": case "openaiapi": diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx index 0f7cc7ee89..a037d978e5 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx @@ -234,8 +234,26 @@ export default function UsageSpendTab(_props: TabProps) { {(summary?.rows ?? []).map((row) => ( {row.displayName} - {formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, t("UsageSpendTokens"))} - {formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, t("UsageSpendTokens"))} + + {formatSpendMetric( + row.sevenDay, + row.sevenDayTokens, + row.currency, + t("UsageSpendTokens"), + row.sevenDayEstimate?.knownSubtotalUsd, + t("UsageSpendKnownSubtotal"), + )} + + + {formatSpendMetric( + row.thirtyDay, + row.thirtyDayTokens, + row.currency, + t("UsageSpendTokens"), + row.thirtyDayEstimate?.knownSubtotalUsd, + t("UsageSpendKnownSubtotal"), + )} + {row.currency || "USD"} {row.source} diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index dda4c24441..be41e7b35e 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -48,7 +48,6 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["mimo", "Xiaomi MiMo"], ["doubao", "Doubao"], ["commandcode", "Command Code"], - ["crof", "Crof"], ["stepfun", "StepFun"], ["venice", "Venice"], ["openaiapi", "OpenAI API"], diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 88754baa7d..fd8dd55109 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -12,7 +12,7 @@ export type SettingsTabId = // ── Narrowed string-literal unions (persisted settings enums) ───────── -export type TrayIconMode = "single" | "perProvider"; +export type TrayIconMode = "single" | "perProvider" | "stacked"; export type NotificationSoundTheme = "windows" | "codexBar"; @@ -190,6 +190,8 @@ export interface SettingsSnapshot { predictivePaceWarningEnabled: boolean; showPace?: boolean; trayIconMode: TrayIconMode; + stackedTrayTopProvider?: string | null; + stackedTrayBottomProvider?: string | null; switcherShowsIcons: boolean; menuBarShowsHighestUsage: boolean; menuBarShowsPercent: boolean; @@ -299,6 +301,8 @@ export interface SettingsUpdate { predictivePaceWarningEnabled?: boolean; showPace?: boolean; trayIconMode?: TrayIconMode; + stackedTrayTopProvider?: string; + stackedTrayBottomProvider?: string; switcherShowsIcons?: boolean; menuBarShowsHighestUsage?: boolean; menuBarShowsPercent?: boolean; @@ -379,6 +383,8 @@ export interface UsageSpendRow { displayName: string; sevenDay: number | null; thirtyDay: number | null; + sevenDayEstimate?: LocalCostEstimate; + thirtyDayEstimate?: LocalCostEstimate; sevenDayTokens?: number | null; thirtyDayTokens?: number | null; currency: string; @@ -391,6 +397,11 @@ export interface UsageSpendRow { staleUpdatedAt?: string; } +export interface LocalCostEstimate { + knownSubtotalUsd: number | null; + coverage: CostCoverageCounts; +} + export interface UsageSpendSummary { rows: UsageSpendRow[]; contract: SpendContract; diff --git a/rust/assets/icons/ProviderIcon-crof.svg b/rust/assets/icons/ProviderIcon-crof.svg deleted file mode 100644 index fdde018b8f..0000000000 --- a/rust/assets/icons/ProviderIcon-crof.svg +++ /dev/null @@ -1,3 +0,0 @@ - - - diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 8e9be16dc2..819b7abb59 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -286,7 +286,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b println!("{title}"); } - if let Some(history) = result.token_history { + if let Some(history) = result.token_history.as_ref() { print_local_token_history(history, days); } else if group_by == CostGroupBy::Session && result.provider == "codex" { print_codex_session_output(result, days); @@ -372,7 +372,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b } } -fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySummary, days: u32) { +fn print_local_token_history(history: &crate::spend_contract::LocalTokenHistorySummary, days: u32) { use crate::spend_contract::LocalHistoryCoverage; match history.coverage { LocalHistoryCoverage::Complete if history.total_tokens == 0 => { @@ -386,7 +386,20 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - println!(" Local token history; dollar costs unavailable"); + if let Some(cost) = history.total_usd() { + println!(" API list-price estimate: ${cost:.2} (not billed spend)"); + } else if let Some(cost) = history.cost_estimate.known_subtotal_usd { + if history.coverage == LocalHistoryCoverage::Complete { + println!( + " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", + history.cost_estimate.coverage.unpriced + ); + } else { + println!(" Known API list-price subtotal: ${cost:.2} (history incomplete)"); + } + } else { + println!(" Local token history; dollar costs unavailable"); + } } fn print_codex_session_output(result: &CostResult, days: u32) { @@ -459,7 +472,7 @@ fn build_json_payloads(results: &[CostResult], days: u32) -> Vec - - diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index d1841db6bc..53e0f62d07 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -71,17 +71,16 @@ pub async fn cost_response(provider: Option<&str>) -> String { let history = crate::providers::antigravity::local_sessions::summarize(30); results.push(crate::spend_contract::local_token_history_json( "antigravity", - history, + &history, 30, )); continue; } if provider_id == ProviderId::Muse { let report = crate::providers::muse::local_usage::scan(30, None); + let history = report.into(); results.push(crate::spend_contract::local_token_history_json( - "muse", - report.into(), - 30, + "muse", &history, 30, )); continue; } @@ -160,10 +159,11 @@ mod tests { use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; let complete = crate::spend_contract::local_token_history_json( "antigravity", - LocalTokenHistorySummary { + &LocalTokenHistorySummary { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, + cost_estimate: Default::default(), }, 30, ); @@ -173,10 +173,11 @@ mod tests { let partial = crate::spend_contract::local_token_history_json( "antigravity", - LocalTokenHistorySummary { + &LocalTokenHistorySummary { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, + cost_estimate: Default::default(), }, 30, ); diff --git a/rust/src/codex_accounts/account_manager/tests.rs b/rust/src/codex_accounts/account_manager/tests.rs index bfb8584f6c..aefd6f9718 100644 --- a/rust/src/codex_accounts/account_manager/tests.rs +++ b/rust/src/codex_accounts/account_manager/tests.rs @@ -5,9 +5,13 @@ mod tests { /// Write an auth.json carrying a JWT identity for the given account id. fn write_auth(home_path: &Path, email: &str, account_id: &str) { + write_user_auth(home_path, email, account_id, &format!("auth0|{account_id}")); + } + + fn write_user_auth(home_path: &Path, email: &str, account_id: &str, subject: &str) { let payload = serde_json::json!({ "email": email, - "sub": format!("auth0|{account_id}"), + "sub": subject, "https://api.openai.com/auth": { "chatgpt_plan_type": "team", "chatgpt_account_id": account_id, @@ -88,6 +92,50 @@ mod tests { super::super::file_locations::clear_app_support_directory_override(); } + #[test] + fn shared_team_users_keep_separate_discovery_and_managed_homes() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path(); + super::super::file_locations::with_app_support_directory(root.to_path_buf()); + let first_home = root.join("managed-homes").join("first"); + let second_home = root.join("managed-homes").join("second"); + let ambient_home = root.join("ambient"); + for home in [&first_home, &second_home, &ambient_home] { + std::fs::create_dir_all(home).unwrap(); + } + write_user_auth(&first_home, "a@x.test", "shared-team", "user-a"); + write_user_auth(&second_home, "b@x.test", "shared-team", "user-b"); + write_user_auth(&ambient_home, "a@x.test", "shared-team", "user-a"); + let second_auth = std::fs::read(second_home.join("auth.json")).unwrap(); + let mut first = make_account(first_home.clone(), "a@x.test", "shared-team"); + first.auth_subject = Some("user-a".into()); + let manager = CodexAccountManager::new(); + let discovered = manager + .discover_managed_accounts(std::slice::from_ref(&first)) + .unwrap(); + assert_eq!(discovered.len(), 2); + let second = discovered + .iter() + .find(|a| a.codex_home_path == second_home) + .unwrap(); + assert_ne!(second.id, first.id); + assert!(!first.matches(second)); + + let mut ambient = first.clone(); + ambient.codex_home_path = ambient_home; + manager.remove_managed_files_if_owned(&first).unwrap(); + assert!(!first_home.exists()); + assert!(second_home.exists()); + let materialized = manager.materialize_as_managed(&ambient).unwrap(); + assert_ne!(materialized.codex_home_path, second_home); + assert_eq!( + std::fs::read(second_home.join("auth.json")).unwrap(), + second_auth + ); + assert_eq!(managed_home_count(root), 2); + super::super::file_locations::clear_app_support_directory_override(); + } + /// Write an auth.json whose credentials carry an explicit refresh time. fn write_auth_refreshed_at( home_path: &Path, diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs index 40c29abb60..53757fb844 100644 --- a/rust/src/codex_accounts/models.rs +++ b/rust/src/codex_accounts/models.rs @@ -61,6 +61,25 @@ fn normalize_identifier(value: Option<&str>) -> Option { .filter(|v| !v.is_empty()) } +/// A shared workspace or home cannot override conflicting user evidence. +fn user_identity_conflicts( + subject: Option<&str>, + other_subject: Option<&str>, + email: Option<&str>, + other_email: Option<&str>, +) -> bool { + if let (Some(a), Some(b)) = ( + normalize_identifier(subject), + normalize_identifier(other_subject), + ) { + return a != b; + } + matches!( + (normalize_identifier(email), normalize_identifier(other_email)), + (Some(a), Some(b)) if a != b + ) +} + /// Where an account's `CODEX_HOME` lives. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -206,6 +225,14 @@ impl CodexAccount { /// Whether two accounts refer to the same identity. pub fn matches(&self, other: &CodexAccount) -> bool { + if user_identity_conflicts( + self.auth_subject.as_deref(), + other.auth_subject.as_deref(), + self.email_hint.as_deref(), + other.email_hint.as_deref(), + ) { + return false; + } if let (Some(a), Some(b)) = ( self.effective_workspace_account_id(), other.effective_workspace_account_id(), @@ -323,6 +350,14 @@ impl RemovedAccountIdentity { } pub fn matches(&self, account: &CodexAccount) -> bool { + if user_identity_conflicts( + self.auth_subject.as_deref(), + account.auth_subject.as_deref(), + self.email_hint.as_deref(), + account.email_hint.as_deref(), + ) { + return false; + } if self.standardized_home_path() == account.standardized_home_path() { return true; } @@ -573,6 +608,74 @@ mod tests { assert!(a.matches(&b)); } + #[test] + fn different_subjects_do_not_match_in_a_shared_workspace_or_home() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/shared", + CodexAccountSource::ManagedByApp, + Some("shared-team"), + ); + a.auth_subject = Some("user-a".into()); + a.email_hint = Some("same@x.test".into()); + let mut b = a.clone(); + b.id = Uuid::new_v4(); + b.auth_subject = Some("user-b".into()); + + assert!(!a.matches(&b)); + assert!(!b.matches(&a)); + assert!(!RemovedAccountIdentity::from_account(&a).matches(&b)); + } + + #[test] + fn different_emails_do_not_match_in_a_shared_workspace_without_subjects() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/a", + CodexAccountSource::ManagedByApp, + Some("shared-team"), + ); + a.email_hint = Some("user-a@x.test".into()); + let mut b = a.clone(); + b.id = Uuid::new_v4(); + b.codex_home_path = PathBuf::from("/managed/b"); + b.email_hint = Some("user-b@x.test".into()); + + assert!(!a.matches(&b)); + assert!(!RemovedAccountIdentity::from_account(&a).matches(&b)); + } + + #[test] + fn matching_subject_remains_authoritative_when_email_changes() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/a", + CodexAccountSource::ManagedByApp, + Some("shared-team"), + ); + a.auth_subject = Some("user-a".into()); + a.email_hint = Some("old@x.test".into()); + let mut b = a.clone(); + b.auth_subject = Some("USER-A".into()); + b.email_hint = Some("new@x.test".into()); + assert!(a.matches(&b)); + } + + #[test] + fn same_user_in_different_workspaces_remains_separate() { + let mut a = account( + "11111111-1111-1111-1111-111111111111", + "/managed/a", + CodexAccountSource::ManagedByApp, + Some("team-a"), + ); + a.auth_subject = Some("same-user".into()); + let mut b = a.clone(); + b.provider_account_id = Some("team-b".into()); + assert!(!a.matches(&b)); + assert!(!b.matches(&a)); + } + #[test] fn disambiguates_different_provider_ids() { let a = account( diff --git a/rust/src/core/claude_routed_pricing.rs b/rust/src/core/claude_routed_pricing.rs index 1076992c9f..1f6a750b69 100644 --- a/rust/src/core/claude_routed_pricing.rs +++ b/rust/src/core/claude_routed_pricing.rs @@ -147,13 +147,52 @@ pub fn cost_usd_from_pricing_with_threshold( cache_write: i32, output: i32, ) -> f64 { - let input = input.max(0); - let cache_read = cache_read.max(0); - let cache_write = cache_write.max(0); - let output = output.max(0); + cost_usd_from_u64_counts_with_threshold( + pricing, + threshold_tokens, + input.max(0) as u64, + cache_read.max(0) as u64, + cache_write.max(0) as u64, + output.max(0) as u64, + ) +} + +/// Calculate routed cost for local history counters without narrowing them to +/// the signed API token-count type. +pub(crate) fn cost_usd_from_u64_counts_with_threshold( + pricing: models_dev_pricing::DynamicModelPricing, + threshold_tokens: Option, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> f64 { let use_tier = threshold_tokens.is_some_and(|threshold| { - (input as u64) + (cache_read as u64) + (cache_write as u64) > threshold + input + .checked_add(cache_read) + .and_then(|total| total.checked_add(cache_write)) + .is_none_or(|total| total > threshold) }); + let rates = selected_cost_rates(pricing, use_tier); + + (input as f64) * rates.input + + (cache_read as f64) * rates.cache_read + + (cache_write as f64) * rates.cache_write + + (output as f64) * rates.output +} + +#[derive(Debug, Clone, Copy)] +struct SelectedCostRates { + input: f64, + cache_read: f64, + cache_write: f64, + output: f64, +} + +fn selected_cost_rates( + pricing: models_dev_pricing::DynamicModelPricing, + use_tier: bool, +) -> SelectedCostRates { let pick = |base: f64, above: Option| { if use_tier { above.unwrap_or(base) @@ -161,39 +200,34 @@ pub fn cost_usd_from_pricing_with_threshold( base } }; - let input_rate = pick( + let input = pick( pricing.input_cost_per_token, pricing.input_cost_per_token_above_threshold, ); - let cache_read_rate = if use_tier { - pricing - .cache_read_input_cost_per_token_above_threshold - .or(pricing.cache_read_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_read_input_cost_per_token - .unwrap_or(input_rate) - }; - let cache_write_rate = if use_tier { - pricing - .cache_write_input_cost_per_token_above_threshold - .or(pricing.cache_write_input_cost_per_token) - .unwrap_or(input_rate) - } else { - pricing - .cache_write_input_cost_per_token - .unwrap_or(input_rate) - }; - let output_rate = pick( - pricing.output_cost_per_token, - pricing.output_cost_per_token_above_threshold, - ); - (input as f64) * input_rate - + (cache_read as f64) * cache_read_rate - + (cache_write as f64) * cache_write_rate - + (output as f64) * output_rate + SelectedCostRates { + input, + cache_read: if use_tier { + pricing + .cache_read_input_cost_per_token_above_threshold + .or(pricing.cache_read_input_cost_per_token) + .unwrap_or(input) + } else { + pricing.cache_read_input_cost_per_token.unwrap_or(input) + }, + cache_write: if use_tier { + pricing + .cache_write_input_cost_per_token_above_threshold + .or(pricing.cache_write_input_cost_per_token) + .unwrap_or(input) + } else { + pricing.cache_write_input_cost_per_token.unwrap_or(input) + }, + output: pick( + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_threshold, + ), + } } fn effective_threshold(provider: &str, model: &str, catalog_threshold: Option) -> Option { diff --git a/rust/src/core/cost_pricing/claude.rs b/rust/src/core/cost_pricing/claude.rs index 8cdab4f34d..32e4aa4e36 100644 --- a/rust/src/core/cost_pricing/claude.rs +++ b/rust/src/core/cost_pricing/claude.rs @@ -86,20 +86,38 @@ impl CostUsagePricing { cache_read_input_tokens: i32, cache_creation_input_tokens: i32, output_tokens: i32, + ) -> f64 { + Self::claude_cost_usd_u64_from_resolution( + resolution, + u64::try_from(input_tokens).unwrap_or(0), + u64::try_from(cache_read_input_tokens).unwrap_or(0), + u64::try_from(cache_creation_input_tokens).unwrap_or(0), + u64::try_from(output_tokens).unwrap_or(0), + ) + } + + /// Calculate cost from a resolved Claude pricing source without narrowing + /// untrusted local-history counters to the API-oriented signed type. + pub(crate) fn claude_cost_usd_u64_from_resolution( + resolution: ClaudePricingResolution, + input_tokens: u64, + cache_read_input_tokens: u64, + cache_creation_input_tokens: u64, + output_tokens: u64, ) -> f64 { match resolution { ClaudePricingResolution::BuiltIn(pricing) => { fn tiered( - tokens: i32, + tokens: u64, base: f64, above: Option, threshold: Option, ) -> f64 { - let tokens = tokens.max(0); match (threshold, above) { (Some(thresh), Some(above_rate)) => { + let thresh = u64::try_from(thresh).unwrap_or(0); let below = tokens.min(thresh); - let over = (tokens - thresh).max(0); + let over = tokens.saturating_sub(thresh); (below as f64) * base + (over as f64) * above_rate } _ => (tokens as f64) * base, @@ -131,7 +149,7 @@ impl CostUsagePricing { ClaudePricingResolution::ModelsDev { pricing, threshold_tokens, - } => claude_routed_pricing::cost_usd_from_pricing_with_threshold( + } => claude_routed_pricing::cost_usd_from_u64_counts_with_threshold( pricing, threshold_tokens, input_tokens, @@ -196,3 +214,85 @@ impl CostUsagePricing { claude_routed_pricing::input_cost_per_token(model, Self::normalize_claude_model(model)) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn routed_pricing() -> (models_dev_pricing::DynamicModelPricing, Option) { + let snapshot = models_dev_pricing::ModelsDevPricingSnapshot::from_catalog_json_for_tests( + r#"{ + "anthropic": {"models": {"threshold-fixture": {"id": "threshold-fixture", "cost": { + "input": 2, "output": 4, "cache_read": 0.25, "cache_write": 3, + "context_over_200k": {"input": 7, "output": 11, "cache_read": 0.5, "cache_write": 9} + }}}} + }"#, + ) + .expect("pricing fixture"); + let resolution = CostUsagePricing::resolve_claude_pricing( + "anthropic/threshold-fixture", + "anthropic/threshold-fixture", + Some(&snapshot), + ) + .expect("Models.dev pricing"); + let ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens, + } = resolution + else { + panic!("expected Models.dev pricing"); + }; + (pricing, threshold_tokens) + } + + #[test] + fn models_dev_u64_cost_uses_routed_rates_for_every_token_field() { + let (pricing, threshold) = routed_pricing(); + assert_eq!(threshold, Some(200_000)); + + for (input, cache_read, cache_write, output, expected_usd) in [ + (10_000, 2_000, 1_000, 500, 0.0255), + (199_999, 1, 0, 25, 0.400_098_25), + (200_000, 1, 0, 25, 1.400_275_5), + (220_000, 10_000, 2_000, 50, 1.563_55), + ] { + let actual = CostUsagePricing::claude_cost_usd_u64_from_resolution( + ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens: threshold, + }, + input, + cache_read, + cache_write, + output, + ); + let routed = claude_routed_pricing::cost_usd_from_u64_counts_with_threshold( + pricing, + threshold, + input, + cache_read, + cache_write, + output, + ); + assert!((actual - expected_usd).abs() < 1e-12); + assert!((actual - routed).abs() < 1e-12); + } + } + + #[test] + fn models_dev_u64_counter_overflow_selects_above_threshold_rates() { + let (pricing, threshold) = routed_pricing(); + let actual = CostUsagePricing::claude_cost_usd_u64_from_resolution( + ClaudePricingResolution::ModelsDev { + pricing, + threshold_tokens: threshold, + }, + u64::MAX, + 1, + 0, + 0, + ); + let expected = (u64::MAX as f64) * 7e-6 + 0.5e-6; + assert!((actual - expected).abs() < 1e-6); + } +} diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 883a31d89f..13e78bcfe9 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -373,6 +373,8 @@ pub(crate) struct CodexSessionMetadata { pub lineage: CodexSessionLineage, pub fork_timestamp: Option, pub history_base_thread_id: Option, + pub is_subagent: bool, + pub subagent_history_start_ordinal: Option, } /// Running totals for Codex token counting @@ -396,6 +398,10 @@ pub struct CodexForkAccountingState { pub inherited_totals: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub remaining_inherited_totals: Option, + /// True when the child log itself supplied enough copied-prefix history to + /// establish the inherited baseline without consulting a parent cache row. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub locally_resolved: bool, } /// Snapshot of the last validated cost report, persisted so spend surfaces keep @@ -456,6 +462,7 @@ pub struct CodexParseResult { pub fork_baseline: Option, /// Remaining inherited counters used when a fork emits last-only rows. pub remaining_inherited_totals: Option, + pub fork_baseline_locally_resolved: bool, } /// A billable Codex token-count delta. diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 30f159c0ca..108d4284d7 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -8,13 +8,14 @@ use helpers::{ BoundedJsonlLine, CODEX_JSONL_MAX_LINE_BYTES, nonempty_json_string, parse_rfc3339_timestamp, read_bounded_jsonl_line, read_bounded_jsonl_line_until, session_meta_field, }; -use parser::CodexParserState; +use parser::{CodexParseMode, CodexParserState}; /// Persisted Codex cache schema version. Version 0 predates 64-bit totals; /// version 1 can retain a terminal pause after treating a paginated v2 /// subagent's independent counters as an inherited fork. Version 3 adds -/// persisted paginated-fork accounting state. Rebuild older artifacts. -pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 3; +/// persisted paginated-fork accounting state. Version 4 reparses copied-prefix +/// subagents with locally inferred component baselines. Rebuild older artifacts. +pub(crate) const CODEX_CACHE_SCHEMA_VERSION: u32 = 4; /// Whether a persisted Codex cache artifact matches the current schema. /// A mismatched artifact (e.g. a pre-64-bit cache from an older release) is @@ -186,6 +187,12 @@ impl JsonlScanner { .pointer("/source/subagent/thread_spawn") .is_some_and(Value::is_object)) }); + let is_subagent = payload.is_some_and(|value| { + value.get("thread_source").and_then(Value::as_str) == Some("subagent") + || value + .pointer("/source/subagent/thread_spawn") + .is_some_and(Value::is_object) + }); let history_base_thread_id = payload .and_then(|value| value.get("history_base")) .filter(|value| value.is_object()) @@ -220,6 +227,10 @@ impl JsonlScanner { payload.and_then(|value| nonempty_json_string(value.get("timestamp"))) }), history_base_thread_id, + is_subagent, + subagent_history_start_ordinal: payload + .and_then(|value| value.get("subagent_history_start_ordinal")) + .and_then(Value::as_i64), }); } @@ -353,17 +364,16 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - start_offset, - initial_model, - initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, cancel, - false, - false, - None, None, max_bytes_to_read, + CodexParseMode::Standard { + start_offset, + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + }, ) } @@ -389,17 +399,16 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - start_offset, - initial_model, - initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, cancel, - false, - false, - None, scan_target_size, max_bytes_to_read, + CodexParseMode::Standard { + start_offset, + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + }, ) } @@ -420,17 +429,14 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - 0, - None, - Some(initial_totals), - None, - None, cancel, - true, - false, - None, None, max_bytes_to_read, + CodexParseMode::ParentBaseline { + baseline: initial_totals, + paginated_continuation: false, + remaining_inherited_totals: None, + }, ) } @@ -459,6 +465,26 @@ impl JsonlScanner { ) } + pub(crate) fn parse_codex_file_with_inferred_fork_baseline( + file_path: &Path, + range: &CostUsageDayRange, + subagent_history_start_ordinal: Option, + cancel: Option<&AtomicBool>, + scan_target_size: Option, + max_bytes_to_read: Option, + ) -> std::io::Result { + Self::parse_codex_file_with_state_bounded_internal( + file_path, + range, + cancel, + scan_target_size, + max_bytes_to_read, + CodexParseMode::InferSubagent { + start_ordinal: subagent_history_start_ordinal, + }, + ) + } + /// Fork equivalent with persisted paginated-continuation accounting. #[allow( clippy::too_many_arguments, @@ -477,38 +503,24 @@ impl JsonlScanner { Self::parse_codex_file_with_state_bounded_internal( file_path, range, - 0, - None, - Some(initial_totals), - None, - None, cancel, - true, - paginated_continuation, - remaining_inherited_totals, scan_target_size, max_bytes_to_read, + CodexParseMode::ParentBaseline { + baseline: initial_totals, + paginated_continuation, + remaining_inherited_totals, + }, ) } - #[allow( - clippy::too_many_arguments, - reason = "resume state mirrors the persisted parser cache" - )] fn parse_codex_file_with_state_bounded_internal( file_path: &Path, range: &CostUsageDayRange, - start_offset: i64, - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, cancel: Option<&AtomicBool>, - fork_baseline_mode: bool, - paginated_continuation: bool, - remaining_inherited_totals: Option, scan_target_size: Option, max_bytes_to_read: Option, + mode: CodexParseMode, ) -> std::io::Result { let file = File::open(file_path)?; // Session JSONL files are bounded by the cache budget; sizes fit i64. @@ -518,7 +530,7 @@ impl JsonlScanner { )] let file_size = file.metadata()?.len() as i64; - let safe_start_offset = start_offset.clamp(0, file_size); + let safe_start_offset = mode.start_offset().clamp(0, file_size); let requested_target_size = scan_target_size .unwrap_or(file_size) .max(safe_start_offset) @@ -529,15 +541,7 @@ impl JsonlScanner { reader.seek(SeekFrom::Start(safe_start_offset as u64))?; } - let mut parser = CodexParserState::with_timestamp_state_and_fork_options( - initial_model, - initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, - fork_baseline_mode, - paginated_continuation, - remaining_inherited_totals, - ); + let mut parser = CodexParserState::from_mode(mode); let mut parsed_bytes = safe_start_offset; let mut committed_bytes = safe_start_offset; let mut cancelled = false; @@ -626,6 +630,7 @@ impl JsonlScanner { }; let is_complete = !cancelled && !budget_exhausted && parsed_bytes >= effective_target_size; let bytes_read = parsed_bytes.saturating_sub(safe_start_offset).max(0); + let fork_baseline_locally_resolved = parser.fork_baseline_locally_resolved(); Ok(CodexParseResult { records: parser.records, parsed_bytes, @@ -644,6 +649,7 @@ impl JsonlScanner { fork_baseline_ambiguous: parser.fork_baseline_ambiguous, fork_baseline: parser.fork_baseline, remaining_inherited_totals: parser.remaining_inherited_totals, + fork_baseline_locally_resolved, }) } diff --git a/rust/src/core/jsonl_scanner/codex/parser.rs b/rust/src/core/jsonl_scanner/codex/parser.rs index 2ec70f9cf1..bdce9c68d8 100644 --- a/rust/src/core/jsonl_scanner/codex/parser.rs +++ b/rust/src/core/jsonl_scanner/codex/parser.rs @@ -22,62 +22,255 @@ pub(super) struct CodexParserState { paginated_continuation: bool, paginated_baseline_checked: bool, pub(super) fork_baseline_ambiguous: bool, + fork_baseline_inference: Option, } -impl CodexParserState { - pub(super) fn new(initial_model: Option, initial_totals: Option) -> Self { - Self::with_timestamp_state(initial_model, initial_totals, None, None) - } - - fn with_timestamp_state( +pub(super) enum CodexParseMode { + Standard { + start_offset: i64, initial_model: Option, initial_totals: Option, previous_token_timestamp: Option, token_timestamps_monotonic: Option, - ) -> Self { - Self::with_timestamp_state_and_fork_mode( + }, + ParentBaseline { + baseline: CodexTotals, + paginated_continuation: bool, + remaining_inherited_totals: Option, + }, + InferSubagent { + start_ordinal: Option, + }, +} + +impl CodexParseMode { + pub(super) fn start_offset(&self) -> i64 { + match self { + Self::Standard { start_offset, .. } => *start_offset, + Self::ParentBaseline { .. } | Self::InferSubagent { .. } => 0, + } + } +} + +#[derive(Debug)] +struct ForkBaselineInference { + explicit_start_ordinal: Option, + baseline: Option, + boundary_open: bool, + inherited_opening: bool, + missing_explicit_ordinal: bool, + locally_confirmed: bool, + resolved: bool, +} + +enum ForkBaselineDecision { + SkipCopiedPrefix, + ProcessWithBaseline(CodexTotals), +} + +impl ForkBaselineInference { + fn new(explicit_start_ordinal: Option) -> Self { + Self { + explicit_start_ordinal, + baseline: explicit_start_ordinal.map(|_| CodexTotals { + input: 0, + cached: 0, + output: 0, + reasoning: None, + }), + boundary_open: false, + inherited_opening: false, + missing_explicit_ordinal: false, + locally_confirmed: false, + resolved: false, + } + } + + fn confirm_local_resolution(&mut self) { + if !self.missing_explicit_ordinal { + self.locally_confirmed = true; + } + } + + fn mark_missing_explicit_ordinal(&mut self) { + self.missing_explicit_ordinal = true; + self.locally_confirmed = false; + } + + fn observe_non_token(&mut self, obj: &Value) { + if obj.get("type").and_then(Value::as_str) == Some("turn_context") && self.inherited_opening + { + self.boundary_open = true; + } + } + + fn observe_token(&mut self, obj: &Value) -> ForkBaselineDecision { + let Some(payload) = token_count_payload(obj) else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let Some(info) = payload.get("info") else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let Some(total_usage) = info.get("total_token_usage") else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let Some(last_usage) = info.get("last_token_usage") else { + return ForkBaselineDecision::SkipCopiedPrefix; + }; + let total = read_token_totals(total_usage); + let last = read_token_totals(last_usage); + let ordinal = obj.get("ordinal").and_then(Value::as_i64); + + if self.explicit_start_ordinal.is_some() && ordinal.is_none() { + self.mark_missing_explicit_ordinal(); + if !self.boundary_open { + self.baseline = Some(total); + } + return ForkBaselineDecision::SkipCopiedPrefix; + } + + if let Some(start) = self.explicit_start_ordinal + && !self.boundary_open + { + let ordinal = ordinal.expect("missing explicit ordinals return above"); + if ordinal < start { + self.baseline = Some(total); + return ForkBaselineDecision::SkipCopiedPrefix; + } + self.boundary_open = true; + } else if self.baseline.is_none() { + if totals_contain_usage(&total) && !totals_contain_usage(&last) { + self.baseline = Some(total); + self.inherited_opening = true; + self.confirm_local_resolution(); + } + return ForkBaselineDecision::SkipCopiedPrefix; + } else if !self.boundary_open { + let changed = self + .baseline + .as_ref() + .is_some_and(|baseline| baseline != &total); + if self.inherited_opening && changed && totals_contain_usage(&last) { + self.boundary_open = true; + } else { + return ForkBaselineDecision::SkipCopiedPrefix; + } + } + + let baseline = self.baseline.clone().unwrap_or(CodexTotals { + input: 0, + cached: 0, + output: 0, + reasoning: None, + }); + if total == baseline { + return ForkBaselineDecision::SkipCopiedPrefix; + } + let copied_snapshot = + totals_contain_usage(&baseline) && total == last && totals_at_least(&total, &baseline); + if copied_snapshot { + self.baseline = Some(total); + self.confirm_local_resolution(); + return ForkBaselineDecision::SkipCopiedPrefix; + } + + let owned_baseline = totals_delta(&last, &total); + self.baseline = Some(owned_baseline.clone()); + self.confirm_local_resolution(); + self.resolved = true; + ForkBaselineDecision::ProcessWithBaseline(owned_baseline) + } +} + +fn totals_contain_usage(totals: &CodexTotals) -> bool { + totals.input > 0 || totals.cached > 0 || totals.output > 0 +} + +fn totals_at_least(total: &CodexTotals, baseline: &CodexTotals) -> bool { + total.input >= baseline.input + && total.cached >= baseline.cached + && total.output >= baseline.output +} + +fn totals_delta(last: &CodexTotals, total: &CodexTotals) -> CodexTotals { + CodexTotals { + input: total.input.saturating_sub(last.input).max(0), + cached: total.cached.saturating_sub(last.cached).max(0), + output: total.output.saturating_sub(last.output).max(0), + reasoning: subtract_optional(total.reasoning, last.reasoning), + } +} + +impl CodexParserState { + pub(super) fn new(initial_model: Option, initial_totals: Option) -> Self { + Self::from_mode(CodexParseMode::Standard { + start_offset: 0, initial_model, initial_totals, - previous_token_timestamp, - token_timestamps_monotonic, - false, - ) + previous_token_timestamp: None, + token_timestamps_monotonic: None, + }) } - pub(super) fn with_timestamp_state_and_fork_mode( - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, - fork_baseline_mode: bool, - ) -> Self { - Self::with_timestamp_state_and_fork_options( + pub(super) fn from_mode(mode: CodexParseMode) -> Self { + let ( initial_model, initial_totals, previous_token_timestamp, token_timestamps_monotonic, - fork_baseline_mode, - false, - None, - ) - } - - pub(super) fn with_timestamp_state_and_fork_options( - initial_model: Option, - initial_totals: Option, - previous_token_timestamp: Option, - token_timestamps_monotonic: Option, - fork_baseline_mode: bool, - paginated_continuation: bool, - remaining_inherited_totals: Option, - ) -> Self { + fork_baseline, + paginated_continuation, + remaining_inherited_totals, + fork_baseline_inference, + ) = match mode { + CodexParseMode::Standard { + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + .. + } => ( + initial_model, + initial_totals, + previous_token_timestamp, + token_timestamps_monotonic, + None, + false, + None, + None, + ), + CodexParseMode::ParentBaseline { + baseline, + paginated_continuation, + remaining_inherited_totals, + } => { + let remaining_inherited_totals = + remaining_inherited_totals.or_else(|| Some(baseline.clone())); + ( + None, + Some(baseline.clone()), + None, + None, + Some(baseline), + paginated_continuation, + remaining_inherited_totals, + None, + ) + } + CodexParseMode::InferSubagent { start_ordinal } => ( + None, + None, + None, + None, + None, + false, + None, + Some(ForkBaselineInference::new(start_ordinal)), + ), + }; let previous_token_timestamp_parsed = previous_token_timestamp .as_deref() .and_then(parse_rfc3339_timestamp); - let fork_baseline = fork_baseline_mode.then(|| initial_totals.clone()).flatten(); - let remaining_inherited_totals = fork_baseline - .as_ref() - .and_then(|baseline| remaining_inherited_totals.or_else(|| Some(baseline.clone()))); Self { current_model: initial_model, previous_totals: initial_totals.clone(), @@ -95,9 +288,16 @@ impl CodexParserState { paginated_continuation, paginated_baseline_checked: false, fork_baseline_ambiguous: false, + fork_baseline_inference, } } + pub(super) fn fork_baseline_locally_resolved(&self) -> bool { + self.fork_baseline_inference + .as_ref() + .is_some_and(|inference| inference.locally_confirmed) + } + pub(super) fn process_line(&mut self, line: &str, range: &CostUsageDayRange) { self.process_line_with_source_offset(line, range, 0); } @@ -108,7 +308,61 @@ impl CodexParserState { range: &CostUsageDayRange, source_end_offset: i64, ) { + if self + .fork_baseline_inference + .as_ref() + .is_some_and(|inference| !inference.resolved) + { + let Ok(obj) = serde_json::from_str::(line) else { + return; + }; + if token_count_payload(&obj).is_some() { + let decision = self + .fork_baseline_inference + .as_mut() + .expect("inference exists") + .observe_token(&obj); + let ForkBaselineDecision::ProcessWithBaseline(baseline) = decision else { + return; + }; + self.fork_baseline = Some(baseline.clone()); + self.remaining_inherited_totals = Some(baseline.clone()); + self.previous_totals = Some(baseline.clone()); + self.totals_watermark = Some(baseline); + } else { + self.fork_baseline_inference + .as_mut() + .expect("inference exists") + .observe_non_token(&obj); + if obj.get("type").and_then(Value::as_str) == Some("turn_context") { + self.update_current_model(&obj); + } + return; + } + } + let event_candidate = is_candidate_codex_line(line); + if event_candidate + && self + .fork_baseline_inference + .as_ref() + .is_some_and(|inference| { + inference.resolved && inference.explicit_start_ordinal.is_some() + }) + { + let Ok(obj) = serde_json::from_str::(line) else { + return; + }; + if token_count_payload(&obj).is_some() + && obj.get("ordinal").and_then(Value::as_i64).is_none() + { + self.fork_baseline_inference + .as_mut() + .expect("resolved inference exists") + .mark_missing_explicit_ordinal(); + return; + } + } let bare_candidate = !event_candidate && line.contains("\"usage\""); if !event_candidate && !bare_candidate { return; diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index bcb1129cb1..1aebd89e20 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -77,13 +77,11 @@ fn fork_baseline_subtracts_known_reasoning_without_affecting_core_tokens() { output: 10, reasoning: Some(4), }; - let mut state = CodexParserState::with_timestamp_state_and_fork_mode( - None, - Some(baseline), - None, - None, - true, - ); + let mut state = CodexParserState::from_mode(CodexParseMode::ParentBaseline { + baseline, + paginated_continuation: false, + remaining_inherited_totals: None, + }); assert_eq!( state.apply_totals_delta(CodexTotals { input: 20, @@ -100,13 +98,11 @@ fn fork_baseline_subtracts_known_reasoning_without_affecting_core_tokens() { output: 10, reasoning: None, }; - let mut state = CodexParserState::with_timestamp_state_and_fork_mode( - None, - Some(baseline_without_reasoning), - None, - None, - true, - ); + let mut state = CodexParserState::from_mode(CodexParseMode::ParentBaseline { + baseline: baseline_without_reasoning, + paginated_continuation: false, + remaining_inherited_totals: None, + }); assert_eq!( state.apply_totals_delta(CodexTotals { input: 20, @@ -119,6 +115,109 @@ fn fork_baseline_subtracts_known_reasoning_without_affecting_core_tokens() { assert!(!state.fork_baseline_ambiguous); } +#[test] +fn inferred_fork_waits_for_present_explicit_start_ordinal() { + let range = CostUsageDayRange::new( + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + ); + let mut state = CodexParserState::from_mode(CodexParseMode::InferSubagent { + start_ordinal: Some(10), + }); + + state.process_line( + r#"{"timestamp":"2026-09-22T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.6-sol","total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":10},"last_token_usage":{"input_tokens":0,"cached_input_tokens":0,"output_tokens":0}}}}"#, + &range, + ); + + assert!(state.records.is_empty()); + assert!(state.fork_baseline.is_none()); + assert!(!state.fork_baseline_locally_resolved()); + + state.process_line( + r#"{"ordinal":10,"timestamp":"2026-09-22T10:00:01Z","type":"event_msg","payload":{"type":"token_count","info":{"model":"gpt-5.6-sol","total_token_usage":{"input_tokens":110,"cached_input_tokens":22,"output_tokens":11},"last_token_usage":{"input_tokens":10,"cached_input_tokens":2,"output_tokens":1}}}}"#, + &range, + ); + + assert_eq!(state.records.len(), 1); + assert_eq!(state.records[0].0.input, 10); + assert_eq!(state.records[0].0.cached, 2); + assert_eq!(state.records[0].0.output, 1); + assert!(!state.fork_baseline_locally_resolved()); +} + +#[test] +fn inferred_fork_keeps_missing_ordinal_unresolved_after_boundary_opens() { + let range = CostUsageDayRange::new( + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + ); + let mut state = CodexParserState::from_mode(CodexParseMode::InferSubagent { + start_ordinal: Some(10), + }); + let token_line = |ordinal: Option, total: i64, last: i64| { + let mut value = serde_json::json!({ + "timestamp": "2026-09-22T10:00:00Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": { + "model": "gpt-5.6-sol", + "total_token_usage": {"input_tokens": total, "cached_input_tokens": 0, "output_tokens": 0}, + "last_token_usage": {"input_tokens": last, "cached_input_tokens": 0, "output_tokens": 0} + }} + }); + if let Some(ordinal) = ordinal { + value["ordinal"] = serde_json::json!(ordinal); + } + value.to_string() + }; + + state.process_line(&token_line(Some(9), 100, 0), &range); + state.process_line(&token_line(Some(10), 100, 0), &range); + state.process_line(&token_line(Some(11), 110, 110), &range); + assert!(state.fork_baseline_locally_resolved()); + state.process_line(&token_line(None, 120, 10), &range); + assert!(!state.fork_baseline_locally_resolved()); + state.process_line(&token_line(Some(12), 130, 10), &range); + + assert_eq!(state.records.len(), 1); + assert_eq!(state.records[0].0.input, 10); + assert!(!state.fork_baseline_locally_resolved()); +} + +#[test] +fn inferred_fork_keeps_missing_ordinal_unresolved_after_local_resolution() { + let range = CostUsageDayRange::new( + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + NaiveDate::from_ymd_opt(2026, 9, 22).unwrap(), + ); + let mut state = CodexParserState::from_mode(CodexParseMode::InferSubagent { + start_ordinal: Some(10), + }); + let token_line = |ordinal: Option, total: i64, last: i64| { + let mut value = serde_json::json!({ + "timestamp": "2026-09-22T10:00:00Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": { + "model": "gpt-5.6-sol", + "total_token_usage": {"input_tokens": total, "cached_input_tokens": 0, "output_tokens": 0}, + "last_token_usage": {"input_tokens": last, "cached_input_tokens": 0, "output_tokens": 0} + }} + }); + if let Some(ordinal) = ordinal { + value["ordinal"] = serde_json::json!(ordinal); + } + value.to_string() + }; + + state.process_line(&token_line(Some(9), 100, 0), &range); + state.process_line(&token_line(Some(10), 100, 0), &range); + state.process_line(&token_line(Some(11), 110, 10), &range); + assert!(state.fork_baseline_locally_resolved()); + state.process_line(&token_line(None, 120, 10), &range); + + assert!(!state.fork_baseline_locally_resolved()); +} + #[test] fn codex_token_pipeline_preserves_counts_above_i32_max() { let parsed = read_token_totals(&serde_json::json!({ @@ -1043,6 +1142,8 @@ fn session_meta_pre_read_accepts_snake_and_camel_fork_identity() { lineage: CodexSessionLineage::Child, fork_timestamp: Some("2026-05-31T10:00:00Z".to_string()), history_base_thread_id: Some("history-snake".to_string()), + is_subagent: false, + subagent_history_start_ordinal: None, } ); diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..218901b770 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -55,7 +55,6 @@ pub enum ProviderId { MiMo, Doubao, CommandCode, - Crof, StepFun, Venice, OpenAIApi, @@ -142,7 +141,6 @@ impl ProviderId { ProviderId::MiMo, ProviderId::Doubao, ProviderId::CommandCode, - ProviderId::Crof, ProviderId::StepFun, ProviderId::Venice, ProviderId::OpenAIApi, @@ -232,7 +230,6 @@ impl ProviderId { ProviderId::MiMo => "mimo", ProviderId::Doubao => "doubao", ProviderId::CommandCode => "commandcode", - ProviderId::Crof => "crof", ProviderId::StepFun => "stepfun", ProviderId::Venice => "venice", ProviderId::OpenAIApi => "openaiapi", @@ -319,7 +316,6 @@ impl ProviderId { ProviderId::MiMo => "Xiaomi MiMo", ProviderId::Doubao => "Doubao", ProviderId::CommandCode => "Command Code", - ProviderId::Crof => "Crof", ProviderId::StepFun => "StepFun", ProviderId::Venice => "Venice", ProviderId::OpenAIApi => "OpenAI API", @@ -418,7 +414,6 @@ impl ProviderId { ProviderId::AiAnd => None, ProviderId::Windsurf => None, ProviderId::Doubao => None, - ProviderId::Crof => None, ProviderId::StepFun => None, ProviderId::OpenAIApi => None, ProviderId::ElevenLabs => None, @@ -505,7 +500,6 @@ impl ProviderId { } "doubao" | "ark" | "volcengine" => Some(ProviderId::Doubao), "commandcode" | "command-code" | "command code" => Some(ProviderId::CommandCode), - "crof" => Some(ProviderId::Crof), "stepfun" | "step-fun" | "step fun" => Some(ProviderId::StepFun), "venice" => Some(ProviderId::Venice), "openaiapi" | "openai-api" | "openai api" | "openai-balance" => { @@ -1069,7 +1063,6 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::MiMo => "#FF6900", ProviderId::Doubao => "#2563EB", ProviderId::CommandCode => "#44FF00", - ProviderId::Crof => "#7C3AED", ProviderId::StepFun => "#999999", ProviderId::Venice => "#111827", ProviderId::OpenAIApi => "#10A37F", @@ -1120,7 +1113,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 80); + assert_eq!(all.len(), 79); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Pi)); @@ -1145,7 +1138,6 @@ mod tests { assert!(all.contains(&ProviderId::MiMo)); assert!(all.contains(&ProviderId::Doubao)); assert!(all.contains(&ProviderId::CommandCode)); - assert!(all.contains(&ProviderId::Crof)); assert!(all.contains(&ProviderId::StepFun)); assert!(all.contains(&ProviderId::Venice)); assert!(all.contains(&ProviderId::OpenAIApi)); @@ -1258,6 +1250,7 @@ mod tests { Some(ProviderId::Antigravity) ); assert_eq!(ProviderId::from_cli_name("zed"), Some(ProviderId::Zed)); + assert_eq!(ProviderId::from_cli_name("crof"), None); assert_eq!(ProviderId::from_cli_name("unknown"), None); assert_eq!( ProviderId::from_cli_name("code-rabbit"), diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 712c1d4fb9..d801961af9 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -10,17 +10,17 @@ use crate::providers::{ AbacusProvider, AiAndProvider, AlibabaProvider, AlibabaTokenPlanProvider, AmpProvider, AntigravityProvider, AugmentProvider, AzureOpenAIProvider, BedrockProvider, ChutesProvider, ClaudeProvider, ClinePassProvider, CodeBuddyProvider, CodeRabbitProvider, CodebuffProvider, - CodexProvider, CommandCodeProvider, CopilotProvider, CrofProvider, CrossModelProvider, - CursorProvider, DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevinProvider, - DoubaoProvider, ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, - GrokProvider, GroqProvider, HelmcodeProvider, HuggingFaceProvider, InfiniProvider, - JetBrainsProvider, KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, - LiteLLMProvider, LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, - MistralProvider, MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, - NousProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, - OpenRouterProvider, PerplexityProvider, PiProvider, PoeProvider, QoderProvider, - QwenCloudProvider, ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, - T3ChatProvider, TypeSafeProvider, V0Provider, VeniceProvider, VertexAIProvider, WarpProvider, + CodexProvider, CommandCodeProvider, CopilotProvider, CrossModelProvider, CursorProvider, + DeepInfraProvider, DeepSeekProvider, DeepgramProvider, DevinProvider, DoubaoProvider, + ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, GrokProvider, + GroqProvider, HelmcodeProvider, HuggingFaceProvider, InfiniProvider, JetBrainsProvider, + KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, LiteLLMProvider, + LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, + MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, NousProvider, + OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, + PerplexityProvider, PiProvider, PoeProvider, QoderProvider, QwenCloudProvider, + ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, + TypeSafeProvider, V0Provider, VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, }; @@ -74,7 +74,6 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::MiMo => Box::new(MiMoProvider::new()), ProviderId::Doubao => Box::new(DoubaoProvider::new()), ProviderId::CommandCode => Box::new(CommandCodeProvider::new()), - ProviderId::Crof => Box::new(CrofProvider::new()), ProviderId::StepFun => Box::new(StepFunProvider::new()), ProviderId::Venice => Box::new(VeniceProvider::new()), ProviderId::OpenAIApi => Box::new(OpenAIApiProvider::new()), diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index 6dedc0fbde..dbbb929d83 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -354,7 +354,6 @@ impl TokenAccountSupport { | ProviderId::DeepSeek | ProviderId::Windsurf | ProviderId::Doubao - | ProviderId::Crof | ProviderId::StepFun | ProviderId::Venice | ProviderId::OpenAIApi diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 619f9ba7b0..55444c6a74 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -435,7 +435,9 @@ struct ClaudeUsageRecord { output: u64, cache_create: u64, cache_read: u64, - cost: f64, + /// `None` means pricing was unavailable or produced a non-finite value. + /// Keep that distinct from a real zero-dollar row. + cost: Option, } /// Exact Claude request rows retained for quota-window projection. @@ -468,6 +470,7 @@ struct ClaudeFileScanResult { malformed_lines: u32, incomplete_requests: u32, read_failures: u32, + aggregation_failures: u32, } impl ClaudeFileScanResult { @@ -478,10 +481,16 @@ impl ClaudeFileScanResult { .incomplete_requests .saturating_add(other.incomplete_requests); self.read_failures = self.read_failures.saturating_add(other.read_failures); + self.aggregation_failures = self + .aggregation_failures + .saturating_add(other.aggregation_failures); } fn is_complete(self) -> bool { - self.malformed_lines == 0 && self.incomplete_requests == 0 && self.read_failures == 0 + self.malformed_lines == 0 + && self.incomplete_requests == 0 + && self.read_failures == 0 + && self.aggregation_failures == 0 } } @@ -590,23 +599,35 @@ impl CostScanner { if projects_dir.exists() { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); - let mut handle_file = |path: &Path| { - let file_result = scan_claude_file_with_pricing( - path, - &cutoff, - &mut seen, - cancel, - &mut pricing, - |record| { - add_claude_record_to_summary(&mut summary, record); - }, - ); - if file_result.counted > 0 { - summary.sessions_count += 1; - } - claude_scan.absorb(file_result); + let traversal_read_failures = { + let mut handle_file = |path: &Path| { + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( + path, + &cutoff, + &mut seen, + cancel, + &mut pricing, + |record| { + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= + add_claude_record_to_summary(&mut summary, record); + }, + ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } + if file_result.counted > 0 { + summary.sessions_count += 1; + } + claude_scan.absorb(file_result); + }; + self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut handle_file) }; - self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut handle_file); + claude_scan.read_failures = claude_scan + .read_failures + .saturating_add(traversal_read_failures); } // OMP / pi-compatible anthropic rows, deduped across shared files. @@ -652,6 +673,7 @@ impl CostScanner { }; let mut daily_cost = HashMap::new(); let mut daily_tokens = HashMap::new(); + let mut unknown_cost_dates = HashSet::new(); for days_ago in 0..self.days { let date = today - Duration::days(days_ago as i64); let key = date.format("%Y-%m-%d").to_string(); @@ -661,35 +683,48 @@ impl CostScanner { let mut quota_records = Vec::new(); let mut scan_result = ClaudeFileScanResult::default(); - let mut missing_timestamp = false; if projects_dir.exists() { let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); - self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut |path| { - let mut file_has_usage = false; - let file_result = scan_claude_file_with_pricing( - path, - &cutoff, - &mut seen, - cancel, - &mut pricing, - |record| { - file_has_usage = true; - add_claude_record_to_summary(&mut summary, record); - add_claude_record_to_daily_costs(&mut daily_cost, record); - add_claude_record_to_daily_tokens(&mut daily_tokens, record); - if let Some(quota_record) = quota_history_record_from_usage(record) { - quota_records.push(quota_record); - } else { - missing_timestamp = true; - } - }, - ); - if file_has_usage { - summary.sessions_count += 1; - } - scan_result.absorb(file_result); - }); + let traversal_read_failures = + self.walk_claude_files(&projects_dir, &cutoff, cancel, &mut |path| { + let mut file_has_usage = false; + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( + path, + &cutoff, + &mut seen, + cancel, + &mut pricing, + |record| { + file_has_usage = true; + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= + add_claude_record_to_summary(&mut summary, record); + aggregation_complete &= add_claude_record_to_daily_costs( + &mut daily_cost, + &mut unknown_cost_dates, + record, + ); + aggregation_complete &= + add_claude_record_to_daily_tokens(&mut daily_tokens, record); + if let Some(quota_record) = quota_history_record_from_usage(record) { + quota_records.push(quota_record); + } + }, + ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } + if file_has_usage { + summary.sessions_count += 1; + } + scan_result.absorb(file_result); + }); + scan_result.read_failures = scan_result + .read_failures + .saturating_add(traversal_read_failures); } crate::pi_session_cost::scan_pi_compatible_into( @@ -700,10 +735,7 @@ impl CostScanner { &mut HashSet::new(), ); - let complete = projects_dir.exists() - && !is_cancelled(cancel) - && scan_result.is_complete() - && !missing_timestamp; + let complete = projects_dir.exists() && !is_cancelled(cancel) && scan_result.is_complete(); finalize_claude_summary( &mut summary, projects_dir.exists(), @@ -711,11 +743,7 @@ impl CostScanner { is_cancelled(cancel), ); if complete { - for value in daily_cost.values_mut() { - if value.is_none() { - *value = Some(0.0); - } - } + zero_fill_uninitialized_claude_daily_costs(&mut daily_cost, &unknown_cost_dates); } let mut daily_cost = daily_cost.into_iter().collect::>(); @@ -801,36 +829,52 @@ impl CostScanner { cutoff: &DateTime, cancel: Option<&AtomicBool>, on_file: &mut F, - ) where + ) -> u32 + where F: FnMut(&Path), { if is_cancelled(cancel) { - return; + return 0; } let entries = match fs::read_dir(dir) { Ok(e) => e, - Err(_) => return, + Err(_) => return 1, }; - for entry in entries.flatten() { + let mut read_failures = 0u32; + for entry in entries { if is_cancelled(cancel) { break; } + let entry = match entry { + Ok(entry) => entry, + Err(_) => { + read_failures = read_failures.saturating_add(1); + continue; + } + }; let path = entry.path(); - if path.is_dir() { - self.walk_claude_files(&path, cutoff, cancel, on_file); - } else if path.extension().is_some_and(|e| e == "jsonl") { - // Check file modification time - if let Ok(metadata) = fs::metadata(&path) - && let Ok(modified) = metadata.modified() - { - let modified_dt: DateTime = modified.into(); - if modified_dt >= *cutoff { - on_file(&path); + match fs::metadata(&path) { + Ok(metadata) if metadata.is_dir() => { + read_failures = read_failures + .saturating_add(self.walk_claude_files(&path, cutoff, cancel, on_file)); + } + Ok(metadata) if path.extension().is_some_and(|e| e == "jsonl") => { + match metadata.modified() { + Ok(modified) => { + let modified_dt: DateTime = modified.into(); + if modified_dt >= *cutoff { + on_file(&path); + } + } + Err(_) => read_failures = read_failures.saturating_add(1), } } + Ok(_) => {} + Err(_) => read_failures = read_failures.saturating_add(1), } } + read_failures } } @@ -854,20 +898,6 @@ where scan_claude_file_with_pricing(path, cutoff, seen, cancel, &mut pricing, on_record).counted } -fn for_each_claude_usage_record_with_pricing( - path: &Path, - cutoff: &DateTime, - seen: &mut HashSet, - cancel: Option<&AtomicBool>, - pricing: &mut ClaudeScanPricingResolver, - on_record: F, -) -> usize -where - F: FnMut(&ClaudeUsageRecord), -{ - scan_claude_file_with_pricing(path, cutoff, seen, cancel, pricing, on_record).counted -} - fn scan_claude_file_with_pricing( path: &Path, cutoff: &DateTime, @@ -1007,7 +1037,7 @@ fn claude_usage_record_from_event_with_pricing( let cache_create_1h = usage.one_hour_cache_creation_tokens(cache_create); let pricing_known = pricing.is_known(model); - let cost = pricing.cost_usd_with_cache_ttl( + let computed_cost = pricing.cost_usd_with_cache_ttl( model, input, cache_create, @@ -1015,6 +1045,7 @@ fn claude_usage_record_from_event_with_pricing( cache_read, output, ); + let cost = computed_cost.is_finite().then_some(computed_cost); Some(ClaudeUsageRecord { model: model.to_string(), @@ -1033,25 +1064,52 @@ fn claude_usage_record_from_event_with_pricing( }) } -fn add_claude_record_to_summary(summary: &mut CostSummary, record: &ClaudeUsageRecord) { +fn add_claude_record_to_summary(summary: &mut CostSummary, record: &ClaudeUsageRecord) -> bool { if !record.pricing_known { summary.unknown_models.insert(record.model.clone()); } - summary.input_tokens += record.input; - summary.output_tokens += record.output; - summary.cached_tokens += record.cache_create + record.cache_read; - summary.total_cost_usd += record.cost; + let mut complete = checked_add_assign(&mut summary.input_tokens, record.input); + complete &= checked_add_assign(&mut summary.output_tokens, record.output); + let cached = record.cache_create.checked_add(record.cache_read); + complete &= cached.is_some_and(|value| checked_add_assign(&mut summary.cached_tokens, value)); - *summary.by_model.entry(record.model.clone()).or_insert(0.0) += record.cost; + if let Some(cost) = record.cost { + complete &= checked_add_finite(&mut summary.total_cost_usd, cost); + complete &= checked_add_finite( + summary.by_model.entry(record.model.clone()).or_insert(0.0), + cost, + ); + } else { + complete = false; + } let model_tokens = summary .by_model_tokens .entry(record.model.clone()) .or_default(); - model_tokens.input_tokens += record.input; - model_tokens.output_tokens += record.output; - model_tokens.cached_tokens += record.cache_create + record.cache_read; + complete &= checked_add_assign(&mut model_tokens.input_tokens, record.input); + complete &= checked_add_assign(&mut model_tokens.output_tokens, record.output); + complete &= + cached.is_some_and(|value| checked_add_assign(&mut model_tokens.cached_tokens, value)); + complete +} + +fn checked_add_assign(total: &mut u64, value: u64) -> bool { + let Some(sum) = total.checked_add(value) else { + return false; + }; + *total = sum; + true +} + +fn checked_add_finite(total: &mut f64, value: f64) -> bool { + let sum = *total + value; + if !value.is_finite() || !sum.is_finite() { + return false; + } + *total = sum; + true } fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option { @@ -1080,9 +1138,9 @@ fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option= 0.0, + cost_is_complete: record.pricing_known && record.cost.is_some_and(|cost| cost >= 0.0), dedup_key, attribution: ClaudeHistoryAttribution::Unavailable, }) @@ -1093,10 +1151,11 @@ fn quota_history_record_from_usage(record: &ClaudeUsageRecord) -> Option>, + unknown_cost_dates: &mut HashSet, record: &ClaudeUsageRecord, -) { +) -> bool { let Some(timestamp) = record.timestamp else { - return; + return true; }; let date_str = timestamp .with_timezone(&Local) @@ -1104,7 +1163,33 @@ fn add_claude_record_to_daily_costs( .format("%Y-%m-%d") .to_string(); if let Some(cost) = daily_costs.get_mut(&date_str) { - *cost = Some(cost.unwrap_or(0.0) + record.cost); + if unknown_cost_dates.contains(&date_str) { + return false; + } + let Some(record_cost) = record.cost else { + *cost = None; + unknown_cost_dates.insert(date_str); + return false; + }; + let sum = cost.unwrap_or(0.0) + record_cost; + if !sum.is_finite() { + *cost = None; + unknown_cost_dates.insert(date_str); + return false; + } + *cost = Some(sum); + } + true +} + +fn zero_fill_uninitialized_claude_daily_costs( + daily_costs: &mut HashMap>, + unknown_cost_dates: &HashSet, +) { + for (day, cost) in daily_costs { + if cost.is_none() && !unknown_cost_dates.contains(day) { + *cost = Some(0.0); + } } } @@ -1188,26 +1273,41 @@ pub fn get_daily_cost_history(provider: &str, days: u32) -> Vec<(String, Option< let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); let mut claude_scan = ClaudeFileScanResult::default(); - let mut handle_file = |path: &Path| { - let file_result = scan_claude_file_with_pricing( - path, - &cutoff, - &mut seen, - None, - &mut pricing, - |record| { - add_claude_record_to_daily_costs(&mut daily_costs, record); - }, - ); - claude_scan.absorb(file_result); + let mut unknown_cost_dates = HashSet::new(); + let traversal_read_failures = { + let mut handle_file = |path: &Path| { + let mut aggregation_complete = true; + let mut file_result = scan_claude_file_with_pricing( + path, + &cutoff, + &mut seen, + None, + &mut pricing, + |record| { + aggregation_complete &= record.timestamp.is_some(); + aggregation_complete &= add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + record, + ); + }, + ); + if !aggregation_complete { + file_result.aggregation_failures = + file_result.aggregation_failures.saturating_add(1); + } + claude_scan.absorb(file_result); + }; + scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file) }; - scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); + claude_scan.read_failures = claude_scan + .read_failures + .saturating_add(traversal_read_failures); if claude_scan.is_complete() { - for slot in daily_costs.values_mut() { - if slot.is_none() { - *slot = Some(0.0); - } - } + zero_fill_uninitialized_claude_daily_costs( + &mut daily_costs, + &unknown_cost_dates, + ); } } } @@ -1287,26 +1387,31 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> } "claude" => { // Per-day token breakdown from the same de-duplicated record walk - // as the cost chart. The full walk is authoritative, so the - // Refreshing marker never applies here. + // as the cost chart. Only a complete valid walk establishes + // authoritative coverage of the requested history window. let projects_dir = scanner.get_claude_projects_dir(); if projects_dir.exists() { let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); let mut pricing = ClaudeScanPricingResolver::default(); - let mut handle_file = |path: &Path| { - for_each_claude_usage_record_with_pricing( - path, - &cutoff, - &mut seen, - None, - &mut pricing, - |record| { - add_claude_record_to_daily_tokens(&mut daily_tokens, record); - }, - ); + let mut claude_scan = ClaudeFileScanResult::default(); + let traversal_read_failures = { + let mut handle_file = |path: &Path| { + let file_result = scan_claude_file_for_daily_tokens( + path, + &cutoff, + &mut seen, + &mut pricing, + &mut daily_tokens, + ); + claude_scan.absorb(file_result); + }; + scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file) }; - scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); + claude_scan.read_failures = claude_scan + .read_failures + .saturating_add(traversal_read_failures); + mark_claude_daily_token_coverage(&mut covered_days, &daily_tokens, claude_scan); } } "pi" => { @@ -1327,14 +1432,12 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> let mut result: Vec<(String, u64)> = daily_tokens.into_iter().collect(); result.sort_by(|a, b| a.0.cmp(&b.0)); - // Codex only: the bounded catch-up may not have reached the requested - // depth yet. Incomplete = history exists but the oldest quarter of the - // window has no scanned day. - let incomplete = if provider == "pi" { - // Pi scans are bounded filesystem walks, so a complete parse covers - // the requested window even when the roots contain no sessions. + let incomplete = if matches!(provider, "claude" | "pi") { + // A complete filesystem scan covers the requested window even when + // the roots contain no sessions; failed scans leave coverage empty. covered_days.is_empty() } else { + // Codex catch-up may not have reached the oldest quarter of the window. provider == "codex" && !covered_days.is_empty() && covered_days.len() < days as usize @@ -1349,9 +1452,9 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> fn add_claude_record_to_daily_tokens( daily_tokens: &mut HashMap, record: &ClaudeUsageRecord, -) { +) -> bool { let Some(timestamp) = record.timestamp else { - return; + return true; }; let date_str = timestamp .with_timezone(&Local) @@ -1359,6 +1462,41 @@ fn add_claude_record_to_daily_tokens( .format("%Y-%m-%d") .to_string(); if let Some(slot) = daily_tokens.get_mut(&date_str) { - *slot += record.input + record.output; + let Some(tokens) = record.input.checked_add(record.output) else { + return false; + }; + return checked_add_assign(slot, tokens); + } + true +} + +fn scan_claude_file_for_daily_tokens( + path: &Path, + cutoff: &DateTime, + seen: &mut HashSet, + pricing: &mut ClaudeScanPricingResolver, + daily_tokens: &mut HashMap, +) -> ClaudeFileScanResult { + let mut aggregation_failures = 0u32; + let mut result = scan_claude_file_with_pricing(path, cutoff, seen, None, pricing, |record| { + if record.timestamp.is_none() || !add_claude_record_to_daily_tokens(daily_tokens, record) { + aggregation_failures = aggregation_failures.saturating_add(1); + } + }); + result.aggregation_failures = result + .aggregation_failures + .saturating_add(aggregation_failures); + result +} + +fn mark_claude_daily_token_coverage( + covered_days: &mut HashSet, + daily_tokens: &HashMap, + scan_result: ClaudeFileScanResult, +) { + if scan_result.is_complete() { + covered_days.extend(daily_tokens.keys().cloned()); + } else { + covered_days.clear(); } } diff --git a/rust/src/cost_scanner/claude_pricing.rs b/rust/src/cost_scanner/claude_pricing.rs index 89bcbe138a..5d3d251036 100644 --- a/rust/src/cost_scanner/claude_pricing.rs +++ b/rust/src/cost_scanner/claude_pricing.rs @@ -139,22 +139,16 @@ impl ClaudeScanPricingResolver { let cache_create_1h = cache_create_1h.min(cache_create); let cache_create_5m = cache_create.saturating_sub(cache_create_1h); - #[allow( - clippy::cast_possible_truncation, - reason = "clamped to i32::MAX before casting" - )] - let clamp = |value: u64| value.min(i32::MAX as u64) as i32; - let resolved = self.resolve(model); let billable = resolved.or_else(|| self.resolve(FALLBACK_CLAUDE_MODEL)); let base = billable .map(|pricing| { - CostUsagePricing::claude_cost_usd_from_resolution( + CostUsagePricing::claude_cost_usd_u64_from_resolution( pricing, - clamp(input), - clamp(cache_read), - clamp(cache_create_5m), - clamp(output), + input, + cache_read, + cache_create_5m, + output, ) }) .unwrap_or(0.0); diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index f054c6eb8f..6c0ebc5efa 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -1,5 +1,5 @@ use super::*; -use crate::core::{CodexForkAccountingState, CodexSessionLineage}; +use crate::core::{CodexForkAccountingState, CodexSessionLineage, CodexSessionMetadata}; mod cache_days; mod logical_target; @@ -13,6 +13,49 @@ use pending_range::{ }; use reconciliation::*; +#[derive(Debug)] +enum CodexAccountingMode { + Standard, + Baseline { + baseline: crate::core::CodexTotals, + paginated_continuation: bool, + remaining_inherited_totals: Option, + provenance: CodexBaselineProvenance, + }, + InferSubagent { + start_ordinal: Option, + }, + Unresolved, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CodexBaselineProvenance { + ValidatedParent { replaces_cached_state: bool }, + CachedValidatedParent, +} + +impl CodexAccountingMode { + fn is_unresolved(&self) -> bool { + matches!(self, Self::Unresolved) + } + + fn infers_subagent_baseline(&self) -> bool { + matches!(self, Self::InferSubagent { .. }) + } + + fn requires_cached_reparse(&self) -> bool { + matches!( + self, + Self::Baseline { + provenance: CodexBaselineProvenance::ValidatedParent { + replaces_cached_state: true + }, + .. + } + ) + } +} + fn summary_from_cached_report( report: &CachedCostReport, period_start: NaiveDate, @@ -41,65 +84,17 @@ fn summary_from_cached_report( } } -fn codex_fork_parent_is_safe(cache: &CostUsageCache, usage: &CostUsageFileUsage) -> bool { - let uses_parent_baseline = usage.codex_lineage.uses_parent_baseline() +fn codex_usage_uses_parent(usage: &CostUsageFileUsage) -> bool { + usage.codex_lineage.uses_parent_baseline() || (matches!(usage.codex_lineage, CodexSessionLineage::Root) - && usage.codex_forked_from_id.is_some()); - !uses_parent_baseline - || usage - .codex_forked_from_id - .as_deref() - .is_some_and(|parent_id| { - codex_parent_baseline(cache, parent_id, usage.codex_fork_timestamp.as_deref()) - .is_some() - }) + && usage.codex_forked_from_id.is_some()) } -/// Return a parent cumulative baseline only when exactly one cached session -/// identity is current, complete, timestamp-ordered, and safe to trust. -fn codex_parent_baseline( - cache: &CostUsageCache, - parent_session_id: &str, - child_fork_timestamp: Option<&str>, -) -> Option { - let mut baseline = None; - for (path_key, usage) in &cache.files { - if usage.codex_session_id.as_deref() != Some(parent_session_id) { - continue; - } - if usage.codex_unresolved_fork_parent - || usage.codex_token_timestamps_monotonic != Some(true) - { - return None; - } - let metadata = fs::metadata(path_key).ok()?; - if let (Some(expected), Some(actual)) = ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata), - ) && expected != &actual - { - return None; - } - #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] - let size = metadata.len().min(i64::MAX as u64) as i64; - if usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) - || usage.size != size - || usage.parsed_bytes.unwrap_or(0) < size - { - return None; - } - let last_totals = usage.last_totals.clone()?; - let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; - let child_fork_timestamp = child_fork_timestamp?; - if !JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) { - return None; - } - if baseline.replace(last_totals).is_some() { - // Duplicate identities make the dependency ambiguous. - return None; - } - } - baseline +fn codex_fork_uses_local_inference(usage: &CostUsageFileUsage) -> bool { + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) } fn is_codex_path_in_scan_window( @@ -129,6 +124,13 @@ struct CodexScanCandidate { mtime_unix_ms: i64, } +struct CodexPreparedCandidate { + path: PathBuf, + session_metadata: CodexSessionMetadata, + lineage_gate: CodexLineageGate, + parent_owner_expected: bool, +} + #[derive(Debug, Clone, Copy, Default)] struct CodexFileScanOutcome { bytes_read: i64, @@ -185,6 +187,7 @@ impl CostScanner { sessions_dirs: &[PathBuf], range: &CostUsageDayRange, cache: &CostUsageCache, + planner: &CodexLineagePlanner, cancel: Option<&AtomicBool>, stats: &mut CostScanStats, ) -> (Vec, bool) { @@ -246,7 +249,7 @@ impl CostScanner { }; let mtime_unix_ms = system_time_to_unix_ms(metadata.modified().ok()); let unchanged_complete = - cached_codex_file_is_complete_for_range(cache, &path_key, range); + cached_codex_file_is_complete_for_range(cache, planner, &path_key, range); if unchanged_complete { stats.files_seen = stats.files_seen.saturating_add(1); stats.files_skipped = stats.files_skipped.saturating_add(1); @@ -306,7 +309,10 @@ impl CostScanner { cancel: Option<&AtomicBool>, stats: &mut CostScanStats, ) { - let _ = self.parse_codex_file_bounded(path, range, summary, cache, cancel, stats, None); + let planner = CodexLineagePlanner::new(cache); + let _ = self.parse_codex_file_bounded( + path, range, summary, cache, cancel, stats, None, None, &planner, + ); } #[allow( @@ -322,11 +328,15 @@ impl CostScanner { cancel: Option<&AtomicBool>, stats: &mut CostScanStats, max_bytes_to_read: Option, + prepared_candidate: Option<&CodexPreparedCandidate>, + planner: &CodexLineagePlanner, ) -> CodexFileScanOutcome { if is_cancelled(cancel) { return CodexFileScanOutcome::default(); } - stats.files_seen = stats.files_seen.saturating_add(1); + if prepared_candidate.is_none() { + stats.files_seen = stats.files_seen.saturating_add(1); + } let metadata = match fs::metadata(path) { Ok(metadata) => metadata, @@ -356,20 +366,21 @@ impl CostScanner { }; } let cache_entry_is_fresh = |entry: &CostUsageFileUsage| { - cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) + cached_codex_file_is_fresh(cache, planner, entry, cache_covers_range, mtime_ms, size) }; - let identity_matches_cached = |entry: &CostUsageFileUsage| match ( - entry.codex_file_identity.as_ref(), - file_identity.as_ref(), - ) { - (Some(expected), Some(actual)) => expected == actual, - _ => false, + let identity_matches_cached = |entry: &CostUsageFileUsage| { + codex_file_identity_matches( + entry.codex_file_identity.as_deref(), + file_identity.as_deref(), + ) }; // The compact cache is authoritative for an unchanged file. Do this // before reading even the bounded metadata prefix; raw token history // is only needed after freshness fails or a fork needs reconciliation. if let Some(entry) = cached.as_ref() + && prepared_candidate + .is_none_or(|candidate| candidate.lineage_gate == CodexLineageGate::Eligible) && cache_entry_is_fresh(entry) && identity_matches_cached(entry) { @@ -386,13 +397,17 @@ impl CostScanner { }; } - stats.codex_metadata_read_paths.push(path_key.clone()); - stats.codex_read_receipt.metadata_reads = - stats.codex_read_receipt.metadata_reads.saturating_add(1); - let session_metadata = JsonlScanner::read_codex_session_metadata(path).unwrap_or_default(); - let cached_identity_matches = cached - .as_ref() - .is_some_and(|entry| entry.mtime_unix_ms == mtime_ms && entry.size == size); + let session_metadata = if let Some(prepared) = prepared_candidate { + prepared.session_metadata.clone() + } else { + stats.codex_metadata_read_paths.push(path_key.clone()); + stats.codex_read_receipt.metadata_reads = + stats.codex_read_receipt.metadata_reads.saturating_add(1); + JsonlScanner::read_codex_session_metadata(path).unwrap_or_default() + }; + // Cached lineage metadata belongs to a physical file, not merely a + // path/size/mtime tuple. Missing identity evidence fails closed. + let cached_identity_matches = cached.as_ref().is_some_and(identity_matches_cached); let codex_session_id = session_metadata.session_id.clone().or_else(|| { cached_identity_matches .then(|| cached.as_ref()?.codex_session_id.clone()) @@ -463,29 +478,34 @@ impl CostScanner { && state.history_base_thread_id == history_base_thread_id && state.fork_timestamp == codex_fork_timestamp }); - let fork_baseline = cached_fork_accounting_state - .as_ref() - .filter(|_| cached_fork_state_matches) - .and_then(|state| state.inherited_totals.clone()) - .or_else(|| { - is_fork - .then_some(codex_forked_from_id.as_deref()) - .flatten() - .and_then(|parent_id| { - codex_parent_baseline(cache, parent_id, codex_fork_timestamp.as_deref()) - }) - }); - let remaining_inherited_totals = cached_fork_accounting_state + let matching_cached_fork_state = cached_fork_accounting_state .as_ref() - .filter(|_| cached_fork_state_matches) - .and_then(|state| state.remaining_inherited_totals.clone()); + .filter(|_| cached_fork_state_matches); let paginated_continuation = is_fork && codex_forked_from_id.is_some() && history_base_thread_id .as_deref() .is_some_and(|history_base| Some(history_base) != codex_forked_from_id.as_deref()); + let lineage_gate = prepared_candidate + .map(|candidate| candidate.lineage_gate) + .unwrap_or_default(); + let parent_owner_expected = + prepared_candidate.is_some_and(|candidate| candidate.parent_owner_expected); + let lineage_decision = planner.decision_for_scan( + cache, + is_fork, + lineage_gate, + codex_forked_from_id.as_deref(), + codex_fork_timestamp.as_deref(), + parent_owner_expected, + ); + let accounting_mode = lineage_decision.accounting_mode( + matching_cached_fork_state, + &session_metadata, + paginated_continuation, + ); - if is_fork && fork_baseline.is_none() { + if accounting_mode.is_unresolved() { cache.files.insert( path_key, CostUsageFileUsage { @@ -515,9 +535,10 @@ impl CostScanner { } if let Some(entry) = &cached - && cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) - && (entry.codex_file_identity.is_none() || identity_matches_cached(entry)) + && cached_codex_file_is_fresh(cache, planner, entry, cache_covers_range, mtime_ms, size) + && identity_matches_cached(entry) && !cached_identity_changed + && !accounting_mode.requires_cached_reparse() { let (session_cost, has_tokens) = add_codex_days_map_to_summary(summary, &entry.days, range); @@ -525,11 +546,6 @@ impl CostScanner { summary.total_cost_usd += session_cost; summary.sessions_count += 1; } - if entry.codex_file_identity != file_identity { - let mut refreshed = entry.clone(); - refreshed.codex_file_identity = file_identity.clone(); - cache.files.insert(path_key.clone(), refreshed); - } stats.files_skipped = stats.files_skipped.saturating_add(1); return CodexFileScanOutcome { bytes_read: 0, @@ -623,22 +639,15 @@ impl CostScanner { } } - let parse_target_size = cached - .as_ref() - .and_then(|entry| codex_resumable_scan_target_size(size, entry)); - let parse_result = match if let Some(baseline) = fork_baseline.clone() { - JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( - path, - range, - baseline, - paginated_continuation, - remaining_inherited_totals.clone(), - cancel, - parse_target_size, - max_bytes_to_read, - ) - } else { - JsonlScanner::parse_codex_file_with_state_bounded( + let parse_target_size = (!accounting_mode.requires_cached_reparse()) + .then(|| { + cached + .as_ref() + .and_then(|entry| codex_resumable_scan_target_size(size, entry)) + }) + .flatten(); + let parse_result = match match &accounting_mode { + CodexAccountingMode::Standard => JsonlScanner::parse_codex_file_with_state_bounded( path, range, 0, @@ -648,7 +657,35 @@ impl CostScanner { None, cancel, max_bytes_to_read, - ) + ), + CodexAccountingMode::Baseline { + baseline, + paginated_continuation, + remaining_inherited_totals, + .. + } => JsonlScanner::parse_codex_file_with_state_bounded_fork_target_with_accounting( + path, + range, + baseline.clone(), + *paginated_continuation, + remaining_inherited_totals.clone(), + cancel, + parse_target_size, + max_bytes_to_read, + ), + CodexAccountingMode::InferSubagent { start_ordinal } => { + JsonlScanner::parse_codex_file_with_inferred_fork_baseline( + path, + range, + *start_ordinal, + cancel, + parse_target_size, + max_bytes_to_read, + ) + } + CodexAccountingMode::Unresolved => { + unreachable!("unresolved forks return before parsing") + } } { Ok(result) => result, Err(_) => return CodexFileScanOutcome::default(), @@ -656,7 +693,10 @@ impl CostScanner { stats.token_timestamp_comparisons = stats .token_timestamp_comparisons .saturating_add(parse_result.token_timestamp_comparisons); - if parse_result.fork_baseline_ambiguous { + if parse_result.fork_baseline_ambiguous + || (accounting_mode.infers_subagent_baseline() + && !parse_result.fork_baseline_locally_resolved) + { cache.files.insert( path_key, CostUsageFileUsage { @@ -696,18 +736,19 @@ impl CostScanner { bytes_read: parse_result.bytes_read, is_complete: parse_result.is_complete, }; - let codex_fork_accounting_state = if is_fork { - parse_result - .fork_baseline - .clone() - .map(|inherited_totals| CodexForkAccountingState { - session_id: codex_session_id.clone(), - forked_from_id: codex_forked_from_id.clone(), - history_base_thread_id: history_base_thread_id.clone(), - fork_timestamp: codex_fork_timestamp.clone(), - inherited_totals: Some(inherited_totals), - remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), - }) + let locally_resolved = parse_result.fork_baseline_locally_resolved; + let codex_fork_accounting_state = if is_fork + && (parse_result.fork_baseline.is_some() || parse_result.fork_baseline_locally_resolved) + { + Some(CodexForkAccountingState { + session_id: codex_session_id.clone(), + forked_from_id: codex_forked_from_id.clone(), + history_base_thread_id: history_base_thread_id.clone(), + fork_timestamp: codex_fork_timestamp.clone(), + inherited_totals: parse_result.fork_baseline.clone(), + remaining_inherited_totals: parse_result.remaining_inherited_totals.clone(), + locally_resolved, + }) } else { None }; diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index a98fb1e63e..dd42b96a98 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -1,7 +1,556 @@ use super::*; +#[cfg(test)] +use std::cell::Cell; +use std::collections::VecDeque; + +#[cfg(test)] +thread_local! { static CODEX_LINEAGE_GRAPH_BUILDS: Cell = const { Cell::new(0) }; } + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub(super) enum CodexLineageGate { + #[default] + Eligible, + Unsafe, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(super) enum CodexLineageDecision { + Root, + ParentAbsent, + ParentReady(crate::core::CodexTotals), + Unsafe, +} + +struct CodexLineageNode { + path: String, + session_id: Option, + parent_id: Option, + candidate_index: Option, + may_infer_missing_parent: bool, + may_author_parent: bool, + initially_unsafe: bool, +} + +struct CodexLineageGraph { + nodes: Vec, + session_owners: HashMap>, + parent_indices: Vec>, + gates: Vec, + candidate_node_indices: Vec, + ordered_candidate_indices: Vec, +} + +impl CodexLineageGraph { + fn new(cache: &CostUsageCache, candidates: Option<&[CodexPreparedCandidate]>) -> Self { + #[cfg(test)] + CODEX_LINEAGE_GRAPH_BUILDS.with(|builds| builds.set(builds.get() + 1)); + let candidate_paths = candidates + .into_iter() + .flatten() + .map(|candidate| candidate.path.to_string_lossy().to_string()) + .collect::>(); + let mut cached_paths = cache + .files + .keys() + .filter(|path| !candidate_paths.contains(*path)) + .cloned() + .collect::>(); + cached_paths.sort(); + + let mut nodes = Vec::with_capacity(cached_paths.len() + candidate_paths.len()); + for path in cached_paths { + let usage = &cache.files[&path]; + let uses_parent = codex_usage_uses_parent(usage); + let locally_inferred = codex_fork_uses_local_inference(usage); + nodes.push(CodexLineageNode { + path, + session_id: usage.codex_session_id.clone(), + parent_id: uses_parent + .then(|| usage.codex_forked_from_id.clone()) + .flatten(), + candidate_index: None, + may_infer_missing_parent: locally_inferred, + may_author_parent: !locally_inferred && !usage.codex_unresolved_fork_parent, + initially_unsafe: usage.codex_unresolved_fork_parent, + }); + } + + let mut candidate_node_indices = Vec::new(); + if let Some(candidates) = candidates { + candidate_node_indices.reserve(candidates.len()); + for (candidate_index, candidate) in candidates.iter().enumerate() { + let cached = cache + .files + .get(&candidate.path.to_string_lossy().to_string()); + let cached_identity_matches = cached.is_some_and(|usage| { + let Ok(metadata) = fs::metadata(&candidate.path) else { + return false; + }; + let expected = usage.codex_file_identity.as_deref(); + let actual = JsonlScanner::codex_file_identity(&candidate.path, &metadata); + codex_file_identity_matches(expected, actual.as_deref()) + }); + let metadata_owns_identity = candidate.session_metadata.session_id.is_some(); + let uses_parent = if metadata_owns_identity { + candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some() + } else if cached_identity_matches { + cached.is_some_and(super::codex_usage_uses_parent) + } else { + candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some() + }; + let session_id = candidate.session_metadata.session_id.clone().or_else(|| { + cached_identity_matches + .then(|| cached.and_then(|usage| usage.codex_session_id.clone())) + .flatten() + }); + let parent_id = if metadata_owns_identity { + candidate.session_metadata.forked_from_id.clone() + } else { + candidate + .session_metadata + .forked_from_id + .clone() + .or_else(|| { + (cached_identity_matches && uses_parent) + .then(|| { + cached.and_then(|usage| usage.codex_forked_from_id.clone()) + }) + .flatten() + }) + }; + // Freshly read identity-bearing metadata owns this candidate's + // current lineage state. Retain cached lineage flags only when + // the bounded metadata read could not establish an identity. + let cached_fallback = if metadata_owns_identity { + None + } else { + cached_identity_matches.then_some(cached).flatten() + }; + nodes.push(CodexLineageNode { + path: candidate.path.to_string_lossy().to_string(), + session_id, + parent_id: uses_parent.then_some(parent_id).flatten(), + candidate_index: Some(candidate_index), + may_infer_missing_parent: candidate.session_metadata.is_subagent + || cached_fallback.is_some_and(super::codex_fork_uses_local_inference), + may_author_parent: cached_fallback.is_none_or(|usage| { + !super::codex_fork_uses_local_inference(usage) + && !usage.codex_unresolved_fork_parent + }), + initially_unsafe: cached_fallback + .is_some_and(|usage| usage.codex_unresolved_fork_parent), + }); + candidate_node_indices.push(nodes.len() - 1); + } + } + + let mut session_owners = HashMap::>::new(); + for (index, node) in nodes.iter().enumerate() { + if let Some(session_id) = node.session_id.as_ref() { + session_owners + .entry(session_id.clone()) + .or_default() + .push(index); + } + } + + let mut gates = nodes + .iter() + .map(|node| { + if node.initially_unsafe { + CodexLineageGate::Unsafe + } else { + CodexLineageGate::Eligible + } + }) + .collect::>(); + for owners in session_owners.values().filter(|owners| owners.len() > 1) { + for &index in owners { + gates[index] = CodexLineageGate::Unsafe; + } + } + + let mut parent_indices = vec![None; nodes.len()]; + for (index, node) in nodes.iter().enumerate() { + let Some(parent_id) = node.parent_id.as_ref() else { + continue; + }; + match session_owners.get(parent_id) { + Some(owners) if owners.len() == 1 => parent_indices[index] = Some(owners[0]), + Some(_) => { + gates[index] = CodexLineageGate::Unsafe; + } + None if !node.may_infer_missing_parent => { + gates[index] = CodexLineageGate::Unsafe; + } + None => {} + } + } + + // This single topological pass both rejects cycles/unsafe ancestry and + // orders candidates. Cached-parent validation consumes the same gates. + let mut completed = vec![false; nodes.len()]; + let mut ordered_candidate_indices = Vec::with_capacity(candidate_node_indices.len()); + let mut children = vec![Vec::new(); nodes.len()]; + let mut ready = VecDeque::new(); + for (index, parent) in parent_indices.iter().enumerate() { + match parent { + Some(parent_index) => children[*parent_index].push(index), + None if gates[index] == CodexLineageGate::Eligible => ready.push_back(index), + None => {} + } + } + while let Some(index) = ready.pop_front() { + if completed[index] || gates[index] == CodexLineageGate::Unsafe { + continue; + } + completed[index] = true; + if let Some(candidate_index) = nodes[index].candidate_index { + ordered_candidate_indices.push(candidate_index); + } + if nodes[index].may_author_parent { + for child in &children[index] { + if gates[*child] == CodexLineageGate::Eligible { + ready.push_back(*child); + } + } + } + } + + for index in 0..nodes.len() { + if !completed[index] { + gates[index] = CodexLineageGate::Unsafe; + if let Some(candidate_index) = nodes[index].candidate_index { + ordered_candidate_indices.push(candidate_index); + } + } + } + + Self { + nodes, + session_owners, + parent_indices, + gates, + candidate_node_indices, + ordered_candidate_indices, + } + } + + fn unique_owner(&self, session_id: &str) -> Result, ()> { + match self.session_owners.get(session_id).map(Vec::as_slice) { + None | Some([]) => Ok(None), + Some([index]) => Ok(Some(*index)), + Some(_) => Err(()), + } + } + + fn apply_candidate_plan( + &self, + candidates: &mut Vec, + sessions_dirs: &[PathBuf], + range: &CostUsageDayRange, + ) -> Vec { + if !candidates.is_empty() { + for (candidate_index, candidate) in candidates.iter_mut().enumerate() { + let node_index = self.candidate_node_indices[candidate_index]; + candidate.lineage_gate = self.gates[node_index]; + candidate.parent_owner_expected = self.parent_indices[node_index].is_some(); + } + + let mut remaining = candidates.drain(..).map(Some).collect::>(); + for candidate_index in &self.ordered_candidate_indices { + candidates.push( + remaining[*candidate_index] + .take() + .expect("candidate is ordered once"), + ); + } + } + + // Keep the complete graph for parent resolution, but invalidate only + // unsafe cached nodes in the active range or in the ancestor closure + // required to resolve an active node. + let mut relevant = vec![false; self.nodes.len()]; + let mut pending = VecDeque::new(); + for (index, node) in self.nodes.iter().enumerate() { + if super::is_codex_path_in_scan_window(Path::new(&node.path), sessions_dirs, range) { + relevant[index] = true; + pending.push_back(index); + } + } + while let Some(index) = pending.pop_front() { + let Some(parent_id) = self.nodes[index].parent_id.as_deref() else { + continue; + }; + if let Some(owners) = self.session_owners.get(parent_id) { + for &owner in owners { + if !relevant[owner] { + relevant[owner] = true; + pending.push_back(owner); + } + } + } + } + + self.nodes + .iter() + .zip(&self.gates) + .enumerate() + .filter(|(index, (node, gate))| { + relevant[*index] + && node.candidate_index.is_none() + && **gate == CodexLineageGate::Unsafe + && !node.initially_unsafe + }) + .map(|(_, (node, _))| node.path.clone()) + .collect() + } +} + +pub(super) struct CodexLineagePlanner { + graph: Option, +} + +impl CodexLineagePlanner { + pub(super) fn new(cache: &CostUsageCache) -> Self { + Self { + graph: Self::needs_graph(cache, None).then(|| CodexLineageGraph::new(cache, None)), + } + } + + pub(super) fn plan_candidates_by_lineage( + cache: &CostUsageCache, + candidates: &mut Vec, + sessions_dirs: &[PathBuf], + range: &CostUsageDayRange, + ) -> (Self, Vec) { + let graph = Self::needs_graph(cache, Some(candidates)) + .then(|| CodexLineageGraph::new(cache, Some(candidates))); + let unsafe_paths = graph.as_ref().map_or_else(Vec::new, |graph| { + graph.apply_candidate_plan(candidates, sessions_dirs, range) + }); + (Self { graph }, unsafe_paths) + } + + fn needs_graph(cache: &CostUsageCache, candidates: Option<&[CodexPreparedCandidate]>) -> bool { + cache.files.values().any(super::codex_usage_uses_parent) + || candidates.is_some_and(|items| { + items.iter().any(|candidate| { + candidate.session_metadata.lineage.uses_parent_baseline() + || candidate.session_metadata.forked_from_id.is_some() + }) + }) + } + + #[cfg(test)] + pub(crate) fn reset_graph_build_count() { + CODEX_LINEAGE_GRAPH_BUILDS.with(|count| count.set(0)); + } + #[cfg(test)] + pub(crate) fn graph_build_count() -> usize { + CODEX_LINEAGE_GRAPH_BUILDS.with(Cell::get) + } + + fn graph(&self) -> Option<&CodexLineageGraph> { + self.graph.as_ref() + } + + pub(super) fn cached_usage_is_safe( + &self, + cache: &CostUsageCache, + usage: &CostUsageFileUsage, + ) -> bool { + let locally_resolved = super::codex_fork_uses_local_inference(usage); + match self.decision_for_usage(cache, usage) { + CodexLineageDecision::Root => true, + CodexLineageDecision::ParentAbsent => locally_resolved, + CodexLineageDecision::ParentReady(_) => !locally_resolved, + CodexLineageDecision::Unsafe => false, + } + } + + pub(super) fn decision_for_scan( + &self, + cache: &CostUsageCache, + uses_parent: bool, + gate: CodexLineageGate, + parent_id: Option<&str>, + fork_timestamp: Option<&str>, + parent_owner_expected: bool, + ) -> CodexLineageDecision { + if gate == CodexLineageGate::Unsafe { + return CodexLineageDecision::Unsafe; + } + if !uses_parent { + return CodexLineageDecision::Root; + } + parent_id.map_or(CodexLineageDecision::Unsafe, |parent_id| { + self.resolve_parent(cache, parent_id, fork_timestamp, parent_owner_expected) + }) + } + + pub(super) fn decision_for_usage( + &self, + cache: &CostUsageCache, + usage: &CostUsageFileUsage, + ) -> CodexLineageDecision { + if usage.codex_unresolved_fork_parent { + return CodexLineageDecision::Unsafe; + } + if !super::codex_usage_uses_parent(usage) { + return CodexLineageDecision::Root; + } + usage + .codex_forked_from_id + .as_deref() + .map_or(CodexLineageDecision::Unsafe, |parent_id| { + self.resolve_parent( + cache, + parent_id, + usage.codex_fork_timestamp.as_deref(), + false, + ) + }) + } + + /// Resolve one parent identity through the persisted graph. Absence is + /// distinct from ambiguity and transitive unsafety so local inference is + /// allowed only when no owner exists at all. + fn resolve_parent( + &self, + cache: &CostUsageCache, + parent_session_id: &str, + child_fork_timestamp: Option<&str>, + parent_owner_expected: bool, + ) -> CodexLineageDecision { + let Some(graph) = self.graph() else { + return CodexLineageDecision::Unsafe; + }; + let node_index = match graph.unique_owner(parent_session_id) { + Ok(None) if !parent_owner_expected => return CodexLineageDecision::ParentAbsent, + Ok(None) | Err(()) => return CodexLineageDecision::Unsafe, + Ok(Some(index)) => index, + }; + self.parent_owner_baseline(cache, node_index, child_fork_timestamp) + .map_or( + CodexLineageDecision::Unsafe, + CodexLineageDecision::ParentReady, + ) + } + + fn parent_owner_baseline( + &self, + cache: &CostUsageCache, + node_index: usize, + child_fork_timestamp: Option<&str>, + ) -> Option { + let graph = self.graph()?; + let node = graph.nodes.get(node_index)?; + if graph.gates[node_index] == CodexLineageGate::Unsafe || !node.may_author_parent { + return None; + } + let usage = cache.files.get(&node.path)?; + if usage.codex_unresolved_fork_parent + || usage.codex_token_timestamps_monotonic != Some(true) + || super::codex_fork_uses_local_inference(usage) + { + return None; + } + + if super::codex_usage_uses_parent(usage) { + let inherited = usage + .codex_fork_accounting_state + .as_ref()? + .inherited_totals + .as_ref()?; + let parent_index = graph.parent_indices[node_index]?; + let baseline = self.parent_owner_baseline( + cache, + parent_index, + usage.codex_fork_timestamp.as_deref(), + )?; + if &baseline != inherited { + return None; + } + } + + let metadata = fs::metadata(&node.path).ok()?; + let expected_identity = usage.codex_file_identity.as_ref()?; + let actual_identity = JsonlScanner::codex_file_identity(Path::new(&node.path), &metadata)?; + if expected_identity != &actual_identity { + return None; + } + #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] + let size = metadata.len().min(i64::MAX as u64) as i64; + if usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) + || usage.size != size + || usage.parsed_bytes.unwrap_or(0) < size + { + return None; + } + let last_totals = usage.last_totals.clone()?; + let last_token_timestamp = usage.codex_last_token_timestamp.as_deref()?; + let child_fork_timestamp = child_fork_timestamp?; + JsonlScanner::codex_timestamp_at_or_before(last_token_timestamp, child_fork_timestamp) + .then_some(last_totals) + } +} + +impl CodexLineageDecision { + pub(super) fn accounting_mode( + &self, + matching_cached_state: Option<&CodexForkAccountingState>, + metadata: &CodexSessionMetadata, + paginated_continuation: bool, + ) -> CodexAccountingMode { + match self { + Self::Root => CodexAccountingMode::Standard, + Self::Unsafe => CodexAccountingMode::Unresolved, + Self::ParentReady(baseline) => { + let replaces_cached_state = matching_cached_state.is_some_and(|state| { + state.locally_resolved || state.inherited_totals.as_ref() != Some(baseline) + }); + let cached_parent_state = matching_cached_state.filter(|state| { + !state.locally_resolved && state.inherited_totals.as_ref() == Some(baseline) + }); + CodexAccountingMode::Baseline { + baseline: baseline.clone(), + paginated_continuation, + remaining_inherited_totals: cached_parent_state + .and_then(|state| state.remaining_inherited_totals.clone()), + provenance: CodexBaselineProvenance::ValidatedParent { + replaces_cached_state, + }, + } + } + Self::ParentAbsent => { + if let Some(state) = matching_cached_state.filter(|state| !state.locally_resolved) + && let Some(baseline) = state.inherited_totals.clone() + { + return CodexAccountingMode::Baseline { + baseline, + paginated_continuation, + remaining_inherited_totals: state.remaining_inherited_totals.clone(), + provenance: CodexBaselineProvenance::CachedValidatedParent, + }; + } + if metadata.is_subagent { + CodexAccountingMode::InferSubagent { + start_ordinal: metadata.subagent_history_start_ordinal, + } + } else { + CodexAccountingMode::Unresolved + } + } + } + } +} pub(super) fn cached_codex_file_is_fresh( cache: &CostUsageCache, + planner: &CodexLineagePlanner, entry: &CostUsageFileUsage, cache_covers_range: bool, mtime_unix_ms: i64, @@ -13,11 +562,18 @@ pub(super) fn cached_codex_file_is_fresh( && entry.size == size && codex_scan_target_size(entry) == size && entry.parsed_bytes.unwrap_or(0) >= size - && super::codex_fork_parent_is_safe(cache, entry) + && planner.cached_usage_is_safe(cache, entry) +} + +pub(super) fn codex_file_identity_matches(expected: Option<&str>, actual: Option<&str>) -> bool { + expected + .zip(actual) + .is_some_and(|(expected, actual)| expected == actual) } pub(super) fn cached_codex_file_is_complete_for_range( cache: &CostUsageCache, + planner: &CodexLineagePlanner, path_key: &str, range: &CostUsageDayRange, ) -> bool { @@ -26,14 +582,10 @@ pub(super) fn cached_codex_file_is_complete_for_range( let Ok(metadata) = fs::metadata(path_key) else { return false; }; - let identity_matches = match ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_ref(), - ) { - (Some(expected), Some(actual)) => expected == actual, - (Some(_), None) => false, - (None, _) => true, - }; + let identity_matches = codex_file_identity_matches( + usage.codex_file_identity.as_deref(), + JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_deref(), + ); #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] let size = metadata.len().min(i64::MAX as u64) as i64; identity_matches @@ -42,10 +594,59 @@ pub(super) fn cached_codex_file_is_complete_for_range( && codex_scan_target_size(usage) == size && usage.parsed_bytes.unwrap_or(0) >= size && !usage.codex_unresolved_fork_parent - && super::codex_fork_parent_is_safe(cache, usage) + // Reconsider locally inferred children after this pass has + // had a chance to discover and cache their parent. + && !super::codex_fork_uses_local_inference(usage) + && planner.cached_usage_is_safe(cache, usage) }) } +/// Process cached local-inference children after all other candidates. A +/// parent discovered in this pass must enter the cache before its unchanged +/// child can decide whether the inferred baseline is still authoritative. +pub(super) fn defer_codex_locally_inferred_candidates( + candidates: &mut Vec, + cache: &CostUsageCache, +) { + if candidates.len() < 2 { + return; + } + + let mut other = Vec::with_capacity(candidates.len()); + let mut locally_inferred = Vec::new(); + for candidate in candidates.drain(..) { + let path_key = candidate.path.to_string_lossy(); + if cache + .files + .get(path_key.as_ref()) + .is_some_and(super::codex_fork_uses_local_inference) + { + locally_inferred.push(candidate); + } else { + other.push(candidate); + } + } + other.extend(locally_inferred); + candidates.extend(other); +} + +pub(super) fn invalidate_codex_unsafe_lineage(cache: &mut CostUsageCache, paths: &[String]) { + for path in paths { + let Some(usage) = cache.files.get_mut(path) else { + continue; + }; + usage.days.clear(); + usage.parsed_bytes = Some(0); + usage.codex_scan_target_size = None; + usage.last_model = None; + usage.last_totals = None; + usage.codex_token_timestamps_monotonic = None; + usage.codex_last_token_timestamp = None; + usage.codex_fork_accounting_state = None; + usage.codex_unresolved_fork_parent = true; + } +} + /// Give paths already in the durable queue their saved turn before newly /// discovered dirty paths. The scanner appends unfinished paths after this /// pass, making the queue a round-robin cursor instead of a newest-first loop. diff --git a/rust/src/cost_scanner/codex/reconciliation.rs b/rust/src/cost_scanner/codex/reconciliation.rs index b16d211b3b..8451c3f72b 100644 --- a/rust/src/cost_scanner/codex/reconciliation.rs +++ b/rust/src/cost_scanner/codex/reconciliation.rs @@ -108,14 +108,10 @@ fn codex_pending_path_affects_current_window( if codex_logical_target_has_unconsumed_tail(observed_size, usage) { return true; } - let identity_matches = match ( - usage.codex_file_identity.as_ref(), - JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_ref(), - ) { - (Some(expected), Some(actual)) => expected == actual, - (Some(_), None) => false, - (None, _) => true, - }; + let identity_matches = super::codex_file_identity_matches( + usage.codex_file_identity.as_deref(), + JsonlScanner::codex_file_identity(Path::new(path_key), &metadata).as_deref(), + ); if !identity_matches || usage.mtime_unix_ms != system_time_to_unix_ms(metadata.modified().ok()) || usage.size != observed_size @@ -270,6 +266,7 @@ mod tests { let old_usage = cache.files.get_mut(&old_key).unwrap(); old_usage.mtime_unix_ms = system_time_to_unix_ms(metadata.modified().ok()); old_usage.size = i64::try_from(metadata.len()).unwrap(); + old_usage.codex_file_identity = JsonlScanner::codex_file_identity(&old_path, &metadata); let range = active_range(); assert!(codex_current_window_is_established(&cache, &range)); @@ -279,6 +276,32 @@ mod tests { assert_eq!(report.sessions_count, 1); } + #[test] + fn historical_pending_entry_without_or_mismatched_identity_blocks_publication() { + let root = tempfile::tempdir().unwrap(); + let old_path = root.path().join("old.jsonl"); + let current_path = root.path().join("current.jsonl"); + std::fs::write(&old_path, vec![0_u8; 100]).unwrap(); + std::fs::write(¤t_path, vec![0_u8; 100]).unwrap(); + let old_key = old_path.to_string_lossy().into_owned(); + let current_key = current_path.to_string_lossy().into_owned(); + let metadata = std::fs::metadata(&old_path).unwrap(); + let range = active_range(); + + for cached_identity in [None, Some("different-file".to_string())] { + let mut cache = historical_pending_cache(&old_key, ¤t_key); + let old_usage = cache.files.get_mut(&old_key).unwrap(); + old_usage.mtime_unix_ms = system_time_to_unix_ms(metadata.modified().ok()); + old_usage.size = i64::try_from(metadata.len()).unwrap(); + old_usage.codex_file_identity = cached_identity; + + assert!(codex_pending_path_affects_current_window( + &cache, &old_key, &range + )); + assert!(!codex_current_window_is_established(&cache, &range)); + } + } + #[test] fn metadata_failure_blocks_historical_pending_publication() { let old_path = r"C:\sessions\missing-old.jsonl"; diff --git a/rust/src/cost_scanner/codex/scan.rs b/rust/src/cost_scanner/codex/scan.rs index e983adfd6a..94cb91f701 100644 --- a/rust/src/cost_scanner/codex/scan.rs +++ b/rust/src/cost_scanner/codex/scan.rs @@ -177,8 +177,15 @@ pub(super) fn scan_codex_detailed_with_cache( cache.codex_pending_scan_root_paths = pending_scan.root_paths.clone(); cache.codex_pending_scan_timezone = Some(pending_scan.timezone.clone()); - let (mut candidates, discovery_complete) = - scanner.collect_codex_candidates(&sessions_dirs, scan_range, &cache, cancel, &mut stats); + let cached_lineage = CodexLineagePlanner::new(&cache); + let (mut candidates, discovery_complete) = scanner.collect_codex_candidates( + &sessions_dirs, + scan_range, + &cache, + &cached_lineage, + cancel, + &mut stats, + ); let candidate_limit = if scanner.options.codex_candidate_limit == 0 { usize::MAX } else { @@ -197,43 +204,78 @@ pub(super) fn scan_codex_detailed_with_cache( let mut bytes_read_this_refresh = 0_i64; let mut pending_next = cache.codex_pending_paths.clone(); let pending_paths_before_pass = cache.codex_pending_paths.clone(); + let mut invalidated_unsafe_lineage = false; prioritize_codex_pending_candidates(&mut candidates, &pending_paths_before_pass); + defer_codex_locally_inferred_candidates(&mut candidates, &cache); if discovery_complete && !is_cancelled(cancel) { - pending_next - .retain(|path| !cached_codex_file_is_complete_for_range(&cache, path, scan_range)); + pending_next.retain(|path| { + !cached_codex_file_is_complete_for_range(&cache, &cached_lineage, path, scan_range) + }); } - let mut incomplete_processed = Vec::new(); - for (index, candidate) in candidates.iter().enumerate() { - if is_cancelled(cancel) - || index >= candidate_limit - || bytes_read_this_refresh >= refresh_byte_limit - { - for deferred in &candidates[index..] { - let key = deferred.path.to_string_lossy().to_string(); - if !pending_next.contains(&key) { - pending_next.push(key); - } + // Admit one bounded set, inspect each admitted candidate once, and order + // that set by lineage before reading token history. This makes cold + // child-before-parent scans parent-first without a second parse pass. + let deferred_candidates = candidates.split_off(candidate_limit.min(candidates.len())); + let deferred_paths = deferred_candidates + .into_iter() + .map(|candidate| candidate.path) + .collect::>(); + let mut work_queue = Vec::with_capacity(candidates.len()); + let mut cancelled_during_preparation = Vec::new(); + for candidate in candidates { + if is_cancelled(cancel) { + cancelled_during_preparation.push(candidate.path); + continue; + } + let key = candidate.path.to_string_lossy().to_string(); + stats.files_seen = stats.files_seen.saturating_add(1); + stats.codex_metadata_read_paths.push(key); + stats.codex_read_receipt.metadata_reads = + stats.codex_read_receipt.metadata_reads.saturating_add(1); + work_queue.push(CodexPreparedCandidate { + session_metadata: JsonlScanner::read_codex_session_metadata(&candidate.path) + .unwrap_or_default(), + path: candidate.path, + lineage_gate: CodexLineageGate::Eligible, + parent_owner_expected: false, + }); + } + let mut unprocessed = Vec::new(); + let cancelled_before_plan = !cancelled_during_preparation.is_empty() || is_cancelled(cancel); + let lineage_planner = if cancelled_before_plan { + unprocessed.extend(work_queue.drain(..).map(|candidate| candidate.path)); + unprocessed.extend(cancelled_during_preparation); + cached_lineage + } else { + let (planner, unsafe_cached_paths) = CodexLineagePlanner::plan_candidates_by_lineage( + &cache, + &mut work_queue, + &sessions_dirs, + scan_range, + ); + invalidated_unsafe_lineage = !unsafe_cached_paths.is_empty(); + if invalidated_unsafe_lineage { + cache.previous_report = None; + } + invalidate_codex_unsafe_lineage(&mut cache, &unsafe_cached_paths); + for path in unsafe_cached_paths { + if !pending_next.contains(&path) { + pending_next.push(path); } - stats.files_deferred = stats.files_deferred.saturating_add( - u32::try_from((candidates.len() - index).min(u32::MAX as usize)) - .unwrap_or(u32::MAX), - ); - break; } + planner + }; + let mut incomplete_processed = Vec::new(); + for (index, candidate) in work_queue.iter().enumerate() { let refresh_remaining = refresh_byte_limit.saturating_sub(bytes_read_this_refresh); let allowance = per_file_limit.min(refresh_remaining); - if allowance <= 0 { - for deferred in &candidates[index..] { - let key = deferred.path.to_string_lossy().to_string(); - if !pending_next.contains(&key) { - pending_next.push(key); - } - } - stats.files_deferred = stats.files_deferred.saturating_add( - u32::try_from((candidates.len() - index).min(u32::MAX as usize)) - .unwrap_or(u32::MAX), + if is_cancelled(cancel) || allowance <= 0 { + unprocessed.extend( + work_queue[index..] + .iter() + .map(|candidate| candidate.path.clone()), ); break; } @@ -246,6 +288,8 @@ pub(super) fn scan_codex_detailed_with_cache( cancel, &mut stats, Some(allowance), + Some(candidate), + &lineage_planner, ); bytes_read_this_refresh = bytes_read_this_refresh.saturating_add(outcome.bytes_read.max(0)); stats.codex_bytes_read = stats @@ -275,6 +319,16 @@ pub(super) fn scan_codex_detailed_with_cache( apply_codex_source_row_plan(&mut cache, &key, plan); } } + unprocessed.extend(deferred_paths); + stats.files_deferred = stats.files_deferred.saturating_add( + u32::try_from(unprocessed.len().min(u32::MAX as usize)).unwrap_or(u32::MAX), + ); + for path in unprocessed { + let key = path.to_string_lossy().to_string(); + if !pending_next.contains(&key) { + pending_next.push(key); + } + } pending_next.extend(incomplete_processed); let mut pruned_paths_pending = Vec::new(); @@ -344,7 +398,7 @@ pub(super) fn scan_codex_detailed_with_cache( // the range so unchanged files stay on the cache fast path. cache.scan_since_key = Some(scan_range.scan_since_key.clone()); cache.scan_until_key = Some(scan_range.scan_until_key.clone()); - } else if cache.previous_report.is_none() { + } else if cache.previous_report.is_none() && !invalidated_unsafe_lineage { cache.previous_report = established_report_before_scan; } if !is_cancelled(cancel) { diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 3b16c92a1c..69f0c8dffa 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -395,7 +395,7 @@ fn counts_claude_usage_once_across_duplicate_records() { assert_eq!(record.output, 50); assert_eq!(record.cache_create, 10); assert_eq!(record.cache_read, 20); - assert!(record.cost > 0.0); + assert!(record.cost.is_some_and(|cost| cost > 0.0)); let cutoff = DateTime::parse_from_rfc3339("2026-01-01T00:00:00Z") .unwrap() @@ -495,6 +495,157 @@ fn malformed_claude_history_stays_unknown_while_valid_empty_history_is_known_zer assert!(!malformed_summary.known_zero); } +#[test] +fn claude_daily_token_coverage_requires_a_complete_valid_scan() { + let root = tempfile::tempdir().unwrap(); + let cutoff = Utc::now() - Duration::days(1); + let valid_path = root.path().join("valid.jsonl"); + let timestamp = Utc::now() - Duration::hours(1); + let today = timestamp + .with_timezone(&Local) + .date_naive() + .format("%Y-%m-%d") + .to_string(); + std::fs::write( + &valid_path, + format!( + "{}\n", + claude_transcript_line( + ×tamp.to_rfc3339(), + "requestId", + "req_valid", + "msg_valid" + ) + ), + ) + .unwrap(); + + let mut valid_tokens = HashMap::from([(today.clone(), 0)]); + let valid_result = scan_claude_file_for_daily_tokens( + &valid_path, + &cutoff, + &mut HashSet::new(), + &mut ClaudeScanPricingResolver::default(), + &mut valid_tokens, + ); + assert!(valid_result.is_complete()); + let mut covered_days = HashSet::new(); + mark_claude_daily_token_coverage(&mut covered_days, &valid_tokens, valid_result); + assert!(covered_days.contains(&today)); + + let assert_uncovered = |path: &Path| { + let mut daily_tokens = HashMap::from([(today.clone(), 0)]); + let result = scan_claude_file_for_daily_tokens( + path, + &cutoff, + &mut HashSet::new(), + &mut ClaudeScanPricingResolver::default(), + &mut daily_tokens, + ); + assert!(!result.is_complete()); + let mut covered_days = HashSet::from(["stale-coverage".to_string()]); + mark_claude_daily_token_coverage(&mut covered_days, &daily_tokens, result); + assert!(covered_days.is_empty()); + result + }; + + let malformed_path = root.path().join("malformed.jsonl"); + std::fs::write(&malformed_path, b"{malformed\n").unwrap(); + assert_eq!(assert_uncovered(&malformed_path).malformed_lines, 1); + + let incomplete_path = root.path().join("incomplete.jsonl"); + std::fs::write( + &incomplete_path, + r#"{"type":"assistant","message":{"id":"msg_preliminary","model":"gpt-5.6-sol","stop_reason":null,"usage":{"input_tokens":1000}}}"#, + ) + .unwrap(); + assert_eq!(assert_uncovered(&incomplete_path).incomplete_requests, 1); + + let missing_timestamp_path = root.path().join("missing-timestamp.jsonl"); + std::fs::write( + &missing_timestamp_path, + r#"{"type":"assistant","requestId":"req_no_timestamp","message":{"id":"msg_no_timestamp","model":"claude-sonnet-4-6","usage":{"input_tokens":1000,"output_tokens":500}}}"#, + ) + .unwrap(); + assert_eq!( + assert_uncovered(&missing_timestamp_path).aggregation_failures, + 1 + ); + + let unreadable_path = root.path().join("missing.jsonl"); + assert_eq!(assert_uncovered(&unreadable_path).read_failures, 1); + + let scanner = CostScanner::new(1); + let missing_directory = root.path().join("missing-directory"); + let traversal_read_failures = + scanner.walk_claude_files(&missing_directory, &cutoff, None, &mut |_| {}); + assert_eq!(traversal_read_failures, 1); + let mut covered_days = HashSet::from([today]); + mark_claude_daily_token_coverage( + &mut covered_days, + &valid_tokens, + ClaudeFileScanResult { + read_failures: traversal_read_failures, + ..ClaudeFileScanResult::default() + }, + ); + assert!(covered_days.is_empty()); +} + +#[test] +fn public_claude_daily_token_dispatch_reports_incomplete_fixture_scans() { + const CHILD_MARKER: &str = "CODEXBAR_CLAUDE_DAILY_TOKEN_TEST_CHILD"; + const CHILD_DONE: &str = "isolated Claude daily-history fixture verified"; + if std::env::var_os(CHILD_MARKER).is_some() { + let config_dir = std::env::var_os("CLAUDE_CONFIG_DIR") + .map(PathBuf::from) + .expect("child receives isolated Claude config directory"); + let projects_dir = config_dir.join("projects"); + let project_dir = projects_dir.join("fixture-project"); + let (complete_history, incomplete) = get_daily_token_history("claude", 1); + assert!(!incomplete, "valid fixture scan should establish coverage"); + assert!(complete_history.iter().any(|(_, tokens)| *tokens > 0)); + + std::fs::write(project_dir.join("malformed.jsonl"), b"{malformed\n").unwrap(); + let (partial_history, incomplete) = get_daily_token_history("claude", 1); + assert!( + incomplete, + "malformed fixture should leave coverage incomplete" + ); + assert_eq!(partial_history, complete_history); + println!("{CHILD_DONE}"); + return; + } + + let config_dir = tempfile::tempdir().unwrap(); + let project_dir = config_dir.path().join("projects").join("fixture-project"); + std::fs::create_dir_all(&project_dir).unwrap(); + let timestamp = Utc::now().to_rfc3339(); + std::fs::write( + project_dir.join("valid.jsonl"), + format!( + "{}\n", + claude_transcript_line(×tamp, "requestId", "req_public", "msg_public") + ), + ) + .unwrap(); + + let test_thread = std::thread::current(); + let test_name = test_thread.name().expect("test harness names this thread"); + let output = std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", test_name, "--nocapture", "--test-threads=1"]) + .env(CHILD_MARKER, "1") + .env("CLAUDE_CONFIG_DIR", config_dir.path()) + .output() + .expect("spawn isolated exact-test child"); + assert!( + output.status.success() && String::from_utf8_lossy(&output.stdout).contains(CHILD_DONE), + "fixture child failed\nstdout:\n{}\nstderr:\n{}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); +} + #[test] fn classifies_vertex_ai_claude_metadata_without_changing_anthropic_rows() { let cases = [ @@ -552,9 +703,17 @@ fn shared_claude_reader_excludes_vertex_rows_but_keeps_anthropic_usage() { let anthropic = format!( r#"{{"type":"assistant","timestamp":"{timestamp}","requestId":"req_anthropic","message":{{"id":"msg_anthropic","model":"claude-sonnet-4-6","usage":{{"input_tokens":10,"output_tokens":5}}}}}}"# ); - let vertex = format!( - r#"{{"type":"assistant","timestamp":"{timestamp}","requestId":"req_vrtx_123","message":{{"id":"msg_vrtx_123","model":"claude-sonnet-4-6","usage":{{"input_tokens":1000,"output_tokens":500}}}}}}"# - ); + let vertex = serde_json::json!({ + "type": "assistant", + "timestamp": timestamp, + "requestId": "req_vrtx_123", + "message": { + "id": "msg_vrtx_123", + "model": "claude-sonnet-4-6", + "usage": {"input_tokens": u64::MAX, "output_tokens": u64::MAX} + } + }) + .to_string(); std::fs::write(&path, format!("{anthropic}\n{vertex}\n")).unwrap(); let cutoff = Utc::now() - Duration::days(30); @@ -569,6 +728,86 @@ fn shared_claude_reader_excludes_vertex_rows_but_keeps_anthropic_usage() { let _removed = std::fs::remove_file(&path); } +#[test] +fn oversized_claude_history_preserves_independent_components_and_fails_closed() { + let first: ClaudeEvent = serde_json::from_str(&format!( + r#"{{"type":"assistant","timestamp":"2026-09-20T12:00:00Z","requestId":"req_overflow_1","message":{{"id":"msg_overflow_1","model":"claude-sonnet-4-6","usage":{{"input_tokens":{},"output_tokens":2}}}}}}"#, + u64::MAX + )) + .unwrap(); + let second: ClaudeEvent = serde_json::from_str( + r#"{"type":"assistant","timestamp":"2026-09-20T12:01:00Z","requestId":"req_overflow_2","message":{"id":"msg_overflow_2","model":"claude-sonnet-4-6","usage":{"input_tokens":1,"output_tokens":3}}}"#, + ) + .unwrap(); + let first = claude_usage_record_from_event(&first).expect("first usage row"); + let second = claude_usage_record_from_event(&second).expect("second usage row"); + let mut summary = CostSummary::default(); + + assert!(add_claude_record_to_summary(&mut summary, &first)); + assert!(!add_claude_record_to_summary(&mut summary, &second)); + assert_eq!(summary.input_tokens, u64::MAX); + assert_eq!(summary.output_tokens, 5); + assert!(summary.total_cost_usd.is_finite()); + + finalize_claude_summary( + &mut summary, + true, + ClaudeFileScanResult { + counted: 2, + aggregation_failures: 1, + ..ClaudeFileScanResult::default() + }, + false, + ); + assert!(!summary.history_coverage_established); + assert!(!summary.known_zero); +} + +#[test] +fn oversized_single_claude_row_keeps_cost_but_marks_combined_quota_tokens_unknown() { + let event: ClaudeEvent = serde_json::from_str(&format!( + r#"{{"type":"assistant","timestamp":"2026-09-20T12:00:00Z","requestId":"req_combined_overflow","message":{{"id":"msg_combined_overflow","model":"claude-sonnet-4-6","usage":{{"input_tokens":{},"output_tokens":1}}}}}}"#, + u64::MAX + )) + .unwrap(); + let record = claude_usage_record_from_event(&event).expect("usage row"); + let quota = quota_history_record_from_usage(&record).expect("timestamped quota row"); + + assert!(record.cost.is_some_and(f64::is_finite)); + assert_eq!(quota.tokens, None); + assert!(!quota.tokens_are_complete); + assert!(quota.cost_usd.is_some_and(f64::is_finite)); + assert!(quota.cost_is_complete); +} + +#[test] +fn nonfinite_claude_price_is_unknown_instead_of_zero() { + let snapshot = crate::core::ModelsDevPricingSnapshot::from_catalog_json_for_tests( + r#"{ + "anthropic": {"models": {"claude-test-extreme-price": { + "id": "claude-test-extreme-price", "cost": {"input": 1e308, "output": 1} + }}} + }"#, + ) + .expect("pricing fixture"); + let mut pricing = ClaudeScanPricingResolver::with_snapshot(snapshot); + let event: ClaudeEvent = serde_json::from_str(&format!( + r#"{{"type":"assistant","timestamp":"2026-09-20T12:00:00Z","requestId":"req_nonfinite","message":{{"id":"msg_nonfinite","model":"claude-test-extreme-price","usage":{{"input_tokens":{},"output_tokens":1}}}}}}"#, + u64::MAX + )) + .unwrap(); + let record = + claude_usage_record_from_event_with_pricing(&event, &mut pricing).expect("usage row"); + let mut summary = CostSummary::default(); + + assert_eq!(record.cost, None); + assert!(!add_claude_record_to_summary(&mut summary, &record)); + assert_eq!(summary.input_tokens, u64::MAX); + assert_eq!(summary.output_tokens, 1); + assert_eq!(summary.total_cost_usd, 0.0); + assert!(!summary.known_zero); +} + fn claude_transcript_line( timestamp: &str, request_key: &str, @@ -626,14 +865,15 @@ fn daily_history_dedups_across_files_and_buckets_by_local_day() { .to_string() }; let mut daily_costs = HashMap::new(); - daily_costs.insert(day_key(&day_one), Some(0.0)); - daily_costs.insert(day_key(&day_two), Some(0.0)); + daily_costs.insert(day_key(&day_one), None); + daily_costs.insert(day_key(&day_two), None); + let mut unknown_cost_dates = HashSet::new(); let cutoff = Utc::now() - Duration::days(30); let mut seen = HashSet::new(); for path in [&file_a, &file_b] { for_each_claude_usage_record(path, &cutoff, &mut seen, None, |record| { - add_claude_record_to_daily_costs(&mut daily_costs, record); + add_claude_record_to_daily_costs(&mut daily_costs, &mut unknown_cost_dates, record); }); } @@ -652,6 +892,65 @@ fn daily_history_dedups_across_files_and_buckets_by_local_day() { let _removed_b = std::fs::remove_file(&file_b); } +fn claude_daily_cost_record(timestamp: DateTime, cost: Option) -> ClaudeUsageRecord { + ClaudeUsageRecord { + model: "claude-test".to_string(), + pricing_known: cost.is_some(), + timestamp: Some(timestamp), + dedup_key: None, + input: 1, + output: 1, + cache_create: 0, + cache_read: 0, + cost, + } +} + +#[test] +fn unknown_claude_cost_date_cannot_be_restored_by_later_priced_record() { + let timestamp = Utc::now(); + let day = timestamp + .with_timezone(&Local) + .date_naive() + .format("%Y-%m-%d") + .to_string(); + let mut daily_costs = HashMap::from([(day.clone(), None)]); + let mut unknown_cost_dates = HashSet::new(); + + assert!(add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + &claude_daily_cost_record(timestamp, Some(0.75)), + )); + assert!(!add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + &claude_daily_cost_record(timestamp, None), + )); + assert!(!add_claude_record_to_daily_costs( + &mut daily_costs, + &mut unknown_cost_dates, + &claude_daily_cost_record(timestamp, Some(1.25)), + )); + + assert_eq!(daily_costs[&day], None); + assert!(unknown_cost_dates.contains(&day)); +} + +#[test] +fn claude_daily_zero_fill_preserves_unknown_dates_and_fills_untouched_dates() { + let unknown_day = "2026-09-22".to_string(); + let untouched_day = "2026-09-23".to_string(); + let mut daily_costs = + HashMap::from([(unknown_day.clone(), None), (untouched_day.clone(), None)]); + let unknown_cost_dates = HashSet::from([unknown_day.clone()]); + + zero_fill_uninitialized_claude_daily_costs(&mut daily_costs, &unknown_cost_dates); + + assert_eq!(daily_costs[&unknown_day], None); + assert_eq!(daily_costs[&untouched_day], Some(0.0)); +} + #[test] fn claude_scan_counts_final_incomplete_jsonl_line() { let path = @@ -1449,10 +1748,16 @@ fn cost_scan_second_pass_skips_unchanged_files_via_cache() { // 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. + CodexLineagePlanner::reset_graph_build_count(); 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!( + CodexLineagePlanner::graph_build_count(), + 0, + "warm root-only scan must bypass lineage graph construction" + ); assert!(stats2.codex_metadata_read_paths.is_empty()); assert!(stats2.codex_history_read_paths.is_empty()); assert_eq!(stats2.codex_read_receipt, Default::default()); @@ -3036,6 +3341,12 @@ fn incomplete_or_buffered_empty_codex_fragment_is_not_marked_complete() { assert!(buffered_cache.codex_pending_paths.contains(&buffered_key)); } +#[cfg(test)] +#[path = "tests/copied_prefix.rs"] +mod copied_prefix; +#[cfg(test)] +#[path = "tests/lineage_cache.rs"] +mod lineage_cache; #[cfg(test)] #[path = "tests/paginated.rs"] mod paginated; diff --git a/rust/src/cost_scanner/tests/copied_prefix.rs b/rust/src/cost_scanner/tests/copied_prefix.rs new file mode 100644 index 0000000000..23e1dccb22 --- /dev/null +++ b/rust/src/cost_scanner/tests/copied_prefix.rs @@ -0,0 +1,578 @@ +use super::*; + +fn write_copied_prefix_subagent_fixture( + sessions_root: &Path, + name: &str, + session_id: &str, + parent_id: &str, + base: DateTime, + owned: bool, +) -> PathBuf { + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let mut lines = vec![ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), + "payload": { + "id": session_id, "forked_from_id": parent_id, + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": parent_id}}} + } + }), + token_row(base, 2, [1_000, 900, 100], [0, 0, 0], "gpt-5.6-sol"), + serde_json::json!({ + "type": "turn_context", "ordinal": 10, "timestamp": base.to_rfc3339(), + "payload": {"model": "gpt-5.6-sol"} + }), + token_row( + base, + 12, + [1_000, 900, 100], + [1_000, 900, 100], + "gpt-5.6-sol", + ), + token_row( + base, + 13, + [5_000, 3_900, 500], + [5_000, 3_900, 500], + "gpt-5.6-sol", + ), + ]; + if owned { + lines.extend([ + token_row(base, 19, [5_050, 3_910, 505], [50, 10, 5], "gpt-5.6-sol"), + token_row( + base + Duration::seconds(1), + 20, + [5_070, 3_915, 510], + [20, 5, 5], + "gpt-5.6-sol", + ), + token_row( + base + Duration::seconds(2), + 21, + [5_070, 3_915, 510], + [20, 5, 5], + "gpt-5.6-sol", + ), + ]); + } + let body = lines + .into_iter() + .map(|line| line.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + +fn token_row( + timestamp: DateTime, + ordinal: i64, + total: [i64; 3], + last: [i64; 3], + model: &str, +) -> serde_json::Value { + serde_json::json!({ + "type": "event_msg", "ordinal": ordinal, "timestamp": timestamp.to_rfc3339(), + "payload": {"type": "token_count", "info": { + "model": model, + "total_token_usage": { + "input_tokens": total[0], "cached_input_tokens": total[1], "output_tokens": total[2] + }, + "last_token_usage": { + "input_tokens": last[0], "cached_input_tokens": last[1], "output_tokens": last[2] + } + }} + }) +} + +#[test] +fn copied_prefix_subagent_infers_advancing_baseline_without_parent() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "missing-parent", + Utc::now() - Duration::hours(1), + true, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 70); + assert_eq!(summary.cached_tokens, 15); + assert_eq!(summary.output_tokens, 10); + assert_eq!(summary.sessions_count, 1); + let usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(!usage.codex_unresolved_fork_parent); + assert!( + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) + ); + assert_eq!( + usage.days.values().next().unwrap()["gpt-5.6-sol"], + vec![70, 15, 10] + ); + + let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(cached.input_tokens, 70); + assert!(stats.codex_history_read_paths.is_empty()); +} + +#[test] +fn copied_prefix_subagent_inherited_only_suffix_is_not_billed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "missing-parent", + Utc::now() - Duration::hours(1), + false, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.output_tokens, 0); + assert_eq!(summary.sessions_count, 0); + let usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(usage.days.is_empty()); + assert!(!usage.codex_unresolved_fork_parent); + let state = usage.codex_fork_accounting_state.as_ref().unwrap(); + assert!(state.locally_resolved); + assert!(state.inherited_totals.is_none()); + + let (cached, stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(cached.input_tokens, 0); + assert_eq!(cached.output_tokens, 0); + assert_eq!(cached.sessions_count, 0); + assert!(stats.codex_history_read_paths.is_empty()); +} + +#[test] +fn copied_prefix_subagent_prefers_validated_parent_baseline() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + true, + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(&child) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(20)) + .unwrap(); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let parent_size = std::fs::metadata(&parent).unwrap().len(); + let child_size = std::fs::metadata(&child).unwrap().len(); + options.codex_max_session_file_bytes = + i64::try_from(parent_size.max(child_size)).expect("fixture size fits i64"); + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + let state = cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + + assert_eq!(state.inherited_totals.as_ref().unwrap().input, 1_000); + assert!(!state.locally_resolved); + assert_eq!(stats.files_seen, 2); + assert_eq!(stats.codex_read_receipt.metadata_reads, 2); + assert_eq!(stats.codex_read_receipt.history_reads, 2); + assert_eq!( + stats.codex_bytes_read, + parent_size.saturating_add(child_size), + "one bounded parse per candidate must enforce the per-file allowance" + ); + assert_eq!( + stats.codex_history_read_paths, + vec![ + parent.to_string_lossy().to_string(), + child.to_string_lossy().to_string(), + ] + ); +} + +#[test] +fn candidate_limit_counts_each_child_parent_candidate_once() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + true, + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(&child) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(20)) + .unwrap(); + std::fs::OpenOptions::new() + .write(true) + .open(&parent) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(10)) + .unwrap(); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + options.codex_candidate_limit = 1; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.files_seen, 1); + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 1); + assert_eq!( + stats.codex_metadata_read_paths, + vec![child.to_string_lossy().to_string()] + ); + assert_eq!( + cache.codex_pending_paths, + vec![parent.to_string_lossy().to_string()] + ); +} + +#[test] +fn cold_scan_orders_multi_level_parent_chain_before_children() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let ancestor = write_codex_fork_session_fixture( + &sessions, + "ancestor.jsonl", + "ancestor-id", + None, + base, + base, + &[1_000], + ); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + Some("ancestor-id"), + base + Duration::seconds(10), + base + Duration::seconds(10), + &[1_500], + ); + let child = write_codex_fork_session_fixture( + &sessions, + "child.jsonl", + "child-id", + Some("parent-id"), + base + Duration::seconds(20), + base + Duration::seconds(20), + &[2_000], + ); + let now = std::time::SystemTime::now(); + for (path, age) in [(&child, 30), (&parent, 20), (&ancestor, 10)] { + std::fs::OpenOptions::new() + .write(true) + .open(path) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(age)) + .unwrap(); + } + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.files_seen, 3); + assert_eq!(stats.codex_read_receipt.metadata_reads, 3); + assert_eq!(stats.codex_read_receipt.history_reads, 3); + assert_eq!( + stats.codex_history_read_paths, + vec![ + ancestor.to_string_lossy().to_string(), + parent.to_string_lossy().to_string(), + child.to_string_lossy().to_string(), + ] + ); + let parent_state = cache.files[&parent.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + let child_state = cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert_eq!(parent_state.inherited_totals.as_ref().unwrap().input, 1_000); + assert_eq!(child_state.inherited_totals.as_ref().unwrap().input, 1_500); +} + +fn assert_unsafe_lineage_is_unresolved( + summary: &CostSummary, + stats: &CostScanStats, + cache: &CostUsageCache, + paths: &[&Path], +) { + assert_eq!(summary.sessions_count, 0); + assert_eq!(summary.input_tokens, 0); + assert_eq!( + stats.codex_read_receipt.metadata_reads, + u32::try_from(paths.len()).expect("fixture count fits u32") + ); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert!(stats.codex_history_read_paths.is_empty()); + for path in paths { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.codex_fork_accounting_state.is_none()); + assert!(usage.days.is_empty()); + } +} + +#[test] +fn duplicate_parent_session_ids_fail_closed_with_their_child() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let first_parent = write_codex_fork_session_fixture( + &sessions, + "first-parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let second_parent = write_codex_fork_session_fixture( + &sessions, + "second-parent.jsonl", + "parent-id", + None, + base + Duration::seconds(1), + base + Duration::seconds(1), + &[2_000], + ); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(2), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_unsafe_lineage_is_unresolved( + &summary, + &stats, + &cache, + &[&first_parent, &second_parent, &child], + ); +} + +#[test] +fn two_node_subagent_cycle_fails_closed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let first = write_copied_prefix_subagent_fixture( + &sessions, + "first.jsonl", + "first-id", + "second-id", + base, + true, + ); + let second = write_copied_prefix_subagent_fixture( + &sessions, + "second.jsonl", + "second-id", + "first-id", + base + Duration::seconds(1), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = false; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_unsafe_lineage_is_unresolved(&summary, &stats, &cache, &[&first, &second]); +} + +#[test] +fn self_referential_subagent_fails_closed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let session = write_copied_prefix_subagent_fixture( + &sessions, + "self-cycle.jsonl", + "self-id", + "self-id", + Utc::now() - Duration::hours(1), + true, + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_unsafe_lineage_is_unresolved(&summary, &stats, &cache, &[&session]); +} + +fn assert_cached_inference_is_replaced_when_parent_appears( + prefer_newest_codex_sessions_first: bool, +) { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_copied_prefix_subagent_fixture( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + true, + ); + let mut options = CostScanOptions::app_driven(); + options.prefer_newest_codex_sessions_first = prefer_newest_codex_sessions_first; + let scanner = CostScanner::new(7) + .with_options(options) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + + let (_, _, inferred_cache) = scanner.scan_codex_detailed_with_cache(None); + let inferred_state = inferred_cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert!(inferred_state.locally_resolved); + assert_eq!( + inferred_state.inherited_totals.as_ref().unwrap().input, + 5_000 + ); + + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let now = std::time::SystemTime::now(); + std::fs::OpenOptions::new() + .write(true) + .open(parent) + .unwrap() + .set_modified(now - std::time::Duration::from_secs(10)) + .unwrap(); + + let (_, stats, validated_cache) = scanner.scan_codex_detailed_with_cache(None); + let validated_state = validated_cache.files[&child.to_string_lossy().to_string()] + .codex_fork_accounting_state + .as_ref() + .unwrap(); + assert!(!validated_state.locally_resolved); + assert_eq!( + validated_state.inherited_totals.as_ref().unwrap().input, + 1_000 + ); + assert!( + stats + .codex_history_read_paths + .contains(&child.to_string_lossy().to_string()), + "the unchanged child must be reparsed when baseline provenance changes" + ); +} + +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_parent_is_visited_first() { + assert_cached_inference_is_replaced_when_parent_appears(true); +} + +#[test] +fn copied_prefix_subagent_replaces_cached_inference_when_child_would_be_visited_first() { + assert_cached_inference_is_replaced_when_parent_appears(false); +} diff --git a/rust/src/cost_scanner/tests/lineage_cache.rs b/rust/src/cost_scanner/tests/lineage_cache.rs new file mode 100644 index 0000000000..124733b098 --- /dev/null +++ b/rust/src/cost_scanner/tests/lineage_cache.rs @@ -0,0 +1,535 @@ +use super::*; + +fn write_subagent( + sessions_root: &Path, + name: &str, + session_id: &str, + parent_id: &str, + timestamp: DateTime, +) -> PathBuf { + let day = timestamp.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let rows = [ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": timestamp.to_rfc3339(), + "payload": { + "id": session_id, + "forked_from_id": parent_id, + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": parent_id}}} + } + }), + lineage_token_row(timestamp, 2, 1_000, 0), + serde_json::json!({ + "type": "turn_context", "ordinal": 10, "timestamp": timestamp.to_rfc3339(), + "payload": {"model": "gpt-5.6-sol"} + }), + lineage_token_row(timestamp, 12, 1_000, 1_000), + lineage_token_row(timestamp + Duration::seconds(1), 20, 1_050, 50), + ]; + let body = rows + .into_iter() + .map(|row| row.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + +fn lineage_token_row( + timestamp: DateTime, + ordinal: i64, + total_input: i64, + last_input: i64, +) -> serde_json::Value { + serde_json::json!({ + "type": "event_msg", "ordinal": ordinal, "timestamp": timestamp.to_rfc3339(), + "payload": {"type": "token_count", "info": { + "model": "gpt-5.6-sol", + "total_token_usage": { + "input_tokens": total_input, "cached_input_tokens": 0, "output_tokens": 5 + }, + "last_token_usage": { + "input_tokens": last_input, "cached_input_tokens": 0, "output_tokens": 5 + } + }} + }) +} + +fn write_missing_ordinal_subagent( + sessions_root: &Path, + name: &str, + base: DateTime, + include_owned_usage: bool, +) -> PathBuf { + let day = base.with_timezone(&Local).date_naive(); + let day_dir = sessions_root + .join(day.format("%Y").to_string()) + .join(day.format("%m").to_string()) + .join(day.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let path = day_dir.join(name); + let mut missing_ordinal = lineage_token_row(base, 11, 100, 0); + missing_ordinal.as_object_mut().unwrap().remove("ordinal"); + let tail = if include_owned_usage { + lineage_token_row(base, 12, 120, 10) + } else { + lineage_token_row(base, 12, 100, 0) + }; + let rows = [ + serde_json::json!({ + "type": "session_meta", "ordinal": 0, "timestamp": base.to_rfc3339(), + "payload": { + "id": "child-id", + "forked_from_id": "absent-parent-id", + "subagent_history_start_ordinal": 10, + "thread_source": "subagent", + "source": {"subagent": {"thread_spawn": {"parent_thread_id": "absent-parent-id"}}} + } + }), + lineage_token_row(base, 9, 100, 0), + lineage_token_row(base, 10, 100, 0), + missing_ordinal, + tail, + ]; + let body = rows + .into_iter() + .map(|row| row.to_string()) + .collect::>() + .join("\n") + + "\n"; + std::fs::write(&path, body).unwrap(); + path +} + +fn bounded_scanner(sessions: &Path, cache_root: &Path) -> CostScanner { + let mut options = CostScanOptions::app_driven(); + options.codex_candidate_limit = 1; + options.prefer_newest_codex_sessions_first = false; + CostScanner::new(7) + .with_options(options) + .with_cache_root(cache_root) + .with_sessions_dirs(vec![sessions.to_path_buf()]) +} + +fn assert_locally_inferred(cache: &CostUsageCache, path: &Path) { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!(!usage.codex_unresolved_fork_parent); + assert!( + usage + .codex_fork_accounting_state + .as_ref() + .is_some_and(|state| state.locally_resolved) + ); +} + +fn assert_unresolved(cache: &CostUsageCache, path: &Path) { + let usage = &cache.files[&path.to_string_lossy().to_string()]; + assert!(usage.codex_unresolved_fork_parent); + assert!(usage.days.is_empty()); + assert!(usage.codex_fork_accounting_state.is_none()); +} + +#[test] +fn appended_owned_token_row_reinfers_locally_resolved_subagent_from_start() { + use std::io::Write as _; + + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_subagent(&sessions, "child.jsonl", "child-id", "missing-parent", base); + let scanner = bounded_scanner(&sessions, &cache_root); + + let (initial, _, initial_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(initial.input_tokens, 50); + assert_locally_inferred(&initial_cache, &child); + + let appended_owned_row = lineage_token_row(base + Duration::seconds(2), 22, 1_100, 50); + std::fs::OpenOptions::new() + .append(true) + .open(&child) + .unwrap() + .write_all(format!("{appended_owned_row}\n").as_bytes()) + .unwrap(); + + let (grown, _, grown_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(grown.input_tokens, 100); + assert_locally_inferred(&grown_cache, &child); +} + +#[test] +fn replaced_parent_with_same_path_size_and_mtime_cannot_author_lineage() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[1_000], + ); + let child = write_subagent( + &sessions, + "child.jsonl", + "child-id", + "parent-id", + base + Duration::seconds(10), + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (_, _, cache) = scanner.scan_codex_detailed_with_cache(None); + let parent_key = parent.to_string_lossy().to_string(); + let child_usage = &cache.files[&child.to_string_lossy().to_string()]; + assert!(matches!( + CodexLineagePlanner::new(&cache).decision_for_usage(&cache, child_usage), + CodexLineageDecision::ParentReady(_) + )); + + let old_identity = cache.files[&parent_key] + .codex_file_identity + .clone() + .expect("parent identity persisted"); + let old_metadata = std::fs::metadata(&parent).unwrap(); + let old_mtime = old_metadata.modified().unwrap(); + let old_size = old_metadata.len(); + let rotated = parent.with_extension("old"); + std::fs::rename(&parent, &rotated).unwrap(); + let replacement = write_codex_fork_session_fixture( + &sessions, + "parent.jsonl", + "parent-id", + None, + base, + base, + &[2_000], + ); + std::fs::OpenOptions::new() + .write(true) + .open(&replacement) + .unwrap() + .set_modified(old_mtime) + .unwrap(); + let replacement_metadata = std::fs::metadata(&replacement).unwrap(); + assert_eq!(replacement_metadata.len(), old_size); + let replacement_identity = + JsonlScanner::codex_file_identity(&replacement, &replacement_metadata) + .expect("replacement identity available"); + assert_ne!(replacement_identity, old_identity); + + assert_eq!( + CodexLineagePlanner::new(&cache).decision_for_usage(&cache, child_usage), + CodexLineageDecision::Unsafe + ); +} + +#[test] +fn missing_explicit_ordinal_keeps_subagent_cache_unresolved() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_missing_ordinal_subagent(&sessions, "child.jsonl", base, true); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &child); +} + +#[test] +fn missing_ordinal_cannot_complete_zero_usage_subagent_cache() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_missing_ordinal_subagent(&sessions, "child.jsonl", base, false); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &child); +} + +#[test] +fn missing_ordinal_after_local_resolution_keeps_subagent_cache_unresolved() { + use std::io::Write as _; + + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_subagent( + &sessions, + "child.jsonl", + "child-id", + "missing-parent-id", + base, + ); + let mut missing_ordinal = lineage_token_row(base + Duration::seconds(2), 21, 1_060, 10); + missing_ordinal.as_object_mut().unwrap().remove("ordinal"); + std::fs::OpenOptions::new() + .append(true) + .open(&child) + .unwrap() + .write_all(format!("{missing_ordinal}\n").as_bytes()) + .unwrap(); + let scanner = bounded_scanner(&sessions, &cache_root); + + let (summary, _, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &child); +} + +#[test] +fn legacy_cache_without_file_identity_is_reparsed() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let session = write_codex_fork_session_fixture( + &sessions, + "session.jsonl", + "root-session-id", + None, + base, + base, + &[1_000], + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (_, _, _) = scanner.scan_codex_detailed_with_cache(None); + let session_key = session.to_string_lossy().to_string(); + let mut legacy_cache = JsonlScanner::load_cache(ProviderId::Codex, Some(&cache_root)); + legacy_cache + .files + .get_mut(&session_key) + .unwrap() + .codex_file_identity = None; + JsonlScanner::save_cache(ProviderId::Codex, &mut legacy_cache, Some(&cache_root)); + + let (_, stats, refreshed_cache) = scanner.scan_codex_detailed_with_cache(None); + + assert!(stats.codex_history_read_paths.contains(&session_key)); + assert!( + refreshed_cache.files[&session_key] + .codex_file_identity + .is_some() + ); +} + +#[test] +fn bounded_refresh_detects_duplicate_parent_owners_across_cache_and_candidate() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let child = write_subagent(&sessions, "child.jsonl", "child-id", "parent-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + + let (_, first_stats, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(first_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(first_stats.codex_read_receipt.history_reads, 1); + assert_locally_inferred(&first_cache, &child); + + let first_parent = write_codex_fork_session_fixture( + &sessions, + "parent-a.jsonl", + "parent-id", + None, + base - Duration::seconds(2), + base - Duration::seconds(2), + &[1_000], + ); + let (_, second_stats, second_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(second_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(second_stats.codex_read_receipt.history_reads, 1); + assert_locally_inferred(&second_cache, &child); + + let second_parent = write_codex_fork_session_fixture( + &sessions, + "parent-b.jsonl", + "parent-id", + None, + base - Duration::seconds(1), + base - Duration::seconds(1), + &[2_000], + ); + let (summary, third_stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(third_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(third_stats.codex_read_receipt.history_reads, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &first_parent); + assert_unresolved(&cache, &second_parent); + assert_unresolved(&cache, &child); +} + +#[test] +fn bounded_refresh_detects_equal_timestamp_two_node_cycle() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let first = write_subagent(&sessions, "first.jsonl", "first-id", "second-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + let (_, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_locally_inferred(&first_cache, &first); + + let second = write_subagent(&sessions, "second.jsonl", "second-id", "first-id", base); + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &first); + assert_unresolved(&cache, &second); +} + +#[test] +fn bounded_refresh_rejects_self_cycle_migration() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let session = write_subagent(&sessions, "self.jsonl", "self-id", "missing-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + let (_, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_locally_inferred(&first_cache, &session); + + write_subagent( + &sessions, + "self.jsonl", + "self-id", + "self-id", + base + Duration::seconds(1), + ); + let (summary, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&cache, &session); +} + +#[test] +fn bounded_refresh_rejects_dependent_of_locally_inferred_parent() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let base = Utc::now() - Duration::hours(1); + let parent = write_subagent(&sessions, "parent.jsonl", "parent-id", "missing-id", base); + let scanner = bounded_scanner(&sessions, &cache_root); + let (_, _, first_cache) = scanner.scan_codex_detailed_with_cache(None); + assert_locally_inferred(&first_cache, &parent); + + let dependent = write_subagent( + &sessions, + "dependent.jsonl", + "dependent-id", + "parent-id", + base + Duration::seconds(1), + ); + let (_, stats, cache) = scanner.scan_codex_detailed_with_cache(None); + + assert_eq!(stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(stats.codex_read_receipt.history_reads, 0); + assert_locally_inferred(&cache, &parent); + assert_unresolved(&cache, &dependent); +} + +#[test] +fn current_refresh_scopes_unsafe_cache_invalidation_to_range_and_dependencies() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let active_time = Utc::now() - Duration::hours(1); + let old_date = Local::now().date_naive() - Duration::days(30); + let old_day = old_date.format("%Y-%m-%d").to_string(); + let old_dir = sessions + .join(old_date.format("%Y").to_string()) + .join(old_date.format("%m").to_string()) + .join(old_date.format("%d").to_string()); + let mut cache = CostUsageCache::default(); + + { + let mut add_cached = |name: &str, session_id: &str, parent_id: Option<&str>| { + let path = old_dir.join(name).to_string_lossy().to_string(); + let mut usage = cached_usage_with_packed(&old_day, "gpt-5.6-sol", vec![100, 0, 5, 0]); + usage.codex_session_id = Some(session_id.to_string()); + usage.codex_forked_from_id = parent_id.map(str::to_string); + cache.files.insert(path, usage); + }; + add_cached("unrelated-a.jsonl", "unrelated-a", Some("unrelated-b")); + add_cached("unrelated-b.jsonl", "unrelated-b", Some("unrelated-a")); + add_cached("required-a.jsonl", "required-parent", None); + add_cached("required-b.jsonl", "required-parent", None); + } + JsonlScanner::save_cache(ProviderId::Codex, &mut cache, Some(&cache_root)); + + let active_child = write_codex_fork_session_fixture( + &sessions, + "active-child.jsonl", + "active-child", + Some("required-parent"), + active_time, + active_time + Duration::seconds(1), + &[1_000_000, 1_000_140], + ); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + + let (summary, _, refreshed) = scanner.scan_codex_detailed_with_cache(None); + let cached_path = |name: &str| old_dir.join(name).to_string_lossy().to_string(); + + for path in [ + cached_path("unrelated-a.jsonl"), + cached_path("unrelated-b.jsonl"), + ] { + let usage = refreshed + .files + .get(&path) + .expect("unrelated history retained"); + assert_eq!(usage.days[&old_day]["gpt-5.6-sol"], vec![100, 0, 5, 0]); + assert!(!usage.codex_unresolved_fork_parent); + } + for path in [ + cached_path("required-a.jsonl"), + cached_path("required-b.jsonl"), + ] { + assert_unresolved(&refreshed, Path::new(&path)); + } + assert_eq!(summary.input_tokens, 0); + assert_eq!(summary.sessions_count, 0); + assert_unresolved(&refreshed, &active_child); +} diff --git a/rust/src/locale.rs b/rust/src/locale.rs index 27f536ae93..6b43e88960 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -659,6 +659,7 @@ locale_keys! { UsageSpendHourlyActivity, UsageSpendRequests, UsageSpendTokens, + UsageSpendKnownSubtotal, UsageSpendAllTimeHistory, UsageSpendCustomPricing, OverviewSpendTitle, @@ -806,6 +807,9 @@ locale_keys! { TrayIconModeHelper, TrayIconModeSingle, TrayIconModePerProvider, + TrayIconModeStacked, + StackedTrayTopProvider, + StackedTrayBottomProvider, ShowProviderIcons, ShowProviderIconsHelper, PreferHighestUsage, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index a0865280aa..1a99574daf 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -353,6 +353,7 @@ UsageSpendDefaultPricing = Default pricing UsageSpendHourlyActivity = Hourly activity UsageSpendRequests = requests UsageSpendTokens = tokens +UsageSpendKnownSubtotal = ≥{ "{}" } known UsageSpendAllTimeHistory = All-time is backed by the latest UsageSpendCustomPricing = custom pricing OverviewSpendTitle = Usage & Spend · 30d @@ -477,6 +478,9 @@ TrayIconModeLabel = Tray icon mode TrayIconModeHelper = Single unified icon or one icon per enabled provider. TrayIconModeSingle = Single TrayIconModePerProvider = Per provider +TrayIconModeStacked = Stacked providers +StackedTrayTopProvider = Top provider +StackedTrayBottomProvider = Bottom provider ShowProviderIcons = Show provider icons ShowProviderIconsHelper = Display provider icons in the tray switcher. PreferHighestUsage = Prefer highest usage diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index 8195889ad1..5ba578c95f 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -433,6 +433,9 @@ TrayIconModeLabel = Modo de ícono de bandeja TrayIconModeHelper = Ícono único combinado o un ícono por cada proveedor habilitado. TrayIconModeSingle = Único TrayIconModePerProvider = Por proveedor +TrayIconModeStacked = Proveedores apilados +StackedTrayTopProvider = Proveedor superior +StackedTrayBottomProvider = Proveedor inferior ShowProviderIcons = Mostrar íconos de proveedores ShowProviderIconsHelper = Mostrar íconos de proveedores en el selector de bandeja. PreferHighestUsage = Preferir uso más alto diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index 70e02dec96..c91cef6370 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -415,6 +415,9 @@ TrayIconModeLabel = トレイアイコンモード TrayIconModeHelper = 単一の統合アイコンか、有効なプロバイダーごとのアイコン。 TrayIconModeSingle = 単一 TrayIconModePerProvider = プロバイダー別 +TrayIconModeStacked = プロバイダーを積み重ねる +StackedTrayTopProvider = 上のプロバイダー +StackedTrayBottomProvider = 下のプロバイダー ShowProviderIcons = プロバイダーアイコンを表示 ShowProviderIconsHelper = トレイスイッチャーにプロバイダーアイコンを表示。 PreferHighestUsage = 最も使用量が多いものを優先 diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index 394bf2d3bd..cf67c452cf 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -420,6 +420,9 @@ TrayIconModeLabel = 트레이 아이콘 모드 TrayIconModeHelper = 단일 통합 아이콘 또는 활성화된 제공업체당 하나의 아이콘. TrayIconModeSingle = 단일 TrayIconModePerProvider = 제공업체별 +TrayIconModeStacked = 제공업체 쌓기 +StackedTrayTopProvider = 위쪽 제공업체 +StackedTrayBottomProvider = 아래쪽 제공업체 ShowProviderIcons = 제공업체 아이콘 표시 ShowProviderIconsHelper = 트레이 메뉴에 제공업체 아이콘을 표시합니다. PreferHighestUsage = 가장 높은 사용량 우선 diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index 31e8c9b0b3..e0129c89de 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -399,6 +399,9 @@ TrayIconModeLabel = Режим значков в трее TrayIconModeHelper = Один унифицированный значок или один значок для каждого включенного провайдера. TrayIconModeSingle = Одинокий TrayIconModePerProvider = За провайдера +TrayIconModeStacked = Провайдеры стопкой +StackedTrayTopProvider = Верхний провайдер +StackedTrayBottomProvider = Нижний провайдер ShowProviderIcons = Показать значки провайдеров ShowProviderIconsHelper = Отображать значки провайдеров в переключателе трея. PreferHighestUsage = Предпочитаю максимальное использование diff --git a/rust/src/locale/tests.rs b/rust/src/locale/tests.rs index 9daf93d809..57becd3865 100644 --- a/rust/src/locale/tests.rs +++ b/rust/src/locale/tests.rs @@ -445,6 +445,10 @@ fn test_fluent_preserves_literal_placeholders_and_status_spacing() { get_text(Language::English, LocaleKey::RemainingAmount), "{:.2} remaining" ); + assert_eq!( + get_text(Language::English, LocaleKey::UsageSpendKnownSubtotal), + "≥{} known" + ); } fn resource_key_names(resource: &str) -> HashSet<&str> { diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index 94fa3750d5..2abf387409 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -436,6 +436,9 @@ TrayIconModeLabel = Tepsi simgesi modu TrayIconModeHelper = Tek birleşik simge veya etkin sağlayıcı başına bir simge. TrayIconModeSingle = Tek TrayIconModePerProvider = Sağlayıcı başına +TrayIconModeStacked = Yığılmış sağlayıcılar +StackedTrayTopProvider = Üst sağlayıcı +StackedTrayBottomProvider = Alt sağlayıcı ShowProviderIcons = Sağlayıcı simgelerini göster ShowProviderIconsHelper = Tepsi değiştiricisinde sağlayıcı simgelerini göster. PreferHighestUsage = En yüksek kullanımı tercih et diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index 71c6a6b2ea..60e9e7a74d 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -414,6 +414,9 @@ TrayIconModeLabel = 托盘图标模式 TrayIconModeHelper = 使用单一合并图标，或为每个已启用服务商显示独立图标。 TrayIconModeSingle = 合并 TrayIconModePerProvider = 按服务商 +TrayIconModeStacked = 堆叠服务商 +StackedTrayTopProvider = 上方服务商 +StackedTrayBottomProvider = 下方服务商 ShowProviderIcons = 显示服务商图标 ShowProviderIconsHelper = 在托盘切换器中显示服务商图标。 PreferHighestUsage = 优先显示最高用量 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index 9c04c3a347..5977e37435 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -414,6 +414,9 @@ TrayIconModeLabel = 系統匣圖示模式 TrayIconModeHelper = 使用單一合併圖示，或為每個已啟用提供者顯示獨立圖示。 TrayIconModeSingle = 合併 TrayIconModePerProvider = 按提供者 +TrayIconModeStacked = 堆疊提供者 +StackedTrayTopProvider = 上方提供者 +StackedTrayBottomProvider = 下方提供者 ShowProviderIcons = 顯示提供者圖示 ShowProviderIconsHelper = 在系統匣切換器中顯示提供者圖示。 PreferHighestUsage = 優先顯示最高用量 diff --git a/rust/src/providers/antigravity/cost.rs b/rust/src/providers/antigravity/cost.rs new file mode 100644 index 0000000000..60ac3bfe8d --- /dev/null +++ b/rust/src/providers/antigravity/cost.rs @@ -0,0 +1,88 @@ +use crate::core::CostUsagePricing; + +/// Antigravity records routing variants of a vendor model that bill at the base model's public +/// price. The alias stays provider-local so shared Claude pricing keeps unknown variants unpriced. +const ROUTING_VARIANT_SUFFIXES: [&str; 3] = ["-tiered", "-low", "-thinking"]; + +pub(super) fn estimate_cost_usd( + model: Option<&str>, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> Option { + let model = model.map(str::trim).filter(|value| !value.is_empty())?; + let input = i32::try_from(input).ok()?; + let cache_read = i32::try_from(cache_read).ok()?; + let cache_write = i32::try_from(cache_write).ok()?; + let output = i32::try_from(output).ok()?; + let resolve = |candidate: &str| { + CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) + .filter(|cost| cost.is_finite() && *cost >= 0.0) + }; + resolve(model).or_else(|| pricing_base_model(model).and_then(resolve)) +} + +/// Matches the routing suffix case-insensitively and keeps the recorded spelling of the base. +fn pricing_base_model(model: &str) -> Option<&str> { + let lowered = model.to_ascii_lowercase(); + let suffix = ROUTING_VARIANT_SUFFIXES + .iter() + .find(|suffix| lowered.ends_with(*suffix))?; + model + .get(..model.len() - suffix.len()) + .filter(|base| !base.is_empty()) +} + +#[cfg(test)] +mod tests { + use super::{estimate_cost_usd, pricing_base_model}; + + #[test] + fn prices_known_models_and_provider_local_routing_variants() { + let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) + .expect("known public price"); + let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) + .expect("routing suffix uses the base public price"); + assert!(direct > 0.0); + assert_eq!(direct, routed); + } + + #[test] + fn routing_suffixes_match_case_insensitively() { + let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) + .expect("known public price"); + for model in [ + "claude-sonnet-4-6-Thinking", + "claude-sonnet-4-6-LOW", + "claude-sonnet-4-6-Tiered", + ] { + assert_eq!( + estimate_cost_usd(Some(model), 1_000, 200, 100, 500), + Some(direct), + "{model}" + ); + } + assert_eq!( + pricing_base_model("Claude-Sonnet-4-6-THINKING"), + Some("Claude-Sonnet-4-6") + ); + assert_eq!(pricing_base_model("-thinking"), None); + assert_eq!(pricing_base_model("claude-sonnet-4-6"), None); + } + + #[test] + fn unknown_or_oversized_pricing_inputs_fail_closed() { + assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); + assert_eq!( + estimate_cost_usd(Some("unknown-thinking"), 1, 2, 3, 4), + None + ); + assert_eq!(estimate_cost_usd(Some(" "), 1, 2, 3, 4), None); + assert_eq!(estimate_cost_usd(None, 1, 2, 3, 4), None); + assert_eq!( + estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), + None + ); + } +} diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs new file mode 100644 index 0000000000..00519d93ca --- /dev/null +++ b/rust/src/providers/antigravity/local_history.rs @@ -0,0 +1,219 @@ +use super::{local_sessions_reader as local_sessions, local_sqlite}; +use std::fs; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; + +use crate::spend_contract::LocalTokenHistorySummary; + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) +} + +fn configured_database_roots(home: &Path) -> [PathBuf; 3] { + let gemini_cli_home = std::env::var("GEMINI_CLI_HOME").ok(); + configured_database_roots_from_values(home, gemini_cli_home.as_deref()) +} + +fn configured_database_roots_from_values( + home: &Path, + gemini_cli_home: Option<&str>, +) -> [PathBuf; 3] { + let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); + local_sqlite::database_roots(&gemini_base) +} + +fn summarize_local_usage_from( + roots: &[PathBuf], + now: DateTime, + days: u32, + jsonl_fallback: impl FnOnce() -> LocalTokenHistorySummary, +) -> LocalTokenHistorySummary { + match local_sqlite::summarize(roots, now, days) { + local_sqlite::SQLiteScan::Summary(summary) => summary, + local_sqlite::SQLiteScan::NoDatabases | local_sqlite::SQLiteScan::Unsupported => { + jsonl_fallback() + } + } +} + +pub fn summarize_local_usage(days: u32) -> LocalTokenHistorySummary { + let now = Utc::now(); + let Some(home) = dirs::home_dir() else { + return LocalTokenHistorySummary::default(); + }; + let roots = configured_database_roots(&home); + let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); + summarize_local_usage_from(&roots, now, days, || { + local_sessions::summarize_jsonl_at(&tokscale_sessions, now, days) + }) +} + +/// Count local Antigravity conversation artifacts for the quota provider's +/// offline fallback. Mirrors upstream #3119 without opening SQLite files. +pub fn offline_conversation_count() -> usize { + let Some(home) = dirs::home_dir() else { + return 0; + }; + let roots = configured_database_roots(&home); + let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); + offline_conversation_count_with_roots(&roots, &tokscale_sessions) +} + +#[cfg(test)] +fn offline_conversation_count_in(home: &Path) -> usize { + let roots = local_sqlite::database_roots(&home.join(".gemini")); + let tokscale_sessions = local_sessions::tokscale_sessions_from_values(home, None); + offline_conversation_count_with_roots(&roots, &tokscale_sessions) +} + +fn offline_conversation_count_with_roots( + database_roots: &[PathBuf], + tokscale_sessions: &Path, +) -> usize { + let db_count = database_roots + .iter() + .map(|root| count_extension(root, "db")) + .sum::(); + if db_count > 0 { + return db_count; + } + local_sessions::count_jsonl_sessions_at(tokscale_sessions) +} + +fn count_extension(root: &Path, extension: &str) -> usize { + fs::read_dir(root) + .ok() + .into_iter() + .flatten() + .flatten() + .map(|entry| entry.path()) + .filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension)) + .count() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spend_contract::LocalHistoryCoverage; + use chrono::TimeZone; + use rusqlite::Connection; + + #[test] + fn foreign_database_preserves_valid_tokscale_history() { + let dir = tempfile::tempdir().unwrap(); + let gemini_base = dir.path().join(".gemini"); + let database_root = gemini_base.join("antigravity-cli").join("conversations"); + fs::create_dir_all(&database_root).unwrap(); + let connection = Connection::open(database_root.join("foreign.db")).unwrap(); + connection + .execute( + "CREATE TABLE unrelated(id INTEGER PRIMARY KEY, value TEXT)", + [], + ) + .unwrap(); + + let tokscale_sessions = dir + .path() + .join(".config/tokscale/antigravity-cache/sessions"); + fs::create_dir_all(&tokscale_sessions).unwrap(); + let session_path = tokscale_sessions.join("session-a.jsonl"); + fs::write( + &session_path, + b"{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", + ) + .unwrap(); + + let roots = local_sqlite::database_roots(&gemini_base); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_local_usage_from(&roots, now, 7, || { + local_sessions::summarize_jsonl_paths( + std::slice::from_ref(&session_path), + now, + 7, + false, + ) + }); + + assert_eq!(summary.total_tokens, 120); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } + + #[test] + fn foreign_only_input_does_not_fabricate_known_zero_native_usage() { + let dir = tempfile::tempdir().unwrap(); + let gemini_base = dir.path().join(".gemini"); + let database_root = gemini_base.join("antigravity-cli").join("conversations"); + fs::create_dir_all(&database_root).unwrap(); + let connection = Connection::open(database_root.join("foreign.db")).unwrap(); + connection + .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) + .unwrap(); + + let roots = local_sqlite::database_roots(&gemini_base); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_local_usage_from(&roots, now, 7, LocalTokenHistorySummary::default); + + assert_eq!(summary, LocalTokenHistorySummary::default()); + assert_eq!(summary.coverage, LocalHistoryCoverage::Unavailable); + } + + #[test] + fn scan_context_honors_non_empty_root_overrides() { + let home = Path::new(r"C:\Users\test"); + let tokscale_sessions = + local_sessions::tokscale_sessions_from_values(home, Some(r"E:\tokscale-root")); + let roots = configured_database_roots_from_values(home, Some(r"D:\gemini-root")); + assert_eq!( + roots[0], + PathBuf::from(r"D:\gemini-root") + .join("antigravity-cli") + .join("conversations") + ); + assert_eq!( + tokscale_sessions, + PathBuf::from(r"E:\tokscale-root") + .join("antigravity-cache") + .join("sessions") + ); + let defaults = local_sessions::tokscale_sessions_from_values(home, Some("")); + let default_roots = configured_database_roots_from_values(home, Some(" ")); + assert_eq!(default_roots[1], home.join(".gemini").join("antigravity")); + assert_eq!( + defaults, + home.join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions") + ); + } + + #[test] + fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { + let dir = tempfile::tempdir().unwrap(); + let app = dir + .path() + .join(".gemini") + .join("antigravity") + .join("conversations"); + fs::create_dir_all(&app).unwrap(); + fs::write(app.join("a.db"), b"").unwrap(); + fs::write(app.join("a.db-wal"), b"").unwrap(); + assert_eq!(offline_conversation_count_in(dir.path()), 1); + + fs::remove_file(app.join("a.db")).unwrap(); + let cache = dir + .path() + .join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions"); + fs::create_dir_all(&cache).unwrap(); + fs::write(cache.join("one.jsonl"), b"{}\n").unwrap(); + assert_eq!(offline_conversation_count_in(dir.path()), 1); + } +} diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 458bb97f69..4efb9af063 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -1,482 +1,4 @@ -use std::collections::HashSet; -use std::fs::{self, File}; -use std::io::{BufRead, BufReader}; -use std::path::{Path, PathBuf}; - -use chrono::{DateTime, Duration, Local, TimeZone, Utc}; -use serde_json::Value; - -const MAX_SESSION_FILES: usize = 2048; -const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; -const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; -const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; - -pub use crate::spend_contract::LocalHistoryCoverage; -pub type LocalSessionSummary = crate::spend_contract::LocalTokenHistorySummary; - -#[derive(Debug, Clone, PartialEq, Eq)] -struct ScanContext { - database_roots: [PathBuf; 3], - tokscale_sessions: PathBuf, -} - -impl ScanContext { - fn from_values( - home: &Path, - gemini_cli_home: Option<&str>, - tokscale_config_dir: Option<&str>, - ) -> Self { - let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); - let tokscale_base = clean_env_path(tokscale_config_dir) - .unwrap_or_else(|| home.join(".config").join("tokscale")); - Self { - database_roots: super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions: tokscale_base.join("antigravity-cache").join("sessions"), - } - } - - fn capture() -> Option { - let home = dirs::home_dir()?; - let gemini = std::env::var("GEMINI_CLI_HOME").ok(); - let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); - Some(Self::from_values( - &home, - gemini.as_deref(), - tokscale.as_deref(), - )) - } -} - -fn clean_env_path(value: Option<&str>) -> Option { - value - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) -} - -pub fn summarize(days: u32) -> LocalSessionSummary { - let now = Utc::now(); - let Some(context) = ScanContext::capture() else { - return LocalSessionSummary::default(); - }; - summarize_context(&context, now, days) -} - -fn summarize_context(context: &ScanContext, now: DateTime, days: u32) -> LocalSessionSummary { - match super::local_sqlite::summarize(&context.database_roots, now, days) { - super::local_sqlite::SQLiteScan::Summary(summary) => summary, - super::local_sqlite::SQLiteScan::NoDatabases - | super::local_sqlite::SQLiteScan::Unsupported => { - let (paths, truncated) = tokscale_paths(&context.tokscale_sessions); - if paths.is_empty() { - LocalSessionSummary::default() - } else { - summarize_paths(&paths, now, days, truncated) - } - } - } -} - -/// Count local Antigravity conversation artifacts for the quota provider's -/// offline fallback. Mirrors upstream #3119 without opening SQLite files. -pub fn offline_conversation_count() -> usize { - let Some(context) = ScanContext::capture() else { - return 0; - }; - offline_conversation_count_context(&context) -} - -fn offline_conversation_count_in(home: &Path) -> usize { - offline_conversation_count_context(&ScanContext::from_values(home, None, None)) -} - -fn offline_conversation_count_context(context: &ScanContext) -> usize { - let db_count = context - .database_roots - .iter() - .map(|root| count_extension(root, "db")) - .sum::(); - if db_count > 0 { - return db_count; - } - tokscale_paths(&context.tokscale_sessions).0.len() -} - -fn count_extension(root: &Path, extension: &str) -> usize { - fs::read_dir(root) - .ok() - .into_iter() - .flatten() - .flatten() - .map(|entry| entry.path()) - .filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension)) - .count() -} - -fn tokscale_paths(base: &Path) -> (Vec, bool) { - let Ok(entries) = fs::read_dir(base) else { - return (Vec::new(), false); - }; - let mut paths: Vec<_> = entries - .flatten() - .map(|entry| entry.path()) - .filter(|path| { - path.extension() - .and_then(|value| value.to_str()) - .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")) - }) - .collect(); - paths.sort(); - let truncated = paths.len() > MAX_SESSION_FILES; - if truncated { - paths.drain(..paths.len() - MAX_SESSION_FILES); - } - (paths, truncated) -} - -fn summarize_paths( - paths: &[PathBuf], - now: DateTime, - days: u32, - truncated: bool, -) -> LocalSessionSummary { - let first_day = now.with_timezone(&Local).date_naive() - - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); - let mut total_tokens = 0_u64; - let mut sessions_with_usage = HashSet::new(); - let mut seen_response_ids = HashSet::new(); - let mut complete = !truncated; - - for path in paths.iter().take(MAX_SESSION_FILES) { - let file = match File::open(path) { - Ok(file) => file, - Err(_) => { - complete = false; - continue; - } - }; - match file.metadata() { - Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, - Ok(_) => {} - Err(_) => complete = false, - } - let mut reader = BufReader::new(file); - let mut remaining = MAX_SESSION_FILE_BYTES; - let mut path_had_usage = false; - loop { - let line = match read_bounded_jsonl_line(&mut reader, &mut remaining) { - Ok(Some(line)) => line, - Ok(None) => break, - Err(_) => { - complete = false; - break; - } - }; - if line.is_empty() { - continue; - } - let Ok(value) = serde_json::from_slice::(&line) else { - continue; - }; - let kind = value.get("type").and_then(Value::as_str); - if kind != Some("usage") && value.get("input").is_none() { - continue; - } - if let Some(response_id) = value - .get("responseId") - .or_else(|| value.get("response_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - && !seen_response_ids.insert(response_id.to_string()) - { - continue; - } - - let timestamp_ms = value - .get("timestamp") - .and_then(Value::as_i64) - .unwrap_or_default(); - let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { - continue; - }; - if at > now || at.with_timezone(&Local).date_naive() < first_day { - continue; - } - - let input = token_field(&value, &["input"]); - let output = token_field(&value, &["output"]); - let cache_read = token_field(&value, &["cacheRead", "cache_read"]); - let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); - let total = input - .saturating_add(output) - .saturating_add(cache_read) - .saturating_add(cache_write); - if total == 0 { - continue; - } - total_tokens = total_tokens.saturating_add(total); - path_had_usage = true; - } - if path_had_usage { - sessions_with_usage.insert(path.clone()); - } - } - - LocalSessionSummary { - total_tokens, - session_count: sessions_with_usage.len(), - coverage: if paths.is_empty() { - LocalHistoryCoverage::Unavailable - } else if complete { - LocalHistoryCoverage::Complete - } else { - LocalHistoryCoverage::Partial - }, - } -} - -fn read_bounded_jsonl_line( - reader: &mut R, - remaining_file_bytes: &mut usize, -) -> std::io::Result>> { - if *remaining_file_bytes == 0 { - return Ok(None); - } - let mut line = Vec::new(); - let mut saw_input = false; - let mut discarding = false; - - loop { - let chunk = reader.fill_buf()?; - if chunk.is_empty() { - return Ok(saw_input.then_some(if discarding { Vec::new() } else { line })); - } - let bounded_len = chunk.len().min(*remaining_file_bytes); - if bounded_len == 0 { - return Ok(None); - } - let bounded = &chunk[..bounded_len]; - let newline = bounded.iter().position(|byte| *byte == b'\n'); - let segment_end = newline.unwrap_or(bounded.len()); - let segment = &bounded[..segment_end]; - saw_input = saw_input || !segment.is_empty() || newline.is_some(); - if !discarding { - if line.len().saturating_add(segment.len()) <= MAX_JSONL_LINE_BYTES { - line.extend_from_slice(segment); - } else { - line.clear(); - discarding = true; - } - } - let consumed = segment_end + usize::from(newline.is_some()); - reader.consume(consumed); - *remaining_file_bytes = remaining_file_bytes.saturating_sub(consumed); - if newline.is_some() { - return Ok(Some(if discarding { Vec::new() } else { line })); - } - if *remaining_file_bytes == 0 { - return Ok(Some(Vec::new())); - } - } -} - -fn token_field(value: &Value, keys: &[&str]) -> u64 { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_u64)) - .unwrap_or(0) -} - -#[cfg(test)] -mod tests { - use super::*; - use rusqlite::Connection; - - #[test] - fn scan_context_honors_non_empty_root_overrides() { - let home = Path::new(r"C:\Users\test"); - let context = - ScanContext::from_values(home, Some(r"D:\gemini-root"), Some(r"E:\tokscale-root")); - assert_eq!( - context.database_roots[0], - PathBuf::from(r"D:\gemini-root") - .join("antigravity-cli") - .join("conversations") - ); - assert_eq!( - context.tokscale_sessions, - PathBuf::from(r"E:\tokscale-root") - .join("antigravity-cache") - .join("sessions") - ); - let defaults = ScanContext::from_values(home, Some(" "), Some("")); - assert_eq!( - defaults.database_roots[1], - home.join(".gemini").join("antigravity") - ); - assert_eq!( - defaults.tokscale_sessions, - home.join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions") - ); - } - - #[test] - fn foreign_database_preserves_valid_tokscale_history() { - let dir = tempfile::tempdir().unwrap(); - let gemini_base = dir.path().join(".gemini"); - let database_root = gemini_base.join("antigravity-cli").join("conversations"); - fs::create_dir_all(&database_root).unwrap(); - let connection = Connection::open(database_root.join("foreign.db")).unwrap(); - connection - .execute( - "CREATE TABLE unrelated(id INTEGER PRIMARY KEY, value TEXT)", - [], - ) - .unwrap(); - - let tokscale_sessions = dir - .path() - .join(".config/tokscale/antigravity-cache/sessions"); - fs::create_dir_all(&tokscale_sessions).unwrap(); - fs::write( - tokscale_sessions.join("session-a.jsonl"), - b"{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", - ) - .unwrap(); - - let context = ScanContext { - database_roots: super::super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions, - }; - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_context(&context, now, 7); - - assert_eq!(summary.total_tokens, 120); - assert_eq!(summary.session_count, 1); - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - } - - #[test] - fn foreign_only_input_does_not_fabricate_known_zero_native_usage() { - let dir = tempfile::tempdir().unwrap(); - let gemini_base = dir.path().join(".gemini"); - let database_root = gemini_base.join("antigravity-cli").join("conversations"); - fs::create_dir_all(&database_root).unwrap(); - let connection = Connection::open(database_root.join("foreign.db")).unwrap(); - connection - .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) - .unwrap(); - - let context = ScanContext { - database_roots: super::super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions: dir.path().join("missing-tokscale-sessions"), - }; - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_context(&context, now, 7); - - assert_eq!(summary, LocalSessionSummary::default()); - assert_eq!(summary.coverage, LocalHistoryCoverage::Unavailable); - } - - #[test] - fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write(&path, concat!( - "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", - "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", - "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" - )).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 135); - assert_eq!(summary.session_count, 1); - } - - #[test] - fn truncated_or_unreadable_tokscale_history_is_partial() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); - assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); - - let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); - assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); - } - #[test] - fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { - let dir = tempfile::tempdir().unwrap(); - let app = dir - .path() - .join(".gemini") - .join("antigravity") - .join("conversations"); - fs::create_dir_all(&app).unwrap(); - fs::write(app.join("a.db"), b"").unwrap(); - fs::write(app.join("a.db-wal"), b"").unwrap(); - assert_eq!(offline_conversation_count_in(dir.path()), 1); - - fs::remove_file(app.join("a.db")).unwrap(); - let cache = dir - .path() - .join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions"); - fs::create_dir_all(&cache).unwrap(); - fs::write( - cache.join("one.jsonl"), - b"{} -", - ) - .unwrap(); - assert_eq!(offline_conversation_count_in(dir.path()), 1); - } - - #[test] - fn excludes_usage_outside_requested_window() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - concat!( - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", - "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 15); - assert_eq!(summary.session_count, 1); - } - - #[test] - fn oversized_jsonl_line_is_discarded_and_next_usage_row_is_counted() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - let mut text = format!( - r#"{{"type":"usage","padding":"{}"}}"#, - "x".repeat(MAX_JSONL_LINE_BYTES + 32) - ); - text.push('\n'); - text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); - text.push('\n'); - fs::write(&path, text).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 15); - assert_eq!(summary.session_count, 1); - } -} +pub use super::local_history::{offline_conversation_count, summarize_local_usage as summarize}; +pub use crate::spend_contract::{ + LocalHistoryCoverage, LocalTokenHistorySummary as LocalSessionSummary, +}; diff --git a/rust/src/providers/antigravity/local_sessions_reader.rs b/rust/src/providers/antigravity/local_sessions_reader.rs new file mode 100644 index 0000000000..cbc3297cb4 --- /dev/null +++ b/rust/src/providers/antigravity/local_sessions_reader.rs @@ -0,0 +1,583 @@ +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashSet}; +use std::fs::{self, File}; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Duration, Local, TimeZone, Utc}; +use serde_json::Value; + +use super::cost::estimate_cost_usd; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + +const MAX_SESSION_FILES: usize = 2048; +const MAX_SESSION_DISCOVERY_ENTRIES: usize = 16 * 1024; +const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; +const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; +const MAX_TOTAL_SESSION_BYTES: usize = 128 * 1024 * 1024; +const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; + +enum BoundedJsonlLine { + Record(Vec), + Oversized, + Truncated, +} + +pub(super) fn tokscale_sessions_from_values( + home: &Path, + tokscale_config_dir: Option<&str>, +) -> PathBuf { + let tokscale_base = clean_env_path(tokscale_config_dir) + .unwrap_or_else(|| home.join(".config").join("tokscale")); + tokscale_base.join("antigravity-cache").join("sessions") +} + +pub(super) fn configured_tokscale_sessions(home: &Path) -> PathBuf { + let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); + tokscale_sessions_from_values(home, tokscale.as_deref()) +} + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) +} + +pub(super) fn summarize_jsonl_at( + tokscale_sessions: &Path, + now: DateTime, + days: u32, +) -> LocalTokenHistorySummary { + let (paths, truncated) = tokscale_paths(tokscale_sessions); + summarize_jsonl_paths(&paths, now, days, truncated) +} + +pub(super) fn summarize_jsonl_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalTokenHistorySummary { + if paths.is_empty() { + LocalTokenHistorySummary::default() + } else { + summarize_paths(paths, now, days, truncated) + } +} + +fn tokscale_paths(base: &Path) -> (Vec, bool) { + let Ok(entries) = fs::read_dir(base) else { + return (Vec::new(), false); + }; + bounded_tokscale_paths(entries, MAX_SESSION_DISCOVERY_ENTRIES, MAX_SESSION_FILES) +} + +fn bounded_tokscale_paths( + entries: fs::ReadDir, + max_entries: usize, + max_files: usize, +) -> (Vec, bool) { + let mut paths = BinaryHeap::>::with_capacity(max_files); + let mut truncated = false; + + for (entries_examined, entry) in entries.enumerate() { + if entries_examined == max_entries { + truncated = true; + break; + } + + let Ok(entry) = entry else { + truncated = true; + continue; + }; + let path = entry.path(); + let is_jsonl = path + .extension() + .and_then(|value| value.to_str()) + .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")); + if !is_jsonl { + continue; + } + + if paths.len() < max_files { + paths.push(Reverse(path)); + } else { + truncated = true; + if let Some(Reverse(smallest)) = paths.peek() + && path > *smallest + { + paths.pop(); + paths.push(Reverse(path)); + } + } + } + + let mut paths: Vec<_> = paths.into_iter().map(|Reverse(path)| path).collect(); + paths.sort(); + (paths, truncated) +} + +pub(super) fn count_jsonl_sessions_at(base: &Path) -> usize { + tokscale_paths(base).0.len() +} + +fn summarize_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalTokenHistorySummary { + summarize_paths_with_budget(paths, now, days, truncated, MAX_TOTAL_SESSION_BYTES) +} + +fn summarize_paths_with_budget( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, + total_byte_budget: usize, +) -> LocalTokenHistorySummary { + let first_day = now.with_timezone(&Local).date_naive() + - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); + let mut total_tokens = 0_u64; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); + let mut sessions_with_usage = HashSet::new(); + let mut seen_response_ids = HashSet::new(); + let mut complete = !truncated; + let mut remaining_total_bytes = total_byte_budget; + + for path in paths.iter().take(MAX_SESSION_FILES) { + if remaining_total_bytes == 0 { + complete = false; + break; + } + let file = match File::open(path) { + Ok(file) => file, + Err(_) => { + complete = false; + continue; + } + }; + match file.metadata() { + Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, + Ok(_) => {} + Err(_) => complete = false, + } + let mut reader = BufReader::new(file); + let mut remaining = MAX_SESSION_FILE_BYTES; + let mut path_had_usage = false; + let mut model = None::; + loop { + let line = match read_bounded_jsonl_line( + &mut reader, + &mut remaining, + &mut remaining_total_bytes, + ) { + Ok(Some(BoundedJsonlLine::Record(line))) => line, + Ok(Some(BoundedJsonlLine::Oversized)) => { + complete = false; + continue; + } + Ok(Some(BoundedJsonlLine::Truncated)) => { + complete = false; + break; + } + Ok(None) => break, + Err(_) => { + complete = false; + break; + } + }; + if line.is_empty() { + continue; + } + let Ok(value) = serde_json::from_slice::(&line) else { + complete = false; + continue; + }; + if !value.is_object() { + complete = false; + continue; + } + let kind = value.get("type").and_then(Value::as_str); + if kind == Some("session_meta") { + model = value + .get("modelId") + .or_else(|| value.get("model_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + continue; + } + if kind != Some("usage") && value.get("input").is_none() { + continue; + } + if !has_valid_token_fields(&value) { + complete = false; + continue; + } + + let Some(timestamp_ms) = value.get("timestamp").and_then(Value::as_i64) else { + complete = false; + continue; + }; + let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { + complete = false; + continue; + }; + + if let Some(response_id) = value + .get("responseId") + .or_else(|| value.get("response_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + && !seen_response_ids.insert(response_id.to_string()) + { + continue; + } + + if at > now || at.with_timezone(&Local).date_naive() < first_day { + continue; + } + + let input = token_field(&value, &["input"]); + let output = token_field(&value, &["output"]); + let cache_read = token_field(&value, &["cacheRead", "cache_read"]); + let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); + let reasoning = token_field( + &value, + &["reasoning", "reasoningTokens", "reasoning_tokens"], + ); + let Some(total) = [input, output, cache_read, cache_write, reasoning] + .into_iter() + .try_fold(0_u64, u64::checked_add) + else { + complete = false; + continue; + }; + if total == 0 { + continue; + } + let Some(next_total_tokens) = total_tokens.checked_add(total) else { + complete = false; + continue; + }; + total_tokens = next_total_tokens; + cost_estimate.record_list_price(estimate_cost_usd( + model.as_deref(), + input, + cache_read, + cache_write, + output.saturating_add(reasoning), + )); + path_had_usage = true; + } + if path_had_usage { + sessions_with_usage.insert(path.clone()); + } + } + + LocalTokenHistorySummary { + total_tokens, + session_count: sessions_with_usage.len(), + coverage: if paths.is_empty() { + LocalHistoryCoverage::Unavailable + } else if complete { + LocalHistoryCoverage::Complete + } else { + LocalHistoryCoverage::Partial + }, + cost_estimate, + } +} + +fn read_bounded_jsonl_line( + reader: &mut R, + remaining_file_bytes: &mut usize, + remaining_total_bytes: &mut usize, +) -> std::io::Result> { + if *remaining_file_bytes == 0 { + return Ok(None); + } + if *remaining_total_bytes == 0 { + return Ok(if reader.fill_buf()?.is_empty() { + None + } else { + Some(BoundedJsonlLine::Truncated) + }); + } + let mut line = Vec::new(); + let mut saw_input = false; + let mut discarding = false; + + loop { + let chunk = reader.fill_buf()?; + if chunk.is_empty() { + return Ok(saw_input.then_some(if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + let bounded_len = chunk + .len() + .min(*remaining_file_bytes) + .min(*remaining_total_bytes); + if bounded_len == 0 { + return Ok(None); + } + let bounded = &chunk[..bounded_len]; + let newline = bounded.iter().position(|byte| *byte == b'\n'); + let segment_end = newline.unwrap_or(bounded.len()); + let segment = &bounded[..segment_end]; + saw_input = saw_input || !segment.is_empty() || newline.is_some(); + if !discarding { + if line.len().saturating_add(segment.len()) <= MAX_JSONL_LINE_BYTES { + line.extend_from_slice(segment); + } else { + line.clear(); + discarding = true; + } + } + let consumed = segment_end + usize::from(newline.is_some()); + reader.consume(consumed); + *remaining_file_bytes = remaining_file_bytes.saturating_sub(consumed); + *remaining_total_bytes = remaining_total_bytes.saturating_sub(consumed); + if newline.is_some() { + return Ok(Some(if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + if *remaining_file_bytes == 0 || *remaining_total_bytes == 0 { + let at_eof = reader.fill_buf()?.is_empty(); + return Ok(Some(if !at_eof { + BoundedJsonlLine::Truncated + } else if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + } +} + +fn token_field(value: &Value, keys: &[&str]) -> u64 { + keys.iter() + .find_map(|key| value.get(*key).and_then(Value::as_u64)) + .unwrap_or(0) +} + +fn has_valid_token_fields(value: &Value) -> bool { + let mut has_token_field = false; + for key in [ + "input", + "output", + "cacheRead", + "cache_read", + "cacheWrite", + "cache_write", + "reasoning", + "reasoningTokens", + "reasoning_tokens", + ] { + if let Some(field) = value.get(key) { + if field.as_u64().is_none() { + return false; + } + has_token_field = true; + } + } + has_token_field +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tokscale_discovery_bounds_entries_and_retained_paths() { + let dir = tempfile::tempdir().unwrap(); + for index in 0..5 { + fs::write(dir.path().join(format!("session-{index}.jsonl")), "").unwrap(); + } + + let (paths, truncated) = bounded_tokscale_paths(fs::read_dir(dir.path()).unwrap(), 2, 10); + assert_eq!(paths.len(), 2); + assert!(truncated); + + let (paths, truncated) = bounded_tokscale_paths(fs::read_dir(dir.path()).unwrap(), 10, 2); + assert_eq!(paths.len(), 2); + assert!(truncated); + assert_eq!(paths[0].file_name().unwrap(), "session-3.jsonl"); + assert_eq!(paths[1].file_name().unwrap(), "session-4.jsonl"); + } + + #[test] + fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { + let dir = tempfile::tempdir().unwrap(); + let known = dir.path().join("known.jsonl"); + let unknown = dir.path().join("unknown.jsonl"); + fs::write( + &known, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", + "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" + ), + ) + .unwrap(); + fs::write( + &unknown, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", + "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[known, unknown], now, 7, false); + + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 1); + assert!(summary.cost_estimate.known_subtotal_usd.is_some()); + assert_eq!(summary.total_usd(), None); + } + #[test] + fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write(&path, concat!( + "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", + "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", + "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" + )).unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 135); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } + + #[test] + fn truncated_or_unreadable_tokscale_history_is_partial() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); + assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); + + let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); + assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); + } + #[test] + fn excludes_usage_outside_requested_window() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", + "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.session_count, 1); + } + + #[test] + fn oversized_jsonl_line_marks_coverage_partial_and_next_row_is_counted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + let mut text = format!( + r#"{{"type":"usage","padding":"{}"}}"#, + "x".repeat(MAX_JSONL_LINE_BYTES + 32) + ); + text.push('\n'); + text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); + text.push('\n'); + fs::write(&path, text).unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn malformed_jsonl_record_marks_coverage_partial_and_next_row_is_counted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{malformed json}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":\"invalid\"}\n", + "{\"type\":\"usage\",\"input\":10}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[path], now, 7, false); + + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn total_scan_byte_budget_stops_later_records_and_marks_partial() { + let dir = tempfile::tempdir().unwrap(); + let first_path = dir.path().join("session-a.jsonl"); + let second_path = dir.path().join("session-b.jsonl"); + let first = "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n"; + fs::write(&first_path, first).unwrap(); + fs::write( + &second_path, + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":20}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = + summarize_paths_with_budget(&[first_path, second_path], now, 7, false, first.len()); + + assert_eq!(summary.total_tokens, 10); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn token_sum_overflow_marks_coverage_partial_without_saturation() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615,\"output\":1}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":1}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[path], now, 7, false); + + assert_eq!(summary.total_tokens, u64::MAX); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } +} diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index b050660411..bc71f794d7 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -10,9 +10,12 @@ use chrono::{DateTime, Duration, Local, TimeZone, Utc}; use rusqlite::{Connection, OpenFlags, TransactionBehavior, types::ValueRef}; use self::local_bot_id::{ExactStepTimestamp, embedded_timestamps_agree, record_exact_bot_id}; +use super::cost::estimate_cost_usd; use super::local_proto::{ParsedTurn, parse_step_metadata, parse_turn}; -use super::local_sessions::{LocalHistoryCoverage, LocalSessionSummary}; use super::local_step_resolver::{StepOccurrence, resolve_step_timestamps}; +#[cfg(test)] +use crate::spend_contract::LocalTokenHistorySummary as LocalSessionSummary; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; const MAX_DATABASES: usize = 500; const MAX_DIRECTORY_ENTRIES: usize = 10_000; @@ -34,7 +37,7 @@ pub(super) enum SQLiteScan { /// This is non-authoritative: callers may continue with another local /// history source instead of treating the scan as known-empty history. Unsupported, - Summary(LocalSessionSummary), + Summary(LocalTokenHistorySummary), } #[derive(Debug)] @@ -184,10 +187,26 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } let mut total_tokens = 0_u64; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); let mut sessions = HashSet::new(); let mut rows: HashMap<(String, i64), Event> = HashMap::new(); let mut responses: HashMap<(String, String), Event> = HashMap::new(); + let mut label_models = HashMap::<(String, String), String>::new(); + let mut conflicting_labels = HashSet::<(String, String)>::new(); + for event in &events { + let (Some(label), Some(model)) = (event.turn.label.as_ref(), event.turn.model.as_ref()) + else { + continue; + }; + let key = (event.session.clone(), label.clone()); + if label_models.get(&key).is_some_and(|prior| prior != model) { + conflicting_labels.insert(key); + } else { + label_models.insert(key, model.clone()); + } + } + for event in events { let row_key = (event.session.clone(), event.row); if let Some(prior) = rows.get(&row_key) { @@ -234,10 +253,30 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } + if event.total > 0 { + let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { + let inherited_model = event.turn.label.as_ref().and_then(|label| { + let key = (event.session.clone(), label.clone()); + (!conflicting_labels.contains(&key)) + .then(|| label_models.get(&key)) + .flatten() + .map(String::as_str) + }); + let model = event.turn.model.as_deref().or(inherited_model); + let input = usage.system_prompt.checked_add(usage.new_input); + let output = usage.output.checked_add(usage.reasoning); + if let (Some(input), Some(output)) = (input, output) { + estimate_cost_usd(model, input, usage.cache_read, 0, output) + } else { + None + } + }); + cost_estimate.record_list_price(estimated_cost); + } sessions.insert(event.session); } - SQLiteScan::Summary(LocalSessionSummary { + SQLiteScan::Summary(LocalTokenHistorySummary { total_tokens, session_count: sessions.len(), coverage: if complete { @@ -245,6 +284,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } else { LocalHistoryCoverage::Partial }, + cost_estimate, }) } @@ -354,10 +394,10 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result Vec { + let mut bytes = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + bytes.push(byte); + if value == 0 { + return bytes; + } + } +} + +fn field_varint(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes +} + +fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes +} + +fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { + valid_turn_blob_with_model(input, timestamp_seconds, None) +} + +fn valid_turn_blob_with_model(input: u64, timestamp_seconds: u64, model: Option<&str>) -> Vec { + let mut usage = field_varint(1, 11); + usage.extend(field_varint(2, input)); + usage.extend(field_varint(5, 50)); + usage.extend(field_varint(9, 30)); + usage.extend(field_varint(10, 7)); + + let mut timestamp = field_varint(1, timestamp_seconds); + timestamp.extend(field_varint(2, 0)); + let mut chat = field_bytes(4, &usage); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + if let Some(model) = model { + chat.extend(field_bytes(19, model.as_bytes())); + } + field_bytes(1, &chat) +} + +fn zero_token_turn_blob(timestamp_seconds: u64, model: &str) -> Vec { + let timestamp = field_varint(1, timestamp_seconds); + let mut chat = field_bytes(4, &[]); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + chat.extend(field_bytes(19, model.as_bytes())); + field_bytes(1, &chat) +} + +#[test] +fn missing_databases_falls_through() { + let dir = tempfile::tempdir().unwrap(); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::NoDatabases + )); +} + +#[test] +fn foreign_database_is_non_authoritative() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::Unsupported + )); +} + +#[test] +fn empty_supported_database_is_confirmed_zero() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 0); + assert_eq!(summary.session_count, 0); +} + +#[test] +fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { + let dir = tempfile::tempdir().unwrap(); + let first_root = dir.path().join("first"); + let second_root = dir.path().join("second"); + fs::create_dir_all(&first_root).unwrap(); + fs::create_dir_all(&second_root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); + + for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { + let conn = Connection::open(root.join("session.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob(input, timestamp)], + ) + .unwrap(); + } + + let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); + assert_eq!(summary.session_count, 2); +} + +#[test] +fn zero_token_unknown_model_does_not_poison_priced_history() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let now = Utc::now(); + let timestamp = u64::try_from(now.timestamp()).unwrap(); + + let priced = Connection::open(root.join("priced.db")).unwrap(); + priced + .execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + priced + .execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob_with_model( + 100, + timestamp, + Some("claude-sonnet-4-6"), + )], + ) + .unwrap(); + + let zero = Connection::open(root.join("zero.db")).unwrap(); + zero.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + zero.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [zero_token_turn_blob(timestamp, "unknown-model")], + ) + .unwrap(); + + let SQLiteScan::Summary(summary) = summarize(&[root], now, 30) else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 198); + assert_eq!(summary.session_count, 2); + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 0); + assert!( + summary + .cost_estimate + .known_subtotal_usd + .is_some_and(|cost| cost > 0.0) + ); + assert_eq!( + summary.total_usd(), + summary.cost_estimate.known_subtotal_usd + ); +} + +#[test] +fn non_blob_rows_make_coverage_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", + params![1_i64, "not-a-blob"], + ) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); +} +#[test] +fn discovery_allows_exactly_500_databases_but_marks_501_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + for index in 0..MAX_DATABASES { + fs::write(root.join(format!("{index:03}.db")), b"").unwrap(); + } + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(complete); + + fs::write(root.join("overflow.db"), b"").unwrap(); + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(!complete); +} + +#[test] +fn expired_budget_marks_discovery_incomplete() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("one.db"), b"").unwrap(); + let mut budget = Budget::with_deadline(Instant::now()); + let (_, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert!(!complete); +} + +#[test] +fn extra_columns_and_without_rowid_schema_is_supported() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER PRIMARY KEY, data BLOB, extra TEXT) WITHOUT ROWID", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Supported + ); +} + +#[test] +fn generated_columns_are_rejected() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER, data BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL)", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Unsupported + ); +} + +#[test] +fn schema_entry_budget_is_incomplete_not_foreign() { + let conn = Connection::open_in_memory().unwrap(); + for index in 0..=MAX_SCHEMA_ENTRIES { + conn.execute(&format!("CREATE TABLE unrelated_{index}(value TEXT)"), []) + .unwrap(); + } + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Incomplete + ); +} + +#[test] +fn list_price_uses_prompt_plus_input_and_output_plus_reasoning() { + // Upstream AntigravityLocalReaderTests: a known model gets a list-price estimate, a routing + // variant prices from its base model, and an unknown model stays unpriced. + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let now = Utc::now(); + let timestamp = u64::try_from(now.timestamp()).unwrap(); + for (name, model) in [ + ("direct.db", "claude-sonnet-4-6"), + ("routed.db", "claude-sonnet-4-6-Thinking"), + ] { + let conn = Connection::open(root.join(name)).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob_with_model(100, timestamp, Some(model))], + ) + .unwrap(); + } + + let SQLiteScan::Summary(summary) = summarize(std::slice::from_ref(&root), now, 30) else { + panic!("supported databases should produce coverage"); + }; + // system prompt 11 + input 100, cache read 50, output 30 + reasoning 7. + let per_request = + crate::core::CostUsagePricing::claude_cost_usd("claude-sonnet-4-6", 111, 50, 0, 37) + .expect("built-in public price"); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.cost_estimate.coverage.estimated, 2); + assert_eq!(summary.cost_estimate.coverage.unpriced, 0); + assert_eq!(summary.total_usd(), Some(per_request * 2.0)); + + let conn = Connection::open(root.join("unknown.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob_with_model( + 100, + timestamp, + Some("fixture-unpriced"), + )], + ) + .unwrap(); + drop(conn); + + let SQLiteScan::Summary(summary) = summarize(std::slice::from_ref(&root), now, 30) else { + panic!("supported databases should produce coverage"); + }; + assert_eq!(summary.cost_estimate.coverage.estimated, 2); + assert_eq!(summary.cost_estimate.coverage.unpriced, 1); + assert_eq!(summary.total_usd(), None); + assert_eq!( + summary.cost_estimate.known_subtotal_usd, + Some(per_request * 2.0) + ); +} diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 5f35764663..d24eb8c637 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -4,9 +4,12 @@ //! Uses Windows process detection to find CSRF token mod cli_fallback; +mod cost; mod legacy_status; +mod local_history; mod local_proto; pub mod local_sessions; +mod local_sessions_reader; mod local_sqlite; mod local_step_resolver; mod quota_summary; @@ -50,6 +53,7 @@ const AGY_READY_POLL_INTERVAL: Duration = Duration::from_millis(250); const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const QUOTA_SUMMARY_PATH: &str = "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; + /// Serialize task-owned `agy` launches so concurrent app surfaces never start /// multiple interactive CLI servers at the same time. #[cfg(windows)] @@ -623,7 +627,7 @@ impl AntigravityProvider { } fn offline_usage_result() -> Option { - let count = local_sessions::offline_conversation_count(); + let count = local_history::offline_conversation_count(); if count == 0 { return None; } diff --git a/rust/src/providers/crof/mod.rs b/rust/src/providers/crof/mod.rs deleted file mode 100644 index 4abef6b4ca..0000000000 --- a/rust/src/providers/crof/mod.rs +++ /dev/null @@ -1,249 +0,0 @@ -//! Crof provider implementation. -//! -//! Fetches API key based credit/request quota data from Crof. - -use async_trait::async_trait; -use reqwest::Client; -use serde::Deserialize; - -use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, -}; - -const CROF_USAGE_URL: &str = "https://crof.ai/usage_api/"; -const CROF_CREDENTIAL_TARGET: &str = "codexbar-crof"; -const BROWSER_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"; - -#[derive(Debug, Deserialize)] -struct CrofUsageResponse { - credits: f64, - #[serde(default, rename = "requests_plan")] - requests_plan: Option, - #[serde(default, rename = "usable_requests")] - usable_requests: Option, -} - -pub struct CrofProvider { - metadata: ProviderMetadata, - client: Client, -} - -impl CrofProvider { - pub fn new() -> Self { - Self { - metadata: ProviderMetadata { - id: ProviderId::Crof, - display_name: "Crof", - session_label: "Balance", - weekly_label: "Requests", - supports_opus: false, - supports_credits: true, - default_enabled: false, - is_primary: false, - dashboard_url: Some("https://crof.ai"), - status_page_url: None, - tertiary_label_key: None, - }, - client: crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(15)) - .build() - .unwrap_or_else(|_| Client::new()), - } - } - - fn api_key(api_key: Option<&str>) -> Result { - super_key( - api_key, - CROF_CREDENTIAL_TARGET, - &["CROF_API_KEY", "CROFAI_API_KEY"], - ) - } - - async fn fetch_api(&self, api_key: &str) -> Result { - let response = self - .client - .get(CROF_USAGE_URL) - .bearer_auth(api_key) - .header("Accept", "application/json") - .header("User-Agent", BROWSER_USER_AGENT) - .send() - .await?; - - let status = response.status(); - let body = response.text().await.unwrap_or_default(); - if status == reqwest::StatusCode::UNAUTHORIZED { - return Err(ProviderError::AuthRequired); - } - if status == reqwest::StatusCode::FORBIDDEN { - if body.contains("cloudflare") || body.contains("Error 1010") { - return Err(ProviderError::Other( - "Crof usage API blocked by Cloudflare (1010). Retry from the desktop app." - .into(), - )); - } - return Err(ProviderError::AuthRequired); - } - if !status.is_success() { - return Err(ProviderError::Other(format!( - "Crof API returned status {status}" - ))); - } - - let usage: CrofUsageResponse = serde_json::from_str(&body) - .map_err(|e| ProviderError::Parse(format!("Failed to parse Crof usage: {e}")))?; - Ok(snapshot_from_usage(&usage)) - } -} - -fn snapshot_from_usage(usage: &CrofUsageResponse) -> UsageSnapshot { - let credits = usage.credits.max(0.0); - let display = if credits <= 0.0 { - "$0.00".to_string() - } else if credits >= 0.01 { - format!("${:.2}", (credits * 100.0).floor() / 100.0) - } else { - format!("${credits:.4}") - }; - let mut primary = RateWindow::new(if credits > 0.0 { 0.0 } else { 100.0 }); - primary.reset_description = Some(display.clone()); - - let mut snapshot = UsageSnapshot::new(primary).with_login_method(format!("{display} balance")); - - if let (Some(plan), Some(usable)) = (usage.requests_plan, usage.usable_requests) { - let remaining = usable.max(0.0).min(plan.max(0.0)); - let remaining_percent = if plan > 0.0 { - ((remaining / plan) * 100.0).clamp(0.0, 100.0) - } else { - 0.0 - }; - let mut requests = RateWindow::new(100.0 - remaining_percent); - requests.reset_description = Some(format!("{remaining:.0} requests left")); - snapshot = snapshot.with_secondary(requests); - } - - snapshot -} - -impl Default for CrofProvider { - fn default() -> Self { - Self::new() - } -} - -#[async_trait] -impl Provider for CrofProvider { - fn id(&self) -> ProviderId { - ProviderId::Crof - } - - fn metadata(&self) -> &ProviderMetadata { - &self.metadata - } - - async fn fetch_usage(&self, ctx: &FetchContext) -> Result { - match ctx.source_mode { - SourceMode::Auto | SourceMode::OAuth => { - let api_key = Self::api_key(ctx.api_key.as_deref())?; - Ok(ProviderFetchResult::new( - self.fetch_api(&api_key).await?, - "api", - )) - } - SourceMode::Web | SourceMode::Cli => { - Err(ProviderError::UnsupportedSource(ctx.source_mode)) - } - } - } - - fn available_sources(&self) -> Vec { - vec![SourceMode::Auto, SourceMode::OAuth] - } -} - -fn super_key( - explicit: Option<&str>, - credential_target: &str, - env_names: &[&str], -) -> Result { - if let Some(key) = explicit - && !key.trim().is_empty() - { - return Ok(key.trim().to_string()); - } - if let Ok(entry) = keyring::Entry::new(credential_target, "api_key") - && let Ok(key) = entry.get_password() - && !key.trim().is_empty() - { - return Ok(key); - } - for env in env_names { - if let Ok(key) = std::env::var(env) - && !key.trim().is_empty() - { - return Ok(key); - } - } - Err(ProviderError::NotInstalled(format!( - "API key not found. Set {} in Preferences or environment.", - env_names.join(" / ") - ))) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn crof_snapshot_formats_request_and_credit_windows() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 12.5, - requests_plan: Some(100.0), - usable_requests: Some(25.0), - }); - assert_eq!(snapshot.primary.used_percent, 0.0); - assert_eq!( - snapshot.primary.reset_description.as_deref(), - Some("$12.50") - ); - assert_eq!(snapshot.secondary.unwrap().used_percent, 75.0); - } - - #[test] - fn crof_payg_balance_only_does_not_require_request_quota() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 3.019, - requests_plan: None, - usable_requests: None, - }); - assert_eq!(snapshot.primary.used_percent, 0.0); - assert_eq!(snapshot.primary.reset_description.as_deref(), Some("$3.01")); - assert!(snapshot.secondary.is_none()); - assert_eq!(snapshot.login_method.as_deref(), Some("$3.01 balance")); - } - - #[test] - fn crof_sub_cent_balance_is_not_exhausted() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 0.0073, - requests_plan: None, - usable_requests: None, - }); - assert_eq!(snapshot.primary.used_percent, 0.0); - assert_eq!( - snapshot.primary.reset_description.as_deref(), - Some("$0.0073") - ); - } - - #[test] - fn crof_zero_balance_is_exhausted() { - let snapshot = snapshot_from_usage(&CrofUsageResponse { - credits: 0.0, - requests_plan: None, - usable_requests: None, - }); - assert_eq!(snapshot.primary.used_percent, 100.0); - assert_eq!(snapshot.primary.reset_description.as_deref(), Some("$0.00")); - } -} diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index 907241d5d9..059c0f8d28 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -7,13 +7,12 @@ use reqwest::Url; use std::path::{Path, PathBuf}; -use super::web; use super::{ - FetchContext, KimiCodeApiUsageResponse, KimiProvider, ProviderError, UsageSnapshot, + FetchContext, KimiCodeApiUsageResponse, KimiProvider, KimiRegion, ProviderError, UsageSnapshot, ascii_header_value, cleaned_env, cleaned_owned, kimi_window_minutes, }; +use super::{ratio_pool, web}; -const KIMI_CODE_API_BASE: &str = "https://api.kimi.com"; const KIMI_CODE_API_KEY_ENV: &str = "KIMI_CODE_API_KEY"; const KIMI_CODE_BASE_URL_ENV: &str = "KIMI_CODE_BASE_URL"; const KIMI_CODE_HOME_ENV: &str = "KIMI_CODE_HOME"; @@ -42,12 +41,13 @@ struct KimiCodeCredentialFile { /// to the un-enriched snapshot. pub(crate) async fn fetch_via_code_api( ctx: &FetchContext, + region: KimiRegion, api_key_override: Option<&str>, identity_headers_override: Option<&[(&str, String)]>, login_method: &str, ) -> Result { let api_key = code_api_key(api_key_override.or(ctx.api_key.as_deref()))?; - let base_url = code_api_base_url()?; + let base_url = code_api_base_url(region)?; let endpoint = code_api_usage_endpoint(&base_url)?; let client = crate::core::credentialed_http_client_builder() .timeout(std::time::Duration::from_secs(30)) @@ -88,14 +88,15 @@ pub(crate) async fn fetch_via_code_api( // Upstream #2622: enrich Code API + CLI usage with the monthly membership // pool from a signed-in Kimi Desktop (or browser/manual) session. - for web_token in web::web_auth_tokens(ctx.manual_cookie_header.as_deref()) { - match web::fetch_subscription_for_enrichment_result(&client, &web_token).await { + for web_token in web::web_auth_tokens(ctx.manual_cookie_header.as_deref(), region) { + match web::fetch_subscription_for_enrichment_result(&client, &web_token, region).await { Ok(subscription) => { if let Some(subscription) = subscription { snapshot = super::apply_subscription_windows(snapshot, &subscription); } if !has_plan_name - && let Some(plan) = web::fetch_subscription_plan(&client, &web_token).await + && let Some(plan) = + web::fetch_subscription_plan(&client, &web_token, region).await { snapshot.login_method = Some(plan); } @@ -116,16 +117,35 @@ pub(super) fn snapshot_from_code_api_response( response: KimiCodeApiUsageResponse, ) -> Result { let pools_present = response.usages.is_some(); + let legacy_limit = response.limits.as_ref().and_then(|limits| limits.first()); + let legacy_session_minutes = + legacy_limit.and_then(|limit| limit.window.as_ref().and_then(kimi_window_minutes)); let session_pool = response .usages .as_ref() .and_then(|pools| pools.session.as_ref()) - .and_then(|pool| pool.rate_window(300)); + .and_then(|pool| { + ratio_pool::resolved_ratio_window( + &response, + pool, + legacy_limit.map(|limit| &limit.detail), + 300, + legacy_session_minutes, + ) + }); let weekly_pool = response .usages .as_ref() .and_then(|pools| pools.weekly.as_ref()) - .and_then(|pool| pool.rate_window(10_080)); + .and_then(|pool| { + ratio_pool::resolved_ratio_window( + &response, + pool, + response.usage.as_ref(), + 10_080, + Some(10_080), + ) + }); let monthly_pool = response .usages .as_ref() @@ -139,7 +159,9 @@ pub(super) fn snapshot_from_code_api_response( response .usage .as_ref() - .and_then(|detail| KimiProvider::rate_window_from_usage_detail(detail, None).ok()) + .and_then(|detail| { + KimiProvider::rate_window_from_usage_detail(detail, Some(10_080)).ok() + }) .ok_or_else(|| { ProviderError::Parse("Kimi Code API has no usable quota window".into()) })? @@ -165,6 +187,7 @@ pub(super) fn snapshot_from_code_api_response( } Ok(usage) } + pub(crate) fn code_api_key(explicit: Option<&str>) -> Result { if let Some(key) = explicit.map(str::trim).filter(|key| !key.is_empty()) { return Ok(key.to_string()); @@ -172,8 +195,9 @@ pub(crate) fn code_api_key(explicit: Option<&str>) -> Result Result { - let raw = cleaned_env(KIMI_CODE_BASE_URL_ENV).unwrap_or_else(|| KIMI_CODE_API_BASE.to_string()); +fn code_api_base_url(region: KimiRegion) -> Result { + let raw = cleaned_env(KIMI_CODE_BASE_URL_ENV) + .unwrap_or_else(|| region.code_api_base_url().to_string()); crate::providers::validated_https_url(&raw, "Kimi Code API base") } @@ -210,8 +234,8 @@ pub(crate) fn kimi_code_home() -> Option { /// /// Never refreshes or rewrites CLI-owned `credentials/kimi-code.json`. /// Skips when `KIMI_CODE_BASE_URL` / OAuth host overrides are set. -pub(crate) fn kimi_code_cli_access_token(now_unix: f64) -> Option { - if has_code_endpoint_override() { +pub(crate) fn kimi_code_cli_access_token(region: KimiRegion, now_unix: f64) -> Option { + if region != KimiRegion::China || has_code_endpoint_override() { return None; } let home = kimi_code_home()?; @@ -342,7 +366,7 @@ mod tests { std::env::set_var(KIMI_CODE_HOME_ENV, home.path()); } - let token = kimi_code_cli_access_token(now); + let token = kimi_code_cli_access_token(KimiRegion::China, now); assert_eq!(token.as_deref(), Some("oauth-token")); let after = std::fs::read(&cred_path).unwrap(); @@ -391,7 +415,7 @@ mod tests { std::env::set_var(KIMI_CODE_BASE_URL_ENV, "https://proxy.example.com/kimi"); } assert!(has_code_endpoint_override()); - assert!(kimi_code_cli_access_token(now).is_none()); + assert!(kimi_code_cli_access_token(KimiRegion::China, now).is_none()); // SAFETY: still under the same env_lock() guard; swapping which // override keys are present between assertions. @@ -399,7 +423,7 @@ mod tests { std::env::remove_var(KIMI_CODE_BASE_URL_ENV); std::env::set_var(KIMI_CODE_OAUTH_HOST_ENV, "https://oauth.example.com"); } - assert!(kimi_code_cli_access_token(now).is_none()); + assert!(kimi_code_cli_access_token(KimiRegion::China, now).is_none()); // SAFETY: final cleanup while the env_lock() guard is still alive. unsafe { diff --git a/rust/src/providers/kimi/desktop_token.rs b/rust/src/providers/kimi/desktop_token.rs index 5c2d6bae42..ac40f793ba 100644 --- a/rust/src/providers/kimi/desktop_token.rs +++ b/rust/src/providers/kimi/desktop_token.rs @@ -29,8 +29,6 @@ const DESKTOP_APP_DIR: &str = "kimi-desktop"; const COOKIES_FILE: &str = "Cookies"; const LOCAL_STATE_FILE: &str = "Local State"; const AUTH_COOKIE_NAME: &str = "kimi-auth"; -const AUTH_COOKIE_HOSTS: [&str; 4] = ["www.kimi.com", ".www.kimi.com", ".kimi.com", "kimi.com"]; - impl KimiDesktopAuthToken { /// Cookies database inside a caller-provided `data_root` (upstream /// `cookiesDatabaseURL(homeDirectory:)` shape for test injection). @@ -47,12 +45,20 @@ impl KimiDesktopAuthToken { /// Desktop session, or `None` when the app/database/cookie is absent or /// unreadable. Production entry point. pub fn load() -> Option { + Self::load_for_region(super::KimiRegion::China) + } + + pub fn load_for_region(region: super::KimiRegion) -> Option { let data_root = dirs::data_dir()?; - Self::load_from(&data_root) + Self::load_from_region(&data_root, region) } /// Read from an explicit `data_root` (Electron `userData` parent). pub fn load_from(data_root: &Path) -> Option { + Self::load_from_region(data_root, super::KimiRegion::China) + } + + pub fn load_from_region(data_root: &Path, region: super::KimiRegion) -> Option { let aes_key = crate::browser::cookies::CookieExtractor::get_chromium_encryption_key( &Self::local_state_path(data_root), ) @@ -63,13 +69,21 @@ impl KimiDesktopAuthToken { ); }) .ok(); - Self::load_token(&Self::cookies_database_path(data_root), aes_key.as_deref()) + Self::load_token( + &Self::cookies_database_path(data_root), + aes_key.as_deref(), + region.desktop_cookie_hosts(), + ) } /// Core read (upstream `read(databaseURL:immutable:)`): WAL-safe /// read-only open → newest `kimi-auth` row → decode. `aes_key` is the /// Chromium app cookie key; `None` restricts reads to plaintext rows. - fn load_token(database_path: &Path, aes_key: Option<&[u8]>) -> Option { + fn load_token( + database_path: &Path, + aes_key: Option<&[u8]>, + hosts: &[&str; 4], + ) -> Option { if !database_path.is_file() { return None; } @@ -81,7 +95,7 @@ impl KimiDesktopAuthToken { tracing::debug!(error = %err, "Kimi Desktop Cookies open failed"); }) .ok()?; - read_newest_auth_cookie(&conn) + read_newest_auth_cookie(&conn, hosts) .inspect_err(|err| { tracing::debug!(error = %err, "Kimi Desktop cookies read failed"); }) @@ -145,7 +159,10 @@ fn decode_cookie_value(row: (String, Vec), aes_key: Option<&[u8]>) -> Option .filter(|plain| !plain.is_empty()) } -fn read_newest_auth_cookie(conn: &rusqlite::Connection) -> rusqlite::Result<(String, Vec)> { +fn read_newest_auth_cookie( + conn: &rusqlite::Connection, + hosts: &[&str; 4], +) -> rusqlite::Result<(String, Vec)> { // Upstream query verbatim: newest `kimi-auth` across the registered // kimi.com cookie scopes by last access. let mut statement = conn.prepare( @@ -157,13 +174,7 @@ fn read_newest_auth_cookie(conn: &rusqlite::Connection) -> rusqlite::Result<(Str LIMIT 1", )?; statement.query_row( - rusqlite::params![ - AUTH_COOKIE_NAME, - AUTH_COOKIE_HOSTS[0], - AUTH_COOKIE_HOSTS[1], - AUTH_COOKIE_HOSTS[2], - AUTH_COOKIE_HOSTS[3], - ], + rusqlite::params![AUTH_COOKIE_NAME, hosts[0], hosts[1], hosts[2], hosts[3],], |row| Ok((row.get::<_, String>(0)?, row.get::<_, Vec>(1)?)), ) } @@ -357,12 +368,24 @@ mod tests { insert_cookie_row(&conn, "www.kimi.com", "", &encrypted, 1); assert_eq!( - KimiDesktopAuthToken::load_token(&database, Some(key.as_slice())).as_deref(), + KimiDesktopAuthToken::load_token( + &database, + Some(key.as_slice()), + crate::providers::KimiRegion::China.desktop_cookie_hosts(), + ) + .as_deref(), Some("encrypted-kimi-token") ); // Without a key the encrypted row cannot be used. - assert_eq!(KimiDesktopAuthToken::load_token(&database, None), None); + assert_eq!( + KimiDesktopAuthToken::load_token( + &database, + None, + crate::providers::KimiRegion::China.desktop_cookie_hosts(), + ), + None + ); // `load_from` without a usable `Local State` reads plaintext only and // yields nothing (no panic, no secret in logs). assert_eq!(KimiDesktopAuthToken::load_from(root.path()), None); diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index d96d3a36fc..8418ca6176 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -11,11 +11,17 @@ //! from a signed-in Kimi Desktop session (#2622). //! - [`desktop_token`]: read-only, WAL-safe reader for the Kimi Desktop //! (Electron) Chromium cookie store. +//! - [`ratio_pool`]: zero-ratio placeholder reconciliation against matching +//! legacy counters (upstream 0.63.0). mod code_api; pub mod desktop_token; +mod ratio_pool; +mod region; mod web; +pub use region::KimiRegion; + use async_trait::async_trait; use chrono::{DateTime, Utc}; use reqwest::Client; @@ -27,13 +33,11 @@ use crate::core::{ RateWindow, SourceMode, UsageSnapshot, }; -const KIMI_WEB_USAGE_URL: &str = - "https://www.kimi.com/apiv2/kimi.gateway.billing.v1.BillingService/GetUsages"; -const KIMI_SUBSCRIPTION_STATS_URL: &str = - "https://www.kimi.com/apiv2/kimi.gateway.membership.v2.MembershipService/GetSubscriptionStats"; -const KIMI_SUBSCRIPTION_URL: &str = - "https://www.kimi.com/apiv2/kimi.gateway.membership.v2.MembershipService/GetSubscription"; -const KIMI_COOKIE_DOMAINS: [&str; 2] = ["www.kimi.com", "kimi.moonshot.cn"]; +const KIMI_WEB_USAGE_SERVICE: &str = "kimi.gateway.billing.v1.BillingService/GetUsages"; +const KIMI_SUBSCRIPTION_STATS_SERVICE: &str = + "kimi.gateway.membership.v2.MembershipService/GetSubscriptionStats"; +const KIMI_SUBSCRIPTION_SERVICE: &str = + "kimi.gateway.membership.v2.MembershipService/GetSubscription"; #[derive(Debug, Deserialize)] struct KimiCodeApiUsageResponse { @@ -330,11 +334,12 @@ impl Provider for KimiProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { tracing::debug!("Fetching Kimi usage"); + let region = KimiRegion::from_settings(ctx.api_region.as_deref()); match ctx.source_mode { SourceMode::Auto => { if code_api::code_api_key(ctx.api_key.as_deref()).is_ok() { - match code_api::fetch_via_code_api(ctx, None, None, "Code API").await { + match code_api::fetch_via_code_api(ctx, region, None, None, "Code API").await { Ok(usage) => { return Ok(ProviderFetchResult::new(usage, "code-api")); } @@ -347,11 +352,14 @@ impl Provider for KimiProvider { } } - if let Some(cli_token) = code_api::kimi_code_cli_access_token(unix_now_secs()) { + if let Some(cli_token) = + code_api::kimi_code_cli_access_token(region, unix_now_secs()) + { let home = code_api::kimi_code_home().unwrap_or_default(); let headers = code_api::kimi_code_cli_identity_headers(&home); match code_api::fetch_via_code_api( ctx, + region, Some(&cli_token), Some(&headers), "Kimi Code CLI", @@ -370,15 +378,16 @@ impl Provider for KimiProvider { } } - let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref()).await?; + let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref(), region).await?; Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::OAuth => { - let usage = code_api::fetch_via_code_api(ctx, None, None, "Code API").await?; + let usage = + code_api::fetch_via_code_api(ctx, region, None, None, "Code API").await?; Ok(ProviderFetchResult::new(usage, "code-api")) } SourceMode::Web => { - let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref()).await?; + let usage = web::fetch_via_web(ctx.manual_cookie_header.as_deref(), region).await?; Ok(ProviderFetchResult::new(usage, "web")) } SourceMode::Cli => Err(ProviderError::UnsupportedSource(SourceMode::Cli)), @@ -489,6 +498,7 @@ fn is_equivalent_to_weekly_window(window: &RateWindow, weekly: &RateWindow) -> b async fn kimi_web_post( client: &Client, url: &str, + region: KimiRegion, token: &str, body: serde_json::Value, ) -> Result { @@ -498,6 +508,8 @@ async fn kimi_web_post( .header("Cookie", format!("kimi-auth={token}")) .header("Accept", "application/json") .header("Content-Type", "application/json") + .header("Origin", region.web_base_url()) + .header("Referer", region.console_url()) .header( "User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36", diff --git a/rust/src/providers/kimi/ratio_pool.rs b/rust/src/providers/kimi/ratio_pool.rs new file mode 100644 index 0000000000..5d9a42a138 --- /dev/null +++ b/rust/src/providers/kimi/ratio_pool.rs @@ -0,0 +1,422 @@ +//! Kimi Code API ratio-pool reconciliation (upstream 0.63.0 `fd2414d`, +//! #3755). +//! +//! Mixed legacy responses can carry zero ratio placeholders next to populated +//! counters for the same quota. A zero ratio stays authoritative unless all +//! upstream conditions hold: no monthly ratio pool, reliable legacy weekly +//! counters, a count window of the same duration with nonzero reliable use, +//! and a count reset within two seconds of the ratio reset. + +use chrono::TimeDelta; + +use super::{ + KimiCodeApiUsageResponse, KimiRatioPool, KimiUsageDetail, RateWindow, format_usage_amount, +}; + +/// Upstream observed the legacy and ratio reset clocks about 1.45 s apart. +const MATCHING_RESET_TOLERANCE_SECS: i64 = 2; + +/// Integer usage counters (upstream `KimiUsageSnapshot.usageCounts`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct UsageCounts { + used: i64, + limit: i64, + reliable: bool, +} + +/// `used` is authoritative and may exceed the limit during overage; +/// `remaining` only counts when it describes a valid balance. A valid limit +/// without usable counters is kept but marked unreliable. +fn usage_counts(detail: &KimiUsageDetail) -> Option { + let limit = integer_counter(detail.limit.as_ref()).filter(|limit| *limit > 0)?; + if let Some(used) = integer_counter(detail.used.as_ref()).filter(|used| *used >= 0) { + return Some(UsageCounts { + used, + limit, + reliable: true, + }); + } + if let Some(remaining) = integer_counter(detail.remaining.as_ref()) + .filter(|remaining| (0..=limit).contains(remaining)) + { + return Some(UsageCounts { + used: limit - remaining, + limit, + reliable: true, + }); + } + Some(UsageCounts { + used: 0, + limit, + reliable: false, + }) +} + +/// Upstream keeps counters as strings and reads them with `Int(_)`: integer +/// strings and integral JSON numbers count; fractions, padded or formatted +/// strings do not. +fn integer_counter(value: Option<&serde_json::Value>) -> Option { + match value? { + serde_json::Value::Number(number) => number.as_i64().or_else(|| { + let value = number.as_f64()?; + let integral = value.is_finite() + && value.fract() == 0.0 + && value >= i64::MIN as f64 + && value < i64::MAX as f64; + #[allow( + clippy::cast_possible_truncation, + reason = "finite integral value is bounded to the i64 range above" + )] + integral.then_some(value as i64) + }), + serde_json::Value::String(text) => text.parse().ok(), + _ => None, + } +} + +/// Resolve a ratio pool, replacing a zero placeholder with the matching +/// legacy count window when the response proves the counters are the same +/// quota. `count_window_minutes` is the duration the legacy counters report. +pub(super) fn resolved_ratio_window( + response: &KimiCodeApiUsageResponse, + pool: &KimiRatioPool, + detail: Option<&KimiUsageDetail>, + window_minutes: u32, + count_window_minutes: Option, +) -> Option { + let ratio_window = pool.rate_window(window_minutes)?; + let count_window = matching_count_window( + response, + &ratio_window, + detail, + window_minutes, + count_window_minutes, + ); + Some(count_window.unwrap_or(ratio_window)) +} + +fn matching_count_window( + response: &KimiCodeApiUsageResponse, + ratio_window: &RateWindow, + detail: Option<&KimiUsageDetail>, + window_minutes: u32, + count_window_minutes: Option, +) -> Option { + let has_monthly_pool = response + .usages + .as_ref() + .is_some_and(|pools| pools.monthly.is_some()); + let weekly_counts_reliable = response + .usage + .as_ref() + .and_then(usage_counts) + .is_some_and(|counts| counts.reliable); + if ratio_window.used_percent != 0.0 + || has_monthly_pool + || !weekly_counts_reliable + || count_window_minutes != Some(window_minutes) + { + return None; + } + + let detail = detail?; + let counts = usage_counts(detail).filter(|counts| counts.reliable && counts.used > 0)?; + let count_reset = detail + .reset_time + .as_ref() + .and_then(super::parse_kimi_timestamp)?; + let ratio_reset = ratio_window.resets_at?; + if (count_reset - ratio_reset).abs() > TimeDelta::seconds(MATCHING_RESET_TOLERANCE_SECS) { + return None; + } + + #[allow( + clippy::cast_precision_loss, + reason = "quota counters are far below 2^52; the percent is display-only" + )] + let (used, limit) = (counts.used as f64, counts.limit as f64); + Some(RateWindow::with_details( + used / limit * 100.0, + Some(window_minutes), + Some(count_reset), + Some(format!( + "{}/{} credits", + format_usage_amount(used), + format_usage_amount(limit) + )), + )) +} + +#[cfg(test)] +mod tests { + use super::super::code_api::snapshot_from_code_api_response; + use super::super::{KimiCodeApiUsageResponse, UsageSnapshot}; + use chrono::{DateTime, Utc}; + use serde_json::{Value, json}; + + fn parse(value: Value) -> UsageSnapshot { + let response: KimiCodeApiUsageResponse = + serde_json::from_value(value).expect("fixture parses"); + snapshot_from_code_api_response(response).expect("fixture has usable quota") + } + + fn at(text: &str) -> Option> { + Some( + DateTime::parse_from_rfc3339(text) + .expect("valid fixture timestamp") + .with_timezone(&Utc), + ) + } + + /// Win-CodexBar keeps the session pool primary, so single-lane upstream + /// fixtures add a nonzero session pool that is never reconciled. + fn weekly_fixture(usage: Value, weekly_pool: Value) -> Value { + json!({ + "usage": usage, + "usages": { + "limit_5h": { "used_ratio": 0.5 }, + "limit_7d": weekly_pool + } + }) + } + + fn weekly_percent(snapshot: &UsageSnapshot) -> f64 { + snapshot + .secondary + .as_ref() + .expect("weekly lane is present") + .used_percent + } + + fn mixed_international_response() -> Value { + json!({ + "usage": { + "limit": "100", + "used": "19", + "remaining": "81", + "resetTime": "2026-09-19T16:45:59.449979Z" + }, + "limits": [{ + "window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" }, + "detail": { + "limit": "100", + "used": "1", + "remaining": "99", + "resetTime": "2026-09-19T14:45:59.449979Z" + } + }], + "usages": { + "limit_5h": { "used_ratio": 0, "reset_time": "2026-09-19T14:45:58Z" }, + "limit_7d": { "used_ratio": 0, "reset_time": "2026-09-19T16:45:58Z" } + } + }) + } + + // Upstream: `zero ratio placeholder does not hide matching nonzero counts`. + #[test] + fn zero_ratio_placeholder_does_not_hide_matching_nonzero_counts() { + for reset in ["2026-09-19T16:45:58Z", "2026-09-19T16:45:59Z"] { + let snapshot = parse(weekly_fixture( + json!({ + "limit": "100", + "used": "19", + "remaining": "81", + "resetTime": "2026-09-19T16:45:59.449979Z" + }), + json!({ "used_ratio": 0, "reset_time": reset }), + )); + assert_eq!(weekly_percent(&snapshot), 19.0, "ratio reset {reset}"); + assert_eq!(snapshot.primary.used_percent, 50.0); + } + } + + // Upstream: `zero ratio after a different reset stays authoritative`. + #[test] + fn zero_ratio_after_a_different_reset_stays_authoritative() { + let snapshot = parse(weekly_fixture( + json!({ + "limit": "100", + "used": "19", + "remaining": "81", + "resetTime": "2026-09-19T16:45:59Z" + }), + json!({ "used_ratio": 0, "reset_time": "2026-09-26T16:45:59Z" }), + )); + assert_eq!(weekly_percent(&snapshot), 0.0); + } + + // Upstream: `mixed international response retains the used weekly and + // session quotas` (lanes in Win-CodexBar order). + #[test] + fn mixed_international_response_retains_the_used_weekly_and_session_quotas() { + let snapshot = parse(mixed_international_response()); + let weekly = snapshot.secondary.as_ref().expect("weekly lane"); + assert_eq!(weekly.used_percent, 19.0); + assert_eq!(weekly.window_minutes, Some(10_080)); + assert_eq!(weekly.resets_at, at("2026-09-19T16:45:59.449979Z")); + assert_eq!(snapshot.primary.used_percent, 1.0); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + assert_eq!( + snapshot.primary.resets_at, + at("2026-09-19T14:45:59.449979Z") + ); + assert!(snapshot.tertiary.is_none()); + assert!(snapshot.extra_rate_windows.is_empty()); + } + + // Upstream: `nonzero ratios remain authoritative over legacy counts`. + #[test] + fn nonzero_ratios_remain_authoritative_over_legacy_counts() { + for ratio in [0.1869, 0.5] { + let snapshot = parse(weekly_fixture( + json!({ "limit": "100", "used": "19", "resetTime": "2026-09-19T16:45:59Z" }), + json!({ "used_ratio": ratio, "reset_time": "2026-09-19T16:45:59Z" }), + )); + assert!((weekly_percent(&snapshot) - ratio * 100.0).abs() < 0.000_01); + } + } + + // Upstream: `monthly ratio accounts retain zero ratios even with matching + // legacy counts`. + #[test] + fn monthly_ratio_accounts_retain_zero_ratios_even_with_matching_legacy_counts() { + let snapshot = parse(json!({ + "usage": { "limit": "100", "used": "19", "resetTime": "2026-09-19T16:45:59Z" }, + "usages": { + "limit_5h": { "used_ratio": 0.5 }, + "limit_7d": { "used_ratio": 0, "reset_time": "2026-09-19T16:45:59Z" }, + "limit_month_total": { "used_ratio": 0.0313 } + } + })); + assert_eq!(weekly_percent(&snapshot), 0.0); + let monthly = snapshot.tertiary.expect("monthly pool"); + assert!((monthly.used_percent - 3.13).abs() < 0.000_01); + } + + // Upstream: `unmatched count resets cannot override a zero ratio`. + #[test] + fn unmatched_count_resets_cannot_override_a_zero_ratio() { + for reset in [json!(null), json!("invalid"), json!("2026-09-19T16:46:02Z")] { + let snapshot = parse(weekly_fixture( + json!({ "limit": "100", "used": "19", "resetTime": reset }), + json!({ "used_ratio": 0, "reset_time": "2026-09-19T16:45:59Z" }), + )); + assert_eq!(weekly_percent(&snapshot), 0.0, "count reset {reset}"); + } + } + + // Upstream: `invalid or empty counts cannot override a zero ratio`. + #[test] + fn invalid_or_empty_counts_cannot_override_a_zero_ratio() { + for used in ["0", "-1", "invalid"] { + let snapshot = parse(weekly_fixture( + json!({ "limit": "100", "used": used, "resetTime": "2026-09-19T16:45:59Z" }), + json!({ "used_ratio": 0, "reset_time": "2026-09-19T16:45:59Z" }), + )); + assert_eq!(weekly_percent(&snapshot), 0.0, "used {used}"); + } + } + + // Upstream: `different count window duration cannot override the session + // ratio`. + #[test] + fn different_count_window_duration_cannot_override_the_session_ratio() { + let mut response = mixed_international_response(); + response["limits"][0]["window"]["duration"] = json!(120); + let snapshot = parse(response); + assert_eq!(weekly_percent(&snapshot), 19.0); + assert_eq!(snapshot.primary.used_percent, 0.0); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + } + + // Upstream `usageCounts`: an invalid `used` falls back to a valid + // `remaining` balance, which is reliable evidence for both lanes. + #[test] + fn remaining_balance_recovers_invalid_used_counters() { + let mut response = mixed_international_response(); + response["usage"]["used"] = json!("invalid"); + response["limits"][0]["detail"]["used"] = json!("-1"); + let snapshot = parse(response); + assert_eq!(weekly_percent(&snapshot), 19.0); + assert_eq!(snapshot.primary.used_percent, 1.0); + assert_eq!( + snapshot.primary.reset_description.as_deref(), + Some("1/100 credits") + ); + } + + // Upstream gates every override on reliable legacy weekly counters. + #[test] + fn session_override_requires_reliable_weekly_counters() { + for weekly_usage in [ + None, + Some(json!({ "limit": "100", "used": "invalid", "remaining": "101" })), + Some(json!({ "limit": "0", "used": "19" })), + ] { + let mut response = mixed_international_response(); + match weekly_usage { + Some(usage) => response["usage"] = usage, + None => { + response + .as_object_mut() + .expect("fixture object") + .remove("usage"); + } + } + let snapshot = parse(response); + assert_eq!(snapshot.primary.used_percent, 0.0); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + } + } + + // Upstream reads counters with `Int(_)`: fractional or padded values are + // not reliable evidence, while integral JSON numbers are. + #[test] + fn counters_must_be_integers() { + for used in [json!("19.5"), json!(" 19"), json!(19.5)] { + let snapshot = parse(weekly_fixture( + json!({ "limit": "100", "used": used, "resetTime": "2026-09-19T16:45:59Z" }), + json!({ "used_ratio": 0, "reset_time": "2026-09-19T16:45:59Z" }), + )); + assert_eq!(weekly_percent(&snapshot), 0.0, "used {used}"); + } + let snapshot = parse(weekly_fixture( + json!({ "limit": 100.0, "used": 19, "resetTime": "2026-09-19T16:45:59Z" }), + json!({ "used_ratio": 0, "reset_time": "2026-09-19T16:45:59Z" }), + )); + assert_eq!(weekly_percent(&snapshot), 19.0); + } + + fn zero_session_ratio_with_legacy_window(window: Option) -> UsageSnapshot { + let mut response = mixed_international_response(); + match window { + Some(window) => response["limits"][0]["window"] = window, + None => { + response["limits"][0] + .as_object_mut() + .expect("limit object") + .remove("window"); + } + } + parse(response) + } + + // Win-CodexBar tolerates a legacy limit without a window (upstream fails + // to decode it); without a duration the counters cannot claim the lane. + #[test] + fn missing_legacy_window_does_not_override_zero_session_ratio() { + let snapshot = zero_session_ratio_with_legacy_window(None); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + assert_eq!(snapshot.primary.used_percent, 0.0); + } + + #[test] + fn unrecognized_legacy_window_does_not_override_zero_session_ratio() { + let snapshot = zero_session_ratio_with_legacy_window(Some(json!({ + "duration": 300, + "timeUnit": "TIME_UNIT_FORTNIGHT" + }))); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + assert_eq!(snapshot.primary.used_percent, 0.0); + } +} diff --git a/rust/src/providers/kimi/region.rs b/rust/src/providers/kimi/region.rs new file mode 100644 index 0000000000..2043ce7a93 --- /dev/null +++ b/rust/src/providers/kimi/region.rs @@ -0,0 +1,106 @@ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum KimiRegion { + China, + International, +} + +impl KimiRegion { + pub const ALL: [Self; 2] = [Self::China, Self::International]; + + pub fn from_settings(value: Option<&str>) -> Self { + match value.map(str::trim).map(str::to_ascii_lowercase).as_deref() { + Some("international" | "intl" | "global") => Self::International, + _ => Self::China, + } + } + + pub const fn settings_value(self) -> &'static str { + match self { + Self::China => "china", + Self::International => "international", + } + } + + pub const fn display_name(self) -> &'static str { + match self { + Self::China => "China (kimi.com)", + Self::International => "International (kimi.ai)", + } + } + + pub const fn code_api_base_url(self) -> &'static str { + match self { + Self::China => "https://api.kimi.com", + Self::International => "https://api.kimi.ai", + } + } + + pub const fn web_base_url(self) -> &'static str { + match self { + Self::China => "https://www.kimi.com", + Self::International => "https://www.kimi.ai", + } + } + + pub const fn console_url(self) -> &'static str { + match self { + Self::China => "https://www.kimi.com/code/console", + Self::International => "https://www.kimi.ai/code/console", + } + } + + pub const fn cookie_domains(self) -> &'static [&'static str] { + match self { + Self::China => &["www.kimi.com", "kimi.com"], + Self::International => &["www.kimi.ai", "kimi.ai"], + } + } + + pub const fn desktop_cookie_hosts(self) -> &'static [&'static str; 4] { + match self { + Self::China => &["www.kimi.com", ".www.kimi.com", ".kimi.com", "kimi.com"], + Self::International => &["www.kimi.ai", ".www.kimi.ai", ".kimi.ai", "kimi.ai"], + } + } + + pub fn web_api_url(self, service: &str) -> String { + format!("{}/apiv2/{service}", self.web_base_url()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn unknown_and_missing_settings_preserve_china_default() { + assert_eq!(KimiRegion::from_settings(None), KimiRegion::China); + assert_eq!( + KimiRegion::from_settings(Some("unknown")), + KimiRegion::China + ); + assert_eq!( + KimiRegion::from_settings(Some("international")), + KimiRegion::International + ); + } + + #[test] + fn regional_hosts_remain_coherent() { + for region in KimiRegion::ALL { + let suffix = match region { + KimiRegion::China => "kimi.com", + KimiRegion::International => "kimi.ai", + }; + assert!(region.code_api_base_url().ends_with(suffix)); + assert!(region.web_base_url().ends_with(suffix)); + assert!(region.console_url().ends_with("/code/console")); + assert!( + region + .cookie_domains() + .iter() + .all(|host| host.ends_with(suffix)) + ); + } + } +} diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index 0cdfad9cfa..c5ed5b8305 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -11,9 +11,9 @@ use reqwest::Client; use super::desktop_token::KimiDesktopAuthToken; use super::{ - KIMI_COOKIE_DOMAINS, KIMI_SUBSCRIPTION_STATS_URL, KIMI_SUBSCRIPTION_URL, KIMI_WEB_USAGE_URL, - KimiProvider, KimiSubscriptionResponse, KimiSubscriptionStatsResponse, KimiWebUsageResponse, - apply_subscription_windows, kimi_web_post, + KIMI_SUBSCRIPTION_SERVICE, KIMI_SUBSCRIPTION_STATS_SERVICE, KIMI_WEB_USAGE_SERVICE, + KimiProvider, KimiRegion, KimiSubscriptionResponse, KimiSubscriptionStatsResponse, + KimiWebUsageResponse, apply_subscription_windows, kimi_web_post, }; use crate::browser::cookies::get_cookie_header; use crate::core::{ProviderError, ProviderId, UsageSnapshot}; @@ -26,21 +26,32 @@ pub(crate) fn cookie_source() -> String { .to_string() } -/// Upstream `KimiBrowserImportPolicy.allowsImport`: everything but `off`. +/// Upstream `KimiBrowserImportPolicy.allowsImport`: automatic discovery is +/// allowed only when the user selected the automatic source. fn browser_import_allowed(cookie_source: &str) -> bool { - !cookie_source.eq_ignore_ascii_case("off") + cookie_source.eq_ignore_ascii_case("auto") || cookie_source.eq_ignore_ascii_case("browser") +} + +fn browser_import_error(cookie_source: &str) -> ProviderError { + let message = if cookie_source.eq_ignore_ascii_case("manual") { + "Kimi cookie source is Manual; provide a valid manual cookie header." + } else { + "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + }; + ProviderError::Other(message.into()) } /// Web auth token chain for both the web fetch and the Code-API enrichment /// (upstream `KimiWebEnrichmentTokenResolver.resolve`): /// 1. Manual cookie header (its `kimi-auth`/auth cookie), source-independent. -/// 2. Kimi Desktop session token (skipped when cookie source is `off`). -/// 3. Browser cookie import (skipped when cookie source is `off`). -pub(crate) fn web_auth_tokens(manual_header: Option<&str>) -> Vec { +/// 2. Kimi Desktop session token (automatic source only). +/// 3. Browser cookie import (automatic source only). +pub(crate) fn web_auth_tokens(manual_header: Option<&str>, region: KimiRegion) -> Vec { resolve_web_tokens(WebTokenInput { manual_header, cookie_source: &cookie_source(), - desktop_token: KimiDesktopAuthToken::load, + region, + desktop_token: KimiDesktopAuthToken::load_for_region, browser_token: browser_auth_token, }) .into_iter() @@ -51,8 +62,9 @@ pub(crate) fn web_auth_tokens(manual_header: Option<&str>) -> Vec { struct WebTokenInput<'a> { manual_header: Option<&'a str>, cookie_source: &'a str, - desktop_token: fn() -> Option, - browser_token: fn() -> Option, + region: KimiRegion, + desktop_token: fn(KimiRegion) -> Option, + browser_token: fn(KimiRegion) -> Option, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -83,7 +95,7 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { let mut candidates = Vec::new(); let mut seen = std::collections::HashSet::new(); - if let Some(token) = (input.desktop_token)() + if let Some(token) = (input.desktop_token)(input.region) && seen.insert(token.clone()) { candidates.push(WebTokenCandidate { @@ -91,7 +103,7 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { source: WebTokenSource::Desktop, }); } - if let Some(token) = (input.browser_token)() + if let Some(token) = (input.browser_token)(input.region) && seen.insert(token.clone()) { candidates.push(WebTokenCandidate { @@ -104,8 +116,9 @@ fn resolve_web_tokens(input: WebTokenInput) -> Vec { /// Browser import only: the first usable `kimi-auth`-class token from any of /// the registered Kimi cookie domains. -fn browser_auth_token() -> Option { - KIMI_COOKIE_DOMAINS +fn browser_auth_token(region: KimiRegion) -> Option { + region + .cookie_domains() .iter() .find_map(|domain| { get_cookie_header(domain) @@ -118,6 +131,7 @@ fn browser_auth_token() -> Option { /// Fetch usage via Kimi web API (weekly quota + rate limit + subscription). pub(crate) async fn fetch_via_web( cookie_header: Option<&str>, + region: KimiRegion, ) -> Result { let source = cookie_source(); if let Some(token) = @@ -126,14 +140,11 @@ pub(crate) async fn fetch_via_web( // An explicit manual credential is authoritative. A rejected manual // token must not silently switch accounts underneath the user. let client = client()?; - return fetch_via_web_token(&client, &token).await; + return fetch_via_web_token(&client, &token, region).await; } if !browser_import_allowed(&source) { - return Err(ProviderError::Other( - "Kimi cookie source is Off; provide a manual cookie header or enable browser import." - .into(), - )); + return Err(browser_import_error(&source)); } let client = client()?; @@ -142,20 +153,20 @@ pub(crate) async fn fetch_via_web( // Read and try the desktop session first. Browser cookies are intentionally // read only after the server rejects this automatic session, so a healthy // desktop account never causes another credential store to be touched. - if let Some(token) = KimiDesktopAuthToken::load() + if let Some(token) = KimiDesktopAuthToken::load_for_region(region) && seen.insert(token.clone()) { - match fetch_via_web_token(&client, &token).await { + match fetch_via_web_token(&client, &token, region).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(error), } } - if let Some(token) = browser_auth_token() + if let Some(token) = browser_auth_token(region) && seen.insert(token.clone()) { - match fetch_via_web_token(&client, &token).await { + match fetch_via_web_token(&client, &token, region).await { Ok(usage) => return Ok(usage), Err(ProviderError::AuthRequired) => {} Err(error) => return Err(error), @@ -175,10 +186,13 @@ fn client() -> Result { async fn fetch_via_web_token( client: &reqwest::Client, token: &str, + region: KimiRegion, ) -> Result { + let usage_url = region.web_api_url(KIMI_WEB_USAGE_SERVICE); let resp = kimi_web_post( client, - KIMI_WEB_USAGE_URL, + &usage_url, + region, token, serde_json::json!({ "scope": ["FEATURE_CODING"] }), ) @@ -197,7 +211,7 @@ async fn fetch_via_web_token( .await .map_err(|e| ProviderError::Parse(e.to_string()))?; - let (subscription, plan_name) = fetch_subscription_details(client, token).await; + let (subscription, plan_name) = fetch_subscription_details(client, token, region).await; snapshot_from_web_usage_response_with_plan(usage, subscription, plan_name) } @@ -207,16 +221,17 @@ const SUBSCRIPTION_ENRICHMENT_TIMEOUT: std::time::Duration = std::time::Duration async fn fetch_subscription_details( client: &reqwest::Client, token: &str, + region: KimiRegion, ) -> (Option, Option) { // The quota statistics and the optional title are independent. Keep a // completed statistics response when the plan endpoint is slow or absent. let stats = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_for_enrichment(client, token), + fetch_subscription_for_enrichment(client, token, region), ); let plan = tokio::time::timeout( SUBSCRIPTION_ENRICHMENT_TIMEOUT, - fetch_subscription_plan(client, token), + fetch_subscription_plan(client, token, region), ); let (stats, plan) = tokio::join!(stats, plan); (stats.ok().flatten(), plan.ok().flatten()) @@ -259,8 +274,13 @@ fn snapshot_from_web_usage_response_with_plan( Ok(usage) } -pub(super) async fn fetch_subscription_plan(client: &Client, token: &str) -> Option { - match kimi_web_post(client, KIMI_SUBSCRIPTION_URL, token, serde_json::json!({})).await { +pub(super) async fn fetch_subscription_plan( + client: &Client, + token: &str, + region: KimiRegion, +) -> Option { + let url = region.web_api_url(KIMI_SUBSCRIPTION_SERVICE); + match kimi_web_post(client, &url, region, token, serde_json::json!({})).await { Ok(response) if response.status().is_success() => response .json::() .await @@ -275,8 +295,9 @@ pub(super) async fn fetch_subscription_plan(client: &Client, token: &str) -> Opt pub(super) async fn fetch_subscription_for_enrichment( client: &Client, token: &str, + region: KimiRegion, ) -> Option { - fetch_subscription_for_enrichment_result(client, token) + fetch_subscription_for_enrichment_result(client, token, region) .await .ok() .flatten() @@ -285,15 +306,10 @@ pub(super) async fn fetch_subscription_for_enrichment( pub(super) async fn fetch_subscription_for_enrichment_result( client: &Client, token: &str, + region: KimiRegion, ) -> Result, ProviderError> { - match kimi_web_post( - client, - KIMI_SUBSCRIPTION_STATS_URL, - token, - serde_json::json!({}), - ) - .await - { + let url = region.web_api_url(KIMI_SUBSCRIPTION_STATS_SERVICE); + match kimi_web_post(client, &url, region, token, serde_json::json!({})).await { Ok(response) if response.status().is_success() => response .json() .await @@ -311,33 +327,34 @@ pub(super) async fn fetch_subscription_for_enrichment_result( mod tests { use super::*; - fn static_desktop() -> Option { + fn static_desktop(_: KimiRegion) -> Option { Some("desktop-token".to_string()) } - fn static_browser() -> Option { + fn static_browser(_: KimiRegion) -> Option { Some("browser-token".to_string()) } - fn no_token() -> Option { + fn no_token(_: KimiRegion) -> Option { None } fn input<'a>( manual_header: Option<&'a str>, cookie_source: &'a str, - desktop_token: fn() -> Option, - browser_token: fn() -> Option, + desktop_token: fn(KimiRegion) -> Option, + browser_token: fn(KimiRegion) -> Option, ) -> WebTokenInput<'a> { WebTokenInput { manual_header, cookie_source, + region: KimiRegion::China, desktop_token, browser_token, } } - fn duplicate_browser() -> Option { + fn duplicate_browser(_: KimiRegion) -> Option { Some("desktop-token".to_string()) } @@ -369,22 +386,17 @@ mod tests { } #[test] - fn cookie_source_off_blocks_desktop_and_browser_but_not_manual() { - assert_eq!( - resolve_web_tokens(input(None, "off", static_desktop, static_browser)), - Vec::new() - ); - assert_eq!( - resolve_web_tokens(input(None, "off", no_token, static_browser)), - Vec::new() - ); - assert_eq!( - resolve_web_tokens(input(Some("kimi-auth=manual"), "off", no_token, no_token)), - vec![WebTokenCandidate { - token: "manual".to_string(), - source: WebTokenSource::Manual, - }] - ); + fn off_and_manual_sources_block_automatic_discovery() { + for source in ["off", "manual"] { + assert_eq!( + resolve_web_tokens(input(None, source, static_desktop, static_browser)), + Vec::new() + ); + assert_eq!( + resolve_web_tokens(input(Some("not-a-token"), source, no_token, static_browser)), + Vec::new() + ); + } } #[test] @@ -400,15 +412,18 @@ mod tests { } #[test] - fn manual_default_source_still_allows_desktop_token() { - // Upstream: desktop-session token applies for any non-off source; - // the local default ("manual") must keep desktop sessions working. - let candidates = resolve_web_tokens(input(None, "manual", static_desktop, no_token)); + fn explicit_manual_token_stays_authoritative() { + let candidates = resolve_web_tokens(input( + Some("kimi-auth=manual-token"), + "manual", + static_desktop, + static_browser, + )); assert_eq!( candidates, vec![WebTokenCandidate { - token: "desktop-token".to_string(), - source: WebTokenSource::Desktop, + token: "manual-token".to_string(), + source: WebTokenSource::Manual, }] ); } @@ -429,7 +444,27 @@ mod tests { fn browser_import_gate_is_case_insensitive() { assert!(!browser_import_allowed("OFF")); assert!(browser_import_allowed("browser")); - assert!(browser_import_allowed("manual")); + assert!(browser_import_allowed("AUTO")); + assert!(!browser_import_allowed("manual")); + } + + #[test] + fn browser_import_error_matches_rejected_cookie_source() { + assert!(matches!( + browser_import_error("manual"), + ProviderError::Other(message) + if message == "Kimi cookie source is Manual; provide a valid manual cookie header." + )); + assert!(matches!( + browser_import_error("off"), + ProviderError::Other(message) + if message == "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + )); + assert!(matches!( + browser_import_error("unexpected"), + ProviderError::Other(message) + if message == "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + )); } #[test] diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index de9d265f6e..0b4cdc2bff 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -26,7 +26,6 @@ pub mod coderabbit; pub mod codex; pub mod commandcode; pub mod copilot; -pub mod crof; pub mod crossmodel; pub mod cursor; pub mod deepgram; @@ -109,7 +108,6 @@ pub use coderabbit::CodeRabbitProvider; pub use codex::CodexProvider; pub use commandcode::CommandCodeProvider; pub use copilot::CopilotProvider; -pub use crof::CrofProvider; pub use crossmodel::CrossModelProvider; pub use cursor::CursorProvider; pub use deepgram::DeepgramProvider; @@ -128,7 +126,7 @@ pub use huggingface::HuggingFaceProvider; pub use infini::InfiniProvider; pub use jetbrains::JetBrainsProvider; pub use kilo::KiloProvider; -pub use kimi::KimiProvider; +pub use kimi::{KimiProvider, KimiRegion}; pub use kimik2::KimiK2Provider; pub use kiro::KiroProvider; pub use litellm::LiteLLMProvider; diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 49a2f7ec34..24e84b2473 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -64,6 +64,7 @@ impl From for crate::spend_contract::LocalTokenHistorySummary { total_tokens: report.total_tokens.unwrap_or(0), session_count: report.session_count, coverage: report.coverage, + cost_estimate: Default::default(), } } } diff --git a/rust/src/settings.rs b/rust/src/settings.rs index 9863f194ce..70e18d32ad 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -136,6 +136,16 @@ pub struct Settings { #[serde(default)] pub tray_icon_mode: TrayIconMode, + /// Optional preferred provider for the upper row of a stacked tray icon. + /// Stale or disabled values are retained and ignored until eligible again. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stacked_tray_top_provider: Option, + + /// Optional preferred provider for the lower row of a stacked tray icon. + /// Stale or duplicate values fall back to the next eligible provider. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stacked_tray_bottom_provider: Option, + /// Show provider icons in the merged switcher UI #[serde(default = "default_true")] pub switcher_shows_icons: bool, @@ -541,6 +551,8 @@ impl Default for Settings { provider_usage_thresholds: HashMap::new(), merge_tray_icons: false, // Show single provider by default tray_icon_mode: TrayIconMode::default(), // Single icon by default + stacked_tray_top_provider: None, + stacked_tray_bottom_provider: None, switcher_shows_icons: true, menu_bar_shows_highest_usage: false, menu_bar_shows_percent: false, @@ -933,12 +945,17 @@ impl Settings { self.provider_configs.entry(id).or_default() } - /// Cookie source for `id`, or the default `"manual"` if unset. + /// Cookie source for `id`. Kimi follows upstream's automatic default; + /// providers with no specific default retain the legacy manual default. pub fn cookie_source(&self, id: ProviderId) -> &str { self.provider_configs .get(&id) .and_then(|c| c.cookie_source.as_deref()) - .unwrap_or(DEFAULT_COOKIE_SOURCE) + .unwrap_or(if id == ProviderId::Kimi { + "auto" + } else { + DEFAULT_COOKIE_SOURCE + }) } pub fn set_cookie_source(&mut self, id: ProviderId, source: impl Into) { diff --git a/rust/src/settings/api_keys.rs b/rust/src/settings/api_keys.rs index c2adf3565c..9fa7544ee0 100644 --- a/rust/src/settings/api_keys.rs +++ b/rust/src/settings/api_keys.rs @@ -439,15 +439,6 @@ pub fn get_api_key_providers() -> Vec { config_file_path: None, dashboard_url: Some("https://console.volcengine.com/ark/region:ark+cn-beijing/usage"), }, - ProviderConfigInfo { - id: ProviderId::Crof, - name: "Crof", - requires_api_key: true, - api_key_env_var: Some("CROF_API_KEY"), - api_key_help: Some("Get your API key from Crof."), - config_file_path: None, - dashboard_url: Some("https://crof.ai"), - }, ProviderConfigInfo { id: ProviderId::StepFun, name: "StepFun", diff --git a/rust/src/settings/raw.rs b/rust/src/settings/raw.rs index 850c90e960..8fec108d41 100644 --- a/rust/src/settings/raw.rs +++ b/rust/src/settings/raw.rs @@ -1,4 +1,86 @@ use super::*; +use serde::Deserializer; +use serde::de::{IgnoredAny, MapAccess, Visitor}; +use std::fmt; + +fn canonical_provider_id(raw: &str) -> Option { + ProviderId::from_cli_name(raw).map(|provider| provider.cli_name().to_string()) +} + +fn canonicalize_provider_id_list(ids: impl IntoIterator) -> Vec { + let mut seen = HashSet::new(); + ids.into_iter() + .filter_map(|raw| canonical_provider_id(&raw)) + .filter(|canonical| seen.insert(canonical.clone())) + .collect() +} + +fn canonicalize_provider_metrics( + metrics: HashMap, +) -> HashMap { + let mut entries = metrics + .into_iter() + .filter_map(|(raw, preference)| { + let canonical = canonical_provider_id(&raw)?; + let canonical_spelling = raw.eq_ignore_ascii_case(&canonical); + Some((canonical, canonical_spelling, raw, preference)) + }) + .collect::>(); + + // HashMap iteration order is unstable. Sort before resolving aliases so a + // canonical spelling always wins and alias-only collisions are repeatable. + entries.sort_by(|left, right| { + left.0 + .cmp(&right.0) + .then_with(|| left.1.cmp(&right.1)) + .then_with(|| { + left.2 + .to_ascii_lowercase() + .cmp(&right.2.to_ascii_lowercase()) + }) + .then_with(|| left.2.cmp(&right.2)) + }); + + let mut canonical = HashMap::with_capacity(entries.len()); + for (provider_id, _, _, preference) in entries { + canonical.insert(provider_id, preference); + } + canonical +} + +fn deserialize_provider_configs<'de, D>( + deserializer: D, +) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + struct ProviderConfigsVisitor; + + impl<'de> Visitor<'de> for ProviderConfigsVisitor { + type Value = HashMap; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a map of provider IDs to provider settings") + } + + fn visit_map(self, mut map: M) -> Result + where + M: MapAccess<'de>, + { + let mut configs = HashMap::with_capacity(map.size_hint().unwrap_or(0)); + while let Some(key) = map.next_key::()? { + if let Some(provider_id) = ProviderId::from_cli_name(&key) { + configs.insert(provider_id, map.next_value()?); + } else { + map.next_value::()?; + } + } + Ok(configs) + } + } + + deserializer.deserialize_map(ProviderConfigsVisitor) +} /// Raw on-disk shape of [`Settings`] used purely for deserialization. /// @@ -35,6 +117,8 @@ pub(super) struct RawSettings { provider_usage_thresholds: HashMap, merge_tray_icons: bool, tray_icon_mode: TrayIconMode, + stacked_tray_top_provider: Option, + stacked_tray_bottom_provider: Option, #[serde(default = "default_true")] switcher_shows_icons: bool, menu_bar_shows_highest_usage: bool, @@ -52,6 +136,7 @@ pub(super) struct RawSettings { show_all_token_accounts_in_menu: bool, // ── New unified per-provider map ───────────────────────────────── + #[serde(default, deserialize_with = "deserialize_provider_configs")] provider_configs: HashMap, // ── Legacy flat per-provider fields (migrated on load) ─────────── @@ -206,6 +291,8 @@ impl Default for RawSettings { provider_usage_thresholds: HashMap::new(), merge_tray_icons: s.merge_tray_icons, tray_icon_mode: s.tray_icon_mode, + stacked_tray_top_provider: s.stacked_tray_top_provider, + stacked_tray_bottom_provider: s.stacked_tray_bottom_provider, switcher_shows_icons: s.switcher_shows_icons, menu_bar_shows_highest_usage: s.menu_bar_shows_highest_usage, menu_bar_shows_percent: s.menu_bar_shows_percent, @@ -515,7 +602,11 @@ impl From for Settings { }; Settings { - enabled_providers: raw.enabled_providers, + enabled_providers: raw + .enabled_providers + .into_iter() + .filter_map(|provider_id| canonical_provider_id(&provider_id)) + .collect(), refresh_interval_secs: raw.refresh_interval_secs, adaptive_refresh: raw.adaptive_refresh, refresh_all_providers_on_menu_open: raw.refresh_all_providers_on_menu_open, @@ -533,6 +624,12 @@ impl From for Settings { ), merge_tray_icons: raw.merge_tray_icons, tray_icon_mode: raw.tray_icon_mode, + stacked_tray_top_provider: raw + .stacked_tray_top_provider + .and_then(|provider_id| canonical_provider_id(&provider_id)), + stacked_tray_bottom_provider: raw + .stacked_tray_bottom_provider + .and_then(|provider_id| canonical_provider_id(&provider_id)), switcher_shows_icons: raw.switcher_shows_icons, menu_bar_shows_highest_usage: raw.menu_bar_shows_highest_usage, menu_bar_shows_percent: raw.menu_bar_shows_percent, @@ -549,7 +646,7 @@ impl From for Settings { disable_keychain_access: raw.disable_keychain_access, hide_personal_info: raw.hide_personal_info, update_channel: raw.update_channel, - provider_metrics: raw.provider_metrics, + provider_metrics: canonicalize_provider_metrics(raw.provider_metrics), provider_order: if raw.provider_order.is_empty() { Vec::new() } else { @@ -578,7 +675,7 @@ impl From for Settings { float_bar_orientation: normalize_float_bar_orientation(&raw.float_bar_orientation), float_bar_style: normalize_float_bar_style(&raw.float_bar_style), float_bar_click_through: raw.float_bar_click_through, - float_bar_provider_ids: raw.float_bar_provider_ids, + float_bar_provider_ids: canonicalize_provider_id_list(raw.float_bar_provider_ids), float_bar_dark_text: raw.float_bar_dark_text, float_bar_show_reset_inline: raw.float_bar_show_reset_inline, float_bar_show_cost: raw.float_bar_show_cost, diff --git a/rust/src/settings/tests.rs b/rust/src/settings/tests.rs index a6db4aa05b..6ffec91c6a 100644 --- a/rust/src/settings/tests.rs +++ b/rust/src/settings/tests.rs @@ -31,6 +31,13 @@ fn test_settings_default() { ); } +#[test] +fn kimi_cookie_source_defaults_to_automatic_discovery() { + let settings = Settings::default(); + assert_eq!(settings.cookie_source(ProviderId::Kimi), "auto"); + assert_eq!(settings.cookie_source(ProviderId::Claude), "manual"); +} + #[test] fn overview_layout_defaults_to_compact_and_round_trips() { let defaulted: Settings = serde_json::from_str(r#"{ "enabled_providers": [] }"#) @@ -771,6 +778,38 @@ fn test_settings_with_utf8_bom_parses_perprovider_tray_mode() { assert_eq!(settings.tray_icon_mode, TrayIconMode::PerProvider); } +#[test] +fn stacked_tray_mode_preserves_provider_preferences() { + let json = r#"{ + "tray_icon_mode": "stacked", + "stacked_tray_top_provider": "claude", + "stacked_tray_bottom_provider": "codex" + }"#; + + let settings: Settings = serde_json::from_str(json).unwrap(); + + assert_eq!(settings.tray_icon_mode, TrayIconMode::Stacked); + assert_eq!( + settings.stacked_tray_top_provider.as_deref(), + Some("claude") + ); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("codex") + ); + + let saved = serde_json::to_string(&settings).unwrap(); + let reloaded: Settings = serde_json::from_str(&saved).unwrap(); + assert_eq!( + reloaded.stacked_tray_top_provider.as_deref(), + Some("claude") + ); + assert_eq!( + reloaded.stacked_tray_bottom_provider.as_deref(), + Some("codex") + ); +} + #[test] fn test_language_serde_serialization() { // Test that Language serializes to lowercase string @@ -1045,6 +1084,108 @@ fn test_new_format_provider_configs_only() { assert_eq!(settings.api_region(ProviderId::Zai), "global"); } +#[test] +fn retired_provider_config_is_ignored_until_explicit_save() { + let original = r#"{ + "enabled_providers": ["codex", "crof"], + "refresh_interval_secs": 300, + "provider_metrics": { "codex": "weekly", "crof": "session" }, + "float_bar_provider_ids": ["codex", "crof"], + "stacked_tray_top_provider": "crof", + "stacked_tray_bottom_provider": "crof", + "provider_configs": { + "crof": { "api_token": "retired-fixture-key" }, + "codex": { "cookie_source": "manual", "openai_web_extras": false }, + "alibaba": { "api_region": "cn", "manual_cookie_header": "ali=PLACEHOLDER" } + } + }"#; + let original_bytes = original.as_bytes().to_vec(); + + let settings: Settings = + serde_json::from_str(original).expect("load settings with retired key"); + + assert_eq!(original.as_bytes(), original_bytes); + assert_eq!(settings.cookie_source(ProviderId::Codex), "manual"); + assert!(!settings.openai_web_extras(ProviderId::Codex)); + assert_eq!( + settings.enabled_providers, + HashSet::from(["codex".to_string()]) + ); + assert_eq!(settings.provider_metrics.len(), 1); + assert_eq!(settings.float_bar_provider_ids, ["codex"]); + assert_eq!(settings.stacked_tray_top_provider, None); + assert_eq!(settings.stacked_tray_bottom_provider, None); + assert_eq!(settings.api_region(ProviderId::Alibaba), "cn"); + assert_eq!( + settings.manual_cookie_header(ProviderId::Alibaba), + "ali=PLACEHOLDER" + ); + + let saved = serde_json::to_string(&settings).expect("serialize sanitized settings"); + let saved_value: serde_json::Value = serde_json::from_str(&saved).unwrap(); + let saved_configs = saved_value["provider_configs"].as_object().unwrap(); + assert!(!saved_configs.contains_key("crof")); + assert!(saved_configs.contains_key("codex")); + assert!(saved_configs.contains_key("alibaba")); + assert!( + !saved.contains("\"crof\""), + "saved settings retained Crof: {saved}" + ); +} + +#[test] +fn provider_aliases_are_canonicalized_at_the_load_boundary() { + let settings: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["openai", "ClAuDe", "not-a-provider"], + "provider_metrics": { + "openai": "weekly", + "CoDeX": "session", + "not-a-provider": "weekly" + }, + "float_bar_provider_ids": ["OPENAI", "codex", "ClAuDe", "unknown"], + "stacked_tray_top_provider": "OPENAI", + "stacked_tray_bottom_provider": "ClAuDe" + }"#, + ) + .expect("load settings containing provider aliases"); + + assert_eq!( + settings.enabled_providers, + HashSet::from(["claude".to_string(), "codex".to_string()]) + ); + assert_eq!( + settings.provider_metrics.get("codex"), + Some(&MetricPreference::Session) + ); + assert_eq!(settings.provider_metrics.len(), 1); + assert_eq!(settings.float_bar_provider_ids, ["codex", "claude"]); + assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); + assert_eq!( + settings.stacked_tray_bottom_provider.as_deref(), + Some("claude") + ); +} + +#[test] +fn stacked_preferences_preserve_known_disabled_providers() { + let settings: Settings = serde_json::from_str( + r#"{ + "enabled_providers": ["claude"], + "stacked_tray_top_provider": "OPENAI", + "stacked_tray_bottom_provider": "not-a-provider" + }"#, + ) + .expect("load stacked preferences independently of enablement"); + + assert_eq!(settings.stacked_tray_top_provider.as_deref(), Some("codex")); + assert_eq!(settings.stacked_tray_bottom_provider, None); + assert_eq!( + settings.enabled_providers, + HashSet::from(["claude".to_string()]) + ); +} + /// Default `Settings` should serialize WITHOUT a `provider_configs` /// field (empty map skipped). #[test] diff --git a/rust/src/settings/types.rs b/rust/src/settings/types.rs index 6113969b15..062338076e 100644 --- a/rust/src/settings/types.rs +++ b/rust/src/settings/types.rs @@ -247,6 +247,8 @@ pub enum TrayIconMode { Single, /// One tray icon per enabled provider PerProvider, + /// One tray icon with the selected metrics for two providers stacked vertically + Stacked, } impl TrayIconMode { @@ -255,6 +257,7 @@ impl TrayIconMode { match self { TrayIconMode::Single => "Single Icon", TrayIconMode::PerProvider => "Per Provider", + TrayIconMode::Stacked => "Stacked Providers", } } @@ -263,6 +266,7 @@ impl TrayIconMode { match self { TrayIconMode::Single => "Show one tray icon for all providers", TrayIconMode::PerProvider => "Show a separate tray icon for each enabled provider", + TrayIconMode::Stacked => "Show two providers as stacked usage meters", } } } diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index af9cdfef34..722acec630 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -79,24 +79,89 @@ pub enum LocalHistoryCoverage { Unavailable, } -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct LocalTokenHistorySummary { pub total_tokens: u64, pub session_count: usize, pub coverage: LocalHistoryCoverage, + pub cost_estimate: LocalCostEstimate, +} + +impl LocalTokenHistorySummary { + /// Return a complete list-price total only when both the history scan and + /// pricing coverage are complete. A complete scan with no token usage is + /// a known zero even though there were no requests to price. + pub fn total_usd(&self) -> Option { + if self.coverage != LocalHistoryCoverage::Complete { + return None; + } + if self.total_tokens == 0 { + return Some(0.0); + } + self.cost_estimate.complete_total_usd() + } +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LocalCostEstimate { + /// Sum of requests whose models have public API list prices. This remains + /// a subtotal when one or more requests are unpriced. + pub known_subtotal_usd: Option, + pub coverage: CostCoverageCounts, +} + +impl LocalCostEstimate { + fn complete_total_usd(&self) -> Option { + if self.coverage.unpriced == 0 && self.coverage.unmetered == 0 { + self.known_subtotal_usd + } else { + None + } + } + + pub(crate) fn record_list_price(&mut self, cost: Option) { + let Some(cost) = cost.filter(|value| value.is_finite() && *value >= 0.0) else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + return; + }; + let next = self.known_subtotal_usd.unwrap_or(0.0) + cost; + if next.is_finite() { + self.known_subtotal_usd = Some(next); + self.coverage.estimated = self.coverage.estimated.saturating_add(1); + } else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + } + } } pub fn local_token_history_json( provider: &str, - history: LocalTokenHistorySummary, + history: &LocalTokenHistorySummary, days: u32, ) -> serde_json::Value { let complete = history.coverage == LocalHistoryCoverage::Complete; + let total_usd = history.total_usd(); + let known_subtotal_usd = history.cost_estimate.known_subtotal_usd; + let note = if total_usd.is_some() { + "Local token history estimated at public API list prices; not billed spend" + } else if known_subtotal_usd.is_some() && !complete { + "Known public API list-price subtotal; local history is incomplete" + } else if known_subtotal_usd.is_some() { + "Known public API list-price subtotal; some local requests are unpriced" + } else { + "Local token history; dollar costs unavailable" + }; serde_json::json!({ "provider": provider, "supported": true, "days_scanned": days, - "cost": {"total_usd": serde_json::Value::Null, "currency": serde_json::Value::Null}, + "cost": { + "total_usd": total_usd, + "known_subtotal_usd": known_subtotal_usd, + "currency": total_usd.or(known_subtotal_usd).map(|_| "USD"), + "pricingCoverage": &history.cost_estimate.coverage, + }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, "sessions_count": complete.then_some(history.session_count), @@ -106,7 +171,7 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": "Local token history; dollar costs unavailable" + "note": note, }) } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] diff --git a/rust/src/spend_contract/tests.rs b/rust/src/spend_contract/tests.rs index fb74ff09de..ac32e2e5eb 100644 --- a/rust/src/spend_contract/tests.rs +++ b/rust/src/spend_contract/tests.rs @@ -1,5 +1,58 @@ use super::*; +#[test] +fn local_history_total_requires_complete_scan_and_pricing() { + let priced = LocalCostEstimate { + known_subtotal_usd: Some(1.25), + coverage: CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + }; + let partial_history = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Partial, + cost_estimate: priced.clone(), + }; + assert_eq!(partial_history.total_usd(), None); + assert_eq!(partial_history.cost_estimate.known_subtotal_usd, Some(1.25)); + + let mixed_pricing = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(1.25), + coverage: CostCoverageCounts { + estimated: 1, + unpriced: 1, + ..Default::default() + }, + }, + }; + assert_eq!(mixed_pricing.total_usd(), None); + assert_eq!(mixed_pricing.cost_estimate.known_subtotal_usd, Some(1.25)); + + let complete = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: priced, + }; + assert_eq!(complete.total_usd(), Some(1.25)); +} + +#[test] +fn complete_empty_local_history_has_a_known_zero_total() { + let history = LocalTokenHistorySummary { + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + }; + + assert_eq!(history.total_usd(), Some(0.0)); +} + #[test] fn coverage_ratio_counts_estimated_as_covered() { let coverage = CostCoverageCounts { diff --git a/rust/src/tray/mod.rs b/rust/src/tray/mod.rs index bf4c52ef88..0e2e80afc7 100755 --- a/rust/src/tray/mod.rs +++ b/rust/src/tray/mod.rs @@ -6,4 +6,6 @@ pub mod icon; pub mod render; pub use icon::LoadingPattern; -pub use render::{TRAY_ICON_SIZE, render_bar_icon_rgba, render_percent_icon_rgba}; +pub use render::{ + TRAY_ICON_SIZE, render_bar_icon_rgba, render_percent_icon_rgba, render_stacked_bar_icon_rgba, +}; diff --git a/rust/src/tray/render.rs b/rust/src/tray/render.rs index 5b486a5b6b..d0d464e7b6 100644 --- a/rust/src/tray/render.rs +++ b/rust/src/tray/render.rs @@ -10,6 +10,62 @@ use super::icon::UsageLevel; /// Side length of the generated tray icon in pixels. pub const TRAY_ICON_SIZE: u32 = 32; +const ICON_INSET: u32 = 2; +const BAR_LEFT: u32 = 4; +const BAR_RIGHT: u32 = TRAY_ICON_SIZE - 4; +const ICON_BACKGROUND_RGB: [u8; 3] = [60, 60, 70]; +const BAR_BACKGROUND: Rgba = Rgba([80, 80, 90, 255]); + +fn new_icon_canvas(has_error: bool) -> RgbaImage { + let mut image: RgbaImage = ImageBuffer::new(TRAY_ICON_SIZE, TRAY_ICON_SIZE); + let background = Rgba([ + ICON_BACKGROUND_RGB[0], + ICON_BACKGROUND_RGB[1], + ICON_BACKGROUND_RGB[2], + if has_error { 180 } else { 255 }, + ]); + for y in ICON_INSET..TRAY_ICON_SIZE - ICON_INSET { + for x in ICON_INSET..TRAY_ICON_SIZE - ICON_INSET { + image.put_pixel(x, y, background); + } + } + image +} + +fn usage_color(percent: f64, has_error: bool) -> Rgba { + let (r, g, b) = UsageLevel::from_percent(percent).color(); + if has_error { + #[allow( + clippy::cast_possible_truncation, + reason = "mean of three u8 channels is bounded to 0..=255" + )] + let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; + Rgba([gray, gray, gray, 255]) + } else { + Rgba([r, g, b, 255]) + } +} + +fn draw_bar_row(image: &mut RgbaImage, y_start: u32, y_end: u32, percent: f64, has_error: bool) { + let bar_width = BAR_RIGHT - BAR_LEFT; + #[allow( + clippy::cast_possible_truncation, + reason = "percent is clamped to 0..=100 and scaled to a 24-pixel meter" + )] + let fill = ((percent.clamp(0.0, 100.0) / 100.0) * bar_width as f64) as u32; + let fill_end = (BAR_LEFT + fill).min(BAR_RIGHT); + let color = usage_color(percent, has_error); + + for y in y_start..y_end { + for x in BAR_LEFT..BAR_RIGHT { + image.put_pixel(x, y, BAR_BACKGROUND); + } + for x in BAR_LEFT..fill_end { + image.put_pixel(x, y, color); + } + } +} + /// Render a usage-bar tray icon as raw RGBA bytes. /// /// - `session_percent`: primary bar fill (0–100), colour-coded by [`UsageLevel`] @@ -24,90 +80,41 @@ pub fn render_bar_icon_rgba( weekly_percent: Option, has_error: bool, ) -> (Vec, u32, u32) { - const SZ: u32 = TRAY_ICON_SIZE; - let mut img: RgbaImage = ImageBuffer::new(SZ, SZ); - - for pixel in img.pixels_mut() { - *pixel = Rgba([0, 0, 0, 0]); - } - - let bg_alpha: u8 = if has_error { 180 } else { 255 }; - let bg_color = Rgba([60, 60, 70, bg_alpha]); - for y in 2..SZ - 2 { - for x in 2..SZ - 2 { - img.put_pixel(x, y, bg_color); - } - } - - let color_for = |percent: f64| -> (u8, u8, u8) { - let (r, g, b) = UsageLevel::from_percent(percent).color(); - if has_error { - // Average of three u8 colour channels: sum ≤ 765, so /3 ≤ 255 fits u8. - #[allow( - clippy::cast_possible_truncation, - reason = "mean of three u8 channels; r+g+b ≤ 765, divided by 3 is ≤ 255 and fits u8" - )] - let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; - (gray, gray, gray) - } else { - (r, g, b) - } - }; - - let bar_left = 4u32; - let bar_right = SZ - 4; - let bar_width = bar_right - bar_left; - - // pct is clamped to 0–100, scaled by bar_width (≤ SZ = 32), so the result fits u32. - #[allow( - clippy::cast_possible_truncation, - reason = "pct clamped to 0–100 and scaled by bar_width ≤ 32; result is a small pixel count that fits u32" - )] - let fill_px = |pct: f64| ((pct.clamp(0.0, 100.0) / 100.0) * bar_width as f64) as u32; - - let mut draw_bar = |y_start: u32, y_end: u32, pct: f64| { - let (r, g, b) = color_for(pct); - let fill_end = (bar_left + fill_px(pct)).min(bar_right); - for y in y_start..y_end { - for x in bar_left..bar_right { - img.put_pixel(x, y, Rgba([80, 80, 90, 255])); - } - } - for y in y_start..y_end { - for x in bar_left..fill_end { - img.put_pixel(x, y, Rgba([r, g, b, 255])); - } - } - }; + let mut image = new_icon_canvas(has_error); match weekly_percent { Some(weekly) => { - draw_bar(8, 15, session_percent); // session bar (top, thicker) - draw_bar(18, 23, weekly); // weekly bar (bottom, thinner) + draw_bar_row(&mut image, 8, 15, session_percent, has_error); + draw_bar_row(&mut image, 18, 23, weekly, has_error); } None => { - draw_bar(10, 22, session_percent); // single thick bar (centred) + draw_bar_row(&mut image, 10, 22, session_percent, has_error); } } - (img.into_raw(), SZ, SZ) + (image.into_raw(), TRAY_ICON_SIZE, TRAY_ICON_SIZE) +} + +/// Render two providers as equally prominent stacked usage meters. +/// +/// Unlike [`render_bar_icon_rgba`], both rows represent the selected metric +/// for separate providers. The upper and lower rows therefore use equal +/// height so neither provider is presented as a secondary quota window. +pub fn render_stacked_bar_icon_rgba( + top_percent: f64, + bottom_percent: f64, + has_error: bool, +) -> (Vec, u32, u32) { + let mut image = new_icon_canvas(has_error); + draw_bar_row(&mut image, 6, 14, top_percent, has_error); + draw_bar_row(&mut image, 18, 26, bottom_percent, has_error); + (image.into_raw(), TRAY_ICON_SIZE, TRAY_ICON_SIZE) } /// Render a compact numeric percent tray icon as raw RGBA bytes. pub fn render_percent_icon_rgba(percent: f64, has_error: bool) -> (Vec, u32, u32) { const SZ: u32 = TRAY_ICON_SIZE; - let mut img: RgbaImage = ImageBuffer::new(SZ, SZ); - - for pixel in img.pixels_mut() { - *pixel = Rgba([0, 0, 0, 0]); - } - - let bg_alpha: u8 = if has_error { 180 } else { 255 }; - for y in 2..SZ - 2 { - for x in 2..SZ - 2 { - img.put_pixel(x, y, Rgba([60, 60, 70, bg_alpha])); - } - } + let mut img = new_icon_canvas(has_error); // percent clamped to 0–100 before rounding, so the cast to u32 cannot truncate. #[allow( @@ -134,18 +141,7 @@ pub fn render_percent_icon_rgba(percent: f64, has_error: bool) -> (Vec, u32, let start_x = (SZ.saturating_sub(text_width)) / 2; let start_y = (SZ.saturating_sub(text_height)) / 2; - let (r, g, b) = UsageLevel::from_percent(percent).color(); - let color = if has_error { - // Average of three u8 colour channels: sum ≤ 765, so /3 ≤ 255 fits u8. - #[allow( - clippy::cast_possible_truncation, - reason = "mean of three u8 channels; r+g+b ≤ 765, divided by 3 is ≤ 255 and fits u8" - )] - let gray = ((r as u16 + g as u16 + b as u16) / 3) as u8; - Rgba([gray, gray, gray, 255]) - } else { - Rgba([r, g, b, 255]) - }; + let color = usage_color(percent, has_error); let mut x = start_x; for ch in text.chars() { @@ -304,4 +300,42 @@ mod tests { let (rgba, w, h) = render_percent_icon_rgba(125.0, false); assert_eq!(u32::try_from(rgba.len()).unwrap(), w * h * 4); } + + #[test] + fn stacked_provider_icon_uses_equal_separate_rows() { + let (rgba, width, height) = render_stacked_bar_icon_rgba(100.0, 0.0, false); + assert_eq!((width, height), (TRAY_ICON_SIZE, TRAY_ICON_SIZE)); + + let pixel = |x: u32, y: u32| { + let index = ((y * width + x) * 4) as usize; + [ + rgba[index], + rgba[index + 1], + rgba[index + 2], + rgba[index + 3], + ] + }; + let (r, g, b) = UsageLevel::Critical.color(); + assert_eq!(pixel(8, 8), [r, g, b, 255]); + assert_eq!(pixel(8, 20), [80, 80, 90, 255]); + assert_eq!(pixel(8, 15), [60, 60, 70, 255]); + } + + #[test] + fn normal_and_stacked_bars_share_error_color_policy() { + let (normal, width, _) = render_bar_icon_rgba(100.0, None, true); + let (stacked, _, _) = render_stacked_bar_icon_rgba(100.0, 0.0, true); + let pixel = |rgba: &[u8], x: u32, y: u32| { + let index = ((y * width + x) * 4) as usize; + [ + rgba[index], + rgba[index + 1], + rgba[index + 2], + rgba[index + 3], + ] + }; + + assert_eq!(pixel(&normal, 8, 12), pixel(&stacked, 8, 8)); + assert_eq!(pixel(&normal, 8, 12)[0], pixel(&normal, 8, 12)[1]); + } }