Skip to content

Commit b6f03d5

Browse files
kmg0308kmg0308
committed
Preserve Codex usage after session cleanup
1 parent f894599 commit b6f03d5

5 files changed

Lines changed: 216 additions & 5 deletions

File tree

‎Sources/TokenMeterCore/TokenEventCacheStore.swift‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -322,9 +322,13 @@ public final class TokenEventCacheStore: @unchecked Sendable {
322322
}
323323
}
324324

325-
func removeMissingOrigins(originKind: OriginKind, keeping existingPaths: Set<String>) throws {
325+
func removeMissingOrigins(
326+
originKind: OriginKind,
327+
keeping existingPaths: Set<String>,
328+
pruningSources: Set<TokenSource>? = nil
329+
) throws {
326330
try locked {
327-
let paths = try originPaths(originKind: originKind)
331+
let paths = try originPaths(originKind: originKind, pruningSources: pruningSources)
328332
let missingPaths = paths.filter { !existingPaths.contains($0) }
329333
guard !missingPaths.isEmpty else { return }
330334

@@ -499,17 +503,23 @@ public final class TokenEventCacheStore: @unchecked Sendable {
499503
return events
500504
}
501505

502-
private func originPaths(originKind: OriginKind) throws -> [String] {
506+
private func originPaths(originKind: OriginKind, pruningSources: Set<TokenSource>?) throws -> [String] {
503507
var statement: OpaquePointer?
504508
try prepare(
505-
"SELECT origin_path FROM origin_files WHERE origin_kind = ?",
509+
"SELECT origin_path, source FROM origin_files WHERE origin_kind = ?",
506510
into: &statement
507511
)
508512
defer { sqlite3_finalize(statement) }
509513
bind(originKind.rawValue, to: statement, at: 1)
510514

511515
var paths: [String] = []
512516
while sqlite3_step(statement) == SQLITE_ROW {
517+
if let pruningSources {
518+
guard let source = nullableColumnString(statement, 1).flatMap(TokenSource.init(rawValue:)),
519+
pruningSources.contains(source) else {
520+
continue
521+
}
522+
}
513523
paths.append(columnString(statement, 0))
514524
}
515525
return paths

‎Sources/TokenMeterCore/TokenLogScanner.swift‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,11 @@ public final class TokenLogScanner: @unchecked Sendable {
221221
private func cleanupMissingLocalCacheEntries(roots: [RootScan]) {
222222
guard let cacheStore else { return }
223223
let paths = Set(roots.flatMap(\.allFilePaths))
224-
try? cacheStore.removeMissingOrigins(originKind: .localLog, keeping: paths)
224+
try? cacheStore.removeMissingOrigins(
225+
originKind: .localLog,
226+
keeping: paths,
227+
pruningSources: [.claude]
228+
)
225229
}
226230

227231
private func cachedSyncStatus(

‎Sources/TokenMeterSelfTestSupport/ScannerCacheTests.swift‎

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ extension TokenMeterSelfTest {
99
try scannerReusesCachedLocalFilesUntilTheyChange()
1010
try scannerReturnsCachedResultBeforeRefresh()
1111
try scannerCachedResultRestoresSyncStatus()
12+
try scannerKeepsCodexCacheWhenLocalSessionFileIsRemoved()
13+
try scannerPrunesClaudeCacheWhenLocalLogFileIsRemoved()
1214
try scannerAppendsGrowingLocalFileFromCache()
1315
try scannerIgnoresIncrementalClaudeDuplicateRequests()
1416
try scannerFallsBackToFullParseWhenGrowingFileModificationDateRegresses()
@@ -351,6 +353,66 @@ extension TokenMeterSelfTest {
351353
try expect(windowed?.syncStatus.importedEventCount == 1, "windowed cached sync status counts the same event window")
352354
}
353355

356+
static func scannerKeepsCodexCacheWhenLocalSessionFileIsRemoved() throws {
357+
let directory = FileManager.default.temporaryDirectory
358+
.appendingPathComponent(UUID().uuidString, isDirectory: true)
359+
let home = directory.appendingPathComponent("home", isDirectory: true)
360+
let cache = try temporaryCache(in: directory)
361+
let scanner = TokenLogScanner(
362+
homeDirectory: home,
363+
localDevice: TokenDeviceMetadata(id: "mac-a", name: "Mac A"),
364+
cacheStore: cache
365+
)
366+
let logURL = try writeCodexLog(
367+
homeDirectory: home,
368+
fileName: "cached-codex.jsonl",
369+
timestamp: "2026-01-01T00:00:00.000Z",
370+
cwd: "/tmp/codex-project-a",
371+
input: 10,
372+
output: 5
373+
)
374+
375+
let initial = scanner.scan(modifiedAfter: isoDate("2025-12-31T00:00:00.000Z"))
376+
try expect(Aggregation.totalUsage(events: initial.events).total == 15, "initial Codex cache total")
377+
378+
try FileManager.default.removeItem(at: logURL)
379+
let refreshed = scanner.scan(modifiedAfter: isoDate("2025-12-31T00:00:00.000Z"))
380+
try expect(Aggregation.totalUsage(events: refreshed.events).total == 15, "removed Codex session remains available from cache")
381+
try expect(refreshed.codexFileCount == 1, "removed Codex cache still contributes to source status")
382+
383+
let cached = scanner.cachedResult(eventAfter: isoDate("2025-12-31T00:00:00.000Z"))
384+
try expect(Aggregation.totalUsage(events: cached?.events ?? []).total == 15, "cached dashboard result keeps removed Codex usage")
385+
}
386+
387+
static func scannerPrunesClaudeCacheWhenLocalLogFileIsRemoved() throws {
388+
let directory = FileManager.default.temporaryDirectory
389+
.appendingPathComponent(UUID().uuidString, isDirectory: true)
390+
let home = directory.appendingPathComponent("home", isDirectory: true)
391+
let cache = try temporaryCache(in: directory)
392+
let scanner = TokenLogScanner(
393+
homeDirectory: home,
394+
localDevice: TokenDeviceMetadata(id: "mac-a", name: "Mac A"),
395+
cacheStore: cache
396+
)
397+
let logURL = try writeClaudeLog(
398+
homeDirectory: home,
399+
fileName: "removed-claude.jsonl",
400+
timestamp: "2026-01-01T00:00:00.000Z",
401+
cwd: "/tmp/project-a",
402+
sessionId: "session-a",
403+
requestId: "request-a",
404+
input: 10
405+
)
406+
407+
let initial = scanner.scan(modifiedAfter: isoDate("2025-12-31T00:00:00.000Z"))
408+
try expect(Aggregation.totalUsage(events: initial.events).total == 10, "initial Claude cache total")
409+
410+
try FileManager.default.removeItem(at: logURL)
411+
let refreshed = scanner.scan(modifiedAfter: isoDate("2025-12-31T00:00:00.000Z"))
412+
try expect(refreshed.events.isEmpty, "removed Claude log is pruned from cache")
413+
try expect(refreshed.claudeFileCount == 0, "removed Claude cache does not contribute to source status")
414+
}
415+
354416
static func scannerAppendsGrowingLocalFileFromCache() throws {
355417
let directory = FileManager.default.temporaryDirectory
356418
.appendingPathComponent(UUID().uuidString, isDirectory: true)

‎Sources/TokenMeterSelfTestSupport/SyncFolderTests.swift‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@ extension TokenMeterSelfTest {
1717
try syncFolderDeduplicatesAppendedDeviceLedgerRecords()
1818
try syncFolderCancellationSkipsPartialLocalLedgerWrite()
1919
try scannerMergesFreshWindowedSyncEventsWithCachedHistory()
20+
try scannerKeepsSyncedCodexUsageAfterLocalSessionFilesAreRemoved()
21+
try scannerKeepsMergedDeviceUsageAfterOneMacRemovesLocalSessionFiles()
22+
try scannerRebuildsMissingLocalSyncLedgerAfterCodexSessionsAreRemoved()
2023
try scannerRestoresMissingLocalSyncLedgerFromCachedLocalHistory()
2124
try scannerExcludesCachedSyncEventsWhenSyncFolderIsDisabled()
2225
try scannerPrunesCachedSyncEventsWhenDeviceLedgerDisappears()
@@ -219,6 +222,110 @@ extension TokenMeterSelfTest {
219222
try expect(windowed.status.parseErrorCount == 0, "sync import accepts plain ISO timestamps")
220223
}
221224

225+
static func scannerKeepsSyncedCodexUsageAfterLocalSessionFilesAreRemoved() throws {
226+
let directory = FileManager.default.temporaryDirectory
227+
.appendingPathComponent(UUID().uuidString, isDirectory: true)
228+
let home = directory.appendingPathComponent("home", isDirectory: true)
229+
let syncFolder = directory.appendingPathComponent("sync", isDirectory: true)
230+
try FileManager.default.createDirectory(at: syncFolder, withIntermediateDirectories: true)
231+
232+
let device = TokenDeviceMetadata(id: "mac-a", name: "Mac A")
233+
let scanner = TokenLogScanner(homeDirectory: home, localDevice: device)
234+
let codexLog = try writeCodexLog(
235+
homeDirectory: home,
236+
fileName: "session-a.jsonl",
237+
timestamp: "2026-01-01T00:00:00.000Z",
238+
cwd: "/tmp/codex-project-a",
239+
input: 10,
240+
output: 5
241+
)
242+
243+
let initial = scanner.scan(syncFolder: syncFolder, replaceSyncLedger: true)
244+
try expect(Aggregation.totalUsage(events: initial.events).total == 15, "initial Codex sync total")
245+
try expect(initial.syncStatus.exportedEventCount == 1, "initial Codex sync exports one event")
246+
247+
try FileManager.default.removeItem(at: codexLog)
248+
let removed = scanner.scan(syncFolder: syncFolder)
249+
try expect(Aggregation.totalUsage(events: removed.events).total == 15, "synced Codex total survives removed local session file")
250+
try expect(removed.events.first?.rawFilePath.hasPrefix("sync://") == true, "removed local session is restored from sync ledger")
251+
}
252+
253+
static func scannerKeepsMergedDeviceUsageAfterOneMacRemovesLocalSessionFiles() throws {
254+
let directory = FileManager.default.temporaryDirectory
255+
.appendingPathComponent(UUID().uuidString, isDirectory: true)
256+
let homeA = directory.appendingPathComponent("home-a", isDirectory: true)
257+
let homeB = directory.appendingPathComponent("home-b", isDirectory: true)
258+
let syncFolder = directory.appendingPathComponent("sync", isDirectory: true)
259+
try FileManager.default.createDirectory(at: syncFolder, withIntermediateDirectories: true)
260+
261+
let logA = try writeCodexLog(
262+
homeDirectory: homeA,
263+
fileName: "session-a.jsonl",
264+
timestamp: "2026-01-01T00:00:00.000Z",
265+
cwd: "/tmp/codex-project-a",
266+
input: 10
267+
)
268+
try writeCodexLog(
269+
homeDirectory: homeB,
270+
fileName: "session-b.jsonl",
271+
timestamp: "2026-01-01T00:01:00.000Z",
272+
cwd: "/tmp/codex-project-b",
273+
input: 20
274+
)
275+
276+
let deviceA = TokenDeviceMetadata(id: "mac-a", name: "Mac A")
277+
let deviceB = TokenDeviceMetadata(id: "mac-b", name: "Mac B")
278+
let scannerA = TokenLogScanner(homeDirectory: homeA, localDevice: deviceA)
279+
let scannerB = TokenLogScanner(homeDirectory: homeB, localDevice: deviceB)
280+
281+
let firstA = scannerA.scan(syncFolder: syncFolder, replaceSyncLedger: true)
282+
try expect(Aggregation.totalUsage(events: firstA.events).total == 10, "first Mac exports local Codex total")
283+
let firstB = scannerB.scan(syncFolder: syncFolder, replaceSyncLedger: true)
284+
try expect(Aggregation.totalUsage(events: firstB.events).total == 30, "second Mac sees merged Codex total")
285+
286+
try FileManager.default.removeItem(at: logA)
287+
let mergedA = scannerA.scan(syncFolder: syncFolder)
288+
try expect(Aggregation.totalUsage(events: mergedA.events).total == 30, "merged Codex total survives one Mac pruning local sessions")
289+
try expect(Set(mergedA.events.map(\.deviceId)) == ["mac-a", "mac-b"], "merged Codex events keep both device ids")
290+
}
291+
292+
static func scannerRebuildsMissingLocalSyncLedgerAfterCodexSessionsAreRemoved() throws {
293+
let directory = FileManager.default.temporaryDirectory
294+
.appendingPathComponent(UUID().uuidString, isDirectory: true)
295+
let home = directory.appendingPathComponent("home", isDirectory: true)
296+
let syncFolder = directory.appendingPathComponent("sync", isDirectory: true)
297+
try FileManager.default.createDirectory(at: syncFolder, withIntermediateDirectories: true)
298+
299+
let cache = try temporaryCache(in: directory)
300+
let device = TokenDeviceMetadata(id: "mac-a", name: "Mac A")
301+
let scanner = TokenLogScanner(homeDirectory: home, localDevice: device, cacheStore: cache)
302+
let codexLog = try writeCodexLog(
303+
homeDirectory: home,
304+
fileName: "session-a.jsonl",
305+
timestamp: "2026-01-01T00:00:00.000Z",
306+
cwd: "/tmp/codex-project-a",
307+
input: 10,
308+
output: 5
309+
)
310+
311+
let initial = scanner.scan(syncFolder: syncFolder, replaceSyncLedger: true)
312+
try expect(Aggregation.totalUsage(events: initial.events).total == 15, "initial Codex cache warms local history")
313+
try expect(try syncLedgerLineCount(syncFolder: syncFolder, deviceId: device.id) == 1, "initial Codex local ledger exists")
314+
315+
let ledgerURL = syncFolder
316+
.appendingPathComponent("devices", isDirectory: true)
317+
.appendingPathComponent("\(device.id).jsonl")
318+
try FileManager.default.removeItem(at: codexLog)
319+
try FileManager.default.removeItem(at: ledgerURL)
320+
321+
let restored = scanner.scan(syncFolder: syncFolder)
322+
try expect(Aggregation.totalUsage(events: restored.events).total == 15, "Codex cache survives local session removal")
323+
try expect(
324+
try syncLedgerLineCount(syncFolder: syncFolder, deviceId: device.id) == 1,
325+
"missing Codex local ledger is rebuilt from cached history"
326+
)
327+
}
328+
222329
static func syncFolderCountsMalformedCanonicalTimestampAsParseError() throws {
223330
let directory = FileManager.default.temporaryDirectory
224331
.appendingPathComponent(UUID().uuidString, isDirectory: true)

‎Sources/TokenMeterSelfTestSupport/TestSupport.swift‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,34 @@ extension TokenMeterSelfTest {
4747
return url
4848
}
4949

50+
@discardableResult
51+
static func writeCodexLog(
52+
homeDirectory: URL,
53+
fileName: String,
54+
timestamp: String,
55+
cwd: String,
56+
model: String = "gpt-5.5",
57+
input: Int,
58+
output: Int = 0,
59+
reasoning: Int = 0
60+
) throws -> URL {
61+
let sessionDirectory = homeDirectory
62+
.appendingPathComponent(".codex", isDirectory: true)
63+
.appendingPathComponent("sessions", isDirectory: true)
64+
.appendingPathComponent("2026", isDirectory: true)
65+
.appendingPathComponent("01", isDirectory: true)
66+
.appendingPathComponent("01", isDirectory: true)
67+
try FileManager.default.createDirectory(at: sessionDirectory, withIntermediateDirectories: true)
68+
let total = input + output
69+
let content = """
70+
{"timestamp":"\(timestamp)","payload":{"type":"session_meta","cwd":"\(cwd)","model":"\(model)"}}
71+
{"timestamp":"\(timestamp)","payload":{"info":{"last_token_usage":{"input_tokens":\(input),"cached_input_tokens":0,"output_tokens":\(output),"reasoning_output_tokens":\(reasoning),"total_tokens":\(total)}}}}
72+
"""
73+
let url = sessionDirectory.appendingPathComponent(fileName)
74+
try content.write(to: url, atomically: true, encoding: .utf8)
75+
return url
76+
}
77+
5078
static func appendClaudeLogLine(
5179
to url: URL,
5280
timestamp: String,

0 commit comments

Comments
 (0)