Skip to content

Commit e78f43c

Browse files
ozymandiashhozymandiashh
committed
fix(menubar): clamp unsafe Antigravity token counts
1 parent 163013f commit e78f43c

10 files changed

Lines changed: 841 additions & 83 deletions

File tree

‎mac/Sources/CodeBurnMenubar/Data/MenubarPayload.swift‎

Lines changed: 59 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,37 @@
11
import Foundation
22

3+
private let maxSafeTokenCount = 9_007_199_254_740_991
4+
5+
private func sanitizedTokenCount(_ value: Int) -> Int {
6+
value >= 0 && value <= maxSafeTokenCount ? value : 0
7+
}
8+
9+
private extension KeyedDecodingContainer {
10+
func decodeTokenCount(forKey key: Key) -> Int {
11+
guard contains(key) else { return 0 }
12+
if let value = try? decode(Int.self, forKey: key) {
13+
return sanitizedTokenCount(value)
14+
}
15+
if let value = try? decode(Double.self, forKey: key),
16+
value.isFinite,
17+
value >= 0,
18+
value <= Double(maxSafeTokenCount),
19+
value.rounded(.towardZero) == value {
20+
return Int(value)
21+
}
22+
if let value = try? decode(String.self, forKey: key) {
23+
let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
24+
guard !trimmed.isEmpty,
25+
trimmed.utf8.allSatisfy({ $0 >= 48 && $0 <= 57 }),
26+
let parsed = Int(trimmed) else {
27+
return 0
28+
}
29+
return sanitizedTokenCount(parsed)
30+
}
31+
return 0
32+
}
33+
}
34+
335
/// Shape of `codeburn status --format menubar-json --period <period>`.
436
/// `current` is scoped to the requested period; the whole payload reflects that slice.
537
struct MenubarPayload: Codable, Sendable {
@@ -32,8 +64,8 @@ struct DailyModelBreakdown: Codable, Sendable {
3264
cost = try c.decode(Double.self, forKey: .cost)
3365
savingsUSD = try c.decodeIfPresent(Double.self, forKey: .savingsUSD) ?? 0
3466
calls = try c.decode(Int.self, forKey: .calls)
35-
inputTokens = try c.decode(Int.self, forKey: .inputTokens)
36-
outputTokens = try c.decode(Int.self, forKey: .outputTokens)
67+
inputTokens = c.decodeTokenCount(forKey: .inputTokens)
68+
outputTokens = c.decodeTokenCount(forKey: .outputTokens)
3769
}
3870
}
3971

@@ -66,10 +98,10 @@ extension DailyHistoryEntry {
6698
cost = try c.decode(Double.self, forKey: .cost)
6799
savingsUSD = try c.decodeIfPresent(Double.self, forKey: .savingsUSD) ?? 0
68100
calls = try c.decode(Int.self, forKey: .calls)
69-
inputTokens = try c.decode(Int.self, forKey: .inputTokens)
70-
outputTokens = try c.decode(Int.self, forKey: .outputTokens)
71-
cacheReadTokens = try c.decode(Int.self, forKey: .cacheReadTokens)
72-
cacheWriteTokens = try c.decode(Int.self, forKey: .cacheWriteTokens)
101+
inputTokens = c.decodeTokenCount(forKey: .inputTokens)
102+
outputTokens = c.decodeTokenCount(forKey: .outputTokens)
103+
cacheReadTokens = c.decodeTokenCount(forKey: .cacheReadTokens)
104+
cacheWriteTokens = c.decodeTokenCount(forKey: .cacheWriteTokens)
73105
topModels = try c.decodeIfPresent([DailyModelBreakdown].self, forKey: .topModels) ?? []
74106
}
75107
}
@@ -144,8 +176,8 @@ extension CurrentBlock {
144176
calls = try c.decode(Int.self, forKey: .calls)
145177
sessions = try c.decode(Int.self, forKey: .sessions)
146178
oneShotRate = try c.decodeIfPresent(Double.self, forKey: .oneShotRate)
147-
inputTokens = try c.decode(Int.self, forKey: .inputTokens)
148-
outputTokens = try c.decode(Int.self, forKey: .outputTokens)
179+
inputTokens = c.decodeTokenCount(forKey: .inputTokens)
180+
outputTokens = c.decodeTokenCount(forKey: .outputTokens)
149181
cacheHitPercent = try c.decodeIfPresent(Double.self, forKey: .cacheHitPercent) ?? 0
150182
codexCredits = try c.decodeIfPresent(Double.self, forKey: .codexCredits)
151183
topActivities = try c.decodeIfPresent([ActivityEntry].self, forKey: .topActivities) ?? []
@@ -174,6 +206,23 @@ struct LocalModelSavingsByModel: Codable, Sendable {
174206
let outputTokens: Int
175207
}
176208

209+
extension LocalModelSavingsByModel {
210+
enum CodingKeys: String, CodingKey {
211+
case name, calls, actualUSD, savingsUSD, baselineModel, inputTokens, outputTokens
212+
}
213+
214+
init(from decoder: Decoder) throws {
215+
let c = try decoder.container(keyedBy: CodingKeys.self)
216+
name = try c.decode(String.self, forKey: .name)
217+
calls = try c.decode(Int.self, forKey: .calls)
218+
actualUSD = try c.decode(Double.self, forKey: .actualUSD)
219+
savingsUSD = try c.decode(Double.self, forKey: .savingsUSD)
220+
baselineModel = try c.decode(String.self, forKey: .baselineModel)
221+
inputTokens = c.decodeTokenCount(forKey: .inputTokens)
222+
outputTokens = c.decodeTokenCount(forKey: .outputTokens)
223+
}
224+
}
225+
177226
struct LocalModelSavingsByProvider: Codable, Sendable {
178227
let name: String
179228
let calls: Int
@@ -260,8 +309,8 @@ struct SessionDetailEntry: Codable, Sendable {
260309
cost = try c.decode(Double.self, forKey: .cost)
261310
savingsUSD = try c.decodeIfPresent(Double.self, forKey: .savingsUSD) ?? 0
262311
calls = try c.decode(Int.self, forKey: .calls)
263-
inputTokens = try c.decode(Int.self, forKey: .inputTokens)
264-
outputTokens = try c.decode(Int.self, forKey: .outputTokens)
312+
inputTokens = c.decodeTokenCount(forKey: .inputTokens)
313+
outputTokens = c.decodeTokenCount(forKey: .outputTokens)
265314
date = try c.decode(String.self, forKey: .date)
266315
models = try c.decodeIfPresent([SessionModelEntry].self, forKey: .models) ?? []
267316
}

‎mac/Tests/CodeBurnMenubarTests/MenubarPayloadDecodeTests.swift‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
import Foundation
2+
import Testing
3+
@testable import CodeBurnMenubar
4+
5+
@Suite("MenubarPayload decode")
6+
struct MenubarPayloadDecodeTests {
7+
private func decode(_ json: String) throws -> MenubarPayload {
8+
try JSONDecoder().decode(MenubarPayload.self, from: Data(json.utf8))
9+
}
10+
11+
@Test("huge malformed token fields decode as zero")
12+
func hugeMalformedTokenFieldsDecodeAsZero() throws {
13+
let payload = try decode("""
14+
{
15+
"generated": "2026-06-22T00:00:00Z",
16+
"current": {
17+
"label": "Today",
18+
"cost": 1.0,
19+
"calls": 1,
20+
"sessions": 1,
21+
"inputTokens": 18446744073709527000,
22+
"outputTokens": "221360928884514260000",
23+
"cacheHitPercent": 0,
24+
"localModelSavings": {
25+
"totalUSD": 1,
26+
"calls": 1,
27+
"byModel": [{
28+
"name": "local",
29+
"calls": 1,
30+
"actualUSD": 0,
31+
"savingsUSD": 1,
32+
"baselineModel": "paid",
33+
"inputTokens": 221360928884514260000,
34+
"outputTokens": -1
35+
}],
36+
"byProvider": []
37+
},
38+
"topProjects": [{
39+
"name": "project",
40+
"cost": 1,
41+
"savingsUSD": 0,
42+
"sessions": 1,
43+
"avgCostPerSession": 1,
44+
"sessionDetails": [{
45+
"cost": 1,
46+
"savingsUSD": 0,
47+
"calls": 1,
48+
"inputTokens": 18446744073709527000,
49+
"outputTokens": 1.5,
50+
"date": "2026-06-21",
51+
"models": []
52+
}]
53+
}]
54+
},
55+
"optimize": {
56+
"findingCount": 0,
57+
"savingsUSD": 0,
58+
"topFindings": []
59+
},
60+
"history": {
61+
"daily": [{
62+
"date": "2026-06-21",
63+
"cost": 1,
64+
"savingsUSD": 0,
65+
"calls": 1,
66+
"inputTokens": 18446744073709527000,
67+
"outputTokens": -1,
68+
"cacheReadTokens": 1.5,
69+
"cacheWriteTokens": "221360928884514260000",
70+
"topModels": [{
71+
"name": "Gemini 3.5 Flash",
72+
"cost": 1,
73+
"savingsUSD": 0,
74+
"calls": 1,
75+
"inputTokens": 221360928884514260000,
76+
"outputTokens": -1
77+
}]
78+
}]
79+
}
80+
}
81+
""")
82+
83+
#expect(payload.current.inputTokens == 0)
84+
#expect(payload.current.outputTokens == 0)
85+
#expect(payload.current.localModelSavings.byModel[0].inputTokens == 0)
86+
#expect(payload.current.localModelSavings.byModel[0].outputTokens == 0)
87+
#expect(payload.current.topProjects[0].sessionDetails[0].inputTokens == 0)
88+
#expect(payload.current.topProjects[0].sessionDetails[0].outputTokens == 0)
89+
#expect(payload.history.daily[0].inputTokens == 0)
90+
#expect(payload.history.daily[0].outputTokens == 0)
91+
#expect(payload.history.daily[0].cacheReadTokens == 0)
92+
#expect(payload.history.daily[0].cacheWriteTokens == 0)
93+
#expect(payload.history.daily[0].topModels[0].inputTokens == 0)
94+
#expect(payload.history.daily[0].topModels[0].outputTokens == 0)
95+
}
96+
97+
@Test("huge non-token integers remain strict decode failures")
98+
func hugeNonTokenIntegersRemainStrictDecodeFailures() {
99+
var didThrow = false
100+
do {
101+
_ = try decode("""
102+
{
103+
"generated": "2026-06-22T00:00:00Z",
104+
"current": {
105+
"label": "Today",
106+
"cost": 1.0,
107+
"calls": 18446744073709551615,
108+
"sessions": 1,
109+
"inputTokens": 1,
110+
"outputTokens": 1
111+
},
112+
"optimize": {
113+
"findingCount": 0,
114+
"savingsUSD": 0,
115+
"topFindings": []
116+
},
117+
"history": { "daily": [] }
118+
}
119+
""")
120+
} catch {
121+
didThrow = true
122+
}
123+
124+
#expect(didThrow)
125+
}
126+
}

‎src/daily-cache.ts‎

Lines changed: 76 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -5,14 +5,11 @@ import { homedir } from 'os'
55
import { join } from 'path'
66
import type { DateRange, ProjectSummary } from './types.js'
77

8-
// Bumped to 8: local-model savings accounting is now part of the daily rollup
9-
// (savingsUSD per day / per model / per category / per provider). Stale entries
10-
// computed by older binaries lack those fields, so MIN_SUPPORTED_VERSION is
11-
// also raised to 8 to force a full re-hydration. The `savingsConfigHash` field
12-
// is invalidated separately when the user changes their `localModelSavings`
13-
// mapping so historical "saved" totals stay in sync with the active baseline.
14-
export const DAILY_CACHE_VERSION = 8
15-
const MIN_SUPPORTED_VERSION = 8
8+
// Bumped to 9: token counts are now required to be non-negative safe integers.
9+
// Older daily rollups may contain unsafe Antigravity uint64-underflow values
10+
// that can crash the menubar Swift decoder, so force a full re-hydration.
11+
export const DAILY_CACHE_VERSION = 9
12+
const MIN_SUPPORTED_VERSION = 9
1613
const DAILY_CACHE_FILENAME = 'daily-cache.json'
1714

1815
export type DailyEntry = {
@@ -71,23 +68,80 @@ function isMigratableCache(parsed: unknown): parsed is { version: number; lastCo
7168
return c.version >= MIN_SUPPORTED_VERSION && c.version <= DAILY_CACHE_VERSION
7269
}
7370

74-
function migrateDays(days: Record<string, unknown>[]): DailyEntry[] {
75-
return days.map(d => ({
71+
function safeTokenCount(value: unknown): number | null {
72+
if (value === undefined) return 0
73+
return typeof value === 'number' && Number.isSafeInteger(value) && value >= 0 ? value : null
74+
}
75+
76+
function migrateModels(models: unknown): DailyEntry['models'] | null {
77+
if (models === undefined) return {}
78+
if (!models || typeof models !== 'object' || Array.isArray(models)) return null
79+
const migrated: DailyEntry['models'] = {}
80+
for (const [name, raw] of Object.entries(models as Record<string, unknown>)) {
81+
if (!raw || typeof raw !== 'object' || Array.isArray(raw)) return null
82+
const m = raw as Record<string, unknown>
83+
const inputTokens = safeTokenCount(m.inputTokens)
84+
const outputTokens = safeTokenCount(m.outputTokens)
85+
const cacheReadTokens = safeTokenCount(m.cacheReadTokens)
86+
const cacheWriteTokens = safeTokenCount(m.cacheWriteTokens)
87+
if (
88+
inputTokens === null ||
89+
outputTokens === null ||
90+
cacheReadTokens === null ||
91+
cacheWriteTokens === null
92+
) return null
93+
migrated[name] = {
94+
calls: (m.calls as number) ?? 0,
95+
cost: (m.cost as number) ?? 0,
96+
savingsUSD: (m.savingsUSD as number) ?? 0,
97+
inputTokens,
98+
outputTokens,
99+
cacheReadTokens,
100+
cacheWriteTokens,
101+
}
102+
}
103+
return migrated
104+
}
105+
106+
function migrateDay(d: Record<string, unknown>): DailyEntry | null {
107+
const inputTokens = safeTokenCount(d.inputTokens)
108+
const outputTokens = safeTokenCount(d.outputTokens)
109+
const cacheReadTokens = safeTokenCount(d.cacheReadTokens)
110+
const cacheWriteTokens = safeTokenCount(d.cacheWriteTokens)
111+
const models = migrateModels(d.models)
112+
if (
113+
inputTokens === null ||
114+
outputTokens === null ||
115+
cacheReadTokens === null ||
116+
cacheWriteTokens === null ||
117+
models === null
118+
) return null
119+
return {
76120
date: d.date as string,
77121
cost: (d.cost as number) ?? 0,
78122
savingsUSD: (d.savingsUSD as number) ?? 0,
79123
calls: (d.calls as number) ?? 0,
80124
sessions: (d.sessions as number) ?? 0,
81-
inputTokens: (d.inputTokens as number) ?? 0,
82-
outputTokens: (d.outputTokens as number) ?? 0,
83-
cacheReadTokens: (d.cacheReadTokens as number) ?? 0,
84-
cacheWriteTokens: (d.cacheWriteTokens as number) ?? 0,
125+
inputTokens,
126+
outputTokens,
127+
cacheReadTokens,
128+
cacheWriteTokens,
85129
editTurns: (d.editTurns as number) ?? 0,
86130
oneShotTurns: (d.oneShotTurns as number) ?? 0,
87-
models: (d.models as DailyEntry['models']) ?? {},
131+
models,
88132
categories: (d.categories as DailyEntry['categories']) ?? {},
89133
providers: (d.providers as DailyEntry['providers']) ?? {},
90-
}))
134+
}
135+
}
136+
137+
function migrateDays(days: Record<string, unknown>[]): DailyEntry[] | null {
138+
const migrated: DailyEntry[] = []
139+
for (const day of days) {
140+
const entry = migrateDay(day)
141+
if (!entry) return null
142+
migrated.push(entry)
143+
}
144+
return migrated
91145
}
92146

93147
async function backupOldCache(path: string, version: number): Promise<void> {
@@ -102,11 +156,16 @@ export async function loadDailyCache(): Promise<DailyCache> {
102156
const raw = await readFile(path, 'utf-8')
103157
const parsed: unknown = JSON.parse(raw)
104158
if (isMigratableCache(parsed)) {
159+
const days = migrateDays(parsed.days)
160+
if (!days) {
161+
await backupOldCache(path, parsed.version).catch(() => {})
162+
return emptyCache()
163+
}
105164
const migrated: DailyCache = {
106165
version: DAILY_CACHE_VERSION,
107166
savingsConfigHash: parsed.savingsConfigHash ?? '',
108167
lastComputedDate: parsed.lastComputedDate,
109-
days: migrateDays(parsed.days),
168+
days,
110169
}
111170
if (parsed.version < DAILY_CACHE_VERSION) {
112171
await saveDailyCache(migrated).catch(() => {})

0 commit comments

Comments
 (0)