diff --git a/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs b/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs index 179c84ada1..a0c1dbca34 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/codex_accounts.rs @@ -6,7 +6,7 @@ use uuid::Uuid; use codexbar::codex_accounts::{ AccountStore, CodexAccount, CodexAccountApi, CodexAccountManager, CodexAccountManagerError, - CodexApiError, CodexSwitchResult, SnapshotStore, restart_codex_desktop, + CodexApiError, CodexSwitchResult, SnapshotStore, display_names_by_id, restart_codex_desktop, }; use crate::state::AppState; @@ -310,6 +310,7 @@ fn into_api_message(error: CodexApiError) -> String { #[serde(rename_all = "camelCase")] pub struct CodexAccountsStateBridge { pub accounts: Vec, + pub display_names: HashMap, pub snapshots: HashMap, } @@ -318,8 +319,10 @@ pub fn get_codex_accounts_state( state: tauri::State<'_, Mutex>, ) -> Result { let _guard = state.lock().map_err(|e| e.to_string())?; + let accounts = load_codex_accounts()?; Ok(CodexAccountsStateBridge { - accounts: load_codex_accounts()?, + display_names: display_names_by_id(&accounts), + accounts, snapshots: codex_account_snapshots()?, }) } diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index db2d49458a..04c0b620ef 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -126,7 +126,7 @@ fn build_usage_spend_summary_cached( } // Hold the cache mutex while building: callers for the same app revision // coalesce behind this single scan instead of starting parallel rescans. - let summary = build_usage_spend_summary(cached, selected_days, &settings); + let summary = build_usage_spend_summary(cached, selected_days, &settings, force_refresh); *guard = Some(CachedUsageSpendSummary { key, summary: summary.clone(), @@ -196,6 +196,7 @@ fn build_usage_spend_summary( cached: &[ProviderUsageSnapshot], selected_days: u32, settings: &codexbar::settings::Settings, + force_refresh: bool, ) -> UsageSpendSummary { let include_opencodex = settings.open_codex_usage_logs_enabled; let hide_native = settings.hide_native_codex_cost_when_open_codex_present; @@ -203,12 +204,21 @@ fn build_usage_spend_summary( // Upstream 0.55.0 #3105: independent provider baselines load in parallel. // Keep each provider's 7d/30d scans serial so they can safely share that // provider's incremental cache, while Codex and Claude run concurrently. + let codex_scan_options = if force_refresh { + codexbar::core::CostScanOptions::app_driven() + } else { + codexbar::core::CostScanOptions::default() + }; let ((codex_7_summary, codex_30_summary), (claude_7_summary, claude_30_summary)) = std::thread::scope(|scope| { - let codex = scope.spawn(|| { + let codex = scope.spawn(move || { ( - CostScanner::new(7).scan_codex(), - CostScanner::new(30).scan_codex(), + CostScanner::new(7) + .with_options(codex_scan_options) + .scan_codex(), + CostScanner::new(30) + .with_options(codex_scan_options) + .scan_codex(), ) }); let claude = scope.spawn(|| { @@ -428,7 +438,9 @@ fn build_usage_spend_summary( let selected_summary: CostSummary = match history_days { 7 => codex_7_summary, 30 => codex_30_summary, - days => CostScanner::new(days).scan_codex(), + days => CostScanner::new(days) + .with_options(codex_scan_options) + .scan_codex(), }; let contract = build_local_spend_contract_from_summary( "codex", diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index 4838ee757b..48a84426de 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -259,6 +259,21 @@ mod tests { ); } + #[test] + fn single_meaningful_quota_omits_the_companion_icon_lane() { + let mut snapshot = snapshot(); + snapshot + .secondary + .as_mut() + .expect("fixture has a secondary window") + .is_informational = true; + + let (selected, companion) = selected_usage_icon_windows(&snapshot, &Settings::default()); + + assert_eq!(selected.used_percent, 20.0); + assert!(companion.is_none()); + } + #[test] fn average_preference_derives_the_combined_percentage() { let mut snapshot = snapshot(); diff --git a/apps/desktop-tauri/src/components/CodexAccountsMenu.tsx b/apps/desktop-tauri/src/components/CodexAccountsMenu.tsx index 643d69ee00..b216c8f325 100644 --- a/apps/desktop-tauri/src/components/CodexAccountsMenu.tsx +++ b/apps/desktop-tauri/src/components/CodexAccountsMenu.tsx @@ -8,6 +8,7 @@ import type { import { useLocale } from "../hooks/useLocale"; import { useFormattedResetTime } from "../hooks/useFormattedResetTime"; import { maskEmail } from "./MenuCard"; +import { buildCodexAccountDisplayNames } from "./codexAccountDisplay"; import { codexAccountSwitch, getCodexAccountsState, @@ -35,6 +36,7 @@ export default function CodexAccountsMenu({ const [snapshots, setSnapshots] = useState< Record >({}); + const [displayNames, setDisplayNames] = useState>({}); const [busy, setBusy] = useState(false); const [error, setError] = useState(null); @@ -44,6 +46,7 @@ export default function CodexAccountsMenu({ try { const next: CodexAccountsStateBridge = await getCodexAccountsState(); setAccounts(next.accounts); + setDisplayNames(next.displayNames ?? {}); setSnapshots(next.snapshots); } catch (err: unknown) { setError(err instanceof Error ? err.message : String(err)); @@ -86,6 +89,11 @@ export default function CodexAccountsMenu({ return null; } + const accountDisplayNames = buildCodexAccountDisplayNames( + accounts, + displayNames, + ); + return (
@@ -103,6 +111,7 @@ export default function CodexAccountsMenu({ key={account.id} account={account} snapshot={snapshots[account.id]} + displayName={accountDisplayNames[account.id]} hideEmail={hideEmail} resetTimeRelative={resetTimeRelative} busy={busy} @@ -117,6 +126,7 @@ export default function CodexAccountsMenu({ function CodexAccountRow({ account, snapshot, + displayName, hideEmail, resetTimeRelative, busy, @@ -124,6 +134,7 @@ function CodexAccountRow({ }: { account: CodexAccount; snapshot: CodexAccountUsageSnapshot | undefined; + displayName: string; hideEmail: boolean; resetTimeRelative: boolean; busy: boolean; @@ -147,12 +158,13 @@ function CodexAccountRow({ : `${t("MetricResetsIn")} ${resetText}` : null; const windowLabel = formatWindowLabel(usageWindow?.limitWindowSeconds); - const label = - account.nickname ?? - account.emailHint ?? - account.authSubject ?? - shrink(account.id); - const shown = hideEmail ? maskEmail(label) : label; + // Only mask labels that actually contain an email. Generic/nickname labels + // have no personal data to hide, and masking them would erase their opaque + // workspace suffix and make distinct accounts look identical. + const shown = + hideEmail && displayName.includes("@") + ? maskEmail(displayName) + : displayName; const isAmbient = account.source === "ambient"; return ( @@ -212,7 +224,3 @@ function formatWindowLabel( } return null; } - -function shrink(id: string): string { - return id.length <= 12 ? id : `${id.slice(0, 8)}…`; -} diff --git a/apps/desktop-tauri/src/components/codexAccountDisplay.test.ts b/apps/desktop-tauri/src/components/codexAccountDisplay.test.ts new file mode 100644 index 0000000000..647558f5c7 --- /dev/null +++ b/apps/desktop-tauri/src/components/codexAccountDisplay.test.ts @@ -0,0 +1,85 @@ +import { describe, expect, it } from "vitest"; +import type { CodexAccount } from "../types/bridge"; +import { buildCodexAccountDisplayNames } from "./codexAccountDisplay"; + +function account(id: string, providerAccountId: string): CodexAccount { + return { + id, + nickname: null, + emailHint: "same@example.com", + authSubject: null, + providerAccountId, + codexHomePath: `C:/private/${providerAccountId}`, + source: "managedByApp", + createdAt: "2024-01-01T00:00:00Z", + updatedAt: "2024-01-01T00:00:00Z", + lastAuthenticatedAt: null, + }; +} + +describe("Codex account display labels", () => { + it("keeps same-email workspace labels opaque and stable across ordering", () => { + const first = account( + "11111111-1111-1111-1111-111111111111", + "workspace-alpha", + ); + const second = account( + "22222222-2222-2222-2222-222222222222", + "workspace-beta", + ); + + const labels = buildCodexAccountDisplayNames([first, second]); + const reversed = buildCodexAccountDisplayNames([second, first]); + + expect(labels[first.id]).not.toBe(labels[second.id]); + expect(labels[first.id]).toMatch(/^same@example\.com · [0-9a-f]{8}$/); + expect(labels[second.id]).toMatch(/^same@example\.com · [0-9a-f]{8}$/); + expect(labels[first.id]).not.toContain("workspace-alpha"); + expect(labels[second.id]).not.toContain("workspace-beta"); + expect(reversed[first.id]).toBe(labels[first.id]); + expect(reversed[second.id]).toBe(labels[second.id]); + }); + + it("binds canonical labels to account ids without exposing fallback fields", () => { + const first = account( + "33333333-3333-3333-3333-333333333333", + "internal-one", + ); + const second = account( + "44444444-4444-4444-4444-444444444444", + "internal-two", + ); + const canonical = { + [first.id]: "same@example.com · 1111aaaa", + [second.id]: "same@example.com · 2222bbbb", + }; + + const labels = buildCodexAccountDisplayNames( + [second, first], + canonical, + ); + + expect(labels[first.id]).toBe(canonical[first.id]); + expect(labels[second.id]).toBe(canonical[second.id]); + expect(Object.values(labels).join(" ")).not.toContain("internal-"); + }); + + it("uses a generic privacy-safe fallback when identity fields are absent", () => { + const missingIdentity = { + ...account("55555555-5555-5555-5555-555555555555", "secret-workspace"), + nickname: null, + emailHint: null, + authSubject: "auth0|secret-subject", + codexHomePath: "C:/private/secret-workspace", + }; + + const [label] = Object.values( + buildCodexAccountDisplayNames([missingIdentity]), + ); + + expect(label).toBe("Workspace"); + expect(label).not.toContain("secret-workspace"); + expect(label).not.toContain("secret-subject"); + expect(label).not.toContain("C:/private"); + }); +}); diff --git a/apps/desktop-tauri/src/components/codexAccountDisplay.ts b/apps/desktop-tauri/src/components/codexAccountDisplay.ts new file mode 100644 index 0000000000..6e2179c295 --- /dev/null +++ b/apps/desktop-tauri/src/components/codexAccountDisplay.ts @@ -0,0 +1,56 @@ +import type { CodexAccount } from "../types/bridge"; + +/** + * Return the same account labels in Settings and the tray menu. The backend + * normally supplies the canonical SHA-256-derived labels; the deterministic + * opaque-id fallback keeps older/test bridges collision-safe too. + */ +export function buildCodexAccountDisplayNames( + accounts: readonly CodexAccount[], + canonical: Readonly> = {}, +): Record { + const groups = new Map(); + for (const account of accounts) { + const base = codexAccountBaseName(account); + const key = base.trim().toLowerCase(); + const group = groups.get(key) ?? []; + group.push(account); + groups.set(key, group); + } + + const result: Record = {}; + for (const account of accounts) { + const supplied = canonical[account.id]?.trim(); + if (supplied) { + result[account.id] = supplied; + continue; + } + + const base = codexAccountBaseName(account); + const group = groups.get(base.trim().toLowerCase()) ?? []; + result[account.id] = + group.length > 1 + ? `${base} · ${opaqueAccountSuffix(account.id)}` + : base; + } + return result; +} + +export function codexAccountBaseName(account: CodexAccount): string { + const nickname = account.nickname?.trim(); + const email = account.emailHint?.trim().toLowerCase(); + if (email && nickname) return `${email} — ${nickname}`; + return nickname || email || "Workspace"; +} + +function opaqueAccountSuffix(id: string): string { + // Production account ids are app-owned UUIDs. FNV-1a is only a fallback + // when the Rust bridge did not provide its SHA-256 label; no provider id, + // email, or filesystem path is exposed by this path. + let hash = 0x811c9dc5; + for (let index = 0; index < id.length; index += 1) { + hash ^= id.charCodeAt(index); + hash = Math.imul(hash, 0x01000193); + } + return (hash >>> 0).toString(16).padStart(8, "0"); +} diff --git a/apps/desktop-tauri/src/hooks/useFormattedResetTime.test.tsx b/apps/desktop-tauri/src/hooks/useFormattedResetTime.test.tsx index cc1b715844..52d80db16e 100644 --- a/apps/desktop-tauri/src/hooks/useFormattedResetTime.test.tsx +++ b/apps/desktop-tauri/src/hooks/useFormattedResetTime.test.tsx @@ -3,6 +3,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { LocaleProvider } from "../i18n/LocaleProvider"; import { buildBundle } from "../test/localeHarness"; import { + normalizeResetDescription, useFormattedResetTime, type ResetTimeFormatMode, } from "./useFormattedResetTime"; @@ -61,6 +62,22 @@ describe("useFormattedResetTime", () => { vi.useRealTimers(); }); + it.each([ + ["Reset", "Resets"], + ["Resets", "Resets"], + ["Reset Jul 10 at 2:59am (Europe/Prague)", "Resets Jul 10 at 2:59am (Europe/Prague)"], + ["Reset in 11m", "Resets in 11m"], + ["Resets in 11m", "Resets in 11m"], + ["Reset at 23:30 (UTC)", "Resets at 23:30 (UTC)"], + [" rEsEt In 2h 5m \n", "Resets in 2h 5m"], + ["Reset demain à 23:30", "Resets demain à 23:30"], + ["at 23:30 (UTC)", "Resets at 23:30 (UTC)"], + ["Resetting soon", "Resets Resetting soon"], + [" \n\t", null], + ] as const)("normalizes reset description %j", (description, expected) => { + expect(normalizeResetDescription(description)).toBe(expected); + }); + it("returns a complete localized countdown in relative mode", async () => { const target = new Date("2024-06-01T03:42:00Z").toISOString(); await mountWithLocale( @@ -77,11 +94,19 @@ describe("useFormattedResetTime", () => { expect(screen.getByTestId("reset")).toHaveTextContent("Resets in 40m"); }); - it("leaves fallback text unlabelled in relative mode", async () => { + it("normalizes a fallback reset description in relative mode", async () => { + await mountWithLocale( + , + ); + expect(screen.getByTestId("reset")).toHaveTextContent("Resets in 3h"); + }); + + it("gives a parsed reset timestamp precedence over fallback wording", async () => { + const target = new Date("2024-06-01T03:42:00Z").toISOString(); await mountWithLocale( - , + , ); - expect(screen.getByTestId("reset")).toHaveTextContent("3h"); + expect(screen.getByTestId("reset")).toHaveTextContent("Resets in 3h 42m"); }); it("returns an absolute local time without the reset label", async () => { diff --git a/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts b/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts index 45f1fd8a91..4ddb5befe5 100644 --- a/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts +++ b/apps/desktop-tauri/src/hooks/useFormattedResetTime.ts @@ -3,6 +3,28 @@ import { useLocale } from "./useLocale"; export type ResetTimeFormatMode = "reset" | "expires"; +/** Normalize backend reset descriptions without changing their suffix. */ +export function normalizeResetDescription(description: string | null): string | null { + const trimmed = description?.trim() ?? ""; + if (!trimmed) return null; + + const lowercased = trimmed.toLowerCase(); + if (lowercased === "reset" || lowercased === "resets") { + return "Resets"; + } + for (const prefix of ["resets in ", "reset in "]) { + if (lowercased.startsWith(prefix)) { + return `Resets in ${trimmed.slice(prefix.length)}`; + } + } + for (const prefix of ["resets ", "reset "]) { + if (lowercased.startsWith(prefix)) { + return `Resets ${trimmed.slice(prefix.length)}`; + } + } + return `Resets ${trimmed}`; +} + const absoluteResetFormatter = new Intl.DateTimeFormat(undefined, { month: "short", day: "numeric", @@ -37,12 +59,15 @@ export function useFormattedResetTime( return () => window.clearInterval(id); }, [resetsAt, relative]); + const normalizedFallback = + mode === "reset" ? normalizeResetDescription(fallback) : fallback?.trim() || null; + if (!resetsAt) { - return fallback; + return normalizedFallback; } const target = Date.parse(resetsAt); if (Number.isNaN(target)) { - return fallback; + return normalizedFallback; } if (relative) { @@ -72,6 +97,6 @@ export function useFormattedResetTime( try { return absoluteResetFormatter.format(new Date(target)); } catch { - return fallback; + return normalizedFallback; } } diff --git a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx index a8fe0ea816..02a55b21ed 100644 --- a/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx +++ b/apps/desktop-tauri/src/hooks/useTrayPanelLayout.sizing.test.tsx @@ -173,15 +173,21 @@ describe("useTrayPanelLayout sizing", () => { }); await waitFor(() => expect(feedbackObserverCallbacks).toBeGreaterThan(0)); + // `layoutReady` can precede one already-queued reveal under a loaded CI runner. + // Give that initial work time to drain, but fail if it turns into repeated + // feedback-driven passes. Once drained, the count must remain stable. + const revealsBeforeSettle = + tauriMocks.revealTrayPanelWindow.mock.calls.length; await act(async () => { - await new Promise((resolve) => window.setTimeout(resolve, 300)); + await new Promise((resolve) => window.setTimeout(resolve, 1_000)); }); const settledRevealCount = tauriMocks.revealTrayPanelWindow.mock.calls.length; + expect(settledRevealCount - revealsBeforeSettle).toBeLessThanOrEqual(1); + await act(async () => { await new Promise((resolve) => window.setTimeout(resolve, 500)); }); - expect(tauriMocks.revealTrayPanelWindow.mock.calls.length).toBe( settledRevealCount, ); diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 2eb2140036..a7824ae2c0 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -250,6 +250,8 @@ describe("TrayPanel provider grid", () => { MenuAbout: "About CodexBar", MenuQuit: "Quit", MenuSettings: "Settings...", + ActionUsageDashboard: "Usage Dashboard", + ActionStatusPage: "Status Page", PanelAllProviders: "All providers", PanelAllProvidersShort: "All", PanelLeftSuffix: "left", @@ -337,6 +339,29 @@ describe("TrayPanel provider grid", () => { }); }); + it("scopes the status-page action to the selected provider", async () => { + const { container } = renderTrayPanel([ + provider("claude", "Claude", 35), + provider("codex", "Codex", 45), + ]); + + await waitFor(() => { + expect(container.querySelector(".tray-panel-reveal--ready")).not.toBeNull(); + }); + + fireEvent.click(screen.getByRole("button", { name: /^Claude$/ })); + fireEvent.click(await screen.findByRole("button", { name: /^Usage Dashboard$/ })); + expect(tauriMocks.openProviderDashboard).toHaveBeenLastCalledWith("claude"); + fireEvent.click(await screen.findByRole("button", { name: /^Status Page$/ })); + expect(tauriMocks.openProviderStatusPage).toHaveBeenLastCalledWith("claude"); + + fireEvent.click(screen.getByRole("button", { name: /^Codex$/ })); + fireEvent.click(await screen.findByRole("button", { name: /^Usage Dashboard$/ })); + expect(tauriMocks.openProviderDashboard).toHaveBeenLastCalledWith("codex"); + fireEvent.click(await screen.findByRole("button", { name: /^Status Page$/ })); + expect(tauriMocks.openProviderStatusPage).toHaveBeenLastCalledWith("codex"); + }); + it("localizes static tray panel labels in Japanese", async () => { tauriMocks.getLocaleStrings.mockResolvedValue( buildBundle( diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/CodexAccountsSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/CodexAccountsSection.tsx index b59332530f..7f24a8b54d 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/CodexAccountsSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/credentials/CodexAccountsSection.tsx @@ -15,6 +15,7 @@ import { codexAccountSwitch, getCodexAccountsState, } from "../../../../../lib/tauri"; +import { buildCodexAccountDisplayNames } from "../../../../../components/codexAccountDisplay"; interface Props { t: (key: LocaleKey) => string; @@ -36,6 +37,7 @@ export function CodexAccountsSection({ t }: Props) { const [snapshots, setSnapshots] = useState< Record >({}); + const [displayNames, setDisplayNames] = useState>({}); const [loaded, setLoaded] = useState(false); const [busy, setBusy] = useState(false); const [error, setError] = useState(null); @@ -49,6 +51,7 @@ export function CodexAccountsSection({ t }: Props) { try { const next: CodexAccountsStateBridge = await getCodexAccountsState(); setAccounts(next.accounts); + setDisplayNames(next.displayNames ?? {}); setSnapshots(next.snapshots); setLoaded(true); } catch (err: unknown) { @@ -152,6 +155,11 @@ export function CodexAccountsSection({ t }: Props) { return null; } + const accountDisplayNames = buildCodexAccountDisplayNames( + accounts, + displayNames, + ); + return (
@@ -210,10 +218,7 @@ export function CodexAccountsSection({ t }: Props) {
- {account.nickname ?? - account.emailHint ?? - account.authSubject ?? - shrink(account.id)} + {accountDisplayNames[account.id]} @@ -277,10 +282,6 @@ export function CodexAccountsSection({ t }: Props) { ); } -function shrink(id: string): string { - return id.length <= 12 ? id : `${id.slice(0, 8)}…`; -} - function CodexUsagePill({ snapshot, t, @@ -300,4 +301,4 @@ function CodexUsagePill({ {label || t("CodexAccountsUsageUnavailable")} ); -} \ No newline at end of file +} diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index a668392a23..ef80370a89 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -954,5 +954,7 @@ export interface CodexSwitchResult { export interface CodexAccountsStateBridge { accounts: CodexAccount[]; + /** Canonical privacy-safe account labels, keyed by stable account id. */ + displayNames?: Record; snapshots: Record; } diff --git a/rust/src/codex_accounts/mod.rs b/rust/src/codex_accounts/mod.rs index 494e6e75b3..ef84bb97b6 100644 --- a/rust/src/codex_accounts/mod.rs +++ b/rust/src/codex_accounts/mod.rs @@ -28,6 +28,6 @@ pub use codex_desktop::{ pub use login_runner::{CodexLoginOutcome, CodexLoginResult, ManagedLoginProcess}; pub use models::{ AccountUsageSnapshot, CodexAccount, CodexAccountSource, CreditsBalanceSnapshot, - RemovedAccountIdentity, UsageWindowSnapshot, utc_now, + RemovedAccountIdentity, UsageWindowSnapshot, display_names_by_id, utc_now, }; pub use stores::{AccountStore, SnapshotStore}; diff --git a/rust/src/codex_accounts/models.rs b/rust/src/codex_accounts/models.rs index 11021adbaa..f69029965e 100644 --- a/rust/src/codex_accounts/models.rs +++ b/rust/src/codex_accounts/models.rs @@ -3,6 +3,7 @@ //! Field names intentionally mirror CodexControl's `windows/.../models.py` (MIT) //! so stored data interops with that project. +use std::collections::HashMap; use std::path::{Path, PathBuf}; use chrono::{DateTime, Utc}; @@ -139,21 +140,30 @@ impl CodexAccount { } pub fn display_name(&self) -> String { - if let Some(nickname) = self + self.display_label_base() + } + + /// Return the user-facing account label without falling back to + /// credentials, provider identifiers, or filesystem paths. + fn display_label_base(&self) -> String { + let nickname = self .nickname .as_deref() .map(str::trim) + .filter(|s| !s.is_empty()); + let email = self + .email_hint + .as_deref() + .map(str::trim) .filter(|s| !s.is_empty()) - { - return nickname.to_string(); - } - if let Some(email) = self.email_hint.as_deref().filter(|s| !s.is_empty()) { - return email.to_string(); + .map(str::to_lowercase); + + match (email, nickname) { + (Some(email), Some(nickname)) => format!("{email} — {nickname}"), + (Some(email), None) => email, + (None, Some(nickname)) => nickname.to_string(), + (None, None) => "Workspace".to_string(), } - self.codex_home_path - .file_name() - .map(|name| name.to_string_lossy().to_string()) - .unwrap_or_else(|| self.codex_home_path.display().to_string()) } pub fn normalized_email_hint(&self) -> Option { @@ -175,6 +185,11 @@ impl CodexAccount { .to_lowercase() } + fn display_identity(&self) -> String { + self.normalized_provider_account_id() + .unwrap_or_else(|| self.id.to_string().to_lowercase()) + } + fn source_priority(&self) -> u8 { if self.source.owns_files() { 2 } else { 1 } } @@ -257,6 +272,56 @@ impl CodexAccount { } } +/// Build the stable labels used by account-facing surfaces. +/// +/// A provider account id is a workspace identity, but it must never be shown +/// directly. Only accounts whose privacy-safe display labels collide receive +/// an opaque suffix. The stored account UUID is folded into the hashed +/// identity when two entries claim the same provider identity, keeping +/// separate local profiles distinguishable without exposing a home path or +/// provider id. +pub fn display_names_by_id(accounts: &[CodexAccount]) -> HashMap { + let mut groups: HashMap> = HashMap::new(); + for (index, account) in accounts.iter().enumerate() { + groups + .entry(account.display_label_base().to_lowercase()) + .or_default() + .push(index); + } + + let mut labels = HashMap::with_capacity(accounts.len()); + for indexes in groups.into_values() { + if indexes.len() == 1 { + let index = indexes[0]; + labels.insert(accounts[index].id, accounts[index].display_label_base()); + continue; + } + + let mut identity_counts: HashMap = HashMap::new(); + for &index in &indexes { + *identity_counts + .entry(accounts[index].display_identity()) + .or_default() += 1; + } + + for &index in &indexes { + let account = &accounts[index]; + let identity = account.display_identity(); + let identity = if identity_counts.get(&identity) == Some(&1) { + identity + } else { + format!("{identity}\0{}", account.id) + }; + let suffix = crate::core::sha256_hex(identity.as_bytes()); + labels.insert( + account.id, + format!("{} · {}", account.display_label_base(), &suffix[..8]), + ); + } + } + labels +} + /// Identity of a previously-removed account, kept to avoid re-adding it. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -478,6 +543,21 @@ fn _path_is_trailing(path: &Path) -> bool { mod tests { use super::*; + fn display_account(id: &str, provider_account_id: &str) -> CodexAccount { + CodexAccount::new( + Uuid::parse_str(id).unwrap(), + None, + Some("user@example.com".to_string()), + None, + Some(provider_account_id.to_string()), + PathBuf::from(format!("C:/private/{provider_account_id}")), + CodexAccountSource::ManagedByApp, + utc_now(), + utc_now(), + None, + ) + } + fn account( id: &str, home: &str, @@ -557,6 +637,64 @@ mod tests { assert!(!CodexAccountSource::Ambient.owns_files()); } + #[test] + fn display_names_disambiguate_same_email_without_exposing_workspace_identity() { + let first = display_account("11111111-1111-1111-1111-111111111111", "workspace-alpha"); + let second = display_account("22222222-2222-2222-2222-222222222222", "workspace-beta"); + + let labels = display_names_by_id(&[first.clone(), second.clone()]); + let first_label = labels.get(&first.id).unwrap(); + let second_label = labels.get(&second.id).unwrap(); + assert_ne!(first_label, second_label); + for label in [first_label, second_label] { + assert!(label.starts_with("user@example.com · ")); + assert_eq!(label.rsplit_once(' ').unwrap().1.len(), 8); + assert!(!label.contains("workspace-")); + assert!(!label.contains("C:/private")); + assert!(!label.contains("auth0|")); + } + + let reordered = display_names_by_id(&[second, first.clone()]); + assert_eq!(reordered.get(&first.id), Some(first_label)); + + let mut relaunched_first = first.clone(); + relaunched_first.codex_home_path = PathBuf::from("C:/different-managed-home"); + let mut relaunched_second = first.clone(); + relaunched_second.id = Uuid::parse_str("66666666-6666-6666-6666-666666666666").unwrap(); + relaunched_second.provider_account_id = Some("workspace-beta".to_string()); + let relaunched = display_names_by_id(&[relaunched_first, relaunched_second]); + assert_eq!(relaunched.get(&first.id), Some(first_label)); + } + + #[test] + fn display_names_keep_duplicate_workspace_profiles_distinct() { + let first = display_account("33333333-3333-3333-3333-333333333333", "shared-workspace"); + let second = display_account("44444444-4444-4444-4444-444444444444", "shared-workspace"); + + let labels = display_names_by_id(&[first.clone(), second.clone()]); + assert_ne!(labels.get(&first.id), labels.get(&second.id)); + assert!(labels[&first.id].starts_with("user@example.com · ")); + assert!(labels[&second.id].starts_with("user@example.com · ")); + let reordered = display_names_by_id(&[second, first.clone()]); + assert_eq!(reordered.get(&first.id), labels.get(&first.id)); + } + + #[test] + fn display_names_use_generic_base_when_identity_fields_are_missing() { + let mut account = + display_account("55555555-5555-5555-5555-555555555555", "secret-workspace"); + account.email_hint = None; + account.auth_subject = Some("auth0|secret-subject".to_string()); + account.nickname = None; + + let labels = display_names_by_id(&[account]); + let label = labels.values().next().unwrap(); + assert!(label.starts_with("Workspace")); + assert!(!label.contains("secret-workspace")); + assert!(!label.contains("secret-subject")); + assert!(!label.contains("C:/private")); + } + #[test] fn window_role_classification() { assert_eq!( @@ -663,17 +801,14 @@ mod tests { } #[test] - fn display_name_falls_back_to_home() { + fn display_name_does_not_fall_back_to_home_path() { let acct = account( "11111111-1111-1111-1111-111111111111", "/x/my-home-dir", CodexAccountSource::ManagedByApp, None, ); - assert!( - acct.display_name().ends_with("my-home-dir") - || acct.display_name().contains("my-home-dir") - ); + assert_eq!(acct.display_name(), "Workspace"); let _ = _path_is_trailing(std::path::Path::new("/x/")); } } diff --git a/rust/src/core/claude_routed_pricing.rs b/rust/src/core/claude_routed_pricing.rs index 8d5a35347a..9ea35f9878 100644 --- a/rust/src/core/claude_routed_pricing.rs +++ b/rust/src/core/claude_routed_pricing.rs @@ -71,6 +71,20 @@ fn models_dev_targets(model: &str, normalized: String) -> Vec<(&'static str, Str targets } +/// Resolve a routed Claude model against one invocation-owned models.dev snapshot. +/// +/// Keeping this separate from the arithmetic lets a local scan memoize both positive and +/// negative model resolution without changing the provider-routing rules. +pub fn resolve_with_snapshot( + model: &str, + normalized: &str, + snapshot: &models_dev_pricing::ModelsDevPricingSnapshot, +) -> Option { + models_dev_targets(model, normalized.to_string()) + .into_iter() + .find_map(|(provider, lookup_model)| snapshot.lookup(provider, &lookup_model)) +} + pub fn cost_usd( model: &str, normalized: String, @@ -82,6 +96,23 @@ pub fn cost_usd( let pricing = models_dev_targets(model, normalized) .into_iter() .find_map(|(provider, lookup_model)| models_dev_pricing::lookup(provider, &lookup_model))?; + Some(cost_usd_from_pricing( + pricing, + input, + cache_read, + cache_write, + output, + )) +} + +/// Calculate routed cost after the models.dev resolution has already been memoized. +pub fn cost_usd_from_pricing( + pricing: models_dev_pricing::DynamicModelPricing, + input: i32, + cache_read: i32, + cache_write: i32, + output: i32, +) -> f64 { let input = input.max(0); let cache_read = cache_read.max(0); let cache_write = cache_write.max(0); @@ -125,12 +156,10 @@ pub fn cost_usd( pricing.output_cost_per_token_above_threshold, ); - Some( - (input as f64) * input_rate - + (cache_read as f64) * cache_read_rate - + (cache_write as f64) * cache_write_rate - + (output as f64) * output_rate, - ) + (input as f64) * input_rate + + (cache_read as f64) * cache_read_rate + + (cache_write as f64) * cache_write_rate + + (output as f64) * output_rate } pub fn input_cost_per_token(model: &str, normalized: String) -> Option { diff --git a/rust/src/core/cost_pricing.rs b/rust/src/core/cost_pricing.rs index 2cf6be0f0e..eb58b11063 100755 --- a/rust/src/core/cost_pricing.rs +++ b/rust/src/core/cost_pricing.rs @@ -1,10 +1,13 @@ //! Cost usage pricing — model-specific token pricing for Codex (OpenAI) and Claude (Anthropic). +use super::codex_routed_pricing; use super::models_dev_pricing; -use super::{claude_routed_pricing, codex_routed_pricing}; use chrono::NaiveDate; use std::collections::HashMap; use std::sync::LazyLock; +#[path = "cost_pricing/claude.rs"] +mod claude_pricing; +pub(crate) use claude_pricing::ClaudePricingResolution; /// Whole-request Codex rates for input above the model context threshold. #[derive(Debug, Clone, Copy)] pub struct CodexLongContextRates { @@ -48,6 +51,7 @@ pub struct ClaudePricing { /// Cost per cache read input token above threshold pub cache_read_input_cost_per_token_above_threshold: Option, } + /// Codex model pricing table static CODEX_PRICING: LazyLock> = LazyLock::new(|| { let mut m = HashMap::new(); @@ -661,41 +665,6 @@ impl CostUsagePricing { .and_then(|p| p.display_label) } - /// Normalize a Claude model name for pricing lookup - pub fn normalize_claude_model(raw: &str) -> String { - let mut trimmed = raw.trim().to_string(); - - // Remove "anthropic." prefix - if let Some(rest) = trimmed.strip_prefix("anthropic.") { - trimmed = rest.to_string(); - } - - // Handle nested model names like "anthropic.claude-sonnet-4.claude-sonnet-4-20250514" - if trimmed.contains("claude-") - && let Some(last_dot) = trimmed.rfind('.') - { - let tail = &trimmed[last_dot + 1..]; - if tail.starts_with("claude-") { - trimmed = tail.to_string(); - } - } - - // Remove version suffix like "-v1:0" - let version_pattern = regex_lite::Regex::new(r"-v\d+:\d+$").unwrap(); - trimmed = version_pattern.replace(&trimmed, "").to_string(); - - // Try without date suffix if base exists in pricing - let date_pattern = regex_lite::Regex::new(r"-\d{8}$").unwrap(); - if let Some(mat) = date_pattern.find(&trimmed) { - let base = &trimmed[..mat.start()]; - if CLAUDE_PRICING.contains_key(base) { - return base.to_string(); - } - } - - trimmed - } - /// Strip Fast/priority suffix to find the base model for pricing lookup. /// /// Fast-tier models ("gpt-5.5-fast", "gpt-5.6-sol-priority") price as the @@ -939,80 +908,6 @@ impl CostUsagePricing { )) } - #[cfg(test)] - pub(crate) fn claude_models_dev_target(model: &str) -> Option<(&'static str, String)> { - claude_routed_pricing::models_dev_target(model, Self::normalize_claude_model(model)) - } - - /// Calculate cost for Claude usage in USD - pub fn claude_cost_usd( - model: &str, - input_tokens: i32, - cache_read_input_tokens: i32, - cache_creation_input_tokens: i32, - output_tokens: i32, - ) -> Option { - let key = Self::normalize_claude_model(model); - if let Some(pricing) = CLAUDE_PRICING.get(key.as_str()) { - /// Calculate tiered cost - fn tiered(tokens: i32, base: f64, above: Option, threshold: Option) -> f64 { - let tokens = tokens.max(0); - match (threshold, above) { - (Some(thresh), Some(above_rate)) => { - let below = tokens.min(thresh); - let over = (tokens - thresh).max(0); - (below as f64) * base + (over as f64) * above_rate - } - _ => (tokens as f64) * base, - } - } - - let cost = tiered( - input_tokens, - pricing.input_cost_per_token, - pricing.input_cost_per_token_above_threshold, - pricing.threshold_tokens, - ) + tiered( - cache_read_input_tokens, - pricing.cache_read_input_cost_per_token, - pricing.cache_read_input_cost_per_token_above_threshold, - pricing.threshold_tokens, - ) + tiered( - cache_creation_input_tokens, - pricing.cache_creation_input_cost_per_token, - pricing.cache_creation_input_cost_per_token_above_threshold, - pricing.threshold_tokens, - ) + tiered( - output_tokens, - pricing.output_cost_per_token, - pricing.output_cost_per_token_above_threshold, - pricing.threshold_tokens, - ); - - return Some(cost); - } - - claude_routed_pricing::cost_usd( - model, - Self::normalize_claude_model(model), - input_tokens, - cache_read_input_tokens, - cache_creation_input_tokens, - output_tokens, - ) - } - - /// Base per-token input rate for a Claude model. Exposed for callers that - /// need a rate the standard cost function doesn't model — e.g. the usage - /// scanner's one-hour cache-write premium, billed at 2x the input rate. - pub fn claude_input_cost_per_token(model: &str) -> Option { - let key = Self::normalize_claude_model(model); - if let Some(pricing) = CLAUDE_PRICING.get(key.as_str()) { - return Some(pricing.input_cost_per_token); - } - claude_routed_pricing::input_cost_per_token(model, Self::normalize_claude_model(model)) - } - /// 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/claude.rs b/rust/src/core/cost_pricing/claude.rs new file mode 100644 index 0000000000..46fa56738a --- /dev/null +++ b/rust/src/core/cost_pricing/claude.rs @@ -0,0 +1,177 @@ +use super::super::{claude_routed_pricing, models_dev_pricing}; +use super::{CLAUDE_PRICING, ClaudePricing, CostUsagePricing}; + +/// Resolved Claude pricing source used by the local scanner's invocation memo. +#[derive(Debug, Clone, Copy)] +pub(crate) enum ClaudePricingResolution { + BuiltIn(ClaudePricing), + ModelsDev(models_dev_pricing::DynamicModelPricing), +} + +impl CostUsagePricing { + /// Normalize a Claude model name for pricing lookup + pub fn normalize_claude_model(raw: &str) -> String { + let mut trimmed = raw.trim().to_string(); + + // Remove "anthropic." prefix + if let Some(rest) = trimmed.strip_prefix("anthropic.") { + trimmed = rest.to_string(); + } + + // Handle nested model names like "anthropic.claude-sonnet-4.claude-sonnet-4-20250514" + if trimmed.contains("claude-") + && let Some(last_dot) = trimmed.rfind('.') + { + let tail = &trimmed[last_dot + 1..]; + if tail.starts_with("claude-") { + trimmed = tail.to_string(); + } + } + + // Remove version suffix like "-v1:0" + let version_pattern = regex_lite::Regex::new(r"-v\d+:\d+$").unwrap(); + trimmed = version_pattern.replace(&trimmed, "").to_string(); + + // Try without date suffix if base exists in pricing + let date_pattern = regex_lite::Regex::new(r"-\d{8}$").unwrap(); + if let Some(mat) = date_pattern.find(&trimmed) { + let base = &trimmed[..mat.start()]; + if CLAUDE_PRICING.contains_key(base) { + return base.to_string(); + } + } + + trimmed + } + + /// Resolve one Claude model without rereading the models.dev artifact. + pub(crate) fn resolve_claude_pricing( + model: &str, + normalized: &str, + pricing_snapshot: Option<&models_dev_pricing::ModelsDevPricingSnapshot>, + ) -> Option { + if let Some(pricing) = CLAUDE_PRICING.get(normalized) { + return Some(ClaudePricingResolution::BuiltIn(*pricing)); + } + pricing_snapshot + .and_then(|snapshot| { + claude_routed_pricing::resolve_with_snapshot(model, normalized, snapshot) + }) + .map(ClaudePricingResolution::ModelsDev) + } + + /// Calculate cost from a previously resolved Claude pricing source. + pub(crate) fn claude_cost_usd_from_resolution( + resolution: ClaudePricingResolution, + input_tokens: i32, + cache_read_input_tokens: i32, + cache_creation_input_tokens: i32, + output_tokens: i32, + ) -> f64 { + match resolution { + ClaudePricingResolution::BuiltIn(pricing) => { + fn tiered( + tokens: i32, + base: f64, + above: Option, + threshold: Option, + ) -> f64 { + let tokens = tokens.max(0); + match (threshold, above) { + (Some(thresh), Some(above_rate)) => { + let below = tokens.min(thresh); + let over = (tokens - thresh).max(0); + (below as f64) * base + (over as f64) * above_rate + } + _ => (tokens as f64) * base, + } + } + + tiered( + input_tokens, + pricing.input_cost_per_token, + pricing.input_cost_per_token_above_threshold, + pricing.threshold_tokens, + ) + tiered( + cache_read_input_tokens, + pricing.cache_read_input_cost_per_token, + pricing.cache_read_input_cost_per_token_above_threshold, + pricing.threshold_tokens, + ) + tiered( + cache_creation_input_tokens, + pricing.cache_creation_input_cost_per_token, + pricing.cache_creation_input_cost_per_token_above_threshold, + pricing.threshold_tokens, + ) + tiered( + output_tokens, + pricing.output_cost_per_token, + pricing.output_cost_per_token_above_threshold, + pricing.threshold_tokens, + ) + } + ClaudePricingResolution::ModelsDev(pricing) => { + claude_routed_pricing::cost_usd_from_pricing( + pricing, + input_tokens, + cache_read_input_tokens, + cache_creation_input_tokens, + output_tokens, + ) + } + } + } + + pub(crate) fn claude_input_cost_per_token_from_resolution( + resolution: ClaudePricingResolution, + ) -> f64 { + match resolution { + ClaudePricingResolution::BuiltIn(pricing) => pricing.input_cost_per_token, + ClaudePricingResolution::ModelsDev(pricing) => pricing.input_cost_per_token, + } + } + + #[cfg(test)] + pub(crate) fn claude_models_dev_target(model: &str) -> Option<(&'static str, String)> { + claude_routed_pricing::models_dev_target(model, Self::normalize_claude_model(model)) + } + + /// Calculate cost for Claude usage in USD + pub fn claude_cost_usd( + model: &str, + input_tokens: i32, + cache_read_input_tokens: i32, + cache_creation_input_tokens: i32, + output_tokens: i32, + ) -> Option { + let key = Self::normalize_claude_model(model); + if let Some(pricing) = CLAUDE_PRICING.get(key.as_str()) { + return Some(Self::claude_cost_usd_from_resolution( + ClaudePricingResolution::BuiltIn(*pricing), + input_tokens, + cache_read_input_tokens, + cache_creation_input_tokens, + output_tokens, + )); + } + + claude_routed_pricing::cost_usd( + model, + Self::normalize_claude_model(model), + input_tokens, + cache_read_input_tokens, + cache_creation_input_tokens, + output_tokens, + ) + } + + /// Base per-token input rate for a Claude model. Exposed for callers that + /// need a rate the standard cost function doesn't model — e.g. the usage + /// scanner's one-hour cache-write premium, billed at 2x the input rate. + pub fn claude_input_cost_per_token(model: &str) -> Option { + let key = Self::normalize_claude_model(model); + if let Some(pricing) = CLAUDE_PRICING.get(key.as_str()) { + return Some(pricing.input_cost_per_token); + } + claude_routed_pricing::input_cost_per_token(model, Self::normalize_claude_model(model)) + } +} diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index b6328e29b4..ef0dd20e47 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -128,6 +128,13 @@ impl CostScanOptions { } } + /// Whether this pass was requested by an explicit/app-driven refresh. + /// The zero debounce used by app-driven scans is the existing refresh + /// state, so no second force/resume flag is needed. + pub fn is_app_driven(&self) -> bool { + self.refresh_min_interval_secs == 0 + } + /// Whether a prior scan at `last_scan_unix_ms` is still within the debounce window. pub fn should_skip_scan(&self, last_scan_unix_ms: i64, now_unix_ms: i64) -> bool { // Debounce intervals are seconds-scale config values, far below i64::MAX. @@ -159,6 +166,16 @@ impl CacheStamp { } } +/// Terminal reason for a bounded Codex catch-up pause. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CodexScanPauseReason { + /// A bounded pass left work queued without consuming any new source data. + NoProgress, + /// The source could not be inspected reliably; keep the validated report until retry. + Error(String), +} + /// Cache for scanned file data #[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct CostUsageCache { @@ -185,6 +202,11 @@ pub struct CostUsageCache { /// True while bounded Codex catch-up has not completed for this window. #[serde(default, skip_serializing_if = "std::ops::Not::not")] pub codex_scan_incomplete: bool, + /// Terminal catch-up pause attached to the existing incomplete state. A + /// background scan must not clear or retry this state; an app-driven + /// refresh clears it before starting the next pass. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub codex_scan_pause_reason: Option, /// Content stamp of the decoded on-disk baseline. This is process-local /// and omitted from JSON so a stale reader cannot replace a newer cache. #[serde(skip)] diff --git a/rust/src/core/jsonl_scanner/codex.rs b/rust/src/core/jsonl_scanner/codex.rs index 076272c550..349bc5bb7c 100644 --- a/rust/src/core/jsonl_scanner/codex.rs +++ b/rust/src/core/jsonl_scanner/codex.rs @@ -4,7 +4,7 @@ mod helpers; mod parser; use helpers::{ - CODEX_JSONL_MAX_LINE_BYTES, nonempty_json_string, parse_rfc3339_timestamp, + BoundedJsonlLine, CODEX_JSONL_MAX_LINE_BYTES, nonempty_json_string, parse_rfc3339_timestamp, read_bounded_jsonl_line, session_meta_field, }; use parser::CodexParserState; @@ -104,11 +104,17 @@ impl JsonlScanner { let mut bytes_examined = 0_usize; while bytes_examined < CODEX_JSONL_MAX_LINE_BYTES { - let Some((line_bytes, consumed)) = - read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES)? + let Some(line) = read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES)? else { break; }; + let (line_bytes, consumed) = match line { + BoundedJsonlLine::Retained { bytes, consumed } => (bytes, consumed), + BoundedJsonlLine::Discarded { consumed } => { + bytes_examined = bytes_examined.saturating_add(consumed); + continue; + } + }; bytes_examined = bytes_examined.saturating_add(consumed); if line_bytes.is_empty() { continue; @@ -320,8 +326,7 @@ impl JsonlScanner { cancelled = true; break; } - let Some((line_bytes, consumed)) = - read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES)? + let Some(line) = read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES)? else { break; }; @@ -329,13 +334,15 @@ impl JsonlScanner { cancelled = true; break; } - // Per-line byte counts are capped at 256 KiB, far inside i64::MAX. - #[allow( - clippy::cast_possible_wrap, - reason = "per-line consumed bytes are capped at CODEX_JSONL_MAX_LINE_BYTES" - )] - let consumed_i64 = consumed as i64; - parsed_bytes += consumed_i64; + let (line_bytes, consumed) = match line { + BoundedJsonlLine::Retained { bytes, consumed } => (Some(bytes), consumed), + BoundedJsonlLine::Discarded { consumed } => (None, consumed), + }; + let consumed_i64 = i64::try_from(consumed).unwrap_or(i64::MAX); + parsed_bytes = parsed_bytes.saturating_add(consumed_i64); + let Some(line_bytes) = line_bytes else { + continue; + }; if line_bytes.is_empty() { continue; } diff --git a/rust/src/core/jsonl_scanner/codex/helpers.rs b/rust/src/core/jsonl_scanner/codex/helpers.rs index f0129ff126..fd69046e61 100644 --- a/rust/src/core/jsonl_scanner/codex/helpers.rs +++ b/rust/src/core/jsonl_scanner/codex/helpers.rs @@ -144,12 +144,20 @@ pub(super) fn cumulative_reasoning_delta( ) } +/// A bounded physical JSONL line. +/// +/// Keeping the discarded case separate prevents callers from accidentally +/// treating an oversized prefix as a parseable empty line. +pub(super) enum BoundedJsonlLine { + Retained { bytes: Vec, consumed: usize }, + Discarded { consumed: usize }, +} + /// Read one JSONL line, discarding content when it exceeds `max_bytes`. -/// Returns `(line_without_newline, bytes_consumed_including_newline)`. pub(super) fn read_bounded_jsonl_line( reader: &mut R, max_bytes: usize, -) -> std::io::Result, usize)>> { +) -> std::io::Result> { let mut line = Vec::new(); let mut saw_bytes = false; let mut discarding = false; @@ -158,9 +166,16 @@ pub(super) fn read_bounded_jsonl_line( loop { let chunk = reader.fill_buf()?; if chunk.is_empty() { - return Ok( - saw_bytes.then_some((if discarding { Vec::new() } else { line }, consumed_total)) - ); + return Ok(saw_bytes.then_some(if discarding { + BoundedJsonlLine::Discarded { + consumed: consumed_total, + } + } else { + BoundedJsonlLine::Retained { + bytes: line, + consumed: consumed_total, + } + })); } let newline = chunk.iter().position(|byte| *byte == b'\n'); let segment_end = newline.unwrap_or(chunk.len()); @@ -181,10 +196,16 @@ pub(super) fn read_bounded_jsonl_line( reader.consume(consumed); consumed_total += consumed; if newline.is_some() { - return Ok(Some(( - if discarding { Vec::new() } else { line }, - consumed_total, - ))); + return Ok(Some(if discarding { + BoundedJsonlLine::Discarded { + consumed: consumed_total, + } + } else { + BoundedJsonlLine::Retained { + bytes: line, + consumed: consumed_total, + } + })); } } } diff --git a/rust/src/core/jsonl_scanner/tests.rs b/rust/src/core/jsonl_scanner/tests.rs index fb39fed1cd..cb08f37b57 100644 --- a/rust/src/core/jsonl_scanner/tests.rs +++ b/rust/src/core/jsonl_scanner/tests.rs @@ -541,19 +541,114 @@ fn codex_parser_discards_oversized_line_and_recovers_next_record() { ); } +#[test] +fn codex_parser_validates_a_record_at_the_line_limit() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + let prefix = r#"{"timestamp":"2026-05-31T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":9,"cached_input_tokens":2,"output_tokens":1}}},"padding":""}"#; + let padding_len = CODEX_JSONL_MAX_LINE_BYTES - prefix.len(); + let line = format!( + "{}{}\"}}", + &prefix[..prefix.len() - 2], + "x".repeat(padding_len) + ); + assert_eq!(line.len(), CODEX_JSONL_MAX_LINE_BYTES); + writeln!(file, "{line}").unwrap(); + + let day = NaiveDate::from_ymd_opt(2026, 5, 31).unwrap(); + let parsed = JsonlScanner::parse_codex_file( + file.path(), + &CostUsageDayRange::new(day, day), + 0, + None, + None, + ) + .expect("parse"); + + assert_eq!(parsed.records.len(), 1); + assert_eq!( + ( + parsed.records[0].input, + parsed.records[0].cached, + parsed.records[0].output + ), + (9, 2, 1) + ); +} + +#[test] +fn codex_parser_discards_a_line_at_limit_plus_one_and_keeps_following_record() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + let prefix = r#"{"padding":""}"#; + let padding_len = CODEX_JSONL_MAX_LINE_BYTES + 1 - prefix.len(); + let oversized = format!( + "{}{}\"}}", + &prefix[..prefix.len() - 2], + "x".repeat(padding_len) + ); + assert_eq!(oversized.len(), CODEX_JSONL_MAX_LINE_BYTES + 1); + writeln!(file, "{oversized}").unwrap(); + let valid = r#"{"timestamp":"2026-05-31T10:00:01Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":9,"cached_input_tokens":2,"output_tokens":1}}}}"#; + writeln!(file, "{valid}").unwrap(); + + let day = NaiveDate::from_ymd_opt(2026, 5, 31).unwrap(); + let parsed = JsonlScanner::parse_codex_file( + file.path(), + &CostUsageDayRange::new(day, day), + 0, + None, + None, + ) + .expect("parse"); + + assert_eq!(parsed.records.len(), 1); + assert_eq!(parsed.records[0].input, 9); +} + +#[test] +fn codex_parser_discards_huge_malformed_lines_before_and_after_valid_records() { + let mut file = tempfile::NamedTempFile::new().expect("temp file"); + let valid = r#"{"timestamp":"2026-05-31T10:00:00Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":9,"cached_input_tokens":2,"output_tokens":1}}}}"#; + let malformed = format!("{{{}", "x".repeat(CODEX_JSONL_MAX_LINE_BYTES * 4)); + writeln!(file, "{malformed}").unwrap(); + writeln!(file, "{valid}").unwrap(); + writeln!(file, "{malformed}").unwrap(); + + let day = NaiveDate::from_ymd_opt(2026, 5, 31).unwrap(); + let parsed = JsonlScanner::parse_codex_file( + file.path(), + &CostUsageDayRange::new(day, day), + 0, + None, + None, + ) + .expect("parse"); + + assert_eq!(parsed.records.len(), 1); + assert_eq!(parsed.records[0].input, 9); +} + #[test] fn bounded_jsonl_reader_accepts_exact_limit_without_retaining_larger_input() { let mut input = vec![b'x'; CODEX_JSONL_MAX_LINE_BYTES]; input.push(b'\n'); - input.extend_from_slice(b"{\"type\":\"event_msg\"}\n"); + input.extend_from_slice(br#"{"type":"event_msg"}"#); + input.push(b'\n'); let mut reader = BufReader::with_capacity(64 * 1024, std::io::Cursor::new(input)); - let (exact, _) = read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES) + let exact = match read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES) .expect("read") - .expect("line"); - let (later, _) = read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES) + .expect("line") + { + BoundedJsonlLine::Retained { bytes, .. } => bytes, + BoundedJsonlLine::Discarded { .. } => panic!("exact-limit line was discarded"), + }; + let later = match read_bounded_jsonl_line(&mut reader, CODEX_JSONL_MAX_LINE_BYTES) .expect("read") - .expect("line"); + .expect("line") + { + BoundedJsonlLine::Retained { bytes, .. } => bytes, + BoundedJsonlLine::Discarded { .. } => panic!("following line was discarded"), + }; assert_eq!(exact.len(), CODEX_JSONL_MAX_LINE_BYTES); assert_eq!(later, br#"{"type":"event_msg"}"#); @@ -758,6 +853,8 @@ fn interleaved_lineage_mid_range_climb_below_watermark_does_not_readd() { fn cost_scan_options_app_driven_bypasses_debounce() { let debounced = CostScanOptions::default(); let forced = CostScanOptions::app_driven(); + assert!(!debounced.is_app_driven()); + assert!(forced.is_app_driven()); let last = 1_000_000_i64; let now = last + 1_000; // 1s later, within 60s window diff --git a/rust/src/cost_scanner.rs b/rust/src/cost_scanner.rs index 09862e017b..d6854d1df9 100755 --- a/rust/src/cost_scanner.rs +++ b/rust/src/cost_scanner.rs @@ -26,12 +26,16 @@ use crate::codex_costs::{ }; use crate::codex_sessions::{codex_sessions_dir_candidates, default_wsl_roots}; use crate::core::{ - CachedCostReport, CostScanOptions, CostUsageCache, CostUsageDayRange, CostUsageFileUsage, - CostUsagePricing, JsonlScanner, ProviderId, + CachedCostReport, CodexScanPauseReason, CostScanOptions, CostUsageCache, CostUsageDayRange, + CostUsageFileUsage, JsonlScanner, ProviderId, }; use crate::providers::opencodego::local as opencodego_local; use crate::settings::Settings; +mod claude_pricing; mod codex; +use claude_pricing::ClaudeScanPricingResolver; +#[cfg(test)] +use claude_pricing::{ClaudePricing, FALLBACK_CLAUDE_MODEL}; /// Completeness of the pricing coverage in a [`CostSummary`] (upstream 0.48.0 F18). /// @@ -126,10 +130,6 @@ fn is_cancelled(cancel: Option<&AtomicBool>) -> bool { cancel.is_some_and(|flag| flag.load(Ordering::Relaxed)) } -/// Fallback Claude model used when a scanned model isn't in the canonical -/// pricing table (unknown or retired IDs). Prices as Sonnet 4.6. -const FALLBACK_CLAUDE_MODEL: &str = "claude-sonnet-4-6"; - fn unix_now_ms() -> i64 { // Duration is clamped to i64::MAX before casting, so the value fits i64. #[allow( @@ -156,56 +156,6 @@ fn system_time_to_unix_ms(modified: Option) -> i64 { millis } -/// input rate. -struct ClaudePricing; - -impl ClaudePricing { - fn cost_usd_with_cache_ttl( - model: &str, - input: u64, - cache_create: u64, - cache_create_1h: u64, - cache_read: u64, - output: u64, - ) -> f64 { - let cache_create_1h = cache_create_1h.min(cache_create); - let cache_create_5m = cache_create.saturating_sub(cache_create_1h); - - // Standard buckets (input, cache-read, 5-minute cache-write, output), - // including any long-context tiering, come from the canonical table. - // Unknown/retired models fall back to Sonnet pricing. - #[allow( - clippy::cast_possible_truncation, - reason = "clamped to i32::MAX before casting" - )] - let clamp = |v: u64| v.min(i32::MAX as u64) as i32; - let base = CostUsagePricing::claude_cost_usd( - model, - clamp(input), - clamp(cache_read), - clamp(cache_create_5m), - clamp(output), - ) - .or_else(|| { - CostUsagePricing::claude_cost_usd( - FALLBACK_CLAUDE_MODEL, - clamp(input), - clamp(cache_read), - clamp(cache_create_5m), - clamp(output), - ) - }) - .unwrap_or(0.0); - - // Scanner-specific: one-hour cache writes bill at 2x the input rate. - let input_rate = CostUsagePricing::claude_input_cost_per_token(model) - .or_else(|| CostUsagePricing::claude_input_cost_per_token(FALLBACK_CLAUDE_MODEL)) - .unwrap_or(0.0); - - base + (cache_create_1h as f64) * input_rate * 2.0 - } -} - /// JSONL event structures for Codex #[allow( dead_code, @@ -429,6 +379,7 @@ fn contains_claude_vertex_marker(value: &str, include_gcp: bool) -> bool { #[derive(Debug)] struct ClaudeUsageRecord { model: String, + pricing_known: bool, timestamp: Option>, dedup_key: Option, input: u64, @@ -512,11 +463,18 @@ impl CostScanner { // that appear across multiple files. if projects_dir.exists() { let mut seen = HashSet::new(); + let mut pricing = ClaudeScanPricingResolver::default(); let mut handle_file = |path: &Path| { - let counted = - for_each_claude_usage_record(path, &cutoff, &mut seen, cancel, |record| { + let counted = for_each_claude_usage_record_with_pricing( + path, + &cutoff, + &mut seen, + cancel, + &mut pricing, + |record| { add_claude_record_to_summary(&mut summary, record); - }); + }, + ); if counted > 0 { summary.sessions_count += 1; } @@ -624,11 +582,27 @@ impl CostScanner { /// consume this single reader, so Claude log semantics live in one place. /// Returns the number of records consumed, so callers can tell whether the /// file contributed anything. +#[cfg(test)] fn for_each_claude_usage_record( path: &Path, cutoff: &DateTime, seen: &mut HashSet, cancel: Option<&AtomicBool>, + on_record: F, +) -> usize +where + F: FnMut(&ClaudeUsageRecord), +{ + let mut pricing = ClaudeScanPricingResolver::default(); + for_each_claude_usage_record_with_pricing(path, cutoff, seen, cancel, &mut pricing, on_record) +} + +fn for_each_claude_usage_record_with_pricing( + path: &Path, + cutoff: &DateTime, + seen: &mut HashSet, + cancel: Option<&AtomicBool>, + pricing: &mut ClaudeScanPricingResolver, mut on_record: F, ) -> usize where @@ -648,7 +622,7 @@ where } if let Ok(event) = serde_json::from_str::(line) && !event.is_vertex_ai_usage_entry() - && let Some(record) = claude_usage_record_from_event(&event) + && let Some(record) = claude_usage_record_from_event_with_pricing(&event, pricing) && should_count_claude_record(&record, cutoff, seen) { counted += 1; @@ -686,7 +660,16 @@ where } } +#[cfg(test)] fn claude_usage_record_from_event(event: &ClaudeEvent) -> Option { + let mut pricing = ClaudeScanPricingResolver::default(); + claude_usage_record_from_event_with_pricing(event, &mut pricing) +} + +fn claude_usage_record_from_event_with_pricing( + event: &ClaudeEvent, + pricing: &mut ClaudeScanPricingResolver, +) -> Option { if event.event_type.as_deref() != Some("assistant") { return None; } @@ -705,7 +688,8 @@ fn claude_usage_record_from_event(event: &ClaudeEvent) -> Option Option Vec<(String, Option< if projects_dir.exists() { let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); + let mut pricing = ClaudeScanPricingResolver::default(); let mut handle_file = |path: &Path| { - for_each_claude_usage_record(path, &cutoff, &mut seen, None, |record| { - add_claude_record_to_daily_costs(&mut daily_costs, record); - }); + for_each_claude_usage_record_with_pricing( + path, + &cutoff, + &mut seen, + None, + &mut pricing, + |record| { + add_claude_record_to_daily_costs(&mut daily_costs, record); + }, + ); }; scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); } @@ -946,10 +939,18 @@ pub fn get_daily_token_history(provider: &str, days: u32) -> (Vec<(String, u64)> if projects_dir.exists() { let cutoff = Utc::now() - Duration::days(days as i64); let mut seen = HashSet::new(); + let mut pricing = ClaudeScanPricingResolver::default(); let mut handle_file = |path: &Path| { - for_each_claude_usage_record(path, &cutoff, &mut seen, None, |record| { - add_claude_record_to_daily_tokens(&mut daily_tokens, record); - }); + for_each_claude_usage_record_with_pricing( + path, + &cutoff, + &mut seen, + None, + &mut pricing, + |record| { + add_claude_record_to_daily_tokens(&mut daily_tokens, record); + }, + ); }; scanner.walk_claude_files(&projects_dir, &cutoff, None, &mut handle_file); } diff --git a/rust/src/cost_scanner/claude_pricing.rs b/rust/src/cost_scanner/claude_pricing.rs new file mode 100644 index 0000000000..0bea3778e0 --- /dev/null +++ b/rust/src/cost_scanner/claude_pricing.rs @@ -0,0 +1,159 @@ +use crate::core::{ClaudePricingResolution, CostUsagePricing, ModelsDevPricingSnapshot}; +use std::collections::HashMap; + +pub(super) const FALLBACK_CLAUDE_MODEL: &str = "claude-sonnet-4-6"; + +#[cfg(test)] +pub(super) struct ClaudePricing; + +#[cfg(test)] +impl ClaudePricing { + pub(super) fn cost_usd_with_cache_ttl( + model: &str, + input: u64, + cache_create: u64, + cache_create_1h: u64, + cache_read: u64, + output: u64, + ) -> f64 { + let cache_create_1h = cache_create_1h.min(cache_create); + let cache_create_5m = cache_create.saturating_sub(cache_create_1h); + + // Standard buckets (input, cache-read, 5-minute cache-write, output), + // including any long-context tiering, come from the canonical table. + // Unknown/retired models fall back to Sonnet pricing. + #[allow( + clippy::cast_possible_truncation, + reason = "clamped to i32::MAX before casting" + )] + let clamp = |v: u64| v.min(i32::MAX as u64) as i32; + let base = CostUsagePricing::claude_cost_usd( + model, + clamp(input), + clamp(cache_read), + clamp(cache_create_5m), + clamp(output), + ) + .or_else(|| { + CostUsagePricing::claude_cost_usd( + FALLBACK_CLAUDE_MODEL, + clamp(input), + clamp(cache_read), + clamp(cache_create_5m), + clamp(output), + ) + }) + .unwrap_or(0.0); + + // Scanner-specific: one-hour cache writes bill at 2x the input rate. + let input_rate = CostUsagePricing::claude_input_cost_per_token(model) + .or_else(|| CostUsagePricing::claude_input_cost_per_token(FALLBACK_CLAUDE_MODEL)) + .unwrap_or(0.0); + + base + (cache_create_1h as f64) * input_rate * 2.0 + } +} + +/// Per-scan Claude pricing memo. +/// +/// Claude logs commonly repeat the same model across many files and records. Keep model +/// normalization and positive/negative models.dev resolution scan-local while retaining the +/// canonical pricing arithmetic and provider-routing rules. +#[derive(Default)] +pub(super) struct ClaudeScanPricingResolver { + snapshot: Option, + pub(super) normalized_models: HashMap, + pub(super) resolutions: HashMap>, + #[cfg(test)] + pub(super) normalization_cache_misses: usize, + #[cfg(test)] + pub(super) resolution_cache_misses: usize, +} + +impl ClaudeScanPricingResolver { + pub(super) const MEMO_ENTRY_LIMIT: usize = 1024; + + pub(super) fn normalize(&mut self, model: &str) -> String { + if let Some(normalized) = self.normalized_models.get(model) { + return normalized.clone(); + } + #[cfg(test)] + { + self.normalization_cache_misses += 1; + } + let normalized = CostUsagePricing::normalize_claude_model(model); + if self.normalized_models.len() < Self::MEMO_ENTRY_LIMIT { + self.normalized_models + .insert(model.to_string(), normalized.clone()); + } + normalized + } + + fn resolve(&mut self, model: &str) -> Option { + if let Some(resolution) = self.resolutions.get(model) { + return *resolution; + } + #[cfg(test)] + { + self.resolution_cache_misses += 1; + } + + let normalized = self.normalize(model); + let needs_catalog = self.snapshot.is_none(); + let mut resolution = + CostUsagePricing::resolve_claude_pricing(model, &normalized, self.snapshot.as_ref()); + if resolution.is_none() && needs_catalog { + let snapshot = self + .snapshot + .get_or_insert_with(crate::core::pricing_snapshot); + resolution = + CostUsagePricing::resolve_claude_pricing(model, &normalized, Some(snapshot)); + } + if self.resolutions.len() < Self::MEMO_ENTRY_LIMIT { + self.resolutions.insert(model.to_string(), resolution); + } + resolution + } + + pub(super) fn is_known(&mut self, model: &str) -> bool { + self.resolve(model).is_some() + } + + pub(super) fn cost_usd_with_cache_ttl( + &mut self, + model: &str, + input: u64, + cache_create: u64, + cache_create_1h: u64, + cache_read: u64, + output: u64, + ) -> f64 { + let cache_create_1h = cache_create_1h.min(cache_create); + let cache_create_5m = cache_create.saturating_sub(cache_create_1h); + + #[allow( + clippy::cast_possible_truncation, + reason = "clamped to i32::MAX before casting" + )] + let clamp = |value: u64| value.min(i32::MAX as u64) as i32; + + let resolved = self.resolve(model); + let billable = resolved.or_else(|| self.resolve(FALLBACK_CLAUDE_MODEL)); + let base = billable + .map(|pricing| { + CostUsagePricing::claude_cost_usd_from_resolution( + pricing, + clamp(input), + clamp(cache_read), + clamp(cache_create_5m), + clamp(output), + ) + }) + .unwrap_or(0.0); + let input_rate = billable + .map(CostUsagePricing::claude_input_cost_per_token_from_resolution) + .unwrap_or(0.0); + + base + (cache_create_1h as f64) * input_rate * 2.0 + } +} diff --git a/rust/src/cost_scanner/codex.rs b/rust/src/cost_scanner/codex.rs index 311cef1138..175c6aad43 100644 --- a/rust/src/cost_scanner/codex.rs +++ b/rust/src/cost_scanner/codex.rs @@ -1,5 +1,8 @@ use super::*; +mod reconciliation; +use reconciliation::*; + fn rebuild_cache_days(cache: &mut CostUsageCache) { cache.days.clear(); for usage in cache.files.values() { @@ -86,42 +89,6 @@ fn summary_from_cached_report( } } -/// Remove cached Codex files that are provably gone from the portion of the -/// sessions tree covered by this scan. Entries outside the current roots or -/// date directories are intentionally retained for a later scan. -fn reconcile_missing_codex_cache_files( - cache: &mut CostUsageCache, - sessions_dirs: &[PathBuf], - range: &CostUsageDayRange, -) { - let scanned_date_dirs: Vec = sessions_dirs - .iter() - .flat_map(|sessions_dir| { - codex_scan_dates(range).into_iter().map(|date| { - sessions_dir - .join(date.format("%Y").to_string()) - .join(date.format("%m").to_string()) - .join(date.format("%d").to_string()) - }) - }) - .collect(); - - cache.files.retain(|path_key, _| { - let path = Path::new(path_key); - let in_scanned_root = sessions_dirs - .iter() - .any(|sessions_dir| path.starts_with(sessions_dir)); - let in_scanned_date = scanned_date_dirs - .iter() - .any(|date_dir| path.starts_with(date_dir)); - - !(in_scanned_root && in_scanned_date && !path.exists()) - }); - cache - .codex_pending_paths - .retain(|path| Path::new(path).exists()); -} - fn cached_codex_file_is_complete_for_range( cache: &CostUsageCache, path_key: &str, @@ -265,6 +232,23 @@ impl CostScanner { let cache_root = self.cache_root.as_deref(); let mut cache = JsonlScanner::load_cache(ProviderId::Codex, cache_root); + // A no-progress or source-error catch-up is terminal for background + // synchronization. Keep the resumable queue and last validated report + // intact until the user explicitly requests an app-driven refresh. + if cache.codex_scan_incomplete + && cache.codex_scan_pause_reason.is_some() + && !self.options.is_app_driven() + { + return ( + paused_codex_summary(&cache, start_date, today), + stats, + cache, + ); + } + if self.options.is_app_driven() { + cache.codex_scan_pause_reason = None; + } + // Debounce: rebuild from disk cache without re-walking session files. if JsonlScanner::should_skip_cached_scan(&cache, self.options, now_ms) && !cache.codex_scan_incomplete @@ -403,22 +387,50 @@ impl CostScanner { } } + let mut pruned_paths_pending = Vec::new(); if discovery_complete && !is_cancelled(cancel) { - reconcile_missing_codex_cache_files(&mut cache, &sessions_dirs, &range); - for path in &pending_paths_before_pass { - if !Path::new(path).exists() { - cache.files.remove(path); + pruned_paths_pending = missing_codex_cache_paths(&cache, &sessions_dirs, &range); + if self.options.is_app_driven() { + reconcile_missing_codex_cache_files(&mut cache, &sessions_dirs, &range); + for path in &pending_paths_before_pass { + if !Path::new(path).exists() { + cache.files.remove(path); + } + } + } else { + for path in &pruned_paths_pending { + if !pending_next.contains(path) { + pending_next.push(path.clone()); + } } } } pending_next.retain(|path| { + // An incomplete discovery is a source failure, not proof that a + // queued path was pruned. Preserve the priority cursor verbatim so + // the next explicit refresh can validate the source and resume it. + if !discovery_complete + || is_cancelled(cancel) + || (!self.options.is_app_driven() + && pruned_paths_pending.iter().any(|pending| pending == path)) + { + return true; + } Path::new(path).exists() - && (!discovery_complete - || is_cancelled(cancel) - || is_codex_path_in_scan_window(Path::new(path), &sessions_dirs, &range)) + && is_codex_path_in_scan_window(Path::new(path), &sessions_dirs, &range) }); - pending_next.sort(); - pending_next.dedup(); + if discovery_complete + && !is_cancelled(cancel) + && (pruned_paths_pending.is_empty() || self.options.is_app_driven()) + { + pending_next.sort(); + pending_next.dedup(); + } else { + // Preserve queue order while the source is unavailable or the + // pass is cancelled; this is the durable priority cursor. + let mut seen_pending = HashSet::new(); + pending_next.retain(|path| seen_pending.insert(path.clone())); + } cache.codex_pending_paths = pending_next; cache.codex_scan_incomplete = !discovery_complete || is_cancelled(cancel) || !cache.codex_pending_paths.is_empty(); @@ -428,10 +440,24 @@ impl CostScanner { if cache.previous_report.is_none() { cache.previous_report = established_report_before_scan; } + if !is_cancelled(cancel) { + cache.codex_scan_pause_reason = if !discovery_complete { + Some(CodexScanPauseReason::Error( + "Codex session source unavailable".to_string(), + )) + } else if !pruned_paths_pending.is_empty() + || (bytes_read_this_refresh == 0 && !cache.codex_pending_paths.is_empty()) + { + Some(CodexScanPauseReason::NoProgress) + } else { + None + }; + } } else { cache.scan_since_key = Some(range.since_key.clone()); cache.scan_until_key = Some(range.until_key.clone()); cache.previous_report = None; + cache.codex_scan_pause_reason = None; } JsonlScanner::save_cache(ProviderId::Codex, &mut cache, cache_root); @@ -538,12 +564,27 @@ impl CostScanner { ) -> (Vec, bool) { let mut candidates = Vec::new(); let mut seen = HashSet::new(); + let mut discovery_complete = true; + let cache_has_validated_state = cache.scan_since_key.is_some() + || cache.scan_until_key.is_some() + || cache.previous_report.is_some() + || !cache.files.is_empty() + || !cache.days.is_empty() + || !cache.codex_pending_paths.is_empty(); let mut dates = codex_scan_dates(range); if self.options.prefer_newest_codex_sessions_first { dates.reverse(); } for sessions_dir in sessions_dirs { + if !sessions_dir.is_dir() { + if cache_has_codex_path_under(cache, sessions_dir) + || (sessions_dirs.len() == 1 && cache_has_validated_state) + { + discovery_complete = false; + } + continue; + } for date in &dates { if is_cancelled(cancel) { return (candidates, false); @@ -552,7 +593,16 @@ impl CostScanner { .join(date.format("%Y").to_string()) .join(date.format("%m").to_string()) .join(date.format("%d").to_string()); + if !day_dir.exists() { + if cache_has_codex_path_under(cache, &day_dir) { + discovery_complete = false; + } + continue; + } let Ok(entries) = fs::read_dir(&day_dir) else { + if cache_has_codex_path_under(cache, &day_dir) { + discovery_complete = false; + } continue; }; for entry in entries.flatten() { @@ -622,7 +672,7 @@ impl CostScanner { .then_with(|| lhs.path.cmp(&rhs.path)) }); } - (candidates, true) + (candidates, discovery_complete) } #[cfg(test)] @@ -670,6 +720,18 @@ impl CostScanner { let path_key = path.to_string_lossy().to_string(); 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| entry.size > size && entry.parsed_bytes.unwrap_or(0) > size); + if trace_was_pruned && !self.options.is_app_driven() { + // A shrinking trace invalidates the append cursor. Preserve the + // validated cache and queue the path for an explicit cold refresh + // instead of silently replacing history during synchronization. + return CodexFileScanOutcome { + bytes_read: 0, + is_complete: false, + }; + } let session_metadata = JsonlScanner::read_codex_session_metadata(path).unwrap_or_default(); let cached_identity_matches = cached .as_ref() diff --git a/rust/src/cost_scanner/codex/reconciliation.rs b/rust/src/cost_scanner/codex/reconciliation.rs new file mode 100644 index 0000000000..fa35b1aefa --- /dev/null +++ b/rust/src/cost_scanner/codex/reconciliation.rs @@ -0,0 +1,76 @@ +use super::*; + +pub(super) fn paused_codex_summary( + cache: &CostUsageCache, + start_date: NaiveDate, + today: NaiveDate, +) -> CostSummary { + let report = cache + .previous_report + .clone() + .unwrap_or_else(|| JsonlScanner::cached_cost_report_from_days(cache)); + summary_from_cached_report(&report, start_date, today) +} + +/// Return cached Codex files that are provably gone from the portion of the +/// sessions tree covered by this scan. Entries outside the current roots or +/// date directories are intentionally retained for a later scan. +pub(super) fn missing_codex_cache_paths( + cache: &CostUsageCache, + sessions_dirs: &[PathBuf], + range: &CostUsageDayRange, +) -> Vec { + let scanned_date_dirs: Vec = sessions_dirs + .iter() + .flat_map(|sessions_dir| { + codex_scan_dates(range).into_iter().map(|date| { + sessions_dir + .join(date.format("%Y").to_string()) + .join(date.format("%m").to_string()) + .join(date.format("%d").to_string()) + }) + }) + .collect(); + + cache + .files + .keys() + .filter(|path_key| { + let path = Path::new(path_key.as_str()); + let in_scanned_root = sessions_dirs + .iter() + .any(|sessions_dir| path.starts_with(sessions_dir)); + let in_scanned_date = scanned_date_dirs + .iter() + .any(|date_dir| path.starts_with(date_dir)); + in_scanned_root && in_scanned_date && !path.exists() + }) + .cloned() + .collect() +} + +/// Remove cached Codex files that are provably gone after an explicit refresh. +pub(super) fn reconcile_missing_codex_cache_files( + cache: &mut CostUsageCache, + sessions_dirs: &[PathBuf], + range: &CostUsageDayRange, +) { + for path in missing_codex_cache_paths(cache, sessions_dirs, range) { + cache.files.remove(&path); + } + cache + .codex_pending_paths + .retain(|path| Path::new(path).exists()); +} + +/// Whether the cache contains a path that depends on this source partition. +/// Missing optional roots are normal; only a root/date partition that has +/// previously contributed a cached or queued path can make discovery +/// incomplete. +pub(super) fn cache_has_codex_path_under(cache: &CostUsageCache, parent: &Path) -> bool { + cache + .files + .keys() + .chain(cache.codex_pending_paths.iter()) + .any(|path| Path::new(path).starts_with(parent)) +} diff --git a/rust/src/cost_scanner/tests.rs b/rust/src/cost_scanner/tests.rs index 04ca6ed842..e4eaa1fde1 100644 --- a/rust/src/cost_scanner/tests.rs +++ b/rust/src/cost_scanner/tests.rs @@ -26,6 +26,74 @@ fn records_unknown_claude_model_while_using_fallback_cost() { assert!(summary.unknown_models.contains("claude-retired-unknown")); } +#[test] +fn claude_scan_pricing_resolver_reuses_positive_and_negative_resolution() { + let unknown = format!("claude-scan-unknown-{}", std::process::id()); + let mut resolver = ClaudeScanPricingResolver::default(); + + assert!(resolver.is_known("claude-sonnet-4-6")); + assert!(!resolver.is_known(&unknown)); + assert!(resolver.is_known("claude-sonnet-4-6")); + assert!(!resolver.is_known(&unknown)); + assert_eq!(resolver.normalization_cache_misses, 2); + assert_eq!(resolver.resolution_cache_misses, 2); + assert_eq!(resolver.resolutions.len(), 2); + + let mut cost_resolver = ClaudeScanPricingResolver::default(); + let resolved_unknown = cost_resolver.cost_usd_with_cache_ttl(&unknown, 100, 20, 10, 30, 40); + let fallback = + ClaudePricing::cost_usd_with_cache_ttl(FALLBACK_CLAUDE_MODEL, 100, 20, 10, 30, 40); + assert!((resolved_unknown - fallback).abs() < f64::EPSILON); +} + +#[test] +fn claude_scan_pricing_resolver_preserves_tiered_and_cache_ttl_pricing() { + let mut resolver = ClaudeScanPricingResolver::default(); + let cases = [ + ("claude-sonnet-4-6", 240_000, 0, 0, 0, 0), + ("claude-fable-5", 100, 30, 20, 20, 5), + ]; + + for (model, input, cache_create, cache_create_1h, cache_read, output) in cases { + let actual = resolver.cost_usd_with_cache_ttl( + model, + input, + cache_create, + cache_create_1h, + cache_read, + output, + ); + let expected = ClaudePricing::cost_usd_with_cache_ttl( + model, + input, + cache_create, + cache_create_1h, + cache_read, + output, + ); + assert!((actual - expected).abs() < f64::EPSILON, "{model}"); + } +} + +#[test] +fn claude_scan_pricing_resolver_bounds_normalization_memo() { + let mut resolver = ClaudeScanPricingResolver::default(); + for index in 0..(ClaudeScanPricingResolver::MEMO_ENTRY_LIMIT + 8) { + let model = format!("claude-memo-{index}"); + assert_eq!(resolver.normalize(&model), model); + } + assert_eq!( + resolver.normalized_models.len(), + ClaudeScanPricingResolver::MEMO_ENTRY_LIMIT + ); + + let misses = resolver.normalization_cache_misses; + assert_eq!(resolver.normalize("claude-memo-0"), "claude-memo-0"); + assert_eq!(resolver.normalization_cache_misses, misses); + assert_eq!(resolver.normalize("claude-memo-1024"), "claude-memo-1024"); + assert_eq!(resolver.normalization_cache_misses, misses + 1); +} + #[test] fn test_claude_fable_5_pricing() { let cost = ClaudePricing::cost_usd_with_cache_ttl("claude-fable-5", 100, 10, 0, 20, 5); @@ -1655,6 +1723,251 @@ fn incomplete_summary_preserves_previous_report_and_marks_it_non_authoritative() assert!(saved.codex_scan_incomplete); } +#[test] +fn failed_catch_up_pause_preserves_cursor_and_report_until_explicit_refresh() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let pending = write_codex_session_fixture(&sessions, "pending.jsonl", 100); + let report = CachedCostReport { + total_cost_usd: 42.5, + input_tokens: 11, + cached_tokens: 2, + output_tokens: 3, + reasoning_tokens: Some(7), + sessions_count: 7, + updated_at: Some("2026-09-06T00:00:00Z".to_string()), + partial: false, + }; + let mut cache = CostUsageCache { + previous_report: Some(report.clone()), + codex_pending_paths: vec![pending.to_string_lossy().to_string()], + codex_scan_incomplete: true, + ..Default::default() + }; + JsonlScanner::save_cache(ProviderId::Codex, &mut cache, Some(&cache_root)); + + let failed = CostScanner::new(7) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![root.path().join("temporarily-unavailable")]); + let (failed_summary, _, failed_cache) = failed.scan_codex_detailed_with_cache(None); + assert_eq!( + failed_cache.codex_scan_pause_reason, + Some(CodexScanPauseReason::Error( + "Codex session source unavailable".to_string() + )) + ); + assert_eq!(failed_cache.codex_pending_paths, cache.codex_pending_paths); + assert_eq!( + failed_cache + .previous_report + .as_ref() + .map(|saved| saved.total_cost_usd), + Some(report.total_cost_usd) + ); + assert_eq!(failed_summary.total_cost_usd, report.total_cost_usd); + + let background = CostScanner::new(7) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (background_summary, background_stats, background_cache) = + background.scan_codex_detailed_with_cache(None); + assert_eq!(background_stats.files_parsed, 0); + assert_eq!(background_summary.total_cost_usd, report.total_cost_usd); + assert_eq!( + background_cache.codex_pending_paths, + failed_cache.codex_pending_paths + ); + assert_eq!( + background_cache.codex_scan_pause_reason, + failed_cache.codex_scan_pause_reason + ); + + let explicit = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + let (resumed_summary, _, resumed_cache) = explicit.scan_codex_detailed_with_cache(None); + assert!(resumed_cache.codex_scan_pause_reason.is_none()); + assert!(!resumed_cache.codex_scan_incomplete); + assert!(resumed_cache.codex_pending_paths.is_empty()); + assert!(resumed_summary.history_coverage_established); + assert_eq!(resumed_summary.input_tokens, 100); + assert!(resumed_cache.previous_report.is_none()); +} + +#[test] +fn missing_unobserved_sessions_root_does_not_pause_validated_cache() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let missing = root.path().join("optional-sessions"); + let cache_root = root.path().join("cache"); + write_codex_session_fixture(&sessions, "observed.jsonl", 100); + + let initial = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (initial_summary, _) = initial.scan_codex_detailed(None); + assert!(initial_summary.history_coverage_established); + + let mut cache = JsonlScanner::load_cache(ProviderId::Codex, Some(&cache_root)); + cache.last_scan_unix_ms = 1; + JsonlScanner::save_cache(ProviderId::Codex, &mut cache, Some(&cache_root)); + + let background = CostScanner::new(7) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions, missing]); + let (summary, _, saved) = background.scan_codex_detailed_with_cache(None); + + assert!(summary.history_coverage_established); + assert_eq!(summary.input_tokens, 100); + assert!(!saved.codex_scan_incomplete); + assert!(saved.codex_scan_pause_reason.is_none()); +} + +#[test] +fn trace_pruning_preserves_cursor_and_validated_history_until_refresh() { + 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, "pruned.jsonl", 100); + + let initial = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (initial_summary, _, initial_cache) = initial.scan_codex_detailed_with_cache(None); + assert_eq!(initial_summary.input_tokens, 100); + assert!(!initial_cache.codex_scan_incomplete); + + // Keep the next pass outside the scanner debounce while simulating Codex + // retention pruning the same trace file down to a smaller valid payload. + let mut initial_cache = JsonlScanner::load_cache(ProviderId::Codex, Some(&cache_root)); + initial_cache.last_scan_unix_ms = 1; + JsonlScanner::save_cache(ProviderId::Codex, &mut initial_cache, Some(&cache_root)); + let _ = write_codex_session_fixture(&sessions, "pruned.jsonl", 1); + + let background = CostScanner::new(7) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (paused_summary, _, paused_cache) = background.scan_codex_detailed_with_cache(None); + assert_eq!(paused_summary.input_tokens, 100); + assert_eq!( + paused_cache.codex_scan_pause_reason, + Some(CodexScanPauseReason::NoProgress) + ); + assert_eq!( + paused_cache.codex_pending_paths, + vec![path.to_string_lossy().to_string()] + ); + assert!(paused_cache.previous_report.is_some()); + + let explicit = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + let (resumed_summary, _, resumed_cache) = explicit.scan_codex_detailed_with_cache(None); + assert_eq!(resumed_summary.input_tokens, 1); + assert!(resumed_cache.codex_scan_pause_reason.is_none()); + assert!(resumed_cache.codex_pending_paths.is_empty()); + assert!(resumed_cache.previous_report.is_none()); +} + +#[test] +fn disappeared_trace_path_waits_for_explicit_validation() { + 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, "disappeared.jsonl", 100); + + let initial = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (initial_summary, _, _) = initial.scan_codex_detailed_with_cache(None); + assert_eq!(initial_summary.input_tokens, 100); + + let mut cache = JsonlScanner::load_cache(ProviderId::Codex, Some(&cache_root)); + cache.last_scan_unix_ms = 1; + JsonlScanner::save_cache(ProviderId::Codex, &mut cache, Some(&cache_root)); + std::fs::remove_file(&path).unwrap(); + + let background = CostScanner::new(7) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions.clone()]); + let (paused_summary, _, paused_cache) = background.scan_codex_detailed_with_cache(None); + assert_eq!(paused_summary.input_tokens, 100); + assert_eq!( + paused_cache.codex_scan_pause_reason, + Some(CodexScanPauseReason::NoProgress) + ); + assert_eq!( + paused_cache.codex_pending_paths, + vec![path.to_string_lossy().to_string()] + ); + + let explicit = CostScanner::new(7) + .with_options(CostScanOptions::app_driven()) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + let (resumed_summary, _, resumed_cache) = explicit.scan_codex_detailed_with_cache(None); + assert_eq!(resumed_summary.sessions_count, 0); + assert!(resumed_summary.history_coverage_established); + assert!(resumed_cache.codex_scan_pause_reason.is_none()); + assert!(resumed_cache.codex_pending_paths.is_empty()); + assert!(resumed_cache.previous_report.is_none()); +} + +#[test] +fn paused_catch_up_round_trips_without_retrying_in_background() { + let root = tempfile::tempdir().unwrap(); + let sessions = root.path().join("sessions"); + let cache_root = root.path().join("cache"); + let pending = write_codex_session_fixture(&sessions, "pending.jsonl", 100); + let report = CachedCostReport { + total_cost_usd: 9.25, + input_tokens: 21, + cached_tokens: 1, + output_tokens: 5, + reasoning_tokens: None, + sessions_count: 2, + updated_at: Some("2026-09-06T00:00:00Z".to_string()), + partial: true, + }; + let mut cache = CostUsageCache { + previous_report: Some(report.clone()), + codex_pending_paths: vec![pending.to_string_lossy().to_string()], + codex_scan_incomplete: true, + codex_scan_pause_reason: Some(CodexScanPauseReason::NoProgress), + ..Default::default() + }; + + let encoded = serde_json::to_string(&cache).unwrap(); + let decoded: CostUsageCache = serde_json::from_str(&encoded).unwrap(); + assert_eq!( + decoded.codex_scan_pause_reason, + cache.codex_scan_pause_reason + ); + + JsonlScanner::save_cache(ProviderId::Codex, &mut cache, Some(&cache_root)); + let scanner = CostScanner::new(7) + .with_cache_root(&cache_root) + .with_sessions_dirs(vec![sessions]); + let (summary, stats, saved) = scanner.scan_codex_detailed_with_cache(None); + assert_eq!(stats.files_parsed, 0); + assert_eq!(summary.total_cost_usd, report.total_cost_usd); + assert_eq!(saved.codex_pending_paths, cache.codex_pending_paths); + assert_eq!( + saved + .previous_report + .as_ref() + .map(|saved| saved.total_cost_usd), + Some(report.total_cost_usd) + ); + assert_eq!(saved.codex_scan_pause_reason, cache.codex_scan_pause_reason); +} + #[test] fn pending_and_incomplete_round_trip_through_cache_json() { let cache = CostUsageCache { diff --git a/rust/src/providers/grok/billing.rs b/rust/src/providers/grok/billing.rs new file mode 100644 index 0000000000..b69e9dec3c --- /dev/null +++ b/rust/src/providers/grok/billing.rs @@ -0,0 +1,703 @@ +use chrono::{DateTime, TimeZone, Utc}; +use reqwest::header::HeaderMap; + +use crate::core::ProviderError; + +#[derive(Debug, Clone, Copy)] +pub(super) struct GrokBillingSnapshot { + pub(super) used_percent: Option, + pub(super) used_percent_is_wire_published: bool, + pub(super) used_percent_is_implicit_zero: bool, + pub(super) resets_at: Option>, + pub(super) window_minutes: Option, +} + +pub(super) fn validate_grpc_headers(headers: &HeaderMap) -> Result<(), ProviderError> { + if let Some(status) = headers + .get("grpc-status") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.parse::().ok()) + && status != 0 + { + if status == 16 { + return Err(ProviderError::AuthRequired); + } + return Err(ProviderError::Other(format!( + "Grok RPC failed with status {status}" + ))); + } + Ok(()) +} + +pub(super) fn parse_grpc_web_response(data: &[u8]) -> Result { + parse_grpc_web_response_at(data, Utc::now()) +} + +fn parse_grpc_web_response_at( + data: &[u8], + now: DateTime, +) -> Result { + let mut payloads = grpc_web_data_frames(data); + if payloads.is_empty() && looks_like_protobuf_payload(data) { + payloads.push(data.to_vec()); + } + if payloads.is_empty() { + return Err(ProviderError::Parse( + "Grok web billing returned no payload".to_string(), + )); + } + let mut scan = ProtoScan::default(); + for payload in &payloads { + scan.scan_message(payload, &mut Vec::new(), 0); + } + let percent_fields: Vec<&Fixed32Field> = scan + .fixed32 + .iter() + .filter(|field| field.path.last() == Some(&1)) + .collect(); + let mut valid_percent_fields: Vec<&Fixed32Field> = percent_fields + .iter() + .copied() + .filter(|field| field.value.is_finite() && (0.0..=100.0).contains(&field.value)) + .collect(); + valid_percent_fields.sort_by(|a, b| { + a.path + .len() + .cmp(&b.path.len()) + .then_with(|| a.order.cmp(&b.order)) + }); + let conflicting_percent = percent_fields.len() != valid_percent_fields.len() + || valid_percent_fields.first().is_some_and(|first| { + valid_percent_fields + .iter() + .any(|field| field.value != first.value) + }); + let parsed_percent = if scan.is_complete && !conflicting_percent { + valid_percent_fields.first().map(|field| field.value as f64) + } else { + None + }; + + let reset_fields: Vec<(&VarintField, DateTime)> = scan + .varints + .iter() + .filter_map(|field| varint_timestamp(field).map(|dt| (field, dt))) + .collect(); + let future_resets: Vec> = reset_fields + .iter() + .filter(|(_, dt)| *dt > now) + .map(|(_, dt)| *dt) + .collect(); + let resets_at = reset_fields + .iter() + .filter(|(field, dt)| field.path.as_slice() == [1, 5, 1] && *dt > now) + .map(|(_, dt)| *dt) + .min() + .or_else(|| future_resets.into_iter().min()); + + let has_usage_period = scan.varints.iter().any(|field| { + field.path.starts_with(&[1, 6]) + || (field.path.as_slice() == [1, 8, 1] && (field.value == 1 || field.value == 2)) + }); + let no_usage_yet = parsed_percent.is_none() + && scan.fixed32.is_empty() + && resets_at.is_some() + && has_usage_period; + let window_minutes = current_period_window_minutes(&scan, now); + let has_active_current_period = window_minutes.is_some(); + let used_percent_is_implicit_zero = + no_usage_yet && payloads.len() == 1 && scan.is_complete && has_active_current_period; + let used_percent = parsed_percent.or_else(|| used_percent_is_implicit_zero.then_some(0.0)); + Ok(GrokBillingSnapshot { + used_percent, + used_percent_is_wire_published: parsed_percent.is_some(), + used_percent_is_implicit_zero, + resets_at, + window_minutes, + }) +} + +fn looks_like_protobuf_payload(data: &[u8]) -> bool { + let Some(&first) = data.first() else { + return false; + }; + let field_number = first >> 3; + let wire_type = first & 0x07; + field_number > 0 && matches!(wire_type, 0 | 1 | 2 | 5) +} + +fn varint_timestamp(field: &VarintField) -> Option> { + // Varint timestamps are Unix seconds inside the range checked below. + #[allow( + clippy::cast_possible_wrap, + reason = "varint timestamps are bounded to the Unix-seconds range checked below" + )] + let seconds = field.value as i64; + (1_700_000_000..=2_100_000_000) + .contains(&field.value) + .then(|| Utc.timestamp_opt(seconds, 0).single()) + .flatten() +} + +fn current_period_window_minutes(scan: &ProtoScan, now: DateTime) -> Option { + let period_type = unique_varint_at_path(scan, &[1, 8, 1])?; + if period_type != 1 && period_type != 2 { + return None; + } + let timestamp_at = |path: &[u64]| { + unique_varint_at_path(scan, path).and_then(|value| { + varint_timestamp(&VarintField { + path: path.to_vec(), + value, + }) + }) + }; + let start = timestamp_at(&[1, 8, 2, 1])?; + let end = timestamp_at(&[1, 8, 3, 1])?; + if start > now || end <= now || end <= start { + return None; + } + u32::try_from((end - start).num_minutes()) + .ok() + .filter(|minutes| *minutes > 0) +} + +fn unique_varint_at_path(scan: &ProtoScan, path: &[u64]) -> Option { + let mut values = scan + .varints + .iter() + .filter(|field| field.path.as_slice() == path) + .map(|field| field.value); + let first = values.next()?; + values.all(|value| value == first).then_some(first) +} + +fn grpc_web_data_frames(data: &[u8]) -> Vec> { + let mut frames = Vec::new(); + let mut index = 0; + while index < data.len() { + if index + 5 > data.len() { + return Vec::new(); + } + let flags = data[index]; + let len = ((data[index + 1] as usize) << 24) + | ((data[index + 2] as usize) << 16) + | ((data[index + 3] as usize) << 8) + | (data[index + 4] as usize); + let start = index + 5; + let Some(end) = start.checked_add(len) else { + return Vec::new(); + }; + if end > data.len() { + return Vec::new(); + } + if flags & 0x80 == 0 { + frames.push(data[start..end].to_vec()); + } + index = end; + } + frames +} + +struct ProtoScan { + fixed32: Vec, + varints: Vec, + order: usize, + is_complete: bool, +} + +impl Default for ProtoScan { + fn default() -> Self { + Self { + fixed32: Vec::new(), + varints: Vec::new(), + order: 0, + is_complete: true, + } + } +} + +struct Fixed32Field { + path: Vec, + value: f32, + order: usize, +} + +struct VarintField { + path: Vec, + value: u64, +} + +impl ProtoScan { + fn scan_message(&mut self, data: &[u8], path: &mut Vec, depth: usize) { + if depth > 8 { + self.is_complete = false; + return; + } + let mut i = 0; + while i < data.len() { + let field_start = i; + let Some((field, wire, next)) = read_key(data, i) else { + self.is_complete = false; + i = field_start.saturating_add(1); + continue; + }; + i = next; + path.push(field); + let Some(next) = self.scan_field(data, i, path, depth, wire) else { + self.is_complete = false; + path.pop(); + i = field_start.saturating_add(1); + continue; + }; + i = next; + path.pop(); + } + } + + fn scan_field( + &mut self, + data: &[u8], + i: usize, + path: &mut Vec, + depth: usize, + wire: u64, + ) -> Option { + if (path.as_slice() == [1, 1] && wire != 5) || (is_known_billing_message(path) && wire != 2) + { + return None; + } + match wire { + 0 => self.scan_varint(data, i, path), + 2 => self.scan_length_delimited(data, i, path, depth), + 5 => self.scan_fixed32(data, i, path), + 1 => i.checked_add(8).filter(|end| *end <= data.len()), + _ => None, + } + } + + fn scan_varint(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option { + let (value, next) = read_varint(data, i)?; + self.varints.push(VarintField { + path: path.to_vec(), + value, + }); + Some(next) + } + + fn scan_length_delimited( + &mut self, + data: &[u8], + i: usize, + path: &mut Vec, + depth: usize, + ) -> Option { + let (len, next) = read_varint(data, i)?; + let start = next; + let len_usize = usize::try_from(len).ok()?; + let end = start.checked_add(len_usize)?; + if end > data.len() { + return None; + } + if depth < 4 && is_known_billing_message(path) { + self.scan_message(&data[start..end], path, depth + 1); + } + Some(end) + } + + fn scan_fixed32(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option { + let end = i.checked_add(4)?; + if end > data.len() { + return None; + } + let bytes = [data[i], data[i + 1], data[i + 2], data[i + 3]]; + self.fixed32.push(Fixed32Field { + path: path.to_vec(), + value: f32::from_le_bytes(bytes), + order: self.order, + }); + self.order += 1; + Some(end) + } +} + +fn read_key(data: &[u8], i: usize) -> Option<(u64, u64, usize)> { + let (key, next) = read_varint(data, i)?; + let field = key >> 3; + (field > 0 && field <= 536_870_911).then_some((field, key & 0x07, next)) +} + +fn read_varint(data: &[u8], mut i: usize) -> Option<(u64, usize)> { + let mut value = 0u64; + let mut shift = 0; + while i < data.len() && shift < 64 { + let b = data[i]; + i += 1; + if shift == 63 && b > 1 { + return None; + } + value |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Some((value, i)); + } + shift += 7; + } + None +} + +fn is_known_billing_message(path: &[u64]) -> bool { + // Only these descriptor-declared messages are recursively decoded. Other + // length-delimited fields may be opaque bytes and must not affect billing. + matches!( + path, + [1] | [1, 2] + | [1, 3] + | [1, 4] + | [1, 5] + | [1, 6] + | [1, 7] + | [1, 8] + | [1, 12] + | [1, 6, 1] + | [1, 6, 2] + | [1, 6, 3] + | [1, 8, 2] + | [1, 8, 3] + | [1, 6, 3, 2] + | [1, 6, 3, 3] + ) +} + +#[cfg(test)] +mod tests { + use super::super::{primary_label_for_cycle_minutes, result_from_billing}; + use super::*; + + #[test] + fn splits_grpc_web_data_frames() { + let data = [0, 0, 0, 0, 2, 1, 2, 0x80, 0, 0, 0, 1, b'x']; + assert_eq!(grpc_web_data_frames(&data), vec![vec![1, 2]]); + } + + #[test] + fn current_period_paths_define_the_full_cycle() { + let now = Utc.timestamp_opt(1_800_000_000, 0).single().unwrap(); + let start = now - chrono::Duration::days(25); + let end = now + chrono::Duration::days(6); + let scan = ProtoScan { + fixed32: Vec::new(), + varints: vec![ + VarintField { + path: vec![1, 8, 1], + value: 1, + }, + VarintField { + path: vec![1, 8, 2, 1], + value: u64::try_from(start.timestamp()).unwrap(), + }, + VarintField { + path: vec![1, 8, 3, 1], + value: u64::try_from(end.timestamp()).unwrap(), + }, + ], + order: 0, + is_complete: true, + }; + + assert_eq!( + current_period_window_minutes(&scan, now), + Some(31 * 24 * 60) + ); + assert_eq!( + primary_label_for_cycle_minutes(current_period_window_minutes(&scan, now).unwrap()), + Some("Monthly") + ); + } + + #[test] + fn captured_active_period_with_omitted_usage_is_zero() { + let frame = hex_bytes( + "00000000440a4212001a00220b0887a8c6d40610f0d7dd142a0b08879debd40610f0d7dd14".to_owned() + + "421c0802120b0887a8c6d40610f0d7dd141a0b08879debd40610f0d7dd14580162006801" + + "800000000f677270632d7374617475733a300d0a", + ); + let parsed = parse_grpc_web_response_at(&frame, fixed_time(1_788_000_000)).unwrap(); + + assert_eq!(parsed.used_percent, Some(0.0)); + assert!(parsed.used_percent_is_implicit_zero); + assert!(parsed.resets_at.is_some()); + assert!(parsed.window_minutes.is_some()); + } + + #[test] + fn complete_monthly_and_weekly_periods_support_implicit_zero() { + for period_type in [1, 2] { + let parsed = parse_grpc_web_response_at( + &payload(period_type, Some(1_787_000_000), true, &[]), + fixed_time(1_788_000_000), + ) + .unwrap(); + + assert_eq!(parsed.used_percent, Some(0.0)); + assert!(parsed.used_percent_is_implicit_zero); + let result = result_from_billing(parsed, "grok-web", None, None, None); + assert_eq!(result.usage.primary.used_percent, 0.0); + } + } + + #[test] + fn malformed_frames_cannot_turn_unknown_usage_into_zero() { + let malformed = [ + vec![0x00], + fixed64_field(&[0x01]), + vec![0x02, 0x00], + fixed64_field(&[0x81, 0x80, 0x80, 0x80, 0x10]), + fixed64_field(&[0x89, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x02]), + vec![0x0d, 0x00], + vec![0x08, 0x01], + vec![0x72, 0x04, 0x08], + vec![0x70, 0x80], + ]; + + for suffix in malformed { + let parsed = parse_grpc_web_response_at( + &payload(2, Some(1_787_000_000), true, &suffix), + fixed_time(1_788_000_000), + ) + .unwrap(); + assert_eq!(parsed.used_percent, None); + assert!(!parsed.used_percent_is_implicit_zero); + } + } + + #[test] + fn valid_unknown_fields_preserve_implicit_zero() { + let unknown_fields = [fixed64_field(&[0xf9, 0xff, 0xff, 0xff, 0x0f]), { + let mut field = vec![0x70]; + field.extend(varint(u64::MAX)); + field + }]; + + for suffix in unknown_fields { + let parsed = parse_grpc_web_response_at( + &payload(2, Some(1_787_000_000), true, &suffix), + fixed_time(1_788_000_000), + ) + .unwrap(); + assert_eq!(parsed.used_percent, Some(0.0)); + assert!(parsed.used_percent_is_implicit_zero); + } + } + + #[test] + fn opaque_unknown_fields_cannot_invalidate_or_invent_billing_values() { + let opaque_payloads = [ + fixed64_field(&[0x01]), + vec![0x0d, 0x00, 0x00, 0x14, 0x42], + { + let mut field = vec![0x08]; + field.extend(varint(1_788_500_000)); + field + }, + ]; + + for at_root in [false, true] { + for bytes in &opaque_payloads { + let opaque_field = length_field(14, bytes); + let mut raw = payload( + 2, + Some(1_787_000_000), + true, + if at_root { &[] } else { &opaque_field }, + ); + if at_root { + raw.extend(opaque_field); + } + let parsed = parse_grpc_web_response_at(&raw, fixed_time(1_788_000_000)).unwrap(); + + assert_eq!(parsed.used_percent, Some(0.0)); + assert!(parsed.used_percent_is_implicit_zero); + assert_eq!(parsed.resets_at, Some(fixed_time(1_789_000_000))); + } + } + } + + #[test] + fn malformed_known_messages_still_prevent_implicit_zero() { + let paths = [ + vec![1], + vec![1, 2], + vec![1, 3], + vec![1, 4], + vec![1, 5], + vec![1, 6], + vec![1, 7], + vec![1, 8], + vec![1, 12], + vec![1, 6, 1], + vec![1, 6, 2], + vec![1, 6, 3], + vec![1, 8, 2], + vec![1, 8, 3], + vec![1, 6, 3, 2], + vec![1, 6, 3, 3], + ]; + + for path in paths { + let mut raw = payload(2, Some(1_787_000_000), true, &[]); + raw.extend(message(&path, &fixed64_field(&[0x01]))); + let parsed = parse_grpc_web_response_at(&raw, fixed_time(1_788_000_000)).unwrap(); + + assert_eq!(parsed.used_percent, None, "path {path:?}"); + assert!(!parsed.used_percent_is_implicit_zero, "path {path:?}"); + } + } + + #[test] + fn historical_period_timestamps_remain_readable_without_current_usage() { + let mut timestamp = vec![0x08]; + timestamp.extend(varint(1_789_000_000)); + let parsed = parse_grpc_web_response_at( + &message(&[1, 6, 3, 3], ×tamp), + fixed_time(1_788_000_000), + ) + .unwrap(); + + assert_eq!(parsed.resets_at, Some(fixed_time(1_789_000_000))); + assert_eq!(parsed.used_percent, None); + assert!(!parsed.used_percent_is_implicit_zero); + } + + #[test] + fn conflicting_percent_tags_remain_unknown() { + let mut suffix = fixed32_field(20.0); + suffix.extend(fixed32_field(30.0)); + let parsed = parse_grpc_web_response_at( + &payload(2, Some(1_787_000_000), true, &suffix), + fixed_time(1_788_000_000), + ) + .unwrap(); + + assert_eq!(parsed.used_percent, None); + assert!(!parsed.used_percent_is_implicit_zero); + } + + #[test] + fn malformed_payload_does_not_publish_a_partial_percent() { + let mut suffix = fixed32_field(23.5); + suffix.extend([0x70, 0x80]); + let parsed = parse_grpc_web_response_at( + &payload(2, Some(1_787_000_000), true, &suffix), + fixed_time(1_788_000_000), + ) + .unwrap(); + + assert_eq!(parsed.used_percent, None); + assert!(!parsed.used_percent_is_wire_published); + assert!(!parsed.used_percent_is_implicit_zero); + } + + #[test] + fn explicit_percent_remains_unchanged() { + let parsed = parse_grpc_web_response_at( + &payload(2, Some(1_787_000_000), true, &fixed32_field(23.5)), + fixed_time(1_788_000_000), + ) + .unwrap(); + + assert_eq!(parsed.used_percent, Some(23.5)); + assert!(parsed.used_percent_is_wire_published); + assert!(!parsed.used_percent_is_implicit_zero); + } + + #[test] + fn incomplete_or_non_current_periods_remain_unknown() { + for period_type in [0, 3] { + let parsed = parse_grpc_web_response_at( + &payload(period_type, Some(1_787_000_000), true, &[]), + fixed_time(1_788_000_000), + ) + .unwrap(); + assert_eq!(parsed.used_percent, None); + assert!(!parsed.used_percent_is_implicit_zero); + } + + for (start, include_start) in [(Some(1_800_000_000), true), (Some(1_787_000_000), false)] { + let parsed = parse_grpc_web_response_at( + &payload(2, start, include_start, &[]), + fixed_time(1_788_000_000), + ) + .unwrap(); + assert_eq!(parsed.used_percent, None); + assert!(!parsed.used_percent_is_implicit_zero); + } + } + + fn fixed_time(seconds: i64) -> DateTime { + Utc.timestamp_opt(seconds, 0).single().unwrap() + } + + fn payload(period_type: u8, start: Option, include_start: bool, suffix: &[u8]) -> Vec { + let mut period = vec![0x08, period_type]; + if include_start { + let start = start.expect("start timestamp required"); + let mut timestamp = vec![0x08]; + timestamp.extend(varint(start)); + period.extend(length_field(2, ×tamp)); + } + let mut end_timestamp = vec![0x08]; + end_timestamp.extend(varint(1_789_000_000)); + period.extend(length_field(3, &end_timestamp)); + + let mut config = length_field(8, &period); + config.extend(suffix); + length_field(1, &config) + } + + fn message(path: &[u64], contents: &[u8]) -> Vec { + path.iter().rev().fold(contents.to_vec(), |payload, field| { + length_field(*field, &payload) + }) + } + + fn length_field(field: u64, contents: &[u8]) -> Vec { + let mut encoded = varint((field << 3) | 2); + encoded.extend(varint(contents.len() as u64)); + encoded.extend(contents); + encoded + } + + fn fixed32_field(value: f32) -> Vec { + let mut encoded = vec![0x0d]; + encoded.extend(value.to_le_bytes()); + encoded + } + + fn fixed64_field(tag: &[u8]) -> Vec { + let mut encoded = tag.to_vec(); + encoded.extend([0; 8]); + encoded + } + + fn varint(mut value: u64) -> Vec { + let mut encoded = Vec::new(); + while value >= 0x80 { + encoded.push(u8::try_from(value & 0x7f).expect("masked varint byte fits u8") | 0x80); + value >>= 7; + } + encoded.push(u8::try_from(value).expect("terminal varint byte fits u8")); + encoded + } + + fn hex_bytes(hex: String) -> Vec { + let (pairs, remainder) = hex.as_bytes().as_chunks::<2>(); + assert!(remainder.is_empty(), "hex fixture length must be even"); + pairs + .iter() + .map(|chunk| { + let text = std::str::from_utf8(chunk).unwrap(); + u8::from_str_radix(text, 16).unwrap() + }) + .collect() + } +} diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index bce4212a59..021d915eb9 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -3,10 +3,11 @@ //! Uses the grok.com billing gRPC-web endpoint via either browser cookies or //! `~/.grok/auth.json` produced by `grok login`. +mod billing; pub mod local_sessions; use async_trait::async_trait; -use chrono::{DateTime, TimeZone, Utc}; +use chrono::{DateTime, Utc}; use reqwest::Client; use serde_json::Value; use std::path::PathBuf; @@ -19,6 +20,8 @@ use crate::core::{ RateWindow, SourceMode, UsageSnapshot, }; +use self::billing::GrokBillingSnapshot; + const BILLING_ENDPOINT: &str = "https://grok.com/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig"; const CLI_SETTINGS_ENDPOINT: &str = "https://cli-chat-proxy.grok.com/v1/settings"; @@ -194,8 +197,8 @@ impl GrokProvider { "Grok web billing returned status {status}" ))); } - validate_grpc_headers(&headers)?; - parse_grpc_web_response(&bytes) + billing::validate_grpc_headers(&headers)?; + billing::parse_grpc_web_response(&bytes) } fn detect_cli_version() -> Option { @@ -419,13 +422,6 @@ fn text_field(value: &Value, key: &str) -> Option { .map(ToOwned::to_owned) } -#[derive(Debug, Clone, Copy)] -struct GrokBillingSnapshot { - used_percent: Option, - resets_at: Option>, - window_minutes: Option, -} - /// Classify Grok from the full billing-cycle duration, not time remaining. /// This preserves the upstream #2431/#2566 invariant that a monthly plan near /// its reset must not become a weekly plan. @@ -460,7 +456,10 @@ fn result_from_billing( let primary_label = billing .window_minutes .and_then(primary_label_for_cycle_minutes); - let primary = match billing.used_percent { + let published_percent = billing.used_percent.filter(|_| { + billing.used_percent_is_wire_published || billing.used_percent_is_implicit_zero + }); + let primary = match published_percent { Some(used_percent) => RateWindow::with_details( used_percent, billing.window_minutes, @@ -484,23 +483,6 @@ fn result_from_billing( ProviderFetchResult::new(usage, source_label) } -fn validate_grpc_headers(headers: &reqwest::header::HeaderMap) -> Result<(), ProviderError> { - if let Some(status) = headers - .get("grpc-status") - .and_then(|value| value.to_str().ok()) - .and_then(|value| value.parse::().ok()) - && status != 0 - { - if status == 16 { - return Err(ProviderError::AuthRequired); - } - return Err(ProviderError::Other(format!( - "Grok RPC failed with status {status}" - ))); - } - Ok(()) -} - /// Whether a cookie-path error should invalidate the cached browser session. fn is_cookie_authentication_failure(err: &ProviderError) -> bool { matches!(err, ProviderError::AuthRequired) @@ -528,452 +510,6 @@ fn cookie_refresh_action( } } -fn parse_grpc_web_response(data: &[u8]) -> Result { - let frames = grpc_web_data_frames(data); - if frames.is_empty() { - return Err(ProviderError::Parse( - "Grok web billing returned no payload".to_string(), - )); - } - let mut scan = ProtoScan::default(); - for frame in frames { - scan.scan_message(&frame, &mut Vec::new(), 0); - } - let used_percent = scan - .fixed32 - .iter() - .filter(|field| { - field.path.last() == Some(&1) - && field.value.is_finite() - && field.value >= 0.0 - && field.value <= 100.0 - }) - .min_by(|a, b| { - a.path - .len() - .cmp(&b.path.len()) - .then_with(|| a.order.cmp(&b.order)) - }) - .map(|field| field.value as f64); - - let now = Utc::now(); - let resets_at = scan - .varints - .iter() - .filter_map(varint_timestamp) - .filter(|dt| *dt > now) - .min(); - Ok(GrokBillingSnapshot { - used_percent, - resets_at, - window_minutes: current_period_window_minutes(&scan, now), - }) -} - -fn varint_timestamp(field: &VarintField) -> Option> { - // Varint timestamps are Unix seconds inside the range checked below. - #[allow( - clippy::cast_possible_wrap, - reason = "varint timestamps are bounded to the Unix-seconds range checked below" - )] - let seconds = field.value as i64; - (1_700_000_000..=2_100_000_000) - .contains(&field.value) - .then(|| Utc.timestamp_opt(seconds, 0).single()) - .flatten() -} - -fn current_period_window_minutes(scan: &ProtoScan, now: DateTime) -> Option { - let timestamp_at = |path: &[u64]| { - scan.varints - .iter() - .find(|field| field.path.as_slice() == path) - .and_then(varint_timestamp) - }; - let start = timestamp_at(&[1, 8, 2, 1])?; - let end = timestamp_at(&[1, 8, 3, 1])?; - if start > now || end <= now || end <= start { - return None; - } - u32::try_from((end - start).num_minutes()) - .ok() - .filter(|minutes| *minutes > 0) -} - -fn grpc_web_data_frames(data: &[u8]) -> Vec> { - let mut frames = Vec::new(); - let mut index = 0; - while index + 5 <= data.len() { - let flags = data[index]; - let len = ((data[index + 1] as usize) << 24) - | ((data[index + 2] as usize) << 16) - | ((data[index + 3] as usize) << 8) - | (data[index + 4] as usize); - let start = index + 5; - let end = start.saturating_add(len); - if end > data.len() { - break; - } - if flags & 0x80 == 0 { - frames.push(data[start..end].to_vec()); - } - index = end; - } - frames -} - -#[derive(Default)] -struct ProtoScan { - fixed32: Vec, - varints: Vec, - order: usize, -} - -struct Fixed32Field { - path: Vec, - value: f32, - order: usize, -} - -struct VarintField { - path: Vec, - value: u64, -} - -impl ProtoScan { - fn scan_message(&mut self, data: &[u8], path: &mut Vec, depth: usize) { - if depth > 8 { - return; - } - let mut i = 0; - while i < data.len() { - let Some((field, wire, next)) = read_key(data, i) else { - break; - }; - i = next; - path.push(field); - let Some(next) = self.scan_field(data, i, path, depth, wire) else { - path.pop(); - break; - }; - i = next; - path.pop(); - } - } - - fn scan_field( - &mut self, - data: &[u8], - i: usize, - path: &mut Vec, - depth: usize, - wire: u64, - ) -> Option { - match wire { - 0 => self.scan_varint(data, i, path), - 2 => self.scan_length_delimited(data, i, path, depth), - 5 => self.scan_fixed32(data, i, path), - 1 => Some(i.saturating_add(8)), - _ => None, - } - } - - fn scan_varint(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option { - let (value, next) = read_varint(data, i)?; - self.varints.push(VarintField { - path: path.to_vec(), - value, - }); - Some(next) - } - - fn scan_length_delimited( - &mut self, - data: &[u8], - i: usize, - path: &mut Vec, - depth: usize, - ) -> Option { - let (len, next) = read_varint(data, i)?; - let start = next; - // Varint field lengths are bounded by the containing message buffer. - #[allow( - clippy::cast_possible_truncation, - reason = "varint field lengths are bounded by the containing message buffer" - )] - let len_usize = len as usize; - let end = start.saturating_add(len_usize); - if end <= data.len() { - self.scan_message(&data[start..end], path, depth + 1); - Some(end) - } else { - None - } - } - - fn scan_fixed32(&mut self, data: &[u8], i: usize, path: &[u64]) -> Option { - if i + 4 > data.len() { - return None; - } - let bytes = [data[i], data[i + 1], data[i + 2], data[i + 3]]; - self.fixed32.push(Fixed32Field { - path: path.to_vec(), - value: f32::from_le_bytes(bytes), - order: self.order, - }); - self.order += 1; - Some(i + 4) - } -} - -fn read_key(data: &[u8], i: usize) -> Option<(u64, u64, usize)> { - let (key, next) = read_varint(data, i)?; - Some((key >> 3, key & 0x07, next)) -} - -fn read_varint(data: &[u8], mut i: usize) -> Option<(u64, usize)> { - let mut value = 0u64; - let mut shift = 0; - while i < data.len() && shift < 64 { - let b = data[i]; - i += 1; - value |= u64::from(b & 0x7f) << shift; - if b & 0x80 == 0 { - return Some((value, i)); - } - shift += 7; - } - None -} - #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn grok_plan_prefers_subscription_tier_display_names() { - assert_eq!( - grok_plan_display_name(Some("SuperGrok Heavy")), - Some("SuperGrok Heavy".to_string()) - ); - assert_eq!( - grok_plan_display_name(Some("heavy")), - Some("SuperGrok Heavy".to_string()) - ); - assert_eq!( - grok_plan_display_name(Some("SuperGrok")), - Some("SuperGrok".to_string()) - ); - assert_eq!( - grok_plan_display_name(Some(" custom ")), - Some("custom".to_string()) - ); - } - - #[test] - fn parses_auth_file_prefer_oidc() { - let auth = r#"{ - "https://accounts.x.ai/sign-in": {"key": "legacy"}, - "https://auth.x.ai::abc": {"key": "oidc", "auth_mode": "oidc", "email": "u@example.com"} - }"#; - let parsed = GrokCredentials::parse_for_kind(auth, GrokAuthKind::OAuth).unwrap(); - assert_eq!(parsed.access_token, "oidc"); - assert_eq!(parsed.login_method().as_deref(), Some("SuperGrok")); - } - - #[test] - fn cli_and_oauth_select_distinct_auth_entries() { - let auth = r#"{ - "https://accounts.x.ai/sign-in": {"key": "cli-token", "auth_mode": "session"}, - "https://auth.x.ai::abc": {"key": "oauth-token", "auth_mode": "oidc"} - }"#; - assert_eq!( - GrokCredentials::parse_for_kind(auth, GrokAuthKind::Cli) - .unwrap() - .access_token, - "cli-token" - ); - assert_eq!( - GrokCredentials::parse_for_kind(auth, GrokAuthKind::OAuth) - .unwrap() - .access_token, - "oauth-token" - ); - } - #[test] - fn splits_grpc_web_data_frames() { - let data = [0, 0, 0, 0, 2, 1, 2, 0x80, 0, 0, 0, 1, b'x']; - assert_eq!(grpc_web_data_frames(&data), vec![vec![1, 2]]); - } - - #[test] - fn cookie_refresh_uses_cache_when_present() { - assert_eq!( - cookie_refresh_action(true, None), - CookieRefreshAction::UseCached - ); - } - - #[test] - fn cookie_refresh_reimports_on_auth_failure() { - assert_eq!( - cookie_refresh_action(true, Some(&ProviderError::AuthRequired)), - CookieRefreshAction::ReimportBrowser - ); - assert_eq!( - cookie_refresh_action(false, None), - CookieRefreshAction::ReimportBrowser - ); - } - - #[test] - fn cookie_refresh_gives_up_on_non_auth_errors() { - assert_eq!( - cookie_refresh_action(true, Some(&ProviderError::Other("network down".into()))), - CookieRefreshAction::GiveUp - ); - } - - #[test] - fn is_cookie_auth_failure_only_auth_required() { - assert!(is_cookie_authentication_failure( - &ProviderError::AuthRequired - )); - assert!(!is_cookie_authentication_failure(&ProviderError::NoCookies)); - } - - #[test] - fn cookie_billing_stays_siloed_from_auth_file_identity() { - let result = result_from_cookie_billing(GrokBillingSnapshot { - used_percent: Some(23.0), - resets_at: None, - window_minutes: None, - }); - assert_eq!(result.source_label, "grok-browser"); - assert!(result.usage.account_email.is_none()); - assert!(result.usage.account_organization.is_none()); - assert!(result.usage.login_method.is_none()); - } - #[test] - fn billing_snapshot_uses_full_weekly_cycle_for_pace() { - let now = Utc::now(); - let resets = now + chrono::Duration::days(2); - let result = result_from_billing( - GrokBillingSnapshot { - used_percent: Some(12.0), - resets_at: Some(resets), - window_minutes: Some(crate::core::WEEKLY_WINDOW_MINUTES), - }, - "web", - None, - None, - Some("SuperGrok".into()), - ); - assert_eq!( - result.usage.primary.window_minutes, - Some(crate::core::WEEKLY_WINDOW_MINUTES) - ); - assert_eq!(result.usage.primary_label.as_deref(), Some("Weekly")); - let pace = crate::core::UsagePace::weekly( - &result.usage.primary, - Some(now), - crate::core::WEEKLY_WINDOW_MINUTES, - ); - assert!(pace.is_some(), "weekly window + reset must yield pace"); - } - - #[test] - fn monthly_cycle_stays_monthly_with_six_days_remaining() { - let now = Utc::now(); - let resets = now + chrono::Duration::days(6); - let monthly_minutes = 31 * 24 * 60; - let result = result_from_billing( - GrokBillingSnapshot { - used_percent: Some(40.0), - resets_at: Some(resets), - window_minutes: Some(monthly_minutes), - }, - "cli", - None, - None, - Some("SuperGrok Heavy".into()), - ); - assert_eq!(result.usage.primary_label.as_deref(), Some("Monthly")); - assert_eq!(result.usage.primary.window_minutes, Some(monthly_minutes)); - } - - #[test] - fn reset_distance_alone_does_not_invent_a_cadence() { - let resets = Utc::now() + chrono::Duration::days(6); - let result = result_from_billing( - GrokBillingSnapshot { - used_percent: Some(80.0), - resets_at: Some(resets), - window_minutes: None, - }, - "web", - None, - None, - Some("SuperGrok".into()), - ); - assert_eq!(result.usage.primary.window_minutes, None); - assert_eq!(result.usage.primary_label, None); - } - - #[test] - fn current_period_paths_define_the_full_cycle() { - let now = Utc.timestamp_opt(1_800_000_000, 0).single().unwrap(); - let start = now - chrono::Duration::days(25); - let end = now + chrono::Duration::days(6); - let scan = ProtoScan { - fixed32: Vec::new(), - varints: vec![ - VarintField { - path: vec![1, 8, 2, 1], - value: u64::try_from(start.timestamp()).unwrap(), - }, - VarintField { - path: vec![1, 8, 3, 1], - value: u64::try_from(end.timestamp()).unwrap(), - }, - ], - order: 0, - }; - - assert_eq!( - current_period_window_minutes(&scan, now), - Some(31 * 24 * 60) - ); - assert_eq!( - primary_label_for_cycle_minutes(current_period_window_minutes(&scan, now).unwrap()), - Some("Monthly") - ); - } - - #[test] - fn period_only_billing_is_informational_not_zero_usage() { - let resets = Utc::now() + chrono::Duration::days(6); - let result = result_from_billing( - GrokBillingSnapshot { - used_percent: None, - resets_at: Some(resets), - window_minutes: None, - }, - "cli", - Some("user@example.com".into()), - None, - Some("SuperGrok Heavy".into()), - ); - - assert!(result.usage.primary.is_informational); - assert_eq!(result.usage.primary.resets_at, Some(resets)); - assert_eq!( - result.usage.account_email.as_deref(), - Some("user@example.com") - ); - assert_eq!( - result.usage.login_method.as_deref(), - Some("SuperGrok Heavy") - ); - } -} +#[path = "tests.rs"] +mod tests; diff --git a/rust/src/providers/grok/tests.rs b/rust/src/providers/grok/tests.rs new file mode 100644 index 0000000000..dc64970f55 --- /dev/null +++ b/rust/src/providers/grok/tests.rs @@ -0,0 +1,222 @@ +use super::*; +use chrono::TimeZone; + +#[test] +fn grok_plan_prefers_subscription_tier_display_names() { + assert_eq!( + grok_plan_display_name(Some("SuperGrok Heavy")), + Some("SuperGrok Heavy".to_string()) + ); + assert_eq!( + grok_plan_display_name(Some("heavy")), + Some("SuperGrok Heavy".to_string()) + ); + assert_eq!( + grok_plan_display_name(Some("SuperGrok")), + Some("SuperGrok".to_string()) + ); + assert_eq!( + grok_plan_display_name(Some(" custom ")), + Some("custom".to_string()) + ); +} + +#[test] +fn parses_auth_file_prefer_oidc() { + let auth = r#"{ + "https://accounts.x.ai/sign-in": {"key": "legacy"}, + "https://auth.x.ai::abc": {"key": "oidc", "auth_mode": "oidc", "email": "u@example.com"} + }"#; + let parsed = GrokCredentials::parse_for_kind(auth, GrokAuthKind::OAuth).unwrap(); + assert_eq!(parsed.access_token, "oidc"); + assert_eq!(parsed.login_method().as_deref(), Some("SuperGrok")); +} + +#[test] +fn cli_and_oauth_select_distinct_auth_entries() { + let auth = r#"{ + "https://accounts.x.ai/sign-in": {"key": "cli-token", "auth_mode": "session"}, + "https://auth.x.ai::abc": {"key": "oauth-token", "auth_mode": "oidc"} + }"#; + assert_eq!( + GrokCredentials::parse_for_kind(auth, GrokAuthKind::Cli) + .unwrap() + .access_token, + "cli-token" + ); + assert_eq!( + GrokCredentials::parse_for_kind(auth, GrokAuthKind::OAuth) + .unwrap() + .access_token, + "oauth-token" + ); +} +#[test] +fn cookie_refresh_uses_cache_when_present() { + assert_eq!( + cookie_refresh_action(true, None), + CookieRefreshAction::UseCached + ); +} + +#[test] +fn cookie_refresh_reimports_on_auth_failure() { + assert_eq!( + cookie_refresh_action(true, Some(&ProviderError::AuthRequired)), + CookieRefreshAction::ReimportBrowser + ); + assert_eq!( + cookie_refresh_action(false, None), + CookieRefreshAction::ReimportBrowser + ); +} + +#[test] +fn cookie_refresh_gives_up_on_non_auth_errors() { + assert_eq!( + cookie_refresh_action(true, Some(&ProviderError::Other("network down".into()))), + CookieRefreshAction::GiveUp + ); +} + +#[test] +fn is_cookie_auth_failure_only_auth_required() { + assert!(is_cookie_authentication_failure( + &ProviderError::AuthRequired + )); + assert!(!is_cookie_authentication_failure(&ProviderError::NoCookies)); +} + +#[test] +fn cookie_billing_stays_siloed_from_auth_file_identity() { + let result = result_from_cookie_billing(GrokBillingSnapshot { + used_percent: Some(23.0), + used_percent_is_wire_published: true, + used_percent_is_implicit_zero: false, + resets_at: None, + window_minutes: None, + }); + assert_eq!(result.source_label, "grok-browser"); + assert!(result.usage.account_email.is_none()); + assert!(result.usage.account_organization.is_none()); + assert!(result.usage.login_method.is_none()); +} +#[test] +fn billing_snapshot_uses_full_weekly_cycle_for_pace() { + let now = Utc::now(); + let resets = now + chrono::Duration::days(2); + let result = result_from_billing( + GrokBillingSnapshot { + used_percent: Some(12.0), + used_percent_is_wire_published: true, + used_percent_is_implicit_zero: false, + resets_at: Some(resets), + window_minutes: Some(crate::core::WEEKLY_WINDOW_MINUTES), + }, + "web", + None, + None, + Some("SuperGrok".into()), + ); + assert_eq!( + result.usage.primary.window_minutes, + Some(crate::core::WEEKLY_WINDOW_MINUTES) + ); + assert_eq!(result.usage.primary_label.as_deref(), Some("Weekly")); + let pace = crate::core::UsagePace::weekly( + &result.usage.primary, + Some(now), + crate::core::WEEKLY_WINDOW_MINUTES, + ); + assert!(pace.is_some(), "weekly window + reset must yield pace"); +} + +#[test] +fn monthly_cycle_stays_monthly_with_six_days_remaining() { + let now = Utc::now(); + let resets = now + chrono::Duration::days(6); + let monthly_minutes = 31 * 24 * 60; + let result = result_from_billing( + GrokBillingSnapshot { + used_percent: Some(40.0), + used_percent_is_wire_published: true, + used_percent_is_implicit_zero: false, + resets_at: Some(resets), + window_minutes: Some(monthly_minutes), + }, + "cli", + None, + None, + Some("SuperGrok Heavy".into()), + ); + assert_eq!(result.usage.primary_label.as_deref(), Some("Monthly")); + assert_eq!(result.usage.primary.window_minutes, Some(monthly_minutes)); +} + +#[test] +fn reset_distance_alone_does_not_invent_a_cadence() { + let resets = Utc::now() + chrono::Duration::days(6); + let result = result_from_billing( + GrokBillingSnapshot { + used_percent: Some(80.0), + used_percent_is_wire_published: true, + used_percent_is_implicit_zero: false, + resets_at: Some(resets), + window_minutes: None, + }, + "web", + None, + None, + Some("SuperGrok".into()), + ); + assert_eq!(result.usage.primary.window_minutes, None); + assert_eq!(result.usage.primary_label, None); +} + +#[test] +fn period_only_billing_is_informational_not_zero_usage() { + let resets = Utc::now() + chrono::Duration::days(6); + let result = result_from_billing( + GrokBillingSnapshot { + used_percent: None, + used_percent_is_wire_published: false, + used_percent_is_implicit_zero: false, + resets_at: Some(resets), + window_minutes: None, + }, + "cli", + Some("user@example.com".into()), + None, + Some("SuperGrok Heavy".into()), + ); + + assert!(result.usage.primary.is_informational); + assert_eq!(result.usage.primary.resets_at, Some(resets)); + assert_eq!( + result.usage.account_email.as_deref(), + Some("user@example.com") + ); + assert_eq!( + result.usage.login_method.as_deref(), + Some("SuperGrok Heavy") + ); +} + +#[test] +fn unpublished_zero_does_not_reach_the_usage_surface() { + let result = result_from_billing( + GrokBillingSnapshot { + used_percent: Some(0.0), + used_percent_is_wire_published: false, + used_percent_is_implicit_zero: false, + resets_at: Some(Utc.timestamp_opt(1_789_000_000, 0).single().unwrap()), + window_minutes: None, + }, + "grok-web", + None, + None, + None, + ); + + assert!(result.usage.primary.is_informational); +} diff --git a/rust/src/providers/ollama/mod.rs b/rust/src/providers/ollama/mod.rs index 3862b3456f..bd33693ec6 100755 --- a/rust/src/providers/ollama/mod.rs +++ b/rust/src/providers/ollama/mod.rs @@ -23,6 +23,8 @@ use crate::settings::ApiKeys; const OLLAMA_SETTINGS_URL: &str = "https://ollama.com/settings"; const OLLAMA_TAGS_URL: &str = "https://ollama.com/api/tags"; const OLLAMA_VALIDATION_URL: &str = "https://ollama.com/api/web_search"; +const OLLAMA_MONTHLY_WINDOW_MINUTES: u32 = 30 * 24 * 60; +const OLLAMA_MONTHLY_USAGE_LABEL: &str = "Monthly usage"; /// Ollama provider pub struct OllamaProvider { @@ -233,23 +235,34 @@ impl OllamaProvider { // Check if we're signed out if html.contains("Sign in") && !html.contains("Cloud Usage") + && !html.contains("Included usage") + && !html.contains(OLLAMA_MONTHLY_USAGE_LABEL) && !html.contains("Session usage") { return Err(ProviderError::AuthRequired); } + let monthly_block = self.parse_usage_block( + &[OLLAMA_MONTHLY_USAGE_LABEL], + html, + Some(OLLAMA_MONTHLY_WINDOW_MINUTES), + ); let session_block = self.parse_usage_block(&["Session usage", "Hourly usage"], html, Some(5 * 60)); let weekly_block = self.parse_usage_block(&["Weekly usage"], html, Some(7 * 24 * 60)); - if session_block.is_none() && weekly_block.is_none() { + if monthly_block.is_none() && session_block.is_none() && weekly_block.is_none() { return Err(ProviderError::Parse( "Could not find usage data on Ollama settings page".to_string(), )); } - let primary = rate_window_from_usage_block(session_block.as_ref()); + let primary = + rate_window_from_usage_block(monthly_block.as_ref().or(session_block.as_ref())); let mut usage = UsageSnapshot::new(primary); + if monthly_block.is_some() { + usage = usage.with_primary_label("Monthly"); + } // Parse plan name if let Some(plan) = self.parse_plan_name(html) { @@ -294,6 +307,15 @@ impl OllamaProvider { }); } + if let Some(val) = parse_dollar_used_percent(window) { + return Some(UsageBlock { + used_percent: val, + window_minutes, + resets_at: parse_first_datetime(window), + reset_description: parse_reset_description(window), + }); + } + // Try "width: XX%" pattern (progress bar CSS) let width_re = Regex::new(r"width:\s*(\d+(?:\.\d+)?)%").ok()?; if let Some(caps) = width_re.captures(window) @@ -311,12 +333,23 @@ impl OllamaProvider { None } - /// Parse plan name from "Cloud Usage" section + /// Parse plan name from the current "Included usage" or legacy "Cloud Usage" section. fn parse_plan_name(&self, html: &str) -> Option { - let re = Regex::new(r#"Cloud Usage\s*\s*]*>([^<]+)"#).ok()?; - re.captures(html) - .and_then(|caps| caps.get(1)) - .map(|m| m.as_str().trim().to_string()) + for pattern in [ + r#"Included usage\s*\s*]*>([^<]+)\s*]*>([^<]+)"#, + ] { + let re = Regex::new(pattern).ok()?; + if let Some(plan) = re + .captures(html) + .and_then(|caps| caps.get(1)) + .map(|m| m.as_str().trim().to_string()) + .filter(|value| !value.is_empty()) + { + return Some(plan); + } + } + None } /// Parse account email from the page @@ -392,13 +425,33 @@ fn clean_secret(raw: Option<&str>) -> Option { } fn usage_block_end(tail: &str, current_label: &str) -> Option { - ["Session usage", "Hourly usage", "Weekly usage"] - .iter() - .filter(|label| **label != current_label) - .filter_map(|label| tail.get(current_label.len()..)?.find(label)) - .map(|idx| idx + current_label.len()) - .min() - .map(|idx| idx.min(4000)) + [ + OLLAMA_MONTHLY_USAGE_LABEL, + "Session usage", + "Hourly usage", + "Weekly usage", + ] + .iter() + .filter(|label| **label != current_label) + .filter_map(|label| tail.get(current_label.len()..)?.find(label)) + .map(|idx| idx + current_label.len()) + .min() + .map(|idx| idx.min(4000)) +} + +/// Convert included dollar credits to quota utilization, not a spend estimate. +fn parse_dollar_used_percent(text: &str) -> Option { + let amount = r"([0-9]{1,3}(?:,[0-9]{3})*(?:\.[0-9]+)?|[0-9]+(?:\.[0-9]+)?)"; + let pattern = format!(r"(?i)\$\s*{amount}\s+of\s+\$\s*{amount}\s+used"); + let re = Regex::new(&pattern).ok()?; + let caps = re.captures(text)?; + let used = caps.get(1)?.as_str().replace(',', "").parse::().ok()?; + let limit = caps.get(2)?.as_str().replace(',', "").parse::().ok()?; + if !used.is_finite() || !limit.is_finite() || limit <= 0.0 { + return None; + } + let percent = used / limit * 100.0; + percent.is_finite().then_some(percent) } fn rate_window_from_usage_block(block: Option<&UsageBlock>) -> RateWindow { @@ -703,6 +756,50 @@ mod tests { ); } + #[test] + fn preserves_legacy_ollama_session_and_weekly_payloads() { + let provider = OllamaProvider::new(); + let snapshot = provider + .parse_usage_html( + r#" + Cloud UsageFree +
Session usage 42% used
+
Weekly usage 84% used
+ "#, + ) + .unwrap(); + + assert_eq!(snapshot.primary.used_percent, 42.0); + assert_eq!(snapshot.primary.window_minutes, Some(5 * 60)); + assert_eq!(snapshot.primary_label, None); + assert_eq!(snapshot.secondary.unwrap().used_percent, 84.0); + assert_eq!(snapshot.login_method.as_deref(), Some("Free")); + } + + #[test] + fn parses_monthly_included_credits_and_keeps_weekly_window() { + let provider = OllamaProvider::new(); + let snapshot = provider + .parse_usage_html( + r#" + Included usagePro +
Monthly usage $7.50 of $60 used +
+
Weekly usage 25% used
+ "#, + ) + .unwrap(); + + assert_eq!(snapshot.primary.used_percent, 12.5); + assert_eq!( + snapshot.primary.window_minutes, + Some(OLLAMA_MONTHLY_WINDOW_MINUTES) + ); + assert_eq!(snapshot.primary_label.as_deref(), Some("Monthly")); + assert_eq!(snapshot.secondary.unwrap().used_percent, 25.0); + assert_eq!(snapshot.login_method.as_deref(), Some("Pro")); + } + #[test] fn parses_ollama_usage_blocks_with_window_bounds() { let provider = OllamaProvider::new(); diff --git a/rust/src/providers/poe/mod.rs b/rust/src/providers/poe/mod.rs index 5be5a9006b..d4548ced6f 100644 --- a/rust/src/providers/poe/mod.rs +++ b/rust/src/providers/poe/mod.rs @@ -1,14 +1,31 @@ use async_trait::async_trait; +use chrono::{DateTime, Duration, Utc}; use reqwest::Client; use serde_json::Value; +use std::collections::BTreeMap; use crate::core::{ - FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, UsageSnapshot, + CostDailyPoint, CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, + ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, }; const POE_BALANCE_URL: &str = "https://api.poe.com/usage/current_balance"; +const POE_HISTORY_URL: &str = "https://api.poe.com/usage/points_history"; const CREDENTIAL_TARGET: &str = "codexbar-poe"; +const POE_HISTORY_PAGE_LIMIT: usize = 100; +const POE_HISTORY_MAX_PAGES: usize = 5; + +#[derive(Debug, Clone, PartialEq)] +struct PoeHistoryEntry { + date: DateTime, + points: f64, +} + +#[derive(Debug, Clone, PartialEq)] +struct PoeHistorySummary { + total_points: f64, + daily: Vec, +} pub struct PoeProvider { metadata: ProviderMetadata, @@ -62,10 +79,12 @@ impl Provider for PoeProvider { CREDENTIAL_TARGET, &["POE_API_KEY"], )?; + // Capture one refresh instant for the complete balance/history snapshot. + let refresh_now = Utc::now(); let response = self .client .get(POE_BALANCE_URL) - .bearer_auth(key) + .bearer_auth(&key) .header("Accept", "application/json") .send() .await?; @@ -83,10 +102,24 @@ impl Provider for PoeProvider { let value: Value = response.json().await.map_err(|e| { ProviderError::Parse(format!("Failed to parse Poe balance: {e}")) })?; - Ok(ProviderFetchResult::new( - snapshot_from_balance(&value), - "api", - )) + let balance = first_number( + &value, + &[ + "current_point_balance", + "currentPointBalance", + "balance", + "points", + ], + ); + let mut result = + ProviderFetchResult::new(snapshot_from_balance_at(&value, refresh_now), "api"); + if ctx.include_credits + && let Ok(entries) = fetch_points_history(&self.client, &key, refresh_now).await + && let Some(cost) = cost_snapshot_from_history(&entries, balance, refresh_now) + { + result = result.with_cost(cost); + } + Ok(result) } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -99,7 +132,67 @@ impl Provider for PoeProvider { } } -fn snapshot_from_balance(value: &Value) -> UsageSnapshot { +async fn fetch_points_history( + client: &Client, + key: &str, + refresh_now: DateTime, +) -> Result, ProviderError> { + let cutoff = refresh_now - Duration::days(30); + let mut cursor = None; + let mut entries = Vec::new(); + + for _ in 0..POE_HISTORY_MAX_PAGES { + let mut url = reqwest::Url::parse(POE_HISTORY_URL) + .map_err(|e| ProviderError::Other(e.to_string()))?; + { + let mut query = url.query_pairs_mut(); + query.append_pair("limit", &POE_HISTORY_PAGE_LIMIT.to_string()); + if let Some(cursor) = cursor.as_deref() { + query.append_pair("starting_after", cursor); + } + } + + let response = client + .get(url) + .bearer_auth(key) + .header("Accept", "application/json") + .send() + .await?; + if !response.status().is_success() { + return Err(ProviderError::Other(format!( + "Poe history returned status {}", + response.status() + ))); + } + let root: Value = response + .json() + .await + .map_err(|e| ProviderError::Parse(format!("Failed to parse Poe history: {e}")))?; + let rows = history_rows(&root); + for row in &rows { + if let Some(entry) = parse_history_entry(row) + && entry.date >= cutoff + { + entries.push(entry); + } + } + + let last_date = rows + .iter() + .rev() + .find_map(|row| history_timestamp(row).and_then(parse_entry_date)); + if last_date.is_some_and(|date| date < cutoff) { + break; + } + cursor = history_cursor(&root, &rows); + if cursor.is_none() { + break; + } + } + Ok(entries) +} + +fn snapshot_from_balance_at(value: &Value, updated_at: DateTime) -> UsageSnapshot { let balance = first_number( value, &[ @@ -115,30 +208,223 @@ fn snapshot_from_balance(value: &Value) -> UsageSnapshot { .reset_description .clone() .unwrap_or_else(|| "Poe API".into()); - UsageSnapshot::new(primary).with_login_method(label) + let mut snapshot = UsageSnapshot::new(primary).with_login_method(label); + snapshot.updated_at = updated_at; + snapshot +} + +fn history_rows(root: &Value) -> Vec<&Value> { + ["data", "items", "results"] + .iter() + .find_map(|key| root.get(*key).and_then(Value::as_array)) + .map(|rows| rows.iter().collect()) + .unwrap_or_default() +} + +fn history_timestamp(row: &Value) -> Option<&Value> { + ["creation_time", "timestamp", "created_at"] + .iter() + .find_map(|key| row.get(*key)) +} + +fn parse_entry_date(value: &Value) -> Option> { + match value { + Value::Number(number) => number.as_f64().and_then(timestamp_to_date), + Value::String(raw) => { + let raw = raw.trim(); + if raw.is_empty() { + return None; + } + raw.parse::() + .ok() + .and_then(timestamp_to_date) + .or_else(|| { + DateTime::parse_from_rfc3339(raw) + .ok() + .map(|date| date.with_timezone(&Utc)) + }) + } + _ => None, + } +} + +fn timestamp_to_date(value: f64) -> Option> { + if !value.is_finite() { + return None; + } + let millis = if value > 100_000_000_000_000.0 { + value / 1_000.0 + } else if value > 1_000_000_000_000.0 { + value + } else { + value * 1_000.0 + }; + if !millis.is_finite() || millis < i64::MIN as f64 || millis > i64::MAX as f64 { + return None; + } + #[allow( + clippy::cast_possible_truncation, + reason = "millis is finite and bounds-checked immediately above" + )] + DateTime::from_timestamp_millis(millis.round() as i64) +} + +fn parse_history_entry(row: &Value) -> Option { + let date = history_timestamp(row).and_then(parse_entry_date)?; + let points = direct_number(row, &["cost_points", "points", "point_cost"])?; + points.is_finite().then_some(PoeHistoryEntry { + date, + points: points.max(0.0), + }) +} + +fn history_cursor(root: &Value, rows: &[&Value]) -> Option { + root.get("next_cursor") + .and_then(Value::as_str) + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .map(str::to_owned) + .or_else(|| { + if root.get("has_more").and_then(Value::as_bool) == Some(true) { + rows.last() + .and_then(|row| row.get("query_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .map(str::to_owned) + } else { + None + } + }) +} + +fn summarize_history(entries: &[PoeHistoryEntry], refresh_now: DateTime) -> PoeHistorySummary { + let cutoff = refresh_now - Duration::days(30); + let mut daily = BTreeMap::::new(); + let mut total_points = 0.0; + + for entry in entries.iter().filter(|entry| entry.date >= cutoff) { + let day = entry.date.format("%Y-%m-%d").to_string(); + *daily.entry(day.clone()).or_default() += entry.points; + total_points += entry.points; + } + + PoeHistorySummary { + total_points, + daily: daily + .into_iter() + .map(|(day, amount)| CostDailyPoint { day, amount }) + .collect(), + } +} + +fn cost_snapshot_from_history( + entries: &[PoeHistoryEntry], + balance: Option, + refresh_now: DateTime, +) -> Option { + let summary = summarize_history(entries, refresh_now); + if summary.daily.is_empty() { + return None; + } + let mut cost = + CostSnapshot::new(summary.total_points, "points", "Last 30 days").with_daily(summary.daily); + if let Some(balance) = balance { + cost = cost.with_balance(balance); + } + cost.updated_at = refresh_now; + Some(cost) } fn first_number(value: &Value, keys: &[&str]) -> Option { match value { Value::Object(map) => keys .iter() - .find_map(|key| map.get(*key).and_then(Value::as_f64)) - .or_else(|| map.values().find_map(|v| first_number(v, keys))), - Value::Array(items) => items.iter().find_map(|v| first_number(v, keys)), + .find_map(|key| map.get(*key).and_then(value_as_number)) + .or_else(|| map.values().find_map(|value| first_number(value, keys))), + Value::Array(items) => items.iter().find_map(|value| first_number(value, keys)), _ => None, } } +fn direct_number(value: &Value, keys: &[&str]) -> Option { + keys.iter() + .find_map(|key| value.get(*key).and_then(value_as_number)) +} + +fn value_as_number(value: &Value) -> Option { + value + .as_f64() + .or_else(|| value.as_str()?.trim().parse::().ok()) + .filter(|value| value.is_finite()) +} + #[cfg(test)] mod tests { use super::*; + use chrono::TimeZone; #[test] fn parses_balance_label() { - let snapshot = snapshot_from_balance(&serde_json::json!({"current_point_balance": 1234})); + let snapshot = snapshot_from_balance_at( + &serde_json::json!({"current_point_balance": 1234}), + Utc.with_ymd_and_hms(2026, 9, 1, 0, 0, 0).unwrap(), + ); assert_eq!( snapshot.login_method.as_deref(), Some("Balance: 1234 points") ); } + + #[test] + fn history_uses_one_refresh_clock_for_midnight_bucket_and_retention() { + let refresh_now = Utc.with_ymd_and_hms(2026, 9, 1, 0, 0, 30).unwrap(); + let entries = vec![ + PoeHistoryEntry { + date: Utc.with_ymd_and_hms(2026, 8, 31, 23, 59, 59).unwrap(), + points: 5.0, + }, + PoeHistoryEntry { + date: Utc.with_ymd_and_hms(2026, 9, 1, 0, 0, 0).unwrap(), + points: 3.0, + }, + PoeHistoryEntry { + date: Utc.with_ymd_and_hms(2026, 8, 2, 0, 0, 29).unwrap(), + points: 99.0, + }, + ]; + + let summary = summarize_history(&entries, refresh_now); + let balance_snapshot = snapshot_from_balance_at( + &serde_json::json!({"current_point_balance": 321}), + refresh_now, + ); + let cost = cost_snapshot_from_history(&entries, Some(321.0), refresh_now).unwrap(); + + assert_eq!(balance_snapshot.updated_at, refresh_now); + assert_eq!(cost.updated_at, refresh_now); + assert_eq!(summary.total_points, 8.0); + let today = refresh_now.format("%Y-%m-%d").to_string(); + assert_eq!( + summary + .daily + .iter() + .find(|point| point.day == today) + .map(|point| point.amount), + Some(3.0) + ); + assert_eq!( + summary.daily, + vec![ + CostDailyPoint { + day: "2026-08-31".to_string(), + amount: 5.0, + }, + CostDailyPoint { + day: "2026-09-01".to_string(), + amount: 3.0, + }, + ] + ); + } } diff --git a/rust/src/tray/render.rs b/rust/src/tray/render.rs index 20f127bec8..5b486a5b6b 100644 --- a/rust/src/tray/render.rs +++ b/rust/src/tray/render.rs @@ -219,6 +219,34 @@ mod tests { assert_eq!(u32::try_from(rgba.len()).unwrap(), w * h * 4); } + #[test] + fn single_quota_uses_one_centered_prominent_meter() { + let (single, _, _) = render_bar_icon_rgba(50.0, None, false); + let (multiple, _, _) = render_bar_icon_rgba(50.0, Some(25.0), false); + + let pixel = |rgba: &[u8], x: u32, y: u32| { + let index = ((y * TRAY_ICON_SIZE + x) * 4) as usize; + [ + rgba[index], + rgba[index + 1], + rgba[index + 2], + rgba[index + 3], + ] + }; + + // The single-quota layout occupies one centered, thick lane. + assert_eq!(pixel(&single, 20, 9), [60, 60, 70, 255]); + assert_eq!(pixel(&single, 20, 10), [80, 80, 90, 255]); + assert_eq!(pixel(&single, 20, 21), [80, 80, 90, 255]); + assert_eq!(pixel(&single, 20, 22), [60, 60, 70, 255]); + + // Multiple quotas retain distinct upper and lower lanes. + assert_eq!(pixel(&multiple, 20, 8), [80, 80, 90, 255]); + assert_eq!(pixel(&multiple, 20, 15), [60, 60, 70, 255]); + assert_eq!(pixel(&multiple, 20, 18), [80, 80, 90, 255]); + assert_eq!(pixel(&multiple, 20, 23), [60, 60, 70, 255]); + } + #[test] fn zero_fill_gives_gray_only_bar() { let (rgba, w, _h) = render_bar_icon_rgba(0.0, None, false);