fix(hermes): harden attribution and database reads

Sources/TokiAgentCore/AgentSnapshotBuilder.swift

@@ -85,15 +85,19 @@ struct AgentSnapshotBuilder: AgentSnapshotBuilding {
             key: configuration.encryptionKey)
         let tokenReplacementCoverages = readerUsages.flatMap(\.usage.tokenReplacementCoverages)
 
-        let tokenEvents = readerUsages
+        let usageEvents = readerUsages
             .flatMap(\.usage.tokenEvents)
             .filter { event in
                 event.timestamp >= coveredFrom
                     && event.timestamp < coveredTo
                     && !tokenReplacementCoverages.contains { $0.replaces(event) }
             }
+        let tokenEvents = usageEvents
             .compactMap(remoteTokenEvent)
             .sorted(by: tokenEventSort)
+        let costEvents = usageEvents
+            .compactMap(remoteCostEvent)
+            .sorted(by: costEventSort)
 
         let activityEvents = readerUsages
             .flatMap { readerUsage in
@@ -120,6 +124,7 @@ struct AgentSnapshotBuilder: AgentSnapshotBuilding {
             coveredFrom: coveredFrom,
             coveredTo: coveredTo,
             tokenEvents: tokenEvents,
+            costEvents: costEvents.isEmpty ? nil : costEvents,
             activityEvents: activityEvents)
     }
 
@@ -129,6 +134,7 @@ struct AgentSnapshotBuilder: AgentSnapshotBuilding {
             coveredFrom: snapshot.coveredFrom,
             coveredTo: snapshot.coveredTo,
             tokenEvents: snapshot.tokenEvents,
+            costEvents: snapshot.costEvents,
             activityEvents: snapshot.activityEvents)
         return try SnapshotCipher.digest(TokiSyncCoding.makeEncoder().encode(content))
     }
@@ -411,7 +417,15 @@ private extension AgentSnapshotBuilder {
         if lhs.outputTokens != rhs.outputTokens { return lhs.outputTokens < rhs.outputTokens }
         if lhs.cacheReadTokens != rhs.cacheReadTokens { return lhs.cacheReadTokens < rhs.cacheReadTokens }
         if lhs.cacheWriteTokens != rhs.cacheWriteTokens { return lhs.cacheWriteTokens < rhs.cacheWriteTokens }
-        return lhs.reasoningTokens < rhs.reasoningTokens
+        if lhs.reasoningTokens != rhs.reasoningTokens { return lhs.reasoningTokens < rhs.reasoningTokens }
+        return (lhs.cost ?? -1) < (rhs.cost ?? -1)
+    }
+
+    private func costEventSort(_ lhs: RemoteCostEvent, _ rhs: RemoteCostEvent) -> Bool {
+        if lhs.timestamp != rhs.timestamp { return lhs.timestamp < rhs.timestamp }
+        if lhs.source != rhs.source { return lhs.source < rhs.source }
+        if lhs.model != rhs.model { return (lhs.model ?? "") < (rhs.model ?? "") }
+        return lhs.cost < rhs.cost
     }
 
     private func activityEventSort(_ lhs: RemoteActivityEvent, _ rhs: RemoteActivityEvent) -> Bool {
@@ -424,6 +438,7 @@ private extension AgentSnapshotBuilder {
 
     private func remoteModel(_ model: String?) -> String? {
         guard let model,
+              model != UsageModelGrouping.mixedOrUnattributedKey,
               TokiSyncValidation.isSafeDisplayText(
                   model,
                   maximumLength: RemoteUsageSnapshotValidator.maximumModelLength) else {
@@ -441,7 +456,10 @@ private extension AgentSnapshotBuilder {
             event.reasoningTokens,
         ]
         let validRange = 0...RemoteUsageSnapshotValidator.maximumTokenCountPerBucket
+        let validCostRange = 0...RemoteUsageSnapshotValidator.maximumCostPerEvent
         guard counts.allSatisfy(validRange.contains),
+              event.cost.isFinite,
+              validCostRange.contains(event.cost),
               counts.contains(where: { $0 > 0 }) else {
             return nil
         }
@@ -453,7 +471,30 @@ private extension AgentSnapshotBuilder {
             outputTokens: event.outputTokens,
             cacheReadTokens: event.cacheReadTokens,
             cacheWriteTokens: event.cacheWriteTokens,
-            reasoningTokens: event.reasoningTokens)
+            reasoningTokens: event.reasoningTokens,
+            cost: event.cost > 0 ? event.cost : nil)
+    }
+
+    private func remoteCostEvent(_ event: TokenUsageEvent) -> RemoteCostEvent? {
+        let counts = [
+            event.inputTokens,
+            event.outputTokens,
+            event.cacheReadTokens,
+            event.cacheWriteTokens,
+            event.reasoningTokens,
+        ]
+        let validCostRange = 0...RemoteUsageSnapshotValidator.maximumCostPerEvent
+        guard counts.allSatisfy({ $0 == 0 }),
+              event.cost.isFinite,
+              event.cost > 0,
+              validCostRange.contains(event.cost) else {
+            return nil
+        }
+        return RemoteCostEvent(
+            timestamp: event.timestamp,
+            source: event.source,
+            model: remoteModel(event.model),
+            cost: event.cost)
     }
 
     private var platformName: String {
@@ -478,6 +519,7 @@ private struct AgentSnapshotContent: Encodable {
     let coveredFrom: Date
     let coveredTo: Date
     let tokenEvents: [RemoteTokenEvent]
+    let costEvents: [RemoteCostEvent]?
     let activityEvents: [RemoteActivityEvent]
 }
 

Sources/TokiSyncProtocol/SnapshotValidation.swift

@@ -2,8 +2,10 @@ import Foundation
 
 public enum RemoteUsageSnapshotValidator {
     public static let maximumTokenEventCount = 200_000
+    public static let maximumCostEventCount = 200_000
     public static let maximumActivityEventCount = 200_000
     public static let maximumTokenCountPerBucket = 1_000_000_000
+    public static let maximumCostPerEvent = 1_000_000_000.0
     public static let maximumModelLength = 200
 
     public static func validate(_ snapshot: RemoteUsageSnapshot, now: Date = Date()) throws {
@@ -28,6 +30,9 @@ public enum RemoteUsageSnapshotValidator {
         guard snapshot.tokenEvents.count <= maximumTokenEventCount else {
             throw RemoteUsageSnapshotValidationError.tooManyEvents
         }
+        guard (snapshot.costEvents?.count ?? 0) <= maximumCostEventCount else {
+            throw RemoteUsageSnapshotValidationError.tooManyEvents
+        }
         guard snapshot.activityEvents.count <= maximumActivityEventCount else {
             throw RemoteUsageSnapshotValidationError.tooManyEvents
         }
@@ -41,11 +46,24 @@ public enum RemoteUsageSnapshotValidator {
                   validTokenCount(event.outputTokens),
                   validTokenCount(event.cacheReadTokens),
                   validTokenCount(event.cacheWriteTokens),
-                  validTokenCount(event.reasoningTokens) else {
+                  validTokenCount(event.reasoningTokens),
+                  event.cost.map(validCost) ?? true,
+                  event.totalTokens > 0 else {
                 throw RemoteUsageSnapshotValidationError.invalidTokenEvent
             }
         }
 
+        for event in snapshot.costEvents ?? [] {
+            guard event.timestamp >= snapshot.coveredFrom,
+                  event.timestamp < snapshot.coveredTo,
+                  TokiSyncValidation.isSafeDisplayText(event.source, maximumLength: 40),
+                  isOptionalBoundedText(event.model, maximumLength: maximumModelLength),
+                  validCost(event.cost),
+                  event.cost > 0 else {
+                throw RemoteUsageSnapshotValidationError.invalidCostEvent
+            }
+        }
+
         for event in snapshot.activityEvents {
             guard event.timestamp >= snapshot.coveredFrom,
                   event.timestamp < snapshot.coveredTo,
@@ -61,6 +79,10 @@ public enum RemoteUsageSnapshotValidator {
         (0...maximumTokenCountPerBucket).contains(value)
     }
 
+    private static func validCost(_ value: Double) -> Bool {
+        value.isFinite && (0...maximumCostPerEvent).contains(value)
+    }
+
     private static func isFinite(_ date: Date) -> Bool {
         date.timeIntervalSince1970.isFinite
     }
@@ -77,6 +99,7 @@ public enum RemoteUsageSnapshotValidationError: LocalizedError {
     case invalidDateRange
     case tooManyEvents
     case invalidTokenEvent
+    case invalidCostEvent
     case invalidActivityEvent
 
     public var errorDescription: String? {
@@ -91,6 +114,8 @@ public enum RemoteUsageSnapshotValidationError: LocalizedError {
             "The remote snapshot contains too many events."
         case .invalidTokenEvent:
             "The remote snapshot contains an invalid token event."
+        case .invalidCostEvent:
+            "The remote snapshot contains an invalid cost event."
         case .invalidActivityEvent:
             "The remote snapshot contains an invalid activity event."
         }

Sources/TokiSyncProtocol/UsageSnapshot.swift

@@ -30,6 +30,7 @@ public struct RemoteTokenEvent: Codable, Equatable, Sendable {
     public let cacheReadTokens: Int
     public let cacheWriteTokens: Int
     public let reasoningTokens: Int
+    public let cost: Double?
 
     public init(
         timestamp: Date,
@@ -39,7 +40,8 @@ public struct RemoteTokenEvent: Codable, Equatable, Sendable {
         outputTokens: Int,
         cacheReadTokens: Int,
         cacheWriteTokens: Int,
-        reasoningTokens: Int) {
+        reasoningTokens: Int,
+        cost: Double? = nil) {
         self.timestamp = timestamp
         self.source = source
         self.model = model
@@ -48,6 +50,7 @@ public struct RemoteTokenEvent: Codable, Equatable, Sendable {
         self.cacheReadTokens = max(0, cacheReadTokens)
         self.cacheWriteTokens = max(0, cacheWriteTokens)
         self.reasoningTokens = max(0, reasoningTokens)
+        self.cost = cost
     }
 
     public var totalTokens: Int {
@@ -56,6 +59,24 @@ public struct RemoteTokenEvent: Codable, Equatable, Sendable {
     }
 }
 
+public struct RemoteCostEvent: Codable, Equatable, Sendable {
+    public let timestamp: Date
+    public let source: String
+    public let model: String?
+    public let cost: Double
+
+    public init(
+        timestamp: Date,
+        source: String,
+        model: String?,
+        cost: Double) {
+        self.timestamp = timestamp
+        self.source = source
+        self.model = model
+        self.cost = cost
+    }
+}
+
 private func saturatingTokenSum(_ total: Int, _ value: Int) -> Int {
     let (sum, overflow) = total.addingReportingOverflow(value)
     guard overflow else { return sum }
@@ -92,6 +113,7 @@ public struct RemoteUsageSnapshot: Codable, Equatable, Sendable {
     public let coveredFrom: Date
     public let coveredTo: Date
     public let tokenEvents: [RemoteTokenEvent]
+    public let costEvents: [RemoteCostEvent]?
     public let activityEvents: [RemoteActivityEvent]
 
     public init(
@@ -101,13 +123,15 @@ public struct RemoteUsageSnapshot: Codable, Equatable, Sendable {
         coveredFrom: Date,
         coveredTo: Date,
         tokenEvents: [RemoteTokenEvent],
+        costEvents: [RemoteCostEvent]? = nil,
         activityEvents: [RemoteActivityEvent]) {
         self.schemaVersion = schemaVersion
         self.device = device
         self.generatedAt = generatedAt
         self.coveredFrom = coveredFrom
         self.coveredTo = coveredTo
         self.tokenEvents = tokenEvents
+        self.costEvents = costEvents
         self.activityEvents = activityEvents
     }
 }

Sources/TokiUsageCore/RawTokenUsage.swift

@@ -12,6 +12,11 @@ public enum UsageQuality: String {
     case derived
 }
 
+public enum UsageModelGrouping {
+    public static let mixedOrUnattributedKey = "\u{0}toki:mixed-or-unattributed"
+    public static let mixedOrUnattributedLabel = "Mixed / Unattributed"
+}
+
 public enum AttributionQuality: String, Codable {
     case exact
     case inferred
@@ -333,7 +338,7 @@ public struct RawTokenUsage {
             reasoningTokens: reasoningTokens,
             cost: cost,
             attribution: attribution)
-        guard event.totalTokens > 0 else { return }
+        guard event.totalTokens > 0 || event.cost > 0 else { return }
         tokenEvents.append(event)
     }
 }