diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index d24d087870..683a35443e 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -169,6 +169,15 @@ pub struct SessionEquivalentForecastSnapshot { pub weekly_used_percent: f64, } +/// Subscription dates from an authenticated OpenAI dashboard/API response. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SubscriptionMetadataSnapshot { + pub starts_at: Option, + pub expires_at: Option, + pub renews_at: Option, +} + /// A frontend-friendly snapshot of one provider's usage data. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -197,6 +206,8 @@ pub struct ProviderUsageSnapshot { pub plan_name: Option, #[serde(default)] pub account_email: Option, + #[serde(default)] + pub subscription: Option, #[serde(default = "default_source_label")] pub source_label: String, #[serde(default)] @@ -399,6 +410,13 @@ impl ProviderUsageSnapshot { }), plan_name: usage.login_method.clone(), account_email: usage.account_email.clone(), + subscription: usage.subscription.as_ref().map(|subscription| { + SubscriptionMetadataSnapshot { + starts_at: subscription.starts_at.map(|date| date.to_rfc3339()), + expires_at: subscription.expires_at.map(|date| date.to_rfc3339()), + renews_at: subscription.renews_at.map(|date| date.to_rfc3339()), + } + }), source_label: result.source_label.clone(), has_successful_claude_cli_quota: result.has_successful_claude_cli_quota, updated_at: usage.updated_at.to_rfc3339(), @@ -446,6 +464,7 @@ impl ProviderUsageSnapshot { cost: None, plan_name: None, account_email: None, + subscription: None, source_label: String::new(), has_successful_claude_cli_quota: false, updated_at: chrono::Utc::now().to_rfc3339(), diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 6d39d6ca24..d842169776 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -1151,6 +1151,7 @@ mod reset_backfill_tests { cost: None, plan_name: None, account_email: None, + subscription: None, source_label: String::new(), has_successful_claude_cli_quota: false, updated_at: "2026-01-01T00:00:00Z".into(), diff --git a/apps/desktop-tauri/src-tauri/src/powertoys.rs b/apps/desktop-tauri/src-tauri/src/powertoys.rs index 5c6352157e..4f5ce88465 100644 --- a/apps/desktop-tauri/src-tauri/src/powertoys.rs +++ b/apps/desktop-tauri/src-tauri/src/powertoys.rs @@ -198,6 +198,7 @@ mod tests { cost: None, plan_name: Some("Team".to_string()), account_email: Some("dev@example.com".to_string()), + subscription: None, source_label: "web".to_string(), has_successful_claude_cli_quota: false, updated_at: "2026-07-09T00:00:00Z".to_string(), @@ -245,6 +246,7 @@ mod tests { cost: None, plan_name: None, account_email: None, + subscription: None, source_label: "web".to_string(), has_successful_claude_cli_quota: false, updated_at: "2026-07-09T00:00:00Z".to_string(), diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index 5c83e2cac1..5d223c3881 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -674,14 +674,22 @@ fn selected_tray_percents( let (selected, companion) = crate::usage_metric::selected_usage_icon_windows(snapshot, settings); ( - display_metric_percent(selected.used_percent, settings.show_as_used), + display_metric_percent(&selected, settings.show_as_used), companion .as_ref() - .map(|window| display_metric_percent(window.used_percent, settings.show_as_used)), + .map(|window| display_metric_percent(window, settings.show_as_used)), ) } -fn display_metric_percent(used_percent: f64, show_as_used: bool) -> f64 { +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 } } @@ -1085,6 +1093,7 @@ mod tests { }), 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(), @@ -1380,6 +1389,76 @@ mod tests { 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(); diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index 9d762c392f..eca2f3642b 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -107,20 +107,29 @@ fn automatic_window( } } - highest_window( - std::iter::once(&snapshot.primary) - .chain(snapshot.secondary.iter()) - .chain(snapshot.model_specific.iter()) - .chain(snapshot.tertiary.iter()) - .chain( - snapshot - .extra_rate_windows - .iter() - .map(|extra| &extra.window), - ) - .filter(|window| !window.is_informational), - ) - .cloned() + let windows = std::iter::once(&snapshot.primary) + .chain(snapshot.secondary.iter()) + .chain(snapshot.model_specific.iter()) + .chain(snapshot.tertiary.iter()) + .chain( + snapshot + .extra_rate_windows + .iter() + .map(|extra| &extra.window), + ) + .filter(|window| !window.is_informational); + let prioritize_exhausted = provider + .map(|id| { + codexbar::core::instantiate_provider(id).automatic_metric_prioritizes_exhausted_window() + }) + .unwrap_or(true); + let selected = if prioritize_exhausted { + highest_automatic_window(windows) + } else { + highest_window(windows) + }; + + selected.cloned() } fn average_window(snapshot: &ProviderUsageSnapshot) -> Option { @@ -187,6 +196,24 @@ fn highest_window<'a>( }) } +fn highest_automatic_window<'a>( + windows: impl Iterator, +) -> Option<&'a RateWindowSnapshot> { + windows.max_by(|a, b| { + automatic_window_is_exhausted(a) + .cmp(&automatic_window_is_exhausted(b)) + .then_with(|| { + a.used_percent + .partial_cmp(&b.used_percent) + .unwrap_or(Ordering::Equal) + }) + }) +} + +fn automatic_window_is_exhausted(window: &RateWindowSnapshot) -> bool { + window.is_exhausted || window.used_percent >= 100.0 +} + #[cfg(test)] mod tests { use super::*; @@ -210,6 +237,7 @@ mod tests { cost: None, plan_name: None, account_email: None, + subscription: None, source_label: "test".to_string(), has_successful_claude_cli_quota: false, updated_at: "2026-08-16T00:00:00Z".to_string(), @@ -260,6 +288,59 @@ mod tests { ); } + #[test] + fn opencodego_automatic_prefers_explicitly_exhausted_window_over_higher_percentage() { + let mut snapshot = snapshot(); + snapshot.provider_id = "opencodego".to_string(); + snapshot.primary.is_exhausted = true; + + let selected = selected_usage_window(&snapshot, &Settings::default()); + + assert_eq!(selected.used_percent, 20.0); + assert!(selected.is_exhausted); + } + + #[test] + fn claude_and_codex_automatic_keep_highest_used_window() { + for provider_id in ["claude", "codex"] { + let mut snapshot = snapshot(); + snapshot.provider_id = provider_id.to_string(); + snapshot.primary.is_exhausted = true; + + let selected = selected_usage_window(&snapshot, &Settings::default()); + + assert_eq!( + selected.used_percent, 60.0, + "{provider_id} should keep highest-used automatic selection" + ); + assert!(!selected.is_exhausted); + } + } + + #[test] + fn automatic_treats_a_full_window_as_exhausted_even_without_the_flag() { + let mut snapshot = snapshot(); + let mut full = window(100.0); + full.is_exhausted = false; + snapshot.tertiary = Some(full); + + let selected = selected_usage_window(&snapshot, &Settings::default()); + + assert_eq!(selected.used_percent, 100.0); + assert!(!selected.is_exhausted); + } + + #[test] + fn non_automatic_highest_window_keeps_percentage_order() { + let healthy = window(80.0); + let mut exhausted = window(20.0); + exhausted.is_exhausted = true; + + let selected = highest_window([&healthy, &exhausted].into_iter()).expect("window"); + + assert_eq!(selected.used_percent, 80.0); + } + #[test] fn single_meaningful_quota_omits_the_companion_icon_lane() { let mut snapshot = snapshot(); diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx index e3ac4dd2a6..119d2b77af 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/AdvancedTab.test.tsx @@ -123,4 +123,22 @@ describe("AdvancedTab", () => { ).toBeGreaterThan(0); }); }); + + it("keeps the Windows Hooks settings surface to one master label and toggle", () => { + render(); + + const hooksSection = screen + .getByRole("heading", { name: "HooksTitle" }) + .closest("section"); + + expect(hooksSection).not.toBeNull(); + expect(hooksSection?.querySelectorAll(".settings-field__label")).toHaveLength(1); + expect(hooksSection?.querySelectorAll('input[type="checkbox"]')).toHaveLength(1); + expect( + hooksSection?.querySelectorAll( + 'input[type="text"], input[type="number"], textarea', + ), + ).toHaveLength(0); + expect(screen.getAllByText("HooksEnableLabel")).toHaveLength(1); + }); }); diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 952db5537a..4c6ca7175c 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -577,6 +577,12 @@ export interface SessionEquivalentForecastSnapshot { weeklyUsedPercent: number; } +export interface SubscriptionMetadataSnapshot { + startsAt: string | null; + expiresAt: string | null; + renewsAt: string | null; +} + /** Backend-classified provider availability state (camelCase serde on the bridge). */ export type ProviderStateKind = | "ready" @@ -606,6 +612,7 @@ export interface ProviderUsageSnapshot { cost: CostSnapshotBridge | null; planName: string | null; accountEmail: string | null; + subscription?: SubscriptionMetadataSnapshot | null; sourceLabel: string; /** Backend proof of a live successful Claude CLI quota fetch; only true is proof. */ hasSuccessfulClaudeCliQuota?: boolean; @@ -948,6 +955,7 @@ export interface CodexAccountUsageSnapshot { credits: CodexCreditsBalance | null; /** Persisted account-scoped extra-usage cost, when available. */ cost?: CostSnapshotBridge | null; + subscription?: SubscriptionMetadataSnapshot | null; updatedAt: string; } diff --git a/rust/src/codex_accounts/api.rs b/rust/src/codex_accounts/api.rs index ac3c368ea1..74c59b2463 100644 --- a/rust/src/codex_accounts/api.rs +++ b/rust/src/codex_accounts/api.rs @@ -15,6 +15,10 @@ use super::models::{ WindowRole, }; use crate::core::credentialed_http_client_builder; +use crate::providers::openai::OpenAISubscriptionFetchResult; + +#[path = "subscription.rs"] +mod subscription; pub const REFRESH_ENDPOINT: &str = "https://auth.openai.com/oauth/token"; pub const USAGE_DEFAULT_BASE: &str = "https://chatgpt.com/backend-api"; @@ -350,14 +354,14 @@ impl CodexAccountApi { workspace_account_id: Option<&str>, verify_live_data: bool, ) -> Result { - if verify_live_data { + let snapshot = if verify_live_data { self.fetch_verified( codex_home_path, credentials, email_hint, workspace_account_id, ) - .await + .await? } else { self.fetch_single( codex_home_path, @@ -365,8 +369,50 @@ impl CodexAccountApi { email_hint, workspace_account_id, ) + .await? + }; + Ok(self + .enrich_subscription_metadata( + codex_home_path, + credentials, + email_hint, + workspace_account_id, + snapshot, + ) + .await) + } + + /// Fetch subscription dates only after the selected account's quota data + /// has been obtained. The request is scoped with the same workspace account + /// header, and the optional result never turns a successful quota read into + /// an error. + async fn enrich_subscription_metadata( + &self, + codex_home_path: &Path, + credentials: &AuthCredentials, + email_hint: Option<&str>, + workspace_account_id: Option<&str>, + snapshot: AccountUsageSnapshot, + ) -> AccountUsageSnapshot { + subscription::enrich_subscription_metadata( + self, + codex_home_path, + credentials, + email_hint, + workspace_account_id, + snapshot, + ) + .await + } + + async fn fetch_subscription_metadata( + &self, + codex_home_path: &Path, + credentials: &AuthCredentials, + account_id: Option<&str>, + ) -> OpenAISubscriptionFetchResult { + subscription::fetch_subscription_metadata(self, codex_home_path, credentials, account_id) .await - } } /// Fetch three reads and require equivalence (CodexControl accuracy model). @@ -460,6 +506,7 @@ impl CodexAccountApi { None, ), updated_at: Utc::now(), + subscription: None, }) } @@ -622,6 +669,13 @@ pub fn resolve_usage_url(codex_home_path: &Path) -> String { format!("{base}{path}") } +/// Resolve the subscription endpoint only for the real OpenAI dashboard host. +/// Custom Codex backends may reuse the usage URL shape but must never receive +/// a ChatGPT subscription probe or be treated as its authority. +pub fn resolve_subscription_url(codex_home_path: &Path) -> Option { + subscription::resolve_subscription_url(codex_home_path) +} + /// Extract `chatgpt_base_url` from a Codex `config.toml`. pub fn parse_chatgpt_base_url(contents: &str) -> Option { for raw_line in contents.lines() { @@ -902,6 +956,7 @@ mod tests { secondary_window: None, credits: None, cost: None, + subscription: None, updated_at: Utc::now(), }; assert!(is_equivalent(&mk(), &mk())); diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs index 4586901654..a43b270fe0 100644 --- a/rust/src/codex_accounts/models.rs +++ b/rust/src/codex_accounts/models.rs @@ -540,6 +540,10 @@ pub struct AccountUsageSnapshot { /// Account-scoped extra-usage cost persisted with the account lane. #[serde(default, skip_serializing_if = "Option::is_none")] pub cost: Option, + /// Subscription dates observed from the same account-scoped OpenAI + /// dashboard/API request as this quota snapshot. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subscription: Option, pub updated_at: DateTime, } @@ -831,6 +835,7 @@ mod tests { secondary_window: None, credits: None, cost: None, + subscription: None, updated_at: utc_now(), }; assert!(snapshot.is_quota_blocked()); diff --git a/rust/src/codex_accounts/stores.rs b/rust/src/codex_accounts/stores.rs index 9ca2fbf117..7b56f5a297 100644 --- a/rust/src/codex_accounts/stores.rs +++ b/rust/src/codex_accounts/stores.rs @@ -248,6 +248,7 @@ mod tests { .with_balance_observation(Some(0.0), utc_now()) .with_account_id("acct-1"), ), + subscription: None, updated_at: utc_now(), }; let mut map = HashMap::new(); diff --git a/rust/src/codex_accounts/subscription.rs b/rust/src/codex_accounts/subscription.rs new file mode 100644 index 0000000000..991246567a --- /dev/null +++ b/rust/src/codex_accounts/subscription.rs @@ -0,0 +1,125 @@ +use std::path::Path; + +use crate::providers::openai::{ + OpenAISubscriptionFetchResult, account_identity_matches, parse_subscription_http_response, +}; + +use super::{ + AuthCredentials, CodexAccountApi, identity_from_credentials, normalize_string, + resolve_usage_url, +}; +use crate::codex_accounts::models::AccountUsageSnapshot; + +const SUBSCRIPTION_PATH: &str = "/subscriptions"; + +pub(super) async fn enrich_subscription_metadata( + api: &CodexAccountApi, + codex_home_path: &Path, + credentials: &AuthCredentials, + email_hint: Option<&str>, + workspace_account_id: Option<&str>, + snapshot: AccountUsageSnapshot, +) -> AccountUsageSnapshot { + if !crate::settings::Settings::load().codex_openai_web_extras() { + return snapshot; + } + let identity = identity_from_credentials(credentials); + if !account_identity_matches(email_hint, identity.email.as_deref()) { + // A managed account hint and the current credential identity + // disagree; retaining quota while dropping optional dates avoids + // attaching a dashboard answer to the wrong lane. + return snapshot; + } + let account_id = workspace_account_id + .and_then(|id| normalize_string(Some(id))) + .or(identity.provider_account_id); + match api + .fetch_subscription_metadata(codex_home_path, credentials, account_id.as_deref()) + .await + { + OpenAISubscriptionFetchResult::Success(metadata) => AccountUsageSnapshot { + subscription: metadata, + ..snapshot + }, + OpenAISubscriptionFetchResult::Unavailable => snapshot, + } +} + +pub(super) async fn fetch_subscription_metadata( + api: &CodexAccountApi, + codex_home_path: &Path, + credentials: &AuthCredentials, + account_id: Option<&str>, +) -> OpenAISubscriptionFetchResult { + let Some(url) = resolve_subscription_url(codex_home_path) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let mut request = api + .client + .get(url) + .header( + "Authorization", + format!("Bearer {}", credentials.access_token), + ) + .header("User-Agent", "codex-cli") + .header("Accept", "application/json") + .header("Cache-Control", "no-cache, no-store, max-age=0") + .header("Pragma", "no-cache") + .timeout(std::time::Duration::from_secs(8)); + if let Some(account_id) = account_id.filter(|id| !id.is_empty()) { + request = request.header("ChatGPT-Account-Id", account_id); + } + let Ok(response) = request.send().await else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let status = response.status().as_u16(); + let Ok(body) = response.bytes().await else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let Ok(body) = std::str::from_utf8(&body) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + parse_subscription_http_response(status, body) +} + +pub(super) fn resolve_subscription_url(codex_home_path: &Path) -> Option { + let usage_url = resolve_usage_url(codex_home_path); + let mut url = reqwest::Url::parse(&usage_url).ok()?; + if !matches!( + url.host_str() + .map(|host| host.to_ascii_lowercase()) + .as_deref(), + Some("chatgpt.com") | Some("chat.openai.com") + ) || !url.path().ends_with("/wham/usage") + { + return None; + } + let base_path = url.path().trim_end_matches("/wham/usage"); + url.set_path(&format!("{base_path}{SUBSCRIPTION_PATH}")); + Some(url.to_string()) +} + +#[cfg(test)] +mod tests { + use super::resolve_subscription_url; + + #[test] + fn resolve_subscription_url_default() { + let dir = tempfile::tempdir().unwrap(); + assert_eq!( + resolve_subscription_url(dir.path()).as_deref(), + Some("https://chatgpt.com/backend-api/subscriptions") + ); + } + + #[test] + fn custom_backend_never_becomes_subscription_authority() { + let dir = tempfile::tempdir().unwrap(); + std::fs::write( + dir.path().join("config.toml"), + "chatgpt_base_url = \"https://gateway.example/backend-api\"\n", + ) + .unwrap(); + assert_eq!(resolve_subscription_url(dir.path()), None); + } +} diff --git a/rust/src/core/cost_cache_budget.rs b/rust/src/core/cost_cache_budget.rs index 2495fd0957..81b40681f7 100644 --- a/rust/src/core/cost_cache_budget.rs +++ b/rust/src/core/cost_cache_budget.rs @@ -76,7 +76,11 @@ fn touches_window(entry: &CostUsageFileUsage, since_key: &str, until_key: &str) /// upstream `estimatedCodexCacheBytes`'s per-entry shape; it deliberately /// overestimates so pruning triggers at or before the real byte budget. fn estimated_entry_bytes(entry: &CostUsageFileUsage) -> usize { - let mut bytes = 240; + let mut bytes = 240 + + entry + .codex_file_identity + .as_ref() + .map_or(0, |identity| identity.len() + 32); for (day, models) in &entry.days { bytes += day.len() + 32; for (model, packed) in models { @@ -321,6 +325,7 @@ mod tests { CostUsageFileUsage { mtime_unix_ms: 0, size, + codex_file_identity: None, days: day_map, parsed_bytes: parsed, codex_scan_target_size: None, diff --git a/rust/src/core/cost_pricing.rs b/rust/src/core/cost_pricing.rs index eb58b11063..e019647a36 100755 --- a/rust/src/core/cost_pricing.rs +++ b/rust/src/core/cost_pricing.rs @@ -7,6 +7,8 @@ use std::collections::HashMap; use std::sync::LazyLock; #[path = "cost_pricing/claude.rs"] mod claude_pricing; +#[path = "cost_pricing/codex.rs"] +mod codex_pricing; pub(crate) use claude_pricing::ClaudePricingResolution; /// Whole-request Codex rates for input above the model context threshold. #[derive(Debug, Clone, Copy)] @@ -337,11 +339,26 @@ static CODEX_PRICING: LazyLock> = LazyLock:: }), }, ); + // GPT-6 Astra pricing (OpenAI model card and pricing table). + // Long-context rates apply to the whole request above 272K input tokens. + m.insert( + "gpt-6-astra", + CodexPricing { + input_cost_per_token: 1e-5, + output_cost_per_token: 5e-5, + cache_read_input_cost_per_token: 1e-6, + display_label: None, + long_context: Some(CodexLongContextRates { + input_cost_per_token: 2e-5, + output_cost_per_token: 7.5e-5, + cache_read_input_cost_per_token: 2e-6, + }), + }, + ); m }); -const CODEX_LONG_CONTEXT_THRESHOLD: u64 = 272_000; /// Claude model pricing table static CLAUDE_PRICING: LazyLock> = LazyLock::new(|| { let mut m = HashMap::new(); @@ -579,20 +596,6 @@ static CLAUDE_PRICING: LazyLock> = LazyLock m }); -fn codex_cost_from_rates( - input_tokens: u64, - cached_input_tokens: u64, - output_tokens: u64, - input_rate: f64, - cache_read_rate: f64, - output_rate: f64, -) -> f64 { - let cached = cached_input_tokens.min(input_tokens); - let non_cached = input_tokens.saturating_sub(cached); - (non_cached as f64) * input_rate - + (cached as f64) * cache_read_rate - + (output_tokens as f64) * output_rate -} /// Cost usage pricing utilities pub struct CostUsagePricing; @@ -687,9 +690,8 @@ impl CostUsagePricing { pub fn codex_api_fast_multiplier(model: &str) -> Option { let base = Self::codex_fast_base_model(model); match base.as_str() { - "gpt-5.4" | "gpt-5.4-mini" | "gpt-5.6-sol" | "gpt-5.6-terra" | "gpt-5.6-luna" => { - Some(2.0) - } + "gpt-5.4" | "gpt-5.4-mini" | "gpt-5.6-sol" | "gpt-5.6-terra" | "gpt-5.6-luna" + | "gpt-6-astra" => Some(2.0), "gpt-5.5" => Some(2.5), _ => None, } @@ -699,12 +701,15 @@ impl CostUsagePricing { /// /// Computes the standard cost for the BASE model (stripping fast/priority /// suffixes), then applies the Fast multiplier. Returns `None` when the - /// model has no Fast lane or when long-context input exceeds the 272 000 - /// threshold guard (Fast is not offered above that). + /// model has no Fast lane or when a model without Astra's published + /// long-context Fast rates exceeds the 272 000 threshold. pub fn codex_fast_cost_usd(model: &str, input: i32, cached: i32, output: i32) -> Option { let multiplier = Self::codex_api_fast_multiplier(model)?; - // Long-context guard: Fast is not offered above the threshold. - if (input as u64) > CODEX_LONG_CONTEXT_THRESHOLD { + // Older models do not offer Fast for long-context requests. Astra + // publishes a Fast rate for the same whole-request long-context tier. + if (input.max(0) as u64) > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD + && !codex_pricing::codex_fast_allows_long_context(model) + { return None; } let base = Self::codex_fast_base_model(model); @@ -747,7 +752,7 @@ impl CostUsagePricing { let key = Self::normalize_codex_model(model); let cutoff = NaiveDate::from_ymd_opt(2026, 7, 30).expect("valid pricing cutoff"); if pricing_date < cutoff { - let long = input_tokens > CODEX_LONG_CONTEXT_THRESHOLD; + let long = input_tokens > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD; let rates = match (key.as_str(), long) { ("gpt-5.6-terra", false) => Some((2.5e-6, 2.5e-7, 1.5e-5)), ("gpt-5.6-terra", true) => Some((5e-6, 5e-7, 2.25e-5)), @@ -756,7 +761,7 @@ impl CostUsagePricing { _ => None, }; if let Some((input_rate, cache_rate, output_rate)) = rates { - return Some(codex_cost_from_rates( + return Some(codex_pricing::codex_cost_from_rates( input_tokens, cached_input_tokens, output_tokens, @@ -783,7 +788,9 @@ impl CostUsagePricing { pricing_date: NaiveDate, ) -> Option { let multiplier = Self::codex_api_fast_multiplier(model)?; - if (input.max(0) as u64) > CODEX_LONG_CONTEXT_THRESHOLD { + if (input.max(0) as u64) > codex_pricing::CODEX_LONG_CONTEXT_THRESHOLD + && !codex_pricing::codex_fast_allows_long_context(model) + { return None; } let base = Self::codex_fast_base_model(model); @@ -804,110 +811,15 @@ impl CostUsagePricing { cached_input_tokens: u64, output_tokens: u64, ) -> Option { - Self::codex_cost_usd_with_pricing_snapshot( + Self::codex_cost_usd_with_cache_write( model, input_tokens, cached_input_tokens, + 0, output_tokens, - None, ) } - pub fn codex_cost_usd_with_pricing_snapshot( - model: &str, - input_tokens: u64, - cached_input_tokens: u64, - output_tokens: u64, - pricing_snapshot: Option<&models_dev_pricing::ModelsDevPricingSnapshot>, - ) -> Option { - let key = Self::normalize_codex_model(model); - // Model-less / deliberately unattributed usage stays unpriced even if a - // pricing catalog later contains a colliding generic entry. - if key == Self::CODEX_UNATTRIBUTED_MODEL { - return None; - } - if let Some(pricing) = CODEX_PRICING.get(key.as_str()) { - let (input_rate, cache_read_rate, output_rate) = - if input_tokens > CODEX_LONG_CONTEXT_THRESHOLD { - if let Some(long_context) = pricing.long_context { - ( - long_context.input_cost_per_token, - long_context.cache_read_input_cost_per_token, - long_context.output_cost_per_token, - ) - } else { - ( - pricing.input_cost_per_token, - pricing.cache_read_input_cost_per_token, - pricing.output_cost_per_token, - ) - } - } else { - ( - pricing.input_cost_per_token, - pricing.cache_read_input_cost_per_token, - pricing.output_cost_per_token, - ) - }; - return Some(codex_cost_from_rates( - input_tokens, - cached_input_tokens, - output_tokens, - input_rate, - cache_read_rate, - output_rate, - )); - } - - // Upstream 0.50.1 #2946: provider-qualified routed models are priced - // against the matching models.dev provider, not OpenAI. Unknown - // `provider/` prefixes are left unpriced (not guessed as OpenAI). - let (provider_id, lookup_model) = match codex_routed_pricing::codex_routed_provider(model) { - Some(routed) => (routed, codex_routed_pricing::strip_route_prefix(model)), - None if model.trim().contains('/') && !model.trim().starts_with("openai/") => { - // Unknown route prefix — do not guess. Leave unpriced. - return None; - } - None => ("openai", model), - }; - let pricing = match pricing_snapshot { - Some(snapshot) => snapshot.lookup(provider_id, lookup_model), - None => models_dev_pricing::lookup(provider_id, lookup_model), - }?; - let use_tier = pricing - .threshold_tokens - .is_some_and(|threshold| input_tokens > threshold); - Some(codex_cost_from_rates( - input_tokens, - cached_input_tokens, - output_tokens, - if use_tier { - pricing - .input_cost_per_token_above_threshold - .unwrap_or(pricing.input_cost_per_token) - } else { - pricing.input_cost_per_token - }, - if use_tier { - pricing - .cache_read_input_cost_per_token_above_threshold - .or(pricing.cache_read_input_cost_per_token) - .unwrap_or(pricing.input_cost_per_token) - } else { - pricing - .cache_read_input_cost_per_token - .unwrap_or(pricing.input_cost_per_token) - }, - if use_tier { - pricing - .output_cost_per_token_above_threshold - .unwrap_or(pricing.output_cost_per_token) - } else { - pricing.output_cost_per_token - }, - )) - } - /// Format model name for display (e.g., "claude-3.5-sonnet" → "Sonnet 3.5") pub fn format_model_name(model: &str) -> String { let lower = model.to_lowercase(); diff --git a/rust/src/core/cost_pricing/codex.rs b/rust/src/core/cost_pricing/codex.rs new file mode 100644 index 0000000000..8b9a20baaa --- /dev/null +++ b/rust/src/core/cost_pricing/codex.rs @@ -0,0 +1,209 @@ +use super::super::{codex_routed_pricing, models_dev_pricing}; +use super::{CODEX_PRICING, CostUsagePricing}; + +pub(super) const CODEX_LONG_CONTEXT_THRESHOLD: u64 = 272_000; +const CODEX_ASTRA_CACHE_WRITE_RATE: f64 = 1.25e-5; +const CODEX_ASTRA_LONG_CACHE_WRITE_RATE: f64 = 2.5e-5; + +pub(super) fn codex_cost_from_rates( + input_tokens: u64, + cached_input_tokens: u64, + output_tokens: u64, + input_rate: f64, + cache_read_rate: f64, + output_rate: f64, +) -> f64 { + let cached = cached_input_tokens.min(input_tokens); + let non_cached = input_tokens.saturating_sub(cached); + (non_cached as f64) * input_rate + + (cached as f64) * cache_read_rate + + (output_tokens as f64) * output_rate +} + +#[allow( + clippy::too_many_arguments, + reason = "Arguments mirror independent token classes and their corresponding pricing rates." +)] +fn codex_cost_from_rates_with_cache_write( + input_tokens: u64, + cached_input_tokens: u64, + cache_write_input_tokens: u64, + output_tokens: u64, + input_rate: f64, + cache_read_rate: f64, + cache_write_rate: f64, + output_rate: f64, +) -> f64 { + if cache_write_input_tokens == 0 { + return codex_cost_from_rates( + input_tokens, + cached_input_tokens, + output_tokens, + input_rate, + cache_read_rate, + output_rate, + ); + } + + let cached = cached_input_tokens.min(input_tokens); + let remaining_input = input_tokens.saturating_sub(cached); + let cache_write = cache_write_input_tokens.min(remaining_input); + let non_cached = remaining_input.saturating_sub(cache_write); + (non_cached as f64) * input_rate + + (cached as f64) * cache_read_rate + + (cache_write as f64) * cache_write_rate + + (output_tokens as f64) * output_rate +} + +pub(super) fn codex_fast_allows_long_context(model: &str) -> bool { + CostUsagePricing::codex_fast_base_model(model) == "gpt-6-astra" +} + +impl CostUsagePricing { + /// Calculate Codex cost in USD when input includes cache-write tokens. + pub fn codex_cost_usd_with_cache_write( + model: &str, + input_tokens: u64, + cached_input_tokens: u64, + cache_write_input_tokens: u64, + output_tokens: u64, + ) -> Option { + Self::codex_cost_usd_with_cache_write_and_pricing_snapshot( + model, + input_tokens, + cached_input_tokens, + cache_write_input_tokens, + output_tokens, + None, + ) + } + + pub fn codex_cost_usd_with_pricing_snapshot( + model: &str, + input_tokens: u64, + cached_input_tokens: u64, + output_tokens: u64, + pricing_snapshot: Option<&models_dev_pricing::ModelsDevPricingSnapshot>, + ) -> Option { + Self::codex_cost_usd_with_cache_write_and_pricing_snapshot( + model, + input_tokens, + cached_input_tokens, + 0, + output_tokens, + pricing_snapshot, + ) + } + + fn codex_cost_usd_with_cache_write_and_pricing_snapshot( + model: &str, + input_tokens: u64, + cached_input_tokens: u64, + cache_write_input_tokens: u64, + output_tokens: u64, + pricing_snapshot: Option<&models_dev_pricing::ModelsDevPricingSnapshot>, + ) -> Option { + let key = Self::normalize_codex_model(model); + // Model-less / deliberately unattributed usage stays unpriced even if a + // pricing catalog later contains a colliding generic entry. + if key == Self::CODEX_UNATTRIBUTED_MODEL { + return None; + } + if let Some(pricing) = CODEX_PRICING.get(key.as_str()) { + let long = input_tokens > CODEX_LONG_CONTEXT_THRESHOLD; + let (input_rate, cache_read_rate, output_rate) = if long { + if let Some(long_context) = pricing.long_context { + ( + long_context.input_cost_per_token, + long_context.cache_read_input_cost_per_token, + long_context.output_cost_per_token, + ) + } else { + ( + pricing.input_cost_per_token, + pricing.cache_read_input_cost_per_token, + pricing.output_cost_per_token, + ) + } + } else { + ( + pricing.input_cost_per_token, + pricing.cache_read_input_cost_per_token, + pricing.output_cost_per_token, + ) + }; + let cache_write_rate = if key == "gpt-6-astra" { + if long { + CODEX_ASTRA_LONG_CACHE_WRITE_RATE + } else { + CODEX_ASTRA_CACHE_WRITE_RATE + } + } else { + input_rate + }; + return Some(codex_cost_from_rates_with_cache_write( + input_tokens, + cached_input_tokens, + cache_write_input_tokens, + output_tokens, + input_rate, + cache_read_rate, + cache_write_rate, + output_rate, + )); + } + + // Upstream 0.50.1 #2946: provider-qualified routed models are priced + // against the matching models.dev provider, not OpenAI. Unknown + // `provider/` prefixes are left unpriced (not guessed as OpenAI). + let (provider_id, lookup_model) = match codex_routed_pricing::codex_routed_provider(model) { + Some(routed) => (routed, codex_routed_pricing::strip_route_prefix(model)), + None if model.trim().contains('/') && !model.trim().starts_with("openai/") => { + // Unknown route prefix — do not guess. Leave unpriced. + return None; + } + None => ("openai", model), + }; + let pricing = match pricing_snapshot { + Some(snapshot) => snapshot.lookup(provider_id, lookup_model), + None => models_dev_pricing::lookup(provider_id, lookup_model), + }?; + let use_tier = pricing + .threshold_tokens + .is_some_and(|threshold| input_tokens > threshold); + let input_rate = if use_tier { + pricing + .input_cost_per_token_above_threshold + .unwrap_or(pricing.input_cost_per_token) + } else { + pricing.input_cost_per_token + }; + 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(pricing.input_cost_per_token) + } else { + pricing + .cache_read_input_cost_per_token + .unwrap_or(pricing.input_cost_per_token) + }; + let output_rate = if use_tier { + pricing + .output_cost_per_token_above_threshold + .unwrap_or(pricing.output_cost_per_token) + } else { + pricing.output_cost_per_token + }; + Some(codex_cost_from_rates_with_cache_write( + input_tokens, + cached_input_tokens, + cache_write_input_tokens, + output_tokens, + input_rate, + cache_read_rate, + input_rate, + output_rate, + )) + } +} diff --git a/rust/src/core/cost_pricing_tests.rs b/rust/src/core/cost_pricing_tests.rs index dc6755fa3d..9984061644 100644 --- a/rust/src/core/cost_pricing_tests.rs +++ b/rust/src/core/cost_pricing_tests.rs @@ -442,3 +442,74 @@ fn gpt56_historical_long_context_uses_pre_cut_rates() { let expected = 270_000.0 * 5e-6 + 30_000.0 * 5e-7 + 1_000.0 * 2.25e-5; assert!((terra - expected).abs() < 1e-10); } + +#[test] +fn gpt6_astra_aliases_use_standard_rates_and_preserve_cached_semantics() { + for model in [ + "gpt-6-astra", + "openai/gpt-6-astra", + "gpt-6-astra-2099-01-01", + ] { + let cost = CostUsagePricing::codex_cost_usd(model, 1000, 300, 100).unwrap(); + // This API receives cache-read tokens only. The remaining 700 input + // tokens are standard input; explicit cache-write tokens use the + // 1.25x Astra rate in the adjacent cache-write regression. + let expected = 700.0 * 10e-6 + 300.0 * 1e-6 + 100.0 * 50e-6; + assert!( + (cost - expected).abs() < 1e-12, + "{model}: expected {expected}, got {cost}" + ); + } +} + +#[test] +fn gpt6_astra_short_context_prices_cache_writes_at_125_percent() { + let cost = + CostUsagePricing::codex_cost_usd_with_cache_write("gpt-6-astra", 1_000, 200, 300, 100) + .unwrap(); + let expected = 500.0 * 1e-5 + 200.0 * 1e-6 + 300.0 * 1.25e-5 + 100.0 * 5e-5; + assert!((cost - expected).abs() < 1e-12); +} + +#[test] +fn gpt6_astra_long_context_prices_cache_writes_at_125_percent() { + let cost = CostUsagePricing::codex_cost_usd_with_cache_write( + "gpt-6-astra", + 272_001, + 100_000, + 50_000, + 1_000, + ) + .unwrap(); + let expected = 122_001.0 * 2e-5 + 100_000.0 * 2e-6 + 50_000.0 * 2.5e-5 + 1_000.0 * 7.5e-5; + assert!((cost - expected).abs() < 1e-12); +} + +#[test] +fn codex_cache_writes_preserve_non_astra_pricing() { + let without_cache_write = CostUsagePricing::codex_cost_usd("gpt-5", 1_000, 200, 100).unwrap(); + let with_cache_write = + CostUsagePricing::codex_cost_usd_with_cache_write("gpt-5", 1_000, 200, 300, 100).unwrap(); + assert!((with_cache_write - without_cache_write).abs() < 1e-12); +} + +#[test] +fn gpt6_astra_switches_the_whole_request_at_long_context_boundary() { + let standard = CostUsagePricing::codex_cost_usd("gpt-6-astra", 272_000, 100_000, 1000).unwrap(); + let long = CostUsagePricing::codex_cost_usd("gpt-6-astra", 272_001, 100_000, 1000).unwrap(); + let expected_standard = 172_000.0 * 10e-6 + 100_000.0 * 1e-6 + 1000.0 * 50e-6; + let expected_long = 172_001.0 * 20e-6 + 100_000.0 * 2e-6 + 1000.0 * 75e-6; + assert!((standard - expected_standard).abs() < 1e-12); + assert!((long - expected_long).abs() < 1e-12); + + let fast = CostUsagePricing::codex_fast_cost_usd("openai/gpt-6-astra", 272_001, 100_000, 1000) + .unwrap(); + assert!((fast - expected_long * 2.0).abs() < 1e-12); +} + +#[test] +fn gpt6_astra_unknown_models_fail_closed() { + for model in ["gpt-6", "other-provider/gpt-6-astra"] { + assert!(CostUsagePricing::codex_cost_usd(model, 1000, 0, 100).is_none()); + } +} diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 4d2c504c3e..d86ca9d720 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -239,6 +239,10 @@ pub struct CostUsageFileUsage { pub mtime_unix_ms: i64, /// File size in bytes pub size: i64, + /// Stable source identity used to detect same-path replacement without + /// opening the raw token history. Legacy entries may omit this field. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub codex_file_identity: Option, /// Daily usage data extracted from this file pub days: HashMap>>, /// Bytes parsed so far (for incremental parsing) diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index e4c2186359..4bf8b07584 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -154,6 +154,50 @@ impl JsonlScanner { Ok(CodexSessionMetadata::default()) } + /// Return the platform file identity used by the cost-cache freshness + /// receipt. This is metadata-only; it never reads token history bytes. + #[cfg(windows)] + pub(crate) fn codex_file_identity( + file_path: &Path, + _metadata: &fs::Metadata, + ) -> Option { + use std::os::windows::io::AsRawHandle; + + use windows::Win32::Foundation::HANDLE; + use windows::Win32::Storage::FileSystem::{ + BY_HANDLE_FILE_INFORMATION, GetFileInformationByHandle, + }; + + let file = File::open(file_path).ok()?; + let mut info = BY_HANDLE_FILE_INFORMATION::default(); + // SAFETY: `file` is an open file handle and `info` is valid for writes + // for the duration of the call. + let ok = unsafe { GetFileInformationByHandle(HANDLE(file.as_raw_handle()), &mut info) }; + if ok.is_err() { + return None; + } + let file_index = ((info.nFileIndexHigh as u64) << 32) | info.nFileIndexLow as u64; + Some(format!("{}:{file_index}", info.dwVolumeSerialNumber)) + } + + #[cfg(unix)] + pub(crate) fn codex_file_identity( + _file_path: &Path, + metadata: &fs::Metadata, + ) -> Option { + use std::os::unix::fs::MetadataExt; + + Some(format!("{}:{}", metadata.dev(), metadata.ino())) + } + + #[cfg(not(any(unix, windows)))] + pub(crate) fn codex_file_identity( + _file_path: &Path, + metadata: &fs::Metadata, + ) -> Option { + Some(format!("{:?}:{}", metadata.modified().ok(), metadata.len())) + } + /// Compare RFC3339 timestamps using parsed instants. Malformed timestamps /// are unsafe for fork-baseline reconciliation and therefore fail closed. pub(crate) fn codex_timestamp_at_or_before(earlier: &str, later: &str) -> bool { diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index 733d1a361c..3e056bbc95 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -1047,6 +1047,7 @@ fn catch_up_snapshot_preserves_established_codex_cost_and_tokens() { CostUsageFileUsage { mtime_unix_ms: 0, size: 100, + codex_file_identity: None, days: HashMap::from([( "2026-08-20".to_string(), HashMap::from([("gpt-5.6-sol".to_string(), vec![1_000, 250, 100])]), @@ -1068,6 +1069,7 @@ fn catch_up_snapshot_preserves_established_codex_cost_and_tokens() { CostUsageFileUsage { mtime_unix_ms: 0, size: 10, + codex_file_identity: None, days: HashMap::new(), parsed_bytes: Some(10), codex_scan_target_size: None, @@ -1120,6 +1122,7 @@ fn save_cache_persists_small_codex_artifact() { CostUsageFileUsage { mtime_unix_ms: 0, size: 100, + codex_file_identity: None, days: HashMap::from([( "2026-01-10".to_string(), HashMap::from([("gpt-5.6-sol".to_string(), vec![10, 0, 1])]), @@ -1190,6 +1193,7 @@ fn save_cache_refuses_non_bounded_provider_oversize() { CostUsageFileUsage { mtime_unix_ms: 0, size: 100, + codex_file_identity: None, days: HashMap::new(), parsed_bytes: None, codex_scan_target_size: None, @@ -1223,6 +1227,7 @@ fn save_cache_refusal_removes_preexisting_destination_artifact() { CostUsageFileUsage { mtime_unix_ms: 0, size: 100, + codex_file_identity: None, days: HashMap::from([( "2026-01-10".to_string(), HashMap::from([("gpt-5.6-sol".to_string(), vec![10, 0, 1])]), diff --git a/rust/src/core/openai_dashboard.rs b/rust/src/core/openai_dashboard.rs index 3e7381c53a..b2d1d88372 100755 --- a/rust/src/core/openai_dashboard.rs +++ b/rust/src/core/openai_dashboard.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::fs; use std::path::PathBuf; -use crate::core::RateWindow; +use crate::core::{RateWindow, SubscriptionMetadata}; /// OpenAI dashboard snapshot with usage and credits data #[derive(Debug, Clone, Serialize, Deserialize)] @@ -41,6 +41,9 @@ pub struct OpenAIDashboardSnapshot { pub credits_remaining: Option, /// Account plan name pub account_plan: Option, + /// Subscription lifecycle dates reported by the authenticated dashboard. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subscription: Option, /// When this snapshot was taken pub updated_at: DateTime, } @@ -63,6 +66,7 @@ impl OpenAIDashboardSnapshot { secondary_limit: None, credits_remaining: None, account_plan: None, + subscription: None, updated_at, } } @@ -149,6 +153,12 @@ impl OpenAIDashboardSnapshot { self.account_plan = Some(plan.into()); self } + + /// Set explicitly observed subscription lifecycle dates. + pub fn with_subscription(mut self, subscription: Option) -> Self { + self.subscription = subscription; + self + } } /// Credit event (purchase, usage, adjustment, etc.) diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 9d97116ca4..c7273d14c3 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -687,6 +687,11 @@ pub trait Provider: Send + Sync { None } + /// Whether Automatic metric selection should prefer an exhausted quota lane. + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + true + } + /// Whether browser-cookie discovery/recovery is owned by the provider. fn owns_browser_cookie_resolution(&self) -> bool { false diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 422c887597..fc9ff8ee12 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -5,6 +5,41 @@ use serde::{Deserialize, Serialize}; use super::RateWindow; +/// Subscription dates explicitly reported by an authenticated provider +/// dashboard or subscription endpoint. +/// +/// These values are deliberately independent from quota-window reset times: +/// a reset is not evidence of a subscription boundary, and a missing date is +/// kept missing rather than inferred. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct SubscriptionMetadata { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub starts_at: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub expires_at: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub renews_at: Option>, +} + +impl SubscriptionMetadata { + pub const fn new( + starts_at: Option>, + expires_at: Option>, + renews_at: Option>, + ) -> Self { + Self { + starts_at, + expires_at, + renews_at, + } + } + + pub const fn is_empty(&self) -> bool { + self.starts_at.is_none() && self.expires_at.is_none() && self.renews_at.is_none() + } +} + /// Provider-specific operational data reported by a Wayfinder gateway. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct WayfinderUsageSnapshot { @@ -116,6 +151,11 @@ pub struct UsageSnapshot { /// Login method/plan info (e.g., "Claude Pro", "Claude Max") #[serde(skip_serializing_if = "Option::is_none")] pub login_method: Option, + + /// Subscription dates explicitly reported by the provider's authenticated + /// dashboard/API. These are not derived from quota reset windows. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub subscription: Option, } impl UsageSnapshot { @@ -133,6 +173,7 @@ impl UsageSnapshot { account_email: None, account_organization: None, login_method: None, + subscription: None, } } @@ -196,6 +237,14 @@ impl UsageSnapshot { self } + /// Attach an explicitly observed subscription payload. Passing `None` + /// clears a previously attached payload when the provider has positively + /// reported that no subscription dates are available. + pub fn with_subscription(mut self, subscription: Option) -> Self { + self.subscription = subscription; + self + } + /// Get the most restrictive (highest used) rate window pub fn most_restrictive(&self) -> &RateWindow { let mut most = &self.primary; diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index d6854d1df9..68aac95a45 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -33,9 +33,13 @@ use crate::providers::opencodego::local as opencodego_local; use crate::settings::Settings; mod claude_pricing; mod codex; +mod read_receipt; +mod stats; use claude_pricing::ClaudeScanPricingResolver; #[cfg(test)] use claude_pricing::{ClaudePricing, FALLBACK_CLAUDE_MODEL}; +pub use read_receipt::CodexScanReadReceipt; +pub use stats::CostScanStats; /// Completeness of the pricing coverage in a [`CostSummary`] (upstream 0.48.0 F18). /// @@ -389,22 +393,6 @@ struct ClaudeUsageRecord { cost: f64, } -/// Per-pass counters for cache/resume behavior (tests + diagnostics). -#[derive(Debug, Clone, Default, PartialEq, Eq)] -pub struct CostScanStats { - pub files_seen: u32, - pub files_parsed: u32, - pub files_skipped: u32, - pub files_resumed: u32, - /// Files deferred to a later bounded Codex catch-up pass. - pub files_deferred: u32, - /// Newly consumed Codex JSONL bytes in this refresh. - pub codex_bytes_read: u64, - /// Timestamp comparisons performed while validating Codex append history. - pub token_timestamp_comparisons: u64, - pub used_cache_debounce: bool, -} - #[derive(Debug, Clone)] pub struct CostScanner { days: u32, diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index dd4bc64928..9f244926d6 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -1,70 +1,14 @@ use super::*; +mod cache_days; mod logical_target; mod pending_range; mod reconciliation; +use cache_days::rebuild_cache_days; use logical_target::*; use pending_range::{CodexPendingScanContext, codex_cache_has_validated_state}; use reconciliation::*; -fn rebuild_cache_days(cache: &mut CostUsageCache) { - cache.days.clear(); - for usage in cache.files.values() { - for (day, models) in &usage.days { - let day_entry = cache.days.entry(day.clone()).or_default(); - for (model, packed) in models { - let dest = day_entry - .entry(model.clone()) - .or_insert_with(|| vec![0, 0, 0]); - if dest.len() < 3 { - dest.resize(3, 0); - } - - let had_core_tokens = dest[0] != 0 || dest[1] != 0 || dest[2] != 0; - let source_input = packed.first().copied().unwrap_or(0); - let source_cached = packed.get(1).copied().unwrap_or(0); - let source_output = packed.get(2).copied().unwrap_or(0); - let source_has_tokens = - source_input != 0 || source_cached != 0 || source_output != 0; - let source_reasoning = packed - .get(3) - .copied() - .map(|reasoning| reasoning.max(0).min(source_output.max(0))); - - dest[0] = dest[0].saturating_add(source_input); - dest[1] = dest[1].saturating_add(source_cached); - dest[2] = dest[2].saturating_add(source_output); - - if !source_has_tokens { - continue; - } - - if !had_core_tokens { - match source_reasoning { - Some(reasoning) => { - if dest.len() >= 4 { - dest[3] = reasoning.min(dest[2].max(0)); - } else { - dest.push(reasoning.min(dest[2].max(0))); - } - } - None => dest.truncate(3), - } - continue; - } - - match (dest.get(3).copied(), source_reasoning) { - (Some(previous), Some(reasoning)) => { - let merged = previous.saturating_add(reasoning).min(dest[2].max(0)); - dest[3] = merged; - } - _ => dest.truncate(3), - } - } - } - } -} - fn summary_from_cached_report( report: &CachedCostReport, period_start: NaiveDate, @@ -121,6 +65,13 @@ fn codex_parent_baseline( 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()) @@ -729,6 +680,7 @@ impl CostScanner { let size = metadata.len().min(i64::MAX as u64) as i64; let mtime_ms = system_time_to_unix_ms(metadata.modified().ok()); let path_key = path.to_string_lossy().to_string(); + let file_identity = JsonlScanner::codex_file_identity(path, &metadata); let cached = cache.files.get(&path_key).cloned(); let cache_covers_range = JsonlScanner::cache_covers_range(cache, range); let trace_was_pruned = cached.as_ref().is_some_and(|entry| { @@ -744,6 +696,40 @@ impl CostScanner { is_complete: false, }; } + let cache_entry_is_fresh = |entry: &CostUsageFileUsage| { + cached_codex_file_is_fresh(cache, 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, + }; + + // 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() + && cache_entry_is_fresh(entry) + && identity_matches_cached(entry) + { + let (session_cost, has_tokens) = + add_codex_days_map_to_summary(summary, &entry.days, range); + if has_tokens { + summary.total_cost_usd += session_cost; + summary.sessions_count += 1; + } + stats.files_skipped = stats.files_skipped.saturating_add(1); + return CodexFileScanOutcome { + bytes_read: 0, + is_complete: true, + }; + } + + 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() @@ -786,6 +772,7 @@ impl CostScanner { CostUsageFileUsage { mtime_unix_ms: mtime_ms, size, + codex_file_identity: file_identity.clone(), days: HashMap::new(), parsed_bytes: Some(0), codex_scan_target_size: None, @@ -807,12 +794,8 @@ impl CostScanner { } if let Some(entry) = &cached - && cache_covers_range - && !entry.codex_unresolved_fork_parent - && entry.mtime_unix_ms == mtime_ms - && entry.size == size - && codex_scan_target_size(entry) == size - && entry.parsed_bytes.unwrap_or(0) >= size + && cached_codex_file_is_fresh(cache, entry, cache_covers_range, mtime_ms, size) + && (entry.codex_file_identity.is_none() || identity_matches_cached(entry)) { let (session_cost, has_tokens) = add_codex_days_map_to_summary(summary, &entry.days, range); @@ -820,6 +803,11 @@ 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, @@ -827,6 +815,10 @@ impl CostScanner { }; } + stats.codex_history_read_paths.push(path_key.clone()); + stats.codex_read_receipt.history_reads = + stats.codex_read_receipt.history_reads.saturating_add(1); + if !is_fork && !cached_identity_changed && let Some(entry) = &cached @@ -880,6 +872,9 @@ impl CostScanner { CostUsageFileUsage { mtime_unix_ms: mtime_ms, size, + codex_file_identity: file_identity + .clone() + .or(entry.codex_file_identity.clone()), days, parsed_bytes: Some(parse_result.parsed_bytes), codex_scan_target_size: Some(parse_result.scan_target_size), @@ -941,6 +936,7 @@ impl CostScanner { CostUsageFileUsage { mtime_unix_ms: mtime_ms, size, + codex_file_identity: file_identity.clone(), days: HashMap::new(), parsed_bytes: Some(0), codex_scan_target_size: None, @@ -977,6 +973,7 @@ impl CostScanner { CostUsageFileUsage { mtime_unix_ms: mtime_ms, size, + codex_file_identity: file_identity, days, parsed_bytes: Some(parse_result.parsed_bytes), codex_scan_target_size: Some(parse_result.scan_target_size), diff --git a/rust/src/cost_scanner/codex/cache_days.rs b/rust/src/cost_scanner/codex/cache_days.rs new file mode 100644 index 0000000000..aea9aff004 --- /dev/null +++ b/rust/src/cost_scanner/codex/cache_days.rs @@ -0,0 +1,59 @@ +use super::*; + +pub(super) fn rebuild_cache_days(cache: &mut CostUsageCache) { + cache.days.clear(); + for usage in cache.files.values() { + for (day, models) in &usage.days { + let day_entry = cache.days.entry(day.clone()).or_default(); + for (model, packed) in models { + let dest = day_entry + .entry(model.clone()) + .or_insert_with(|| vec![0, 0, 0]); + if dest.len() < 3 { + dest.resize(3, 0); + } + + let had_core_tokens = dest[0] != 0 || dest[1] != 0 || dest[2] != 0; + let source_input = packed.first().copied().unwrap_or(0); + let source_cached = packed.get(1).copied().unwrap_or(0); + let source_output = packed.get(2).copied().unwrap_or(0); + let source_has_tokens = + source_input != 0 || source_cached != 0 || source_output != 0; + let source_reasoning = packed + .get(3) + .copied() + .map(|reasoning| reasoning.max(0).min(source_output.max(0))); + + dest[0] = dest[0].saturating_add(source_input); + dest[1] = dest[1].saturating_add(source_cached); + dest[2] = dest[2].saturating_add(source_output); + + if !source_has_tokens { + continue; + } + + if !had_core_tokens { + match source_reasoning { + Some(reasoning) => { + if dest.len() >= 4 { + dest[3] = reasoning.min(dest[2].max(0)); + } else { + dest.push(reasoning.min(dest[2].max(0))); + } + } + None => dest.truncate(3), + } + continue; + } + + match (dest.get(3).copied(), source_reasoning) { + (Some(previous), Some(reasoning)) => { + let merged = previous.saturating_add(reasoning).min(dest[2].max(0)); + dest[3] = merged; + } + _ => dest.truncate(3), + } + } + } + } +} diff --git a/rust/src/cost_scanner/codex/logical_target.rs b/rust/src/cost_scanner/codex/logical_target.rs index 20f27abd14..a98fb1e63e 100644 --- a/rust/src/cost_scanner/codex/logical_target.rs +++ b/rust/src/cost_scanner/codex/logical_target.rs @@ -1,5 +1,21 @@ use super::*; +pub(super) fn cached_codex_file_is_fresh( + cache: &CostUsageCache, + entry: &CostUsageFileUsage, + cache_covers_range: bool, + mtime_unix_ms: i64, + size: i64, +) -> bool { + cache_covers_range + && !entry.codex_unresolved_fork_parent + && entry.mtime_unix_ms == mtime_unix_ms + && entry.size == size + && codex_scan_target_size(entry) == size + && entry.parsed_bytes.unwrap_or(0) >= size + && super::codex_fork_parent_is_safe(cache, entry) +} + pub(super) fn cached_codex_file_is_complete_for_range( cache: &CostUsageCache, path_key: &str, @@ -10,9 +26,18 @@ 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, + }; #[allow(clippy::cast_possible_wrap, reason = "session file sizes fit i64")] let size = metadata.len().min(i64::MAX as u64) as i64; - usage.mtime_unix_ms == system_time_to_unix_ms(metadata.modified().ok()) + identity_matches + && usage.mtime_unix_ms == system_time_to_unix_ms(metadata.modified().ok()) && usage.size == size && codex_scan_target_size(usage) == size && usage.parsed_bytes.unwrap_or(0) >= size diff --git a/rust/src/cost_scanner/read_receipt.rs b/rust/src/cost_scanner/read_receipt.rs new file mode 100644 index 0000000000..1876795dfc --- /dev/null +++ b/rust/src/cost_scanner/read_receipt.rs @@ -0,0 +1,5 @@ +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CodexScanReadReceipt { + pub metadata_reads: u32, + pub history_reads: u32, +} diff --git a/rust/src/cost_scanner/stats.rs b/rust/src/cost_scanner/stats.rs new file mode 100644 index 0000000000..3cc9ef155c --- /dev/null +++ b/rust/src/cost_scanner/stats.rs @@ -0,0 +1,23 @@ +use super::read_receipt::CodexScanReadReceipt; + +/// Per-pass counters for cache/resume behavior (tests + diagnostics). +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CostScanStats { + pub files_seen: u32, + pub files_parsed: u32, + pub files_skipped: u32, + pub files_resumed: u32, + /// Files deferred to a later bounded Codex catch-up pass. + pub files_deferred: u32, + /// Newly consumed Codex JSONL bytes in this refresh. + pub codex_bytes_read: u64, + /// Timestamp comparisons performed while validating Codex append history. + pub token_timestamp_comparisons: u64, + /// Source-read receipt: JSONL paths whose identity prefix was inspected. + pub codex_metadata_read_paths: Vec, + /// Source-read receipt: JSONL paths whose token history was parsed. + pub codex_history_read_paths: Vec, + /// Lazy-read state kept separate so callers can prove cache-only refreshes. + pub codex_read_receipt: CodexScanReadReceipt, + pub used_cache_debounce: bool, +} diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 6f238c21f7..3b5bdc5263 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -214,6 +214,52 @@ fn parses_current_codex_payload_token_count_events() { let _removed = std::fs::remove_file(&path); } +#[test] +fn scans_gpt6_astra_usage_with_cached_and_reasoning_tokens() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let today = Local::now().date_naive(); + let day = today.format("%Y-%m-%d").to_string(); + let day_dir = sessions + .join(today.format("%Y").to_string()) + .join(today.format("%m").to_string()) + .join(today.format("%d").to_string()); + std::fs::create_dir_all(&day_dir).unwrap(); + let line = serde_json::json!({ + "timestamp": Local::now().to_rfc3339(), + "type": "event_msg", + "payload": { + "type": "token_count", + "info": { + "model": "gpt-6-astra", + "total_token_usage": { + "input_tokens": 1000, + "cached_input_tokens": 300, + "output_tokens": 100, + "reasoning_output_tokens": 7 + } + } + } + }); + std::fs::write(day_dir.join("astra.jsonl"), format!("{line}\n")).unwrap(); + + 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, 1000); + assert_eq!(summary.cached_tokens, 300); + assert_eq!(summary.output_tokens, 100); + assert_eq!(summary.reasoning_tokens, Some(7)); + // Codex token-count rows expose cache reads, not cache writes. The 700 + // non-cached input tokens therefore use Astra's standard input rate. + assert!((summary.total_cost_usd - 0.0123).abs() < 1e-12); + assert_eq!(cache.days[&day]["gpt-6-astra"], vec![1000, 300, 100, 7]); +} + #[test] fn derives_claude_dedup_key_from_message_and_request_ids() { assert_eq!( @@ -512,6 +558,7 @@ fn cached_usage_with_packed(day: &str, model: &str, packed: Vec) -> CostUsa CostUsageFileUsage { mtime_unix_ms: 0, size: 1, + codex_file_identity: None, days: HashMap::from([( day.to_string(), HashMap::from([(model.to_string(), packed)]), @@ -1159,6 +1206,10 @@ fn cost_scan_second_pass_skips_unchanged_files_via_cache() { let (summary1, stats1) = scanner.scan_codex_detailed(None); assert_eq!(stats1.files_parsed, 2, "first pass parses both files"); assert_eq!(stats1.files_skipped, 0); + assert_eq!(stats1.codex_metadata_read_paths.len(), 2); + assert_eq!(stats1.codex_history_read_paths.len(), 2); + assert_eq!(stats1.codex_read_receipt.metadata_reads, 2); + assert_eq!(stats1.codex_read_receipt.history_reads, 2); assert!(summary1.total_cost_usd > 0.0); assert_eq!(summary1.sessions_count, 2); @@ -1168,6 +1219,9 @@ fn cost_scan_second_pass_skips_unchanged_files_via_cache() { assert_eq!(stats2.files_seen, 2); assert_eq!(stats2.files_skipped, 2, "cache hit skips re-parse"); assert_eq!(stats2.files_parsed, 0); + assert!(stats2.codex_metadata_read_paths.is_empty()); + assert!(stats2.codex_history_read_paths.is_empty()); + assert_eq!(stats2.codex_read_receipt, Default::default()); assert_eq!(summary2.input_tokens, summary1.input_tokens); assert!((summary2.total_cost_usd - summary1.total_cost_usd).abs() < 1e-9); @@ -1193,6 +1247,101 @@ fn cost_scan_second_pass_skips_unchanged_files_via_cache() { assert!(!stats4.used_cache_debounce); assert_eq!(stats4.files_skipped, 2); assert_eq!(stats4.files_parsed, 0); + assert!(stats4.codex_history_read_paths.is_empty()); +} + +#[test] +fn codex_lazy_history_receipt_reads_only_changed_file_and_matches_fresh_parse() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let first_path = write_codex_session_fixture_with_inputs(&sessions, "first.jsonl", &[100]); + let second_path = write_codex_session_fixture_with_inputs(&sessions, "second.jsonl", &[200]); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + + let (initial, _, _) = scanner.scan_codex_detailed_with_cache(None); + let (unchanged, unchanged_stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(unchanged.input_tokens, initial.input_tokens); + assert!(unchanged_stats.codex_metadata_read_paths.is_empty()); + assert!(unchanged_stats.codex_history_read_paths.is_empty()); + assert_eq!(unchanged_stats.codex_read_receipt, Default::default()); + + use std::io::Write as _; + let timestamp = Utc::now().format("%Y-%m-%dT%H:%M:%S%.3fZ").to_string(); + let extra = format!( + r#"{{"timestamp":"{timestamp}","type":"event_msg","payload":{{"type":"token_count","info":{{"model":"gpt-5","total_token_usage":{{"input_tokens":150,"cached_input_tokens":0,"output_tokens":5}}}}}}}} +"# + ); + std::fs::OpenOptions::new() + .append(true) + .open(&first_path) + .unwrap() + .write_all(extra.as_bytes()) + .unwrap(); + + let (incremental, incremental_stats, _) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!( + incremental_stats.codex_metadata_read_paths, + vec![first_path.to_string_lossy().to_string()] + ); + assert_eq!( + incremental_stats.codex_history_read_paths, + vec![first_path.to_string_lossy().to_string()] + ); + assert_eq!(incremental_stats.codex_read_receipt.metadata_reads, 1); + assert_eq!(incremental_stats.codex_read_receipt.history_reads, 1); + assert_eq!(incremental.input_tokens, 350); + + let fresh = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(root.path().join("fresh-cache")) + .with_sessions_dirs(vec![sessions]); + let (full, full_stats) = fresh.scan_codex_detailed(None); + assert_eq!(full_stats.codex_history_read_paths.len(), 2); + assert_eq!(incremental.input_tokens, full.input_tokens); + assert_eq!(incremental.output_tokens, full.output_tokens); + assert_eq!(incremental.cached_tokens, full.cached_tokens); + assert_eq!(incremental.by_model_tokens, full.by_model_tokens); + assert!((incremental.total_cost_usd - full.total_cost_usd).abs() < 1e-12); + assert!(second_path.exists()); +} + +#[test] +fn codex_file_identity_invalidates_same_path_cache_without_eager_history_read() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let path = write_codex_session_fixture(&sessions, "replacement.jsonl", 100); + let scanner = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (_, _, _) = scanner.scan_codex_detailed_with_cache(None); + let old_mtime = std::fs::metadata(&path).unwrap().modified().unwrap(); + let rotated = path.with_extension("old"); + std::fs::rename(&path, &rotated).unwrap(); + let replacement = write_codex_session_fixture(&sessions, "replacement.jsonl", 200); + // Windows requires a handle with write-attribute access for set_modified; + // keep the replacement's mtime equal to the original without opening it + // read-only. The file contents have the same length, so path/mtime/size + // remain unchanged while the file identity changes. + std::fs::OpenOptions::new() + .write(true) + .open(&replacement) + .unwrap() + .set_modified(old_mtime) + .unwrap(); + + let (summary, stats) = scanner.scan_codex_detailed(None); + assert_eq!(summary.input_tokens, 200); + assert_eq!( + stats.codex_history_read_paths, + vec![replacement.to_string_lossy().to_string()] + ); + assert_eq!(stats.codex_read_receipt.history_reads, 1); } #[test] @@ -1211,6 +1360,7 @@ fn cancelled_fresh_cache_hit_is_not_authoritative() { CostUsageFileUsage { mtime_unix_ms: 0, size: 100, + codex_file_identity: None, days: usage.clone(), parsed_bytes: Some(100), codex_scan_target_size: None, diff --git a/rust/src/pi_session_cost.rs b/rust/src/pi_session_cost.rs index 5d65abd7c9..e372bec322 100644 --- a/rust/src/pi_session_cost.rs +++ b/rust/src/pi_session_cost.rs @@ -272,9 +272,14 @@ fn parse_pi_assistant_entry(value: &Value, target: PiMappedProvider) -> Option

{ - CostUsagePricing::codex_cost_usd(&model, input, cache_read, output).unwrap_or(0.0) - } + PiMappedProvider::Codex => CostUsagePricing::codex_cost_usd_with_cache_write( + &model, + input, + cache_read, + cache_create, + output, + ) + .unwrap_or(0.0), PiMappedProvider::Claude => { // Token counts come from API usage records and fit within i32; // the canonical Claude pricing table takes i32 per-token counts. @@ -377,6 +382,26 @@ mod tests { assert_eq!(entry.model, "gpt-5"); } + #[test] + fn parses_astra_cache_write_and_prices_it() { + let raw = serde_json::json!({ + "id": "astra-msg-1", + "role": "assistant", + "provider": "openai-codex", + "model": "gpt-6-astra", + "usage": { + "input": 1_000, + "output": 100, + "cacheRead": 200, + "cacheWrite": 300 + } + }); + let entry = parse_pi_assistant_entry(&raw, PiMappedProvider::Codex).unwrap(); + let expected = 500.0 * 1e-5 + 200.0 * 1e-6 + 300.0 * 1.25e-5 + 100.0 * 5e-5; + assert_eq!(entry.cache_create, 300); + assert!((entry.cost - expected).abs() < 1e-12); + } + #[test] fn dedupes_shared_entry_ids_across_files() { let dir = tempdir().unwrap(); diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 4ab5bbaada..43592b6f16 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -610,6 +610,10 @@ impl Default for AntigravityProvider { #[async_trait] impl Provider for AntigravityProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Antigravity } diff --git a/rust/src/providers/claude/mod.rs b/rust/src/providers/claude/mod.rs index 7677c80ee6..af52fbbe92 100755 --- a/rust/src/providers/claude/mod.rs +++ b/rust/src/providers/claude/mod.rs @@ -433,6 +433,10 @@ fn last_good_failure_policy_for_error(error: &str) -> LastGoodFailurePolicy { #[async_trait] impl Provider for ClaudeProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Claude } diff --git a/rust/src/providers/codex/api.rs b/rust/src/providers/codex/api.rs index d36c545d4e..3d2763e7b5 100755 --- a/rust/src/providers/codex/api.rs +++ b/rust/src/providers/codex/api.rs @@ -6,6 +6,7 @@ use super::{pat, weekly_reset}; use crate::core::{ CostSnapshot, NamedRateWindow, ProviderError, RateWindow, RateWindowCadence, UsageSnapshot, }; +use crate::providers::openai::OpenAISubscriptionFetchResult; use base64::Engine; use chrono::{DateTime, TimeZone, Utc}; use serde::Deserialize; @@ -13,6 +14,9 @@ use std::path::PathBuf; use std::sync::{Mutex, OnceLock}; use std::time::{Duration, Instant, SystemTime}; +#[path = "subscription.rs"] +mod subscription; + const DEFAULT_BASE_URL: &str = "https://chatgpt.com/backend-api"; const USAGE_PATH: &str = "/wham/usage"; const RESET_CREDITS_PATH: &str = "/wham/rate-limit-reset-credits"; @@ -79,6 +83,7 @@ impl CodexApi { let token = pat::load_token(&self.get_auth_path())?; let (json, whoami) = pat::fetch_usage(&self.client, &self.resolve_base_url(), &token, cli_version).await?; + let account_id = whoami.account_id.clone(); let (mut usage, cost) = self.build_result_from_json(&json)?; if let Some(email) = whoami.email { usage = usage.with_email(email); @@ -88,6 +93,14 @@ impl CodexApi { { usage = usage.with_login_method(format_plan_type(&plan_type)); } + let usage = self + .enrich_subscription_metadata( + &self.resolve_base_url(), + &token, + account_id.as_deref(), + usage, + ) + .await; Ok((usage, cost)) } @@ -111,7 +124,7 @@ impl CodexApi { self.fetch_usage_once(&creds, &base_url).await?; let observed_at = Utc::now(); let first_inventory = weekly_reset::inventory(first_credits.as_ref(), observed_at); - match weekly_reset::initial_decision( + let (usage, cost) = match weekly_reset::initial_decision( &mut state, &first_usage, first_inventory.as_ref(), @@ -121,12 +134,12 @@ impl CodexApi { weekly_reset::InitialDecision::Publish => { weekly_reset::commit_publication(&mut state, &first_usage, first_inventory); weekly_reset::save(&scope, &state); - Ok((first_usage, first_cost)) + (first_usage, first_cost) } weekly_reset::InitialDecision::Preserve => { let usage = weekly_reset::preserve_weekly(&state, first_usage); weekly_reset::save(&scope, &state); - Ok((usage, first_cost)) + (usage, first_cost) } weekly_reset::InitialDecision::RequiresConfirmation => { let confirmation = self.fetch_usage_once(&creds, &base_url).await; @@ -144,7 +157,16 @@ impl CodexApi { first_cost, ); weekly_reset::save(&scope, &state); - return Ok(result); + let (usage, cost) = result; + let usage = self + .enrich_subscription_metadata( + &base_url, + &creds.access_token, + creds.account_id.as_deref(), + usage, + ) + .await; + return Ok((usage, cost)); } }; let confirmation_inventory = @@ -165,16 +187,49 @@ impl CodexApi { confirmation_inventory, ); weekly_reset::save(&scope, &state); - Ok((confirmation_usage, confirmation_cost)) + (confirmation_usage, confirmation_cost) } weekly_reset::ConfirmationDecision::Preserve => { let usage = weekly_reset::preserve_weekly(&state, first_usage); weekly_reset::save(&scope, &state); - Ok((usage, first_cost)) + (usage, first_cost) } } } - } + }; + let usage = self + .enrich_subscription_metadata( + &base_url, + &creds.access_token, + creds.account_id.as_deref(), + usage, + ) + .await; + Ok((usage, cost)) + } + + /// Subscription metadata is optional enrichment. Usage remains usable when + /// the endpoint is unavailable, malformed, unauthorized, or points at a + /// custom backend. A successful empty cancellation response is the only + /// result allowed to clear dates on the fresh snapshot. + async fn enrich_subscription_metadata( + &self, + base_url: &str, + access_token: &str, + account_id: Option<&str>, + usage: UsageSnapshot, + ) -> UsageSnapshot { + subscription::enrich_subscription_metadata(self, base_url, access_token, account_id, usage) + .await + } + + async fn fetch_subscription_metadata( + &self, + base_url: &str, + access_token: &str, + account_id: Option<&str>, + ) -> OpenAISubscriptionFetchResult { + subscription::fetch_subscription_metadata(self, base_url, access_token, account_id).await } fn preserve_after_confirmation_failure( diff --git a/rust/src/providers/codex/mod.rs b/rust/src/providers/codex/mod.rs index 399bc51577..090c5ee265 100755 --- a/rust/src/providers/codex/mod.rs +++ b/rust/src/providers/codex/mod.rs @@ -80,6 +80,10 @@ impl Default for CodexProvider { #[async_trait] impl Provider for CodexProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Codex } diff --git a/rust/src/providers/codex/subscription.rs b/rust/src/providers/codex/subscription.rs new file mode 100644 index 0000000000..8a127ac445 --- /dev/null +++ b/rust/src/providers/codex/subscription.rs @@ -0,0 +1,72 @@ +use std::time::Duration; + +use crate::core::UsageSnapshot; +use crate::providers::openai::{OpenAISubscriptionFetchResult, parse_subscription_http_response}; + +use super::CodexApi; + +const SUBSCRIPTION_PATH: &str = "/subscriptions"; + +pub(super) async fn enrich_subscription_metadata( + api: &CodexApi, + base_url: &str, + access_token: &str, + account_id: Option<&str>, + usage: UsageSnapshot, +) -> UsageSnapshot { + if !crate::settings::Settings::load().codex_openai_web_extras() { + return usage; + } + match api + .fetch_subscription_metadata(base_url, access_token, account_id) + .await + { + OpenAISubscriptionFetchResult::Success(metadata) => usage.with_subscription(metadata), + OpenAISubscriptionFetchResult::Unavailable => usage, + } +} + +pub(super) async fn fetch_subscription_metadata( + api: &CodexApi, + base_url: &str, + access_token: &str, + account_id: Option<&str>, +) -> OpenAISubscriptionFetchResult { + let Some(host) = reqwest::Url::parse(base_url) + .ok() + .and_then(|url| url.host_str().map(str::to_ascii_lowercase)) + else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + if !matches!(host.as_str(), "chatgpt.com" | "chat.openai.com") { + return OpenAISubscriptionFetchResult::Unavailable; + } + + let mut request = api + .client + .get(format!( + "{}{}", + base_url.trim_end_matches('/'), + SUBSCRIPTION_PATH + )) + .header("Authorization", format!("Bearer {access_token}")) + .header("User-Agent", "CodexBar") + .header("Accept", "application/json") + .header("Cache-Control", "no-cache, no-store, max-age=0") + .header("Pragma", "no-cache") + .timeout(Duration::from_secs(8)); + if let Some(account_id) = account_id.filter(|id| !id.is_empty()) { + request = request.header("ChatGPT-Account-Id", account_id); + } + let Ok(response) = request.send().await else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let status = response.status().as_u16(); + let Ok(body) = response.bytes().await else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let Ok(body) = std::str::from_utf8(&body) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + parse_subscription_http_response(status, body) +} diff --git a/rust/src/providers/copilot/mod.rs b/rust/src/providers/copilot/mod.rs index 09f6ddebf6..ffca7d978e 100755 --- a/rust/src/providers/copilot/mod.rs +++ b/rust/src/providers/copilot/mod.rs @@ -48,6 +48,10 @@ impl Default for CopilotProvider { #[async_trait] impl Provider for CopilotProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Copilot } diff --git a/rust/src/providers/cursor/mod.rs b/rust/src/providers/cursor/mod.rs index 206dc3c84e..c6991cecd4 100755 --- a/rust/src/providers/cursor/mod.rs +++ b/rust/src/providers/cursor/mod.rs @@ -192,6 +192,10 @@ impl Default for CursorProvider { #[async_trait] impl Provider for CursorProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Cursor } diff --git a/rust/src/providers/kimi/code_api.rs b/rust/src/providers/kimi/code_api.rs index 58677f5e42..b8e08852cc 100644 --- a/rust/src/providers/kimi/code_api.rs +++ b/rust/src/providers/kimi/code_api.rs @@ -81,18 +81,30 @@ pub(crate) async fn fetch_via_code_api( let json: KimiCodeApiUsageResponse = resp.json().await.map_err(|e| { ProviderError::Parse(format!("Failed to parse Kimi Code API response: {e}")) })?; + let plan_name = json.plan_name(); + let has_plan_name = plan_name.is_some(); let mut snapshot = snapshot_from_code_api_response(json)?; - snapshot.login_method = Some(login_method.to_string()); + snapshot.login_method = Some(plan_name.unwrap_or_else(|| login_method.to_string())); // Upstream #2622: enrich Code API + CLI usage with the monthly membership // pool from a signed-in Kimi Desktop (or browser/manual) session. - if let Some(web_token) = web::web_auth_token(ctx.manual_cookie_header.as_deref()) { - match web::fetch_subscription_for_enrichment(&client, &web_token).await { - Some(subscription) => { - snapshot = super::apply_subscription_windows(snapshot, &subscription); + 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 { + 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 + { + snapshot.login_method = Some(plan); + } + break; } - None => { - tracing::debug!("Kimi Code monthly enrichment unavailable"); + Err(ProviderError::AuthRequired) => continue, + Err(error) => { + tracing::debug!(error = %error, "Kimi Code monthly enrichment unavailable"); + break; } } } @@ -104,7 +116,11 @@ pub(super) fn snapshot_from_code_api_response( response: KimiCodeApiUsageResponse, ) -> Result { let primary = KimiProvider::rate_window_from_usage_detail(&response.usage, None)?; - let mut usage = UsageSnapshot::new(primary).with_login_method("Code API"); + let mut usage = UsageSnapshot::new(primary).with_login_method( + response + .plan_name() + .unwrap_or_else(|| "Code API".to_string()), + ); if let Some(limit) = response.limits.unwrap_or_default().into_iter().next() { let window_minutes = limit.window.as_ref().and_then(kimi_window_minutes); diff --git a/rust/src/providers/kimi/desktop_token.rs b/rust/src/providers/kimi/desktop_token.rs index cb0adc6e60..5c2d6bae42 100644 --- a/rust/src/providers/kimi/desktop_token.rs +++ b/rust/src/providers/kimi/desktop_token.rs @@ -18,7 +18,9 @@ //! //! Auth cookies are secrets: token values are never logged. +use base64::Engine; use std::path::{Path, PathBuf}; +use std::time::{SystemTime, UNIX_EPOCH}; pub struct KimiDesktopAuthToken; @@ -85,9 +87,45 @@ impl KimiDesktopAuthToken { }) .ok() .and_then(|row| decode_cookie_value(row, aes_key)) + .filter(|token| !is_expired_jwt(token, unix_now_secs())) } } +/// Cookie expiry and JWT expiry can differ. A stale desktop JWT must not +/// shadow a live browser session, but opaque/non-JWT credentials remain +/// usable because there is no local expiry claim to inspect. +fn is_expired_jwt(token: &str, now_unix: f64) -> bool { + let mut parts = token.split('.'); + let _header = parts.next(); + let Some(payload) = parts.next() else { + return false; + }; + if parts.next().is_none() || parts.next().is_some() { + return false; + } + + let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(payload) + .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(payload)); + let Ok(decoded) = decoded else { + return false; + }; + let Ok(claims) = serde_json::from_slice::(&decoded) else { + return false; + }; + let Some(expiry) = claims.get("exp").and_then(serde_json::Value::as_f64) else { + return false; + }; + expiry.is_finite() && expiry <= now_unix +} + +fn unix_now_secs() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or(0.0) +} + /// Decode a `(value, encrypted_value)` pair: plaintext first, AES-256-GCM /// fallback (upstream reads `value` only; Windows Chromium rows are usually /// encrypted). @@ -263,6 +301,21 @@ mod tests { assert_eq!(KimiDesktopAuthToken::load_from(root.path()), None); } + #[test] + fn expired_jwt_is_skipped_but_future_and_opaque_tokens_are_kept() { + let encode = |payload: &str| { + format!( + "header.{}.signature", + base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(payload) + ) + }; + + assert!(is_expired_jwt(&encode(r#"{"exp":100}"#), 101.0)); + assert!(!is_expired_jwt(&encode(r#"{"exp":100}"#), 99.0)); + assert!(!is_expired_jwt("opaque-token", 101.0)); + assert!(!is_expired_jwt(&encode(r#"{"exp":"100"}"#), 101.0)); + } + #[test] fn malformed_database_returns_none() { let (root, database) = make_environment(); diff --git a/rust/src/providers/kimi/mod.rs b/rust/src/providers/kimi/mod.rs index 129c27f9b7..ef74d214b7 100755 --- a/rust/src/providers/kimi/mod.rs +++ b/rust/src/providers/kimi/mod.rs @@ -31,6 +31,8 @@ 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"]; #[derive(Debug, Deserialize)] @@ -38,6 +40,49 @@ struct KimiCodeApiUsageResponse { usage: KimiUsageDetail, #[serde(default)] limits: Option>, + /// Optional membership metadata is deliberately kept as JSON. The API has + /// added fields and changed types without changing the usage payload; a + /// malformed membership section must not discard valid quota statistics. + #[serde(default)] + user: Option, + #[serde(default)] + version: Option, +} + +impl KimiCodeApiUsageResponse { + fn plan_name(&self) -> Option { + let level = self + .user + .as_ref() + .and_then(|user| user.get("membership")) + .and_then(|membership| membership.get("level")) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|level| !level.is_empty() && *level != "LEVEL_UNSPECIFIED")?; + + // Only the known V1 catalog gets friendly names. Unknown catalogs and + // malformed versions remain visible as their raw level value. + let known_catalog = match self.version.as_ref() { + None => true, + Some(serde_json::Value::String(version)) => version == "GOODS_VERSION_V1", + Some(_) => false, + }; + if !known_catalog { + return Some(level.to_string()); + } + + Some( + match level { + "LEVEL_FREE" => "Adagio", + "LEVEL_TRIAL" => "Andante", + "LEVEL_BASIC" => "Moderato", + "LEVEL_INTERMEDIATE" => "Allegretto", + "LEVEL_ADVANCED" => "Allegro", + other => other, + } + .to_string(), + ) + } } #[derive(Debug, Deserialize)] @@ -81,6 +126,36 @@ struct KimiSubscriptionRateLimit { reset_time: Option, } +#[derive(Debug, Deserialize)] +struct KimiSubscriptionResponse { + #[serde(default)] + subscription: Option, +} + +impl KimiSubscriptionResponse { + fn plan_name(&self) -> Option { + let subscription = self.subscription.as_ref()?; + if subscription + .get("active") + .and_then(serde_json::Value::as_bool) + != Some(true) + || subscription + .get("status") + .and_then(serde_json::Value::as_str) + != Some("SUBSCRIPTION_STATUS_ACTIVE") + { + return None; + } + subscription + .get("goods") + .and_then(|goods| goods.get("title")) + .and_then(serde_json::Value::as_str) + .map(str::trim) + .filter(|title| !title.is_empty()) + .map(str::to_string) + } +} + #[derive(Debug, Deserialize)] struct KimiUsageDetail { #[serde(default)] @@ -551,6 +626,34 @@ mod tests { assert!(snapshot.secondary.is_none()); } + #[test] + fn parses_membership_level_without_making_optional_metadata_required() { + let response: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { "limit": "100", "used": "25" }, + "user": { "membership": { "level": "LEVEL_ADVANCED" } }, + "version": "GOODS_VERSION_V1" + })) + .unwrap(); + assert_eq!(response.plan_name().as_deref(), Some("Allegro")); + + let unknown_catalog: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { "limit": "100", "used": "25" }, + "user": { "membership": { "level": "LEVEL_CUSTOM" } }, + "version": "GOODS_VERSION_V2" + })) + .unwrap(); + assert_eq!(unknown_catalog.plan_name().as_deref(), Some("LEVEL_CUSTOM")); + + let malformed_optional: KimiCodeApiUsageResponse = serde_json::from_value(json!({ + "usage": { "limit": "100", "used": "25" }, + "user": "not-an-object", + "version": { "unexpected": true } + })) + .unwrap(); + assert_eq!(malformed_optional.plan_name(), None); + assert!(code_api::snapshot_from_code_api_response(malformed_optional).is_ok()); + } + #[test] fn parses_web_usage_with_subscription_windows() { let usage: KimiWebUsageResponse = serde_json::from_value(json!({ diff --git a/rust/src/providers/kimi/web.rs b/rust/src/providers/kimi/web.rs index 04429862e0..0cdfad9cfa 100644 --- a/rust/src/providers/kimi/web.rs +++ b/rust/src/providers/kimi/web.rs @@ -11,8 +11,9 @@ use reqwest::Client; use super::desktop_token::KimiDesktopAuthToken; use super::{ - KIMI_COOKIE_DOMAINS, KIMI_SUBSCRIPTION_STATS_URL, KIMI_WEB_USAGE_URL, KimiProvider, - KimiSubscriptionStatsResponse, KimiWebUsageResponse, apply_subscription_windows, kimi_web_post, + KIMI_COOKIE_DOMAINS, KIMI_SUBSCRIPTION_STATS_URL, KIMI_SUBSCRIPTION_URL, KIMI_WEB_USAGE_URL, + KimiProvider, KimiSubscriptionResponse, KimiSubscriptionStatsResponse, KimiWebUsageResponse, + apply_subscription_windows, kimi_web_post, }; use crate::browser::cookies::get_cookie_header; use crate::core::{ProviderError, ProviderId, UsageSnapshot}; @@ -35,13 +36,16 @@ fn browser_import_allowed(cookie_source: &str) -> bool { /// 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_token(manual_header: Option<&str>) -> Option { - resolve_web_token(WebTokenInput { +pub(crate) fn web_auth_tokens(manual_header: Option<&str>) -> Vec { + resolve_web_tokens(WebTokenInput { manual_header, cookie_source: &cookie_source(), desktop_token: KimiDesktopAuthToken::load, browser_token: browser_auth_token, }) + .into_iter() + .map(|candidate| candidate.token) + .collect() } struct WebTokenInput<'a> { @@ -51,19 +55,51 @@ struct WebTokenInput<'a> { browser_token: fn() -> Option, } -fn resolve_web_token(input: WebTokenInput) -> Option { +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum WebTokenSource { + Manual, + Desktop, + Browser, +} + +#[derive(Debug, PartialEq, Eq)] +struct WebTokenCandidate { + token: String, + source: WebTokenSource, +} + +fn resolve_web_tokens(input: WebTokenInput) -> Vec { if let Some(header) = input.manual_header && let Ok(token) = KimiProvider::auth_token_from_cookie_header(header) { - return Some(token); + return vec![WebTokenCandidate { + token, + source: WebTokenSource::Manual, + }]; } if !browser_import_allowed(input.cookie_source) { - return None; + return Vec::new(); + } + + let mut candidates = Vec::new(); + let mut seen = std::collections::HashSet::new(); + if let Some(token) = (input.desktop_token)() + && seen.insert(token.clone()) + { + candidates.push(WebTokenCandidate { + token, + source: WebTokenSource::Desktop, + }); } - if let Some(token) = (input.desktop_token)() { - return Some(token); + if let Some(token) = (input.browser_token)() + && seen.insert(token.clone()) + { + candidates.push(WebTokenCandidate { + token, + source: WebTokenSource::Browser, + }); } - (input.browser_token)() + candidates } /// Browser import only: the first usable `kimi-auth`-class token from any of @@ -83,26 +119,67 @@ fn browser_auth_token() -> Option { pub(crate) async fn fetch_via_web( cookie_header: Option<&str>, ) -> Result { - let token = web_auth_token(cookie_header).ok_or_else(|| { - if browser_import_allowed(&cookie_source()) { - ProviderError::AuthRequired - } else { - ProviderError::Other( - "Kimi cookie source is Off; provide a manual cookie header or enable browser import." - .into(), - ) + let source = cookie_source(); + if let Some(token) = + cookie_header.and_then(|header| KimiProvider::auth_token_from_cookie_header(header).ok()) + { + // 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; + } + + if !browser_import_allowed(&source) { + return Err(ProviderError::Other( + "Kimi cookie source is Off; provide a manual cookie header or enable browser import." + .into(), + )); + } + + let client = client()?; + let mut seen = std::collections::HashSet::new(); + + // 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() + && seen.insert(token.clone()) + { + match fetch_via_web_token(&client, &token).await { + Ok(usage) => return Ok(usage), + Err(ProviderError::AuthRequired) => {} + Err(error) => return Err(error), + } + } + + if let Some(token) = browser_auth_token() + && seen.insert(token.clone()) + { + match fetch_via_web_token(&client, &token).await { + Ok(usage) => return Ok(usage), + Err(ProviderError::AuthRequired) => {} + Err(error) => return Err(error), } - })?; + } + + Err(ProviderError::AuthRequired) +} - let client = crate::core::credentialed_http_client_builder() +fn client() -> Result { + crate::core::credentialed_http_client_builder() .timeout(std::time::Duration::from_secs(30)) .build() - .map_err(|e| ProviderError::Other(e.to_string()))?; + .map_err(|e| ProviderError::Other(e.to_string())) +} +async fn fetch_via_web_token( + client: &reqwest::Client, + token: &str, +) -> Result { let resp = kimi_web_post( - &client, + client, KIMI_WEB_USAGE_URL, - &token, + token, serde_json::json!({ "scope": ["FEATURE_CODING"] }), ) .await?; @@ -120,24 +197,42 @@ pub(crate) async fn fetch_via_web( .await .map_err(|e| ProviderError::Parse(e.to_string()))?; - let subscription = match kimi_web_post( - &client, - KIMI_SUBSCRIPTION_STATS_URL, - &token, - serde_json::json!({}), - ) - .await - { - Ok(response) if response.status().is_success() => response.json().await.ok(), - _ => None, - }; + let (subscription, plan_name) = fetch_subscription_details(client, token).await; + + snapshot_from_web_usage_response_with_plan(usage, subscription, plan_name) +} + +const SUBSCRIPTION_ENRICHMENT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(2); - snapshot_from_web_usage_response(usage, subscription) +async fn fetch_subscription_details( + client: &reqwest::Client, + token: &str, +) -> (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), + ); + let plan = tokio::time::timeout( + SUBSCRIPTION_ENRICHMENT_TIMEOUT, + fetch_subscription_plan(client, token), + ); + let (stats, plan) = tokio::join!(stats, plan); + (stats.ok().flatten(), plan.ok().flatten()) } pub(super) fn snapshot_from_web_usage_response( response: KimiWebUsageResponse, subscription: Option, +) -> Result { + snapshot_from_web_usage_response_with_plan(response, subscription, None) +} + +fn snapshot_from_web_usage_response_with_plan( + response: KimiWebUsageResponse, + subscription: Option, + plan_name: Option, ) -> Result { let coding = response .usages @@ -157,16 +252,40 @@ pub(super) fn snapshot_from_web_usage_response( if let Some(subscription) = subscription.as_ref() { usage = apply_subscription_windows(usage, subscription); } + if let Some(plan_name) = plan_name { + usage = usage.with_login_method(plan_name); + } 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 { + Ok(response) if response.status().is_success() => response + .json::() + .await + .ok() + .and_then(|response| response.plan_name()), + _ => None, + } +} + // Kept for `code_api`: resolve the subscription stats snapshot with a web // token; any failure means "no enrichment", never an error. pub(super) async fn fetch_subscription_for_enrichment( client: &Client, token: &str, ) -> Option { + fetch_subscription_for_enrichment_result(client, token) + .await + .ok() + .flatten() +} + +pub(super) async fn fetch_subscription_for_enrichment_result( + client: &Client, + token: &str, +) -> Result, ProviderError> { match kimi_web_post( client, KIMI_SUBSCRIPTION_STATS_URL, @@ -175,8 +294,16 @@ pub(super) async fn fetch_subscription_for_enrichment( ) .await { - Ok(response) if response.status().is_success() => response.json().await.ok(), - _ => None, + Ok(response) if response.status().is_success() => response + .json() + .await + .map(Some) + .map_err(|error| ProviderError::Parse(error.to_string())), + Ok(response) if response.status().as_u16() == 401 || response.status().as_u16() == 403 => { + Err(ProviderError::AuthRequired) + } + Ok(_) => Ok(None), + Err(error) => Err(error), } } @@ -199,80 +326,103 @@ mod tests { fn input<'a>( manual_header: Option<&'a str>, cookie_source: &'a str, - desktop_token: Option<&'static str>, - browser_token: Option<&'static str>, + desktop_token: fn() -> Option, + browser_token: fn() -> Option, ) -> WebTokenInput<'a> { WebTokenInput { manual_header, cookie_source, - desktop_token: if desktop_token.is_some() { - static_desktop - } else { - no_token - }, - browser_token: if browser_token.is_some() { - static_browser - } else { - no_token - }, + desktop_token, + browser_token, } } + fn duplicate_browser() -> Option { + Some("desktop-token".to_string()) + } + #[test] fn manual_cookie_header_wins_regardless_of_source() { - let token = resolve_web_token(input( + let candidates = resolve_web_tokens(input( Some("kimi-auth=manual-token"), "off", - Some("desktop-token"), - Some("browser-token"), + static_desktop, + static_browser, )); - assert_eq!(token.as_deref(), Some("manual-token")); + assert_eq!( + candidates, + vec![WebTokenCandidate { + token: "manual-token".to_string(), + source: WebTokenSource::Manual, + }] + ); } #[test] fn desktop_token_precedes_browser_import() { - let token = resolve_web_token(input( - None, - "browser", - Some("desktop-token"), - Some("browser-token"), - )); - assert_eq!(token.as_deref(), Some("desktop-token")); + let candidates = resolve_web_tokens(input(None, "browser", static_desktop, static_browser)); + assert_eq!(candidates.len(), 2); + assert_eq!(candidates[0].source, WebTokenSource::Desktop); + assert_eq!(candidates[0].token, "desktop-token"); + assert_eq!(candidates[1].source, WebTokenSource::Browser); + assert_eq!(candidates[1].token, "browser-token"); } #[test] fn cookie_source_off_blocks_desktop_and_browser_but_not_manual() { assert_eq!( - resolve_web_token(input( - None, - "off", - Some("desktop-token"), - Some("browser-token") - )), - None + resolve_web_tokens(input(None, "off", static_desktop, static_browser)), + Vec::new() ); assert_eq!( - resolve_web_token(input(None, "off", None, Some("browser-token"))), - None + resolve_web_tokens(input(None, "off", no_token, static_browser)), + Vec::new() ); assert_eq!( - resolve_web_token(input(Some("kimi-auth=manual"), "off", None, None)).as_deref(), - Some("manual") + resolve_web_tokens(input(Some("kimi-auth=manual"), "off", no_token, no_token)), + vec![WebTokenCandidate { + token: "manual".to_string(), + source: WebTokenSource::Manual, + }] ); } #[test] fn browser_token_used_when_desktop_absent() { - let token = resolve_web_token(input(None, "auto", None, Some("browser-token"))); - assert_eq!(token.as_deref(), Some("browser-token")); + let candidates = resolve_web_tokens(input(None, "auto", no_token, static_browser)); + assert_eq!( + candidates, + vec![WebTokenCandidate { + token: "browser-token".to_string(), + source: WebTokenSource::Browser, + }] + ); } #[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 token = resolve_web_token(input(None, "manual", Some("desktop-token"), None)); - assert_eq!(token.as_deref(), Some("desktop-token")); + let candidates = resolve_web_tokens(input(None, "manual", static_desktop, no_token)); + assert_eq!( + candidates, + vec![WebTokenCandidate { + token: "desktop-token".to_string(), + source: WebTokenSource::Desktop, + }] + ); + } + + #[test] + fn duplicate_automatic_tokens_are_deduplicated() { + let candidates = resolve_web_tokens(input(None, "auto", static_desktop, duplicate_browser)); + assert_eq!( + candidates, + vec![WebTokenCandidate { + token: "desktop-token".to_string(), + source: WebTokenSource::Desktop, + }] + ); } #[test] @@ -281,4 +431,61 @@ mod tests { assert!(browser_import_allowed("browser")); assert!(browser_import_allowed("manual")); } + + #[test] + fn subscription_stats_do_not_invent_a_membership_label() { + let usage: KimiWebUsageResponse = serde_json::from_value(serde_json::json!({ + "usages": [{ + "scope": "FEATURE_CODING", + "detail": { "limit": "1000", "used": "125" } + }] + })) + .unwrap(); + let subscription: KimiSubscriptionStatsResponse = + serde_json::from_value(serde_json::json!({ + "subscriptionBalance": { + "amountUsedRatio": 0.25, + "expireTime": "2026-09-30T00:00:00Z" + }, + "ratelimitCode7d": { + "ratio": 0.1, + "enabled": true, + "resetTime": "2026-09-14T00:00:00Z" + } + })) + .unwrap(); + + let snapshot = snapshot_from_web_usage_response(usage, Some(subscription)).unwrap(); + + assert_eq!(snapshot.login_method.as_deref(), Some("Kimi")); + assert!( + snapshot + .extra_rate_windows + .iter() + .any(|window| window.id == "kimi-monthly") + ); + } + + #[test] + fn active_subscription_title_is_used_but_inactive_title_is_ignored() { + let active: KimiSubscriptionResponse = serde_json::from_value(serde_json::json!({ + "subscription": { + "active": true, + "status": "SUBSCRIPTION_STATUS_ACTIVE", + "goods": { "title": " Allegro " } + } + })) + .unwrap(); + assert_eq!(active.plan_name().as_deref(), Some("Allegro")); + + let inactive: KimiSubscriptionResponse = serde_json::from_value(serde_json::json!({ + "subscription": { + "active": false, + "status": "SUBSCRIPTION_STATUS_EXPIRED", + "goods": { "title": "Allegro" } + } + })) + .unwrap(); + assert_eq!(inactive.plan_name(), None); + } } diff --git a/rust/src/providers/kiro/mod.rs b/rust/src/providers/kiro/mod.rs index 5f36caf109..7b4bb6cc88 100755 --- a/rust/src/providers/kiro/mod.rs +++ b/rust/src/providers/kiro/mod.rs @@ -3,6 +3,8 @@ //! Fetches usage data from Kiro (Amazon's AI coding assistant) //! Uses kiro-cli for authentication and usage fetching +#[cfg(test)] +mod tests; mod usage_limits; pub mod version; @@ -31,10 +33,10 @@ use crate::core::{ pub struct KiroProvider { metadata: ProviderMetadata, } - struct KiroCliUsage { plan_name: String, matched_new_format: bool, + is_summary: bool, is_managed_plan: bool, reset_date: Option>, credits_percent: f64, @@ -180,11 +182,16 @@ impl KiroProvider { let lowered = stripped.to_lowercase(); let parsed = Self::parse_usage_fields(&stripped, &lowered); - if let Some(usage) = Self::usage_without_metrics(&parsed) { - return Ok(usage); - } - - let mut usage = Self::usage_with_metrics(&parsed); + let mut usage = if let Some(usage) = Self::usage_without_metrics(&parsed) { + usage + } else { + if !parsed.matched_percent && !parsed.matched_credits { + return Err(ProviderError::Parse( + "Kiro CLI output did not include usable usage metrics".to_string(), + )); + } + Self::usage_with_metrics(&parsed) + }; usage = Self::apply_overage_windows(usage, &parsed); if let Some(bonus) = parsed.bonus_window { @@ -214,7 +221,10 @@ impl KiroProvider { } fn usage_without_metrics(parsed: &KiroCliUsage) -> Option { - if parsed.matched_percent || parsed.matched_credits { + if parsed.matched_percent + || parsed.matched_credits + || !(parsed.is_summary || (parsed.matched_new_format && parsed.is_managed_plan)) + { return None; } @@ -224,7 +234,10 @@ impl KiroProvider { "Kiro (installed)" }; - Some(UsageSnapshot::new(RateWindow::new(0.0)).with_login_method(method)) + Some( + UsageSnapshot::new(RateWindow::informational("Usage unavailable")) + .with_login_method(method), + ) } fn usage_with_metrics(parsed: &KiroCliUsage) -> UsageSnapshot { @@ -234,7 +247,7 @@ impl KiroProvider { } fn parse_usage_fields(stripped: &str, lowered: &str) -> KiroCliUsage { - let (plan_name, matched_new_format) = Self::parse_plan_name(stripped); + let (plan_name, matched_new_format, is_summary) = Self::parse_plan_name(stripped); let (credits_percent, matched_percent, matched_credits) = Self::parse_credit_usage(stripped); let (overages_enabled, overage_credits_used, estimated_overage_cost) = @@ -243,6 +256,7 @@ impl KiroProvider { KiroCliUsage { plan_name, matched_new_format, + is_summary, is_managed_plan: lowered.contains("managed by admin") || lowered.contains("managed by organization"), reset_date: Self::capture_text(stripped, r"resets on (\d{2}/\d{2})") @@ -258,16 +272,24 @@ impl KiroProvider { } } - fn parse_plan_name(stripped: &str) -> (String, bool) { - if let Some(plan_line) = Self::capture_text(stripped, r"Plan:\s*(.+)") - && let Some(first_line) = plan_line.lines().next() + fn parse_plan_name(stripped: &str) -> (String, bool, bool) { + if let Some(summary_name) = Self::capture_text( + stripped, + r"(?m)^[ \t]*Plan:[ \t]*([^|\r\n]+?)[ \t]*\|[ \t]*[0-9]+[ \t]+usage breakdowns?[ \t]*\r?$", + ) && !summary_name.trim().is_empty() { - return (first_line.trim().to_string(), true); + return (summary_name.trim().to_string(), true, true); + } + + if let Some(plan_line) = + Self::capture_text(stripped, r"(?m)^[ \t]*Plan:[ \t]*([^\r\n]+?)[ \t]*\r?$") + { + return (plan_line.trim().to_string(), true, false); } let legacy = Self::capture_text(stripped, r"\|\s*(KIRO\s+\w+)") .unwrap_or_else(|| "Kiro".to_string()); - (legacy, false) + (legacy, false, false) } fn parse_credit_usage(stripped: &str) -> (f64, bool, bool) { @@ -324,14 +346,14 @@ impl KiroProvider { usage = usage.with_extra_rate_window( "kiro-overage-credits", "Overage usage", - RateWindow::with_details(0.0, None, None, Some(format!("{credits:.2} credits"))), + RateWindow::informational(format!("{credits:.2} credits")), ); } if let Some(cost) = parsed.estimated_overage_cost { usage = usage.with_extra_rate_window( "kiro-overage-cost", "Overage cost", - RateWindow::with_details(0.0, None, None, Some(format!("${cost:.2} USD"))), + RateWindow::informational(format!("${cost:.2} USD")), ); } @@ -495,27 +517,3 @@ impl Provider for KiroProvider { } } } - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn cli_presence_maps_to_local_runtime_offline_but_state_db_stays_default() { - assert_eq!( - KiroProvider::new().error_state_kind(&ProviderError::NotInstalled( - "kiro-cli not found. Install from https://kiro.dev".to_string(), - )), - crate::core::ProviderStateKind::LocalRuntimeOffline - ); - // The state-database token lookup is auth-flavored and keeps the - // default mapping. - assert_eq!( - KiroProvider::new().error_state_kind(&ProviderError::NotInstalled( - "Kiro CLI state database not found at C:\\Users\\x\\Kiro-Cli\\data.sqlite3" - .to_string(), - )), - crate::core::ProviderStateKind::NeedsAuthentication - ); - } -} diff --git a/rust/src/providers/kiro/tests.rs b/rust/src/providers/kiro/tests.rs new file mode 100644 index 0000000000..1ac02921d3 --- /dev/null +++ b/rust/src/providers/kiro/tests.rs @@ -0,0 +1,180 @@ +use chrono::{TimeZone, Utc}; + +use super::{KiroProvider, usage_limits}; +use crate::core::{Provider, ProviderError, ProviderFetchResult, ProviderId}; + +fn parse(output: &str) -> crate::core::UsageSnapshot { + KiroProvider::new() + .parse_cli_output(output) + .expect("Kiro CLI output should parse") +} + +fn limits(plan_limit: f64, plan_used: f64) -> usage_limits::KiroUsageLimits { + usage_limits::KiroUsageLimits { + plan_limit, + plan_used, + overage_used: 0.0, + overage_cap: None, + overage_enabled: Some(false), + overage_charges: None, + overage_rate: None, + currency_code: "USD".to_string(), + resets_at: Utc.timestamp_opt(1_790_812_800, 0).single().unwrap(), + has_unseparated_bonus: false, + } +} + +#[test] +fn summary_preserves_plan_without_inventing_usage() { + let usage = parse("\u{1b}[32mPlan: KIRO PRO MAX | 1 usage breakdowns\u{1b}[0m\n"); + + assert_eq!(usage.login_method.as_deref(), Some("KIRO PRO MAX")); + assert!(usage.primary.is_informational); + assert_eq!(usage.primary.used_percent, 0.0); + assert_eq!( + usage.primary.reset_description.as_deref(), + Some("Usage unavailable") + ); +} + +#[test] +fn managed_plan_without_metrics_is_unavailable() { + let usage = parse("Plan: Q Developer Pro\nYour plan is managed by admin\n"); + + assert_eq!(usage.login_method.as_deref(), Some("Q Developer Pro")); + assert!(usage.primary.is_informational); +} + +#[test] +fn managed_marker_without_plan_is_not_a_valid_summary() { + assert!( + KiroProvider::new() + .parse_cli_output("Your plan is managed by admin\n") + .is_err() + ); +} + +#[test] +fn malformed_summaries_do_not_become_plan_only_usage() { + for output in [ + "Plan: KIRO PRO MAX", + "Plan: KIRO PRO MAX | usage breakdowns", + "Plan: KIRO PRO MAX | -1 usage breakdowns", + "Plan: KIRO PRO MAX | 1.5 usage breakdowns", + "Plan: KIRO PRO MAX | 1 usage breakdowns failed", + "Plan: | 1 usage breakdowns", + "echo Plan: KIRO PRO MAX | 1 usage breakdowns", + "Plan: KIRO PRO MAX |\n1 usage breakdowns", + ] { + assert!( + KiroProvider::new().parse_cli_output(output).is_err(), + "expected parse failure for {output:?}" + ); + } +} + +#[test] +fn summary_with_real_zero_usage_keeps_available_allowance() { + let usage = parse( + "Plan: KIRO PRO MAX | 1 usage breakdowns\nCredits (0 of 5000 covered in plan)\nresets on 12/31\n", + ); + + assert_eq!(usage.login_method.as_deref(), Some("KIRO PRO MAX")); + assert!(!usage.primary.is_informational); + assert_eq!(usage.primary.used_percent, 0.0); + assert!(usage.primary.resets_at.is_some()); +} + +#[test] +fn cli_brief_summary_reports_unavailable_instead_of_zero() { + let result = + ProviderFetchResult::new(parse("Plan: KIRO PRO MAX | 1 usage breakdowns\n"), "test"); + let brief = crate::cli::usage::render_brief_text(ProviderId::Kiro, &result); + + assert!(brief.contains("Session unavailable")); + assert!(!brief.contains("Session 0%")); + assert!(brief.contains("KIRO PRO MAX")); +} + +#[test] +fn plan_only_summary_keeps_bonus_and_overage_metadata() { + let usage = parse( + "Plan: KIRO PRO MAX | 1 usage breakdowns\n\ + Bonus credits: 10/100 credits used, expires in 3 days\n\ + Overages: Enabled\n\ + Credits used: 4.5\n\ + Est. cost: $1.25 USD\n", + ); + + assert!(usage.primary.is_informational); + assert_eq!( + usage.secondary.as_ref().map(|window| window.used_percent), + Some(10.0) + ); + assert_eq!(usage.extra_rate_windows.len(), 2); + assert!( + usage + .extra_rate_windows + .iter() + .any(|row| row.id == "kiro-overage-credits" && row.window.is_informational) + ); + assert!( + usage + .extra_rate_windows + .iter() + .any(|row| row.id == "kiro-overage-cost" && row.window.is_informational) + ); +} + +#[test] +fn positive_api_allowance_enriches_plan_only_summary() { + let usage = usage_limits::apply_usage_limits( + parse("Plan: KIRO PRO MAX | 1 usage breakdowns\n"), + &limits(5000.0, 282.49), + ); + + assert!(!usage.primary.is_informational); + assert!((usage.primary.used_percent - 5.6498).abs() < 0.0001); + assert_eq!( + usage.primary.resets_at, + Some(limits(5000.0, 282.49).resets_at) + ); + assert!(usage.primary.reset_description.is_none()); +} + +#[test] +fn zero_api_allowance_preserves_unknown_or_cli_metrics() { + let unknown = usage_limits::apply_usage_limits( + parse("Plan: KIRO PRO MAX | 1 usage breakdowns\n"), + &limits(0.0, 0.0), + ); + assert!(unknown.primary.is_informational); + + let known = usage_limits::apply_usage_limits( + parse( + "Plan: KIRO PRO MAX | 1 usage breakdowns\nCredits (20 of 50 covered in plan)\nresets on 12/31\n", + ), + &limits(0.0, 0.0), + ); + assert!(!known.primary.is_informational); + assert_eq!(known.primary.used_percent, 40.0); + assert!(known.primary.resets_at.is_some()); +} + +#[test] +fn cli_presence_maps_to_local_runtime_offline_but_state_db_stays_default() { + assert_eq!( + KiroProvider::new().error_state_kind(&ProviderError::NotInstalled( + "kiro-cli not found. Install from https://kiro.dev".to_string(), + )), + crate::core::ProviderStateKind::LocalRuntimeOffline + ); + // The state-database token lookup is auth-flavored and keeps the + // default mapping. + assert_eq!( + KiroProvider::new().error_state_kind(&ProviderError::NotInstalled( + "Kiro CLI state database not found at C:\\Users\\x\\Kiro-Cli\\data.sqlite3".to_string(), + )), + crate::core::ProviderStateKind::NeedsAuthentication + ); +} diff --git a/rust/src/providers/kiro/usage_limits.rs b/rust/src/providers/kiro/usage_limits.rs index a4a345e399..fcac60011c 100644 --- a/rust/src/providers/kiro/usage_limits.rs +++ b/rust/src/providers/kiro/usage_limits.rs @@ -97,6 +97,8 @@ pub(super) fn apply_usage_limits( usage.primary.used_percent = (limits.plan_used / limits.plan_limit * 100.0).clamp(0.0, 100.0); usage.primary.resets_at = Some(limits.resets_at); + usage.primary.reset_description = None; + usage.primary.is_informational = false; } if limits.overage_enabled == Some(false) { diff --git a/rust/src/providers/minimax/mod.rs b/rust/src/providers/minimax/mod.rs index c6d02d7154..b47992664e 100755 --- a/rust/src/providers/minimax/mod.rs +++ b/rust/src/providers/minimax/mod.rs @@ -1017,6 +1017,10 @@ impl Default for MiniMaxProvider { #[async_trait] impl Provider for MiniMaxProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::MiniMax } diff --git a/rust/src/providers/openai/mod.rs b/rust/src/providers/openai/mod.rs index 03e9085c21..23ac841233 100755 --- a/rust/src/providers/openai/mod.rs +++ b/rust/src/providers/openai/mod.rs @@ -4,6 +4,7 @@ pub mod friendly_errors; pub mod scraper; +pub mod subscription; // Re-exports for error handling and dashboard scraping #[allow( @@ -21,3 +22,7 @@ pub use scraper::{ CreditsHistoryEntry, OPENAI_DASHBOARD_SCRAPE_SCRIPT, OpenAIDashboardData, UsageBreakdown, parse_dashboard_json, }; +pub use subscription::{ + OPENAI_SUBSCRIPTION_CAPTURE_SCRIPT, OpenAISubscriptionFetchResult, account_identity_matches, + parse_subscription_http_response, parse_subscription_json, +}; diff --git a/rust/src/providers/openai/scraper.rs b/rust/src/providers/openai/scraper.rs index 4a8bf95e24..6e0e63aae6 100755 --- a/rust/src/providers/openai/scraper.rs +++ b/rust/src/providers/openai/scraper.rs @@ -1,43 +1,47 @@ -//! OpenAI Dashboard Scraper -//! -//! JavaScript-based scraper for extracting usage data from the OpenAI/ChatGPT dashboard. -//! Uses React Fiber inspection to extract data from chart components. - -use serde::{Deserialize, Serialize}; - -/// Usage breakdown by service (e.g., GPT-4, DALL-E) -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct UsageBreakdown { - /// Service name - pub service: String, - /// Hex color for the service in charts - pub color: String, - /// Usage amount in dollars - pub amount: f64, -} - -/// Credits usage history entry -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct CreditsHistoryEntry { - /// Date string - pub date: String, - /// Description of usage - pub description: String, - /// Amount in dollars (positive = credit, negative = usage) - pub amount: f64, -} - -/// Scraped dashboard data -#[derive(Debug, Clone, Default, Serialize, Deserialize)] -pub struct OpenAIDashboardData { - /// Remaining credits balance - pub credits_remaining: Option, - /// Total credits limit - pub credits_limit: Option, - /// Usage breakdown by service - pub usage_breakdown: Vec, - /// Credits usage history - pub credits_history: Vec, +//! OpenAI Dashboard Scraper +//! +//! JavaScript-based scraper for extracting usage data from the OpenAI/ChatGPT dashboard. +//! Uses React Fiber inspection to extract data from chart components. + +use serde::{Deserialize, Serialize}; + +use super::subscription::{ + OpenAISubscriptionFetchResult, parse_subscription_http_response, parse_subscription_value, +}; + +/// Usage breakdown by service (e.g., GPT-4, DALL-E) +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UsageBreakdown { + /// Service name + pub service: String, + /// Hex color for the service in charts + pub color: String, + /// Usage amount in dollars + pub amount: f64, +} + +/// Credits usage history entry +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CreditsHistoryEntry { + /// Date string + pub date: String, + /// Description of usage + pub description: String, + /// Amount in dollars (positive = credit, negative = usage) + pub amount: f64, +} + +/// Scraped dashboard data +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct OpenAIDashboardData { + /// Remaining credits balance + pub credits_remaining: Option, + /// Total credits limit + pub credits_limit: Option, + /// Usage breakdown by service + pub usage_breakdown: Vec, + /// Credits usage history + pub credits_history: Vec, /// Account email pub email: Option, /// Authentication status from ChatGPT bootstrap data @@ -46,168 +50,223 @@ pub struct OpenAIDashboardData { pub account_plan: Option, /// Organization name pub organization: Option, - /// Purchase credits URL - pub purchase_url: Option, -} - -impl OpenAIDashboardData { - /// Calculate used percentage - pub fn used_percent(&self) -> Option { - let remaining = self.credits_remaining?; - let limit = self.credits_limit?; - if limit <= 0.0 { - return None; - } - let used = limit - remaining; - Some((used / limit) * 100.0) - } - - /// Get total usage across all services - pub fn total_usage(&self) -> f64 { - self.usage_breakdown.iter().map(|b| b.amount).sum() - } -} - -/// JavaScript scrape script for OpenAI dashboard -/// -/// This script is injected into the ChatGPT dashboard page to extract usage data. -/// It uses React Fiber inspection to access chart data that isn't directly in the DOM. -pub const OPENAI_DASHBOARD_SCRAPE_SCRIPT: &str = r#" -(() => { - const textOf = el => { - const raw = el && (el.innerText || el.textContent) ? String(el.innerText || el.textContent) : ''; - return raw.trim(); - }; - - const parseHexColor = (color) => { - if (!color) return null; - const c = String(color).trim().toLowerCase(); - if (c.startsWith('#')) { - if (c.length === 4) { - return '#' + c[1] + c[1] + c[2] + c[2] + c[3] + c[3]; - } - if (c.length === 7) return c; - return c; - } - const m = c.match(/^rgba?\(([^)]+)\)$/); - if (m) { - const parts = m[1].split(',').map(x => parseFloat(x.trim())).filter(x => Number.isFinite(x)); - if (parts.length >= 3) { - const r = Math.max(0, Math.min(255, Math.round(parts[0]))); - const g = Math.max(0, Math.min(255, Math.round(parts[1]))); - const b = Math.max(0, Math.min(255, Math.round(parts[2]))); - const toHex = n => n.toString(16).padStart(2, '0'); - return '#' + toHex(r) + toHex(g) + toHex(b); - } - } - return c; - }; - - // React Fiber inspection for extracting chart data - const reactPropsOf = (el) => { - if (!el) return null; - try { - const keys = Object.keys(el); - const propsKey = keys.find(k => k.startsWith('__reactProps$')); - if (propsKey) return el[propsKey] || null; - const fiberKey = keys.find(k => k.startsWith('__reactFiber$')); - if (fiberKey) { - const fiber = el[fiberKey]; - return (fiber && (fiber.memoizedProps || fiber.pendingProps)) || null; - } - } catch {} - return null; - }; - - const reactFiberOf = (el) => { - if (!el) return null; - try { - const keys = Object.keys(el); - const fiberKey = keys.find(k => k.startsWith('__reactFiber$')); - return fiberKey ? (el[fiberKey] || null) : null; - } catch { - return null; - } - }; - - // Traverse React Fiber tree to find chart payload data - const nestedBarMetaOf = (root) => { - if (!root || typeof root !== 'object') return null; - const queue = [root]; - const seen = typeof WeakSet !== 'undefined' ? new WeakSet() : null; - let steps = 0; - while (queue.length && steps < 250) { - const cur = queue.shift(); - steps++; - if (!cur || typeof cur !== 'object') continue; - if (seen) { - if (seen.has(cur)) continue; - seen.add(cur); - } - if (cur.payload && (cur.dataKey || cur.name || cur.value !== undefined)) return cur; - const values = Array.isArray(cur) ? cur : Object.values(cur); - for (const v of values) { - if (v && typeof v === 'object') queue.push(v); - } - } - return null; - }; - - // Extract chart metadata from DOM element via React Fiber - const barMetaFromElement = (el) => { - const direct = reactPropsOf(el); - if (direct && direct.payload && (direct.dataKey || direct.name || direct.value !== undefined)) return direct; - - const fiber = reactFiberOf(el); - if (fiber) { - let cur = fiber; - for (let i = 0; i < 10 && cur; i++) { - const props = (cur.memoizedProps || cur.pendingProps) || null; - if (props && props.payload && (props.dataKey || props.name || props.value !== undefined)) return props; - const nested = props ? nestedBarMetaOf(props) : null; - if (nested) return nested; - cur = cur.return || null; - } - } - - if (direct) { - const nested = nestedBarMetaOf(direct); - if (nested) return nested; - } - return null; - }; - - // Parse dollar amounts from text - const parseDollarAmount = (text) => { - if (!text) return null; - const cleaned = String(text).replace(/[^0-9.,\-]/g, ''); - const num = parseFloat(cleaned.replace(',', '')); - return Number.isFinite(num) ? num : null; - }; - - // Find credits remaining - const findCreditsRemaining = () => { - const patterns = [ - /\$?(\d+(?:\.\d+)?)\s*(?:credits?)?\s*(?:remaining|left|available)/i, - /(?:remaining|left|available)[:\s]*\$?(\d+(?:\.\d+)?)/i, - /balance[:\s]*\$?(\d+(?:\.\d+)?)/i, - ]; - - const textNodes = document.querySelectorAll('*'); - for (const node of textNodes) { - const text = textOf(node); - for (const pattern of patterns) { - const match = text.match(pattern); - if (match) { - const num = parseFloat(match[1]); - if (Number.isFinite(num)) return num; - } - } - } - return null; - }; - - // Find account email + /// Purchase credits URL + pub purchase_url: Option, + /// Explicit subscription payload captured by the dashboard page. This is + /// kept as a small JSON object so missing fields remain distinguishable + /// from explicit nulls until the strict subscription parser runs. + #[serde(default)] + pub subscription: Option, +} + +impl OpenAIDashboardData { + /// Calculate used percentage + pub fn used_percent(&self) -> Option { + let remaining = self.credits_remaining?; + let limit = self.credits_limit?; + if limit <= 0.0 { + return None; + } + let used = limit - remaining; + Some((used / limit) * 100.0) + } + + /// Get total usage across all services + pub fn total_usage(&self) -> f64 { + self.usage_breakdown.iter().map(|b| b.amount).sum() + } + + /// Convert the page-captured subscription payload only when its lifecycle + /// fields are present and correctly typed. + pub fn subscription_metadata(&self) -> OpenAISubscriptionFetchResult { + let Some(subscription) = self.subscription.as_ref() else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let Some(captured) = subscription.as_object() else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + if let (Some(status), Some(payload)) = ( + captured.get("status").and_then(serde_json::Value::as_u64), + captured.get("payload"), + ) { + let Ok(status) = u16::try_from(status) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + let Ok(payload) = serde_json::to_string(payload) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + return parse_subscription_http_response(status, &payload); + } + // Accept direct payloads for callers that already checked the HTTP + // status before constructing the dashboard DTO. + parse_subscription_value(subscription) + } + + /// Return subscription metadata only when the dashboard identity is + /// unambiguous for the account the caller is refreshing. A target account + /// without a matching page email is deliberately fail-closed. + pub fn authorized_subscription_metadata( + &self, + target_email: Option<&str>, + ) -> OpenAISubscriptionFetchResult { + let target = target_email + .map(str::trim) + .filter(|email| !email.is_empty()) + .map(str::to_ascii_lowercase); + let dashboard = self + .email + .as_deref() + .map(str::trim) + .filter(|email| !email.is_empty()) + .map(str::to_ascii_lowercase); + match (target.as_deref(), dashboard.as_deref()) { + (Some(target), Some(dashboard)) if target == dashboard => self.subscription_metadata(), + (None, Some(_)) => self.subscription_metadata(), + _ => OpenAISubscriptionFetchResult::Unavailable, + } + } +} + +/// JavaScript scrape script for OpenAI dashboard +/// +/// This script is injected into the ChatGPT dashboard page to extract usage data. +/// It uses React Fiber inspection to access chart data that isn't directly in the DOM. +pub const OPENAI_DASHBOARD_SCRAPE_SCRIPT: &str = r#" +(() => { + const textOf = el => { + const raw = el && (el.innerText || el.textContent) ? String(el.innerText || el.textContent) : ''; + return raw.trim(); + }; + + const parseHexColor = (color) => { + if (!color) return null; + const c = String(color).trim().toLowerCase(); + if (c.startsWith('#')) { + if (c.length === 4) { + return '#' + c[1] + c[1] + c[2] + c[2] + c[3] + c[3]; + } + if (c.length === 7) return c; + return c; + } + const m = c.match(/^rgba?\(([^)]+)\)$/); + if (m) { + const parts = m[1].split(',').map(x => parseFloat(x.trim())).filter(x => Number.isFinite(x)); + if (parts.length >= 3) { + const r = Math.max(0, Math.min(255, Math.round(parts[0]))); + const g = Math.max(0, Math.min(255, Math.round(parts[1]))); + const b = Math.max(0, Math.min(255, Math.round(parts[2]))); + const toHex = n => n.toString(16).padStart(2, '0'); + return '#' + toHex(r) + toHex(g) + toHex(b); + } + } + return c; + }; + + // React Fiber inspection for extracting chart data + const reactPropsOf = (el) => { + if (!el) return null; + try { + const keys = Object.keys(el); + const propsKey = keys.find(k => k.startsWith('__reactProps$')); + if (propsKey) return el[propsKey] || null; + const fiberKey = keys.find(k => k.startsWith('__reactFiber$')); + if (fiberKey) { + const fiber = el[fiberKey]; + return (fiber && (fiber.memoizedProps || fiber.pendingProps)) || null; + } + } catch {} + return null; + }; + + const reactFiberOf = (el) => { + if (!el) return null; + try { + const keys = Object.keys(el); + const fiberKey = keys.find(k => k.startsWith('__reactFiber$')); + return fiberKey ? (el[fiberKey] || null) : null; + } catch { + return null; + } + }; + + // Traverse React Fiber tree to find chart payload data + const nestedBarMetaOf = (root) => { + if (!root || typeof root !== 'object') return null; + const queue = [root]; + const seen = typeof WeakSet !== 'undefined' ? new WeakSet() : null; + let steps = 0; + while (queue.length && steps < 250) { + const cur = queue.shift(); + steps++; + if (!cur || typeof cur !== 'object') continue; + if (seen) { + if (seen.has(cur)) continue; + seen.add(cur); + } + if (cur.payload && (cur.dataKey || cur.name || cur.value !== undefined)) return cur; + const values = Array.isArray(cur) ? cur : Object.values(cur); + for (const v of values) { + if (v && typeof v === 'object') queue.push(v); + } + } + return null; + }; + + // Extract chart metadata from DOM element via React Fiber + const barMetaFromElement = (el) => { + const direct = reactPropsOf(el); + if (direct && direct.payload && (direct.dataKey || direct.name || direct.value !== undefined)) return direct; + + const fiber = reactFiberOf(el); + if (fiber) { + let cur = fiber; + for (let i = 0; i < 10 && cur; i++) { + const props = (cur.memoizedProps || cur.pendingProps) || null; + if (props && props.payload && (props.dataKey || props.name || props.value !== undefined)) return props; + const nested = props ? nestedBarMetaOf(props) : null; + if (nested) return nested; + cur = cur.return || null; + } + } + + if (direct) { + const nested = nestedBarMetaOf(direct); + if (nested) return nested; + } + return null; + }; + + // Parse dollar amounts from text + const parseDollarAmount = (text) => { + if (!text) return null; + const cleaned = String(text).replace(/[^0-9.,\-]/g, ''); + const num = parseFloat(cleaned.replace(',', '')); + return Number.isFinite(num) ? num : null; + }; + + // Find credits remaining + const findCreditsRemaining = () => { + const patterns = [ + /\$?(\d+(?:\.\d+)?)\s*(?:credits?)?\s*(?:remaining|left|available)/i, + /(?:remaining|left|available)[:\s]*\$?(\d+(?:\.\d+)?)/i, + /balance[:\s]*\$?(\d+(?:\.\d+)?)/i, + ]; + + const textNodes = document.querySelectorAll('*'); + for (const node of textNodes) { + const text = textOf(node); + for (const pattern of patterns) { + const match = text.match(pattern); + if (match) { + const num = parseFloat(match[1]); + if (Number.isFinite(num)) return num; + } + } + } + return null; + }; + + // Find account email const findEmail = () => { const bootstrap = parseJsonScript('client-bootstrap'); const bootstrapEmail = bootstrap?.session?.user?.email || bootstrap?.user?.email || null; @@ -216,15 +275,15 @@ pub const OPENAI_DASHBOARD_SCRAPE_SCRIPT: &str = r#" const next = parseJsonScript('__NEXT_DATA__'); const nextEmail = next?.props?.pageProps?.user?.email || next?.props?.session?.user?.email || null; if (nextEmail && String(nextEmail).includes('@')) return String(nextEmail); - - // Look for email patterns in the page - const emailPattern = /[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}/; - const textNodes = document.querySelectorAll('[class*="email"], [class*="user"], [data-testid*="email"]'); - for (const node of textNodes) { - const text = textOf(node); - const match = text.match(emailPattern); - if (match) return match[0]; - } + + // Look for email patterns in the page + const emailPattern = /[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}/; + const textNodes = document.querySelectorAll('[class*="email"], [class*="user"], [data-testid*="email"]'); + for (const node of textNodes) { + const text = textOf(node); + const match = text.match(emailPattern); + if (match) return match[0]; + } return null; }; @@ -319,35 +378,35 @@ pub const OPENAI_DASHBOARD_SCRAPE_SCRIPT: &str = r#" } return null; }; - - // Extract usage breakdown from chart - const extractUsageBreakdown = () => { - const breakdown = []; - - // Find Recharts bar elements - const bars = document.querySelectorAll('.recharts-bar-rectangle, [class*="bar"]'); - for (const bar of bars) { - const meta = barMetaFromElement(bar); - if (meta && meta.payload) { - const name = meta.name || meta.dataKey || 'Unknown'; - const value = meta.value || meta.payload[meta.dataKey] || 0; - const color = parseHexColor(bar.getAttribute('fill')) || '#888888'; - if (value > 0) { - breakdown.push({ service: name, color, amount: value }); - } - } - } - - // Dedupe by service name - const seen = new Set(); - return breakdown.filter(b => { - if (seen.has(b.service)) return false; - seen.add(b.service); - return true; - }); - }; - - // Main scrape function + + // Extract usage breakdown from chart + const extractUsageBreakdown = () => { + const breakdown = []; + + // Find Recharts bar elements + const bars = document.querySelectorAll('.recharts-bar-rectangle, [class*="bar"]'); + for (const bar of bars) { + const meta = barMetaFromElement(bar); + if (meta && meta.payload) { + const name = meta.name || meta.dataKey || 'Unknown'; + const value = meta.value || meta.payload[meta.dataKey] || 0; + const color = parseHexColor(bar.getAttribute('fill')) || '#888888'; + if (value > 0) { + breakdown.push({ service: name, color, amount: value }); + } + } + } + + // Dedupe by service name + const seen = new Set(); + return breakdown.filter(b => { + if (seen.has(b.service)) return false; + seen.add(b.service); + return true; + }); + }; + + // Main scrape function const result = { credits_remaining: findCreditsRemaining(), credits_limit: null, @@ -360,67 +419,79 @@ pub const OPENAI_DASHBOARD_SCRAPE_SCRIPT: &str = r#" })(), account_plan: findPlan(parseJsonScript('client-bootstrap')) || findPlan(parseJsonScript('__NEXT_DATA__')), organization: null, - purchase_url: null + purchase_url: null, + subscription: (() => { + const captured = window.__codexbarSubscriptionResponse; + const payload = captured && captured.payload; + if (!payload || typeof payload !== 'object' || Array.isArray(payload)) return null; + const keys = ['active_until', 'activeUntil', 'will_renew', 'willRenew', + 'starts_at', 'startsAt', 'active_from', 'activeFrom']; + const output = {}; + for (const key of keys) { + if (Object.prototype.hasOwnProperty.call(payload, key)) output[key] = payload[key]; + } + return Object.keys(output).length ? {status: captured.status, payload: output} : null; + })() }; - - return JSON.stringify(result); -})(); -"#; - -/// Parse scraped JSON data into structured format -pub fn parse_dashboard_json(json: &str) -> Result { - serde_json::from_str(json) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_usage_breakdown() { - let breakdown = UsageBreakdown { - service: "GPT-4".to_string(), - color: "#10a37f".to_string(), - amount: 15.50, - }; - - assert_eq!(breakdown.service, "GPT-4"); - assert_eq!(breakdown.amount, 15.50); - } - - #[test] - fn test_dashboard_data_used_percent() { - let data = OpenAIDashboardData { - credits_remaining: Some(75.0), - credits_limit: Some(100.0), - ..Default::default() - }; - - assert_eq!(data.used_percent(), Some(25.0)); - } - - #[test] - fn test_dashboard_data_total_usage() { - let data = OpenAIDashboardData { - usage_breakdown: vec![ - UsageBreakdown { - service: "GPT-4".to_string(), - color: "#10a37f".to_string(), - amount: 10.0, - }, - UsageBreakdown { - service: "DALL-E".to_string(), - color: "#ff6b6b".to_string(), - amount: 5.0, - }, - ], - ..Default::default() - }; - - assert_eq!(data.total_usage(), 15.0); - } - - #[test] + + return JSON.stringify(result); +})(); +"#; + +/// Parse scraped JSON data into structured format +pub fn parse_dashboard_json(json: &str) -> Result { + serde_json::from_str(json) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_usage_breakdown() { + let breakdown = UsageBreakdown { + service: "GPT-4".to_string(), + color: "#10a37f".to_string(), + amount: 15.50, + }; + + assert_eq!(breakdown.service, "GPT-4"); + assert_eq!(breakdown.amount, 15.50); + } + + #[test] + fn test_dashboard_data_used_percent() { + let data = OpenAIDashboardData { + credits_remaining: Some(75.0), + credits_limit: Some(100.0), + ..Default::default() + }; + + assert_eq!(data.used_percent(), Some(25.0)); + } + + #[test] + fn test_dashboard_data_total_usage() { + let data = OpenAIDashboardData { + usage_breakdown: vec![ + UsageBreakdown { + service: "GPT-4".to_string(), + color: "#10a37f".to_string(), + amount: 10.0, + }, + UsageBreakdown { + service: "DALL-E".to_string(), + color: "#ff6b6b".to_string(), + amount: 5.0, + }, + ], + ..Default::default() + }; + + assert_eq!(data.total_usage(), 15.0); + } + + #[test] fn test_parse_dashboard_json() { let json = r#"{"credits_remaining":50.0,"credits_limit":100.0,"usage_breakdown":[],"credits_history":[],"email":"test@example.com","auth_status":"logged_in","account_plan":"Pro 5x","organization":null,"purchase_url":null}"#; @@ -430,4 +501,27 @@ mod tests { assert_eq!(data.auth_status, Some("logged_in".to_string())); assert_eq!(data.account_plan, Some("Pro 5x".to_string())); } + + #[test] + fn dashboard_subscription_requires_matching_account_identity() { + let subscription = serde_json::json!({ + "active_until": "2026-09-20T14:30:07Z", + "will_renew": true + }); + let data = OpenAIDashboardData { + email: Some("current@example.com".to_string()), + subscription: Some(subscription), + ..Default::default() + }; + assert!( + data.authorized_subscription_metadata(Some("old@example.com")) + .metadata() + .is_none() + ); + assert!(data.authorized_subscription_metadata(None).succeeded()); + assert!( + data.authorized_subscription_metadata(Some("CURRENT@example.com")) + .succeeded() + ); + } } diff --git a/rust/src/providers/openai/subscription.rs b/rust/src/providers/openai/subscription.rs new file mode 100644 index 0000000000..1b9be905a2 --- /dev/null +++ b/rust/src/providers/openai/subscription.rs @@ -0,0 +1,324 @@ +//! Authenticated OpenAI subscription metadata parsing and page-capture support. +//! +//! The dashboard is the authority for these dates. The parser accepts only +//! explicit, typed values from the subscription response and never derives a +//! subscription boundary from a quota reset, plan name, or current time. + +use chrono::{DateTime, Utc}; +use serde_json::Value; + +use crate::core::SubscriptionMetadata; + +/// Result of a subscription request. `Success(None)` is distinct from an +/// unavailable/invalid response so callers may clear dates only after a +/// valid, authoritative response. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum OpenAISubscriptionFetchResult { + Unavailable, + Success(Option), +} + +impl OpenAISubscriptionFetchResult { + pub fn metadata(&self) -> Option<&SubscriptionMetadata> { + match self { + Self::Success(metadata) => metadata.as_ref(), + Self::Unavailable => None, + } + } + + pub fn cloned_metadata(&self) -> Option { + self.metadata().cloned() + } + + pub const fn succeeded(&self) -> bool { + matches!(self, Self::Success(_)) + } +} + +/// Page-side capture hook for WebView2/Tauri dashboard sessions. +/// +/// It is intentionally limited to same-origin `/backend-api/subscriptions` +/// responses. Generation and request guards prevent a late response from a +/// previous navigation or an older concurrent request from being published. +pub const OPENAI_SUBSCRIPTION_CAPTURE_SCRIPT: &str = r#" +(() => { + if (window.__codexbarSubscriptionCaptureInstalled) return; + window.__codexbarSubscriptionCaptureInstalled = true; + window.__codexbarSubscriptionCaptureGeneration = 0; + window.__codexbarSubscriptionResponse = null; + let latestRequest = 0; + const originalFetch = window.fetch.bind(window); + window.fetch = async (...args) => { + const generation = Number(window.__codexbarSubscriptionCaptureGeneration || 0); + let request = null; + try { + const input = args[0]; + const rawUrl = input && input.url ? input.url : input; + const url = new URL(String(rawUrl), window.location.href); + if (url.origin === window.location.origin && + url.pathname === '/backend-api/subscriptions') { + request = ++latestRequest; + } + } catch (_) {} + const response = await originalFetch(...args); + if (request !== null) { + try { + const payload = await response.clone().json(); + if (generation !== Number(window.__codexbarSubscriptionCaptureGeneration || 0) || + request !== latestRequest) return response; + window.__codexbarSubscriptionResponse = { + status: response.status, + payload: payload + }; + } catch (_) { + if (generation === Number(window.__codexbarSubscriptionCaptureGeneration || 0) && + request === latestRequest) { + window.__codexbarSubscriptionResponse = { + status: response.status, + payload: null + }; + } + } + } + return response; + }; +})(); +"#; + +/// Bump the capture generation before a dashboard navigation or account +/// change. A caller should evaluate this in the same WebView2 page context. +pub const OPENAI_SUBSCRIPTION_RESET_SCRIPT: &str = r#" +(() => { + window.__codexbarSubscriptionCaptureGeneration = + Number(window.__codexbarSubscriptionCaptureGeneration || 0) + 1; + window.__codexbarSubscriptionResponse = null; + return window.__codexbarSubscriptionCaptureGeneration; +})(); +"#; + +/// Read the latest same-origin response captured by +/// [`OPENAI_SUBSCRIPTION_CAPTURE_SCRIPT`]. +pub const OPENAI_SUBSCRIPTION_READ_SCRIPT: &str = r#" +(() => window.__codexbarSubscriptionResponse || null)(); +"#; + +/// Compare a managed-account hint with the identity observed from the same +/// authenticated credentials. A supplied hint without a matching observed +/// email is ambiguous and cannot authorize subscription metadata. +pub fn account_identity_matches( + expected_email: Option<&str>, + observed_email: Option<&str>, +) -> bool { + let expected = expected_email + .map(str::trim) + .filter(|email| !email.is_empty()) + .map(str::to_ascii_lowercase); + let observed = observed_email + .map(str::trim) + .filter(|email| !email.is_empty()) + .map(str::to_ascii_lowercase); + match (expected.as_deref(), observed.as_deref()) { + (Some(expected), Some(observed)) => expected == observed, + (Some(_), None) => false, + (None, _) => true, + } +} + +pub fn parse_subscription_json(json: &str) -> OpenAISubscriptionFetchResult { + let Ok(value) = serde_json::from_str::(json) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + parse_subscription_value(&value) +} + +/// Parse a response captured by a dashboard page or fetched directly from the +/// authenticated API. Non-success responses are never interpreted as an +/// empty subscription. +pub fn parse_subscription_http_response(status: u16, body: &str) -> OpenAISubscriptionFetchResult { + if !(200..300).contains(&status) { + return OpenAISubscriptionFetchResult::Unavailable; + } + parse_subscription_json(body) +} + +/// Parse a successful JSON response. Both lifecycle fields must be present and +/// correctly typed; otherwise the result is unavailable and the caller must +/// retain any prior dates. +pub fn parse_subscription_value(value: &Value) -> OpenAISubscriptionFetchResult { + let Some(object) = value.as_object() else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + + let active_until = object + .get("active_until") + .or_else(|| object.get("activeUntil")); + let will_renew = object.get("will_renew").or_else(|| object.get("willRenew")); + let (Some(active_until), Some(will_renew)) = (active_until, will_renew) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + + let active_until = match active_until { + Value::Null => None, + Value::String(value) => Some(value.as_str()), + _ => return OpenAISubscriptionFetchResult::Unavailable, + }; + let will_renew = match will_renew { + Value::Null => None, + Value::Bool(value) => Some(*value), + _ => return OpenAISubscriptionFetchResult::Unavailable, + }; + + let starts_at = parse_optional_date( + object, + &["starts_at", "startsAt", "active_from", "activeFrom"], + ); + if has_any( + object, + &["starts_at", "startsAt", "active_from", "activeFrom"], + ) && starts_at.is_err() + { + return OpenAISubscriptionFetchResult::Unavailable; + } + let starts_at = starts_at.ok().flatten(); + + match (active_until, will_renew) { + (Some(active_until), Some(true)) => { + let Some(renews_at) = parse_date(active_until) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + OpenAISubscriptionFetchResult::Success(Some(SubscriptionMetadata::new( + starts_at, + None, + Some(renews_at), + ))) + } + (Some(active_until), Some(false)) => { + let Some(expires_at) = parse_date(active_until) else { + return OpenAISubscriptionFetchResult::Unavailable; + }; + OpenAISubscriptionFetchResult::Success(Some(SubscriptionMetadata::new( + starts_at, + Some(expires_at), + None, + ))) + } + (None, Some(false)) => OpenAISubscriptionFetchResult::Success( + (!starts_at.is_none()).then(|| SubscriptionMetadata::new(starts_at, None, None)), + ), + // A renewal without an explicit active-until date is not safe to + // represent as a date, even if a plan is known. + (None, Some(true)) | (_, None) => OpenAISubscriptionFetchResult::Unavailable, + } +} + +fn has_any(object: &serde_json::Map, keys: &[&str]) -> bool { + keys.iter().any(|key| object.contains_key(*key)) +} + +fn parse_optional_date( + object: &serde_json::Map, + keys: &[&str], +) -> Result>, ()> { + let Some(value) = keys.iter().find_map(|key| object.get(*key)) else { + return Ok(None); + }; + match value { + Value::Null => Ok(None), + Value::String(value) => parse_date(value).ok_or(()).map(Some), + _ => Err(()), + } +} + +fn parse_date(value: &str) -> Option> { + DateTime::parse_from_rfc3339(value) + .ok() + .map(|value| value.with_timezone(&Utc)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_renewal_and_preserves_explicit_start() { + let result = parse_subscription_json( + r#"{"active_until":"2026-09-20T14:30:07.123Z","will_renew":true,"active_from":"2026-08-20T14:30:07Z"}"#, + ); + let OpenAISubscriptionFetchResult::Success(Some(metadata)) = result else { + panic!("expected subscription metadata") + }; + assert_eq!( + metadata.starts_at.unwrap().to_rfc3339(), + "2026-08-20T14:30:07+00:00" + ); + assert_eq!(metadata.expires_at, None); + assert_eq!( + metadata + .renews_at + .unwrap() + .to_rfc3339_opts(chrono::SecondsFormat::Millis, true), + "2026-09-20T14:30:07.123Z" + ); + } + + #[test] + fn cancellation_maps_active_until_to_expiration() { + let result = parse_subscription_json( + r#"{"active_until":"2026-09-20T14:30:07Z","will_renew":false}"#, + ); + let metadata = result.metadata().expect("metadata"); + assert!(metadata.renews_at.is_none()); + assert_eq!( + metadata.expires_at.unwrap().to_rfc3339(), + "2026-09-20T14:30:07+00:00" + ); + } + + #[test] + fn rejects_missing_or_malformed_dates_and_flags_without_fallback() { + for json in [ + r#"{"active_until":null,"will_renew":true}"#, + r#"{"active_until":"not-a-date","will_renew":false}"#, + r#"{"active_until":null,"will_renew":0}"#, + r#"{"active_until":"2026-09-20T14:30:07Z"}"#, + ] { + assert_eq!( + parse_subscription_json(json), + OpenAISubscriptionFetchResult::Unavailable + ); + } + } + + #[test] + fn valid_empty_cancellation_is_success_without_inventing_dates() { + assert_eq!( + parse_subscription_json(r#"{"active_until":null,"will_renew":false}"#), + OpenAISubscriptionFetchResult::Success(None) + ); + } + + #[test] + fn non_success_dashboard_response_cannot_clear_dates() { + assert_eq!( + parse_subscription_http_response(403, r#"{"active_until":null,"will_renew":false}"#), + OpenAISubscriptionFetchResult::Unavailable + ); + } + + #[test] + fn managed_identity_mismatch_fails_closed() { + assert!(account_identity_matches( + Some("A@EXAMPLE.COM"), + Some("a@example.com") + )); + assert!(!account_identity_matches( + Some("old@example.com"), + Some("new@example.com") + )); + assert!(!account_identity_matches( + Some("expected@example.com"), + None + )); + assert!(account_identity_matches(None, Some("observed@example.com"))); + } +} diff --git a/rust/src/providers/perplexity/mod.rs b/rust/src/providers/perplexity/mod.rs index f244388cbf..3577166dcb 100644 --- a/rust/src/providers/perplexity/mod.rs +++ b/rust/src/providers/perplexity/mod.rs @@ -215,6 +215,10 @@ impl Default for PerplexityProvider { #[async_trait] impl Provider for PerplexityProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Perplexity } diff --git a/rust/src/providers/zai/mod.rs b/rust/src/providers/zai/mod.rs index 7e5a11cd8b..aca7e29702 100755 --- a/rust/src/providers/zai/mod.rs +++ b/rust/src/providers/zai/mod.rs @@ -575,6 +575,10 @@ impl Default for ZaiProvider { #[async_trait] impl Provider for ZaiProvider { + fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { + false + } + fn id(&self) -> ProviderId { ProviderId::Zai }