fix: make cost usage scans cancellation-aware

Author: steipete

CHANGELOG.md

@@ -12,6 +12,7 @@
 - Codex: show Codex Spark model-specific usage as an optional extra quota lane (#1195, fixes #1177). Thanks @LeoLin990405!
 
 ### Fixed
+- Cost history: make token-cost JSONL scans cancellation-aware so quitting, forced refreshes, and account switches can stop stale scans sooner.
 - Codex: show captured `codex login` output when managed Add Account fails so users can recover from account-selection or OAuth failures (#1199). Thanks @chapati23!
 - Claude: hide the obsolete Design quota lane now that Claude Design shares the main Claude usage limit (#1197).
 - Menu bar: coalesce visible-menu rebuilds and reduce hover highlight work so the dropdown stays responsive on macOS 26.5 (#1196).

Sources/CodexBarCore/CostUsageFetcher.swift

@@ -98,12 +98,18 @@ public struct CostUsageFetcher: Sendable {
         if forceRefresh {
             options.refreshMinIntervalSeconds = 0
         }
-        var daily = CostUsageScanner.loadDailyReport(
+        let checkCancellation: CostUsageScanner.CancellationCheck = {
+            try Task.checkCancellation()
+        }
+        try Task.checkCancellation()
+        var daily = try CostUsageScanner.loadDailyReportCancellable(
             provider: provider,
             since: since,
             until: until,
             now: now,
-            options: options)
+            options: options,
+            checkCancellation: checkCancellation)
+        try Task.checkCancellation()
 
         if provider == .vertexai,
            !allowVertexClaudeFallback,
@@ -112,12 +118,14 @@ public struct CostUsageFetcher: Sendable {
         {
             var fallback = options
             fallback.claudeLogProviderFilter = .all
-            daily = CostUsageScanner.loadDailyReport(
+            daily = try CostUsageScanner.loadDailyReportCancellable(
                 provider: provider,
                 since: since,
                 until: until,
                 now: now,
-                options: fallback)
+                options: fallback,
+                checkCancellation: checkCancellation)
+            try Task.checkCancellation()
         }
 
         if provider == .codex || provider == .claude {
@@ -128,12 +136,14 @@ public struct CostUsageFetcher: Sendable {
             if forceRefresh {
                 piOptions.refreshMinIntervalSeconds = 0
             }
-            let piReport = PiSessionCostScanner.loadDailyReport(
+            let piReport = try PiSessionCostScanner.loadDailyReportCancellable(
                 provider: provider,
                 since: since,
                 until: until,
                 now: now,
-                options: piOptions)
+                options: piOptions,
+                checkCancellation: checkCancellation)
+            try Task.checkCancellation()
             daily = CostUsageDailyReport.merged([daily, piReport])
         }
 

Sources/CodexBarCore/Generated/CodexParserHash.generated.swift

@@ -1,5 +1,5 @@
 // Generated by Scripts/regenerate-codex-parser-hash.sh. Do not edit by hand.
 
 enum CodexParserHash {
-    static let value = "5387f958e9e06f7d"
+    static let value = "0d6c12b99dd77d4e"
 }

Sources/CodexBarCore/PiSessionCostScanner.swift

@@ -36,6 +36,13 @@ enum PiSessionCostScanner {
         let cacheRoot: URL?
     }
 
+    private struct ScanContext {
+        let range: CostUsageScanner.CostUsageDayRange
+        let forceRescan: Bool
+        let pricingContext: ModelsDevPricingContext
+        let checkCancellation: CostUsageScanner.CancellationCheck?
+    }
+
     private static let costScale = 1_000_000_000.0
     private static let maxLineBytes = 16 * 1024 * 1024
     private static let maxSafeRoundedInt = Double(Int.max) - 1
@@ -46,6 +53,24 @@ enum PiSessionCostScanner {
         until: Date,
         now: Date = Date(),
         options: Options = Options()) -> CostUsageDailyReport
+    {
+        (
+            try? self.loadDailyReportCancellable(
+                provider: provider,
+                since: since,
+                until: until,
+                now: now,
+                options: options,
+                checkCancellation: nil)) ?? CostUsageDailyReport(data: [], summary: nil)
+    }
+
+    static func loadDailyReportCancellable(
+        provider: UsageProvider,
+        since: Date,
+        until: Date,
+        now: Date = Date(),
+        options: Options = Options(),
+        checkCancellation: CostUsageScanner.CancellationCheck?) throws -> CostUsageDailyReport
     {
         guard provider == .codex || provider == .claude else {
             return CostUsageDailyReport(data: [], summary: nil)
@@ -66,19 +91,23 @@ enum PiSessionCostScanner {
             || nowMs - cache.lastScanUnixMs > refreshMs
 
         if shouldRefresh {
+            try checkCancellation?()
             let root = self.defaultPiSessionsRoot(options: options)
             let startCutoff = self.dateFromDayKey(range.scanSinceKey) ?? since
             let files = self.listPiSessionFiles(root: root, startCutoffLocal: startCutoff)
             let filePathsInScan = Set(files.map(\.path))
 
             for fileURL in files {
-                self.scanPiSessionFile(
+                try self.scanPiSessionFile(
                     fileURL: fileURL,
-                    range: range,
-                    forceRescan: options.forceRescan || windowExpanded,
-                    pricingContext: pricingContext,
-                    cache: &cache)
+                    cache: &cache,
+                    context: ScanContext(
+                        range: range,
+                        forceRescan: options.forceRescan || windowExpanded,
+                        pricingContext: pricingContext,
+                        checkCancellation: checkCancellation))
             }
+            try checkCancellation?()
 
             for key in cache.files.keys where !filePathsInScan.contains(key) {
                 if let old = cache.files[key] {
@@ -93,6 +122,7 @@ enum PiSessionCostScanner {
             cache.scanSinceKey = range.scanSinceKey
             cache.scanUntilKey = range.scanUntilKey
             cache.lastScanUnixMs = nowMs
+            try checkCancellation?()
             PiSessionCostCacheIO.save(cache: cache, cacheRoot: options.cacheRoot)
         }
 
@@ -176,11 +206,11 @@ enum PiSessionCostScanner {
 
     private static func scanPiSessionFile(
         fileURL: URL,
-        range: CostUsageScanner.CostUsageDayRange,
-        forceRescan: Bool,
-        pricingContext: ModelsDevPricingContext,
-        cache: inout PiSessionCostCache)
+        cache: inout PiSessionCostCache,
+        context: ScanContext)
+        throws
     {
+        try context.checkCancellation?()
         let path = fileURL.path
         let attrs = (try? FileManager.default.attributesOfItem(atPath: path)) ?? [:]
         let mtime = (attrs[.modificationDate] as? Date)?.timeIntervalSince1970 ?? 0
@@ -192,26 +222,27 @@ enum PiSessionCostScanner {
         }
 
         let cached = cache.files[path]
-        if !forceRescan,
+        if !context.forceRescan,
            let cached,
            cached.mtimeUnixMs == mtimeMs,
            cached.size == size
         {
             return
         }
 
-        if !forceRescan,
+        if !context.forceRescan,
            let cached,
            size > cached.size,
            cached.parsedBytes > 0,
            cached.parsedBytes <= size
         {
-            let delta = self.parsePiSessionFile(
+            let delta = try self.parsePiSessionFile(
                 fileURL: fileURL,
-                range: range,
+                range: context.range,
                 startOffset: cached.parsedBytes,
                 initialModelContext: cached.lastModelContext,
-                pricingContext: pricingContext)
+                pricingContext: context.pricingContext,
+                checkCancellation: context.checkCancellation)
             if !delta.contributions.isEmpty {
                 self.applyContributions(
                     daysByProvider: &cache.daysByProvider,
@@ -235,10 +266,11 @@ enum PiSessionCostScanner {
                 sign: -1)
         }
 
-        let parsed = self.parsePiSessionFile(
+        let parsed = try self.parsePiSessionFile(
             fileURL: fileURL,
-            range: range,
-            pricingContext: pricingContext)
+            range: context.range,
+            pricingContext: context.pricingContext,
+            checkCancellation: context.checkCancellation)
         if !parsed.contributions.isEmpty {
             self.applyContributions(daysByProvider: &cache.daysByProvider, contributions: parsed.contributions, sign: 1)
         }
@@ -256,7 +288,8 @@ enum PiSessionCostScanner {
         range: CostUsageScanner.CostUsageDayRange,
         startOffset: Int64 = 0,
         initialModelContext: PiModelContext? = nil,
-        pricingContext: ModelsDevPricingContext? = nil) -> ParseResult
+        pricingContext: ModelsDevPricingContext? = nil,
+        checkCancellation: CostUsageScanner.CancellationCheck? = nil) throws -> ParseResult
     {
         var currentModelContext = initialModelContext
         var contributions: [String: [String: [String: PiPackedUsage]]] = [:]
@@ -292,41 +325,49 @@ enum PiSessionCostScanner {
             }
         }
 
-        let parsedBytes = (try? CostUsageJsonl.scan(
-            fileURL: fileURL,
-            offset: startOffset,
-            maxLineBytes: Self.maxLineBytes,
-            prefixBytes: Self.maxLineBytes,
-            onLine: { line in
-                guard !line.bytes.isEmpty, !line.wasTruncated else { return }
-                autoreleasepool {
-                    guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any]
-                    else { return }
-                    guard let type = object["type"] as? String else { return }
-
-                    if type == "model_change" {
-                        currentModelContext = self.modelContext(from: object)
-                        return
+        let parsedBytes: Int64
+        do {
+            parsedBytes = try CostUsageJsonl.scan(
+                fileURL: fileURL,
+                offset: startOffset,
+                maxLineBytes: Self.maxLineBytes,
+                prefixBytes: Self.maxLineBytes,
+                checkCancellation: checkCancellation,
+                onLine: { line in
+                    guard !line.bytes.isEmpty, !line.wasTruncated else { return }
+                    autoreleasepool {
+                        guard let object = (try? JSONSerialization.jsonObject(with: line.bytes)) as? [String: Any]
+                        else { return }
+                        guard let type = object["type"] as? String else { return }
+
+                        if type == "model_change" {
+                            currentModelContext = self.modelContext(from: object)
+                            return
+                        }
+
+                        guard type == "message", let message = object["message"] as? [String: Any] else { return }
+                        guard (message["role"] as? String) == "assistant" else { return }
+
+                        let identity = self.resolveAssistantIdentity(
+                            entry: object,
+                            message: message,
+                            fallback: currentModelContext)
+                        guard let identity else { return }
+                        guard let date = self.timestampDate(entry: object, message: message) else { return }
+                        let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date)
+                        let usage = self.extractUsage(
+                            provider: identity.provider,
+                            modelName: identity.modelName,
+                            message: message,
+                            pricingContext: pricingContext)
+                        add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage)
                     }
-
-                    guard type == "message", let message = object["message"] as? [String: Any] else { return }
-                    guard (message["role"] as? String) == "assistant" else { return }
-
-                    let identity = self.resolveAssistantIdentity(
-                        entry: object,
-                        message: message,
-                        fallback: currentModelContext)
-                    guard let identity else { return }
-                    guard let date = self.timestampDate(entry: object, message: message) else { return }
-                    let dayKey = CostUsageScanner.CostUsageDayRange.dayKey(from: date)
-                    let usage = self.extractUsage(
-                        provider: identity.provider,
-                        modelName: identity.modelName,
-                        message: message,
-                        pricingContext: pricingContext)
-                    add(provider: identity.provider, dayKey: dayKey, modelName: identity.modelName, usage: usage)
-                }
-            })) ?? startOffset
+                })
+        } catch is CancellationError {
+            throw CancellationError()
+        } catch {
+            parsedBytes = startOffset
+        }
 
         return ParseResult(
             contributions: contributions,

Sources/CodexBarCore/Vendored/CostUsage/CostUsageJsonl.swift

@@ -14,6 +14,25 @@ enum CostUsageJsonl {
         prefixBytes: Int,
         onLine: (Line) -> Void) throws
         -> Int64
+    {
+        try self.scan(
+            fileURL: fileURL,
+            offset: offset,
+            maxLineBytes: maxLineBytes,
+            prefixBytes: prefixBytes,
+            checkCancellation: nil,
+            onLine: onLine)
+    }
+
+    @discardableResult
+    static func scan(
+        fileURL: URL,
+        offset: Int64 = 0,
+        maxLineBytes: Int,
+        prefixBytes: Int,
+        checkCancellation: (() throws -> Void)? = nil,
+        onLine: (Line) -> Void) throws
+        -> Int64
     {
         let handle = try FileHandle(forReadingFrom: fileURL)
         defer { try? handle.close() }
@@ -53,12 +72,14 @@ enum CostUsageJsonl {
         }
 
         while true {
+            try checkCancellation?()
             let chunk = try handle.read(upToCount: 256 * 1024) ?? Data()
             if chunk.isEmpty {
                 flushLine()
                 break
             }
 
+            try checkCancellation?()
             bytesRead += Int64(chunk.count)
             chunk.withUnsafeBytes { rawBuffer in
                 guard let base = rawBuffer.bindMemory(to: UInt8.self).baseAddress else { return }
@@ -76,6 +97,7 @@ enum CostUsageJsonl {
                     appendSegment(base.advanced(by: segmentStart), count: rawBuffer.count - segmentStart)
                 }
             }
+            try checkCancellation?()
         }
 
         return startOffset + bytesRead