From e6f115a5f4eb7889f2f45e1d6614d6f9ab12549a Mon Sep 17 00:00:00 2001 From: ozymandiashh <234437643+ozymandiashh@users.noreply.github.com> Date: Mon, 22 Jun 2026 05:53:05 +0300 Subject: [PATCH] fix(menubar): clamp unsafe Antigravity token counts --- .../CodeBurnMenubar/Data/MenubarPayload.swift | 69 +++++- .../MenubarPayloadDecodeTests.swift | 126 ++++++++++ src/daily-cache.ts | 84 ++++++- src/menubar-json.ts | 52 +++- src/providers/antigravity.ts | 173 ++++++++++---- src/session-cache.ts | 89 ++++++- tests/daily-cache.test.ts | 69 ++++++ tests/menubar-json.test.ts | 76 ++++++ tests/providers/antigravity.test.ts | 225 +++++++++++++++++- tests/session-cache.test.ts | 159 +++++++++++++ 10 files changed, 1040 insertions(+), 82 deletions(-) create mode 100644 mac/Tests/CodeBurnMenubarTests/MenubarPayloadDecodeTests.swift diff --git a/mac/Sources/CodeBurnMenubar/Data/MenubarPayload.swift b/mac/Sources/CodeBurnMenubar/Data/MenubarPayload.swift index 50507c43..3033258f 100644 --- a/mac/Sources/CodeBurnMenubar/Data/MenubarPayload.swift +++ b/mac/Sources/CodeBurnMenubar/Data/MenubarPayload.swift @@ -1,5 +1,37 @@ import Foundation +private let maxSafeTokenCount = 9_007_199_254_740_991 + +private func sanitizedTokenCount(_ value: Int) -> Int { + value >= 0 && value <= maxSafeTokenCount ? value : 0 +} + +private extension KeyedDecodingContainer { + func decodeTokenCount(forKey key: Key) -> Int { + guard contains(key) else { return 0 } + if let value = try? decode(Int.self, forKey: key) { + return sanitizedTokenCount(value) + } + if let value = try? decode(Double.self, forKey: key), + value.isFinite, + value >= 0, + value <= Double(maxSafeTokenCount), + value.rounded(.towardZero) == value { + return Int(value) + } + if let value = try? decode(String.self, forKey: key) { + let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines) + guard !trimmed.isEmpty, + trimmed.utf8.allSatisfy({ $0 >= 48 && $0 <= 57 }), + let parsed = Int(trimmed) else { + return 0 + } + return sanitizedTokenCount(parsed) + } + return 0 + } +} + /// Shape of `codeburn status --format menubar-json --period `. /// `current` is scoped to the requested period; the whole payload reflects that slice. struct MenubarPayload: Codable, Sendable { @@ -32,8 +64,8 @@ struct DailyModelBreakdown: Codable, Sendable { cost = try c.decode(Double.self, forKey: .cost) savingsUSD = try c.decodeIfPresent(Double.self, forKey: .savingsUSD) ?? 0 calls = try c.decode(Int.self, forKey: .calls) - inputTokens = try c.decode(Int.self, forKey: .inputTokens) - outputTokens = try c.decode(Int.self, forKey: .outputTokens) + inputTokens = c.decodeTokenCount(forKey: .inputTokens) + outputTokens = c.decodeTokenCount(forKey: .outputTokens) } } @@ -66,10 +98,10 @@ extension DailyHistoryEntry { cost = try c.decode(Double.self, forKey: .cost) savingsUSD = try c.decodeIfPresent(Double.self, forKey: .savingsUSD) ?? 0 calls = try c.decode(Int.self, forKey: .calls) - inputTokens = try c.decode(Int.self, forKey: .inputTokens) - outputTokens = try c.decode(Int.self, forKey: .outputTokens) - cacheReadTokens = try c.decode(Int.self, forKey: .cacheReadTokens) - cacheWriteTokens = try c.decode(Int.self, forKey: .cacheWriteTokens) + inputTokens = c.decodeTokenCount(forKey: .inputTokens) + outputTokens = c.decodeTokenCount(forKey: .outputTokens) + cacheReadTokens = c.decodeTokenCount(forKey: .cacheReadTokens) + cacheWriteTokens = c.decodeTokenCount(forKey: .cacheWriteTokens) topModels = try c.decodeIfPresent([DailyModelBreakdown].self, forKey: .topModels) ?? [] } } @@ -144,8 +176,8 @@ extension CurrentBlock { calls = try c.decode(Int.self, forKey: .calls) sessions = try c.decode(Int.self, forKey: .sessions) oneShotRate = try c.decodeIfPresent(Double.self, forKey: .oneShotRate) - inputTokens = try c.decode(Int.self, forKey: .inputTokens) - outputTokens = try c.decode(Int.self, forKey: .outputTokens) + inputTokens = c.decodeTokenCount(forKey: .inputTokens) + outputTokens = c.decodeTokenCount(forKey: .outputTokens) cacheHitPercent = try c.decodeIfPresent(Double.self, forKey: .cacheHitPercent) ?? 0 codexCredits = try c.decodeIfPresent(Double.self, forKey: .codexCredits) topActivities = try c.decodeIfPresent([ActivityEntry].self, forKey: .topActivities) ?? [] @@ -174,6 +206,23 @@ struct LocalModelSavingsByModel: Codable, Sendable { let outputTokens: Int } +extension LocalModelSavingsByModel { + enum CodingKeys: String, CodingKey { + case name, calls, actualUSD, savingsUSD, baselineModel, inputTokens, outputTokens + } + + init(from decoder: Decoder) throws { + let c = try decoder.container(keyedBy: CodingKeys.self) + name = try c.decode(String.self, forKey: .name) + calls = try c.decode(Int.self, forKey: .calls) + actualUSD = try c.decode(Double.self, forKey: .actualUSD) + savingsUSD = try c.decode(Double.self, forKey: .savingsUSD) + baselineModel = try c.decode(String.self, forKey: .baselineModel) + inputTokens = c.decodeTokenCount(forKey: .inputTokens) + outputTokens = c.decodeTokenCount(forKey: .outputTokens) + } +} + struct LocalModelSavingsByProvider: Codable, Sendable { let name: String let calls: Int @@ -260,8 +309,8 @@ struct SessionDetailEntry: Codable, Sendable { cost = try c.decode(Double.self, forKey: .cost) savingsUSD = try c.decodeIfPresent(Double.self, forKey: .savingsUSD) ?? 0 calls = try c.decode(Int.self, forKey: .calls) - inputTokens = try c.decode(Int.self, forKey: .inputTokens) - outputTokens = try c.decode(Int.self, forKey: .outputTokens) + inputTokens = c.decodeTokenCount(forKey: .inputTokens) + outputTokens = c.decodeTokenCount(forKey: .outputTokens) date = try c.decode(String.self, forKey: .date) models = try c.decodeIfPresent([SessionModelEntry].self, forKey: .models) ?? [] } diff --git a/mac/Tests/CodeBurnMenubarTests/MenubarPayloadDecodeTests.swift b/mac/Tests/CodeBurnMenubarTests/MenubarPayloadDecodeTests.swift new file mode 100644 index 00000000..075da76c --- /dev/null +++ b/mac/Tests/CodeBurnMenubarTests/MenubarPayloadDecodeTests.swift @@ -0,0 +1,126 @@ +import Foundation +import Testing +@testable import CodeBurnMenubar + +@Suite("MenubarPayload decode") +struct MenubarPayloadDecodeTests { + private func decode(_ json: String) throws -> MenubarPayload { + try JSONDecoder().decode(MenubarPayload.self, from: Data(json.utf8)) + } + + @Test("huge malformed token fields decode as zero") + func hugeMalformedTokenFieldsDecodeAsZero() throws { + let payload = try decode(""" + { + "generated": "2026-06-22T00:00:00Z", + "current": { + "label": "Today", + "cost": 1.0, + "calls": 1, + "sessions": 1, + "inputTokens": 18446744073709527000, + "outputTokens": "221360928884514260000", + "cacheHitPercent": 0, + "localModelSavings": { + "totalUSD": 1, + "calls": 1, + "byModel": [{ + "name": "local", + "calls": 1, + "actualUSD": 0, + "savingsUSD": 1, + "baselineModel": "paid", + "inputTokens": 221360928884514260000, + "outputTokens": -1 + }], + "byProvider": [] + }, + "topProjects": [{ + "name": "project", + "cost": 1, + "savingsUSD": 0, + "sessions": 1, + "avgCostPerSession": 1, + "sessionDetails": [{ + "cost": 1, + "savingsUSD": 0, + "calls": 1, + "inputTokens": 18446744073709527000, + "outputTokens": 1.5, + "date": "2026-06-21", + "models": [] + }] + }] + }, + "optimize": { + "findingCount": 0, + "savingsUSD": 0, + "topFindings": [] + }, + "history": { + "daily": [{ + "date": "2026-06-21", + "cost": 1, + "savingsUSD": 0, + "calls": 1, + "inputTokens": 18446744073709527000, + "outputTokens": -1, + "cacheReadTokens": 1.5, + "cacheWriteTokens": "221360928884514260000", + "topModels": [{ + "name": "Gemini 3.5 Flash", + "cost": 1, + "savingsUSD": 0, + "calls": 1, + "inputTokens": 221360928884514260000, + "outputTokens": -1 + }] + }] + } + } + """) + + #expect(payload.current.inputTokens == 0) + #expect(payload.current.outputTokens == 0) + #expect(payload.current.localModelSavings.byModel[0].inputTokens == 0) + #expect(payload.current.localModelSavings.byModel[0].outputTokens == 0) + #expect(payload.current.topProjects[0].sessionDetails[0].inputTokens == 0) + #expect(payload.current.topProjects[0].sessionDetails[0].outputTokens == 0) + #expect(payload.history.daily[0].inputTokens == 0) + #expect(payload.history.daily[0].outputTokens == 0) + #expect(payload.history.daily[0].cacheReadTokens == 0) + #expect(payload.history.daily[0].cacheWriteTokens == 0) + #expect(payload.history.daily[0].topModels[0].inputTokens == 0) + #expect(payload.history.daily[0].topModels[0].outputTokens == 0) + } + + @Test("huge non-token integers remain strict decode failures") + func hugeNonTokenIntegersRemainStrictDecodeFailures() { + var didThrow = false + do { + _ = try decode(""" + { + "generated": "2026-06-22T00:00:00Z", + "current": { + "label": "Today", + "cost": 1.0, + "calls": 18446744073709551615, + "sessions": 1, + "inputTokens": 1, + "outputTokens": 1 + }, + "optimize": { + "findingCount": 0, + "savingsUSD": 0, + "topFindings": [] + }, + "history": { "daily": [] } + } + """) + } catch { + didThrow = true + } + + #expect(didThrow) + } +} diff --git a/src/daily-cache.ts b/src/daily-cache.ts index 64843457..52e7f529 100644 --- a/src/daily-cache.ts +++ b/src/daily-cache.ts @@ -11,6 +11,10 @@ import type { DateRange, ProjectSummary } from './types.js' // forces a one-time full re-hydration so newly supported providers backfill // without a manual cache clear. // +// v9 also requires token counts to be non-negative safe integers. Older daily +// rollups may contain unsafe Antigravity uint64-underflow values that can crash +// the menubar Swift decoder, so the same re-hydration clears unsafe token data. +// // v8 added local-model savings to the daily rollup (savingsUSD per day / model / // category / provider). The `savingsConfigHash` field is invalidated separately // when the user changes their `localModelSavings` mapping so historical "saved" @@ -75,23 +79,80 @@ function isMigratableCache(parsed: unknown): parsed is { version: number; lastCo return c.version >= MIN_SUPPORTED_VERSION && c.version <= DAILY_CACHE_VERSION } -function migrateDays(days: Record[]): DailyEntry[] { - return days.map(d => ({ +function safeTokenCount(value: unknown): number | null { + if (value === undefined) return 0 + return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : null +} + +function migrateModels(models: unknown): DailyEntry['models'] | null { + if (models === undefined) return {} + if (!models || typeof models !== 'object' || Array.isArray(models)) return null + const migrated: DailyEntry['models'] = {} + for (const [name, raw] of Object.entries(models as Record)) { + if (!raw || typeof raw !== 'object' || Array.isArray(raw)) return null + const m = raw as Record + const inputTokens = safeTokenCount(m.inputTokens) + const outputTokens = safeTokenCount(m.outputTokens) + const cacheReadTokens = safeTokenCount(m.cacheReadTokens) + const cacheWriteTokens = safeTokenCount(m.cacheWriteTokens) + if ( + inputTokens === null || + outputTokens === null || + cacheReadTokens === null || + cacheWriteTokens === null + ) return null + migrated[name] = { + calls: (m.calls as number) ?? 0, + cost: (m.cost as number) ?? 0, + savingsUSD: (m.savingsUSD as number) ?? 0, + inputTokens, + outputTokens, + cacheReadTokens, + cacheWriteTokens, + } + } + return migrated +} + +function migrateDay(d: Record): DailyEntry | null { + const inputTokens = safeTokenCount(d.inputTokens) + const outputTokens = safeTokenCount(d.outputTokens) + const cacheReadTokens = safeTokenCount(d.cacheReadTokens) + const cacheWriteTokens = safeTokenCount(d.cacheWriteTokens) + const models = migrateModels(d.models) + if ( + inputTokens === null || + outputTokens === null || + cacheReadTokens === null || + cacheWriteTokens === null || + models === null + ) return null + return { date: d.date as string, cost: (d.cost as number) ?? 0, savingsUSD: (d.savingsUSD as number) ?? 0, calls: (d.calls as number) ?? 0, sessions: (d.sessions as number) ?? 0, - inputTokens: (d.inputTokens as number) ?? 0, - outputTokens: (d.outputTokens as number) ?? 0, - cacheReadTokens: (d.cacheReadTokens as number) ?? 0, - cacheWriteTokens: (d.cacheWriteTokens as number) ?? 0, + inputTokens, + outputTokens, + cacheReadTokens, + cacheWriteTokens, editTurns: (d.editTurns as number) ?? 0, oneShotTurns: (d.oneShotTurns as number) ?? 0, - models: (d.models as DailyEntry['models']) ?? {}, + models, categories: (d.categories as DailyEntry['categories']) ?? {}, providers: (d.providers as DailyEntry['providers']) ?? {}, - })) + } +} + +function migrateDays(days: Record[]): DailyEntry[] | null { + const migrated: DailyEntry[] = [] + for (const day of days) { + const entry = migrateDay(day) + if (!entry) return null + migrated.push(entry) + } + return migrated } async function backupOldCache(path: string, version: number): Promise { @@ -106,11 +167,16 @@ export async function loadDailyCache(): Promise { const raw = await readFile(path, 'utf-8') const parsed: unknown = JSON.parse(raw) if (isMigratableCache(parsed)) { + const days = migrateDays(parsed.days) + if (!days) { + await backupOldCache(path, parsed.version).catch(() => {}) + return emptyCache() + } const migrated: DailyCache = { version: DAILY_CACHE_VERSION, savingsConfigHash: parsed.savingsConfigHash ?? '', lastComputedDate: parsed.lastComputedDate, - days: migrateDays(parsed.days), + days, } if (parsed.version < DAILY_CACHE_VERSION) { await saveDailyCache(migrated).catch(() => {}) diff --git a/src/menubar-json.ts b/src/menubar-json.ts index 9fe3bec6..3eb88cff 100644 --- a/src/menubar-json.ts +++ b/src/menubar-json.ts @@ -204,6 +204,41 @@ function cacheHitPercent(inputTokens: number, cacheReadTokens: number): number { return (cacheReadTokens / denom) * 100 } +function safeMenubarToken(value: unknown): number { + return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : 0 +} + +function sanitizeDailyModelBreakdown(model: DailyModelBreakdown): DailyModelBreakdown { + return { + ...model, + inputTokens: safeMenubarToken(model.inputTokens), + outputTokens: safeMenubarToken(model.outputTokens), + } +} + +function sanitizeDailyHistoryEntry(entry: DailyHistoryEntry): DailyHistoryEntry { + return { + ...entry, + inputTokens: safeMenubarToken(entry.inputTokens), + outputTokens: safeMenubarToken(entry.outputTokens), + cacheReadTokens: safeMenubarToken(entry.cacheReadTokens), + cacheWriteTokens: safeMenubarToken(entry.cacheWriteTokens), + topModels: entry.topModels.map(sanitizeDailyModelBreakdown), + } +} + +function sanitizeLocalModelSavings(savings: LocalModelSavings | undefined): LocalModelSavings { + if (!savings) return { totalUSD: 0, calls: 0, byModel: [], byProvider: [] } + return { + ...savings, + byModel: savings.byModel.map(model => ({ + ...model, + inputTokens: safeMenubarToken(model.inputTokens), + outputTokens: safeMenubarToken(model.outputTokens), + })), + } +} + function buildTopActivities(categories: PeriodData['categories']): MenubarPayload['current']['topActivities'] { return categories.slice(0, TOP_ACTIVITIES_LIMIT).map(cat => ({ name: cat.name, @@ -252,7 +287,7 @@ function buildHistory(daily: DailyHistoryEntry[] | undefined): MenubarPayload['h if (!daily || daily.length === 0) return { daily: [] } const sorted = [...daily].sort((a, b) => a.date.localeCompare(b.date)) const trimmed = sorted.slice(-HISTORY_DAYS_LIMIT) - return { daily: trimmed } + return { daily: trimmed.map(sanitizeDailyHistoryEntry) } } function buildTopProjects(projects: PeriodData['projects']): MenubarPayload['current']['topProjects'] { @@ -270,8 +305,8 @@ function buildTopProjects(projects: PeriodData['projects']): MenubarPayload['cur cost: s.cost, savingsUSD: s.savingsUSD, calls: s.calls, - inputTokens: s.inputTokens, - outputTokens: s.outputTokens, + inputTokens: safeMenubarToken(s.inputTokens), + outputTokens: safeMenubarToken(s.outputTokens), date: s.date, models: s.models, })), @@ -315,6 +350,9 @@ export function buildMenubarPayload( routingWaste?: MenubarPayload['current']['routingWaste'], breakdowns?: BreakdownArrays, ): MenubarPayload { + const inputTokens = safeMenubarToken(current.inputTokens) + const outputTokens = safeMenubarToken(current.outputTokens) + const cacheReadTokens = safeMenubarToken(current.cacheReadTokens) return { generated: new Date().toISOString(), current: { @@ -323,13 +361,13 @@ export function buildMenubarPayload( calls: current.calls, sessions: current.sessions, oneShotRate: aggregateOneShotRate(current.categories), - inputTokens: current.inputTokens, - outputTokens: current.outputTokens, - cacheHitPercent: cacheHitPercent(current.inputTokens, current.cacheReadTokens), + inputTokens, + outputTokens, + cacheHitPercent: cacheHitPercent(inputTokens, cacheReadTokens), codexCredits: current.codexCredits ?? 0, topActivities: buildTopActivities(current.categories), topModels: buildTopModels(current.models), - localModelSavings: breakdowns?.localModelSavings ?? { totalUSD: 0, calls: 0, byModel: [], byProvider: [] }, + localModelSavings: sanitizeLocalModelSavings(breakdowns?.localModelSavings), providers: buildProviders(providers), topProjects: buildTopProjects(current.projects ?? []), modelEfficiency: buildModelEfficiency(current.modelEfficiency ?? []), diff --git a/src/providers/antigravity.ts b/src/providers/antigravity.ts index 710b4bf7..5debbf12 100644 --- a/src/providers/antigravity.ts +++ b/src/providers/antigravity.ts @@ -43,7 +43,7 @@ const CONVERSATION_ROOTS: readonly AntigravityConversationRoot[] = [ extensions: ['.pb'], }, ] as const -const CACHE_VERSION = 2 +const CACHE_VERSION = 4 const RPC_TIMEOUT_MS = 5000 const MAX_RESPONSE_BYTES = 16 * 1024 * 1024 @@ -95,10 +95,10 @@ type GeneratorMetadataResponse = { } type StatusLineCurrentUsage = { - input_tokens?: number - output_tokens?: number - cache_creation_input_tokens?: number - cache_read_input_tokens?: number + input_tokens?: number | string + output_tokens?: number | string + cache_creation_input_tokens?: number | string + cache_read_input_tokens?: number | string } type StatusLinePayload = { @@ -129,6 +129,7 @@ type StatusLineEvent = { type CachedCascade = { mtimeMs: number sizeBytes: number + fingerprint: string calls: ParsedProviderCall[] } @@ -644,8 +645,7 @@ function protoFieldText(field: ProtoField | undefined): string | undefined { function protoFieldPositiveInteger(field: ProtoField | undefined): number { if (field?.value === undefined) return 0 - const value = Number(field.value) - return Number.isSafeInteger(value) && value > 0 ? value : 0 + return bigintToSafeTokenCount(field.value) } function protoFieldBytes(field: ProtoField | undefined): Uint8Array | undefined { @@ -702,19 +702,24 @@ function buildCallFromSqliteGenMetadataRow(cascadeId: string, row: AntigravityGe let responseTokens = protoFieldPositiveInteger(firstProtoField(usageFields, 9)) let thinkingTokens = protoFieldPositiveInteger(firstProtoField(usageFields, 10)) + let splitOutputTokens = safeTokenSum(responseTokens, thinkingTokens) if (responseTokens === 0 && thinkingTokens === 0) { responseTokens = totalOutputTokens - } else if (totalOutputTokens > 0 && responseTokens + thinkingTokens !== totalOutputTokens) { + splitOutputTokens = responseTokens + } else if (totalOutputTokens > 0 && splitOutputTokens !== totalOutputTokens) { const adjustedResponseTokens = totalOutputTokens - thinkingTokens - if (adjustedResponseTokens >= 0) responseTokens = adjustedResponseTokens + if (adjustedResponseTokens >= 0) { + responseTokens = adjustedResponseTokens + splitOutputTokens = safeTokenSum(responseTokens, thinkingTokens) + } } - if (inputTokens === 0 && totalOutputTokens === 0) return null + if (inputTokens === 0 && splitOutputTokens === 0) return null const responseId = antigravitySqliteResponseId(usageFields, String(row.idx)) const model = antigravitySqliteModel(chatFields) const pricingModel = normalizePricingModel(model) - const costUSD = calculateCost(pricingModel, inputTokens, responseTokens + thinkingTokens, 0, 0, 0) + const costUSD = calculateCost(pricingModel, inputTokens, splitOutputTokens, 0, 0, 0) return { provider: 'antigravity', @@ -771,10 +776,38 @@ async function parseSqliteGenMetadataCalls(filePath: string, cascadeId: string): } } -function parseFiniteToken(value: unknown): number { - return typeof value === 'number' && Number.isFinite(value) && value > 0 - ? Math.floor(value) - : 0 +function bigintToSafeTokenCount(value: bigint): number { + if (value < 0n || value > BigInt(Number.MAX_SAFE_INTEGER)) return 0 + return Number(value) +} + +export function parseAntigravityTokenCount(value: unknown): number { + if (typeof value === 'number') { + return Number.isSafeInteger(value) && value >= 0 ? value : 0 + } + if (typeof value === 'bigint') { + return bigintToSafeTokenCount(value) + } + if (typeof value === 'string') { + const trimmed = value.trim() + if (!/^\d+$/.test(trimmed)) return 0 + try { + return bigintToSafeTokenCount(BigInt(trimmed)) + } catch { + return 0 + } + } + return 0 +} + +function safeTokenSum(...values: number[]): number { + let total = 0 + for (const value of values) { + if (!Number.isSafeInteger(value) || value < 0) return 0 + if (total > Number.MAX_SAFE_INTEGER - value) return 0 + total += value + } + return total } function usageSignature(event: StatusLineEvent): string { @@ -819,7 +852,7 @@ export function antigravityCascadeIdFromPath(path: string): string { return basename(path).replace(/\.(pb|db)$/i, '') } -function buildCallsFromGeneratorMetadata( +export function buildCallsFromGeneratorMetadata( cascadeId: string, metadata: GeneratorMetadata[], modelMap: ModelMap, @@ -831,12 +864,29 @@ function buildCallsFromGeneratorMetadata( const usage = entry.chatModel?.usage if (!usage) continue - const inputTokens = parseInt(usage.inputTokens ?? '0', 10) - const outputTokens = parseInt(usage.outputTokens ?? '0', 10) - const thinkingTokens = parseInt(usage.thinkingOutputTokens ?? '0', 10) - const responseTokens = parseInt(usage.responseOutputTokens ?? '0', 10) + const inputTokens = parseAntigravityTokenCount(usage.inputTokens) + const outputTokens = parseAntigravityTokenCount(usage.outputTokens) + const thinkingTokens = parseAntigravityTokenCount(usage.thinkingOutputTokens) + let assertionResponseTokens = parseAntigravityTokenCount(usage.responseOutputTokens) + + // Mirror SQLite invariant: when total outputTokens exists and the + // response+thinking split does not equal it, infer the missing + // response tokens from the gap so billing covers the full output. + if (assertionResponseTokens === 0 && thinkingTokens === 0) { + assertionResponseTokens = outputTokens + } else if (outputTokens > 0) { + const splitSum = safeTokenSum(assertionResponseTokens, thinkingTokens) + if (splitSum !== outputTokens) { + const adjusted = outputTokens - thinkingTokens + if (adjusted >= 0) { + assertionResponseTokens = adjusted + } + } + } - if (inputTokens === 0 && outputTokens === 0) continue + const finalOutputCostTokens = safeTokenSum(assertionResponseTokens, thinkingTokens) + + if (inputTokens === 0 && finalOutputCostTokens === 0) continue const responseId = usage.responseId || String(i) const dedupKey = `antigravity:${cascadeId}:${responseId}` @@ -844,13 +894,13 @@ function buildCallsFromGeneratorMetadata( const model = modelMap[usage.model] ?? usage.model const pricingModel = normalizePricingModel(model) const timestamp = entry.chatModel?.chatStartMetadata?.createdAt ?? '' - const costUSD = calculateCost(pricingModel, inputTokens, responseTokens + thinkingTokens, 0, 0, 0) + const costUSD = calculateCost(pricingModel, inputTokens, finalOutputCostTokens, 0, 0, 0) results.push({ provider: 'antigravity', model, inputTokens, - outputTokens: responseTokens, + outputTokens: assertionResponseTokens, cacheCreationInputTokens: 0, cacheReadInputTokens: 0, cachedInputTokens: 0, @@ -935,10 +985,10 @@ function parseStatusLinePayload(input: unknown): StatusLineEvent | null { ? payload.model : payload.model?.id ?? payload.model?.display_name ?? 'unknown', usage: { - inputTokens: parseFiniteToken(usage.input_tokens), - outputTokens: parseFiniteToken(usage.output_tokens), - cacheCreationInputTokens: parseFiniteToken(usage.cache_creation_input_tokens), - cacheReadInputTokens: parseFiniteToken(usage.cache_read_input_tokens), + inputTokens: parseAntigravityTokenCount(usage.input_tokens), + outputTokens: parseAntigravityTokenCount(usage.output_tokens), + cacheCreationInputTokens: parseAntigravityTokenCount(usage.cache_creation_input_tokens), + cacheReadInputTokens: parseAntigravityTokenCount(usage.cache_read_input_tokens), }, } @@ -974,10 +1024,10 @@ function parseStatusLineEvent(input: unknown): StatusLineEvent | null { if (!event.usage || typeof event.usage !== 'object') return null const usage = { - inputTokens: parseFiniteToken(event.usage.inputTokens), - outputTokens: parseFiniteToken(event.usage.outputTokens), - cacheCreationInputTokens: parseFiniteToken(event.usage.cacheCreationInputTokens), - cacheReadInputTokens: parseFiniteToken(event.usage.cacheReadInputTokens), + inputTokens: parseAntigravityTokenCount(event.usage.inputTokens), + outputTokens: parseAntigravityTokenCount(event.usage.outputTokens), + cacheCreationInputTokens: parseAntigravityTokenCount(event.usage.cacheCreationInputTokens), + cacheReadInputTokens: parseAntigravityTokenCount(event.usage.cacheReadInputTokens), } if ( @@ -1099,6 +1149,40 @@ export function shouldReparseAntigravitySource(path: string, cachedTurnCount: nu return isAntigravityStatusLineEventsPath(path) } +export async function computeAntigravityCacheFingerprint(filePath: string): Promise<{ + fingerprint: string + baseMtimeMs: number + baseSize: number +} | null> { + const base = await stat(filePath).catch(() => null) + if (!base) return null + + // Legacy .pb: fingerprint is base mtime + size alone + if (!filePath.toLowerCase().endsWith('.db')) { + return { + fingerprint: `${base.mtimeMs}:${base.size}`, + baseMtimeMs: base.mtimeMs, + baseSize: base.size, + } + } + + // SQLite WAL-mode .db: include sidecars so a WAL checkpoint or write + // that touches -wal/-shm invalidates the cache even when the base .db + // mtime and size haven't changed yet. + const wal = await stat(`${filePath}-wal`).catch(() => null) + const shm = await stat(`${filePath}-shm`).catch(() => null) + + return { + fingerprint: [ + base.mtimeMs, base.size, + wal?.mtimeMs ?? 0, wal?.size ?? 0, + shm?.mtimeMs ?? 0, shm?.size ?? 0, + ].join(':'), + baseMtimeMs: base.mtimeMs, + baseSize: base.size, + } +} + async function findCascadeSource(cascadeId: string): Promise { const sources = await discoverAntigravitySessionSources() return sources.find(source => { @@ -1116,12 +1200,12 @@ export async function snapshotAntigravityStatusLinePayload(input: unknown): Prom const source = await findCascadeSource(cascadeId) if (!source) return false - const s = await stat(source.path).catch(() => null) - if (!s) return false + const fp = await computeAntigravityCacheFingerprint(source.path) + if (!fp) return false const cache = await loadCache() const cached = cache.cascades[cascadeId] - if (cached && cached.mtimeMs === s.mtimeMs && cached.sizeBytes === s.size && cached.calls.length > 0) { + if (cached && cached.fingerprint === fp.fingerprint && cached.calls.length > 0) { return true } @@ -1135,8 +1219,9 @@ export async function snapshotAntigravityStatusLinePayload(input: unknown): Prom await rpc(server, 'GetCascadeTrajectoryGeneratorMetadata', { cascadeId }), ) cache.cascades[cascadeId] = { - mtimeMs: s.mtimeMs, - sizeBytes: s.size, + mtimeMs: fp.baseMtimeMs, + sizeBytes: fp.baseSize, + fingerprint: fp.fingerprint, calls: buildCallsFromGeneratorMetadata(cascadeId, metadata, modelMap), } cacheDirty = true @@ -1211,13 +1296,13 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars const cascadeId = antigravityCascadeIdFromPath(source.path) const cache = await loadCache() - const s = await stat(source.path).catch(() => null) - if (!s) return + const fp = await computeAntigravityCacheFingerprint(source.path) + if (!fp) return const projectPath = await extractWorkspacePath(source.path) const cached = cache.cascades[cascadeId] - if (cached && cached.mtimeMs === s.mtimeMs && cached.sizeBytes === s.size && cached.calls.length > 0) { + if (cached && cached.fingerprint === fp.fingerprint && cached.calls.length > 0) { for (const call of cached.calls) { applyAntigravityProject(call, source, projectPath) if (seenKeys.has(call.deduplicationKey)) continue @@ -1234,8 +1319,9 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars } cache.cascades[cascadeId] = { - mtimeMs: s.mtimeMs, - sizeBytes: s.size, + mtimeMs: fp.baseMtimeMs, + sizeBytes: fp.baseSize, + fingerprint: fp.fingerprint, calls: sqliteResults, } cacheDirty = true @@ -1286,8 +1372,9 @@ function createParser(source: SessionSource, seenKeys: Set): SessionPars } cache.cascades[cascadeId] = { - mtimeMs: s.mtimeMs, - sizeBytes: s.size, + mtimeMs: fp.baseMtimeMs, + sizeBytes: fp.baseSize, + fingerprint: fp.fingerprint, calls: results, } cacheDirty = true diff --git a/src/session-cache.ts b/src/session-cache.ts index 1748b964..e64c175c 100644 --- a/src/session-cache.ts +++ b/src/session-cache.ts @@ -48,6 +48,10 @@ export type FileFingerprint = { ino: number mtimeMs: number sizeBytes: number + walMtimeMs?: number + walSizeBytes?: number + shmMtimeMs?: number + shmSizeBytes?: number } export type CachedFile = { @@ -117,7 +121,7 @@ const PROVIDER_PARSE_VERSIONS: Record = { 'kilo-code': 'worktree-project-grouping-v1', 'roo-code': 'worktree-project-grouping-v1', warp: 'worktree-project-grouping-v1', - antigravity: 'worktree-project-grouping-v3', + antigravity: 'safe-token-counts-v1', } // ── Cache Dir ────────────────────────────────────────────────────────── @@ -150,6 +154,10 @@ function isNum(v: unknown): v is number { return typeof v === 'number' && Number.isFinite(v) } +function isSafeTokenCount(v: unknown): v is number { + return typeof v === 'number' && Number.isSafeInteger(v) && v >= 0 +} + function isStringArray(v: unknown): v is string[] { return Array.isArray(v) && v.every(e => typeof e === 'string') } @@ -183,10 +191,10 @@ function validateFingerprint(fp: unknown): fp is FileFingerprint { function validateUsage(u: unknown): u is CachedUsage { if (!u || typeof u !== 'object') return false const o = u as Record - return isNum(o['inputTokens']) && isNum(o['outputTokens']) - && isNum(o['cacheCreationInputTokens']) && isNum(o['cacheReadInputTokens']) - && isNum(o['cachedInputTokens']) && isNum(o['reasoningTokens']) - && isNum(o['webSearchRequests']) && isNum(o['cacheCreationOneHourTokens']) + return isSafeTokenCount(o['inputTokens']) && isSafeTokenCount(o['outputTokens']) + && isSafeTokenCount(o['cacheCreationInputTokens']) && isSafeTokenCount(o['cacheReadInputTokens']) + && isSafeTokenCount(o['cachedInputTokens']) && isSafeTokenCount(o['reasoningTokens']) + && isNum(o['webSearchRequests']) && isSafeTokenCount(o['cacheCreationOneHourTokens']) } function validateCall(c: unknown): c is CachedCall { @@ -238,6 +246,25 @@ function validateProviderSection(s: unknown): s is ProviderSection { return Object.values(o['files'] as Record).every(validateCachedFile) } +function sanitizeProviderSection(s: unknown): ProviderSection | null { + if (!s || typeof s !== 'object') return null + const o = s as Record + if (typeof o['envFingerprint'] !== 'string') return null + if (!o['files'] || typeof o['files'] !== 'object' || Array.isArray(o['files'])) return null + + const files: Record = {} + for (const [path, file] of Object.entries(o['files'] as Record)) { + if (validateCachedFile(file)) files[path] = file + } + if (Object.keys(files).length === 0) return null + + return { + envFingerprint: o['envFingerprint'], + files, + ...(o['durable'] === true ? { durable: true } : {}), + } +} + function validateCache(raw: unknown): raw is SessionCache { if (!raw || typeof raw !== 'object') return false const o = raw as Record @@ -246,11 +273,25 @@ function validateCache(raw: unknown): raw is SessionCache { return Object.values(o['providers'] as Record).every(validateProviderSection) } +function sanitizeCache(raw: unknown): SessionCache | null { + if (!raw || typeof raw !== 'object') return null + const o = raw as Record + if (o['version'] !== CACHE_VERSION) return null + if (!o['providers'] || typeof o['providers'] !== 'object' || Array.isArray(o['providers'])) return null + + const providers: Record = {} + for (const [provider, section] of Object.entries(o['providers'] as Record)) { + const sanitized = sanitizeProviderSection(section) + if (sanitized) providers[provider] = sanitized + } + return { version: CACHE_VERSION, providers } +} + export async function loadCache(): Promise { try { const raw = await readFile(getCachePath(), 'utf-8') const parsed = JSON.parse(raw) - if (!validateCache(parsed)) return emptyCache() + if (!validateCache(parsed)) return sanitizeCache(parsed) ?? emptyCache() return parsed } catch { return emptyCache() @@ -285,9 +326,27 @@ export async function saveCache(cache: SessionCache): Promise { // ── File Fingerprinting ──────────────────────────────────────────────── export async function fingerprintFile(filePath: string): Promise { + async function fingerprintBasePath(basePath: string): Promise { + const s = await stat(basePath) + const fp: FileFingerprint = { dev: s.dev, ino: s.ino, mtimeMs: s.mtimeMs, sizeBytes: s.size } + + const wal = await stat(`${basePath}-wal`).catch(() => null) + if (wal) { + fp.walMtimeMs = wal.mtimeMs + fp.walSizeBytes = wal.size + } + + const shm = await stat(`${basePath}-shm`).catch(() => null) + if (shm) { + fp.shmMtimeMs = shm.mtimeMs + fp.shmSizeBytes = shm.size + } + + return fp + } + try { - const s = await stat(filePath) - return { dev: s.dev, ino: s.ino, mtimeMs: s.mtimeMs, sizeBytes: s.size } + return await fingerprintBasePath(filePath) } catch { // Providers encode extra context into source paths using virtual suffixes: // - Cursor: `#cursor-ws=` (workspace-aware routing) @@ -298,8 +357,7 @@ export async function fingerprintFile(filePath: string): Promise 0) { try { - const s = await stat(filePath.slice(0, hashIdx)) - return { dev: s.dev, ino: s.ino, mtimeMs: s.mtimeMs, sizeBytes: s.size } + return await fingerprintBasePath(filePath.slice(0, hashIdx)) } catch { // fall through to colon check } @@ -307,8 +365,7 @@ export async function fingerprintFile(filePath: string): Promise 0) { try { - const s = await stat(filePath.slice(0, colonIdx)) - return { dev: s.dev, ino: s.ino, mtimeMs: s.mtimeMs, sizeBytes: s.size } + return await fingerprintBasePath(filePath.slice(0, colonIdx)) } catch { return null } @@ -332,6 +389,14 @@ export function reconcileFile( if (!cached) return { action: 'new' } const fp = cached.fingerprint + if ( + fp.walMtimeMs !== current.walMtimeMs || + fp.walSizeBytes !== current.walSizeBytes || + fp.shmMtimeMs !== current.shmMtimeMs || + fp.shmSizeBytes !== current.shmSizeBytes + ) { + return { action: 'modified' } + } if ( fp.dev === current.dev && diff --git a/tests/daily-cache.test.ts b/tests/daily-cache.test.ts index 98923192..639fc555 100644 --- a/tests/daily-cache.test.ts +++ b/tests/daily-cache.test.ts @@ -38,6 +38,10 @@ function emptyDay(date: string, cost = 0, calls = 0): DailyEntry { } const TMP_CACHE_ROOT = join(tmpdir(), `codeburn-cache-test-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`) +const EXACT_UNSAFE_TOKEN_COUNTS = [ + '221360928884514260000', + '18446744073709527000', +] as const beforeEach(() => { process.env['CODEBURN_CACHE_DIR'] = TMP_CACHE_ROOT @@ -149,6 +153,71 @@ describe('loadDailyCache', () => { const loaded = await loadDailyCache() expect(loaded).toEqual(saved) }) + + it('rejects a current cache with an unsafe day token count', async () => { + const saved = { + version: DAILY_CACHE_VERSION, + savingsConfigHash: '', + lastComputedDate: '2026-04-10', + days: [{ ...emptyDay('2026-04-10'), inputTokens: Number(EXACT_UNSAFE_TOKEN_COUNTS[1]) }], + } + const { writeFile, mkdir } = await import('fs/promises') + await mkdir(TMP_CACHE_ROOT, { recursive: true }) + await writeFile(join(TMP_CACHE_ROOT, 'daily-cache.json'), JSON.stringify(saved), 'utf-8') + + const cache = await loadDailyCache() + expect(cache.days).toEqual([]) + expect(cache.lastComputedDate).toBeNull() + expect(existsSync(join(TMP_CACHE_ROOT, `daily-cache.json.v${DAILY_CACHE_VERSION}.bak`))).toBe(true) + }) + + it('rejects a current cache with an unsafe model token count', async () => { + const saved = { + version: DAILY_CACHE_VERSION, + savingsConfigHash: '', + lastComputedDate: '2026-04-10', + days: [{ + ...emptyDay('2026-04-10'), + models: { + 'Gemini 3.5 Flash': { + calls: 1, + cost: 0.01, + savingsUSD: 0, + inputTokens: 10, + outputTokens: EXACT_UNSAFE_TOKEN_COUNTS[0], + cacheReadTokens: 0, + cacheWriteTokens: 0, + }, + }, + }], + } + const { writeFile, mkdir } = await import('fs/promises') + await mkdir(TMP_CACHE_ROOT, { recursive: true }) + await writeFile(join(TMP_CACHE_ROOT, 'daily-cache.json'), JSON.stringify(saved), 'utf-8') + + const cache = await loadDailyCache() + expect(cache.days).toEqual([]) + expect(cache.lastComputedDate).toBeNull() + expect(existsSync(join(TMP_CACHE_ROOT, `daily-cache.json.v${DAILY_CACHE_VERSION}.bak`))).toBe(true) + }) + + it('backs up a pre-v9 cache so stale token rollups are recomputed', async () => { + const staleVersion = DAILY_CACHE_VERSION - 1 + const saved = { + version: staleVersion, + savingsConfigHash: '', + lastComputedDate: '2026-04-10', + days: [{ ...emptyDay('2026-04-10'), inputTokens: Number(EXACT_UNSAFE_TOKEN_COUNTS[1]) }], + } + const { writeFile, mkdir } = await import('fs/promises') + await mkdir(TMP_CACHE_ROOT, { recursive: true }) + await writeFile(join(TMP_CACHE_ROOT, 'daily-cache.json'), JSON.stringify(saved), 'utf-8') + + const cache = await loadDailyCache() + expect(cache.days).toEqual([]) + expect(cache.lastComputedDate).toBeNull() + expect(existsSync(join(TMP_CACHE_ROOT, `daily-cache.json.v${staleVersion}.bak`))).toBe(true) + }) }) describe('saveDailyCache', () => { diff --git a/tests/menubar-json.test.ts b/tests/menubar-json.test.ts index f7493d0b..8fa07582 100644 --- a/tests/menubar-json.test.ts +++ b/tests/menubar-json.test.ts @@ -3,6 +3,11 @@ import { describe, expect, it } from 'vitest' import { buildMenubarPayload, type PeriodData, type ProviderCost } from '../src/menubar-json.js' import type { OptimizeResult } from '../src/optimize.js' +const EXACT_UNSAFE_TOKEN_COUNTS = [ + '221360928884514260000', + '18446744073709527000', +] as const + function emptyPeriod(label: string): PeriodData { return { label, @@ -231,4 +236,75 @@ describe('buildMenubarPayload', () => { const payload = buildMenubarPayload(emptyPeriod('Today'), providers, null) expect(payload.current.providers).toEqual({ claude: 76.45 }) }) + + it('sanitizes unsafe token fields before emitting menubar JSON', () => { + const unsafe = Number(EXACT_UNSAFE_TOKEN_COUNTS[1]) + const unsafeRuntimeString = EXACT_UNSAFE_TOKEN_COUNTS[0] as unknown as number + const period: PeriodData = { + ...emptyPeriod('Today'), + inputTokens: unsafe, + outputTokens: unsafeRuntimeString, + cacheReadTokens: 100, + projects: [{ + name: 'project', + cost: 1, + savingsUSD: 0, + sessions: 1, + sessionDetails: [{ + cost: 1, + savingsUSD: 0, + calls: 1, + inputTokens: unsafe, + outputTokens: Number(EXACT_UNSAFE_TOKEN_COUNTS[0]), + date: '2026-06-21', + models: [], + }], + }], + } + const history = [{ + date: '2026-06-21', + cost: 1, + savingsUSD: 0, + calls: 1, + inputTokens: unsafe, + outputTokens: -1, + cacheReadTokens: unsafeRuntimeString, + cacheWriteTokens: 10, + topModels: [{ name: 'Gemini 3.5 Flash', cost: 1, savingsUSD: 0, calls: 1, inputTokens: unsafeRuntimeString, outputTokens: -1 }], + }] + + const payload = buildMenubarPayload(period, [], null, history, undefined, undefined, { + localModelSavings: { + totalUSD: 1, + calls: 1, + byModel: [{ + name: 'local', + calls: 1, + actualUSD: 0, + savingsUSD: 1, + baselineModel: 'paid', + inputTokens: unsafeRuntimeString, + outputTokens: -1, + }], + byProvider: [], + }, + }) + + expect(payload.current.inputTokens).toBe(0) + expect(payload.current.outputTokens).toBe(0) + expect(payload.current.cacheHitPercent).toBe(100) + expect(payload.current.topProjects[0]!.sessionDetails[0]!.inputTokens).toBe(0) + expect(payload.current.topProjects[0]!.sessionDetails[0]!.outputTokens).toBe(0) + expect(payload.history.daily[0]!.inputTokens).toBe(0) + expect(payload.history.daily[0]!.outputTokens).toBe(0) + expect(payload.history.daily[0]!.cacheReadTokens).toBe(0) + expect(payload.history.daily[0]!.cacheWriteTokens).toBe(10) + expect(payload.history.daily[0]!.topModels[0]!.inputTokens).toBe(0) + expect(payload.history.daily[0]!.topModels[0]!.outputTokens).toBe(0) + expect(payload.current.localModelSavings.byModel[0]!.inputTokens).toBe(0) + expect(payload.current.localModelSavings.byModel[0]!.outputTokens).toBe(0) + const json = JSON.stringify(payload) + expect(json).not.toContain(EXACT_UNSAFE_TOKEN_COUNTS[0]) + expect(json).not.toContain(EXACT_UNSAFE_TOKEN_COUNTS[1]) + }) }) diff --git a/tests/providers/antigravity.test.ts b/tests/providers/antigravity.test.ts index e98fd591..42657b58 100644 --- a/tests/providers/antigravity.test.ts +++ b/tests/providers/antigravity.test.ts @@ -1,4 +1,4 @@ -import { mkdtemp, mkdir, readFile, rm, writeFile } from 'fs/promises' +import { mkdtemp, mkdir, readFile, rm, stat, writeFile } from 'fs/promises' import { tmpdir } from 'os' import { join } from 'path' import { createRequire } from 'node:module' @@ -8,12 +8,15 @@ import { isSqliteAvailable } from '../../src/sqlite.js' import { antigravityAppDataDirFromSourcePath, antigravityCascadeIdFromPath, + buildCallsFromGeneratorMetadata, + computeAntigravityCacheFingerprint, createAntigravityProvider, discoverAntigravitySessionSources, extractAntigravityAppDataDirFromLine, extractAntigravityGeneratorMetadata, extractAntigravityModelMap, getAntigravityStatusLineEventsPath, + parseAntigravityTokenCount, parseAntigravityServerInfo, parseAntigravityServerInfoFromLine, recordAntigravityStatusLinePayload, @@ -22,6 +25,10 @@ import { import type { ParsedProviderCall } from '../../src/providers/types.js' const requireForTest = createRequire(import.meta.url) +const EXACT_UNSAFE_TOKEN_COUNTS = [ + '221360928884514260000', + '18446744073709527000', +] as const type CurrentCliFixture = { conversationId: string @@ -206,6 +213,83 @@ describe('antigravity provider helpers', () => { expect(extractAntigravityGeneratorMetadata(null)).toEqual([]) }) + it('keeps output-only generator metadata calls when split tokens recover output', () => { + const calls = buildCallsFromGeneratorMetadata('split-output-cascade', [{ + chatModel: { + usage: { + model: 'gemini-3.5-flash-high', + inputTokens: '0', + outputTokens: '0', + responseOutputTokens: '7', + thinkingOutputTokens: '3', + responseId: 'split-output-only', + }, + chatStartMetadata: { + createdAt: '2026-06-22T00:00:00Z', + }, + }, + }], {}) + + expect(calls).toHaveLength(1) + expect(calls[0]).toMatchObject({ + provider: 'antigravity', + model: 'gemini-3.5-flash-high', + inputTokens: 0, + outputTokens: 7, + reasoningTokens: 3, + sessionId: 'split-output-cascade', + deduplicationKey: 'antigravity:split-output-cascade:split-output-only', + }) + expect(calls[0]!.costUSD).toBeGreaterThan(0) + }) + + it('infers missing responseOutputTokens from total outputTokens when thinking tokens are present', () => { + const calls = buildCallsFromGeneratorMetadata('infer-response', [{ + chatModel: { + usage: { + model: 'gemini-3.5-flash-high', + inputTokens: '5', + outputTokens: '10', + thinkingOutputTokens: '3', + responseId: 'infer-response-id', + }, + chatStartMetadata: { + createdAt: '2026-06-22T00:00:00Z', + }, + }, + }], {}) + + expect(calls).toHaveLength(1) + expect(calls[0]).toMatchObject({ + provider: 'antigravity', + model: 'gemini-3.5-flash-high', + inputTokens: 5, + outputTokens: 7, + reasoningTokens: 3, + sessionId: 'infer-response', + deduplicationKey: 'antigravity:infer-response:infer-response-id', + }) + expect(calls[0]!.costUSD).toBeGreaterThan(0) + }) + + it('parses token counts via BigInt and rejects unsafe values', () => { + expect(parseAntigravityTokenCount('42')).toBe(42) + expect(parseAntigravityTokenCount(42n)).toBe(42) + expect(parseAntigravityTokenCount(`${Number.MAX_SAFE_INTEGER}`)).toBe(Number.MAX_SAFE_INTEGER) + expect(parseAntigravityTokenCount(BigInt(Number.MAX_SAFE_INTEGER))).toBe(Number.MAX_SAFE_INTEGER) + for (const value of EXACT_UNSAFE_TOKEN_COUNTS) { + expect(parseAntigravityTokenCount(value)).toBe(0) + expect(parseAntigravityTokenCount(Number(value))).toBe(0) + expect(parseAntigravityTokenCount(BigInt(value))).toBe(0) + } + expect(parseAntigravityTokenCount('18446744073709551615')).toBe(0) + expect(parseAntigravityTokenCount(Number.MAX_SAFE_INTEGER + 1)).toBe(0) + expect(parseAntigravityTokenCount(-1)).toBe(0) + expect(parseAntigravityTokenCount(1.5)).toBe(0) + expect(parseAntigravityTokenCount('1.5')).toBe(0) + expect(parseAntigravityTokenCount('10tokens')).toBe(0) + }) + it('derives cascade ids from legacy .pb and Antigravity 2 .db files', () => { expect(antigravityCascadeIdFromPath('/tmp/123.pb')).toBe('123') expect(antigravityCascadeIdFromPath('/tmp/456.db')).toBe('456') @@ -362,6 +446,53 @@ describe('antigravity provider helpers', () => { } }) + it('sanitizes unsafe statusLine token counts before recording fallback calls', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-statusline-overflow-')) + process.env['CODEBURN_CACHE_DIR'] = dir + + try { + const payload = { + conversation_id: 'overflow-statusline', + session_id: 'session-1', + model: 'Gemini 3.5 Flash (High)', + context_window: { + current_usage: { + input_tokens: EXACT_UNSAFE_TOKEN_COUNTS[1], + output_tokens: 7, + cache_creation_input_tokens: EXACT_UNSAFE_TOKEN_COUNTS[0], + cache_read_input_tokens: Number(EXACT_UNSAFE_TOKEN_COUNTS[1]), + }, + }, + } + + expect(await recordAntigravityStatusLinePayload(payload)).toBe(true) + expect(await recordAntigravityStatusLinePayload(payload)).toBe(true) + + const recorded = await readFile(getAntigravityStatusLineEventsPath(), 'utf-8') + expect(recorded).not.toContain(EXACT_UNSAFE_TOKEN_COUNTS[0]) + expect(recorded).not.toContain(EXACT_UNSAFE_TOKEN_COUNTS[1]) + + const parser = createAntigravityProvider().createSessionParser({ + path: getAntigravityStatusLineEventsPath(), + project: 'antigravity-cli', + provider: 'antigravity', + }, new Set()) + + const calls = [] + for await (const call of parser.parse()) calls.push(call) + + expect(calls).toHaveLength(1) + expect(calls[0]).toMatchObject({ + inputTokens: 0, + outputTokens: 7, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + }) + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + it('skips statusLine fallback calls when RPC cache already covered the conversation', async () => { const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-statusline-rpc-dedup-')) process.env['CODEBURN_CACHE_DIR'] = dir @@ -605,4 +736,96 @@ describe('antigravity provider helpers', () => { await rm(tempHome, { recursive: true, force: true }) } }) + + it('includes WAL and SHM sidecar stats in .db fingerprint but not in .pb', async () => { + const dir = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-fingerprint-')) + + try { + // .pb: only base stats, no sidecar fields + const pbPath = join(dir, 'session.pb') + await writeFile(pbPath, 'legacy pb bytes') + const pbStat = await stat(pbPath) + const pbFp = await computeAntigravityCacheFingerprint(pbPath) + expect(pbFp).not.toBeNull() + expect(pbFp!.fingerprint).toBe(`${pbStat.mtimeMs}:${pbStat.size}`) + expect(pbFp!.fingerprint).not.toContain(':0:0:0:0') + + // .db without sidecars: base + four zeros + const dbPath = join(dir, 'session.db') + const { DatabaseSync: Database } = requireForTest('node:sqlite') + const db = new Database(dbPath) + db.exec('CREATE TABLE gen_metadata (idx integer, data blob, size integer NOT NULL DEFAULT 0, PRIMARY KEY (idx))') + db.close() + const dbStat = await stat(dbPath) + const dbFp = await computeAntigravityCacheFingerprint(dbPath) + expect(dbFp).not.toBeNull() + expect(dbFp!.fingerprint).toBe(`${dbStat.mtimeMs}:${dbStat.size}:0:0:0:0`) + + // .db with WAL sidecar present: fingerprint differs from no-WAL case + await writeFile(`${dbPath}-wal`, 'wal journal bytes') + const walStat = await stat(`${dbPath}-wal`) + const dbWalFp = await computeAntigravityCacheFingerprint(dbPath) + expect(dbWalFp).not.toBeNull() + expect(dbWalFp!.fingerprint).toBe( + `${dbStat.mtimeMs}:${dbStat.size}:${walStat.mtimeMs}:${walStat.size}:0:0`, + ) + expect(dbFp!.fingerprint).not.toBe(dbWalFp!.fingerprint) + } finally { + await rm(dir, { recursive: true, force: true }) + } + }) + + it('reparses a cached .db source when its WAL sidecar mtime changes', async () => { + if (!isSqliteAvailable()) return + + const tempHome = await mkdtemp(join(tmpdir(), 'codeburn-antigravity-wal-reparse-')) + const cacheDir = join(tempHome, 'cache') + const previousCacheDir = process.env['CODEBURN_CACHE_DIR'] + process.env['CODEBURN_CACHE_DIR'] = cacheDir + + try { + const fixture = JSON.parse(await readFile( + new URL('../fixtures/antigravity-cli-current/gen-metadata.json', import.meta.url), + 'utf-8', + )) as CurrentCliFixture + const conversationsDir = join(tempHome, '.gemini', 'antigravity-cli', 'conversations') + await mkdir(conversationsDir, { recursive: true }) + + const dbPath = join(conversationsDir, `${fixture.conversationId}.db`) + createCurrentAntigravityCliDb(dbPath, fixture) + + // First parse: populates the in-memory cache with fingerprint A + const sources = await discoverAntigravitySessionSources([{ + dir: conversationsDir, + project: 'antigravity-cli', + extensions: ['.pb', '.db'], + }]) + expect(sources).toHaveLength(1) + + const firstCalls = await collectAntigravityCalls(sources[0]!) + expect(firstCalls.length).toBeGreaterThanOrEqual(1) + + // Touch the -wal sidecar so its mtime changes while the base .db is unchanged + const walPath = `${dbPath}-wal` + await writeFile(walPath, 'touched wal journal') + // Ensure mtime actually advanced (some filesystems have coarse resolution) + await new Promise(resolve => setTimeout(resolve, 10)) + const walTouchStat = await stat(walPath) + + // Compute fingerprint now → must differ from the cached fingerprint + const fpAfter = await computeAntigravityCacheFingerprint(dbPath) + expect(fpAfter).not.toBeNull() + expect(fpAfter!.fingerprint).toContain(`${walTouchStat.mtimeMs}:${walTouchStat.size}`) + + // Second parse with same source: fingerprint changed → cache miss → re-read + const secondCalls = await collectAntigravityCalls(sources[0]!) + expect(secondCalls.length).toBeGreaterThanOrEqual(1) + // The re-read should produce equivalent calls (same underlying DB rows) + expect(secondCalls.length).toBe(firstCalls.length) + } finally { + if (previousCacheDir === undefined) delete process.env['CODEBURN_CACHE_DIR'] + else process.env['CODEBURN_CACHE_DIR'] = previousCacheDir + await rm(tempHome, { recursive: true, force: true }) + } + }) }) diff --git a/tests/session-cache.test.ts b/tests/session-cache.test.ts index b015322b..76919622 100644 --- a/tests/session-cache.test.ts +++ b/tests/session-cache.test.ts @@ -22,6 +22,10 @@ import { } from '../src/session-cache.js' const TMP_DIR = join(tmpdir(), `codeburn-scache-test-${Date.now()}-${Math.random().toString(36).slice(2, 8)}`) +const EXACT_UNSAFE_TOKEN_COUNTS = [ + '221360928884514260000', + '18446744073709527000', +] as const beforeEach(() => { process.env['CODEBURN_CACHE_DIR'] = TMP_DIR @@ -244,6 +248,41 @@ describe('fingerprintFile', () => { expect(fp).not.toBeNull() expect(fp!.sizeBytes).toBe(9) }) + + it('records WAL and SHM sidecar metadata when present', async () => { + await mkdir(TMP_DIR, { recursive: true }) + const dbPath = join(TMP_DIR, 'antigravity.db') + await writeFile(dbPath, 'sqlite-base') + await writeFile(`${dbPath}-wal`, 'wal-data') + await writeFile(`${dbPath}-shm`, 'shm-data') + + const fp = await fingerprintFile(dbPath) + expect(fp).not.toBeNull() + expect(fp!.walMtimeMs).toBeGreaterThan(0) + expect(fp!.walSizeBytes).toBe(8) + expect(fp!.shmMtimeMs).toBeGreaterThan(0) + expect(fp!.shmSizeBytes).toBe(8) + }) + + it('records sidecar metadata for compound DB paths', async () => { + await mkdir(TMP_DIR, { recursive: true }) + const hashDbPath = join(TMP_DIR, 'state.vscdb') + const colonDbPath = join(TMP_DIR, 'opencode.db') + await writeFile(hashDbPath, 'cursor-data') + await writeFile(`${hashDbPath}-wal`, 'cursor-wal') + await writeFile(colonDbPath, 'opencode-data') + await writeFile(`${colonDbPath}-shm`, 'opencode-shm') + + const hashFp = await fingerprintFile(`${hashDbPath}#cursor-ws=__orphan__`) + expect(hashFp).not.toBeNull() + expect(hashFp!.walSizeBytes).toBe(10) + expect(hashFp!.shmSizeBytes).toBeUndefined() + + const colonFp = await fingerprintFile(`${colonDbPath}:ses_abc123`) + expect(colonFp).not.toBeNull() + expect(colonFp!.shmSizeBytes).toBe(12) + expect(colonFp!.walSizeBytes).toBeUndefined() + }) }) // ── reconcileFile ────────────────────────────────────────────────────── @@ -327,6 +366,49 @@ describe('reconcileFile', () => { const current: FileFingerprint = { dev: 2, ino: 100, mtimeMs: 2000, sizeBytes: 8000 } expect(reconcileFile(current, cached)).toEqual({ action: 'modified' }) }) + + it('returns "modified" when WAL sidecar size changes but base is unchanged', () => { + const cached = makeCachedFile({ + fingerprint: { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000, walMtimeMs: 2000, walSizeBytes: 100 }, + }) + const current: FileFingerprint = { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000, walMtimeMs: 2000, walSizeBytes: 120 } + expect(reconcileFile(current, cached)).toEqual({ action: 'modified' }) + }) + + it('returns "modified" when SHM sidecar mtime changes but base is unchanged', () => { + const cached = makeCachedFile({ + fingerprint: { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000, shmMtimeMs: 2000, shmSizeBytes: 100 }, + }) + const current: FileFingerprint = { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000, shmMtimeMs: 3000, shmSizeBytes: 100 } + expect(reconcileFile(current, cached)).toEqual({ action: 'modified' }) + }) + + it('returns "modified" when a WAL sidecar appears or disappears with unchanged base', () => { + const cachedWithoutWal = makeCachedFile({ + fingerprint: { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000 }, + }) + const currentWithWal: FileFingerprint = { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000, walMtimeMs: 2000, walSizeBytes: 100 } + expect(reconcileFile(currentWithWal, cachedWithoutWal)).toEqual({ action: 'modified' }) + + const cachedWithWal = makeCachedFile({ fingerprint: currentWithWal }) + const currentWithoutWal: FileFingerprint = { dev: 1, ino: 100, mtimeMs: 1000, sizeBytes: 5000 } + expect(reconcileFile(currentWithoutWal, cachedWithWal)).toEqual({ action: 'modified' }) + }) + + it('returns "unchanged" when sidecar metadata and base fields match', () => { + const fp: FileFingerprint = { + dev: 1, + ino: 100, + mtimeMs: 1000, + sizeBytes: 5000, + walMtimeMs: 2000, + walSizeBytes: 100, + shmMtimeMs: 1500, + shmSizeBytes: 64, + } + const cached = makeCachedFile({ fingerprint: { ...fp } }) + expect(reconcileFile({ ...fp }, cached)).toEqual({ action: 'unchanged' }) + }) }) // ── mergeCallByDedupKey ──────────────────────────────────────────────── @@ -433,6 +515,60 @@ describe('loadCache validation', () => { expect((await loadCache()).providers).toEqual({}) }) + it('drops a provider with unsafe token counts without discarding other valid providers', async () => { + const unsafeUsage = { + inputTokens: Number.MAX_SAFE_INTEGER + 1, + outputTokens: 0, + cacheCreationInputTokens: 0, + cacheReadInputTokens: 0, + cachedInputTokens: 0, + reasoningTokens: 0, + webSearchRequests: 0, + cacheCreationOneHourTokens: 0, + } + const unsafeCall = { + ...validCallJson(), + provider: 'antigravity', + usage: unsafeUsage, + } + const validCopilotCall = { + ...validCallJson(), + provider: 'copilot', + model: 'gpt-5.5', + } + await writeRawCache({ + version: CACHE_VERSION, + providers: { + antigravity: { + envFingerprint: 'bad', + files: { + '/bad.pb': { + fingerprint: { dev: 1, ino: 2, mtimeMs: 3, sizeBytes: 4 }, + mcpInventory: [], + turns: [{ timestamp: 'x', sessionId: 'bad', userMessage: 'bad', calls: [unsafeCall] }], + }, + }, + }, + copilot: { + envFingerprint: 'good', + durable: true, + files: { + '/durable': { + fingerprint: { dev: 5, ino: 6, mtimeMs: 7, sizeBytes: 8 }, + mcpInventory: [], + turns: [{ timestamp: 'x', sessionId: 'good', userMessage: 'good', calls: [validCopilotCall] }], + }, + }, + }, + }, + }) + + const loaded = await loadCache() + expect(loaded.providers['antigravity']).toBeUndefined() + expect(Object.keys(loaded.providers['copilot']!.files)).toEqual(['/durable']) + expect(loaded.providers['copilot']!.durable).toBe(true) + }) + function validCallJson() { return { provider: 'claude', model: 'm', deduplicationKey: 'k', timestamp: 't', speed: 'standard', @@ -461,6 +597,29 @@ describe('loadCache validation', () => { } } + it('rejects exact unsafe token counts in usage', async () => { + for (const value of EXACT_UNSAFE_TOKEN_COUNTS) { + await writeRawCache(wrapCall({ + usage: { ...validCallJson().usage, inputTokens: Number(value) }, + })) + expect((await loadCache()).providers).toEqual({}) + } + }) + + it('rejects negative token counts in usage', async () => { + await writeRawCache(wrapCall({ + usage: { ...validCallJson().usage, cacheReadInputTokens: -1 }, + })) + expect((await loadCache()).providers).toEqual({}) + }) + + it('rejects fractional token counts in usage', async () => { + await writeRawCache(wrapCall({ + usage: { ...validCallJson().usage, outputTokens: 1.5 }, + })) + expect((await loadCache()).providers).toEqual({}) + }) + it('rejects tools containing non-string element', async () => { await writeRawCache(wrapCall({ tools: ['Read', 42] })) expect((await loadCache()).providers).toEqual({})