Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
69 changes: 59 additions & 10 deletions mac/Sources/CodeBurnMenubar/Data/MenubarPayload.swift
Original file line number Diff line number Diff line change
@@ -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 <period>`.
/// `current` is scoped to the requested period; the whole payload reflects that slice.
struct MenubarPayload: Codable, Sendable {
Expand Down Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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) ?? []
}
}
Expand Down Expand Up @@ -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) ?? []
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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) ?? []
}
Expand Down
126 changes: 126 additions & 0 deletions mac/Tests/CodeBurnMenubarTests/MenubarPayloadDecodeTests.swift
Original file line number Diff line number Diff line change
@@ -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)
}
}
84 changes: 75 additions & 9 deletions src/daily-cache.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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<string, unknown>[]): 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<string, unknown>)) {
if (!raw || typeof raw !== 'object' || Array.isArray(raw)) return null
const m = raw as Record<string, unknown>
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<string, unknown>): 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<string, unknown>[]): 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<void> {
Expand All @@ -106,11 +167,16 @@ export async function loadDailyCache(): Promise<DailyCache> {
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(() => {})
Expand Down
Loading
Loading