diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 043e2ec199..2edf24acd4 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -335,10 +335,12 @@ impl ProviderUsageSnapshot { .unwrap_or_else(|| metadata.session_label.to_string()), ), secondary: secondary_snap, - secondary_label: usage - .secondary - .as_ref() - .map(|_| metadata.weekly_label.to_string()), + secondary_label: usage.secondary.as_ref().map(|_| { + usage + .secondary_label + .clone() + .unwrap_or_else(|| metadata.weekly_label.to_string()) + }), model_specific: usage .model_specific .as_ref() 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 23384b5aae..28c74bab5b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -179,15 +179,15 @@ fn build_usage_spend_summary( let include_opencodex = settings.open_codex_usage_logs_enabled; let hide_native = settings.hide_native_codex_cost_when_open_codex_present; - let codex_cache = - codexbar::core::JsonlScanner::load_cache(codexbar::core::ProviderId::Codex, None); - let codex_stale = !codex_cache.days.is_empty() && codex_cache.previous_report.is_some(); + let codex_cache_status = + codexbar::core::JsonlScanner::load_cache_status(codexbar::core::ProviderId::Codex, None); + let codex_stale = codex_cache_status.has_days && codex_cache_status.previous_report.is_some(); let codex_stale_updated_at = codex_stale .then(|| { - codex_cache + codex_cache_status .previous_report .as_ref() - .and_then(|r| r.updated_at.clone()) + .and_then(|report| report.updated_at.clone()) }) .flatten(); @@ -351,13 +351,16 @@ fn build_usage_spend_summary( spend } "antigravity" => { + use codexbar::providers::antigravity::local_sessions::LocalHistoryCoverage; let seven = codexbar::providers::antigravity::local_sessions::summarize(7); let thirty = codexbar::providers::antigravity::local_sessions::summarize(30); let mut spend = cached_spend(cached_snapshot); - spend.seven_day_tokens = (seven.session_count > 0).then_some(seven.total_tokens); - spend.thirty_day_tokens = (thirty.session_count > 0).then_some(thirty.total_tokens); - if thirty.session_count > 0 { - spend.source = "local Antigravity sessions".to_string(); + spend.seven_day_tokens = matches!(seven.coverage, LocalHistoryCoverage::Complete) + .then_some(seven.total_tokens); + spend.thirty_day_tokens = matches!(thirty.coverage, LocalHistoryCoverage::Complete) + .then_some(thirty.total_tokens); + if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { + spend.source = "local Antigravity history".to_string(); } spend } diff --git a/apps/desktop-tauri/src/components/MenuCard.test.tsx b/apps/desktop-tauri/src/components/MenuCard.test.tsx index 10af70ace2..bed4d2ca6f 100644 --- a/apps/desktop-tauri/src/components/MenuCard.test.tsx +++ b/apps/desktop-tauri/src/components/MenuCard.test.tsx @@ -113,6 +113,7 @@ describe("MenuCard", () => { tauriMocks.getLocaleStrings.mockResolvedValue( buildBundle({ ActionCopyError: "Copy error", + ApiSpendTitle: "API spend", DetailPaceRunsOutIn: "Runs out in", PanelEstimatedFromLocalLogs: "Estimated from local logs", PanelLeftSuffix: "left", diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 2537d17f71..0e1ed9f738 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -546,7 +546,7 @@ export default function MenuCardDetails({
{provider.cost.alwaysVisible === true && (provider.cost.limit ?? 0) <= 0 - ? "API spend" + ? t("ApiSpendTitle") : provider.cost.balance != null && provider.cost.limit == null ? provider.cost.period || t("CreditsLabel") : `${t("DetailCostTitle")} — ${provider.cost.period}`} diff --git a/apps/desktop-tauri/src/i18n/keys.ts b/apps/desktop-tauri/src/i18n/keys.ts index ac2b06c784..2da4c36537 100644 --- a/apps/desktop-tauri/src/i18n/keys.ts +++ b/apps/desktop-tauri/src/i18n/keys.ts @@ -570,6 +570,7 @@ export const ALL_LOCALE_KEYS = [ "DetailPaceRunsOutIn", "DetailPaceWillLastToReset", "DetailCostTitle", + "ApiSpendTitle", "DetailCostUsed", "DetailCostLimit", "DetailCostRemaining", diff --git a/docs/adr/0006-dependency-consolidation-candidates.md b/docs/adr/0006-dependency-consolidation-candidates.md new file mode 100644 index 0000000000..df45f9f388 --- /dev/null +++ b/docs/adr/0006-dependency-consolidation-candidates.md @@ -0,0 +1,245 @@ +# ADR 0006: Dependency consolidation candidates + +Date: 2026-09-05 +Status: Proposed + +## Context + +Win-CodexBar has grown across roughly 70 providers, a Rust backend, a Tauri 2 +shell, and a React frontend. The codebase now contains several places where the +same domain information or infrastructure policy is expressed repeatedly. +Adding dependencies is only worthwhile when a crate removes measurable +boilerplate, reduces synchronization points, or materially improves +correctness. This ADR records candidates for review before implementation. It +does not approve adding every crate listed below. + +The September 2026 audit found these concrete pressure points: + +- `rust/src/core/provider.rs` is 1,131 lines. `ProviderId` has 70 variants, with + repeated handwritten tables for iteration, CLI names, display names, cookie + domains, aliases, and brand colors. The largest repeated metadata blocks span + roughly 540 lines. +- `apps/desktop-tauri/src/types/bridge.ts` is 958 lines with 101 exported + types/interfaces that must stay aligned with Rust DTOs. +- `apps/desktop-tauri/src/lib/tauri.ts` is 534 lines with 91 handwritten invoke + wrappers, while the Rust shell exposes 98 `#[tauri::command]` functions. +- `FetchContext` is manually constructed at 26 call sites. +- The Rust code contains about 20 `lock().unwrap*` sites, including eight + explicit poison-recovery patterns using `into_inner()`. +- There are many manual `sort`, `dedup`, `HashSet` uniqueness, grouping, and + collection reshaping paths across provider, session, cost, and dashboard + code. +- URL-like configuration and command inputs are commonly represented as raw + `String`/`&str`, including gateway, proxy, dashboard, status, release, and + provider endpoint URLs. `validate_external_url` currently performs manual + HTTP(S) prefix validation. +- Provider and configuration parsing has many generic Serde error mappings that + report the parser error but not the precise failing field path. +- Several in-memory caches hand-roll `Mutex>`, timestamps, and TTL + checks for provider, chart/local-usage, status, token, and pricing data. +- Cost and spend calculations are extensively represented as `f64`. + +## Decision + +Use a tiered adoption policy. A dependency is added only in the PR that also +uses it to remove or centralize existing code. Do not add dormant dependencies +for hypothetical future use. + +### Tier 1: highest-confidence candidates + +#### `strum` + +Use derives such as `EnumIter`, `EnumString`, `Display`, `VariantArray`, and +variant metadata where they eliminate repeated enum plumbing. + +Primary target: `ProviderId` and other enums that manually implement `all`, +`parse`, `as_str`, `display_name`, or alias tables. + +Acceptance criteria: + +- provider metadata has one obvious source of truth per property; +- aliases and compatibility names remain behaviorally identical; +- exhaustive provider factory behavior remains compile-time checked; +- net handwritten LOC decreases materially. + +#### `ts-rs` + +Generate TypeScript DTO definitions from Rust for bridge types that are true +wire contracts. Keep frontend-only unions/types handwritten when Rust is not +the source of truth. + +Primary target: the mirrored portions of `src/types/bridge.ts`. + +Acceptance criteria: + +- generated files are deterministic and checked in or generated by an explicit + reproducible step; +- Serde rename/tag behavior is covered by tests; +- the generated boundary removes manual Rust/TypeScript synchronization rather + than introducing a second schema layer; +- no runtime dependency is added to the frontend. + +Do not adopt `tauri-specta` in the same step unless its Tauri 2 release line is +independently proven stable and it replaces enough of `src/lib/tauri.ts` to +justify the additional integration surface. + +#### `serde_with` + +Use only where it replaces custom Serde adapters or repetitive conversion +logic, for example display/from-string bridges or optional/default handling. + +Acceptance criteria: each use deletes or simplifies existing serialization +code; avoid cosmetic attribute churn. + +#### `url` + +Promote URL values to `url::Url` at validation/security boundaries where the +value is semantically a URL. The crate is already present transitively through +`reqwest`, so making it direct should not materially enlarge the dependency +graph at the current lockfile state. + +Primary targets: external URL validation and provider/gateway/proxy endpoints. + +Acceptance criteria: + +- scheme allowlists stay explicit; +- file/javascript and other unsafe schemes remain rejected where applicable; +- persisted wire compatibility is preserved when settings currently store + strings. + +#### `parking_lot` + +Use for internal synchronization where standard-library poisoning semantics are +not part of the intended behavior and where it removes repeated poison recovery. + +Acceptance criteria: + +- no lock is held across `.await`; +- migration is focused rather than repo-wide mechanical churn; +- test-only locks do not justify the dependency by themselves. + +#### `itertools` + +Use for collection transforms that become materially clearer, especially +unique/sorted/grouped pipelines that currently require temporary sets or +multi-step mutation. + +Acceptance criteria: the replacement is shorter and easier to read than the +standard-library version; simple iterator chains remain standard library only. + +#### `bon` + +Prototype on `FetchContext`, which has 26 literal construction sites, before +using it elsewhere. + +Acceptance criteria: + +- call sites communicate only intentional overrides; +- defaults remain centralized; +- builder generation does not obscure required fields or make diagnostics + worse; +- measurable net LOC/readability improvement is shown in the prototype diff. + +#### `serde_path_to_error` + +Use at user-facing or provider-facing JSON boundaries where a precise field path +would materially improve diagnostics. + +Acceptance criteria: preserve the existing provider-specific context while +adding the failing data path; do not wrap silent best-effort parsers whose +failure is intentionally ignored. + +### Tier 2: architecture-specific candidates + +These need a focused prototype and should not be added merely because the crate +is useful in general. + +#### `tokio-util` + +Candidate for `CancellationToken` in refresh/probe/background-task ownership. +It is already present transitively through the current HTTP stack. Adopt only +where it replaces bespoke cancellation flags/channels and clarifies task +lifetime. + +#### `secrecy` + +Candidate for API keys, tokens, and other credentials that currently travel as +ordinary `String`s. Adoption must integrate with existing DPAPI, keyring, +redaction, and persistence behavior rather than creating a parallel secret +model. + +#### `moka` + +Candidate for simple in-memory TTL caches currently implemented as +`Mutex>` plus timestamps. Do not replace the JSONL/disk cost cache +or other caches with domain-specific persistence/invalidation semantics unless a +prototype proves equivalence. + +#### `reqwest-middleware` and `reqwest-retry` + +Candidate for genuinely shared idempotent HTTP retry/backoff policy. The repo is +currently on `reqwest 0.12`, so compatible middleware versions must be verified +before adoption. Provider-specific authentication refresh, WAF fallback, +alternate-host logic, and semantic retries must remain provider-owned. + +#### `rust_decimal` + +Candidate for monetary totals if floating-point rounding is demonstrated to be +a product-level correctness problem. Migration would be broad because cost, +pricing, bridge DTOs, formatting, tests, and persisted/exported values currently +use `f64`. + +Do not migrate merely for theoretical precision. First add a failing test or +real rounding example that changes an observable result. + +#### `phf` + +Candidate for immutable lookup tables that are currently runtime-built +`LazyLock>` values. Prefer ordinary `match`, slices, or arrays when +those are already clearer and compile-time cheap. + +## Rejected default: dependency accumulation + +Do not add all candidates in one dependency-only PR. That would increase build, +supply-chain, maintenance, and audit surface while proving none of the claimed +benefits. + +The preferred sequence is small implementation PRs, each with before/after +evidence: + +1. `strum` for provider/enum metadata. +2. `ts-rs` for Rust-to-TypeScript wire types. +3. `url` for URL validation boundaries. +4. `bon` prototype for `FetchContext`. +5. `serde_path_to_error` for selected provider/config diagnostics. +6. `parking_lot` and `itertools` where the diff clearly pays for them. +7. Tier 2 only after a concrete problem demonstrates the need. + +`serde_with` may be folded into one of the above only when the same PR has a +clear serialization target. + +## Review questions + +Reviewers should challenge each candidate with the same questions: + +1. Which current code is deleted or centralized by this crate? +2. How many call sites or synchronization points disappear? +3. Does the dependency become part of a public/wire/persistence contract? +4. What new failure mode, compile-time cost, runtime cost, or supply-chain + surface does it add? +5. Could a small local helper or derive already in the graph solve the same + problem more simply? +6. Is the migration independently reversible? +7. What test proves behavior stayed identical? + +## Consequences + +- Dependency additions become evidence-driven rather than preference-driven. +- The highest-value work focuses on duplicated sources of truth first: + `ProviderId` metadata and the Rust/TypeScript bridge contract. +- Large architectural migrations are split from low-risk boilerplate removal. +- Reviewers can reject individual candidates without blocking the rest of the + strategy. +- This ADR is intentionally `Proposed`; accepting it does not automatically + authorize every Tier 1 or Tier 2 dependency. Each implementation PR still + needs its own measured justification and tests. diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 62bd63f61d..d701dfa6dc 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -576,7 +576,10 @@ fn append_usage_window_lines( append_secondary_window_line( lines, usage.secondary.as_ref(), - metadata.weekly_label, + usage + .secondary_label + .as_deref() + .unwrap_or(metadata.weekly_label), use_color, ); append_model_specific_line(lines, usage.model_specific.as_ref(), use_color); @@ -671,7 +674,10 @@ pub fn render_brief_text(provider: ProviderId, result: &ProviderFetchResult) -> if let Some(secondary) = &usage.secondary { parts.push(format!( "{} {}", - metadata.weekly_label, + usage + .secondary_label + .as_deref() + .unwrap_or(metadata.weekly_label), format_percent(secondary.used_percent) )); } @@ -810,6 +816,18 @@ mod tests { ); } + #[test] + fn secondary_label_override_is_shared_by_full_and_brief_renderers() { + let result = fetch_result( + UsageSnapshot::new(RateWindow::new(10.0)) + .with_secondary(RateWindow::new(20.0)) + .with_secondary_label("Weekly"), + ); + let full = render_text_with_status(ProviderId::Antigravity, &result, None, false); + let brief = render_brief_text(ProviderId::Antigravity, &result); + assert!(full.contains("Weekly:")); + assert!(brief.contains("Weekly 20%")); + } #[test] fn primary_label_override_is_shared_by_full_and_brief_renderers() { let result = diff --git a/rust/src/core/jsonl_scanner.rs b/rust/src/core/jsonl_scanner.rs index 5a1ed45dc9..56032b8d4d 100755 --- a/rust/src/core/jsonl_scanner.rs +++ b/rust/src/core/jsonl_scanner.rs @@ -17,6 +17,53 @@ use std::fs::{self, File}; use std::io::{BufRead, BufReader, Seek, SeekFrom}; use std::path::{Path, PathBuf}; +#[derive(Debug, Clone, Default)] +pub struct CachedCostReadStatus { + pub has_days: bool, + pub previous_report: Option, +} + +#[derive(Deserialize, Default)] +struct CachedCostReadStatusProjection { + #[serde( + default, + rename = "days", + deserialize_with = "deserialize_nonempty_object" + )] + has_days: bool, + #[serde(default)] + previous_report: Option, +} + +fn deserialize_nonempty_object<'de, D>(deserializer: D) -> Result +where + D: serde::Deserializer<'de>, +{ + use serde::de::{IgnoredAny, MapAccess, Visitor}; + + struct NonemptyObjectVisitor; + + impl<'de> Visitor<'de> for NonemptyObjectVisitor { + type Value = bool; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("a JSON object") + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut nonempty = false; + while map.next_entry::()?.is_some() { + nonempty = true; + } + Ok(nonempty) + } + } + + deserializer.deserialize_map(NonemptyObjectVisitor) +} /// Maximum retained Codex JSONL line size (upstream session-metadata bound). const CODEX_JSONL_MAX_LINE_BYTES: usize = 256 * 1024; @@ -1059,6 +1106,39 @@ impl JsonlScanner { CostUsageCache::default() } + /// Read only the cache metadata needed by presentation surfaces. + /// + /// v0.56.0 performance parity: skip raw per-file scanner state and day + /// payloads when callers only need stale/catch-up status. + pub fn load_cache_status( + provider: ProviderId, + cache_root: Option<&Path>, + ) -> CachedCostReadStatus { + let cache_path = Self::cache_path(provider, cache_root); + if crate::core::is_bounded_provider(provider) { + #[allow( + clippy::cast_possible_truncation, + reason = "bounded artifacts fit usize on any supported target" + )] + let file_bytes = crate::core::artifact_file_size(&cache_path) as usize; + if file_bytes > crate::core::CostUsageCacheBudget::MAX_LOAD_BYTES { + return CachedCostReadStatus::default(); + } + } + + let Ok(file) = File::open(cache_path) else { + return CachedCostReadStatus::default(); + }; + let Ok(projection) = + serde_json::from_reader::<_, CachedCostReadStatusProjection>(BufReader::new(file)) + else { + return CachedCostReadStatus::default(); + }; + CachedCostReadStatus { + has_days: projection.has_days, + previous_report: projection.previous_report, + } + } fn cached_cost_report_from_days(cache: &CostUsageCache) -> CachedCostReport { let mut total_cost_usd = 0.0; let mut input_tokens = 0_i32; @@ -1076,9 +1156,11 @@ impl JsonlScanner { cached_tokens = cached_tokens.saturating_add(cached); output_tokens = output_tokens.saturating_add(output); - if CostUsagePricing::is_codex_unattributed_model(model) - || !CostUsagePricing::counts_toward_codex_subscription(model) - { + if CostUsagePricing::is_codex_unattributed_model(model) { + partial = true; + continue; + } + if !CostUsagePricing::counts_toward_codex_subscription(model) { continue; } let priced = pricing_day diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 6a0a4a5ef6..f8894609eb 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -85,6 +85,11 @@ pub struct UsageSnapshot { #[serde(skip_serializing_if = "Option::is_none")] pub secondary: Option, + /// Provider-resolved label for the secondary rate window when metadata + /// describes a model family rather than this snapshot's cadence. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub secondary_label: Option, + /// Model-specific rate window (e.g., Opus quota for Claude) #[serde(skip_serializing_if = "Option::is_none")] pub model_specific: Option, @@ -120,6 +125,7 @@ impl UsageSnapshot { primary, primary_label: None, secondary: None, + secondary_label: None, model_specific: None, tertiary: None, extra_rate_windows: Vec::new(), @@ -142,6 +148,12 @@ impl UsageSnapshot { self } + /// Builder pattern: override the secondary window label for this snapshot. + pub fn with_secondary_label(mut self, label: impl Into) -> Self { + self.secondary_label = Some(label.into()); + self + } + /// Builder pattern: set model-specific window pub fn with_model_specific(mut self, model_specific: RateWindow) -> Self { self.model_specific = Some(model_specific); diff --git a/rust/src/locale.rs b/rust/src/locale.rs index 7191b9a46c..59a6ecf6b0 100644 --- a/rust/src/locale.rs +++ b/rust/src/locale.rs @@ -844,6 +844,7 @@ locale_keys! { DetailPaceRunsOutIn, DetailPaceWillLastToReset, DetailCostTitle, + ApiSpendTitle, DetailCostUsed, DetailCostLimit, DetailCostRemaining, diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index 68c125c3c1..d9a0f26041 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -544,6 +544,7 @@ DetailPaceFarBehind = Far behind DetailPaceRunsOutIn = Runs out in DetailPaceWillLastToReset = Will last to reset DetailCostTitle = Cost +ApiSpendTitle = API spend DetailCostUsed = Used DetailCostLimit = Limit DetailCostRemaining = Remaining diff --git a/rust/src/locale/es-MX.ftl b/rust/src/locale/es-MX.ftl index 5bc94d8d4b..52591aab01 100644 --- a/rust/src/locale/es-MX.ftl +++ b/rust/src/locale/es-MX.ftl @@ -505,6 +505,7 @@ DetailPaceFarBehind = Muy atrasado DetailPaceRunsOutIn = Se agota en DetailPaceWillLastToReset = Durará hasta el reinicio DetailCostTitle = Costo +ApiSpendTitle = Gasto de API DetailCostUsed = Usado DetailCostLimit = Límite DetailCostRemaining = Restante diff --git a/rust/src/locale/ja-JP.ftl b/rust/src/locale/ja-JP.ftl index 5b3a386812..9b0c9b671c 100644 --- a/rust/src/locale/ja-JP.ftl +++ b/rust/src/locale/ja-JP.ftl @@ -487,6 +487,7 @@ DetailPaceFarBehind = 大きく遅れ DetailPaceRunsOutIn = 残り DetailPaceWillLastToReset = リセットまで持ちます DetailCostTitle = コスト +ApiSpendTitle = API 費用 DetailCostUsed = 使用済み DetailCostLimit = 上限 DetailCostRemaining = 残り diff --git a/rust/src/locale/ko-KR.ftl b/rust/src/locale/ko-KR.ftl index c36aef1cf2..5d98586961 100644 --- a/rust/src/locale/ko-KR.ftl +++ b/rust/src/locale/ko-KR.ftl @@ -492,6 +492,7 @@ DetailPaceFarBehind = 매우 느림 DetailPaceRunsOutIn = 소진까지 DetailPaceWillLastToReset = 초기화까지 유지 예상 DetailCostTitle = 비용 +ApiSpendTitle = API 비용 DetailCostUsed = 사용량 DetailCostLimit = 한도 DetailCostRemaining = 남음 diff --git a/rust/src/locale/ru-RU.ftl b/rust/src/locale/ru-RU.ftl index d7d6788734..f619853bdc 100644 --- a/rust/src/locale/ru-RU.ftl +++ b/rust/src/locale/ru-RU.ftl @@ -471,6 +471,7 @@ DetailPaceFarBehind = Далеко позади DetailPaceRunsOutIn = заканчивается в DetailPaceWillLastToReset = Продлится сброс DetailCostTitle = Стоимость +ApiSpendTitle = Расходы на API DetailCostUsed = Б/у DetailCostLimit = Лимит DetailCostRemaining = Осталось diff --git a/rust/src/locale/tr-TR.ftl b/rust/src/locale/tr-TR.ftl index 66eb354368..5a7d60b2ff 100644 --- a/rust/src/locale/tr-TR.ftl +++ b/rust/src/locale/tr-TR.ftl @@ -188,7 +188,7 @@ ActionRefresh = Yenile ActionSwitchAccount = Hesap değiştir... ActionUsageDashboard = Kullanım panosu ActionStatusPage = Durum sayfası -ActionCopyError = Kopyalama hatası +ActionCopyError = Hatayı kopyala ActionBuyCredits = Kredi satın al... PaceOnTrack = Planlandığı gibi PaceBehind = Geride @@ -508,6 +508,7 @@ DetailPaceFarBehind = Çok geride DetailPaceRunsOutIn = Tükenmesine kalan DetailPaceWillLastToReset = Sıfırlamaya kadar yeter DetailCostTitle = Maliyet +ApiSpendTitle = API harcaması DetailCostUsed = Kullanılan DetailCostLimit = Sınır DetailCostRemaining = Kalan diff --git a/rust/src/locale/zh-CN.ftl b/rust/src/locale/zh-CN.ftl index b447c5afa4..34b05f4a1b 100644 --- a/rust/src/locale/zh-CN.ftl +++ b/rust/src/locale/zh-CN.ftl @@ -486,6 +486,7 @@ DetailPaceFarBehind = 远落后 DetailPaceRunsOutIn = 预计耗尽时间 DetailPaceWillLastToReset = 足以支撑到重置 DetailCostTitle = 费用 +ApiSpendTitle = API 花费 DetailCostUsed = 已用 DetailCostLimit = 限额 DetailCostRemaining = 剩余 diff --git a/rust/src/locale/zh-TW.ftl b/rust/src/locale/zh-TW.ftl index 4353a2fc1b..a14dcf9920 100644 --- a/rust/src/locale/zh-TW.ftl +++ b/rust/src/locale/zh-TW.ftl @@ -486,6 +486,7 @@ DetailPaceFarBehind = 遠落後 DetailPaceRunsOutIn = 預計耗盡時間 DetailPaceWillLastToReset = 足以支撐到重置 DetailCostTitle = 費用 +ApiSpendTitle = API 花費 DetailCostUsed = 已用 DetailCostLimit = 限額 DetailCostRemaining = 剩餘 diff --git a/rust/src/providers/alibabatokenplan/mod.rs b/rust/src/providers/alibabatokenplan/mod.rs index b9158c7640..1a96decbd0 100644 --- a/rust/src/providers/alibabatokenplan/mod.rs +++ b/rust/src/providers/alibabatokenplan/mod.rs @@ -375,6 +375,13 @@ impl Provider for AlibabaTokenPlanProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { + SourceMode::Auto if ctx.auto_prefer_web => match self.fetch_via_web(ctx).await { + Ok(usage) => Ok(ProviderFetchResult::new(usage, "web")), + Err(_) => { + let usage = self.fetch_via_cli(ctx).await?; + Ok(ProviderFetchResult::new(usage, "cli")) + } + }, SourceMode::Auto => match self.fetch_via_cli(ctx).await { Ok(usage) => Ok(ProviderFetchResult::new(usage, "cli")), Err(_) => { @@ -398,6 +405,15 @@ impl Provider for AlibabaTokenPlanProvider { vec![SourceMode::Auto, SourceMode::Cli, SourceMode::Web] } + fn error_state_kind(&self, error: &ProviderError) -> crate::core::ProviderStateKind { + match error { + ProviderError::NotInstalled(message) if message.contains("Bailian CLI 'bl'") => { + crate::core::ProviderStateKind::LocalRuntimeOffline + } + _ => error.state_kind(), + } + } + fn supports_web(&self) -> bool { true } @@ -1006,6 +1022,17 @@ fn payload_diagnostics(value: &Value) -> String { mod tests { use super::*; + #[test] + fn missing_bailian_cli_is_local_runtime_offline() { + let provider = AlibabaTokenPlanProvider::new(); + let error = ProviderError::NotInstalled( + "Bailian CLI 'bl' is not installed or not on PATH.".to_string(), + ); + assert_eq!( + provider.error_state_kind(&error), + crate::core::ProviderStateKind::LocalRuntimeOffline + ); + } #[test] fn parses_token_plan_instance_payload() { let payload = serde_json::json!({ diff --git a/rust/src/providers/antigravity/local_proto.rs b/rust/src/providers/antigravity/local_proto.rs new file mode 100644 index 0000000000..6fa92e4df6 --- /dev/null +++ b/rust/src/providers/antigravity/local_proto.rs @@ -0,0 +1,267 @@ +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(super) struct ParsedUsage { + pub system_prompt: u64, + pub new_input: u64, + pub cache_read: u64, + pub output: u64, + pub reasoning: u64, + pub response_id: Option, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub(super) struct ParsedTurn { + pub usage: Option, + pub timestamp_ms: Option, + pub model: Option, + pub label: Option, +} + +#[derive(Clone, Copy)] +enum FieldValue<'a> { + Varint(u64), + Bytes(&'a [u8]), +} +#[derive(Clone, Copy)] +struct Field<'a> { + number: u32, + value: FieldValue<'a>, +} +struct Reader<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Reader<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + fn read_varint(&mut self) -> Option { + let mut result = 0_u64; + for index in 0..10 { + let byte = *self.bytes.get(self.offset)?; + self.offset += 1; + if index == 9 && byte > 1 { + return None; + } + result |= u64::from(byte & 0x7f) << (index * 7); + if byte & 0x80 == 0 { + return Some(result); + } + } + None + } + fn next_field(&mut self) -> Option> { + if self.offset >= self.bytes.len() { + return None; + } + let tag = self.read_varint()?; + let number = u32::try_from(tag >> 3).ok()?; + if number == 0 || number > 536_870_911 { + return None; + } + match tag & 7 { + 0 => Some(Field { + number, + value: FieldValue::Varint(self.read_varint()?), + }), + 1 | 5 => { + let width = if tag & 7 == 1 { 8 } else { 4 }; + let end = self.offset.checked_add(width)?; + let data = self.bytes.get(self.offset..end)?; + self.offset = end; + Some(Field { + number, + value: FieldValue::Bytes(data), + }) + } + 2 => { + let count = usize::try_from(self.read_varint()?).ok()?; + let end = self.offset.checked_add(count)?; + let data = self.bytes.get(self.offset..end)?; + self.offset = end; + Some(Field { + number, + value: FieldValue::Bytes(data), + }) + } + _ => None, + } + } +} + +fn fields(bytes: &[u8], mut visit: impl FnMut(Field<'_>) -> Option<()>) -> Option<()> { + let mut reader = Reader::new(bytes); + while reader.offset < bytes.len() { + visit(reader.next_field()?)?; + } + Some(()) +} +fn message(field: Field<'_>) -> Option<&[u8]> { + match field.value { + FieldValue::Bytes(v) => Some(v), + _ => None, + } +} +fn integer(field: Field<'_>) -> Option { + match field.value { + FieldValue::Varint(v) => Some(v), + _ => None, + } +} +fn text(field: Field<'_>) -> Option> { + let value = std::str::from_utf8(message(field)?).ok()?.trim(); + Some((!value.is_empty()).then(|| value.to_string())) +} + +pub(super) fn parse_turn(root: &[u8]) -> Option { + let mut turn = ParsedTurn::default(); + let mut seconds = None; + let mut nanos = 0_u64; + let mut found_chat = false; + fields(root, |field| { + if field.number != 1 { + return Some(()); + } + found_chat = true; + parse_chat(message(field)?, &mut turn, &mut seconds, &mut nanos) + })?; + if !found_chat { + return None; + } + turn.timestamp_ms = match seconds { + Some(value) if value > 0 && value <= 253_402_300_799 && nanos <= 999_999_999 => { + let seconds = i64::try_from(value).ok()?; + let nanos = i64::try_from(nanos).ok()?; + seconds.checked_mul(1000)?.checked_add(nanos / 1_000_000) + } + Some(_) => return None, + None => None, + }; + Some(turn) +} + +fn parse_chat( + bytes: &[u8], + turn: &mut ParsedTurn, + seconds: &mut Option, + nanos: &mut u64, +) -> Option<()> { + fields(bytes, |field| { + match field.number { + 4 => { + let mut usage = turn.usage.take().unwrap_or_default(); + parse_usage(message(field)?, &mut usage)?; + turn.usage = Some(usage); + } + 9 => parse_generation(message(field)?, seconds, nanos)?, + 19 => turn.model = text(field)?, + 21 => turn.label = text(field)?, + _ => {} + } + Some(()) + }) +} +fn parse_usage(bytes: &[u8], usage: &mut ParsedUsage) -> Option<()> { + fields(bytes, |field| { + match field.number { + 1 => usage.system_prompt = integer(field)?, + 2 => usage.new_input = integer(field)?, + 5 => usage.cache_read = integer(field)?, + 9 => usage.output = integer(field)?, + 10 => usage.reasoning = integer(field)?, + 11 => usage.response_id = text(field)?, + _ => {} + } + Some(()) + }) +} +fn parse_generation(bytes: &[u8], seconds: &mut Option, nanos: &mut u64) -> Option<()> { + fields(bytes, |field| { + if field.number != 4 { + return Some(()); + } + fields(message(field)?, |stamp| { + match stamp.number { + 1 => { + let value = integer(stamp)?; + if value == 0 || value > 253_402_300_799 { + return None; + } + *seconds = Some(value); + } + 2 => { + let value = integer(stamp)?; + if value > 999_999_999 { + return None; + } + *nanos = value; + } + _ => {} + } + Some(()) + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + fn varint(mut value: u64) -> Vec { + let mut bytes = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + bytes.push(byte); + if value == 0 { + return bytes; + } + } + } + fn field_varint(number: u64, value: u64) -> Vec { + let mut b = varint(number << 3); + b.extend(varint(value)); + b + } + fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut b = varint((number << 3) | 2); + b.extend(varint(value.len() as u64)); + b.extend(value); + b + } + #[test] + fn decodes_generation_usage_and_timestamp() { + let mut usage = Vec::new(); + usage.extend(field_varint(1, 10)); + usage.extend(field_varint(2, 20)); + usage.extend(field_varint(5, 30)); + usage.extend(field_varint(9, 40)); + usage.extend(field_varint(10, 50)); + usage.extend(field_bytes(11, b"response-1")); + let mut stamp = Vec::new(); + stamp.extend(field_varint(1, 1_787_572_800)); + stamp.extend(field_varint(2, 123_000_000)); + let generation = field_bytes(4, &stamp); + let mut chat = Vec::new(); + chat.extend(field_bytes(4, &usage)); + chat.extend(field_bytes(9, &generation)); + chat.extend(field_bytes(19, b"test-model-antigravity-a")); + chat.extend(field_bytes(21, b"label-a")); + let root = field_bytes(1, &chat); + let turn = parse_turn(&root).unwrap(); + let usage = turn.usage.unwrap(); + assert_eq!(usage.system_prompt, 10); + assert_eq!(usage.new_input, 20); + assert_eq!(usage.cache_read, 30); + assert_eq!(usage.output, 40); + assert_eq!(usage.reasoning, 50); + assert_eq!(usage.response_id.as_deref(), Some("response-1")); + assert_eq!(turn.timestamp_ms, Some(1_787_572_800_123)); + } + #[test] + fn rejects_malformed_varint() { + assert!(parse_turn(&[0x0a, 0x80]).is_none()); + } +} diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index ae79bb5447..650da25b86 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -8,42 +8,105 @@ use serde_json::Value; const MAX_SESSION_FILES: usize = 2048; const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; +const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub enum LocalHistoryCoverage { + Complete, + Partial, + #[default] + Unavailable, +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct LocalSessionSummary { pub total_tokens: u64, pub session_count: usize, + pub coverage: LocalHistoryCoverage, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ScanContext { + database_roots: [PathBuf; 3], + tokscale_sessions: PathBuf, +} + +impl ScanContext { + fn from_values( + home: &Path, + gemini_cli_home: Option<&str>, + tokscale_config_dir: Option<&str>, + ) -> Self { + let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); + let tokscale_base = clean_env_path(tokscale_config_dir) + .unwrap_or_else(|| home.join(".config").join("tokscale")); + Self { + database_roots: super::local_sqlite::database_roots(&gemini_base), + tokscale_sessions: tokscale_base.join("antigravity-cache").join("sessions"), + } + } + + fn capture() -> Option { + let home = dirs::home_dir()?; + let gemini = std::env::var("GEMINI_CLI_HOME").ok(); + let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); + Some(Self::from_values( + &home, + gemini.as_deref(), + tokscale.as_deref(), + )) + } +} + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) } pub fn summarize(days: u32) -> LocalSessionSummary { - summarize_paths(&tokscale_paths(None), Utc::now(), days) + let now = Utc::now(); + let Some(context) = ScanContext::capture() else { + return LocalSessionSummary::default(); + }; + match super::local_sqlite::summarize(&context.database_roots, now, days) { + super::local_sqlite::SQLiteScan::Summary(summary) => summary, + super::local_sqlite::SQLiteScan::NoDatabases => { + let (paths, truncated) = tokscale_paths(&context.tokscale_sessions); + if paths.is_empty() { + LocalSessionSummary::default() + } else { + summarize_paths(&paths, now, days, truncated) + } + } + } } /// Count local Antigravity conversation artifacts for the quota provider's /// offline fallback. Mirrors upstream #3119 without opening SQLite files. pub fn offline_conversation_count() -> usize { - let Some(home) = dirs::home_dir() else { + let Some(context) = ScanContext::capture() else { return 0; }; - offline_conversation_count_in(&home) + offline_conversation_count_context(&context) } fn offline_conversation_count_in(home: &Path) -> usize { - let gemini = home.join(".gemini"); - let roots = [ - gemini.join("antigravity-cli").join("conversations"), - gemini.join("antigravity"), - gemini.join("antigravity").join("conversations"), - ]; - let db_count = roots + offline_conversation_count_context(&ScanContext::from_values(home, None, None)) +} + +fn offline_conversation_count_context(context: &ScanContext) -> usize { + let db_count = context + .database_roots .iter() .map(|root| count_extension(root, "db")) .sum::(); if db_count > 0 { return db_count; } - tokscale_paths(Some(home)).len() + tokscale_paths(&context.tokscale_sessions).0.len() } fn count_extension(root: &Path, extension: &str) -> usize { @@ -57,55 +120,65 @@ fn count_extension(root: &Path, extension: &str) -> usize { .count() } -fn tokscale_paths(home: Option<&Path>) -> Vec { - let base = if let Some(home) = home { - home.join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions") - } else if let Ok(root) = std::env::var("TOKSCALE_CONFIG_DIR") { - PathBuf::from(root) - .join("antigravity-cache") - .join("sessions") - } else { - let Some(home) = dirs::home_dir() else { - return Vec::new(); - }; - home.join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions") - }; +fn tokscale_paths(base: &Path) -> (Vec, bool) { let Ok(entries) = fs::read_dir(base) else { - return Vec::new(); + return (Vec::new(), false); }; let mut paths: Vec<_> = entries .flatten() .map(|entry| entry.path()) - .filter(|path| path.extension().and_then(|value| value.to_str()) == Some("jsonl")) + .filter(|path| { + path.extension() + .and_then(|value| value.to_str()) + .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")) + }) .collect(); paths.sort(); - if paths.len() > MAX_SESSION_FILES { + let truncated = paths.len() > MAX_SESSION_FILES; + if truncated { paths.drain(..paths.len() - MAX_SESSION_FILES); } - paths + (paths, truncated) } -fn summarize_paths(paths: &[PathBuf], now: DateTime, days: u32) -> LocalSessionSummary { +fn summarize_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalSessionSummary { let first_day = now.with_timezone(&Local).date_naive() - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); let mut total_tokens = 0_u64; let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); + let mut complete = !truncated; for path in paths.iter().take(MAX_SESSION_FILES) { - let Ok(file) = File::open(path) else { - continue; + let file = match File::open(path) { + Ok(file) => file, + Err(_) => { + complete = false; + continue; + } }; + match file.metadata() { + Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, + Ok(_) => {} + Err(_) => complete = false, + } let mut reader = BufReader::new(file); let mut remaining = MAX_SESSION_FILE_BYTES; let mut path_had_usage = false; - while let Ok(Some(line)) = read_bounded_jsonl_line(&mut reader, &mut remaining) { + loop { + let line = match read_bounded_jsonl_line(&mut reader, &mut remaining) { + Ok(Some(line)) => line, + Ok(None) => break, + Err(_) => { + complete = false; + break; + } + }; if line.is_empty() { continue; } @@ -160,6 +233,13 @@ fn summarize_paths(paths: &[PathBuf], now: DateTime, days: u32) -> LocalSes LocalSessionSummary { total_tokens, session_count: sessions_with_usage.len(), + coverage: if paths.is_empty() { + LocalHistoryCoverage::Unavailable + } else if complete { + LocalHistoryCoverage::Complete + } else { + LocalHistoryCoverage::Partial + }, } } @@ -218,6 +298,37 @@ fn token_field(value: &Value, keys: &[&str]) -> u64 { mod tests { use super::*; + #[test] + fn scan_context_honors_non_empty_root_overrides() { + let home = Path::new(r"C:\Users\test"); + let context = + ScanContext::from_values(home, Some(r"D:\gemini-root"), Some(r"E:\tokscale-root")); + assert_eq!( + context.database_roots[0], + PathBuf::from(r"D:\gemini-root") + .join("antigravity-cli") + .join("conversations") + ); + assert_eq!( + context.tokscale_sessions, + PathBuf::from(r"E:\tokscale-root") + .join("antigravity-cache") + .join("sessions") + ); + let defaults = ScanContext::from_values(home, Some(" "), Some("")); + assert_eq!( + defaults.database_roots[1], + home.join(".gemini").join("antigravity") + ); + assert_eq!( + defaults.tokscale_sessions, + home.join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions") + ); + } + #[test] fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { let dir = tempfile::tempdir().unwrap(); @@ -228,11 +339,27 @@ mod tests { "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" )).unwrap(); let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7); + let summary = summarize_paths(&[path], now, 7, false); assert_eq!(summary.total_tokens, 135); assert_eq!(summary.session_count, 1); } + #[test] + fn truncated_or_unreadable_tokscale_history_is_partial() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); + assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); + + let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); + assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); + } #[test] fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { let dir = tempfile::tempdir().unwrap(); @@ -276,7 +403,7 @@ mod tests { ) .unwrap(); let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7); + let summary = summarize_paths(&[path], now, 7, false); assert_eq!(summary.total_tokens, 15); assert_eq!(summary.session_count, 1); } @@ -294,7 +421,7 @@ mod tests { text.push('\n'); fs::write(&path, text).unwrap(); let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7); + let summary = summarize_paths(&[path], now, 7, false); assert_eq!(summary.total_tokens, 15); assert_eq!(summary.session_count, 1); } diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs new file mode 100644 index 0000000000..52cf904a91 --- /dev/null +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -0,0 +1,570 @@ +use std::collections::{HashMap, HashSet}; +use std::fs; +use std::path::{Path, PathBuf}; +use std::time::{Duration as StdDuration, Instant}; + +use chrono::{DateTime, Duration, Local, TimeZone, Utc}; +use rusqlite::{Connection, OpenFlags, TransactionBehavior, types::ValueRef}; + +use super::local_proto::{ParsedTurn, parse_turn}; +use super::local_sessions::{LocalHistoryCoverage, LocalSessionSummary}; + +const MAX_DATABASES: usize = 500; +const MAX_DIRECTORY_ENTRIES: usize = 10_000; +const MAX_ROWS_PER_DATABASE: usize = 10_000; +const MAX_ROWS: usize = 50_000; +const MAX_BLOB_BYTES: usize = 16 * 1024 * 1024; +const MAX_DATABASE_BYTES: usize = 64 * 1024 * 1024; +const MAX_TOTAL_BYTES: usize = 128 * 1024 * 1024; +const MAX_SCHEMA_ENTRIES: usize = 128; +const MAX_SCHEMA_COLUMNS: usize = 64; +const MAX_SCHEMA_BYTES: usize = 64 * 1024; +const MAX_SCAN_DURATION: StdDuration = StdDuration::from_secs(5); + +#[derive(Debug)] +pub(super) enum SQLiteScan { + NoDatabases, + Summary(LocalSessionSummary), +} + +struct Budget { + directory_entries: usize, + databases: usize, + rows: usize, + bytes: usize, + schema_bytes: usize, + deadline: Instant, +} + +impl Budget { + fn new() -> Self { + Self::with_deadline(Instant::now() + MAX_SCAN_DURATION) + } + + fn with_deadline(deadline: Instant) -> Self { + Self { + directory_entries: 0, + databases: 0, + rows: 0, + bytes: 0, + schema_bytes: 0, + deadline, + } + } + + fn check(&self) -> bool { + Instant::now() < self.deadline + } + + fn charge_schema_text(&mut self, value: &str) -> bool { + let Some(next) = self.schema_bytes.checked_add(value.len()) else { + return false; + }; + self.schema_bytes = next; + next <= MAX_SCHEMA_BYTES + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct Event { + session: String, + row: i64, + turn: ParsedTurn, + total: u64, +} + +pub(super) fn database_roots(gemini_base: &Path) -> [PathBuf; 3] { + [ + gemini_base.join("antigravity-cli").join("conversations"), + gemini_base.join("antigravity"), + gemini_base.join("antigravity").join("conversations"), + ] +} + +pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQLiteScan { + let mut budget = Budget::new(); + let (paths, discovery_complete) = discover_databases(roots, &mut budget); + if paths.is_empty() && discovery_complete { + return SQLiteScan::NoDatabases; + } + + let first_day = now.with_timezone(&Local).date_naive() + - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); + let mut complete = discovery_complete && budget.check(); + let mut events = Vec::new(); + + for path in &paths { + if !budget.check() { + complete = false; + break; + } + budget.databases += 1; + if budget.databases > MAX_DATABASES { + complete = false; + break; + } + match read_database(path, &mut budget) { + Ok((mut rows, is_complete)) => { + events.append(&mut rows); + complete &= is_complete; + } + Err(_) => complete = false, + } + if budget.rows >= MAX_ROWS || budget.bytes >= MAX_TOTAL_BYTES { + complete = false; + break; + } + } + + let mut total_tokens = 0_u64; + let mut sessions = HashSet::new(); + let mut rows: HashMap<(String, i64), Event> = HashMap::new(); + let mut responses: HashMap<(String, String), Event> = HashMap::new(); + + for event in events { + let row_key = (event.session.clone(), event.row); + if let Some(prior) = rows.get(&row_key) { + if prior != &event { + complete = false; + } + continue; + } + + if let Some(response_id) = event + .turn + .usage + .as_ref() + .and_then(|usage| usage.response_id.as_ref()) + { + let response_key = (event.session.clone(), response_id.clone()); + if let Some(prior) = responses.get(&response_key) { + if prior.turn != event.turn { + complete = false; + } else { + rows.insert(row_key, event); + } + continue; + } + responses.insert(response_key, event.clone()); + } + + let Some(timestamp_ms) = event.turn.timestamp_ms else { + complete = false; + continue; + }; + let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { + complete = false; + continue; + }; + rows.insert(row_key, event.clone()); + if at > now || at.with_timezone(&Local).date_naive() < first_day { + continue; + } + match total_tokens.checked_add(event.total) { + Some(total) => total_tokens = total, + None => { + complete = false; + continue; + } + } + sessions.insert(event.session); + } + + SQLiteScan::Summary(LocalSessionSummary { + total_tokens, + session_count: sessions.len(), + coverage: if complete { + LocalHistoryCoverage::Complete + } else { + LocalHistoryCoverage::Partial + }, + }) +} + +fn discover_databases(roots: &[PathBuf], budget: &mut Budget) -> (Vec, bool) { + let mut paths = Vec::new(); + let mut complete = true; + + for root in roots { + if !budget.check() { + return (paths, false); + } + let resolved_root = match fs::canonicalize(root) { + Ok(path) => path, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => continue, + Err(_) => { + complete = false; + continue; + } + }; + match fs::metadata(&resolved_root) { + Ok(metadata) if metadata.is_dir() => {} + Ok(_) | Err(_) => { + complete = false; + continue; + } + } + let entries = match fs::read_dir(&resolved_root) { + Ok(entries) => entries, + Err(_) => { + complete = false; + continue; + } + }; + for entry in entries { + if !budget.check() { + return (paths, false); + } + budget.directory_entries += 1; + if budget.directory_entries > MAX_DIRECTORY_ENTRIES { + return (paths, false); + } + let entry = match entry { + Ok(entry) => entry, + Err(_) => { + complete = false; + continue; + } + }; + let path = entry.path(); + let Some(name) = path.file_name().and_then(|value| value.to_str()) else { + complete = false; + continue; + }; + if name.starts_with('.') { + continue; + } + let is_db = path + .extension() + .and_then(|value| value.to_str()) + .is_some_and(|value| value.eq_ignore_ascii_case("db")); + if !is_db { + continue; + } + let resolved = match fs::canonicalize(&path) { + Ok(path) => path, + Err(_) => { + complete = false; + continue; + } + }; + match fs::metadata(&resolved) { + Ok(metadata) if metadata.is_file() => {} + Ok(_) | Err(_) => { + complete = false; + continue; + } + } + if paths.len() >= MAX_DATABASES { + return (paths, false); + } + paths.push(resolved); + } + } + paths.sort(); + paths.dedup(); + (paths, complete) +} + +fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result<(Vec, bool)> { + if !budget.check() { + return Ok((Vec::new(), false)); + } + let mut conn = Connection::open_with_flags(path, OpenFlags::SQLITE_OPEN_READ_ONLY)?; + let tx = conn.transaction_with_behavior(TransactionBehavior::Deferred)?; + if !supported_schema(&tx, budget)? { + return Ok((Vec::new(), false)); + } + + let session = path + .file_stem() + .and_then(|value| value.to_str()) + .unwrap_or("unknown") + .to_string(); + let mut statement = tx.prepare( + "SELECT idx, CASE WHEN typeof(data) = 'blob' THEN length(data) END, CASE WHEN typeof(data) = 'blob' AND length(data) <= ?2 THEN data END FROM main.gen_metadata NOT INDEXED LIMIT ?1", + )?; + let row_limit = i64::try_from(MAX_ROWS_PER_DATABASE + 1).unwrap_or(i64::MAX); + let blob_limit = i64::try_from(MAX_BLOB_BYTES).unwrap_or(i64::MAX); + let mut query = statement.query(rusqlite::params![row_limit, blob_limit])?; + let mut database_bytes = 0usize; + let mut database_rows = 0usize; + let mut complete = true; + let mut events = Vec::new(); + + while let Some(row) = query.next()? { + if !budget.check() { + complete = false; + break; + } + database_rows += 1; + budget.rows += 1; + if database_rows > MAX_ROWS_PER_DATABASE || budget.rows > MAX_ROWS { + complete = false; + break; + } + + let idx: i64 = match row.get(0) { + Ok(value) if value >= 0 => value, + _ => { + complete = false; + continue; + } + }; + let declared: Option = row.get(1).ok(); + let Some(declared) = declared.and_then(|value| usize::try_from(value).ok()) else { + complete = false; + continue; + }; + database_bytes = match database_bytes.checked_add(declared) { + Some(value) if value <= MAX_DATABASE_BYTES => value, + _ => { + complete = false; + break; + } + }; + budget.bytes = match budget.bytes.checked_add(declared) { + Some(value) if value <= MAX_TOTAL_BYTES => value, + _ => { + complete = false; + break; + } + }; + if declared == 0 || declared > MAX_BLOB_BYTES { + complete = false; + continue; + } + + let blob = match row.get_ref(2)? { + ValueRef::Blob(bytes) if bytes.len() == declared => bytes, + _ => { + complete = false; + continue; + } + }; + let Some(turn) = parse_turn(blob) else { + complete = false; + continue; + }; + let Some(usage) = turn.usage.as_ref() else { + complete = false; + continue; + }; + if turn.timestamp_ms.is_none() { + complete = false; + continue; + } + let Some(input) = usage.system_prompt.checked_add(usage.new_input) else { + complete = false; + continue; + }; + let Some(total) = input + .checked_add(usage.output) + .and_then(|value| value.checked_add(usage.cache_read)) + .and_then(|value| value.checked_add(usage.reasoning)) + else { + complete = false; + continue; + }; + events.push(Event { + session: session.clone(), + row: idx, + turn, + total, + }); + } + + Ok((events, complete)) +} + +fn supported_schema(conn: &Connection, budget: &mut Budget) -> rusqlite::Result { + let mut statement = + conn.prepare("SELECT name, type, rootpage FROM main.sqlite_master LIMIT ?1")?; + let mut rows = statement.query([i64::try_from(MAX_SCHEMA_ENTRIES + 1).unwrap_or(i64::MAX)])?; + let mut found = false; + let mut schema_entries = 0usize; + while let Some(row) = rows.next()? { + if !budget.check() { + return Ok(false); + } + schema_entries += 1; + if schema_entries > MAX_SCHEMA_ENTRIES { + return Ok(false); + } + let name: String = row.get(0)?; + let kind: String = row.get(1)?; + if !budget.charge_schema_text(&name) || !budget.charge_schema_text(&kind) { + return Ok(false); + } + if !name.eq_ignore_ascii_case("gen_metadata") { + continue; + } + let rootpage: i64 = row.get(2)?; + if kind != "table" || rootpage <= 0 || found { + return Ok(false); + } + found = true; + } + if !found { + return Ok(false); + } + + let mut columns = HashSet::new(); + let mut schema_columns = 0usize; + let mut info = conn.prepare("PRAGMA main.table_xinfo('gen_metadata')")?; + let mut rows = info.query([])?; + while let Some(row) = rows.next()? { + if !budget.check() { + return Ok(false); + } + schema_columns += 1; + if schema_columns > MAX_SCHEMA_COLUMNS { + return Ok(false); + } + let hidden: i64 = row.get(6)?; + if hidden != 0 { + return Ok(false); + } + let name: String = row.get(1)?; + let column_type: String = row.get(2).unwrap_or_default(); + let default_value: Option = row.get(4).ok(); + if !budget.charge_schema_text(&name) || !budget.charge_schema_text(&column_type) { + return Ok(false); + } + if let Some(default_value) = default_value.as_deref() + && !budget.charge_schema_text(default_value) + { + return Ok(false); + } + columns.insert(name.to_ascii_lowercase()); + } + Ok(columns.contains("idx") && columns.contains("data")) +} + +#[cfg(test)] +mod tests { + use super::*; + use rusqlite::params; + + #[test] + fn missing_databases_falls_through() { + let dir = tempfile::tempdir().unwrap(); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::NoDatabases + )); + } + + #[test] + fn unsupported_database_is_partial_not_zero() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("database should be attempted"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + assert_eq!(summary.total_tokens, 0); + } + + #[test] + fn empty_supported_database_is_confirmed_zero() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 0); + assert_eq!(summary.session_count, 0); + } + + #[test] + fn non_blob_rows_make_coverage_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", + params![1_i64, "not-a-blob"], + ) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + #[test] + fn discovery_allows_exactly_500_databases_but_marks_501_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + for index in 0..MAX_DATABASES { + fs::write(root.join(format!("{index:03}.db")), b"").unwrap(); + } + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(complete); + + fs::write(root.join("overflow.db"), b"").unwrap(); + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(!complete); + } + + #[test] + fn expired_budget_marks_discovery_incomplete() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("one.db"), b"").unwrap(); + let mut budget = Budget::with_deadline(Instant::now()); + let (_, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert!(!complete); + } + + #[test] + fn extra_columns_and_without_rowid_schema_is_supported() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER PRIMARY KEY, data BLOB, extra TEXT) WITHOUT ROWID", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert!(supported_schema(&conn, &mut budget).unwrap()); + } + + #[test] + fn generated_columns_are_rejected() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER, data BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL)", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert!(!supported_schema(&conn, &mut budget).unwrap()); + } +} diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 685c628d72..359b235b32 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -3,7 +3,10 @@ //! Fetches usage data from Antigravity's local language server probe //! Uses Windows process detection to find CSRF token +mod local_proto; pub mod local_sessions; +mod local_sqlite; +mod quota_summary; use async_trait::async_trait; use regex_lite::Regex; @@ -20,6 +23,9 @@ use crate::core::{ const NOT_RUNNING_MESSAGE: &str = "Antigravity language server not running. Start Google Antigravity and sign in, then retry."; +const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; +const QUOTA_SUMMARY_PATH: &str = + "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; /// Antigravity provider pub struct AntigravityProvider { @@ -303,14 +309,27 @@ impl AntigravityProvider { Vec::new() } - /// Fetch user status from Antigravity API + /// Fetch user status from Antigravity API. + /// + /// v0.56.0: prefer the quota-summary endpoint so the 5-hour and weekly + /// lanes can be resolved independently across model families. The legacy + /// model-quota payload remains the compatibility fallback. + fn with_cadence_labels(mut usage: UsageSnapshot) -> UsageSnapshot { + if usage + .secondary + .as_ref() + .is_some_and(|window| window.window_minutes == Some(7 * 24 * 60)) + { + usage.secondary_label = Some("Weekly".to_string()); + } + usage + } + async fn fetch_user_status(&self) -> Result { let process_info = Self::detect_process_info()?; let api_port = Self::find_api_port(process_info.extension_port, process_info.pid).await?; - // SECURITY: TLS verification disabled for local language server (see find_api_port) - // The language server is a local loopback endpoint. Do not route it - // through the app-wide outbound proxy. + // SECURITY: TLS verification disabled only for this loopback language server. let client = crate::core::credentialed_http_client_builder() .no_proxy() .timeout(std::time::Duration::from_secs(8)) @@ -319,10 +338,55 @@ impl AntigravityProvider { .build() .map_err(|e| ProviderError::Other(e.to_string()))?; - let url = format!( - "https://127.0.0.1:{}/exa.language_server_pb.LanguageServerService/GetUserStatus", - api_port - ); + let quota_body = serde_json::json!({ "forceRefresh": true }); + match Self::fetch_local_payload( + &client, + &process_info, + api_port, + QUOTA_SUMMARY_PATH, + "a_body, + std::time::Duration::from_secs(4), + ) + .await + { + Ok(bytes) => match quota_summary::parse_usage_snapshot(&bytes) { + Ok(mut snapshot) => { + // Identity is best-effort enrichment and must not displace a + // successful quota-summary result. + let identity_body = serde_json::json!({ + "metadata": { + "ideName": "antigravity", + "extensionName": "antigravity", + "ideVersion": "unknown", + "locale": "en" + } + }); + if let Ok(identity_bytes) = Self::fetch_local_payload( + &client, + &process_info, + api_port, + GET_USER_STATUS_PATH, + &identity_body, + std::time::Duration::from_secs(1), + ) + .await + && let Ok(identity) = + serde_json::from_slice::(&identity_bytes) + { + Self::apply_user_identity(&mut snapshot, &identity); + } + return Ok(snapshot); + } + Err(error) => tracing::debug!( + %error, + "Antigravity quota summary unusable; falling back to model quotas" + ), + }, + Err(error) => tracing::debug!( + %error, + "Antigravity quota summary unavailable; falling back to model quotas" + ), + } let body = serde_json::json!({ "metadata": { @@ -332,78 +396,106 @@ impl AntigravityProvider { "locale": "en" } }); + let bytes = Self::fetch_local_payload( + &client, + &process_info, + api_port, + GET_USER_STATUS_PATH, + &body, + std::time::Duration::from_secs(8), + ) + .await?; + let response: UserStatusResponse = serde_json::from_slice(&bytes) + .map_err(|e| ProviderError::Parse(format!("Failed to parse response: {e}")))?; + self.parse_user_status(response) + } - // The `agy` CLI serves the quota endpoints without a CSRF token; the - // desktop IDE/app server requires one. Only attach the CSRF header when - // the matched process is the desktop server (and a token was found). + async fn fetch_local_payload( + client: &reqwest::Client, + process_info: &ProcessInfo, + api_port: u16, + path: &str, + body: &serde_json::Value, + timeout: std::time::Duration, + ) -> Result, ProviderError> { + let url = format!("https://127.0.0.1:{api_port}{path}"); let requires_csrf = process_info.source == ProcessSource::Ide; - let csrf_token = if requires_csrf { - process_info - .extension_server_csrf_token - .as_deref() - .unwrap_or(&process_info.csrf_token) - } else { - "" - }; + let csrf_token = process_info + .extension_server_csrf_token + .as_deref() + .unwrap_or(&process_info.csrf_token); let mut request = client .post(&url) .header("Content-Type", "application/json") .header("Connect-Protocol-Version", "1") - .json(&body); + .timeout(timeout) + .json(body); if requires_csrf { request = request.header("X-Codeium-Csrf-Token", csrf_token); } - let resp = request + let response = request .send() .await - .map_err(|e| ProviderError::Other(format!("API request failed: {}", e)))?; - - if !resp.status().is_success() { - // Retry with language server CSRF token if extension server token failed - if process_info.extension_server_csrf_token.is_some() { - let retry_resp = client - .post(&url) - .header("Content-Type", "application/json") - .header("Connect-Protocol-Version", "1") - .header("X-Codeium-Csrf-Token", &process_info.csrf_token) - .json(&body) - .send() - .await; - - if let Ok(retry) = retry_resp - && retry.status().is_success() - { - let json: UserStatusResponse = retry - .json() - .await - .map_err(|e| ProviderError::Parse(e.to_string()))?; - return self.parse_user_status(json); - } - } + .map_err(|e| ProviderError::Other(format!("API request failed: {e}")))?; + if response.status().is_success() { + return response + .bytes() + .await + .map(|bytes| bytes.to_vec()) + .map_err(|e| ProviderError::Other(format!("Failed to read response: {e}"))); + } - let status = resp.status(); - let text = resp.text().await.unwrap_or_default(); - if process_info.source == ProcessSource::Cli - && (status == reqwest::StatusCode::UNAUTHORIZED - || status == reqwest::StatusCode::FORBIDDEN - || text.to_ascii_lowercase().contains("not logged") - || text.to_ascii_lowercase().contains("login method") - || text.to_ascii_lowercase().contains("keyring")) + let status = response.status(); + let text = response.text().await.unwrap_or_default(); + if requires_csrf && process_info.extension_server_csrf_token.is_some() { + let retry = client + .post(&url) + .header("Content-Type", "application/json") + .header("Connect-Protocol-Version", "1") + .header("X-Codeium-Csrf-Token", &process_info.csrf_token) + .timeout(timeout) + .json(body) + .send() + .await; + if let Ok(retry) = retry + && retry.status().is_success() { - return Err(ProviderError::AuthRequired); + return retry + .bytes() + .await + .map(|bytes| bytes.to_vec()) + .map_err(|e| ProviderError::Other(format!("Failed to read response: {e}"))); } - return Err(ProviderError::Other(format!( - "API error {}: {}", - status, text - ))); } - let json: UserStatusResponse = resp - .json() - .await - .map_err(|e| ProviderError::Other(format!("Failed to parse response: {}", e)))?; + if process_info.source == ProcessSource::Cli + && (status == reqwest::StatusCode::UNAUTHORIZED + || status == reqwest::StatusCode::FORBIDDEN + || text.to_ascii_lowercase().contains("not logged") + || text.to_ascii_lowercase().contains("login method") + || text.to_ascii_lowercase().contains("keyring")) + { + return Err(ProviderError::AuthRequired); + } + Err(ProviderError::Other(format!("API error {status}: {text}"))) + } - self.parse_user_status(json) + fn apply_user_identity(snapshot: &mut UsageSnapshot, response: &UserStatusResponse) { + let Some(status) = response.user_status.as_ref() else { + return; + }; + snapshot.account_email = status + .email + .as_deref() + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + snapshot.login_method = status + .plan_status + .as_ref() + .and_then(|plan_status| plan_status.plan_info.as_ref()) + .and_then(|plan| plan.plan_display_name.as_ref().or(plan.plan_name.as_ref())) + .cloned(); } fn parse_user_status( @@ -538,7 +630,10 @@ impl Provider for AntigravityProvider { tracing::debug!("Fetching Antigravity usage via local probe"); match self.fetch_user_status().await { - Ok(usage) => Ok(ProviderFetchResult::new(usage, "local")), + Ok(usage) => Ok(ProviderFetchResult::new( + Self::with_cadence_labels(usage), + "local", + )), Err(e) => { let count = local_sessions::offline_conversation_count(); if count > 0 { diff --git a/rust/src/providers/antigravity/quota_summary.rs b/rust/src/providers/antigravity/quota_summary.rs new file mode 100644 index 0000000000..0cceb1a9f3 --- /dev/null +++ b/rust/src/providers/antigravity/quota_summary.rs @@ -0,0 +1,317 @@ +use chrono::{DateTime, TimeZone, Utc}; +use serde::Deserialize; + +use crate::core::{NamedRateWindow, ProviderError, RateWindow, UsageSnapshot}; + +const WINDOW_ID_PREFIX: &str = "antigravity-quota-summary-"; +const SESSION_MINUTES: u32 = 300; +const WEEKLY_MINUTES: u32 = 7 * 24 * 60; + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct QuotaSummaryEnvelope { + response: Option, + summary: Option, + description: Option, + groups: Option>, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct QuotaSummaryPayload { + #[allow(dead_code, reason = "mirrors the local quota-summary response")] + description: Option, + #[serde(default)] + groups: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct QuotaSummaryGroup { + display_name: Option, + #[allow(dead_code, reason = "mirrors the local quota-summary response")] + description: Option, + #[serde(default)] + buckets: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct QuotaSummaryBucket { + bucket_id: Option, + display_name: Option, + description: Option, + disabled: Option, + remaining_fraction: Option, + remaining: Option, + reset_time: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct QuotaSummaryRemaining { + remaining_fraction: Option, + #[serde(rename = "case")] + oneof_case: Option, + value: Option, +} + +impl QuotaSummaryBucket { + fn resolved_remaining_fraction(&self) -> Option { + self.remaining_fraction.or_else(|| { + let remaining = self.remaining.as_ref()?; + remaining.remaining_fraction.or_else(|| { + (remaining.oneof_case.as_deref() == Some("remainingFraction")) + .then_some(remaining.value) + .flatten() + }) + }) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum BucketKind { + Session, + Weekly, + Other, +} + +pub(super) fn parse_usage_snapshot(data: &[u8]) -> Result { + let envelope: QuotaSummaryEnvelope = serde_json::from_slice(data) + .map_err(|error| ProviderError::Parse(format!("Antigravity quota summary: {error}")))?; + let payload = envelope + .response + .or(envelope.summary) + .or_else(|| { + envelope.groups.map(|groups| QuotaSummaryPayload { + description: envelope.description, + groups, + }) + }) + .ok_or_else(|| ProviderError::Parse("Antigravity quota summary missing payload".into()))?; + + let windows = quota_windows(payload.groups); + if !windows.iter().any(|window| window.usage_known) { + return Err(ProviderError::Parse( + "Antigravity quota summary has no usable quota buckets".into(), + )); + } + + let primary = + most_constrained(&windows, SESSION_MINUTES).unwrap_or_else(RateWindow::no_active_session); + let secondary = most_constrained(&windows, WEEKLY_MINUTES); + let mut snapshot = UsageSnapshot::new(primary).with_primary_label("Session"); + if let Some(weekly) = secondary { + snapshot = snapshot.with_secondary(weekly); + } + snapshot.extra_rate_windows = windows; + Ok(snapshot) +} + +fn quota_windows(groups: Vec) -> Vec { + let mut indexed_groups = groups.into_iter().enumerate().collect::>(); + indexed_groups.sort_by_key(|(index, group)| (group_rank(group), *index)); + + let mut windows = Vec::new(); + for (_, group) in indexed_groups { + let group_title = group_title(&group); + let mut buckets = group.buckets.into_iter().enumerate().collect::>(); + buckets.sort_by_key(|(index, bucket)| (bucket_kind(bucket), *index)); + for (_, bucket) in buckets { + let Some(bucket_id) = non_empty(bucket.bucket_id.as_deref()) else { + continue; + }; + let kind = bucket_kind(&bucket); + let title = format!("{} {}", group_title, bucket_title(&bucket, kind)); + let remaining = bucket.resolved_remaining_fraction(); + let usage_known = !bucket.disabled.unwrap_or(false) && remaining.is_some(); + let used_percent = remaining + .map(|fraction| 100.0 - (fraction * 100.0).clamp(0.0, 100.0)) + .unwrap_or(0.0); + let window_minutes = match kind { + BucketKind::Session => Some(SESSION_MINUTES), + BucketKind::Weekly => Some(WEEKLY_MINUTES), + BucketKind::Other => None, + }; + let reset = bucket.reset_time.as_deref().and_then(parse_reset_time); + let window = RateWindow::with_details( + used_percent, + window_minutes, + reset, + bucket.description.clone(), + ); + windows.push( + NamedRateWindow::new(format!("{WINDOW_ID_PREFIX}{bucket_id}"), title, window) + .with_usage_known(usage_known), + ); + } + } + windows +} + +fn most_constrained(windows: &[NamedRateWindow], minutes: u32) -> Option { + windows + .iter() + .filter(|row| row.usage_known && row.window.window_minutes == Some(minutes)) + .max_by(|left, right| { + left.window + .used_percent + .partial_cmp(&right.window.used_percent) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| right.id.cmp(&left.id)) + }) + .map(|row| row.window.clone()) +} + +fn group_rank(group: &QuotaSummaryGroup) -> u8 { + let title = group + .display_name + .as_deref() + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if title.contains("gemini") { + 0 + } else if title.contains("claude") || title.contains("gpt") { + 1 + } else { + 2 + } +} + +fn group_title(group: &QuotaSummaryGroup) -> String { + let title = non_empty(group.display_name.as_deref()).unwrap_or("Quota"); + let lower = title.to_ascii_lowercase(); + if lower.contains("gemini") { + "Gemini".into() + } else if lower.contains("claude") || lower.contains("gpt") { + "Claude/GPT".into() + } else { + title.to_string() + } +} + +fn bucket_kind(bucket: &QuotaSummaryBucket) -> BucketKind { + let mut candidates = Vec::new(); + for raw in [bucket.bucket_id.as_deref(), bucket.display_name.as_deref()] + .into_iter() + .flatten() + { + let normalized = raw.trim().to_ascii_lowercase().replace('_', "-"); + if normalized.is_empty() { + continue; + } + candidates.push(normalized.clone()); + if let Some(stripped) = normalized.strip_suffix(" limit") { + candidates.push(stripped.to_string()); + } + } + const SESSION_ALIASES: [&str; 5] = ["session", "5h", "5-hour", "five hour", "five-hour"]; + if candidates.iter().any(|candidate| { + SESSION_ALIASES + .iter() + .any(|alias| candidate == alias || candidate.ends_with(&format!("-{alias}"))) + }) { + BucketKind::Session + } else if candidates + .iter() + .any(|candidate| candidate == "weekly" || candidate.ends_with("-weekly")) + { + BucketKind::Weekly + } else { + BucketKind::Other + } +} + +fn bucket_title(bucket: &QuotaSummaryBucket, kind: BucketKind) -> String { + match kind { + BucketKind::Session => "5-hour".into(), + BucketKind::Weekly => "weekly".into(), + BucketKind::Other => non_empty(bucket.display_name.as_deref()) + .or_else(|| non_empty(bucket.bucket_id.as_deref())) + .unwrap_or("quota") + .to_string(), + } +} + +fn non_empty(value: Option<&str>) -> Option<&str> { + value.map(str::trim).filter(|value| !value.is_empty()) +} + +fn parse_reset_time(raw: &str) -> Option> { + let raw = raw.trim(); + DateTime::parse_from_rfc3339(raw) + .ok() + .map(|value| value.with_timezone(&Utc)) + .or_else(|| { + let seconds = raw.parse::().ok()?; + if !seconds.is_finite() { + return None; + } + let whole = seconds.trunc(); + if whole < i64::MIN as f64 || whole > i64::MAX as f64 { + return None; + } + #[allow( + clippy::cast_possible_truncation, + reason = "finite epoch seconds are range-checked before this conversion" + )] + let whole_seconds = whole as i64; + Utc.timestamp_opt(whole_seconds, 0).single() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn picks_most_constrained_session_and_weekly_independently() { + let data = br#"{ + "response": {"groups": [ + {"displayName":"Gemini Models","buckets":[ + {"bucketId":"gemini-5h","displayName":"5h","remainingFraction":0.9}, + {"bucketId":"gemini-weekly","displayName":"Weekly","remainingFraction":0.1} + ]}, + {"displayName":"Claude and GPT","buckets":[ + {"bucketId":"3p-5h","displayName":"5-hour limit","remainingFraction":0.2}, + {"bucketId":"3p-weekly","displayName":"Weekly","remainingFraction":0.8} + ]} + ]} + }"#; + let snapshot = parse_usage_snapshot(data).unwrap(); + assert!((snapshot.primary.used_percent - 80.0).abs() < 0.001); + assert_eq!(snapshot.primary.window_minutes, Some(300)); + let weekly = snapshot.secondary.unwrap(); + assert!((weekly.used_percent - 90.0).abs() < 0.001); + assert_eq!(weekly.window_minutes, Some(10_080)); + assert_eq!(snapshot.extra_rate_windows.len(), 4); + } + + #[test] + fn nested_oneof_remaining_and_disabled_buckets_preserve_unknown_state() { + let data = br#"{ + "groups": [{"displayName":"Gemini","buckets":[ + {"bucketId":"gemini_session","displayName":"Session","remaining":{"case":"remainingFraction","value":0.25}}, + {"bucketId":"gemini-weekly","displayName":"Weekly","disabled":true,"remainingFraction":0.01}, + {"bucketId":"future","displayName":"Daily"} + ]}] + }"#; + let snapshot = parse_usage_snapshot(data).unwrap(); + assert!((snapshot.primary.used_percent - 75.0).abs() < 0.001); + assert!(snapshot.secondary.is_none()); + assert_eq!(snapshot.extra_rate_windows.len(), 3); + assert!(snapshot.extra_rate_windows[0].usage_known); + assert!(!snapshot.extra_rate_windows[1].usage_known); + assert!(!snapshot.extra_rate_windows[2].usage_known); + } + + #[test] + fn rejects_empty_or_unusable_quota_summary() { + assert!(parse_usage_snapshot(br#"{}"#).is_err()); + assert!(parse_usage_snapshot( + br#"{"groups":[{"displayName":"Gemini","buckets":[{"bucketId":"weekly","displayName":"Weekly"}]}]}"# + ) + .is_err()); + } +} diff --git a/rust/src/providers/antigravity/tests.rs b/rust/src/providers/antigravity/tests.rs index 4ed757fb79..9e6545bf24 100644 --- a/rust/src/providers/antigravity/tests.rs +++ b/rust/src/providers/antigravity/tests.rs @@ -1,5 +1,14 @@ use super::*; +#[test] +fn cadence_labels_are_owned_by_antigravity_snapshot() { + let mut secondary = RateWindow::new(20.0); + secondary.window_minutes = Some(7 * 24 * 60); + let usage = UsageSnapshot::new(RateWindow::new(10.0)).with_secondary(secondary); + let usage = AntigravityProvider::with_cadence_labels(usage); + assert_eq!(usage.secondary_label.as_deref(), Some("Weekly")); +} + #[test] fn test_classify_model_families() { assert_eq!(classify_model("Claude 3.5 Sonnet"), ModelFamily::Claude); diff --git a/rust/src/providers/codex/api.rs b/rust/src/providers/codex/api.rs index 554c545291..8a86f52de2 100755 --- a/rust/src/providers/codex/api.rs +++ b/rust/src/providers/codex/api.rs @@ -6,6 +6,7 @@ use super::{pat, weekly_reset}; use crate::core::{ CostSnapshot, NamedRateWindow, ProviderError, RateWindow, RateWindowCadence, UsageSnapshot, }; +use base64::Engine; use chrono::{DateTime, TimeZone, Utc}; use serde::Deserialize; use std::path::PathBuf; @@ -20,6 +21,7 @@ const CREDENTIAL_CACHE_TTL: Duration = Duration::from_secs(5); /// refreshed it. Matches the CLI's own `needs_refresh` window (8 days) so a /// token the CLI considers fresh is also trusted here (upstream 0.50.1 #2944). const EXTERNAL_OAUTH_STALENESS_WINDOW: chrono::TimeDelta = chrono::Duration::days(8); +const EXTERNAL_OAUTH_REFRESH_WINDOW: chrono::TimeDelta = chrono::Duration::minutes(5); static CREDENTIAL_CACHE: OnceLock>> = OnceLock::new(); @@ -127,8 +129,24 @@ impl CodexApi { Ok((usage, first_cost)) } weekly_reset::InitialDecision::RequiresConfirmation => { + let confirmation = self.fetch_usage_once(&creds, &base_url).await; let (confirmation_usage, confirmation_cost, confirmation_credits) = - self.fetch_usage_once(&creds, &base_url).await?; + match confirmation { + Ok(value) => value, + Err(error) => { + tracing::debug!( + %error, + "Codex weekly reset confirmation failed; preserving first successful usage" + ); + let result = Self::preserve_after_confirmation_failure( + &state, + first_usage, + first_cost, + ); + weekly_reset::save(&scope, &state); + return Ok(result); + } + }; let confirmation_inventory = weekly_reset::inventory(confirmation_credits.as_ref(), Utc::now()); match weekly_reset::confirmation_decision( @@ -159,6 +177,17 @@ impl CodexApi { } } + fn preserve_after_confirmation_failure( + state: &weekly_reset::AccountState, + first_usage: UsageSnapshot, + first_cost: Option, + ) -> (UsageSnapshot, Option) { + ( + weekly_reset::preserve_weekly(state, first_usage), + first_cost, + ) + } + async fn fetch_usage_once( &self, creds: &CodexCredentials, @@ -280,6 +309,7 @@ impl CodexApi { access_token: trimmed.to_string(), account_id: None, is_external_oauth: false, + access_token_expires_at: None, last_refresh: None, }); } @@ -317,10 +347,13 @@ impl CodexApi { .and_then(|v| v.as_str()) .and_then(parse_timestamp); + let access_token_expires_at = parse_access_token_expiry(&access_token); + Ok(CodexCredentials { access_token, account_id, is_external_oauth: has_refresh_token, + access_token_expires_at, last_refresh, }) } @@ -332,17 +365,32 @@ impl CodexApi { /// not an API key). "Stale" means the CLI has not refreshed the token /// recently (no `last_refresh`, or older than the staleness window). fn enforce_external_oauth_gate(credentials: &CodexCredentials) -> Result<(), ProviderError> { + Self::enforce_external_oauth_gate_at( + credentials, + crate::settings::Settings::load().codex_external_oauth_sources_allowed, + Utc::now(), + ) + } + + fn enforce_external_oauth_gate_at( + credentials: &CodexCredentials, + external_sources_allowed: bool, + now: DateTime, + ) -> Result<(), ProviderError> { if !credentials.is_external_oauth { return Ok(()); } - if crate::settings::Settings::load().codex_external_oauth_sources_allowed { - return Ok(()); + if !external_sources_allowed { + let is_stale = credentials + .last_refresh + .is_none_or(|last| now - last > EXTERNAL_OAUTH_STALENESS_WINDOW); + if is_stale { + return Err(ProviderError::AuthRequired); + } } - let now = Utc::now(); - let is_stale = credentials - .last_refresh - .is_none_or(|last| now - last > EXTERNAL_OAUTH_STALENESS_WINDOW); - if is_stale { + if let Some(expires_at) = credentials.access_token_expires_at + && expires_at - now <= EXTERNAL_OAUTH_REFRESH_WINDOW + { return Err(ProviderError::AuthRequired); } Ok(()) @@ -881,6 +929,9 @@ struct CodexCredentials { /// `refresh_token`), as opposed to an `OPENAI_API_KEY`. Used by the /// `codex_external_oauth_sources_allowed` gate (upstream 0.50.1 #2944). is_external_oauth: bool, + /// Native access-token JWT expiry. When available, this is authoritative + /// for refresh scheduling; the CLI still owns the refresh lifecycle. + access_token_expires_at: Option>, /// `last_refresh` timestamp from auth.json, when present. Used to detect /// stale external OAuth tokens that should fail closed when the opt-in /// setting is OFF. @@ -1039,6 +1090,16 @@ fn timestamp_to_datetime(timestamp: Option) -> Option> { /// Parse an ISO-8601 / RFC-3339 timestamp from the `last_refresh` field of /// auth.json. Accepts the same formats the Codex CLI writes. +fn parse_access_token_expiry(token: &str) -> Option> { + let payload = token.split('.').nth(1)?; + let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(payload) + .or_else(|_| base64::engine::general_purpose::URL_SAFE.decode(payload)) + .ok()?; + let json: serde_json::Value = serde_json::from_slice(&decoded).ok()?; + let exp = json.get("exp")?.as_i64()?; + Utc.timestamp_opt(exp, 0).single() +} fn parse_timestamp(raw: &str) -> Option> { let trimmed = raw.trim(); if trimmed.is_empty() { @@ -1752,11 +1813,21 @@ mod tests { // ── Upstream 0.50.1 #2944: external OAuth source gate ────────────────── + #[test] + fn confirmation_failure_fallback_keeps_first_successful_usage_and_cost() { + let state = weekly_reset::AccountState::default(); + let first = UsageSnapshot::new(RateWindow::new(10.0)).with_secondary(RateWindow::new(0.5)); + let cost = Some(CostSnapshot::new(3.25, "USD", "Monthly")); + let (usage, kept_cost) = CodexApi::preserve_after_confirmation_failure(&state, first, cost); + assert!((usage.secondary.expect("weekly").used_percent - 0.5).abs() < f64::EPSILON); + assert_eq!(kept_cost.expect("cost").used, 3.25); + } #[test] fn api_key_credentials_are_not_external_oauth() { let creds = CodexApi::parse_credentials_json(r#"{"OPENAI_API_KEY": "sk-test"}"#) .expect("credentials"); assert!(!creds.is_external_oauth); + assert!(creds.access_token_expires_at.is_none()); assert!(creds.last_refresh.is_none()); assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); } @@ -1774,6 +1845,7 @@ mod tests { ) .expect("credentials"); assert!(creds.is_external_oauth); + assert!(creds.access_token_expires_at.is_none()); assert!(creds.last_refresh.is_none()); } @@ -1797,6 +1869,7 @@ mod tests { access_token: "access".to_string(), account_id: None, is_external_oauth: true, + access_token_expires_at: None, last_refresh: None, }; let err = CodexApi::enforce_external_oauth_gate(&creds) @@ -1811,6 +1884,7 @@ mod tests { access_token: "access".to_string(), account_id: None, is_external_oauth: true, + access_token_expires_at: None, last_refresh: Some(old), }; let err = CodexApi::enforce_external_oauth_gate(&creds) @@ -1825,11 +1899,56 @@ mod tests { access_token: "access".to_string(), account_id: None, is_external_oauth: true, + access_token_expires_at: None, last_refresh: Some(fresh), }; assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); } + #[test] + fn external_oauth_staleness_gate_precedes_future_jwt_expiry() { + let now = Utc::now(); + let future = now + chrono::Duration::hours(2); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(format!(r#"{{"exp":{}}}"#, future.timestamp())); + let token = format!("header.{payload}.signature"); + let json = format!( + r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"2026-01-01T00:00:00Z"}}"# + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + assert!(creds.access_token_expires_at.is_some()); + let err = CodexApi::enforce_external_oauth_gate_at(&creds, false, now) + .expect_err("stale external OAuth must not be revived by JWT expiry"); + assert!(matches!(err, ProviderError::AuthRequired)); + assert!(CodexApi::enforce_external_oauth_gate_at(&creds, true, now).is_ok()); + } + + #[test] + fn external_oauth_gate_requires_cli_refresh_when_jwt_is_near_expiry() { + let soon = Utc::now() + chrono::Duration::minutes(2); + let payload = base64::engine::general_purpose::URL_SAFE_NO_PAD + .encode(format!(r#"{{"exp":{}}}"#, soon.timestamp())); + let token = format!("header.{payload}.signature"); + let json = format!( + r#"{{"tokens":{{"access_token":"{token}","refresh_token":"refresh"}},"last_refresh":"{}"}}"#, + Utc::now().to_rfc3339() + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + let err = CodexApi::enforce_external_oauth_gate(&creds) + .expect_err("near-expiry native OAuth must refresh through the CLI"); + assert!(matches!(err, ProviderError::AuthRequired)); + } + + #[test] + fn malformed_or_opaque_jwt_falls_back_to_last_refresh() { + let fresh = Utc::now().to_rfc3339(); + let json = format!( + r#"{{"tokens":{{"access_token":"opaque-token","refresh_token":"refresh"}},"last_refresh":"{fresh}"}}"# + ); + let creds = CodexApi::parse_credentials_json(&json).expect("credentials"); + assert!(creds.access_token_expires_at.is_none()); + assert!(CodexApi::enforce_external_oauth_gate(&creds).is_ok()); + } #[test] fn parse_timestamp_reads_iso8601() { assert!(parse_timestamp("2026-08-17T10:00:00Z").is_some()); diff --git a/rust/src/providers/codex/weekly_reset.rs b/rust/src/providers/codex/weekly_reset.rs index f0548f47e1..8ef0b4924b 100644 --- a/rust/src/providers/codex/weekly_reset.rs +++ b/rust/src/providers/codex/weekly_reset.rs @@ -162,7 +162,6 @@ pub(super) fn inventory( let mut credits = reset .credits .iter() - .filter(|credit| credit_is_available(credit)) .filter_map(|credit| credit_identity(credit, observed_at)) .collect::>(); credits.sort_by(|left, right| { @@ -172,7 +171,11 @@ pub(super) fn inventory( .then_with(|| left.status.cmp(&right.status)) .then_with(|| left.expires_at.cmp(&right.expires_at)) }); - if reset.available_count > 0 && credits.len() != usize::try_from(reset.available_count).ok()? { + let available = credits + .iter() + .filter(|credit| credit_identity_is_available(credit)) + .count(); + if available != usize::try_from(reset.available_count).ok()? { return None; } Some(CreditInventory { @@ -181,11 +184,8 @@ pub(super) fn inventory( }) } -fn credit_is_available(credit: &ResetCredit) -> bool { - credit - .status - .as_deref() - .is_none_or(|status| status.is_empty() || status.eq_ignore_ascii_case("available")) +fn credit_identity_is_available(credit: &CreditIdentity) -> bool { + credit.status.is_empty() || credit.status.eq_ignore_ascii_case("available") } fn credit_identity(credit: &ResetCredit, observed_at: DateTime) -> Option { @@ -578,6 +578,47 @@ mod tests { } } + #[test] + fn inventory_retains_consumed_status_rows_but_counts_only_available_credits() { + let reset = ResetCredits { + available_count: 1, + credits: vec![ + ResetCredit { + id: Some("available-a".into()), + reset_type: Some("weekly".into()), + status: Some("available".into()), + expires_at: None, + }, + ResetCredit { + id: Some("redeeming-b".into()), + reset_type: Some("weekly".into()), + status: Some("redeeming".into()), + expires_at: None, + }, + ResetCredit { + id: Some("redeemed-c".into()), + reset_type: Some("weekly".into()), + status: Some("redeemed".into()), + expires_at: None, + }, + ], + }; + let inventory = super::inventory(Some(&reset), now()).expect("credit inventory"); + assert_eq!(inventory.available_count, 1); + assert_eq!(inventory.credits.len(), 3); + assert!( + inventory + .credits + .iter() + .any(|credit| credit.status == "redeeming") + ); + assert!( + inventory + .credits + .iter() + .any(|credit| credit.status == "redeemed") + ); + } #[test] fn early_low_usage_requires_confirmation_without_spending_credit() { let mut state = baseline(); diff --git a/rust/src/providers/cursor/token_cost.rs b/rust/src/providers/cursor/token_cost.rs index 0e316e7901..07796b419f 100644 --- a/rust/src/providers/cursor/token_cost.rs +++ b/rust/src/providers/cursor/token_cost.rs @@ -3,25 +3,83 @@ //! `POST /api/dashboard/get-filtered-usage-events` — per-model API-rate totals //! from `tokenUsage.totalCents` and plan-metered totals from `chargedCents`. -use chrono::{DateTime, Duration, Utc}; -use serde::Deserialize; -use serde_json::json; +use chrono::{DateTime, Duration, TimeZone, Utc}; +use serde::de::Error as _; +use serde::{Deserialize, Deserializer}; +use serde_json::{Value, json}; -use crate::core::{CostSnapshot, NamedRateWindow, ProviderError, RateWindow}; +use crate::core::{CostSnapshot, CostUsagePricing, NamedRateWindow, ProviderError, RateWindow}; const EVENTS_PATH: &str = "/api/dashboard/get-filtered-usage-events"; const PAGE_SIZE: usize = 200; /// Keep fetches bounded for menu-bar refresh latency. const MAX_PAGES: usize = 5; -#[derive(Debug, Deserialize)] -#[serde(rename_all = "camelCase")] +#[derive(Debug)] struct UsageEventsPage { total_usage_events_count: Option, - #[serde(default)] usage_events_display: Vec, } +impl<'de> Deserialize<'de> for UsageEventsPage { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = Value::deserialize(deserializer)?; + let object = value + .as_object() + .ok_or_else(|| D::Error::custom("Cursor usage events response must be an object"))?; + if object.is_empty() { + return Ok(Self { + total_usage_events_count: Some(0), + usage_events_display: Vec::new(), + }); + } + if object.contains_key("error") { + return Err(D::Error::custom( + "Cursor usage events error envelope is not an empty result", + )); + } + + let total_usage_events_count = match object.get("totalUsageEventsCount") { + Some(value) => { + let count = strict_i64(value).ok_or_else(|| { + D::Error::custom("Cursor usage event count must be a finite integer") + })?; + if count < 0 { + return Err(D::Error::custom( + "Cursor usage event count cannot be negative", + )); + } + Some(count) + } + None => None, + }; + if object.len() == 1 && total_usage_events_count.is_some() { + return Ok(Self { + total_usage_events_count, + usage_events_display: Vec::new(), + }); + } + + let events = object + .get("usageEventsDisplay") + .ok_or_else(|| D::Error::custom("Cursor usage events array is missing"))?; + if !events.is_array() { + return Err(D::Error::custom( + "Cursor usage events field must be an array", + )); + } + let usage_events_display = serde_json::from_value(events.clone()) + .map_err(|error| D::Error::custom(format!("Cursor usage events: {error}")))?; + Ok(Self { + total_usage_events_count, + usage_events_display, + }) + } +} + #[derive(Debug, Deserialize, Clone)] #[serde(rename_all = "camelCase")] struct UsageEvent { @@ -35,23 +93,63 @@ struct UsageEvent { token_usage: Option, #[serde(default, deserialize_with = "deserialize_opt_f64")] charged_cents: Option, - #[serde(default, deserialize_with = "deserialize_opt_f64")] - cursor_token_fee: Option, } -#[derive(Debug, Deserialize, Clone)] -#[serde(rename_all = "camelCase")] +#[derive(Debug, Clone, Copy, PartialEq)] +enum EventCost { + Valid(f64), + Omitted, + Invalid, +} + +#[derive(Debug, Clone)] struct EventTokenUsage { - #[serde(default, deserialize_with = "deserialize_i64")] input_tokens: i64, - #[serde(default, deserialize_with = "deserialize_i64")] output_tokens: i64, - #[serde(default, deserialize_with = "deserialize_i64")] cache_write_tokens: i64, - #[serde(default, deserialize_with = "deserialize_i64")] cache_read_tokens: i64, - #[serde(default, deserialize_with = "deserialize_opt_f64")] - total_cents: Option, + cost: EventCost, +} + +impl<'de> Deserialize<'de> for EventTokenUsage { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = Value::deserialize(deserializer)?; + let object = value + .as_object() + .ok_or_else(|| D::Error::custom("Cursor token usage must be an object"))?; + let count = |key: &str| object.get(key).and_then(strict_i64).unwrap_or(0); + let cost = match object.get("totalCents") { + None | Some(Value::Null) => EventCost::Omitted, + Some(value) => match finite_f64(value) { + Some(cents) if cents >= 0.0 => EventCost::Valid(cents), + _ => EventCost::Invalid, + }, + }; + Ok(Self { + input_tokens: count("inputTokens"), + output_tokens: count("outputTokens"), + cache_write_tokens: count("cacheWriteTokens"), + cache_read_tokens: count("cacheReadTokens"), + cost, + }) + } +} + +impl EventTokenUsage { + fn validated_counts(&self) -> Option<(u64, u64, u64, u64)> { + let input = u64::try_from(self.input_tokens).ok()?; + let output = u64::try_from(self.output_tokens).ok()?; + let cache_write = u64::try_from(self.cache_write_tokens).ok()?; + let cache_read = u64::try_from(self.cache_read_tokens).ok()?; + let total = input + .checked_add(output)? + .checked_add(cache_write)? + .checked_add(cache_read)?; + (total > 0).then_some((input, output, cache_write, cache_read)) + } } #[derive(Debug, Clone, Default)] @@ -62,6 +160,12 @@ pub struct CursorTokenCostReport { pub metered_usd: Option, /// Per-model API-rate spend for extra tray windows. pub by_model_usd: Vec<(String, f64)>, + /// Vendor-reported list-price requests. + pub priced_requests: u32, + /// Requests estimated from model pricing because Cursor omitted totalCents. + pub estimated_requests: u32, + /// Requests whose API-rate cost is unknown or invalid. + pub unpriced_requests: u32, } impl CursorTokenCostReport { @@ -97,6 +201,10 @@ impl CursorTokenCostReport { .unwrap_or(self.api_rate_usd); let period = if self.metered_usd.is_some() { "Token cost (metered, billing window)" + } else if self.unpriced_requests > 0 { + "Token cost (API-rate partial, billing window)" + } else if self.estimated_requests > 0 { + "Token cost (API-rate estimate, billing window)" } else { "Token cost (API-rate, billing window)" }; @@ -196,61 +304,185 @@ async fn fetch_page( fn summarize_events(events: &[UsageEvent]) -> CursorTokenCostReport { use std::collections::{HashMap, HashSet}; + let mut by_model: HashMap = HashMap::new(); let mut invalid_models: HashSet = HashSet::new(); let mut metered_cents = 0.0; - let mut metered_complete = !events.is_empty(); + let mut metered_complete = true; + let mut saw_metered_event = false; + let mut priced_requests = 0u32; + let mut estimated_requests = 0u32; + let mut unpriced_requests = 0u32; for event in events { + let Some(timestamp_ms) = event.timestamp_ms.filter(|value| *value > 0) else { + continue; + }; + + saw_metered_event = true; + match event + .charged_cents + .filter(|value| value.is_finite() && *value >= 0.0) + { + Some(cents) => { + let next = metered_cents + cents; + if next.is_finite() { + metered_cents = next; + } else { + metered_complete = false; + } + } + None => metered_complete = false, + } + + let Some(usage) = event.token_usage.as_ref() else { + continue; + }; + if usage.validated_counts().is_none() { + continue; + } let model = event .model .as_deref() .map(str::trim) - .filter(|s| !s.is_empty()) + .filter(|value| !value.is_empty()) .unwrap_or("unknown") .to_string(); - if let Some(list_cents) = event.token_usage.as_ref().and_then(|u| u.total_cents) { - if !list_cents.is_finite() || list_cents < 0.0 { - invalid_models.insert(model.clone()); - by_model.remove(&model); - } else if !invalid_models.contains(&model) { - *by_model.entry(model.clone()).or_insert(0.0) += list_cents; - } - } - match event.charged_cents.filter(|v| v.is_finite() && *v >= 0.0) { - Some(c) => metered_cents += c, - None => { - // cursorTokenFee is sometimes the only metered field. - if let Some(fee) = event - .cursor_token_fee - .filter(|v| v.is_finite() && *v >= 0.0) - { - metered_cents += fee; + match usage.cost { + EventCost::Valid(cents) if cents.is_finite() && cents >= 0.0 => { + priced_requests = priced_requests.saturating_add(1); + add_model_cost(&mut by_model, &mut invalid_models, &model, cents); + } + EventCost::Omitted => { + if let Some(cents) = estimated_list_price_cents(usage, &model, timestamp_ms) { + estimated_requests = estimated_requests.saturating_add(1); + add_model_cost(&mut by_model, &mut invalid_models, &model, cents); } else { - metered_complete = false; + unpriced_requests = unpriced_requests.saturating_add(1); } } + EventCost::Invalid | EventCost::Valid(_) => { + unpriced_requests = unpriced_requests.saturating_add(1); + invalid_models.insert(model.clone()); + by_model.remove(&model); + } } } let mut by_model_usd: Vec<(String, f64)> = by_model .into_iter() .filter(|(model, _)| !invalid_models.contains(model)) - .map(|(m, cents)| (m, cents / 100.0)) - .filter(|(_, usd)| *usd > 0.0) + .map(|(model, cents)| (model, cents / 100.0)) + .filter(|(_, usd)| usd.is_finite() && *usd > 0.0) .collect(); - by_model_usd.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + by_model_usd.sort_by(|left, right| { + right + .1 + .partial_cmp(&left.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| left.0.cmp(&right.0)) + }); let api_rate_usd = by_model_usd.iter().map(|(_, usd)| *usd).sum(); CursorTokenCostReport { api_rate_usd, - metered_usd: if metered_complete && metered_cents > 0.0 { + metered_usd: if saw_metered_event && metered_complete { Some(metered_cents / 100.0) } else { None }, by_model_usd, + priced_requests, + estimated_requests, + unpriced_requests, + } +} + +fn add_model_cost( + by_model: &mut std::collections::HashMap, + invalid_models: &mut std::collections::HashSet, + model: &str, + cents: f64, +) { + if invalid_models.contains(model) { + return; + } + let next = by_model.get(model).copied().unwrap_or(0.0) + cents; + if next.is_finite() && next >= 0.0 { + by_model.insert(model.to_string(), next); + } else { + invalid_models.insert(model.to_string()); + by_model.remove(model); + } +} + +fn estimated_list_price_cents( + usage: &EventTokenUsage, + model: &str, + timestamp_ms: i64, +) -> Option { + if usage.cost != EventCost::Omitted { + return None; + } + let (input, output, cache_write, cache_read) = usage.validated_counts()?; + let pricing_day = Utc + .timestamp_millis_opt(timestamp_ms) + .single()? + .date_naive(); + + let codex_input = input.checked_add(cache_read)?.checked_add(cache_write)?; + if let Some(usd) = CostUsagePricing::codex_cost_usd_at_date( + model, + codex_input, + cache_read, + output, + pricing_day, + ) { + let cents = usd * 100.0; + return cents.is_finite().then_some(cents); + } + + let alias = cursor_claude_catalog_model(model); + let input = i32::try_from(input).ok()?; + let output = i32::try_from(output).ok()?; + let cache_write = i32::try_from(cache_write).ok()?; + let cache_read = i32::try_from(cache_read).ok()?; + CostUsagePricing::claude_cost_usd(&alias, input, cache_read, cache_write, output).and_then( + |usd| { + let cents = usd * 100.0; + cents.is_finite().then_some(cents) + }, + ) +} + +fn cursor_claude_catalog_model(model: &str) -> String { + let Some(rest) = model.strip_prefix("claude-") else { + return model.to_string(); + }; + let mut parts = rest.splitn(3, '-'); + let Some(version) = parts.next() else { + return model.to_string(); + }; + let Some(family) = parts.next() else { + return model.to_string(); + }; + if !matches!(family, "sonnet" | "opus" | "haiku") { + return model.to_string(); + } + let Some((major, minor)) = version.split_once('.') else { + return model.to_string(); + }; + if major.is_empty() + || minor.is_empty() + || !major.chars().all(|ch| ch.is_ascii_digit()) + || !minor.chars().all(|ch| ch.is_ascii_digit()) + { + return model.to_string(); + } + match parts.next() { + Some(suffix) if !suffix.is_empty() => format!("claude-{family}-{major}-{minor}-{suffix}"), + _ => format!("claude-{family}-{major}-{minor}"), } } @@ -259,52 +491,50 @@ pub fn default_since() -> DateTime { Utc::now() - Duration::days(30) } -fn deserialize_i64<'de, D>(deserializer: D) -> Result -where - D: serde::Deserializer<'de>, -{ - use serde::de::{self, Visitor}; - struct V; - impl<'de> Visitor<'de> for V { - type Value = i64; - fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - f.write_str("int or string int") - } - fn visit_i64(self, v: i64) -> Result { - Ok(v) - } - fn visit_u64(self, v: u64) -> Result { - // Usage-event counts from the Cursor API are small non-negative - // integers, far below i64::MAX. - #[expect( - clippy::cast_possible_wrap, - reason = "usage-event counts are small non-negative integers" - )] - let count = v as i64; - Ok(count) - } - fn visit_f64(self, v: f64) -> Result { - // Token counts are whole numbers; the fractional part is rounding - // noise from JSON float parsing. - #[expect( - clippy::cast_possible_truncation, - reason = "token counts are whole numbers; fractional part is rounding noise" - )] - let count = v as i64; - Ok(count) - } - fn visit_str(self, v: &str) -> Result { - v.parse().map_err(E::custom) +fn finite_f64(value: &Value) -> Option { + let parsed = value + .as_f64() + .or_else(|| value.as_i64().map(|number| number as f64)) + .or_else(|| value.as_u64().map(|number| number as f64)) + .or_else(|| value.as_str()?.trim().parse::().ok())?; + parsed.is_finite().then_some(parsed) +} + +fn strict_i64(value: &Value) -> Option { + if let Some(number) = value.as_i64() { + return Some(number); + } + if let Some(number) = value.as_u64() { + return i64::try_from(number).ok(); + } + if let Some(text) = value.as_str() { + let text = text.trim(); + if let Ok(number) = text.parse::() { + return Some(number); } } - deserializer.deserialize_any(V) + let parsed = value + .as_f64() + .or_else(|| value.as_str()?.trim().parse::().ok())?; + if !parsed.is_finite() + || parsed.fract() != 0.0 + || parsed < i64::MIN as f64 + || parsed > i64::MAX as f64 + { + return None; + } + #[allow( + clippy::cast_possible_truncation, + reason = "range and integrality are checked before the cast" + )] + Some(parsed as i64) } - fn deserialize_opt_i64<'de, D>(deserializer: D) -> Result, D::Error> where D: serde::Deserializer<'de>, { - Ok(Some(deserialize_i64(deserializer)?)) + let value = Value::deserialize(deserializer)?; + Ok(strict_i64(&value)) } fn deserialize_opt_f64<'de, D>(deserializer: D) -> Result, D::Error> @@ -355,10 +585,9 @@ mod tests { output_tokens: 5, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: Some(25.0), + cost: EventCost::Valid(25.0), }), charged_cents: Some(10.0), - cursor_token_fee: None, }, UsageEvent { timestamp_ms: Some(2), @@ -368,10 +597,9 @@ mod tests { output_tokens: 1, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: Some(75.0), + cost: EventCost::Valid(75.0), }), charged_cents: Some(40.0), - cursor_token_fee: None, }, ]; let report = summarize_events(&events); @@ -394,10 +622,9 @@ mod tests { output_tokens: 1, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: Some(25.0), + cost: EventCost::Valid(25.0), }), charged_cents: Some(0.0), - cursor_token_fee: None, }, UsageEvent { timestamp_ms: Some(2), @@ -407,15 +634,16 @@ mod tests { output_tokens: 1, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: None, + cost: EventCost::Omitted, }), charged_cents: Some(0.0), - cursor_token_fee: None, }, ]; let report = summarize_events(&events); - assert!((report.api_rate_usd - 0.25).abs() < 0.001); - assert_eq!(report.by_model_usd, vec![("gpt-5".to_string(), 0.25)]); + assert!(report.api_rate_usd > 0.25); + assert_eq!(report.priced_requests, 1); + assert_eq!(report.estimated_requests, 1); + assert_eq!(report.unpriced_requests, 0); } #[test] @@ -429,10 +657,9 @@ mod tests { output_tokens: 1, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: Some(-1.0), + cost: EventCost::Invalid, }), charged_cents: Some(0.0), - cursor_token_fee: None, }, UsageEvent { timestamp_ms: Some(2), @@ -442,10 +669,9 @@ mod tests { output_tokens: 1, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: Some(50.0), + cost: EventCost::Valid(50.0), }), charged_cents: Some(0.0), - cursor_token_fee: None, }, ]; let report = summarize_events(&events); @@ -463,13 +689,110 @@ mod tests { output_tokens: 1, cache_write_tokens: 0, cache_read_tokens: 0, - total_cents: Some(10.0), + cost: EventCost::Valid(10.0), }), charged_cents: None, - cursor_token_fee: None, }]; let report = summarize_events(&events); assert!(report.metered_usd.is_none()); assert!((report.api_rate_usd - 0.1).abs() < 0.001); } + + #[test] + fn strict_page_decoder_accepts_literal_and_count_only_empty_results() { + let empty: UsageEventsPage = serde_json::from_str("{}").unwrap(); + assert_eq!(empty.total_usage_events_count, Some(0)); + assert!(empty.usage_events_display.is_empty()); + + let counted: UsageEventsPage = + serde_json::from_str(r#"{"totalUsageEventsCount":2}"#).unwrap(); + assert_eq!(counted.total_usage_events_count, Some(2)); + assert!(counted.usage_events_display.is_empty()); + } + + #[test] + fn strict_page_decoder_rejects_error_and_malformed_empty_envelopes() { + for json in [ + r#"{"error":"temporarily unavailable"}"#, + r#"{"usageEventsDisplay":null}"#, + r#"{"usageEventsDisplay":{}}"#, + r#"{"unknown":null}"#, + r#"{"totalUsageEventsCount":-1,"usageEventsDisplay":[]}"#, + r#"{"totalUsageEventsCount":"Infinity","usageEventsDisplay":[]}"#, + "[]", + "null", + ] { + assert!( + serde_json::from_str::(json).is_err(), + "must reject {json}" + ); + } + } + + #[test] + fn omitted_known_cost_is_estimated_but_unknown_model_stays_unpriced() { + let timestamp = 1_700_000_000_000; + let events = vec![ + UsageEvent { + timestamp_ms: Some(timestamp), + model: Some("gpt-5".into()), + token_usage: Some(EventTokenUsage { + input_tokens: 200, + output_tokens: 20, + cache_write_tokens: 0, + cache_read_tokens: 0, + cost: EventCost::Omitted, + }), + charged_cents: Some(10.0), + }, + UsageEvent { + timestamp_ms: Some(timestamp + 1), + model: Some("fixture-model".into()), + token_usage: Some(EventTokenUsage { + input_tokens: 7, + output_tokens: 0, + cache_write_tokens: 0, + cache_read_tokens: 0, + cost: EventCost::Omitted, + }), + charged_cents: Some(10.0), + }, + ]; + let report = summarize_events(&events); + assert!((report.api_rate_usd - 0.00045).abs() < 1e-9); + assert_eq!(report.priced_requests, 0); + assert_eq!(report.estimated_requests, 1); + assert_eq!(report.unpriced_requests, 1); + assert_eq!(report.metered_usd, Some(0.20)); + } + + #[test] + fn cursor_claude_alias_estimate_keeps_cache_buckets_disjoint() { + let usage = EventTokenUsage { + input_tokens: 100, + output_tokens: 50, + cache_write_tokens: 300, + cache_read_tokens: 200, + cost: EventCost::Omitted, + }; + let cents = estimated_list_price_cents(&usage, "claude-4.5-sonnet", 1_700_000_000_000) + .expect("known Claude alias price"); + let expected = CostUsagePricing::claude_cost_usd("claude-sonnet-4-5", 100, 200, 300, 50) + .unwrap() + * 100.0; + assert!((cents - expected).abs() < 1e-12); + } + + #[test] + fn invalid_total_cents_is_never_reestimated() { + let page: UsageEventsPage = serde_json::from_str( + r#"{"totalUsageEventsCount":1,"usageEventsDisplay":[{"timestamp":"1700000000000","model":"gpt-5","chargedCents":5,"tokenUsage":{"inputTokens":200,"outputTokens":20,"totalCents":"NaN"}}]}"#, + ) + .unwrap(); + let report = summarize_events(&page.usage_events_display); + assert_eq!(report.api_rate_usd, 0.0); + assert_eq!(report.estimated_requests, 0); + assert_eq!(report.unpriced_requests, 1); + assert_eq!(report.metered_usd, Some(0.05)); + } } diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 3134876379..bce4212a59 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -127,21 +127,10 @@ impl GrokProvider { let billing = self .fetch_billing(None, Some(cookie_header.to_string())) .await?; - // Upstream 0.52 (#2991): the browser billing response does not carry - // the paid SuperGrok tier. If the local Grok principal is available, - // use its settings endpoint only as identity enrichment, never as a - // replacement for the validated browser usage result. - let plan = match Self::load_credentials(GrokAuthKind::Cli) { - Ok(credentials) => self.fetch_cli_subscription_tier(&credentials).await, - Err(_) => None, - }; - Ok(result_from_billing( - billing, - "grok-browser", - None, - None, - plan, - )) + // v0.56.0: a browser session is its own principal. Never enrich a + // successful cookie billing result from ambient auth.json metadata, + // which may belong to a different account or change during the fetch. + Ok(result_from_cookie_billing(billing)) } /// Cookie refresh path (upstream #2458): @@ -455,6 +444,9 @@ fn primary_label_for_cycle_minutes(minutes: u32) -> Option<&'static str> { } } +fn result_from_cookie_billing(billing: GrokBillingSnapshot) -> ProviderFetchResult { + result_from_billing(billing, "grok-browser", None, None, None) +} fn result_from_billing( billing: GrokBillingSnapshot, source_label: &str, @@ -850,6 +842,18 @@ mod tests { 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(); diff --git a/rust/src/providers/openrouter/mod.rs b/rust/src/providers/openrouter/mod.rs index 14709894be..0691fecd9a 100755 --- a/rust/src/providers/openrouter/mod.rs +++ b/rust/src/providers/openrouter/mod.rs @@ -117,7 +117,7 @@ impl OpenRouterProvider { id: ProviderId::OpenRouter, display_name: "OpenRouter", session_label: "Credits", - weekly_label: "Usage", + weekly_label: "API key limit", supports_opus: false, supports_credits: true, default_enabled: false, @@ -375,7 +375,9 @@ impl OpenRouterProvider { let key_percent = ((used / limit) * 100.0).clamp(0.0, 100.0); let mut key_window = RateWindow::new(key_percent); - key_window.reset_description = Some(format!("${used:.2}/${limit:.2} key quota")); + key_window.reset_description = Some(format!( + "${used:.2}/${limit:.2} spending cap · Spending cap, not balance" + )); *usage = usage.clone().with_secondary(key_window); } @@ -501,6 +503,36 @@ mod tests { usage.secondary.map(|window| window.used_percent) } + #[test] + fn key_limit_copy_stays_distinct_from_account_balance() { + let provider = OpenRouterProvider::new(); + assert_eq!(provider.metadata.weekly_label, "API key limit"); + + let credits = CreditsData { + total_credits: 5.0, + total_usage: 3.1, + }; + let mut usage = OpenRouterProvider::build_credits_usage(&credits); + OpenRouterProvider::add_key_quota( + &mut usage, + &key_data( + Some(30.0), + Some(30.0), + Some("monthly"), + Some(0.0), + None, + None, + Some(0.0), + ), + ); + assert_eq!(usage.login_method.as_deref(), Some("$1.90 balance")); + let key = usage.secondary.expect("key spending cap"); + assert_eq!(key.used_percent, 0.0); + assert_eq!( + key.reset_description.as_deref(), + Some("$0.00/$30.00 spending cap · Spending cap, not balance") + ); + } #[test] fn server_remaining_replaces_lifetime_usage_for_meter() { // limit 50, server says 12.50 left this period → 75% used, even though diff --git a/rust/src/spend_contract/opencodex/cache.rs b/rust/src/spend_contract/opencodex/cache.rs index ded1c07266..81c23b6743 100644 --- a/rust/src/spend_contract/opencodex/cache.rs +++ b/rust/src/spend_contract/opencodex/cache.rs @@ -55,8 +55,11 @@ pub(super) fn load_entries_with_cache( source_path: &Path, cache_path: &Path, ) -> Option> { + let mut latest_visible = None; for _ in 0..2 { - let identity = log_identity(source_path)?; + let Some(identity) = log_identity(source_path) else { + return latest_visible; + }; let state = read_cache(cache_path); if let Some(state) = state.as_ref() && cursor_matches_source(&state.cursor, &identity, source_path) @@ -77,6 +80,10 @@ pub(super) fn load_entries_with_cache( parsed_offset: parsed.next_offset, prefix_digest: prefix_digest(source_path, parsed.next_offset)?, }; + if !source_matches_snapshot(source_path, &identity) { + continue; + } + latest_visible = Some(visible.clone()); match write_incremental_cache( cache_path, &state.cursor, @@ -99,14 +106,13 @@ pub(super) fn load_entries_with_cache( parsed_offset: parsed.next_offset, prefix_digest: prefix_digest(source_path, parsed.next_offset)?, }; - let current = log_identity(source_path)?; - if current.file_identity != identity.file_identity || current.size < identity.size { + if !source_matches_snapshot(source_path, &identity) { continue; } write_full_cache(cache_path, &cursor, &parsed.committed); return Some(visible); } - None + latest_visible } fn parse_segment( @@ -155,6 +161,12 @@ fn parse_segment( }) } +fn source_matches_snapshot(source_path: &Path, identity: &LogIdentity) -> bool { + log_identity(source_path).is_some_and(|current| { + current.file_identity == identity.file_identity && current.size >= identity.size + }) +} + fn cursor_matches_source(cursor: &ParseCursor, identity: &LogIdentity, source_path: &Path) -> bool { cursor.source_path == identity.source_path && cursor.file_identity == identity.file_identity @@ -179,7 +191,7 @@ fn dedup_entries(entries: Vec) -> Vec { fn cache_path() -> Option { dirs::cache_dir().map(|root| { - root.join("openCodexBar") + root.join("CodexBar") .join("opencodex") .join("usage-cache-v2.sqlite") }) @@ -402,6 +414,10 @@ mod tests { identity_before.file_identity, identity_after.file_identity, "re-created file must get a new file identity so the cache cursor is invalidated" ); + assert!( + !source_matches_snapshot(&log_path, &identity_before), + "a replacement must be rejected before incremental cache commit" + ); assert!(identity_after.size >= identity_before.size); } }