From 54b69cba351bb6c0bba75c44805a55db1a0de628 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:25:26 +0700 Subject: [PATCH 1/6] Estimate Antigravity local history costs --- .../src-tauri/src/commands/usage_spend.rs | 12 ++- rust/src/cli/cost.rs | 27 ++++++- rust/src/cli/serve/data.rs | 2 + .../providers/antigravity/local_sessions.rs | 80 ++++++++++++++++++- .../src/providers/antigravity/local_sqlite.rs | 41 ++++++++++ rust/src/providers/muse/local_usage/mod.rs | 1 + rust/src/spend_contract.rs | 16 +++- 7 files changed, 172 insertions(+), 7 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs index 01f562e111..f2dab7fb08 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -590,12 +590,20 @@ fn build_usage_spend_summary( 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 = seven.estimated_cost_usd; + spend.thirty_day = thirty.estimated_cost_usd; 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(); + if thirty.estimated_cost_usd.is_some() { + spend.source = if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { + "local Antigravity history · API list-price estimate".to_string() + } else { + "partial local Antigravity history · API list-price estimate".to_string() + }; + } else if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { + spend.source = "local Antigravity history · unpriced".to_string(); } spend } diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 8e9be16dc2..86dd742b05 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -386,7 +386,11 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - println!(" Local token history; dollar costs unavailable"); + if let Some(cost) = history.estimated_cost_usd { + println!(" API list-price estimate: ${cost:.2} (not billed spend)"); + } else { + println!(" Local token history; dollar costs unavailable"); + } } fn print_codex_session_output(result: &CostResult, days: u32) { @@ -625,6 +629,7 @@ mod tests { total_tokens: 12_345, session_count: 2, coverage: LocalHistoryCoverage::Complete, + estimated_cost_usd: None, }, 30, ); @@ -639,6 +644,7 @@ mod tests { total_tokens: 999, session_count: 1, coverage: LocalHistoryCoverage::Partial, + estimated_cost_usd: None, }, 30, ); @@ -646,6 +652,25 @@ mod tests { assert!(partial["tokens"]["total"].is_null()); assert_eq!(partial["historyCoverage"], "partial"); } + + #[test] + fn antigravity_json_labels_public_price_estimates() { + use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + LocalTokenHistorySummary { + total_tokens: 1_000, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + estimated_cost_usd: Some(0.0125), + }, + 30, + ); + assert_eq!(payload["cost"]["total_usd"], 0.0125); + assert_eq!(payload["cost"]["currency"], "USD"); + assert!(payload["note"].as_str().unwrap().contains("not billed")); + } + #[test] fn provider_native_only_flag_default_false() { // Default CostArgs has provider_native_only = false (backward compat). diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index d1841db6bc..b615c69150 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -164,6 +164,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, + estimated_cost_usd: None, }, 30, ); @@ -177,6 +178,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, + estimated_cost_usd: None, }, 30, ); diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 458bb97f69..ff1bc5daed 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -6,6 +6,8 @@ use std::path::{Path, PathBuf}; use chrono::{DateTime, Duration, Local, TimeZone, Utc}; use serde_json::Value; +use crate::core::CostUsagePricing; + 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; @@ -143,6 +145,7 @@ fn summarize_paths( 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 estimated_cost_usd = None; let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); let mut complete = !truncated; @@ -163,6 +166,7 @@ fn summarize_paths( let mut reader = BufReader::new(file); let mut remaining = MAX_SESSION_FILE_BYTES; let mut path_had_usage = false; + let mut model = None::; loop { let line = match read_bounded_jsonl_line(&mut reader, &mut remaining) { Ok(Some(line)) => line, @@ -179,6 +183,16 @@ fn summarize_paths( continue; }; let kind = value.get("type").and_then(Value::as_str); + if kind == Some("session_meta") { + model = value + .get("modelId") + .or_else(|| value.get("model_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + continue; + } if kind != Some("usage") && value.get("input").is_none() { continue; } @@ -208,14 +222,28 @@ fn summarize_paths( let output = token_field(&value, &["output"]); let cache_read = token_field(&value, &["cacheRead", "cache_read"]); let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); + let reasoning = token_field( + &value, + &["reasoning", "reasoningTokens", "reasoning_tokens"], + ); let total = input .saturating_add(output) .saturating_add(cache_read) - .saturating_add(cache_write); + .saturating_add(cache_write) + .saturating_add(reasoning); if total == 0 { continue; } total_tokens = total_tokens.saturating_add(total); + if let Some(cost) = estimate_cost_usd( + model.as_deref(), + input, + cache_read, + cache_write, + output.saturating_add(reasoning), + ) { + estimated_cost_usd = checked_cost_sum(estimated_cost_usd, cost); + } path_had_usage = true; } if path_had_usage { @@ -233,9 +261,40 @@ fn summarize_paths( } else { LocalHistoryCoverage::Partial }, + estimated_cost_usd, } } +pub(super) fn estimate_cost_usd( + model: Option<&str>, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> Option { + let model = model.map(str::trim).filter(|value| !value.is_empty())?; + let input = i32::try_from(input).ok()?; + let cache_read = i32::try_from(cache_read).ok()?; + let cache_write = i32::try_from(cache_write).ok()?; + let output = i32::try_from(output).ok()?; + let resolve = |candidate: &str| { + CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) + .filter(|cost| cost.is_finite() && *cost >= 0.0) + }; + resolve(model).or_else(|| { + ["-tiered", "-low", "-thinking"] + .iter() + .find_map(|suffix| model.strip_suffix(suffix)) + .filter(|base| !base.is_empty()) + .and_then(resolve) + }) +} + +pub(super) fn checked_cost_sum(current: Option, cost: f64) -> Option { + let next = current.unwrap_or(0.0) + cost; + next.is_finite().then_some(next) +} + fn read_bounded_jsonl_line( reader: &mut R, remaining_file_bytes: &mut usize, @@ -290,6 +349,25 @@ fn token_field(value: &Value, keys: &[&str]) -> u64 { #[cfg(test)] mod tests { use super::*; + + #[test] + fn prices_known_models_and_provider_local_routing_variants() { + let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) + .expect("known public price"); + let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) + .expect("routing suffix uses the base public price"); + assert!(direct > 0.0); + assert_eq!(direct, routed); + } + + #[test] + fn unknown_or_oversized_pricing_inputs_fail_closed() { + assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); + assert_eq!( + estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), + None + ); + } use rusqlite::Connection; #[test] diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index b050660411..ea8b3d3dd5 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -184,10 +184,26 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } let mut total_tokens = 0_u64; + let mut estimated_cost_usd = None; let mut sessions = HashSet::new(); let mut rows: HashMap<(String, i64), Event> = HashMap::new(); let mut responses: HashMap<(String, String), Event> = HashMap::new(); + let mut label_models = HashMap::<(String, String), String>::new(); + let mut conflicting_labels = HashSet::<(String, String)>::new(); + for event in &events { + let (Some(label), Some(model)) = (event.turn.label.as_ref(), event.turn.model.as_ref()) + else { + continue; + }; + let key = (event.session.clone(), label.clone()); + if label_models.get(&key).is_some_and(|prior| prior != model) { + conflicting_labels.insert(key); + } else { + label_models.insert(key, model.clone()); + } + } + for event in events { let row_key = (event.session.clone(), event.row); if let Some(prior) = rows.get(&row_key) { @@ -234,6 +250,30 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } + if let Some(usage) = event.turn.usage.as_ref() { + let inherited_model = event.turn.label.as_ref().and_then(|label| { + let key = (event.session.clone(), label.clone()); + (!conflicting_labels.contains(&key)) + .then(|| label_models.get(&key)) + .flatten() + .map(String::as_str) + }); + let model = event.turn.model.as_deref().or(inherited_model); + let input = usage.system_prompt.checked_add(usage.new_input); + let output = usage.output.checked_add(usage.reasoning); + if let (Some(input), Some(output)) = (input, output) + && let Some(cost) = super::local_sessions::estimate_cost_usd( + model, + input, + usage.cache_read, + 0, + output, + ) + { + estimated_cost_usd = + super::local_sessions::checked_cost_sum(estimated_cost_usd, cost); + } + } sessions.insert(event.session); } @@ -245,6 +285,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } else { LocalHistoryCoverage::Partial }, + estimated_cost_usd, }) } diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 49a2f7ec34..7ee5b3f093 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -64,6 +64,7 @@ impl From for crate::spend_contract::LocalTokenHistorySummary { total_tokens: report.total_tokens.unwrap_or(0), session_count: report.session_count, coverage: report.coverage, + estimated_cost_usd: None, } } } diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index af9cdfef34..f43bcd0529 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -79,11 +79,14 @@ pub enum LocalHistoryCoverage { Unavailable, } -#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, Default, PartialEq)] pub struct LocalTokenHistorySummary { pub total_tokens: u64, pub session_count: usize, pub coverage: LocalHistoryCoverage, + /// Known subtotal priced at public API list rates. `None` means that no + /// local request in the selected window had a resolvable model price. + pub estimated_cost_usd: Option, } pub fn local_token_history_json( @@ -96,7 +99,10 @@ pub fn local_token_history_json( "provider": provider, "supported": true, "days_scanned": days, - "cost": {"total_usd": serde_json::Value::Null, "currency": serde_json::Value::Null}, + "cost": { + "total_usd": history.estimated_cost_usd, + "currency": history.estimated_cost_usd.map(|_| "USD") + }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, "sessions_count": complete.then_some(history.session_count), @@ -106,7 +112,11 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": "Local token history; dollar costs unavailable" + "note": if history.estimated_cost_usd.is_some() { + "Local token history estimated at public API list prices; not billed spend" + } else { + "Local token history; dollar costs unavailable" + } }) } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] From 7bb133aa070e60296554a2951fdc701530e3bc7a Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:07:37 +0700 Subject: [PATCH 2/6] Track local pricing coverage --- .../src-tauri/src/commands/usage_spend.rs | 27 +++++++--- .../src/lib/usageSpendSharing.test.ts | 5 ++ .../src/lib/usageSpendSharing.ts | 23 +++++++-- .../src/surfaces/TrayPanel.test.tsx | 16 +++++- .../surfaces/settings/tabs/UsageSpendTab.tsx | 20 +++++++- apps/desktop-tauri/src/types/bridge.ts | 7 +++ rust/src/cli/cost.rs | 49 ++++++++++++++++-- rust/src/cli/serve/data.rs | 4 +- .../providers/antigravity/local_sessions.rs | 46 +++++++++++++---- .../src/providers/antigravity/local_sqlite.rs | 24 ++++----- rust/src/providers/muse/local_usage/mod.rs | 2 +- rust/src/spend_contract.rs | 51 ++++++++++++++++--- 12 files changed, 220 insertions(+), 54 deletions(-) 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 f2dab7fb08..49d84b7b34 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -26,6 +26,10 @@ pub struct UsageSpendRow { pub display_name: String, pub seven_day: Option, pub thirty_day: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub seven_day_estimate: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub thirty_day_estimate: Option, pub seven_day_tokens: Option, pub thirty_day_tokens: Option, pub currency: String, @@ -496,6 +500,7 @@ fn build_usage_spend_summary( }) .unwrap_or_else(|| provider_id.clone()); + let mut local_cost_estimates = None; let spend = match provider_id.as_str() { "codex" => SpendValues { seven_day: codex_7_contract.known_cost_usd, @@ -590,21 +595,22 @@ fn build_usage_spend_summary( 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 = seven.estimated_cost_usd; - spend.thirty_day = thirty.estimated_cost_usd; + spend.seven_day = seven.cost_estimate.total_usd(); + spend.thirty_day = thirty.cost_estimate.total_usd(); 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 thirty.estimated_cost_usd.is_some() { - spend.source = if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { - "local Antigravity history · API list-price estimate".to_string() - } else { - "partial local Antigravity history · API list-price estimate".to_string() - }; + if thirty.cost_estimate.total_usd().is_some() { + spend.source = + "local Antigravity history · API list-price estimate".to_string(); + } else if thirty.cost_estimate.known_subtotal_usd.is_some() { + spend.source = + "local Antigravity history · known API list-price subtotal".to_string(); } else if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { spend.source = "local Antigravity history · unpriced".to_string(); } + local_cost_estimates = Some((seven.cost_estimate, thirty.cost_estimate)); spend } _ => cached_spend(cached_snapshot), @@ -626,11 +632,16 @@ fn build_usage_spend_summary( .collect() }) .unwrap_or_default(); + let (seven_day_estimate, thirty_day_estimate) = local_cost_estimates + .map(|(seven, thirty)| (Some(seven), Some(thirty))) + .unwrap_or((None, None)); rows.push(UsageSpendRow { provider_id: provider_id.clone(), display_name, seven_day: spend.seven_day, thirty_day: spend.thirty_day, + seven_day_estimate, + thirty_day_estimate, seven_day_tokens: spend.seven_day_tokens, thirty_day_tokens: spend.thirty_day_tokens, currency, diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index bf693ce7b8..d6748ff079 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from "vitest"; import { formatUsageSpendReportingDay, + formatSpendMetric, filterUsageSpendSummaryForOverview, renderUsageSpendSharePng, usageSpendShareFooter, @@ -10,6 +11,10 @@ import { import type { SpendContract, UsageSpendRow, UsageSpendSummary } from "../types/bridge"; describe("usage spend sharing", () => { + it("labels a mixed-pricing subtotal without presenting it as a total", () => { + expect(formatSpendMetric(null, 1_500, "USD", "tokens", 0.0125)).toMatch(/^≥.* known/); + }); + it.each([ [0, "0 subscriptions"], [1, "1 subscription"], diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.ts index 2246e98505..4045eb9199 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.ts @@ -130,9 +130,14 @@ export function formatSpendMetric( tokens: number | null | undefined, currency: string, tokenLabel: string, + knownSubtotal?: number | null, ): string { const parts: string[] = []; - if (cost != null && Number.isFinite(cost)) parts.push(formatUsd(cost, currency)); + if (cost != null && Number.isFinite(cost)) { + parts.push(formatUsd(cost, currency)); + } else if (knownSubtotal != null && Number.isFinite(knownSubtotal)) { + parts.push(`≥${formatUsd(knownSubtotal, currency)} known`); + } if (tokens != null && Number.isFinite(tokens)) { parts.push(`${Math.max(0, tokens).toLocaleString()} ${tokenLabel}`); } @@ -202,8 +207,20 @@ export function renderUsageSpendSharePng(summary: UsageSpendSummary, title: stri const y = y0 + (index + 1) * rowH; const cells = [ row.displayName, - formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, "tokens"), - formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, "tokens"), + formatSpendMetric( + row.sevenDay, + row.sevenDayTokens, + row.currency, + "tokens", + row.sevenDayEstimate?.knownSubtotalUsd, + ), + formatSpendMetric( + row.thirtyDay, + row.thirtyDayTokens, + row.currency, + "tokens", + row.thirtyDayEstimate?.knownSubtotalUsd, + ), row.currency || "USD", row.source, ]; diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index 09a5b01e1c..ecefece9f2 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -339,13 +339,27 @@ describe("TrayPanel provider grid", () => { source: "hidden", includedInOverview: false, }, + { + providerId: "antigravity", + displayName: "Antigravity", + sevenDay: null, + thirtyDay: null, + thirtyDayEstimate: { + knownSubtotalUsd: 9, + coverage: { priced: 0, unpriced: 1, unmetered: 0, estimated: 1 }, + }, + currency: "USD", + source: "known subtotal", + includedInOverview: true, + }, ], }); renderTrayPanel([provider("codex", "Codex", 35)]); expect(await screen.findByRole("button", { name: "UsageSpendShare" })).toBeInTheDocument(); - expect(screen.getByText(/1 of 1 OverviewSpendProviderCoverage/)).toBeInTheDocument(); + expect(screen.getByText("$2.00")).toBeInTheDocument(); + expect(screen.getByText(/1 of 2 OverviewSpendProviderCoverage/)).toBeInTheDocument(); }); it("dismisses the tray panel on unmodified Escape", async () => { diff --git a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx index 0f7cc7ee89..e7ee282ece 100644 --- a/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/tabs/UsageSpendTab.tsx @@ -234,8 +234,24 @@ export default function UsageSpendTab(_props: TabProps) { {(summary?.rows ?? []).map((row) => ( {row.displayName} - {formatSpendMetric(row.sevenDay, row.sevenDayTokens, row.currency, t("UsageSpendTokens"))} - {formatSpendMetric(row.thirtyDay, row.thirtyDayTokens, row.currency, t("UsageSpendTokens"))} + + {formatSpendMetric( + row.sevenDay, + row.sevenDayTokens, + row.currency, + t("UsageSpendTokens"), + row.sevenDayEstimate?.knownSubtotalUsd, + )} + + + {formatSpendMetric( + row.thirtyDay, + row.thirtyDayTokens, + row.currency, + t("UsageSpendTokens"), + row.thirtyDayEstimate?.knownSubtotalUsd, + )} + {row.currency || "USD"} {row.source} diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 5aa0a40f27..d581262477 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -379,6 +379,8 @@ export interface UsageSpendRow { displayName: string; sevenDay: number | null; thirtyDay: number | null; + sevenDayEstimate?: LocalCostEstimate; + thirtyDayEstimate?: LocalCostEstimate; sevenDayTokens?: number | null; thirtyDayTokens?: number | null; currency: string; @@ -391,6 +393,11 @@ export interface UsageSpendRow { staleUpdatedAt?: string; } +export interface LocalCostEstimate { + knownSubtotalUsd: number | null; + coverage: CostCoverageCounts; +} + export interface UsageSpendSummary { rows: UsageSpendRow[]; contract: SpendContract; diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 86dd742b05..30eba40a42 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -386,8 +386,13 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - if let Some(cost) = history.estimated_cost_usd { + if let Some(cost) = history.cost_estimate.total_usd() { println!(" API list-price estimate: ${cost:.2} (not billed spend)"); + } else if let Some(cost) = history.cost_estimate.known_subtotal_usd { + println!( + " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", + history.cost_estimate.coverage.unpriced + ); } else { println!(" Local token history; dollar costs unavailable"); } @@ -629,7 +634,7 @@ mod tests { total_tokens: 12_345, session_count: 2, coverage: LocalHistoryCoverage::Complete, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); @@ -644,7 +649,7 @@ mod tests { total_tokens: 999, session_count: 1, coverage: LocalHistoryCoverage::Partial, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); @@ -662,15 +667,51 @@ mod tests { total_tokens: 1_000, session_count: 1, coverage: LocalHistoryCoverage::Complete, - estimated_cost_usd: Some(0.0125), + cost_estimate: crate::spend_contract::LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: crate::spend_contract::CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + }, }, 30, ); assert_eq!(payload["cost"]["total_usd"], 0.0125); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); assert_eq!(payload["cost"]["currency"], "USD"); assert!(payload["note"].as_str().unwrap().contains("not billed")); } + #[test] + fn antigravity_json_keeps_mixed_pricing_as_a_known_subtotal() { + use crate::spend_contract::{ + CostCoverageCounts, LocalCostEstimate, LocalHistoryCoverage, LocalTokenHistorySummary, + }; + let payload = crate::spend_contract::local_token_history_json( + "antigravity", + LocalTokenHistorySummary { + total_tokens: 1_500, + session_count: 2, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(0.0125), + coverage: CostCoverageCounts { + estimated: 1, + unpriced: 1, + ..Default::default() + }, + }, + }, + 30, + ); + assert!(payload["cost"]["total_usd"].is_null()); + assert_eq!(payload["cost"]["known_subtotal_usd"], 0.0125); + assert_eq!(payload["cost"]["pricingCoverage"]["estimated"], 1); + assert_eq!(payload["cost"]["pricingCoverage"]["unpriced"], 1); + assert!(payload["note"].as_str().unwrap().contains("subtotal")); + } + #[test] fn provider_native_only_flag_default_false() { // Default CostArgs has provider_native_only = false (backward compat). diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index b615c69150..aa12414d6e 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -164,7 +164,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); @@ -178,7 +178,7 @@ mod tests { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, - estimated_cost_usd: None, + cost_estimate: Default::default(), }, 30, ); diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index ff1bc5daed..3ab0e66e9c 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -145,7 +145,7 @@ fn summarize_paths( 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 estimated_cost_usd = None; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); let mut sessions_with_usage = HashSet::new(); let mut seen_response_ids = HashSet::new(); let mut complete = !truncated; @@ -235,15 +235,13 @@ fn summarize_paths( continue; } total_tokens = total_tokens.saturating_add(total); - if let Some(cost) = estimate_cost_usd( + cost_estimate.record_list_price(estimate_cost_usd( model.as_deref(), input, cache_read, cache_write, output.saturating_add(reasoning), - ) { - estimated_cost_usd = checked_cost_sum(estimated_cost_usd, cost); - } + )); path_had_usage = true; } if path_had_usage { @@ -261,7 +259,7 @@ fn summarize_paths( } else { LocalHistoryCoverage::Partial }, - estimated_cost_usd, + cost_estimate, } } @@ -290,11 +288,6 @@ pub(super) fn estimate_cost_usd( }) } -pub(super) fn checked_cost_sum(current: Option, cost: f64) -> Option { - let next = current.unwrap_or(0.0) + cost; - next.is_finite().then_some(next) -} - fn read_bounded_jsonl_line( reader: &mut R, remaining_file_bytes: &mut usize, @@ -368,6 +361,37 @@ mod tests { None ); } + + #[test] + fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { + let dir = tempfile::tempdir().unwrap(); + let known = dir.path().join("known.jsonl"); + let unknown = dir.path().join("unknown.jsonl"); + fs::write( + &known, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", + "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" + ), + ) + .unwrap(); + fs::write( + &unknown, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", + "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[known, unknown], now, 7, false); + + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 1); + assert!(summary.cost_estimate.known_subtotal_usd.is_some()); + assert_eq!(summary.cost_estimate.total_usd(), None); + } use rusqlite::Connection; #[test] diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index ea8b3d3dd5..8e83eb8e33 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -184,7 +184,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } let mut total_tokens = 0_u64; - let mut estimated_cost_usd = None; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); let mut sessions = HashSet::new(); let mut rows: HashMap<(String, i64), Event> = HashMap::new(); let mut responses: HashMap<(String, String), Event> = HashMap::new(); @@ -250,7 +250,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } - if let Some(usage) = event.turn.usage.as_ref() { + let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { let inherited_model = event.turn.label.as_ref().and_then(|label| { let key = (event.session.clone(), label.clone()); (!conflicting_labels.contains(&key)) @@ -261,19 +261,13 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL let model = event.turn.model.as_deref().or(inherited_model); let input = usage.system_prompt.checked_add(usage.new_input); let output = usage.output.checked_add(usage.reasoning); - if let (Some(input), Some(output)) = (input, output) - && let Some(cost) = super::local_sessions::estimate_cost_usd( - model, - input, - usage.cache_read, - 0, - output, - ) - { - estimated_cost_usd = - super::local_sessions::checked_cost_sum(estimated_cost_usd, cost); + if let (Some(input), Some(output)) = (input, output) { + super::local_sessions::estimate_cost_usd(model, input, usage.cache_read, 0, output) + } else { + None } - } + }); + cost_estimate.record_list_price(estimated_cost); sessions.insert(event.session); } @@ -285,7 +279,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL } else { LocalHistoryCoverage::Partial }, - estimated_cost_usd, + cost_estimate, }) } diff --git a/rust/src/providers/muse/local_usage/mod.rs b/rust/src/providers/muse/local_usage/mod.rs index 7ee5b3f093..24e84b2473 100644 --- a/rust/src/providers/muse/local_usage/mod.rs +++ b/rust/src/providers/muse/local_usage/mod.rs @@ -64,7 +64,7 @@ impl From for crate::spend_contract::LocalTokenHistorySummary { total_tokens: report.total_tokens.unwrap_or(0), session_count: report.session_count, coverage: report.coverage, - estimated_cost_usd: None, + cost_estimate: Default::default(), } } } diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index f43bcd0529..54c7d3490a 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -79,14 +79,45 @@ pub enum LocalHistoryCoverage { Unavailable, } -#[derive(Debug, Clone, Copy, Default, PartialEq)] +#[derive(Debug, Clone, Default, PartialEq)] pub struct LocalTokenHistorySummary { pub total_tokens: u64, pub session_count: usize, pub coverage: LocalHistoryCoverage, - /// Known subtotal priced at public API list rates. `None` means that no - /// local request in the selected window had a resolvable model price. - pub estimated_cost_usd: Option, + pub cost_estimate: LocalCostEstimate, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct LocalCostEstimate { + /// Sum of requests whose models have public API list prices. This remains + /// a subtotal when one or more requests are unpriced. + pub known_subtotal_usd: Option, + pub coverage: CostCoverageCounts, +} + +impl LocalCostEstimate { + pub fn total_usd(&self) -> Option { + if self.coverage.unpriced == 0 && self.coverage.unmetered == 0 { + self.known_subtotal_usd + } else { + None + } + } + + pub(crate) fn record_list_price(&mut self, cost: Option) { + let Some(cost) = cost.filter(|value| value.is_finite() && *value >= 0.0) else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + return; + }; + let next = self.known_subtotal_usd.unwrap_or(0.0) + cost; + if next.is_finite() { + self.known_subtotal_usd = Some(next); + self.coverage.estimated = self.coverage.estimated.saturating_add(1); + } else { + self.coverage.unpriced = self.coverage.unpriced.saturating_add(1); + } + } } pub fn local_token_history_json( @@ -95,13 +126,17 @@ pub fn local_token_history_json( days: u32, ) -> serde_json::Value { let complete = history.coverage == LocalHistoryCoverage::Complete; + let total_usd = history.cost_estimate.total_usd(); + let known_subtotal_usd = history.cost_estimate.known_subtotal_usd; serde_json::json!({ "provider": provider, "supported": true, "days_scanned": days, "cost": { - "total_usd": history.estimated_cost_usd, - "currency": history.estimated_cost_usd.map(|_| "USD") + "total_usd": total_usd, + "known_subtotal_usd": known_subtotal_usd, + "currency": known_subtotal_usd.map(|_| "USD"), + "pricingCoverage": history.cost_estimate.coverage, }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, @@ -112,8 +147,10 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": if history.estimated_cost_usd.is_some() { + "note": if total_usd.is_some() { "Local token history estimated at public API list prices; not billed spend" + } else if known_subtotal_usd.is_some() { + "Known public API list-price subtotal; some local requests are unpriced" } else { "Local token history; dollar costs unavailable" } From 33697f174c69a7148aecd6cbb1d7cd65aefa95df Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 21:47:09 +0700 Subject: [PATCH 3/6] Fix local pricing coverage totals --- .../src-tauri/src/commands/usage_spend.rs | 93 +++++++++++++++---- .../src/lib/usageSpendSharing.test.ts | 8 ++ rust/src/cli/cost.rs | 77 ++++++++++++--- rust/src/cli/serve/data.rs | 11 +-- .../providers/antigravity/local_sessions.rs | 2 +- rust/src/spend_contract.rs | 42 ++++++--- rust/src/spend_contract/tests.rs | 53 +++++++++++ 7 files changed, 238 insertions(+), 48 deletions(-) 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 49d84b7b34..ee7ae528c8 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/usage_spend.rs @@ -591,25 +591,10 @@ 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 = seven.cost_estimate.total_usd(); - spend.thirty_day = thirty.cost_estimate.total_usd(); - 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 thirty.cost_estimate.total_usd().is_some() { - spend.source = - "local Antigravity history · API list-price estimate".to_string(); - } else if thirty.cost_estimate.known_subtotal_usd.is_some() { - spend.source = - "local Antigravity history · known API list-price subtotal".to_string(); - } else if matches!(thirty.coverage, LocalHistoryCoverage::Complete) { - spend.source = "local Antigravity history · unpriced".to_string(); - } + let spend = + antigravity_spend_values(cached_spend(cached_snapshot), &seven, &thirty); local_cost_estimates = Some((seven.cost_estimate, thirty.cost_estimate)); spend } @@ -720,6 +705,29 @@ fn total_token_mix(mix: &codexbar::spend_contract::SpendTokenMix) -> Option saw.then_some(total) } +fn antigravity_spend_values( + mut spend: SpendValues, + seven: &codexbar::spend_contract::LocalTokenHistorySummary, + thirty: &codexbar::spend_contract::LocalTokenHistorySummary, +) -> SpendValues { + use codexbar::spend_contract::LocalHistoryCoverage; + + spend.seven_day = seven.total_usd(); + spend.thirty_day = thirty.total_usd(); + spend.seven_day_tokens = + (seven.coverage == LocalHistoryCoverage::Complete).then_some(seven.total_tokens); + spend.thirty_day_tokens = + (thirty.coverage == LocalHistoryCoverage::Complete).then_some(thirty.total_tokens); + if spend.thirty_day.is_some() { + spend.source = "local Antigravity history · API list-price estimate".to_string(); + } else if thirty.cost_estimate.known_subtotal_usd.is_some() { + spend.source = "local Antigravity history · known API list-price subtotal".to_string(); + } else if thirty.coverage == LocalHistoryCoverage::Complete { + spend.source = "local Antigravity history · unpriced".to_string(); + } + spend +} + fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { let Some(snapshot) = snapshot else { return SpendValues { @@ -797,6 +805,27 @@ fn cached_spend(snapshot: Option<&ProviderUsageSnapshot>) -> SpendValues { mod cache_key_tests { use super::*; + fn local_history( + total_tokens: u64, + coverage: codexbar::spend_contract::LocalHistoryCoverage, + known_subtotal_usd: Option, + unpriced: u32, + ) -> codexbar::spend_contract::LocalTokenHistorySummary { + codexbar::spend_contract::LocalTokenHistorySummary { + total_tokens, + session_count: if total_tokens > 0 { 1 } else { 0 }, + coverage, + cost_estimate: codexbar::spend_contract::LocalCostEstimate { + known_subtotal_usd, + coverage: codexbar::spend_contract::CostCoverageCounts { + estimated: if known_subtotal_usd.is_some() { 1 } else { 0 }, + unpriced, + ..Default::default() + }, + }, + } + } + #[test] fn invalidated_owner_clears_orphaned_indexing_activity() { let mut coordinator = UsageSpendCoordinator::default(); @@ -853,4 +882,34 @@ mod cache_key_tests { assert!(include_in_shared_overview("claude", false, true)); assert!(!include_in_shared_overview("codex", false, false)); } + + #[test] + fn antigravity_partial_history_exposes_only_the_known_subtotal() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let seven = local_history(100, LocalHistoryCoverage::Partial, Some(1.25), 0); + let thirty = local_history(200, LocalHistoryCoverage::Partial, Some(2.50), 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &thirty); + + assert_eq!(spend.seven_day, None); + assert_eq!(spend.thirty_day, None); + assert_eq!(spend.seven_day_tokens, None); + assert_eq!(spend.thirty_day_tokens, None); + assert!(spend.source.contains("known API list-price subtotal")); + } + + #[test] + fn antigravity_complete_empty_history_is_a_known_zero() { + use codexbar::spend_contract::LocalHistoryCoverage; + + let seven = local_history(0, LocalHistoryCoverage::Complete, None, 0); + let thirty = local_history(0, LocalHistoryCoverage::Complete, None, 0); + let spend = antigravity_spend_values(cached_spend(None), &seven, &thirty); + + assert_eq!(spend.seven_day, Some(0.0)); + assert_eq!(spend.thirty_day, Some(0.0)); + assert_eq!(spend.seven_day_tokens, Some(0)); + assert_eq!(spend.thirty_day_tokens, Some(0)); + assert!(spend.source.contains("API list-price estimate")); + } } diff --git a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts index d6748ff079..39c95c9a3f 100644 --- a/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts +++ b/apps/desktop-tauri/src/lib/usageSpendSharing.test.ts @@ -15,6 +15,14 @@ describe("usage spend sharing", () => { expect(formatSpendMetric(null, 1_500, "USD", "tokens", 0.0125)).toMatch(/^≥.* known/); }); + it("renders a complete known-zero total instead of a subtotal", () => { + const metric = formatSpendMetric(0, 0, "USD", "tokens", 9); + expect(metric).not.toBe("—"); + expect(metric).not.toContain("≥"); + expect(metric).not.toContain("9.00"); + expect(metric).toContain("0 tokens"); + }); + it.each([ [0, "0 subscriptions"], [1, "1 subscription"], diff --git a/rust/src/cli/cost.rs b/rust/src/cli/cost.rs index 30eba40a42..819b7abb59 100755 --- a/rust/src/cli/cost.rs +++ b/rust/src/cli/cost.rs @@ -286,7 +286,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b println!("{title}"); } - if let Some(history) = result.token_history { + if let Some(history) = result.token_history.as_ref() { print_local_token_history(history, days); } else if group_by == CostGroupBy::Session && result.provider == "codex" { print_codex_session_output(result, days); @@ -372,7 +372,7 @@ fn print_text_output(results: &[CostResult], use_color: bool, days: u32, group_b } } -fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySummary, days: u32) { +fn print_local_token_history(history: &crate::spend_contract::LocalTokenHistorySummary, days: u32) { use crate::spend_contract::LocalHistoryCoverage; match history.coverage { LocalHistoryCoverage::Complete if history.total_tokens == 0 => { @@ -386,13 +386,17 @@ fn print_local_token_history(history: crate::spend_contract::LocalTokenHistorySu println!(" Local token history is unavailable or incomplete"); } } - if let Some(cost) = history.cost_estimate.total_usd() { + if let Some(cost) = history.total_usd() { println!(" API list-price estimate: ${cost:.2} (not billed spend)"); } else if let Some(cost) = history.cost_estimate.known_subtotal_usd { - println!( - " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", - history.cost_estimate.coverage.unpriced - ); + if history.coverage == LocalHistoryCoverage::Complete { + println!( + " Known API list-price subtotal: ${cost:.2} ({} unpriced requests)", + history.cost_estimate.coverage.unpriced + ); + } else { + println!(" Known API list-price subtotal: ${cost:.2} (history incomplete)"); + } } else { println!(" Local token history; dollar costs unavailable"); } @@ -468,7 +472,7 @@ fn build_json_payloads(results: &[CostResult], days: u32) -> Vec) -> String { let history = crate::providers::antigravity::local_sessions::summarize(30); results.push(crate::spend_contract::local_token_history_json( "antigravity", - history, + &history, 30, )); continue; } if provider_id == ProviderId::Muse { let report = crate::providers::muse::local_usage::scan(30, None); + let history = report.into(); results.push(crate::spend_contract::local_token_history_json( - "muse", - report.into(), - 30, + "muse", &history, 30, )); continue; } @@ -160,7 +159,7 @@ mod tests { use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; let complete = crate::spend_contract::local_token_history_json( "antigravity", - LocalTokenHistorySummary { + &LocalTokenHistorySummary { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Complete, @@ -174,7 +173,7 @@ mod tests { let partial = crate::spend_contract::local_token_history_json( "antigravity", - LocalTokenHistorySummary { + &LocalTokenHistorySummary { total_tokens: 42, session_count: 1, coverage: LocalHistoryCoverage::Partial, diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 3ab0e66e9c..3191e88810 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -390,7 +390,7 @@ mod tests { assert_eq!(summary.cost_estimate.coverage.estimated, 1); assert_eq!(summary.cost_estimate.coverage.unpriced, 1); assert!(summary.cost_estimate.known_subtotal_usd.is_some()); - assert_eq!(summary.cost_estimate.total_usd(), None); + assert_eq!(summary.total_usd(), None); } use rusqlite::Connection; diff --git a/rust/src/spend_contract.rs b/rust/src/spend_contract.rs index 54c7d3490a..722acec630 100644 --- a/rust/src/spend_contract.rs +++ b/rust/src/spend_contract.rs @@ -87,6 +87,21 @@ pub struct LocalTokenHistorySummary { pub cost_estimate: LocalCostEstimate, } +impl LocalTokenHistorySummary { + /// Return a complete list-price total only when both the history scan and + /// pricing coverage are complete. A complete scan with no token usage is + /// a known zero even though there were no requests to price. + pub fn total_usd(&self) -> Option { + if self.coverage != LocalHistoryCoverage::Complete { + return None; + } + if self.total_tokens == 0 { + return Some(0.0); + } + self.cost_estimate.complete_total_usd() + } +} + #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct LocalCostEstimate { @@ -97,7 +112,7 @@ pub struct LocalCostEstimate { } impl LocalCostEstimate { - pub fn total_usd(&self) -> Option { + fn complete_total_usd(&self) -> Option { if self.coverage.unpriced == 0 && self.coverage.unmetered == 0 { self.known_subtotal_usd } else { @@ -122,12 +137,21 @@ impl LocalCostEstimate { pub fn local_token_history_json( provider: &str, - history: LocalTokenHistorySummary, + history: &LocalTokenHistorySummary, days: u32, ) -> serde_json::Value { let complete = history.coverage == LocalHistoryCoverage::Complete; - let total_usd = history.cost_estimate.total_usd(); + let total_usd = history.total_usd(); let known_subtotal_usd = history.cost_estimate.known_subtotal_usd; + let note = if total_usd.is_some() { + "Local token history estimated at public API list prices; not billed spend" + } else if known_subtotal_usd.is_some() && !complete { + "Known public API list-price subtotal; local history is incomplete" + } else if known_subtotal_usd.is_some() { + "Known public API list-price subtotal; some local requests are unpriced" + } else { + "Local token history; dollar costs unavailable" + }; serde_json::json!({ "provider": provider, "supported": true, @@ -135,8 +159,8 @@ pub fn local_token_history_json( "cost": { "total_usd": total_usd, "known_subtotal_usd": known_subtotal_usd, - "currency": known_subtotal_usd.map(|_| "USD"), - "pricingCoverage": history.cost_estimate.coverage, + "currency": total_usd.or(known_subtotal_usd).map(|_| "USD"), + "pricingCoverage": &history.cost_estimate.coverage, }, "daily": [], "tokens": {"total": complete.then_some(history.total_tokens)}, @@ -147,13 +171,7 @@ pub fn local_token_history_json( LocalHistoryCoverage::Unavailable => "unavailable", }, "knownZero": complete && history.total_tokens == 0, - "note": if total_usd.is_some() { - "Local token history estimated at public API list prices; not billed spend" - } else if known_subtotal_usd.is_some() { - "Known public API list-price subtotal; some local requests are unpriced" - } else { - "Local token history; dollar costs unavailable" - } + "note": note, }) } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] diff --git a/rust/src/spend_contract/tests.rs b/rust/src/spend_contract/tests.rs index fb74ff09de..ac32e2e5eb 100644 --- a/rust/src/spend_contract/tests.rs +++ b/rust/src/spend_contract/tests.rs @@ -1,5 +1,58 @@ use super::*; +#[test] +fn local_history_total_requires_complete_scan_and_pricing() { + let priced = LocalCostEstimate { + known_subtotal_usd: Some(1.25), + coverage: CostCoverageCounts { + estimated: 1, + ..Default::default() + }, + }; + let partial_history = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Partial, + cost_estimate: priced.clone(), + }; + assert_eq!(partial_history.total_usd(), None); + assert_eq!(partial_history.cost_estimate.known_subtotal_usd, Some(1.25)); + + let mixed_pricing = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: LocalCostEstimate { + known_subtotal_usd: Some(1.25), + coverage: CostCoverageCounts { + estimated: 1, + unpriced: 1, + ..Default::default() + }, + }, + }; + assert_eq!(mixed_pricing.total_usd(), None); + assert_eq!(mixed_pricing.cost_estimate.known_subtotal_usd, Some(1.25)); + + let complete = LocalTokenHistorySummary { + total_tokens: 100, + session_count: 1, + coverage: LocalHistoryCoverage::Complete, + cost_estimate: priced, + }; + assert_eq!(complete.total_usd(), Some(1.25)); +} + +#[test] +fn complete_empty_local_history_has_a_known_zero_total() { + let history = LocalTokenHistorySummary { + coverage: LocalHistoryCoverage::Complete, + ..Default::default() + }; + + assert_eq!(history.total_usd(), Some(0.0)); +} + #[test] fn coverage_ratio_counts_estimated_as_covered() { let coverage = CostCoverageCounts { From 58e984b2067baf7d8043e8823641f7199f6d487f Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Tue, 22 Sep 2026 22:17:54 +0700 Subject: [PATCH 4/6] Fix partial overview spend expectation --- apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx index ecefece9f2..e36f1e8f1e 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.test.tsx @@ -358,7 +358,7 @@ describe("TrayPanel provider grid", () => { renderTrayPanel([provider("codex", "Codex", 35)]); expect(await screen.findByRole("button", { name: "UsageSpendShare" })).toBeInTheDocument(); - expect(screen.getByText("$2.00")).toBeInTheDocument(); + expect(screen.getByText("~$2.00")).toBeInTheDocument(); expect(screen.getByText(/1 of 2 OverviewSpendProviderCoverage/)).toBeInTheDocument(); }); From 042c88b7dfa26d01f15a96c4cd5876b74a96b842 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 16:41:11 +0700 Subject: [PATCH 5/6] Bound Antigravity local history scans --- rust/src/providers/antigravity/cost.rs | 50 ++ .../providers/antigravity/local_history.rs | 218 +++++++ .../providers/antigravity/local_sessions.rs | 588 +----------------- .../antigravity/local_sessions_reader.rs | 583 +++++++++++++++++ .../src/providers/antigravity/local_sqlite.rs | 91 ++- rust/src/providers/antigravity/mod.rs | 6 +- 6 files changed, 943 insertions(+), 593 deletions(-) create mode 100644 rust/src/providers/antigravity/cost.rs create mode 100644 rust/src/providers/antigravity/local_history.rs create mode 100644 rust/src/providers/antigravity/local_sessions_reader.rs diff --git a/rust/src/providers/antigravity/cost.rs b/rust/src/providers/antigravity/cost.rs new file mode 100644 index 0000000000..f8eb38f24a --- /dev/null +++ b/rust/src/providers/antigravity/cost.rs @@ -0,0 +1,50 @@ +use crate::core::CostUsagePricing; + +pub(super) fn estimate_cost_usd( + model: Option<&str>, + input: u64, + cache_read: u64, + cache_write: u64, + output: u64, +) -> Option { + let model = model.map(str::trim).filter(|value| !value.is_empty())?; + let input = i32::try_from(input).ok()?; + let cache_read = i32::try_from(cache_read).ok()?; + let cache_write = i32::try_from(cache_write).ok()?; + let output = i32::try_from(output).ok()?; + let resolve = |candidate: &str| { + CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) + .filter(|cost| cost.is_finite() && *cost >= 0.0) + }; + resolve(model).or_else(|| { + ["-tiered", "-low", "-thinking"] + .iter() + .find_map(|suffix| model.strip_suffix(suffix)) + .filter(|base| !base.is_empty()) + .and_then(resolve) + }) +} + +#[cfg(test)] +mod tests { + use super::estimate_cost_usd; + + #[test] + fn prices_known_models_and_provider_local_routing_variants() { + let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) + .expect("known public price"); + let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) + .expect("routing suffix uses the base public price"); + assert!(direct > 0.0); + assert_eq!(direct, routed); + } + + #[test] + fn unknown_or_oversized_pricing_inputs_fail_closed() { + assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); + assert_eq!( + estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), + None + ); + } +} diff --git a/rust/src/providers/antigravity/local_history.rs b/rust/src/providers/antigravity/local_history.rs new file mode 100644 index 0000000000..055f9b8d97 --- /dev/null +++ b/rust/src/providers/antigravity/local_history.rs @@ -0,0 +1,218 @@ +use super::{local_sessions_reader as local_sessions, local_sqlite}; +use std::fs; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Utc}; + +use crate::spend_contract::LocalTokenHistorySummary; + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) +} + +fn configured_database_roots(home: &Path) -> [PathBuf; 3] { + let gemini_cli_home = std::env::var("GEMINI_CLI_HOME").ok(); + configured_database_roots_from_values(home, gemini_cli_home.as_deref()) +} + +fn configured_database_roots_from_values( + home: &Path, + gemini_cli_home: Option<&str>, +) -> [PathBuf; 3] { + let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); + local_sqlite::database_roots(&gemini_base) +} + +fn summarize_local_usage_from( + roots: &[PathBuf], + now: DateTime, + days: u32, + jsonl_fallback: impl FnOnce() -> LocalTokenHistorySummary, +) -> LocalTokenHistorySummary { + match local_sqlite::summarize(roots, now, days) { + local_sqlite::SQLiteScan::Summary(summary) => summary, + local_sqlite::SQLiteScan::NoDatabases | local_sqlite::SQLiteScan::Unsupported => { + jsonl_fallback() + } + } +} + +pub fn summarize_local_usage(days: u32) -> LocalTokenHistorySummary { + let now = Utc::now(); + let Some(home) = dirs::home_dir() else { + return LocalTokenHistorySummary::default(); + }; + let roots = configured_database_roots(&home); + let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); + summarize_local_usage_from(&roots, now, days, || { + local_sessions::summarize_jsonl_at(&tokscale_sessions, now, days) + }) +} + +/// Count local Antigravity conversation artifacts for the quota provider's +/// offline fallback. Mirrors upstream #3119 without opening SQLite files. +pub fn offline_conversation_count() -> usize { + let Some(home) = dirs::home_dir() else { + return 0; + }; + let roots = configured_database_roots(&home); + let tokscale_sessions = local_sessions::configured_tokscale_sessions(&home); + offline_conversation_count_with_roots(&roots, &tokscale_sessions) +} + +fn offline_conversation_count_in(home: &Path) -> usize { + let roots = local_sqlite::database_roots(&home.join(".gemini")); + let tokscale_sessions = local_sessions::tokscale_sessions_from_values(home, None); + offline_conversation_count_with_roots(&roots, &tokscale_sessions) +} + +fn offline_conversation_count_with_roots( + database_roots: &[PathBuf], + tokscale_sessions: &Path, +) -> usize { + let db_count = database_roots + .iter() + .map(|root| count_extension(root, "db")) + .sum::(); + if db_count > 0 { + return db_count; + } + local_sessions::count_jsonl_sessions_at(tokscale_sessions) +} + +fn count_extension(root: &Path, extension: &str) -> usize { + fs::read_dir(root) + .ok() + .into_iter() + .flatten() + .flatten() + .map(|entry| entry.path()) + .filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension)) + .count() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::spend_contract::LocalHistoryCoverage; + use chrono::TimeZone; + use rusqlite::Connection; + + #[test] + fn foreign_database_preserves_valid_tokscale_history() { + let dir = tempfile::tempdir().unwrap(); + let gemini_base = dir.path().join(".gemini"); + let database_root = gemini_base.join("antigravity-cli").join("conversations"); + fs::create_dir_all(&database_root).unwrap(); + let connection = Connection::open(database_root.join("foreign.db")).unwrap(); + connection + .execute( + "CREATE TABLE unrelated(id INTEGER PRIMARY KEY, value TEXT)", + [], + ) + .unwrap(); + + let tokscale_sessions = dir + .path() + .join(".config/tokscale/antigravity-cache/sessions"); + fs::create_dir_all(&tokscale_sessions).unwrap(); + let session_path = tokscale_sessions.join("session-a.jsonl"); + fs::write( + &session_path, + b"{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", + ) + .unwrap(); + + let roots = local_sqlite::database_roots(&gemini_base); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_local_usage_from(&roots, now, 7, || { + local_sessions::summarize_jsonl_paths( + std::slice::from_ref(&session_path), + now, + 7, + false, + ) + }); + + assert_eq!(summary.total_tokens, 120); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } + + #[test] + fn foreign_only_input_does_not_fabricate_known_zero_native_usage() { + let dir = tempfile::tempdir().unwrap(); + let gemini_base = dir.path().join(".gemini"); + let database_root = gemini_base.join("antigravity-cli").join("conversations"); + fs::create_dir_all(&database_root).unwrap(); + let connection = Connection::open(database_root.join("foreign.db")).unwrap(); + connection + .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) + .unwrap(); + + let roots = local_sqlite::database_roots(&gemini_base); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_local_usage_from(&roots, now, 7, LocalTokenHistorySummary::default); + + assert_eq!(summary, LocalTokenHistorySummary::default()); + assert_eq!(summary.coverage, LocalHistoryCoverage::Unavailable); + } + + #[test] + fn scan_context_honors_non_empty_root_overrides() { + let home = Path::new(r"C:\Users\test"); + let tokscale_sessions = + local_sessions::tokscale_sessions_from_values(home, Some(r"E:\tokscale-root")); + let roots = configured_database_roots_from_values(home, Some(r"D:\gemini-root")); + assert_eq!( + roots[0], + PathBuf::from(r"D:\gemini-root") + .join("antigravity-cli") + .join("conversations") + ); + assert_eq!( + tokscale_sessions, + PathBuf::from(r"E:\tokscale-root") + .join("antigravity-cache") + .join("sessions") + ); + let defaults = local_sessions::tokscale_sessions_from_values(home, Some("")); + let default_roots = configured_database_roots_from_values(home, Some(" ")); + assert_eq!(default_roots[1], home.join(".gemini").join("antigravity")); + assert_eq!( + defaults, + home.join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions") + ); + } + + #[test] + fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { + let dir = tempfile::tempdir().unwrap(); + let app = dir + .path() + .join(".gemini") + .join("antigravity") + .join("conversations"); + fs::create_dir_all(&app).unwrap(); + fs::write(app.join("a.db"), b"").unwrap(); + fs::write(app.join("a.db-wal"), b"").unwrap(); + assert_eq!(offline_conversation_count_in(dir.path()), 1); + + fs::remove_file(app.join("a.db")).unwrap(); + let cache = dir + .path() + .join(".config") + .join("tokscale") + .join("antigravity-cache") + .join("sessions"); + fs::create_dir_all(&cache).unwrap(); + fs::write(cache.join("one.jsonl"), b"{}\n").unwrap(); + assert_eq!(offline_conversation_count_in(dir.path()), 1); + } +} diff --git a/rust/src/providers/antigravity/local_sessions.rs b/rust/src/providers/antigravity/local_sessions.rs index 3191e88810..4efb9af063 100644 --- a/rust/src/providers/antigravity/local_sessions.rs +++ b/rust/src/providers/antigravity/local_sessions.rs @@ -1,584 +1,4 @@ -use std::collections::HashSet; -use std::fs::{self, File}; -use std::io::{BufRead, BufReader}; -use std::path::{Path, PathBuf}; - -use chrono::{DateTime, Duration, Local, TimeZone, Utc}; -use serde_json::Value; - -use crate::core::CostUsagePricing; - -const MAX_SESSION_FILES: usize = 2048; -const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; -const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; -const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; - -pub use crate::spend_contract::LocalHistoryCoverage; -pub type LocalSessionSummary = crate::spend_contract::LocalTokenHistorySummary; - -#[derive(Debug, Clone, PartialEq, Eq)] -struct ScanContext { - database_roots: [PathBuf; 3], - tokscale_sessions: PathBuf, -} - -impl ScanContext { - fn from_values( - home: &Path, - gemini_cli_home: Option<&str>, - tokscale_config_dir: Option<&str>, - ) -> Self { - let gemini_base = clean_env_path(gemini_cli_home).unwrap_or_else(|| home.join(".gemini")); - let tokscale_base = clean_env_path(tokscale_config_dir) - .unwrap_or_else(|| home.join(".config").join("tokscale")); - Self { - database_roots: super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions: tokscale_base.join("antigravity-cache").join("sessions"), - } - } - - fn capture() -> Option { - let home = dirs::home_dir()?; - let gemini = std::env::var("GEMINI_CLI_HOME").ok(); - let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); - Some(Self::from_values( - &home, - gemini.as_deref(), - tokscale.as_deref(), - )) - } -} - -fn clean_env_path(value: Option<&str>) -> Option { - value - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(PathBuf::from) -} - -pub fn summarize(days: u32) -> LocalSessionSummary { - let now = Utc::now(); - let Some(context) = ScanContext::capture() else { - return LocalSessionSummary::default(); - }; - summarize_context(&context, now, days) -} - -fn summarize_context(context: &ScanContext, now: DateTime, days: u32) -> LocalSessionSummary { - match super::local_sqlite::summarize(&context.database_roots, now, days) { - super::local_sqlite::SQLiteScan::Summary(summary) => summary, - super::local_sqlite::SQLiteScan::NoDatabases - | super::local_sqlite::SQLiteScan::Unsupported => { - let (paths, truncated) = tokscale_paths(&context.tokscale_sessions); - if paths.is_empty() { - LocalSessionSummary::default() - } else { - summarize_paths(&paths, now, days, truncated) - } - } - } -} - -/// Count local Antigravity conversation artifacts for the quota provider's -/// offline fallback. Mirrors upstream #3119 without opening SQLite files. -pub fn offline_conversation_count() -> usize { - let Some(context) = ScanContext::capture() else { - return 0; - }; - offline_conversation_count_context(&context) -} - -fn offline_conversation_count_in(home: &Path) -> usize { - offline_conversation_count_context(&ScanContext::from_values(home, None, None)) -} - -fn offline_conversation_count_context(context: &ScanContext) -> usize { - let db_count = context - .database_roots - .iter() - .map(|root| count_extension(root, "db")) - .sum::(); - if db_count > 0 { - return db_count; - } - tokscale_paths(&context.tokscale_sessions).0.len() -} - -fn count_extension(root: &Path, extension: &str) -> usize { - fs::read_dir(root) - .ok() - .into_iter() - .flatten() - .flatten() - .map(|entry| entry.path()) - .filter(|path| path.extension().and_then(|value| value.to_str()) == Some(extension)) - .count() -} - -fn tokscale_paths(base: &Path) -> (Vec, bool) { - let Ok(entries) = fs::read_dir(base) else { - return (Vec::new(), false); - }; - let mut paths: Vec<_> = entries - .flatten() - .map(|entry| entry.path()) - .filter(|path| { - path.extension() - .and_then(|value| value.to_str()) - .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")) - }) - .collect(); - paths.sort(); - let truncated = paths.len() > MAX_SESSION_FILES; - if truncated { - paths.drain(..paths.len() - MAX_SESSION_FILES); - } - (paths, truncated) -} - -fn summarize_paths( - paths: &[PathBuf], - now: DateTime, - days: u32, - truncated: bool, -) -> LocalSessionSummary { - let first_day = now.with_timezone(&Local).date_naive() - - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); - let mut total_tokens = 0_u64; - let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); - let mut sessions_with_usage = HashSet::new(); - let mut seen_response_ids = HashSet::new(); - let mut complete = !truncated; - - for path in paths.iter().take(MAX_SESSION_FILES) { - let file = match File::open(path) { - Ok(file) => file, - Err(_) => { - complete = false; - continue; - } - }; - match file.metadata() { - Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, - Ok(_) => {} - Err(_) => complete = false, - } - let mut reader = BufReader::new(file); - let mut remaining = MAX_SESSION_FILE_BYTES; - let mut path_had_usage = false; - let mut model = None::; - loop { - let line = match read_bounded_jsonl_line(&mut reader, &mut remaining) { - Ok(Some(line)) => line, - Ok(None) => break, - Err(_) => { - complete = false; - break; - } - }; - if line.is_empty() { - continue; - } - let Ok(value) = serde_json::from_slice::(&line) else { - continue; - }; - let kind = value.get("type").and_then(Value::as_str); - if kind == Some("session_meta") { - model = value - .get("modelId") - .or_else(|| value.get("model_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(str::to_string); - continue; - } - if kind != Some("usage") && value.get("input").is_none() { - continue; - } - if let Some(response_id) = value - .get("responseId") - .or_else(|| value.get("response_id")) - .and_then(Value::as_str) - .map(str::trim) - .filter(|value| !value.is_empty()) - && !seen_response_ids.insert(response_id.to_string()) - { - continue; - } - - let timestamp_ms = value - .get("timestamp") - .and_then(Value::as_i64) - .unwrap_or_default(); - let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { - continue; - }; - if at > now || at.with_timezone(&Local).date_naive() < first_day { - continue; - } - - let input = token_field(&value, &["input"]); - let output = token_field(&value, &["output"]); - let cache_read = token_field(&value, &["cacheRead", "cache_read"]); - let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); - let reasoning = token_field( - &value, - &["reasoning", "reasoningTokens", "reasoning_tokens"], - ); - let total = input - .saturating_add(output) - .saturating_add(cache_read) - .saturating_add(cache_write) - .saturating_add(reasoning); - if total == 0 { - continue; - } - total_tokens = total_tokens.saturating_add(total); - cost_estimate.record_list_price(estimate_cost_usd( - model.as_deref(), - input, - cache_read, - cache_write, - output.saturating_add(reasoning), - )); - path_had_usage = true; - } - if path_had_usage { - sessions_with_usage.insert(path.clone()); - } - } - - LocalSessionSummary { - total_tokens, - session_count: sessions_with_usage.len(), - coverage: if paths.is_empty() { - LocalHistoryCoverage::Unavailable - } else if complete { - LocalHistoryCoverage::Complete - } else { - LocalHistoryCoverage::Partial - }, - cost_estimate, - } -} - -pub(super) fn estimate_cost_usd( - model: Option<&str>, - input: u64, - cache_read: u64, - cache_write: u64, - output: u64, -) -> Option { - let model = model.map(str::trim).filter(|value| !value.is_empty())?; - let input = i32::try_from(input).ok()?; - let cache_read = i32::try_from(cache_read).ok()?; - let cache_write = i32::try_from(cache_write).ok()?; - let output = i32::try_from(output).ok()?; - let resolve = |candidate: &str| { - CostUsagePricing::claude_cost_usd(candidate, input, cache_read, cache_write, output) - .filter(|cost| cost.is_finite() && *cost >= 0.0) - }; - resolve(model).or_else(|| { - ["-tiered", "-low", "-thinking"] - .iter() - .find_map(|suffix| model.strip_suffix(suffix)) - .filter(|base| !base.is_empty()) - .and_then(resolve) - }) -} - -fn read_bounded_jsonl_line( - reader: &mut R, - remaining_file_bytes: &mut usize, -) -> std::io::Result>> { - if *remaining_file_bytes == 0 { - return Ok(None); - } - let mut line = Vec::new(); - let mut saw_input = false; - let mut discarding = false; - - loop { - let chunk = reader.fill_buf()?; - if chunk.is_empty() { - return Ok(saw_input.then_some(if discarding { Vec::new() } else { line })); - } - let bounded_len = chunk.len().min(*remaining_file_bytes); - if bounded_len == 0 { - return Ok(None); - } - let bounded = &chunk[..bounded_len]; - let newline = bounded.iter().position(|byte| *byte == b'\n'); - let segment_end = newline.unwrap_or(bounded.len()); - let segment = &bounded[..segment_end]; - saw_input = saw_input || !segment.is_empty() || newline.is_some(); - if !discarding { - if line.len().saturating_add(segment.len()) <= MAX_JSONL_LINE_BYTES { - line.extend_from_slice(segment); - } else { - line.clear(); - discarding = true; - } - } - let consumed = segment_end + usize::from(newline.is_some()); - reader.consume(consumed); - *remaining_file_bytes = remaining_file_bytes.saturating_sub(consumed); - if newline.is_some() { - return Ok(Some(if discarding { Vec::new() } else { line })); - } - if *remaining_file_bytes == 0 { - return Ok(Some(Vec::new())); - } - } -} - -fn token_field(value: &Value, keys: &[&str]) -> u64 { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_u64)) - .unwrap_or(0) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn prices_known_models_and_provider_local_routing_variants() { - let direct = estimate_cost_usd(Some("claude-sonnet-4-6"), 1_000, 200, 100, 500) - .expect("known public price"); - let routed = estimate_cost_usd(Some("claude-sonnet-4-6-thinking"), 1_000, 200, 100, 500) - .expect("routing suffix uses the base public price"); - assert!(direct > 0.0); - assert_eq!(direct, routed); - } - - #[test] - fn unknown_or_oversized_pricing_inputs_fail_closed() { - assert_eq!(estimate_cost_usd(Some("unknown"), 1, 2, 3, 4), None); - assert_eq!( - estimate_cost_usd(Some("claude-sonnet-4-6"), i32::MAX as u64 + 1, 0, 0, 0), - None - ); - } - - #[test] - fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { - let dir = tempfile::tempdir().unwrap(); - let known = dir.path().join("known.jsonl"); - let unknown = dir.path().join("unknown.jsonl"); - fs::write( - &known, - concat!( - "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", - "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" - ), - ) - .unwrap(); - fs::write( - &unknown, - concat!( - "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", - "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_paths(&[known, unknown], now, 7, false); - - assert_eq!(summary.cost_estimate.coverage.estimated, 1); - assert_eq!(summary.cost_estimate.coverage.unpriced, 1); - assert!(summary.cost_estimate.known_subtotal_usd.is_some()); - assert_eq!(summary.total_usd(), None); - } - use rusqlite::Connection; - - #[test] - fn scan_context_honors_non_empty_root_overrides() { - let home = Path::new(r"C:\Users\test"); - let context = - ScanContext::from_values(home, Some(r"D:\gemini-root"), Some(r"E:\tokscale-root")); - assert_eq!( - context.database_roots[0], - PathBuf::from(r"D:\gemini-root") - .join("antigravity-cli") - .join("conversations") - ); - assert_eq!( - context.tokscale_sessions, - PathBuf::from(r"E:\tokscale-root") - .join("antigravity-cache") - .join("sessions") - ); - let defaults = ScanContext::from_values(home, Some(" "), Some("")); - assert_eq!( - defaults.database_roots[1], - home.join(".gemini").join("antigravity") - ); - assert_eq!( - defaults.tokscale_sessions, - home.join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions") - ); - } - - #[test] - fn foreign_database_preserves_valid_tokscale_history() { - let dir = tempfile::tempdir().unwrap(); - let gemini_base = dir.path().join(".gemini"); - let database_root = gemini_base.join("antigravity-cli").join("conversations"); - fs::create_dir_all(&database_root).unwrap(); - let connection = Connection::open(database_root.join("foreign.db")).unwrap(); - connection - .execute( - "CREATE TABLE unrelated(id INTEGER PRIMARY KEY, value TEXT)", - [], - ) - .unwrap(); - - let tokscale_sessions = dir - .path() - .join(".config/tokscale/antigravity-cache/sessions"); - fs::create_dir_all(&tokscale_sessions).unwrap(); - fs::write( - tokscale_sessions.join("session-a.jsonl"), - b"{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n", - ) - .unwrap(); - - let context = ScanContext { - database_roots: super::super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions, - }; - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_context(&context, now, 7); - - assert_eq!(summary.total_tokens, 120); - assert_eq!(summary.session_count, 1); - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - } - - #[test] - fn foreign_only_input_does_not_fabricate_known_zero_native_usage() { - let dir = tempfile::tempdir().unwrap(); - let gemini_base = dir.path().join(".gemini"); - let database_root = gemini_base.join("antigravity-cli").join("conversations"); - fs::create_dir_all(&database_root).unwrap(); - let connection = Connection::open(database_root.join("foreign.db")).unwrap(); - connection - .execute("CREATE TABLE unrelated(id INTEGER PRIMARY KEY)", []) - .unwrap(); - - let context = ScanContext { - database_roots: super::super::local_sqlite::database_roots(&gemini_base), - tokscale_sessions: dir.path().join("missing-tokscale-sessions"), - }; - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - - let summary = summarize_context(&context, now, 7); - - assert_eq!(summary, LocalSessionSummary::default()); - assert_eq!(summary.coverage, LocalHistoryCoverage::Unavailable); - } - - #[test] - fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write(&path, concat!( - "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", - "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", - "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" - )).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 135); - assert_eq!(summary.session_count, 1); - } - - #[test] - fn truncated_or_unreadable_tokscale_history_is_partial() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); - assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); - - let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); - assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); - } - #[test] - fn offline_count_prefers_cli_and_app_db_artifacts_then_tokscale() { - let dir = tempfile::tempdir().unwrap(); - let app = dir - .path() - .join(".gemini") - .join("antigravity") - .join("conversations"); - fs::create_dir_all(&app).unwrap(); - fs::write(app.join("a.db"), b"").unwrap(); - fs::write(app.join("a.db-wal"), b"").unwrap(); - assert_eq!(offline_conversation_count_in(dir.path()), 1); - - fs::remove_file(app.join("a.db")).unwrap(); - let cache = dir - .path() - .join(".config") - .join("tokscale") - .join("antigravity-cache") - .join("sessions"); - fs::create_dir_all(&cache).unwrap(); - fs::write( - cache.join("one.jsonl"), - b"{} -", - ) - .unwrap(); - assert_eq!(offline_conversation_count_in(dir.path()), 1); - } - - #[test] - fn excludes_usage_outside_requested_window() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - fs::write( - &path, - concat!( - "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", - "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" - ), - ) - .unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 15); - assert_eq!(summary.session_count, 1); - } - - #[test] - fn oversized_jsonl_line_is_discarded_and_next_usage_row_is_counted() { - let dir = tempfile::tempdir().unwrap(); - let path = dir.path().join("session-a.jsonl"); - let mut text = format!( - r#"{{"type":"usage","padding":"{}"}}"#, - "x".repeat(MAX_JSONL_LINE_BYTES + 32) - ); - text.push('\n'); - text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); - text.push('\n'); - fs::write(&path, text).unwrap(); - let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); - let summary = summarize_paths(&[path], now, 7, false); - assert_eq!(summary.total_tokens, 15); - assert_eq!(summary.session_count, 1); - } -} +pub use super::local_history::{offline_conversation_count, summarize_local_usage as summarize}; +pub use crate::spend_contract::{ + LocalHistoryCoverage, LocalTokenHistorySummary as LocalSessionSummary, +}; diff --git a/rust/src/providers/antigravity/local_sessions_reader.rs b/rust/src/providers/antigravity/local_sessions_reader.rs new file mode 100644 index 0000000000..cbc3297cb4 --- /dev/null +++ b/rust/src/providers/antigravity/local_sessions_reader.rs @@ -0,0 +1,583 @@ +use std::cmp::Reverse; +use std::collections::{BinaryHeap, HashSet}; +use std::fs::{self, File}; +use std::io::{BufRead, BufReader}; +use std::path::{Path, PathBuf}; + +use chrono::{DateTime, Duration, Local, TimeZone, Utc}; +use serde_json::Value; + +use super::cost::estimate_cost_usd; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; + +const MAX_SESSION_FILES: usize = 2048; +const MAX_SESSION_DISCOVERY_ENTRIES: usize = 16 * 1024; +const MAX_SESSION_FILE_BYTES: usize = 32 * 1024 * 1024; +const MAX_SESSION_FILE_BYTES_U64: u64 = 32 * 1024 * 1024; +const MAX_TOTAL_SESSION_BYTES: usize = 128 * 1024 * 1024; +const MAX_JSONL_LINE_BYTES: usize = 1024 * 1024; + +enum BoundedJsonlLine { + Record(Vec), + Oversized, + Truncated, +} + +pub(super) fn tokscale_sessions_from_values( + home: &Path, + tokscale_config_dir: Option<&str>, +) -> PathBuf { + let tokscale_base = clean_env_path(tokscale_config_dir) + .unwrap_or_else(|| home.join(".config").join("tokscale")); + tokscale_base.join("antigravity-cache").join("sessions") +} + +pub(super) fn configured_tokscale_sessions(home: &Path) -> PathBuf { + let tokscale = std::env::var("TOKSCALE_CONFIG_DIR").ok(); + tokscale_sessions_from_values(home, tokscale.as_deref()) +} + +fn clean_env_path(value: Option<&str>) -> Option { + value + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(PathBuf::from) +} + +pub(super) fn summarize_jsonl_at( + tokscale_sessions: &Path, + now: DateTime, + days: u32, +) -> LocalTokenHistorySummary { + let (paths, truncated) = tokscale_paths(tokscale_sessions); + summarize_jsonl_paths(&paths, now, days, truncated) +} + +pub(super) fn summarize_jsonl_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalTokenHistorySummary { + if paths.is_empty() { + LocalTokenHistorySummary::default() + } else { + summarize_paths(paths, now, days, truncated) + } +} + +fn tokscale_paths(base: &Path) -> (Vec, bool) { + let Ok(entries) = fs::read_dir(base) else { + return (Vec::new(), false); + }; + bounded_tokscale_paths(entries, MAX_SESSION_DISCOVERY_ENTRIES, MAX_SESSION_FILES) +} + +fn bounded_tokscale_paths( + entries: fs::ReadDir, + max_entries: usize, + max_files: usize, +) -> (Vec, bool) { + let mut paths = BinaryHeap::>::with_capacity(max_files); + let mut truncated = false; + + for (entries_examined, entry) in entries.enumerate() { + if entries_examined == max_entries { + truncated = true; + break; + } + + let Ok(entry) = entry else { + truncated = true; + continue; + }; + let path = entry.path(); + let is_jsonl = path + .extension() + .and_then(|value| value.to_str()) + .is_some_and(|value| value.eq_ignore_ascii_case("jsonl")); + if !is_jsonl { + continue; + } + + if paths.len() < max_files { + paths.push(Reverse(path)); + } else { + truncated = true; + if let Some(Reverse(smallest)) = paths.peek() + && path > *smallest + { + paths.pop(); + paths.push(Reverse(path)); + } + } + } + + let mut paths: Vec<_> = paths.into_iter().map(|Reverse(path)| path).collect(); + paths.sort(); + (paths, truncated) +} + +pub(super) fn count_jsonl_sessions_at(base: &Path) -> usize { + tokscale_paths(base).0.len() +} + +fn summarize_paths( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, +) -> LocalTokenHistorySummary { + summarize_paths_with_budget(paths, now, days, truncated, MAX_TOTAL_SESSION_BYTES) +} + +fn summarize_paths_with_budget( + paths: &[PathBuf], + now: DateTime, + days: u32, + truncated: bool, + total_byte_budget: usize, +) -> LocalTokenHistorySummary { + let first_day = now.with_timezone(&Local).date_naive() + - Duration::days(i64::from(days.clamp(1, 365).saturating_sub(1))); + let mut total_tokens = 0_u64; + let mut cost_estimate = crate::spend_contract::LocalCostEstimate::default(); + let mut sessions_with_usage = HashSet::new(); + let mut seen_response_ids = HashSet::new(); + let mut complete = !truncated; + let mut remaining_total_bytes = total_byte_budget; + + for path in paths.iter().take(MAX_SESSION_FILES) { + if remaining_total_bytes == 0 { + complete = false; + break; + } + let file = match File::open(path) { + Ok(file) => file, + Err(_) => { + complete = false; + continue; + } + }; + match file.metadata() { + Ok(metadata) if metadata.len() > MAX_SESSION_FILE_BYTES_U64 => complete = false, + Ok(_) => {} + Err(_) => complete = false, + } + let mut reader = BufReader::new(file); + let mut remaining = MAX_SESSION_FILE_BYTES; + let mut path_had_usage = false; + let mut model = None::; + loop { + let line = match read_bounded_jsonl_line( + &mut reader, + &mut remaining, + &mut remaining_total_bytes, + ) { + Ok(Some(BoundedJsonlLine::Record(line))) => line, + Ok(Some(BoundedJsonlLine::Oversized)) => { + complete = false; + continue; + } + Ok(Some(BoundedJsonlLine::Truncated)) => { + complete = false; + break; + } + Ok(None) => break, + Err(_) => { + complete = false; + break; + } + }; + if line.is_empty() { + continue; + } + let Ok(value) = serde_json::from_slice::(&line) else { + complete = false; + continue; + }; + if !value.is_object() { + complete = false; + continue; + } + let kind = value.get("type").and_then(Value::as_str); + if kind == Some("session_meta") { + model = value + .get("modelId") + .or_else(|| value.get("model_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + .map(str::to_string); + continue; + } + if kind != Some("usage") && value.get("input").is_none() { + continue; + } + if !has_valid_token_fields(&value) { + complete = false; + continue; + } + + let Some(timestamp_ms) = value.get("timestamp").and_then(Value::as_i64) else { + complete = false; + continue; + }; + let Some(at) = Utc.timestamp_millis_opt(timestamp_ms).single() else { + complete = false; + continue; + }; + + if let Some(response_id) = value + .get("responseId") + .or_else(|| value.get("response_id")) + .and_then(Value::as_str) + .map(str::trim) + .filter(|value| !value.is_empty()) + && !seen_response_ids.insert(response_id.to_string()) + { + continue; + } + + if at > now || at.with_timezone(&Local).date_naive() < first_day { + continue; + } + + let input = token_field(&value, &["input"]); + let output = token_field(&value, &["output"]); + let cache_read = token_field(&value, &["cacheRead", "cache_read"]); + let cache_write = token_field(&value, &["cacheWrite", "cache_write"]); + let reasoning = token_field( + &value, + &["reasoning", "reasoningTokens", "reasoning_tokens"], + ); + let Some(total) = [input, output, cache_read, cache_write, reasoning] + .into_iter() + .try_fold(0_u64, u64::checked_add) + else { + complete = false; + continue; + }; + if total == 0 { + continue; + } + let Some(next_total_tokens) = total_tokens.checked_add(total) else { + complete = false; + continue; + }; + total_tokens = next_total_tokens; + cost_estimate.record_list_price(estimate_cost_usd( + model.as_deref(), + input, + cache_read, + cache_write, + output.saturating_add(reasoning), + )); + path_had_usage = true; + } + if path_had_usage { + sessions_with_usage.insert(path.clone()); + } + } + + LocalTokenHistorySummary { + total_tokens, + session_count: sessions_with_usage.len(), + coverage: if paths.is_empty() { + LocalHistoryCoverage::Unavailable + } else if complete { + LocalHistoryCoverage::Complete + } else { + LocalHistoryCoverage::Partial + }, + cost_estimate, + } +} + +fn read_bounded_jsonl_line( + reader: &mut R, + remaining_file_bytes: &mut usize, + remaining_total_bytes: &mut usize, +) -> std::io::Result> { + if *remaining_file_bytes == 0 { + return Ok(None); + } + if *remaining_total_bytes == 0 { + return Ok(if reader.fill_buf()?.is_empty() { + None + } else { + Some(BoundedJsonlLine::Truncated) + }); + } + let mut line = Vec::new(); + let mut saw_input = false; + let mut discarding = false; + + loop { + let chunk = reader.fill_buf()?; + if chunk.is_empty() { + return Ok(saw_input.then_some(if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + let bounded_len = chunk + .len() + .min(*remaining_file_bytes) + .min(*remaining_total_bytes); + if bounded_len == 0 { + return Ok(None); + } + let bounded = &chunk[..bounded_len]; + let newline = bounded.iter().position(|byte| *byte == b'\n'); + let segment_end = newline.unwrap_or(bounded.len()); + let segment = &bounded[..segment_end]; + saw_input = saw_input || !segment.is_empty() || newline.is_some(); + if !discarding { + if line.len().saturating_add(segment.len()) <= MAX_JSONL_LINE_BYTES { + line.extend_from_slice(segment); + } else { + line.clear(); + discarding = true; + } + } + let consumed = segment_end + usize::from(newline.is_some()); + reader.consume(consumed); + *remaining_file_bytes = remaining_file_bytes.saturating_sub(consumed); + *remaining_total_bytes = remaining_total_bytes.saturating_sub(consumed); + if newline.is_some() { + return Ok(Some(if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + if *remaining_file_bytes == 0 || *remaining_total_bytes == 0 { + let at_eof = reader.fill_buf()?.is_empty(); + return Ok(Some(if !at_eof { + BoundedJsonlLine::Truncated + } else if discarding { + BoundedJsonlLine::Oversized + } else { + BoundedJsonlLine::Record(line) + })); + } + } +} + +fn token_field(value: &Value, keys: &[&str]) -> u64 { + keys.iter() + .find_map(|key| value.get(*key).and_then(Value::as_u64)) + .unwrap_or(0) +} + +fn has_valid_token_fields(value: &Value) -> bool { + let mut has_token_field = false; + for key in [ + "input", + "output", + "cacheRead", + "cache_read", + "cacheWrite", + "cache_write", + "reasoning", + "reasoningTokens", + "reasoning_tokens", + ] { + if let Some(field) = value.get(key) { + if field.as_u64().is_none() { + return false; + } + has_token_field = true; + } + } + has_token_field +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tokscale_discovery_bounds_entries_and_retained_paths() { + let dir = tempfile::tempdir().unwrap(); + for index in 0..5 { + fs::write(dir.path().join(format!("session-{index}.jsonl")), "").unwrap(); + } + + let (paths, truncated) = bounded_tokscale_paths(fs::read_dir(dir.path()).unwrap(), 2, 10); + assert_eq!(paths.len(), 2); + assert!(truncated); + + let (paths, truncated) = bounded_tokscale_paths(fs::read_dir(dir.path()).unwrap(), 10, 2); + assert_eq!(paths.len(), 2); + assert!(truncated); + assert_eq!(paths[0].file_name().unwrap(), "session-3.jsonl"); + assert_eq!(paths[1].file_name().unwrap(), "session-4.jsonl"); + } + + #[test] + fn mixed_known_and_unknown_models_keep_only_a_known_subtotal() { + let dir = tempfile::tempdir().unwrap(); + let known = dir.path().join("known.jsonl"); + let unknown = dir.path().join("unknown.jsonl"); + fs::write( + &known, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"claude-sonnet-4-6\"}\n", + "{\"type\":\"usage\",\"responseId\":\"known\",\"timestamp\":1787572800000,\"input\":1000,\"output\":200}\n" + ), + ) + .unwrap(); + fs::write( + &unknown, + concat!( + "{\"type\":\"session_meta\",\"modelId\":\"future-model\"}\n", + "{\"type\":\"usage\",\"responseId\":\"unknown\",\"timestamp\":1787572800000,\"input\":500,\"output\":100}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[known, unknown], now, 7, false); + + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 1); + assert!(summary.cost_estimate.known_subtotal_usd.is_some()); + assert_eq!(summary.total_usd(), None); + } + #[test] + fn summarizes_tokscale_jsonl_and_deduplicates_response_ids() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write(&path, concat!( + "{\"type\":\"session_meta\",\"modelId\":\"test-model-antigravity-a\"}\n", + "{\"type\":\"usage\",\"responseId\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20,\"cacheRead\":10,\"cacheWrite\":5}\n", + "{\"type\":\"usage\",\"response_id\":\"r1\",\"timestamp\":1787572800000,\"input\":100,\"output\":20}\n" + )).unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 135); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + } + + #[test] + fn truncated_or_unreadable_tokscale_history_is_partial() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + b"{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let truncated = summarize_paths(std::slice::from_ref(&path), now, 7, true); + assert_eq!(truncated.coverage, LocalHistoryCoverage::Partial); + + let missing = summarize_paths(&[dir.path().join("missing.jsonl")], now, 7, false); + assert_eq!(missing.coverage, LocalHistoryCoverage::Partial); + } + #[test] + fn excludes_usage_outside_requested_window() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n", + "{\"type\":\"usage\",\"timestamp\":1784894400000,\"input\":99,\"output\":99}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.session_count, 1); + } + + #[test] + fn oversized_jsonl_line_marks_coverage_partial_and_next_row_is_counted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + let mut text = format!( + r#"{{"type":"usage","padding":"{}"}}"#, + "x".repeat(MAX_JSONL_LINE_BYTES + 32) + ); + text.push('\n'); + text.push_str(r#"{"type":"usage","timestamp":1787572800000,"input":10,"output":5}"#); + text.push('\n'); + fs::write(&path, text).unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + let summary = summarize_paths(&[path], now, 7, false); + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn malformed_jsonl_record_marks_coverage_partial_and_next_row_is_counted() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{malformed json}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":\"invalid\"}\n", + "{\"type\":\"usage\",\"input\":10}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10,\"output\":5}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[path], now, 7, false); + + assert_eq!(summary.total_tokens, 15); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn total_scan_byte_budget_stops_later_records_and_marks_partial() { + let dir = tempfile::tempdir().unwrap(); + let first_path = dir.path().join("session-a.jsonl"); + let second_path = dir.path().join("session-b.jsonl"); + let first = "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":10}\n"; + fs::write(&first_path, first).unwrap(); + fs::write( + &second_path, + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":20}\n", + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = + summarize_paths_with_budget(&[first_path, second_path], now, 7, false, first.len()); + + assert_eq!(summary.total_tokens, 10); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } + + #[test] + fn token_sum_overflow_marks_coverage_partial_without_saturation() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("session-a.jsonl"); + fs::write( + &path, + concat!( + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615,\"output\":1}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":18446744073709551615}\n", + "{\"type\":\"usage\",\"timestamp\":1787572800000,\"input\":1}\n" + ), + ) + .unwrap(); + let now = Utc.timestamp_millis_opt(1787576400000).single().unwrap(); + + let summary = summarize_paths(&[path], now, 7, false); + + assert_eq!(summary.total_tokens, u64::MAX); + assert_eq!(summary.session_count, 1); + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); + } +} diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index 8e83eb8e33..4c5c2e5192 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -10,9 +10,12 @@ use chrono::{DateTime, Duration, Local, TimeZone, Utc}; use rusqlite::{Connection, OpenFlags, TransactionBehavior, types::ValueRef}; use self::local_bot_id::{ExactStepTimestamp, embedded_timestamps_agree, record_exact_bot_id}; +use super::cost::estimate_cost_usd; use super::local_proto::{ParsedTurn, parse_step_metadata, parse_turn}; -use super::local_sessions::{LocalHistoryCoverage, LocalSessionSummary}; use super::local_step_resolver::{StepOccurrence, resolve_step_timestamps}; +#[cfg(test)] +use crate::spend_contract::LocalTokenHistorySummary as LocalSessionSummary; +use crate::spend_contract::{LocalHistoryCoverage, LocalTokenHistorySummary}; const MAX_DATABASES: usize = 500; const MAX_DIRECTORY_ENTRIES: usize = 10_000; @@ -34,7 +37,7 @@ pub(super) enum SQLiteScan { /// This is non-authoritative: callers may continue with another local /// history source instead of treating the scan as known-empty history. Unsupported, - Summary(LocalSessionSummary), + Summary(LocalTokenHistorySummary), } #[derive(Debug)] @@ -262,7 +265,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL let input = usage.system_prompt.checked_add(usage.new_input); let output = usage.output.checked_add(usage.reasoning); if let (Some(input), Some(output)) = (input, output) { - super::local_sessions::estimate_cost_usd(model, input, usage.cache_read, 0, output) + estimate_cost_usd(model, input, usage.cache_read, 0, output) } else { None } @@ -271,7 +274,7 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL sessions.insert(event.session); } - SQLiteScan::Summary(LocalSessionSummary { + SQLiteScan::Summary(LocalTokenHistorySummary { total_tokens, session_count: sessions.len(), coverage: if complete { @@ -389,10 +392,10 @@ fn read_database(path: &Path, budget: &mut Budget) -> rusqlite::Result Vec { + let mut bytes = Vec::new(); + loop { + let mut byte = (value & 0x7f) as u8; + value >>= 7; + if value != 0 { + byte |= 0x80; + } + bytes.push(byte); + if value == 0 { + return bytes; + } + } + } + + fn field_varint(number: u64, value: u64) -> Vec { + let mut bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes + } + + fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes + } + + fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { + let mut usage = field_varint(1, 11); + usage.extend(field_varint(2, input)); + usage.extend(field_varint(5, 50)); + usage.extend(field_varint(9, 30)); + usage.extend(field_varint(10, 7)); + + let mut timestamp = field_varint(1, timestamp_seconds); + timestamp.extend(field_varint(2, 0)); + let mut chat = field_bytes(4, &usage); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + field_bytes(1, &chat) + } + #[test] fn missing_databases_falls_through() { let dir = tempfile::tempdir().unwrap(); @@ -892,6 +937,36 @@ mod tests { assert_eq!(summary.session_count, 0); } + #[test] + fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { + let dir = tempfile::tempdir().unwrap(); + let first_root = dir.path().join("first"); + let second_root = dir.path().join("second"); + fs::create_dir_all(&first_root).unwrap(); + fs::create_dir_all(&second_root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); + + for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { + let conn = Connection::open(root.join("session.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob(input, timestamp)], + ) + .unwrap(); + } + + let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) + else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); + assert_eq!(summary.session_count, 2); + } + #[test] fn non_blob_rows_make_coverage_partial() { let dir = tempfile::tempdir().unwrap(); diff --git a/rust/src/providers/antigravity/mod.rs b/rust/src/providers/antigravity/mod.rs index 5f35764663..d24eb8c637 100755 --- a/rust/src/providers/antigravity/mod.rs +++ b/rust/src/providers/antigravity/mod.rs @@ -4,9 +4,12 @@ //! Uses Windows process detection to find CSRF token mod cli_fallback; +mod cost; mod legacy_status; +mod local_history; mod local_proto; pub mod local_sessions; +mod local_sessions_reader; mod local_sqlite; mod local_step_resolver; mod quota_summary; @@ -50,6 +53,7 @@ const AGY_READY_POLL_INTERVAL: Duration = Duration::from_millis(250); const GET_USER_STATUS_PATH: &str = "/exa.language_server_pb.LanguageServerService/GetUserStatus"; const QUOTA_SUMMARY_PATH: &str = "/exa.language_server_pb.LanguageServerService/RetrieveUserQuotaSummary"; + /// Serialize task-owned `agy` launches so concurrent app surfaces never start /// multiple interactive CLI servers at the same time. #[cfg(windows)] @@ -623,7 +627,7 @@ impl AntigravityProvider { } fn offline_usage_result() -> Option { - let count = local_sessions::offline_conversation_count(); + let count = local_history::offline_conversation_count(); if count == 0 { return None; } From b35ec2941e99f730caa39eb495fd9cbc0e83e4de Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Wed, 23 Sep 2026 21:09:39 +0700 Subject: [PATCH 6/6] Skip pricing zero-token Antigravity events --- .../src/providers/antigravity/local_sqlite.rs | 253 ++-------------- .../antigravity/local_sqlite_tests.rs | 280 ++++++++++++++++++ 2 files changed, 301 insertions(+), 232 deletions(-) create mode 100644 rust/src/providers/antigravity/local_sqlite_tests.rs diff --git a/rust/src/providers/antigravity/local_sqlite.rs b/rust/src/providers/antigravity/local_sqlite.rs index 4c5c2e5192..bc71f794d7 100644 --- a/rust/src/providers/antigravity/local_sqlite.rs +++ b/rust/src/providers/antigravity/local_sqlite.rs @@ -253,24 +253,26 @@ pub(super) fn summarize(roots: &[PathBuf], now: DateTime, days: u32) -> SQL continue; } } - let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { - let inherited_model = event.turn.label.as_ref().and_then(|label| { - let key = (event.session.clone(), label.clone()); - (!conflicting_labels.contains(&key)) - .then(|| label_models.get(&key)) - .flatten() - .map(String::as_str) + if event.total > 0 { + let estimated_cost = event.turn.usage.as_ref().and_then(|usage| { + let inherited_model = event.turn.label.as_ref().and_then(|label| { + let key = (event.session.clone(), label.clone()); + (!conflicting_labels.contains(&key)) + .then(|| label_models.get(&key)) + .flatten() + .map(String::as_str) + }); + let model = event.turn.model.as_deref().or(inherited_model); + let input = usage.system_prompt.checked_add(usage.new_input); + let output = usage.output.checked_add(usage.reasoning); + if let (Some(input), Some(output)) = (input, output) { + estimate_cost_usd(model, input, usage.cache_read, 0, output) + } else { + None + } }); - let model = event.turn.model.as_deref().or(inherited_model); - let input = usage.system_prompt.checked_add(usage.new_input); - let output = usage.output.checked_add(usage.reasoning); - if let (Some(input), Some(output)) = (input, output) { - estimate_cost_usd(model, input, usage.cache_read, 0, output) - } else { - None - } - }); - cost_estimate.record_list_price(estimated_cost); + cost_estimate.record_list_price(estimated_cost); + } sessions.insert(event.session); } @@ -848,218 +850,5 @@ fn has_stored_columns( mod synthetic_tests; #[cfg(test)] -mod tests { - use super::*; - use rusqlite::params; - - 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 bytes = varint(number << 3); - bytes.extend(varint(value)); - bytes - } - - fn field_bytes(number: u64, value: &[u8]) -> Vec { - let mut bytes = varint((number << 3) | 2); - bytes.extend(varint(value.len() as u64)); - bytes.extend(value); - bytes - } - - fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { - let mut usage = field_varint(1, 11); - usage.extend(field_varint(2, input)); - usage.extend(field_varint(5, 50)); - usage.extend(field_varint(9, 30)); - usage.extend(field_varint(10, 7)); - - let mut timestamp = field_varint(1, timestamp_seconds); - timestamp.extend(field_varint(2, 0)); - let mut chat = field_bytes(4, &usage); - chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); - field_bytes(1, &chat) - } - - #[test] - fn missing_databases_falls_through() { - let dir = tempfile::tempdir().unwrap(); - assert!(matches!( - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), - SQLiteScan::NoDatabases - )); - } - - #[test] - fn foreign_database_is_non_authoritative() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) - .unwrap(); - drop(conn); - assert!(matches!( - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), - SQLiteScan::Unsupported - )); - } - - #[test] - fn empty_supported_database_is_confirmed_zero() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 0); - assert_eq!(summary.session_count, 0); - } - - #[test] - fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { - let dir = tempfile::tempdir().unwrap(); - let first_root = dir.path().join("first"); - let second_root = dir.path().join("second"); - fs::create_dir_all(&first_root).unwrap(); - fs::create_dir_all(&second_root).unwrap(); - let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); - - for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { - let conn = Connection::open(root.join("session.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", - [valid_turn_blob(input, timestamp)], - ) - .unwrap(); - } - - let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) - else { - panic!("supported databases should produce coverage"); - }; - - assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); - assert_eq!(summary.total_tokens, 496); - assert_eq!(summary.session_count, 2); - } - - #[test] - fn non_blob_rows_make_coverage_partial() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join(".gemini/antigravity-cli/conversations"); - fs::create_dir_all(&root).unwrap(); - let conn = Connection::open(root.join("one.db")).unwrap(); - conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) - .unwrap(); - conn.execute( - "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", - params![1_i64, "not-a-blob"], - ) - .unwrap(); - drop(conn); - let SQLiteScan::Summary(summary) = - summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) - else { - panic!("supported database should produce coverage"); - }; - assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); - } - #[test] - fn discovery_allows_exactly_500_databases_but_marks_501_partial() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join("dbs"); - fs::create_dir_all(&root).unwrap(); - for index in 0..MAX_DATABASES { - fs::write(root.join(format!("{index:03}.db")), b"").unwrap(); - } - let mut budget = Budget::new(); - let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); - assert_eq!(paths.len(), MAX_DATABASES); - assert!(complete); - - fs::write(root.join("overflow.db"), b"").unwrap(); - let mut budget = Budget::new(); - let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); - assert_eq!(paths.len(), MAX_DATABASES); - assert!(!complete); - } - - #[test] - fn expired_budget_marks_discovery_incomplete() { - let dir = tempfile::tempdir().unwrap(); - let root = dir.path().join("dbs"); - fs::create_dir_all(&root).unwrap(); - fs::write(root.join("one.db"), b"").unwrap(); - let mut budget = Budget::with_deadline(Instant::now()); - let (_, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); - assert!(!complete); - } - - #[test] - fn extra_columns_and_without_rowid_schema_is_supported() { - let conn = Connection::open_in_memory().unwrap(); - conn.execute( - "CREATE TABLE gen_metadata(idx INTEGER PRIMARY KEY, data BLOB, extra TEXT) WITHOUT ROWID", - [], - ) - .unwrap(); - let mut budget = Budget::new(); - assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), - SchemaInspection::Supported - ); - } - - #[test] - fn generated_columns_are_rejected() { - let conn = Connection::open_in_memory().unwrap(); - conn.execute( - "CREATE TABLE gen_metadata(idx INTEGER, data BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL)", - [], - ) - .unwrap(); - let mut budget = Budget::new(); - assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), - SchemaInspection::Unsupported - ); - } - - #[test] - fn schema_entry_budget_is_incomplete_not_foreign() { - let conn = Connection::open_in_memory().unwrap(); - for index in 0..=MAX_SCHEMA_ENTRIES { - conn.execute(&format!("CREATE TABLE unrelated_{index}(value TEXT)"), []) - .unwrap(); - } - let mut budget = Budget::new(); - assert_eq!( - supported_schema(&conn, &mut budget).unwrap(), - SchemaInspection::Incomplete - ); - } -} +#[path = "local_sqlite_tests.rs"] +mod tests; diff --git a/rust/src/providers/antigravity/local_sqlite_tests.rs b/rust/src/providers/antigravity/local_sqlite_tests.rs new file mode 100644 index 0000000000..1846bd5e26 --- /dev/null +++ b/rust/src/providers/antigravity/local_sqlite_tests.rs @@ -0,0 +1,280 @@ +use super::*; +use rusqlite::params; + +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 bytes = varint(number << 3); + bytes.extend(varint(value)); + bytes +} + +fn field_bytes(number: u64, value: &[u8]) -> Vec { + let mut bytes = varint((number << 3) | 2); + bytes.extend(varint(value.len() as u64)); + bytes.extend(value); + bytes +} + +fn valid_turn_blob(input: u64, timestamp_seconds: u64) -> Vec { + valid_turn_blob_with_model(input, timestamp_seconds, None) +} + +fn valid_turn_blob_with_model(input: u64, timestamp_seconds: u64, model: Option<&str>) -> Vec { + let mut usage = field_varint(1, 11); + usage.extend(field_varint(2, input)); + usage.extend(field_varint(5, 50)); + usage.extend(field_varint(9, 30)); + usage.extend(field_varint(10, 7)); + + let mut timestamp = field_varint(1, timestamp_seconds); + timestamp.extend(field_varint(2, 0)); + let mut chat = field_bytes(4, &usage); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + if let Some(model) = model { + chat.extend(field_bytes(19, model.as_bytes())); + } + field_bytes(1, &chat) +} + +fn zero_token_turn_blob(timestamp_seconds: u64, model: &str) -> Vec { + let timestamp = field_varint(1, timestamp_seconds); + let mut chat = field_bytes(4, &[]); + chat.extend(field_bytes(9, &field_bytes(4, ×tamp))); + chat.extend(field_bytes(19, model.as_bytes())); + field_bytes(1, &chat) +} + +#[test] +fn missing_databases_falls_through() { + let dir = tempfile::tempdir().unwrap(); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::NoDatabases + )); +} + +#[test] +fn foreign_database_is_non_authoritative() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE wrong(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + assert!(matches!( + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30), + SQLiteScan::Unsupported + )); +} + +#[test] +fn empty_supported_database_is_confirmed_zero() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 0); + assert_eq!(summary.session_count, 0); +} + +#[test] +fn same_named_databases_in_separate_roots_keep_distinct_rows_and_sessions() { + let dir = tempfile::tempdir().unwrap(); + let first_root = dir.path().join("first"); + let second_root = dir.path().join("second"); + fs::create_dir_all(&first_root).unwrap(); + fs::create_dir_all(&second_root).unwrap(); + let timestamp = u64::try_from(Utc::now().timestamp()).unwrap(); + + for (root, input) in [(&first_root, 100_u64), (&second_root, 200_u64)] { + let conn = Connection::open(root.join("session.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob(input, timestamp)], + ) + .unwrap(); + } + + let SQLiteScan::Summary(summary) = summarize(&[first_root, second_root], Utc::now(), 30) else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 496); + assert_eq!(summary.session_count, 2); +} + +#[test] +fn zero_token_unknown_model_does_not_poison_priced_history() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let now = Utc::now(); + let timestamp = u64::try_from(now.timestamp()).unwrap(); + + let priced = Connection::open(root.join("priced.db")).unwrap(); + priced + .execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + priced + .execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [valid_turn_blob_with_model( + 100, + timestamp, + Some("claude-sonnet-4-6"), + )], + ) + .unwrap(); + + let zero = Connection::open(root.join("zero.db")).unwrap(); + zero.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + zero.execute( + "INSERT INTO gen_metadata(idx, data) VALUES(1, ?1)", + [zero_token_turn_blob(timestamp, "unknown-model")], + ) + .unwrap(); + + let SQLiteScan::Summary(summary) = summarize(&[root], now, 30) else { + panic!("supported databases should produce coverage"); + }; + + assert_eq!(summary.coverage, LocalHistoryCoverage::Complete); + assert_eq!(summary.total_tokens, 198); + assert_eq!(summary.session_count, 2); + assert_eq!(summary.cost_estimate.coverage.estimated, 1); + assert_eq!(summary.cost_estimate.coverage.unpriced, 0); + assert!( + summary + .cost_estimate + .known_subtotal_usd + .is_some_and(|cost| cost > 0.0) + ); + assert_eq!( + summary.total_usd(), + summary.cost_estimate.known_subtotal_usd + ); +} + +#[test] +fn non_blob_rows_make_coverage_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join(".gemini/antigravity-cli/conversations"); + fs::create_dir_all(&root).unwrap(); + let conn = Connection::open(root.join("one.db")).unwrap(); + conn.execute("CREATE TABLE gen_metadata(idx INTEGER, data BLOB)", []) + .unwrap(); + conn.execute( + "INSERT INTO gen_metadata(idx,data) VALUES(?1,?2)", + params![1_i64, "not-a-blob"], + ) + .unwrap(); + drop(conn); + let SQLiteScan::Summary(summary) = + summarize(&database_roots(&dir.path().join(".gemini")), Utc::now(), 30) + else { + panic!("supported database should produce coverage"); + }; + assert_eq!(summary.coverage, LocalHistoryCoverage::Partial); +} +#[test] +fn discovery_allows_exactly_500_databases_but_marks_501_partial() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + for index in 0..MAX_DATABASES { + fs::write(root.join(format!("{index:03}.db")), b"").unwrap(); + } + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(complete); + + fs::write(root.join("overflow.db"), b"").unwrap(); + let mut budget = Budget::new(); + let (paths, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert_eq!(paths.len(), MAX_DATABASES); + assert!(!complete); +} + +#[test] +fn expired_budget_marks_discovery_incomplete() { + let dir = tempfile::tempdir().unwrap(); + let root = dir.path().join("dbs"); + fs::create_dir_all(&root).unwrap(); + fs::write(root.join("one.db"), b"").unwrap(); + let mut budget = Budget::with_deadline(Instant::now()); + let (_, complete) = discover_databases(std::slice::from_ref(&root), &mut budget); + assert!(!complete); +} + +#[test] +fn extra_columns_and_without_rowid_schema_is_supported() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER PRIMARY KEY, data BLOB, extra TEXT) WITHOUT ROWID", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Supported + ); +} + +#[test] +fn generated_columns_are_rejected() { + let conn = Connection::open_in_memory().unwrap(); + conn.execute( + "CREATE TABLE gen_metadata(idx INTEGER, data BLOB, derived TEXT GENERATED ALWAYS AS (idx || 'x') VIRTUAL)", + [], + ) + .unwrap(); + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Unsupported + ); +} + +#[test] +fn schema_entry_budget_is_incomplete_not_foreign() { + let conn = Connection::open_in_memory().unwrap(); + for index in 0..=MAX_SCHEMA_ENTRIES { + conn.execute(&format!("CREATE TABLE unrelated_{index}(value TEXT)"), []) + .unwrap(); + } + let mut budget = Budget::new(); + assert_eq!( + supported_schema(&conn, &mut budget).unwrap(), + SchemaInspection::Incomplete + ); +}