From f803fbcf76d0f8dca032fa60641a785e24393895 Mon Sep 17 00:00:00 2001 From: choi138 Date: Thu, 30 Jul 2026 19:18:37 +0900 Subject: [PATCH 01/10] fix(hermes): harden attribution and database reads --- Sources/TokiUsageReaders/HermesReader.swift | 60 ++---- .../HermesSQLiteConnection.swift | 198 ++++++++++++++++++ .../HermesUsageAttribution.swift | 161 ++++++++++++++ .../HermesUsageCostBreakdown.swift | 21 +- .../TokiUsageReaders/HermesUsageLedger.swift | 106 +++++----- .../HermesUsageLedgerSupport.swift | 39 ++++ .../HermesUsageResolver.swift | 15 +- .../AgentSnapshotModelAttributionTests.swift | 50 +++++ Toki/Domain/Usage/UsageReportModelStats.swift | 6 +- TokiTests/HermesPricingRefreshTests.swift | 54 ++++- TokiTests/HermesReaderTests.swift | 152 +++++++++++++- TokiTests/HermesUsageLedgerTests.swift | 137 ++++++++++++ TokiTests/UsageServiceBehaviorTests.swift | 27 +++ 13 files changed, 915 insertions(+), 111 deletions(-) create mode 100644 Sources/TokiUsageReaders/HermesSQLiteConnection.swift create mode 100644 Sources/TokiUsageReaders/HermesUsageAttribution.swift create mode 100644 Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift diff --git a/Sources/TokiUsageReaders/HermesReader.swift b/Sources/TokiUsageReaders/HermesReader.swift index 640d53f..1e16047 100644 --- a/Sources/TokiUsageReaders/HermesReader.swift +++ b/Sources/TokiUsageReaders/HermesReader.swift @@ -32,9 +32,9 @@ public struct HermesReader: TokenReader { } public func readUsage(from startDate: Date, to endDate: Date) async throws -> RawTokenUsage { - if let database = try openDatabase() { - defer { sqlite3_close(database) } - let observations = try readSessionObservations(from: database) + if let observations = try readDatabaseSnapshot({ database in + try readSessionObservations(from: database) + }) { try await usageLedger.refresh( observations: observations, observedAt: now()) @@ -45,25 +45,29 @@ public struct HermesReader: TokenReader { } public func coverageStatus() throws -> HermesUsageCoverageStatus { - guard let database = try openDatabase() else { - return HermesUsageCoverageStatus(unmeteredMainAPICallCount: 0) - } - defer { sqlite3_close(database) } - return try readSessionModelUsage(from: database).coverage + try readDatabaseSnapshot { database in + try readSessionModelUsage(from: database).coverage + } ?? HermesUsageCoverageStatus(unmeteredMainAPICallCount: 0) } - private func openDatabase() throws -> OpaquePointer? { - guard FileManager.default.fileExists(atPath: dbPath) else { return nil } - - var database: OpaquePointer? - guard sqlite3_open_v2(dbPath, &database, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { - let error = HermesSQLiteError(operation: "open", database: database) - sqlite3_close(database) - throw error + private func readDatabaseSnapshot( + _ read: (OpaquePointer) throws -> Value) throws -> Value? { + for attempt in 0..<2 { + guard let connection = try HermesSQLiteConnection.open(atPath: dbPath) else { + return nil + } + let value = try read(connection.database) + if connection.isSourceStateCurrent { + return value + } + guard attempt == 0 else { + throw HermesSQLiteError( + operation: "read snapshot", + message: "database changed during read", + code: SQLITE_BUSY) + } } - - sqlite3_busy_timeout(database, 2000) - return database + return nil } private func readSessionObservations(from database: OpaquePointer) throws -> [HermesSessionObservation] { @@ -430,24 +434,6 @@ private struct HermesSessionUsageRow { } } -private struct HermesSQLiteError: LocalizedError { - let operation: String - let message: String - - init(operation: String, database: OpaquePointer?) { - self.operation = operation - if let database, let errorMessage = sqlite3_errmsg(database) { - message = String(cString: errorMessage) - } else { - message = "unknown SQLite error" - } - } - - var errorDescription: String? { - "Hermes SQLite \(operation) failed: \(message)" - } -} - private let hermesSQLiteTransient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) private func hermesSQLiteText(_ statement: OpaquePointer?, at index: Int32) -> String { diff --git a/Sources/TokiUsageReaders/HermesSQLiteConnection.swift b/Sources/TokiUsageReaders/HermesSQLiteConnection.swift new file mode 100644 index 0000000..d541263 --- /dev/null +++ b/Sources/TokiUsageReaders/HermesSQLiteConnection.swift @@ -0,0 +1,198 @@ +import Foundation + +#if os(Linux) + import CSQLite +#else + import SQLite3 +#endif + +final class HermesSQLiteConnection { + let database: OpaquePointer + + private let immutableSnapshot: HermesDatabaseSourceSnapshot? + + private init( + database: OpaquePointer, + immutableSnapshot: HermesDatabaseSourceSnapshot?) { + self.database = database + self.immutableSnapshot = immutableSnapshot + } + + deinit { + sqlite3_close(database) + } + + var isSourceStateCurrent: Bool { + immutableSnapshot?.isCurrent() ?? true + } + + var isUsingImmutableSnapshot: Bool { + immutableSnapshot != nil + } + + static func open( + atPath path: String, + fileManager: FileManager = .default) throws -> HermesSQLiteConnection? { + guard fileManager.fileExists(atPath: path) else { return nil } + + do { + return try openValidatedDatabase( + path: path, + flags: SQLITE_OPEN_READONLY, + immutableSnapshot: nil) + } catch let error as HermesSQLiteError { + let databaseURL = URL(fileURLWithPath: path) + guard error.code == SQLITE_CANTOPEN, + let snapshot = HermesDatabaseSourceSnapshot.captureForImmutableFallback( + databaseURL: databaseURL, + fileManager: fileManager) else { + throw error + } + + return try openValidatedDatabase( + path: immutableDatabaseURI(for: databaseURL), + flags: SQLITE_OPEN_READONLY | SQLITE_OPEN_URI, + immutableSnapshot: snapshot) + } + } + + private static func openValidatedDatabase( + path: String, + flags: Int32, + immutableSnapshot: HermesDatabaseSourceSnapshot?) throws -> HermesSQLiteConnection { + var database: OpaquePointer? + let openStatus = sqlite3_open_v2(path, &database, flags, nil) + guard openStatus == SQLITE_OK, let database else { + let error = HermesSQLiteError( + operation: "open", + database: database, + code: openStatus) + sqlite3_close(database) + throw error + } + + sqlite3_busy_timeout(database, 2000) + let probeStatus = probeDatabase(database) + guard probeStatus == SQLITE_OK else { + let error = HermesSQLiteError( + operation: "probe", + database: database, + code: probeStatus) + sqlite3_close(database) + throw error + } + + return HermesSQLiteConnection( + database: database, + immutableSnapshot: immutableSnapshot) + } + + private static func probeDatabase(_ database: OpaquePointer) -> Int32 { + var statement: OpaquePointer? + let prepareStatus = sqlite3_prepare_v2( + database, + "PRAGMA schema_version", + -1, + &statement, + nil) + guard prepareStatus == SQLITE_OK else { return prepareStatus } + defer { sqlite3_finalize(statement) } + + let stepStatus = sqlite3_step(statement) + guard stepStatus == SQLITE_ROW || stepStatus == SQLITE_DONE else { + return stepStatus + } + return SQLITE_OK + } + + private static func immutableDatabaseURI(for databaseURL: URL) -> String { + "\(databaseURL.absoluteString)?mode=ro&immutable=1" + } +} + +struct HermesDatabaseSourceSnapshot: Equatable { + let databaseURL: URL + let databaseSignature: HermesDatabaseFileSignature + + static func captureForImmutableFallback( + databaseURL: URL, + fileManager: FileManager = .default) -> HermesDatabaseSourceSnapshot? { + guard !hasSQLiteSidecars(databaseURL: databaseURL, fileManager: fileManager), + let databaseSignature = HermesDatabaseFileSignature.capture( + at: databaseURL, + fileManager: fileManager), + !hasSQLiteSidecars(databaseURL: databaseURL, fileManager: fileManager) else { + return nil + } + return HermesDatabaseSourceSnapshot( + databaseURL: databaseURL, + databaseSignature: databaseSignature) + } + + func isCurrent(fileManager: FileManager = .default) -> Bool { + guard let currentSnapshot = Self.captureForImmutableFallback( + databaseURL: databaseURL, + fileManager: fileManager) else { + return false + } + return currentSnapshot == self + } + + private static func hasSQLiteSidecars( + databaseURL: URL, + fileManager: FileManager) -> Bool { + let path = databaseURL.path + return fileManager.fileExists(atPath: "\(path)-wal") + || fileManager.fileExists(atPath: "\(path)-shm") + } +} + +struct HermesDatabaseFileSignature: Equatable { + let systemNumber: UInt64? + let fileNumber: UInt64? + let size: UInt64 + let modificationDate: Date? + + static func capture( + at fileURL: URL, + fileManager: FileManager = .default) -> HermesDatabaseFileSignature? { + guard let attributes = try? fileManager.attributesOfItem(atPath: fileURL.path), + let size = (attributes[.size] as? NSNumber)?.uint64Value else { + return nil + } + return HermesDatabaseFileSignature( + systemNumber: (attributes[.systemNumber] as? NSNumber)?.uint64Value, + fileNumber: (attributes[.systemFileNumber] as? NSNumber)?.uint64Value, + size: size, + modificationDate: attributes[.modificationDate] as? Date) + } +} + +struct HermesSQLiteError: LocalizedError { + let operation: String + let message: String + let code: Int32 + + init( + operation: String, + database: OpaquePointer?, + code: Int32? = nil) { + self.operation = operation + self.code = code ?? database.map(sqlite3_errcode) ?? SQLITE_ERROR + if let database, let errorMessage = sqlite3_errmsg(database) { + message = String(cString: errorMessage) + } else { + message = "unknown SQLite error" + } + } + + init(operation: String, message: String, code: Int32) { + self.operation = operation + self.message = message + self.code = code + } + + var errorDescription: String? { + "Hermes SQLite \(operation) failed: \(message)" + } +} diff --git a/Sources/TokiUsageReaders/HermesUsageAttribution.swift b/Sources/TokiUsageReaders/HermesUsageAttribution.swift new file mode 100644 index 0000000..0f0b563 --- /dev/null +++ b/Sources/TokiUsageReaders/HermesUsageAttribution.swift @@ -0,0 +1,161 @@ +import Foundation + +private struct HermesUsageEventPart { + let model: String? + let counters: HermesTokenCounters +} + +private struct HermesUsageCostAllocationBasis { + let pricedCosts: [Double] + let remainderTokenCounts: [Int] +} + +func hermesUsageEvents( + identifier: String, + timestamp: Date, + observation: HermesSessionObservation, + previousModelCounters: [String: HermesTokenCounters]?, + previousModelPricingCounters: [String: HermesTokenCounters]?, + counters: HermesTokenCounters, + cost: Double) -> [HermesUsageLedgerEvent] { + let parts = hermesUsageEventParts( + observation: observation, + previousModelCounters: previousModelCounters, + counters: counters) + let modelPricingDeltas = hermesModelCounterDeltas( + current: observation.modelPricingCounters, + previous: previousModelPricingCounters, + maximumDelta: counters) + let costs = hermesAllocatedUsageCosts( + totalCost: cost, + parts: parts, + modelPricingDeltas: modelPricingDeltas, + timestamp: timestamp) + return zip(parts, costs).map { part, allocatedCost in + HermesUsageLedgerEvent( + sessionIdentifier: identifier, + timestamp: timestamp, + model: part.model, + counters: part.counters, + cost: allocatedCost, + projectName: observation.projectName, + attributionQuality: observation.attributionQuality) + } +} + +private func hermesUsageEventParts( + observation: HermesSessionObservation, + previousModelCounters: [String: HermesTokenCounters]?, + counters: HermesTokenCounters) -> [HermesUsageEventPart] { + guard let modelDeltas = hermesModelCounterDeltas( + current: observation.modelCounters, + previous: previousModelCounters, + maximumDelta: counters) else { + return [HermesUsageEventPart(model: observation.model, counters: counters)] + } + + var combinedCounters = HermesTokenCounters.zero + var parts: [HermesUsageEventPart] = [] + for model in modelDeltas.keys.sorted() { + guard let modelCounters = modelDeltas[model] else { continue } + combinedCounters = combinedCounters.adding(modelCounters) + parts.append(HermesUsageEventPart(model: model, counters: modelCounters)) + } + + let residual = counters.subtracting(combinedCounters) + if residual.totalTokens > 0 { + parts.append(HermesUsageEventPart(model: nil, counters: residual)) + } + return parts +} + +private func hermesAllocatedUsageCosts( + totalCost: Double, + parts: [HermesUsageEventPart], + modelPricingDeltas: [String: HermesTokenCounters]?, + timestamp: Date) -> [Double] { + guard totalCost > 0, !parts.isEmpty else { + return Array(repeating: 0, count: parts.count) + } + + let fallbackBasis = HermesUsageCostAllocationBasis( + pricedCosts: Array(repeating: 0, count: parts.count), + remainderTokenCounts: parts.map(\.counters.totalTokens)) + let basis = hermesUsageCostAllocationBasis( + parts: parts, + modelPricingDeltas: modelPricingDeltas, + timestamp: timestamp) ?? fallbackBasis + let pricedTotal = basis.pricedCosts.reduce(0, +) + if pricedTotal > totalCost, pricedTotal > 0 { + let scale = totalCost / pricedTotal + return basis.pricedCosts.map { $0 * scale } + } + + var allocations = basis.pricedCosts + let remainder = totalCost - pricedTotal + guard remainder > 0 else { return allocations } + + var remainderTokenCounts = basis.remainderTokenCounts + if !remainderTokenCounts.contains(where: { $0 > 0 }) { + remainderTokenCounts = parts.map(\.counters.totalTokens) + } + let recipientIndices = remainderTokenCounts.indices.filter { + remainderTokenCounts[$0] > 0 + } + let recipientTokens = recipientIndices.reduce(0) { + $0 + remainderTokenCounts[$1] + } + guard recipientTokens > 0, let lastRecipient = recipientIndices.last else { + return allocations + } + + var remainingCost = remainder + for index in recipientIndices.dropLast() { + let proportionalShare = + remainder * Double(remainderTokenCounts[index]) / Double(recipientTokens) + let share = min(remainingCost, proportionalShare) + allocations[index] += share + remainingCost -= share + } + allocations[lastRecipient] += remainingCost + return allocations +} + +private func hermesUsageCostAllocationBasis( + parts: [HermesUsageEventPart], + modelPricingDeltas: [String: HermesTokenCounters]?, + timestamp: Date) -> HermesUsageCostAllocationBasis? { + guard let modelPricingDeltas else { return nil } + + var unmatchedModels = Set(modelPricingDeltas.keys) + var pricedCosts: [Double] = [] + var remainderTokenCounts: [Int] = [] + for part in parts { + guard let model = part.model else { + pricedCosts.append(0) + remainderTokenCounts.append(part.counters.totalTokens) + continue + } + + let pricingCounters = modelPricingDeltas[model] ?? .zero + guard !part.counters.hasDecrease(comparedTo: pricingCounters) else { + return nil + } + unmatchedModels.remove(model) + + let pricedCost = hermesModelPricedCost( + counters: pricingCounters, + model: model, + timestamp: timestamp) + guard pricedCost.isFinite, pricedCost >= 0 else { return nil } + + pricedCosts.append(pricedCost) + remainderTokenCounts.append( + part.counters.subtracting(pricingCounters).totalTokens) + } + guard unmatchedModels.isEmpty else { return nil } + + return HermesUsageCostAllocationBasis( + pricedCosts: pricedCosts, + remainderTokenCounts: remainderTokenCounts) +} diff --git a/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift b/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift index 2a9a4d8..e7f7be7 100644 --- a/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift +++ b/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift @@ -30,21 +30,15 @@ func hermesModelPricedDeltaCost( previous: [String: HermesTokenCounters]?, maximumDelta: HermesTokenCounters, timestamp: Date) -> Double? { - guard let current, let previous else { return nil } + guard let deltas = hermesModelCounterDeltas( + current: current, + previous: previous, + maximumDelta: maximumDelta) else { + return nil + } - var combinedDelta = HermesTokenCounters.zero var cost = 0.0 - for model in Set(current.keys).union(previous.keys) { - let currentCounters = current[model] ?? .zero - let previousCounters = previous[model] ?? .zero - guard !currentCounters.hasDecrease(comparedTo: previousCounters) else { return nil } - let delta = currentCounters.subtracting(previousCounters) - guard combinedDelta.canAdd( - delta, - maximum: hermesLedgerMaximumCumulativeTokens) else { - return nil - } - combinedDelta = combinedDelta.adding(delta) + for (model, delta) in deltas { let modelCost = hermesModelPricedCost( counters: delta, model: model, @@ -52,7 +46,6 @@ func hermesModelPricedDeltaCost( guard modelCost.isFinite, (cost + modelCost).isFinite else { return nil } cost += modelCost } - guard !maximumDelta.hasDecrease(comparedTo: combinedDelta) else { return nil } return cost } diff --git a/Sources/TokiUsageReaders/HermesUsageLedger.swift b/Sources/TokiUsageReaders/HermesUsageLedger.swift index 75f0a2a..8515145 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedger.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedger.swift @@ -154,31 +154,13 @@ private extension HermesUsageLedger { previous: previous) guard let previous else { - guard observation.counters.totalTokens > 0 else { return false } - candidate.baselines[identifier] = currentBaseline - if initialUsageIsDated( + return try applyInitial( observation, - after: previousSuccessfulObservationAt, - observedAt: observedAt) { - append( - event( - identifier: identifier, - timestamp: initialTimestamp( - observation: observation, - observedAt: observedAt), - observation: observation, - counters: observation.counters, - cost: observation.cost), - to: &candidate.events) - } else { - try addUnattributed( - identifier: identifier, - counters: observation.counters, - cost: observation.cost, - observedAt: observedAt, - to: &candidate.unattributed) - } - return true + currentBaseline: currentBaseline, + identifier: identifier, + previousSuccessfulObservationAt: previousSuccessfulObservationAt, + observedAt: observedAt, + to: &candidate) } if observation.counters.hasDecrease(comparedTo: previous.counters) { @@ -216,15 +198,57 @@ private extension HermesUsageLedger { to: &candidate.unattributed) return true } - append( - event( + for event in hermesUsageEvents( + identifier: identifier, + timestamp: timestamp, + observation: observation, + previousModelCounters: previous.modelCounters, + previousModelPricingCounters: previous.modelPricingCounters, + counters: delta, + cost: cost) { + append( + event, + to: &candidate.events) + } + candidate.baselines[identifier] = currentBaseline + return true + } + + func applyInitial( + _ observation: HermesSessionObservation, + currentBaseline: HermesUsageLedgerBaseline, + identifier: String, + previousSuccessfulObservationAt: Date?, + observedAt: Date, + to candidate: inout HermesUsageLedgerDocument) throws -> Bool { + guard observation.counters.totalTokens > 0 else { return false } + candidate.baselines[identifier] = currentBaseline + if initialUsageIsDated( + observation, + after: previousSuccessfulObservationAt, + observedAt: observedAt) { + for event in hermesUsageEvents( identifier: identifier, - timestamp: timestamp, + timestamp: initialTimestamp( + observation: observation, + observedAt: observedAt), observation: observation, - counters: delta, - cost: cost), - to: &candidate.events) - candidate.baselines[identifier] = currentBaseline + previousModelCounters: observation.modelCounters.map { _ in [:] }, + previousModelPricingCounters: observation.modelPricingCounters.map { _ in [:] }, + counters: observation.counters, + cost: observation.cost) { + append( + event, + to: &candidate.events) + } + } else { + try addUnattributed( + identifier: identifier, + counters: observation.counters, + cost: observation.cost, + observedAt: observedAt, + to: &candidate.unattributed) + } return true } @@ -413,6 +437,7 @@ private extension HermesUsageLedger { lastObservedAt: observedAt, model: observation.model, counters: observation.counters, + modelCounters: observation.modelCounters, cost: observation.cost, reportedCost: observation.reportedCost, modelPricingCounters: observation.modelPricingCounters, @@ -460,22 +485,6 @@ private extension HermesUsageLedger { observedAt: observedAt) ?? observation.startedAt } - private func event( - identifier: String, - timestamp: Date, - observation: HermesSessionObservation, - counters: HermesTokenCounters, - cost: Double) -> HermesUsageLedgerEvent { - HermesUsageLedgerEvent( - sessionIdentifier: identifier, - timestamp: timestamp, - model: observation.model, - counters: counters, - cost: cost, - projectName: observation.projectName, - attributionQuality: observation.attributionQuality) - } - private func append( _ event: HermesUsageLedgerEvent, to events: inout [HermesUsageLedgerEvent]) { @@ -556,6 +565,9 @@ private extension HermesUsageLedger { hermesModelPricingCountersAreValid( observation.modelPricingCounters, within: observation.counters), + hermesModelPricingCountersAreValid( + observation.modelCounters, + within: observation.counters), observation.model?.utf8.count ?? 0 <= 512, observation.projectName?.utf8.count ?? 0 <= 512 else { throw HermesUsageLedgerError.invalidObservation diff --git a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift index 4224af6..2172d10 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift @@ -18,6 +18,7 @@ struct HermesSessionObservation { let latestActivityAt: Date? let model: String? let counters: HermesTokenCounters + let modelCounters: [String: HermesTokenCounters]? let cost: Double let costIsDerivedFromModelPricing: Bool let reportedCost: Double? @@ -32,6 +33,7 @@ struct HermesSessionObservation { latestActivityAt: Date?, model: String?, counters: HermesTokenCounters, + modelCounters: [String: HermesTokenCounters]? = nil, cost: Double, costIsDerivedFromModelPricing: Bool = false, reportedCost: Double? = nil, @@ -44,6 +46,7 @@ struct HermesSessionObservation { self.latestActivityAt = latestActivityAt self.model = model self.counters = counters + self.modelCounters = modelCounters self.cost = cost self.costIsDerivedFromModelPricing = costIsDerivedFromModelPricing self.reportedCost = reportedCost @@ -171,6 +174,34 @@ func hermesModelPricingCountersAreValid( return !totalCounters.hasDecrease(comparedTo: combinedCounters) } +func hermesModelCounterDeltas( + current: [String: HermesTokenCounters]?, + previous: [String: HermesTokenCounters]?, + maximumDelta: HermesTokenCounters) -> [String: HermesTokenCounters]? { + guard let current, let previous else { return nil } + + var combinedDelta = HermesTokenCounters.zero + var deltas: [String: HermesTokenCounters] = [:] + for model in Set(current.keys).union(previous.keys) { + let currentCounters = current[model] ?? .zero + let previousCounters = previous[model] ?? .zero + guard !currentCounters.hasDecrease(comparedTo: previousCounters) else { return nil } + + let delta = currentCounters.subtracting(previousCounters) + guard combinedDelta.canAdd( + delta, + maximum: hermesLedgerMaximumCumulativeTokens) else { + return nil + } + combinedDelta = combinedDelta.adding(delta) + if delta.totalTokens > 0 { + deltas[model] = delta + } + } + guard !maximumDelta.hasDecrease(comparedTo: combinedDelta) else { return nil } + return deltas +} + struct HermesUsageLedgerEvent: Codable, Equatable { let sessionIdentifier: String var timestamp: Date @@ -309,6 +340,7 @@ struct HermesUsageLedgerPrivateBaseline: Codable, Equatable { let lastObservedAt: Date let model: String? let counters: HermesTokenCounters + let modelCounters: [String: HermesTokenCounters]? let cost: Double let reportedCost: Double? let modelPricingCounters: [String: HermesTokenCounters]? @@ -320,6 +352,7 @@ struct HermesUsageLedgerPrivateBaseline: Codable, Equatable { lastObservedAt = baseline.lastObservedAt model = baseline.model counters = baseline.counters + modelCounters = baseline.modelCounters cost = baseline.cost reportedCost = baseline.reportedCost modelPricingCounters = baseline.modelPricingCounters @@ -333,6 +366,7 @@ struct HermesUsageLedgerPrivateBaseline: Codable, Equatable { lastObservedAt: lastObservedAt, model: model, counters: counters, + modelCounters: modelCounters, cost: cost, reportedCost: reportedCost, modelPricingCounters: modelPricingCounters, @@ -387,6 +421,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { let lastObservedAt: Date let model: String? let counters: HermesTokenCounters + let modelCounters: [String: HermesTokenCounters]? let cost: Double let reportedCost: Double? let modelPricingCounters: [String: HermesTokenCounters]? @@ -399,6 +434,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { lastObservedAt: Date, model: String?, counters: HermesTokenCounters, + modelCounters: [String: HermesTokenCounters]? = nil, cost: Double, reportedCost: Double? = nil, modelPricingCounters: [String: HermesTokenCounters]? = nil, @@ -409,6 +445,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { self.lastObservedAt = lastObservedAt self.model = model self.counters = counters + self.modelCounters = modelCounters self.cost = cost self.reportedCost = reportedCost self.modelPricingCounters = modelPricingCounters @@ -423,6 +460,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { && lastActivityAt >= startedAt && lastObservedAt >= startedAt && counters.isValid() + && hermesModelPricingCountersAreValid(modelCounters, within: counters) && cost.isFinite && cost >= 0 && hermesReportedCostBreakdownIsValid( @@ -437,6 +475,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { func metadataDiffers(from previous: Self) -> Bool { model != previous.model || counters != previous.counters + || modelCounters != previous.modelCounters || cost != previous.cost || reportedCost != previous.reportedCost || modelPricingCounters != previous.modelPricingCounters diff --git a/Sources/TokiUsageReaders/HermesUsageResolver.swift b/Sources/TokiUsageReaders/HermesUsageResolver.swift index 3463645..ef2362a 100644 --- a/Sources/TokiUsageReaders/HermesUsageResolver.swift +++ b/Sources/TokiUsageReaders/HermesUsageResolver.swift @@ -36,6 +36,7 @@ enum HermesUsageResolver { var modelCost = 0.0 var modelReportedCost = 0.0 var modelCostIsDerivedFromModelPricing = true + var attributedModelCounters: [String: HermesTokenCounters] = [:] var modelPricingCounters: [String: HermesTokenCounters] = [:] var models: Set = [] @@ -55,14 +56,21 @@ enum HermesUsageResolver { modelReportedCost += usage.cost } if usage.counters.totalTokens > 0, let model = usage.model { + let existingCounters = attributedModelCounters[model] ?? .zero + guard existingCounters.canAdd( + usage.counters, + maximum: hermesLedgerMaximumCumulativeTokens) else { + throw HermesUsageLedgerError.invalidObservation + } + attributedModelCounters[model] = existingCounters.adding(usage.counters) if usage.costIsDerivedFromModelPricing { - let existingCounters = modelPricingCounters[model] ?? .zero - guard existingCounters.canAdd( + let existingPricingCounters = modelPricingCounters[model] ?? .zero + guard existingPricingCounters.canAdd( usage.counters, maximum: hermesLedgerMaximumCumulativeTokens) else { throw HermesUsageLedgerError.invalidObservation } - modelPricingCounters[model] = existingCounters.adding(usage.counters) + modelPricingCounters[model] = existingPricingCounters.adding(usage.counters) } models.insert(model) } @@ -87,6 +95,7 @@ enum HermesUsageResolver { latestActivityAt: session.latestActivityAt, model: resolvedModel, counters: session.counters.maximum(modelCounters), + modelCounters: modelUsage.isEmpty ? nil : attributedModelCounters, cost: resolvedCost.value, costIsDerivedFromModelPricing: resolvedCost.isDerivedFromModelPricing, reportedCost: resolvedCost.reportedValue, diff --git a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift new file mode 100644 index 0000000..bf7d5b9 --- /dev/null +++ b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift @@ -0,0 +1,50 @@ +import Foundation +import TokiUsageCore +import XCTest +@testable import TokiAgentCore +@testable import TokiUsageReaders + +final class AgentSnapshotModelAttributionTests: XCTestCase { + func test_snapshotPreservesMixedHermesModelAttributionAndResidual() async throws { + let now = Date(timeIntervalSince1970: 1_784_200_000) + let eventDate = now.addingTimeInterval(-60) + var usage = RawTokenUsage(inputTokens: 150) + usage.recordTokenEvent( + timestamp: eventDate, + source: "Hermes", + model: "gpt-5.6-sol", + inputTokens: 100, + outputTokens: 0) + usage.recordTokenEvent( + timestamp: eventDate, + source: "Hermes", + model: "kr/claude-opus-5", + inputTokens: 30, + outputTokens: 0) + usage.recordTokenEvent( + timestamp: eventDate, + source: "Hermes", + model: nil, + inputTokens: 20, + outputTokens: 0) + let descriptor = LocalUsageReaderDescriptor( + reader: FixedTokenReader(name: "Hermes", usage: usage), + sourceLocations: []) + let fixture = try AgentSnapshotFixture() + defer { fixture.remove() } + let builder = AgentSnapshotBuilder( + home: fixture.root, + readerDescriptors: [descriptor]) + + let snapshot = try await builder.build(configuration: fixture.configuration, now: now) + let tokensByModel = snapshot.tokenEvents.reduce(into: [String: Int]()) { result, event in + result[event.model ?? "Mixed / Unattributed", default: 0] += event.totalTokens + } + + XCTAssertEqual(tokensByModel, [ + "gpt-5.6-sol": 100, + "kr/claude-opus-5": 30, + "Mixed / Unattributed": 20, + ]) + } +} diff --git a/Toki/Domain/Usage/UsageReportModelStats.swift b/Toki/Domain/Usage/UsageReportModelStats.swift index 5306f43..86feec3 100644 --- a/Toki/Domain/Usage/UsageReportModelStats.swift +++ b/Toki/Domain/Usage/UsageReportModelStats.swift @@ -2,6 +2,8 @@ import Foundation import TokiUsageCore import TokiUsageReaders +private let unattributedModelID = "Mixed / Unattributed" + private struct ModelSourceStatAggregate { var totalTokens = 0 var cost: Double = 0 @@ -124,10 +126,10 @@ extension UsageReportBuilder { var activityEventsBySource: [String: [ActivityTimeEvent]] = [:] for event in events where event.totalTokens > 0 { - guard let modelID = event.model?.trimmedNonEmpty, - let source = event.source.trimmedNonEmpty else { + guard let source = event.source.trimmedNonEmpty else { continue } + let modelID = event.model?.trimmedNonEmpty ?? unattributedModelID let key = ModelSourceUsageKey(modelID: modelID, source: source) aggregates[key, default: ModelSourceStatAggregate()].accumulate(event, modelID: modelID) activityEventsBySource[source, default: []].append( diff --git a/TokiTests/HermesPricingRefreshTests.swift b/TokiTests/HermesPricingRefreshTests.swift index bcaf498..1bf7fbc 100644 --- a/TokiTests/HermesPricingRefreshTests.swift +++ b/TokiTests/HermesPricingRefreshTests.swift @@ -1,3 +1,4 @@ +// swiftlint:disable file_length import TokiUsageCore import XCTest @testable import TokiUsageReaders @@ -151,10 +152,16 @@ final class HermesPricingRefreshTests: XCTestCase { XCTAssertEqual(increment.inputTokens, 100) XCTAssertEqual(increment.totalTokens, 100) XCTAssertEqual(increment.cost, expectedCost, accuracy: 0.000001) + XCTAssertEqual(increment.perModel[firstModel]?.totalTokens, 60) + XCTAssertEqual(increment.perModel[secondModel]?.totalTokens, 40) + XCTAssertEqual( + hermesEventTokenTotalsByModel(in: increment.tokenEvents), + [firstModel: 60, secondModel: 40]) } } extension HermesPricingRefreshTests { + // swiftlint:disable:next function_body_length func test_hermesReader_pricesOnlyDerivedPartOfMixedCostAfterRateChange() async throws { let tempDir = try makeHermesTemporaryDirectory() defer { @@ -163,9 +170,13 @@ extension HermesPricingRefreshTests { } let reportedModel = "reported-hermes-model" let derivedModel = "derived-hermes-model" + let reportedCatalogPrice = uniformHermesModelPrice(perMillion: 50) let initialPrice = uniformHermesModelPrice(perMillion: 1) let updatedPrice = uniformHermesModelPrice(perMillion: 10) - ModelPricingSupplement.install([derivedModel: initialPrice]) + ModelPricingSupplement.install([ + reportedModel: reportedCatalogPrice, + derivedModel: initialPrice, + ]) let dbURL = tempDir.appendingPathComponent("state.db") let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") @@ -187,6 +198,11 @@ extension HermesPricingRefreshTests { XCTAssertEqual(initialUsage.totalTokens, 0) ModelPricingSupplement.install(priceHistories: [ + reportedModel: [ + ModelPricingSupplement.PriceVersion( + effectiveFrom: .distantPast, + price: reportedCatalogPrice), + ], derivedModel: [ ModelPricingSupplement.PriceVersion( effectiveFrom: .distantPast, @@ -201,6 +217,12 @@ extension HermesPricingRefreshTests { sessionID: "mixed-cost-session", task: "derived", inputTokens: 160) + try updateHermesModelUsage( + databaseURL: dbURL, + sessionID: "mixed-cost-session", + task: "reported", + inputTokens: 160, + actualCost: 1.6) let increment = try await HermesReader( dbPathOverride: dbURL.path, @@ -210,11 +232,27 @@ extension HermesPricingRefreshTests { from: tokiTestISODate("2026-04-10T00:00:00Z"), to: tokiTestISODate("2026-04-11T00:00:00Z")) - XCTAssertEqual(increment.inputTokens, 60) - XCTAssertEqual(increment.totalTokens, 60) + let expectedReportedCost = 0.6 + let expectedDerivedCost = updatedPrice.cost( + input: 60, + output: 0, + cacheRead: 0, + cacheWrite: 0) + XCTAssertEqual(increment.inputTokens, 120) + XCTAssertEqual(increment.totalTokens, 120) XCTAssertEqual( increment.cost, - updatedPrice.cost(input: 60, output: 0, cacheRead: 0, cacheWrite: 0), + expectedReportedCost + expectedDerivedCost, + accuracy: 0.000001) + XCTAssertEqual(increment.perModel[reportedModel]?.totalTokens, 60) + XCTAssertEqual(increment.perModel[derivedModel]?.totalTokens, 60) + XCTAssertEqual( + increment.perModel[reportedModel]?.cost ?? -1, + expectedReportedCost, + accuracy: 0.000001) + XCTAssertEqual( + increment.perModel[derivedModel]?.cost ?? -1, + expectedDerivedCost, accuracy: 0.000001) } @@ -402,6 +440,14 @@ private func uniformHermesModelPrice(perMillion: Double) -> ModelPrice { cacheWritePerMillion: perMillion) } +private func hermesEventTokenTotalsByModel( + in events: [TokenUsageEvent]) -> [String: Int] { + events.reduce(into: [String: Int]()) { result, event in + guard let model = event.model else { return } + result[model, default: 0] += event.totalTokens + } +} + private func createHermesPricingRefreshDatabase( at databaseURL: URL, model: String) throws { diff --git a/TokiTests/HermesReaderTests.swift b/TokiTests/HermesReaderTests.swift index f6315fd..4bff706 100644 --- a/TokiTests/HermesReaderTests.swift +++ b/TokiTests/HermesReaderTests.swift @@ -392,6 +392,123 @@ final class HermesReaderTests: XCTestCase { XCTAssertTrue(error.localizedDescription.contains("Hermes SQLite prepare failed")) } } + + func test_hermesReader_readsCheckpointedWALWithoutSidecarsFromReadOnlyDirectory() async throws { + let fileManager = FileManager.default + let tempDir = try makeHermesTemporaryDirectory() + let databaseDirectory = tempDir.appendingPathComponent("database", isDirectory: true) + try fileManager.createDirectory( + at: databaseDirectory, + withIntermediateDirectories: false, + attributes: [.posixPermissions: 0o700]) + defer { + try? fileManager.setAttributes( + [.posixPermissions: 0o700], + ofItemAtPath: databaseDirectory.path) + try? fileManager.removeItem(at: tempDir) + } + + let dbURL = databaseDirectory.appendingPathComponent("state.db") + try createHermesStateDB( + at: dbURL, + rows: [hermesSingleCounterFixture(id: "checkpointed-wal", inputTokens: 123)]) + try convertHermesDatabaseToWALWithoutSidecars(databaseURL: dbURL) + XCTAssertFalse(fileManager.fileExists(atPath: "\(dbURL.path)-wal")) + XCTAssertFalse(fileManager.fileExists(atPath: "\(dbURL.path)-shm")) + + try fileManager.setAttributes( + [.posixPermissions: 0o500], + ofItemAtPath: databaseDirectory.path) + + do { + let connection = try XCTUnwrap(HermesSQLiteConnection.open(atPath: dbURL.path)) + XCTAssertTrue(connection.isUsingImmutableSnapshot) + } + + let ledger = HermesUsageLedger(fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-09T07:00:00Z")) + let reader = HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T12:00:00Z") }) + + let usage = try await reader.readUsage( + from: tokiTestISODate("2026-04-09T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(usage.inputTokens, 123) + } + + func test_hermesReader_readsUncheckpointedUsageFromLiveWAL() async throws { + let fileManager = FileManager.default + let tempDir = try makeHermesTemporaryDirectory() + defer { try? fileManager.removeItem(at: tempDir) } + + let dbURL = tempDir.appendingPathComponent("state.db") + try createHermesStateDB( + at: dbURL, + rows: [hermesSingleCounterFixture(id: "live-wal", inputTokens: 123)]) + + var writer: OpaquePointer? + guard sqlite3_open(dbURL.path, &writer) == SQLITE_OK, let writer else { + throw NSError(domain: "HermesReaderTests", code: 26) + } + defer { sqlite3_close(writer) } + guard sqlite3_exec( + writer, + """ + PRAGMA journal_mode=WAL; + PRAGMA wal_autocheckpoint=0; + UPDATE sessions SET input_tokens = 777 WHERE id = 'live-wal'; + """, + nil, + nil, + nil) == SQLITE_OK, + sqlite3_changes(writer) == 1 else { + throw NSError(domain: "HermesReaderTests", code: 27) + } + XCTAssertTrue(fileManager.fileExists(atPath: "\(dbURL.path)-wal")) + XCTAssertTrue(fileManager.fileExists(atPath: "\(dbURL.path)-shm")) + + do { + let connection = try XCTUnwrap(HermesSQLiteConnection.open(atPath: dbURL.path)) + XCTAssertFalse(connection.isUsingImmutableSnapshot) + } + + let ledger = HermesUsageLedger(fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-09T07:00:00Z")) + let reader = HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T12:00:00Z") }) + + let usage = try await reader.readUsage( + from: tokiTestISODate("2026-04-09T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(usage.inputTokens, 777) + } + + func test_hermesDatabaseSourceSnapshotDetectsSidecarAppearance() throws { + let fileManager = FileManager.default + let tempDir = try makeHermesTemporaryDirectory() + defer { try? fileManager.removeItem(at: tempDir) } + + let dbURL = tempDir.appendingPathComponent("state.db") + try createHermesStateDB( + at: dbURL, + rows: [hermesSingleCounterFixture(id: "snapshot", inputTokens: 1)]) + let snapshot = try XCTUnwrap( + HermesDatabaseSourceSnapshot.captureForImmutableFallback(databaseURL: dbURL)) + + try Data().write(to: URL(fileURLWithPath: "\(dbURL.path)-wal")) + + XCTAssertFalse(snapshot.isCurrent()) + } } extension HermesReaderTests { @@ -909,6 +1026,31 @@ func createHermesStateDB( } } +func convertHermesDatabaseToWALWithoutSidecars(databaseURL: URL) throws { + var database: OpaquePointer? + guard sqlite3_open(databaseURL.path, &database) == SQLITE_OK, let database else { + throw NSError(domain: "HermesReaderTests", code: 28) + } + guard sqlite3_exec( + database, + "PRAGMA journal_mode=WAL; PRAGMA wal_checkpoint(TRUNCATE);", + nil, + nil, + nil) == SQLITE_OK else { + sqlite3_close(database) + throw NSError(domain: "HermesReaderTests", code: 29) + } + sqlite3_close(database) + + let fileManager = FileManager.default + for suffix in ["-wal", "-shm"] { + let sidecarURL = URL(fileURLWithPath: "\(databaseURL.path)\(suffix)") + if fileManager.fileExists(atPath: sidecarURL.path) { + try fileManager.removeItem(at: sidecarURL) + } + } +} + func insertHermesModelUsage( databaseURL: URL, rows: [HermesModelUsageFixture]) throws { @@ -1026,7 +1168,8 @@ func updateHermesModelUsage( databaseURL: URL, sessionID: String, task: String, - inputTokens: Int) throws { + inputTokens: Int, + actualCost: Double? = nil) throws { var database: OpaquePointer? guard sqlite3_open(databaseURL.path, &database) == SQLITE_OK, let database else { throw NSError(domain: "HermesReaderTests", code: 18) @@ -1038,7 +1181,7 @@ func updateHermesModelUsage( database, """ UPDATE session_model_usage - SET input_tokens = ? + SET input_tokens = ?, actual_cost_usd = COALESCE(?, actual_cost_usd) WHERE session_id = ? AND task = ? """, -1, @@ -1049,8 +1192,9 @@ func updateHermesModelUsage( defer { sqlite3_finalize(statement) } guard sqlite3_bind_int64(statement, 1, Int64(inputTokens)) == SQLITE_OK, - sqlite3_bind_text(statement, 2, sessionID, -1, hermesTestSQLiteTransient) == SQLITE_OK, - sqlite3_bind_text(statement, 3, task, -1, hermesTestSQLiteTransient) == SQLITE_OK, + bindHermesDouble(actualCost, at: 2, in: statement), + sqlite3_bind_text(statement, 3, sessionID, -1, hermesTestSQLiteTransient) == SQLITE_OK, + sqlite3_bind_text(statement, 4, task, -1, hermesTestSQLiteTransient) == SQLITE_OK, sqlite3_step(statement) == SQLITE_DONE, sqlite3_changes(database) == 1 else { throw NSError(domain: "HermesReaderTests", code: 20) diff --git a/TokiTests/HermesUsageLedgerTests.swift b/TokiTests/HermesUsageLedgerTests.swift index 1479305..fe3d65d 100644 --- a/TokiTests/HermesUsageLedgerTests.swift +++ b/TokiTests/HermesUsageLedgerTests.swift @@ -1,3 +1,4 @@ +// swiftlint:disable file_length import TokiSyncProtocol import TokiUsageCore import XCTest @@ -61,6 +62,142 @@ final class HermesUsageLedgerTests: XCTestCase { try assertHermesLedgerPrivacy(at: fixture.ledgerURL) } + func test_hermesUsageLedger_splitsMixedModelUsageAndKeepsOnlyResidualUnattributed() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledger = HermesUsageLedger(fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + let baselineAt = tokiTestISODate("2026-04-10T08:00:00Z") + let startedAt = tokiTestISODate("2026-04-10T09:00:00Z") + let observedAt = tokiTestISODate("2026-04-10T10:00:00Z") + try await ledger.refresh(observations: [], observedAt: baselineAt) + + try await ledger.refresh( + observations: [HermesSessionObservation( + sessionID: "mixed-model-session", + startedAt: startedAt, + earliestActivityAt: startedAt, + latestActivityAt: observedAt, + model: nil, + counters: HermesTokenCounters( + inputTokens: 150, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + modelCounters: [ + "gpt-5.6-sol": HermesTokenCounters( + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + "kr/claude-opus-5": HermesTokenCounters( + inputTokens: 30, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + ], + cost: 0, + projectName: nil, + attributionQuality: .unknown)], + observedAt: observedAt) + + let events = try await ledger.events( + from: baselineAt, + to: observedAt.addingTimeInterval(1)) + let tokensByModel = events.reduce(into: [String: Int]()) { result, event in + result[event.model ?? "Mixed / Unattributed", default: 0] += event.counters.totalTokens + } + + XCTAssertEqual(tokensByModel, [ + "gpt-5.6-sol": 100, + "kr/claude-opus-5": 30, + "Mixed / Unattributed": 20, + ]) + XCTAssertEqual(events.reduce(0) { $0 + $1.counters.totalTokens }, 150) + } + + func test_hermesUsageLedger_establishesModelBaselineBeforeSplittingLegacySessionDeltas() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledger = HermesUsageLedger(fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + let baselineAt = tokiTestISODate("2026-04-10T08:00:00Z") + let startedAt = tokiTestISODate("2026-04-09T08:00:00Z") + try await ledger.refresh(observations: [], observedAt: baselineAt) + + func observation( + inputTokens: Int, + modelCounters: [String: HermesTokenCounters]?) -> HermesSessionObservation { + HermesSessionObservation( + sessionID: "legacy-mixed-session", + startedAt: startedAt, + earliestActivityAt: nil, + latestActivityAt: nil, + model: nil, + counters: HermesTokenCounters( + inputTokens: inputTokens, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + modelCounters: modelCounters, + cost: 0, + projectName: nil, + attributionQuality: .unknown) + } + + try await ledger.refresh( + observations: [observation(inputTokens: 100, modelCounters: nil)], + observedAt: tokiTestISODate("2026-04-10T09:00:00Z")) + try await ledger.refresh( + observations: [observation(inputTokens: 130, modelCounters: [ + "gpt-5.6-sol": HermesTokenCounters( + inputTokens: 80, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + "kr/claude-opus-5": HermesTokenCounters( + inputTokens: 50, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + ])], + observedAt: tokiTestISODate("2026-04-10T10:00:00Z")) + try await ledger.refresh( + observations: [observation(inputTokens: 160, modelCounters: [ + "gpt-5.6-sol": HermesTokenCounters( + inputTokens: 90, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + "kr/claude-opus-5": HermesTokenCounters( + inputTokens: 70, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + ])], + observedAt: tokiTestISODate("2026-04-10T11:00:00Z")) + + let events = try await ledger.events( + from: baselineAt, + to: tokiTestISODate("2026-04-10T12:00:00Z")) + let tokensByModel = events.reduce(into: [String: Int]()) { result, event in + result[event.model ?? "Mixed / Unattributed", default: 0] += event.counters.totalTokens + } + + XCTAssertEqual(tokensByModel, [ + "Mixed / Unattributed": 30, + "gpt-5.6-sol": 10, + "kr/claude-opus-5": 20, + ]) + XCTAssertEqual(events.reduce(0) { $0 + $1.counters.totalTokens }, 60) + } + func test_hermesReader_rebaselinesCounterDecreaseWithoutNegativeOrDuplicateUsage() async throws { let tempDir = try makeHermesTemporaryDirectory() defer { try? FileManager.default.removeItem(at: tempDir) } diff --git a/TokiTests/UsageServiceBehaviorTests.swift b/TokiTests/UsageServiceBehaviorTests.swift index 6abfb02..9a8f8aa 100644 --- a/TokiTests/UsageServiceBehaviorTests.swift +++ b/TokiTests/UsageServiceBehaviorTests.swift @@ -229,6 +229,33 @@ final class UsageServiceBehaviorTests: XCTestCase { } } +extension UsageServiceBehaviorTests { + func test_usageService_surfacesUnattributedTokenEventsInByModelBreakdown() async { + let recorder = MockReaderRecorder() + let reader = MockReader(name: "Hermes", recorder: recorder) { _, _ in + var usage = RawTokenUsage(inputTokens: 120) + usage.recordTokenEvent( + timestamp: behaviorTestISODate("2026-04-10T09:00:00Z"), + source: "Hermes", + model: nil, + inputTokens: 120, + outputTokens: 0) + return usage + } + + let service = await MainActor.run { UsageService(readers: [reader]) } + await MainActor.run { service.selectDay(behaviorTestISODate("2026-04-10T12:00:00Z")) } + await service.refresh() + + let usageData = await MainActor.run { service.usageData } + let unattributed = usageData.perModel.first { $0.modelID == "Mixed / Unattributed" } + + XCTAssertEqual(unattributed?.totalTokens, 120) + XCTAssertEqual(unattributed?.sources, ["Hermes"]) + XCTAssertEqual(unattributed?.isPriceKnown, false) + } +} + final class UsageServicePeriodTotalsConcurrencyTests: XCTestCase { func test_usageService_remoteSyncChangeRejectsSupersededPeriodTokenTotalsLoad() async throws { let suiteName = "UsageServicePeriodTotalsConcurrencyTests.\(UUID().uuidString)" From 538acb8a8e89bb870259ba4a2c58f3468aa3d091 Mon Sep 17 00:00:00 2001 From: choi138 Date: Fri, 31 Jul 2026 11:14:25 +0900 Subject: [PATCH 02/10] fix(hermes): address attribution review feedback --- Sources/TokiUsageReaders/HermesReader.swift | 75 +++++-- .../HermesSQLiteConnection.swift | 8 +- .../HermesUsageAttribution.swift | 79 ++++--- .../HermesUsageCostBreakdown.swift | 64 +++++- .../TokiUsageReaders/HermesUsageLedger.swift | 52 ++--- .../HermesUsageLedgerSupport.swift | 53 +---- .../HermesUsageObservation.swift | 88 ++++++++ .../HermesUsageResolver.swift | 195 +++++++++++------ TokiTests/HermesPricingRefreshTests.swift | 206 ++++++++++++++++++ TokiTests/HermesReaderTests.swift | 9 + 10 files changed, 624 insertions(+), 205 deletions(-) create mode 100644 Sources/TokiUsageReaders/HermesUsageObservation.swift diff --git a/Sources/TokiUsageReaders/HermesReader.swift b/Sources/TokiUsageReaders/HermesReader.swift index 1e16047..e5a56cd 100644 --- a/Sources/TokiUsageReaders/HermesReader.swift +++ b/Sources/TokiUsageReaders/HermesReader.swift @@ -32,12 +32,15 @@ public struct HermesReader: TokenReader { } public func readUsage(from startDate: Date, to endDate: Date) async throws -> RawTokenUsage { + let observedAt = now() if let observations = try readDatabaseSnapshot({ database in - try readSessionObservations(from: database) + try readSessionObservations( + from: database, + modelPricingTimestamp: observedAt) }) { try await usageLedger.refresh( observations: observations, - observedAt: now()) + observedAt: observedAt) } let events = try await usageLedger.events(from: startDate, to: endDate) @@ -45,8 +48,11 @@ public struct HermesReader: TokenReader { } public func coverageStatus() throws -> HermesUsageCoverageStatus { - try readDatabaseSnapshot { database in - try readSessionModelUsage(from: database).coverage + let modelPricingTimestamp = now() + return try readDatabaseSnapshot { database in + try readSessionModelUsage( + from: database, + modelPricingTimestamp: modelPricingTimestamp).coverage } ?? HermesUsageCoverageStatus(unmeteredMainAPICallCount: 0) } @@ -70,12 +76,16 @@ public struct HermesReader: TokenReader { return nil } - private func readSessionObservations(from database: OpaquePointer) throws -> [HermesSessionObservation] { + private func readSessionObservations( + from database: OpaquePointer, + modelPricingTimestamp: Date) throws -> [HermesSessionObservation] { guard sqlite3_exec(database, "BEGIN DEFERRED TRANSACTION", nil, nil, nil) == SQLITE_OK else { throw HermesSQLiteError(operation: "begin read transaction", database: database) } do { - let observations = try readSessionObservationsInSnapshot(from: database) + let observations = try readSessionObservationsInSnapshot( + from: database, + modelPricingTimestamp: modelPricingTimestamp) guard sqlite3_exec(database, "COMMIT", nil, nil, nil) == SQLITE_OK else { throw HermesSQLiteError(operation: "commit read transaction", database: database) } @@ -87,8 +97,11 @@ public struct HermesReader: TokenReader { } private func readSessionObservationsInSnapshot( - from database: OpaquePointer) throws -> [HermesSessionObservation] { - let modelUsageBySessionID = try readSessionModelUsage(from: database).usageBySessionID + from database: OpaquePointer, + modelPricingTimestamp: Date) throws -> [HermesSessionObservation] { + let modelUsageBySessionID = try readSessionModelUsage( + from: database, + modelPricingTimestamp: modelPricingTimestamp).usageBySessionID let statement = try preparedUsageStatement(in: database) defer { sqlite3_finalize(statement) } @@ -152,7 +165,8 @@ public struct HermesReader: TokenReader { // swiftlint:disable:next function_body_length private func readSessionModelUsage( - from database: OpaquePointer) throws -> HermesSessionModelUsageReadResult { + from database: OpaquePointer, + modelPricingTimestamp: Date) throws -> HermesSessionModelUsageReadResult { guard try tableExists("session_model_usage", in: database) else { return .empty } let requiredColumns: Set = [ "session_id", @@ -215,13 +229,14 @@ public struct HermesReader: TokenReader { counters: counters, estimatedCost: max(0, sqlite3_column_double(statement, 9)), actualCost: max(0, sqlite3_column_double(statement, 10)), - timestamp: nil) + timestamp: modelPricingTimestamp) usageBySessionID[sessionID, default: []].append( HermesSessionModelUsage( model: model, counters: counters, cost: resolvedCost.value, - costIsDerivedFromModelPricing: resolvedCost.isDerivedFromModelPricing)) + costIsDerivedFromModelPricing: resolvedCost.isDerivedFromModelPricing, + modelPricingTimestamp: resolvedCost.modelPricingTimestamp)) let hasReportedTokens = counters.inputTokens > 0 || counters.outputTokens > 0 || counters.cacheReadTokens > 0 @@ -298,7 +313,6 @@ public struct HermesReader: TokenReader { events: [HermesUsageLedgerEvent], clippingEndDate: Date) -> RawTokenUsage { var result = RawTokenUsage() - var activityEvents: [ActivityTimeEvent] = [] for event in events { let counters = event.counters @@ -315,12 +329,6 @@ public struct HermesReader: TokenReader { result.perModel[model, default: PerModelUsage()].sources.insert(name) } - activityEvents.append( - ActivityTimeEvent( - streamID: event.sessionIdentifier, - timestamp: event.timestamp, - key: event.model)) - result.recordTokenEvent( timestamp: event.timestamp, source: name, @@ -337,9 +345,32 @@ public struct HermesReader: TokenReader { quality: event.attributionQuality)) } + let activityEvents = Self.activityEvents(from: events) result.mergeActivityEvents(activityEvents, source: name, clippingEndDate: clippingEndDate) return result } + + private static func activityEvents( + from events: [HermesUsageLedgerEvent]) -> [ActivityTimeEvent] { + Dictionary(grouping: events) { event in + HermesActivityEventIdentity( + streamID: event.sessionIdentifier, + timestamp: event.timestamp) + } + .map { identity, groupedEvents in + let models = Set(groupedEvents.map(\.model)) + let model = models.count == 1 ? models.first ?? nil : nil + return ActivityTimeEvent( + streamID: identity.streamID, + timestamp: identity.timestamp, + key: model) + } + .sorted { lhs, rhs in + if lhs.timestamp != rhs.timestamp { return lhs.timestamp < rhs.timestamp } + if lhs.streamID != rhs.streamID { return lhs.streamID < rhs.streamID } + return (lhs.key ?? "") < (rhs.key ?? "") + } + } } // swiftlint:enable type_body_length @@ -353,6 +384,11 @@ private struct HermesSessionModelUsageReadResult { let coverage: HermesUsageCoverageStatus } +private struct HermesActivityEventIdentity: Hashable { + let streamID: String + let timestamp: Date +} + private struct HermesSessionUsageRow { let sessionID: String let startedAt: Date @@ -366,6 +402,7 @@ private struct HermesSessionUsageRow { let reasoningTokens: Int let cost: Double let costIsDerivedFromModelPricing: Bool + let modelPricingTimestamp: Date? let projectName: String? let attributionQuality: AttributionQuality @@ -396,6 +433,7 @@ private struct HermesSessionUsageRow { timestamp: startedAt) cost = resolvedCost.value costIsDerivedFromModelPricing = resolvedCost.isDerivedFromModelPricing + modelPricingTimestamp = resolvedCost.modelPricingTimestamp if sqlite3_column_type(statement, 12) == SQLITE_NULL { earliestActivityAt = nil @@ -429,6 +467,7 @@ private struct HermesSessionUsageRow { reasoningTokens: reasoningTokens), cost: cost, costIsDerivedFromModelPricing: costIsDerivedFromModelPricing, + modelPricingTimestamp: modelPricingTimestamp, projectName: projectName, attributionQuality: attributionQuality) } diff --git a/Sources/TokiUsageReaders/HermesSQLiteConnection.swift b/Sources/TokiUsageReaders/HermesSQLiteConnection.swift index d541263..53d861f 100644 --- a/Sources/TokiUsageReaders/HermesSQLiteConnection.swift +++ b/Sources/TokiUsageReaders/HermesSQLiteConnection.swift @@ -42,7 +42,7 @@ final class HermesSQLiteConnection { immutableSnapshot: nil) } catch let error as HermesSQLiteError { let databaseURL = URL(fileURLWithPath: path) - guard error.code == SQLITE_CANTOPEN, + guard hermesSQLiteShouldRetryImmutableFallback(after: error.code), let snapshot = HermesDatabaseSourceSnapshot.captureForImmutableFallback( databaseURL: databaseURL, fileManager: fileManager) else { @@ -110,6 +110,12 @@ final class HermesSQLiteConnection { } } +func hermesSQLiteShouldRetryImmutableFallback(after resultCode: Int32) -> Bool { + let primaryResultCode = resultCode & 0xFF + return primaryResultCode == SQLITE_CANTOPEN + || primaryResultCode == SQLITE_READONLY +} + struct HermesDatabaseSourceSnapshot: Equatable { let databaseURL: URL let databaseSignature: HermesDatabaseFileSignature diff --git a/Sources/TokiUsageReaders/HermesUsageAttribution.swift b/Sources/TokiUsageReaders/HermesUsageAttribution.swift index 0f0b563..c3a5972 100644 --- a/Sources/TokiUsageReaders/HermesUsageAttribution.swift +++ b/Sources/TokiUsageReaders/HermesUsageAttribution.swift @@ -6,7 +6,7 @@ private struct HermesUsageEventPart { } private struct HermesUsageCostAllocationBasis { - let pricedCosts: [Double] + let assignedCosts: [Double] let remainderTokenCounts: [Int] } @@ -15,9 +15,11 @@ func hermesUsageEvents( timestamp: Date, observation: HermesSessionObservation, previousModelCounters: [String: HermesTokenCounters]?, + previousModelReportedCosts: [String: Double]?, previousModelPricingCounters: [String: HermesTokenCounters]?, counters: HermesTokenCounters, - cost: Double) -> [HermesUsageLedgerEvent] { + cost: Double, + pricingTimestamp: Date) -> [HermesUsageLedgerEvent] { let parts = hermesUsageEventParts( observation: observation, previousModelCounters: previousModelCounters, @@ -26,11 +28,15 @@ func hermesUsageEvents( current: observation.modelPricingCounters, previous: previousModelPricingCounters, maximumDelta: counters) + let modelReportedCostDeltas = hermesModelReportedCostDeltas( + current: observation.modelReportedCosts, + previous: previousModelReportedCosts) let costs = hermesAllocatedUsageCosts( totalCost: cost, parts: parts, modelPricingDeltas: modelPricingDeltas, - timestamp: timestamp) + modelReportedCostDeltas: modelReportedCostDeltas, + pricingTimestamp: pricingTimestamp) return zip(parts, costs).map { part, allocatedCost in HermesUsageLedgerEvent( sessionIdentifier: identifier, @@ -73,26 +79,28 @@ private func hermesAllocatedUsageCosts( totalCost: Double, parts: [HermesUsageEventPart], modelPricingDeltas: [String: HermesTokenCounters]?, - timestamp: Date) -> [Double] { + modelReportedCostDeltas: [String: Double]?, + pricingTimestamp: Date) -> [Double] { guard totalCost > 0, !parts.isEmpty else { return Array(repeating: 0, count: parts.count) } let fallbackBasis = HermesUsageCostAllocationBasis( - pricedCosts: Array(repeating: 0, count: parts.count), + assignedCosts: Array(repeating: 0, count: parts.count), remainderTokenCounts: parts.map(\.counters.totalTokens)) let basis = hermesUsageCostAllocationBasis( parts: parts, modelPricingDeltas: modelPricingDeltas, - timestamp: timestamp) ?? fallbackBasis - let pricedTotal = basis.pricedCosts.reduce(0, +) - if pricedTotal > totalCost, pricedTotal > 0 { - let scale = totalCost / pricedTotal - return basis.pricedCosts.map { $0 * scale } + modelReportedCostDeltas: modelReportedCostDeltas, + pricingTimestamp: pricingTimestamp) ?? fallbackBasis + let assignedTotal = basis.assignedCosts.reduce(0, +) + if assignedTotal > totalCost, assignedTotal > 0 { + let scale = totalCost / assignedTotal + return basis.assignedCosts.map { $0 * scale } } - var allocations = basis.pricedCosts - let remainder = totalCost - pricedTotal + var allocations = basis.assignedCosts + let remainder = totalCost - assignedTotal guard remainder > 0 else { return allocations } var remainderTokenCounts = basis.remainderTokenCounts @@ -124,38 +132,57 @@ private func hermesAllocatedUsageCosts( private func hermesUsageCostAllocationBasis( parts: [HermesUsageEventPart], modelPricingDeltas: [String: HermesTokenCounters]?, - timestamp: Date) -> HermesUsageCostAllocationBasis? { - guard let modelPricingDeltas else { return nil } + modelReportedCostDeltas: [String: Double]?, + pricingTimestamp: Date) -> HermesUsageCostAllocationBasis? { + guard modelPricingDeltas != nil || modelReportedCostDeltas != nil else { + return nil + } - var unmatchedModels = Set(modelPricingDeltas.keys) - var pricedCosts: [Double] = [] + var unmatchedPricingModels = Set(modelPricingDeltas?.keys.map { $0 } ?? []) + var unmatchedReportedModels = Set(modelReportedCostDeltas?.keys.map { $0 } ?? []) + var assignedCosts: [Double] = [] var remainderTokenCounts: [Int] = [] for part in parts { guard let model = part.model else { - pricedCosts.append(0) + assignedCosts.append(0) remainderTokenCounts.append(part.counters.totalTokens) continue } - let pricingCounters = modelPricingDeltas[model] ?? .zero - guard !part.counters.hasDecrease(comparedTo: pricingCounters) else { + let pricingCounters = modelPricingDeltas?[model] ?? .zero + if modelPricingDeltas != nil, + part.counters.hasDecrease(comparedTo: pricingCounters) { return nil } - unmatchedModels.remove(model) + unmatchedPricingModels.remove(model) + unmatchedReportedModels.remove(model) let pricedCost = hermesModelPricedCost( counters: pricingCounters, model: model, - timestamp: timestamp) - guard pricedCost.isFinite, pricedCost >= 0 else { return nil } + timestamp: pricingTimestamp) + let reportedCost = modelReportedCostDeltas?[model] ?? 0 + let assignedCost = pricedCost + reportedCost + guard pricedCost.isFinite, + pricedCost >= 0, + reportedCost.isFinite, + reportedCost >= 0, + assignedCost.isFinite else { + return nil + } - pricedCosts.append(pricedCost) + assignedCosts.append(assignedCost) remainderTokenCounts.append( - part.counters.subtracting(pricingCounters).totalTokens) + modelReportedCostDeltas == nil + ? part.counters.subtracting(pricingCounters).totalTokens + : 0) + } + guard unmatchedPricingModels.isEmpty, + unmatchedReportedModels.isEmpty else { + return nil } - guard unmatchedModels.isEmpty else { return nil } return HermesUsageCostAllocationBasis( - pricedCosts: pricedCosts, + assignedCosts: assignedCosts, remainderTokenCounts: remainderTokenCounts) } diff --git a/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift b/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift index e7f7be7..79beb0d 100644 --- a/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift +++ b/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift @@ -3,13 +3,59 @@ import TokiUsageCore func hermesReportedCostBreakdownIsValid( _ reportedCost: Double?, + modelReportedCosts: [String: Double]? = nil, modelPricingCounters: [String: HermesTokenCounters]?, totalCost: Double) -> Bool { - guard let reportedCost else { return true } + guard let reportedCost else { return modelReportedCosts == nil } return reportedCost.isFinite && reportedCost >= 0 && reportedCost <= totalCost && modelPricingCounters != nil + && hermesModelReportedCostsAreValid( + modelReportedCosts, + totalReportedCost: reportedCost) +} + +func hermesModelReportedCostDeltas( + current: [String: Double]?, + previous: [String: Double]?) -> [String: Double]? { + guard let current, let previous else { return nil } + + var deltas: [String: Double] = [:] + for model in Set(current.keys).union(previous.keys) { + let currentCost = current[model] ?? 0 + let previousCost = previous[model] ?? 0 + guard currentCost.isFinite, + previousCost.isFinite, + currentCost >= previousCost else { + return nil + } + let delta = currentCost - previousCost + if delta > 0 { + deltas[model] = delta + } + } + return deltas +} + +private func hermesModelReportedCostsAreValid( + _ modelReportedCosts: [String: Double]?, + totalReportedCost: Double) -> Bool { + guard let modelReportedCosts else { return true } + + var combinedCost = 0.0 + for (model, cost) in modelReportedCosts { + guard !model.isEmpty, + model.utf8.count <= 512, + cost.isFinite, + cost >= 0, + (combinedCost + cost).isFinite else { + return false + } + combinedCost += cost + } + let tolerance = 0.000_000_001 * max(1, totalReportedCost) + return abs(combinedCost - totalReportedCost) <= tolerance } func hermesModelPricedCost( @@ -53,7 +99,7 @@ func hermesIncrementalCost( observation: HermesSessionObservation, previous: HermesUsageLedgerBaseline, delta: HermesTokenCounters, - timestamp: Date) -> Double? { + pricingTimestamp: Date) -> Double? { if observation.reportedCost != nil || previous.reportedCost != nil { if let reportedCost = observation.reportedCost, let previousReportedCost = previous.reportedCost { @@ -62,7 +108,7 @@ func hermesIncrementalCost( current: observation.modelPricingCounters, previous: previous.modelPricingCounters, maximumDelta: delta, - timestamp: timestamp) { + timestamp: pricingTimestamp) { let cost = reportedCostDelta + pricedCost return cost.isFinite ? cost : nil } @@ -80,7 +126,7 @@ func hermesIncrementalCost( return hermesModelPricedCost( counters: delta, model: model, - timestamp: timestamp) + timestamp: pricingTimestamp) } guard let reportedCost = observation.reportedCost, @@ -96,14 +142,14 @@ func hermesIncrementalCost( current: observation.modelPricingCounters, previous: previous.modelPricingCounters, maximumDelta: delta, - timestamp: timestamp) { + timestamp: pricingTimestamp) { return detailedCost } guard let model = observation.model else { return nil } return hermesModelPricedCost( counters: delta, model: model, - timestamp: timestamp) + timestamp: pricingTimestamp) } if observation.costIsDerivedFromModelPricing { @@ -111,13 +157,13 @@ func hermesIncrementalCost( current: observation.modelPricingCounters, previous: previous.modelPricingCounters, maximumDelta: delta, - timestamp: timestamp) { + timestamp: pricingTimestamp) { return detailedCost } return hermesModelPricedCost( counters: delta, model: observation.model, - timestamp: timestamp) + timestamp: pricingTimestamp) } if observation.cost >= previous.cost { return observation.cost - previous.cost @@ -125,5 +171,5 @@ func hermesIncrementalCost( return hermesModelPricedCost( counters: delta, model: observation.model, - timestamp: timestamp) + timestamp: pricingTimestamp) } diff --git a/Sources/TokiUsageReaders/HermesUsageLedger.swift b/Sources/TokiUsageReaders/HermesUsageLedger.swift index 8515145..b686814 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedger.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedger.swift @@ -145,7 +145,7 @@ private extension HermesUsageLedger { previousSuccessfulObservationAt: Date?, identifierHasher: SnapshotOpaqueIdentifierHasher, to candidate: inout HermesUsageLedgerDocument) throws -> Bool { - try validate(observation, observedAt: observedAt) + try validateHermesUsageObservation(observation, observedAt: observedAt) let identifier = identifierHasher.identifier(for: observation.sessionID) let previous = candidate.baselines[identifier] let currentBaseline = baseline( @@ -184,11 +184,12 @@ private extension HermesUsageLedger { observation: observation, previous: previous, observedAt: observedAt) + let pricingTimestamp = observation.modelPricingTimestamp ?? timestamp guard let cost = hermesIncrementalCost( observation: observation, previous: previous, delta: delta, - timestamp: timestamp) else { + pricingTimestamp: pricingTimestamp) else { candidate.baselines[identifier] = currentBaseline try addUnattributed( identifier: identifier, @@ -203,9 +204,11 @@ private extension HermesUsageLedger { timestamp: timestamp, observation: observation, previousModelCounters: previous.modelCounters, + previousModelReportedCosts: previous.modelReportedCosts, previousModelPricingCounters: previous.modelPricingCounters, counters: delta, - cost: cost) { + cost: cost, + pricingTimestamp: pricingTimestamp) { append( event, to: &candidate.events) @@ -227,16 +230,19 @@ private extension HermesUsageLedger { observation, after: previousSuccessfulObservationAt, observedAt: observedAt) { + let timestamp = initialTimestamp( + observation: observation, + observedAt: observedAt) for event in hermesUsageEvents( identifier: identifier, - timestamp: initialTimestamp( - observation: observation, - observedAt: observedAt), + timestamp: timestamp, observation: observation, previousModelCounters: observation.modelCounters.map { _ in [:] }, + previousModelReportedCosts: observation.modelReportedCosts.map { _ in [:] }, previousModelPricingCounters: observation.modelPricingCounters.map { _ in [:] }, counters: observation.counters, - cost: observation.cost) { + cost: observation.cost, + pricingTimestamp: observation.modelPricingTimestamp ?? timestamp) { append( event, to: &candidate.events) @@ -440,6 +446,7 @@ private extension HermesUsageLedger { modelCounters: observation.modelCounters, cost: observation.cost, reportedCost: observation.reportedCost, + modelReportedCosts: observation.modelReportedCosts, modelPricingCounters: observation.modelPricingCounters, projectName: observation.projectName, attributionQuality: observation.attributionQuality) @@ -543,37 +550,6 @@ private extension HermesUsageLedger { firstObservedAt: min(existing.firstObservedAt, observedAt)) } - private func validate( - _ observation: HermesSessionObservation, - observedAt: Date) throws { - guard !observation.sessionID.isEmpty, - observation.sessionID.utf8.count <= 4096, - hermesDateIsValid(observation.startedAt), - observation.startedAt <= observedAt, - observation.earliestActivityAt.map(hermesDateIsValid) ?? true, - observation.latestActivityAt.map(hermesDateIsValid) ?? true, - observation.counters.isValid(), - observation.cost.isFinite, - observation.cost >= 0, - hermesReportedCostBreakdownIsValid( - observation.reportedCost, - modelPricingCounters: observation.modelPricingCounters, - totalCost: observation.cost), - observation.reportedCost != nil - || observation.modelPricingCounters == nil - || observation.costIsDerivedFromModelPricing, - hermesModelPricingCountersAreValid( - observation.modelPricingCounters, - within: observation.counters), - hermesModelPricingCountersAreValid( - observation.modelCounters, - within: observation.counters), - observation.model?.utf8.count ?? 0 <= 512, - observation.projectName?.utf8.count ?? 0 <= 512 else { - throw HermesUsageLedgerError.invalidObservation - } - } - private func validate(_ document: HermesUsageLedgerDocument) throws { guard document.schemaVersion == hermesUsageLedgerSchemaVersion, document.baselines.count <= hermesUsageLedgerMaximumBaselines, diff --git a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift index 2172d10..55801a8 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift @@ -11,51 +11,6 @@ let hermesUsageLedgerMaximumEvents = 100_000 let hermesLedgerMaximumCumulativeTokens = Int.max / 8 let hermesUsageLedgerMaximumEventTokenCount = 1_000_000_000 -struct HermesSessionObservation { - let sessionID: String - let startedAt: Date - let earliestActivityAt: Date? - let latestActivityAt: Date? - let model: String? - let counters: HermesTokenCounters - let modelCounters: [String: HermesTokenCounters]? - let cost: Double - let costIsDerivedFromModelPricing: Bool - let reportedCost: Double? - let modelPricingCounters: [String: HermesTokenCounters]? - let projectName: String? - let attributionQuality: AttributionQuality - - init( - sessionID: String, - startedAt: Date, - earliestActivityAt: Date?, - latestActivityAt: Date?, - model: String?, - counters: HermesTokenCounters, - modelCounters: [String: HermesTokenCounters]? = nil, - cost: Double, - costIsDerivedFromModelPricing: Bool = false, - reportedCost: Double? = nil, - modelPricingCounters: [String: HermesTokenCounters]? = nil, - projectName: String?, - attributionQuality: AttributionQuality) { - self.sessionID = sessionID - self.startedAt = startedAt - self.earliestActivityAt = earliestActivityAt - self.latestActivityAt = latestActivityAt - self.model = model - self.counters = counters - self.modelCounters = modelCounters - self.cost = cost - self.costIsDerivedFromModelPricing = costIsDerivedFromModelPricing - self.reportedCost = reportedCost - self.modelPricingCounters = modelPricingCounters - self.projectName = projectName - self.attributionQuality = attributionQuality - } -} - struct HermesTokenCounters: Codable, Equatable { let inputTokens: Int let outputTokens: Int @@ -343,6 +298,7 @@ struct HermesUsageLedgerPrivateBaseline: Codable, Equatable { let modelCounters: [String: HermesTokenCounters]? let cost: Double let reportedCost: Double? + let modelReportedCosts: [String: Double]? let modelPricingCounters: [String: HermesTokenCounters]? let attributionQuality: AttributionQuality @@ -355,6 +311,7 @@ struct HermesUsageLedgerPrivateBaseline: Codable, Equatable { modelCounters = baseline.modelCounters cost = baseline.cost reportedCost = baseline.reportedCost + modelReportedCosts = baseline.modelReportedCosts modelPricingCounters = baseline.modelPricingCounters attributionQuality = baseline.attributionQuality } @@ -369,6 +326,7 @@ struct HermesUsageLedgerPrivateBaseline: Codable, Equatable { modelCounters: modelCounters, cost: cost, reportedCost: reportedCost, + modelReportedCosts: modelReportedCosts, modelPricingCounters: modelPricingCounters, projectName: nil, attributionQuality: attributionQuality) @@ -424,6 +382,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { let modelCounters: [String: HermesTokenCounters]? let cost: Double let reportedCost: Double? + let modelReportedCosts: [String: Double]? let modelPricingCounters: [String: HermesTokenCounters]? let projectName: String? let attributionQuality: AttributionQuality @@ -437,6 +396,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { modelCounters: [String: HermesTokenCounters]? = nil, cost: Double, reportedCost: Double? = nil, + modelReportedCosts: [String: Double]? = nil, modelPricingCounters: [String: HermesTokenCounters]? = nil, projectName: String?, attributionQuality: AttributionQuality) { @@ -448,6 +408,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { self.modelCounters = modelCounters self.cost = cost self.reportedCost = reportedCost + self.modelReportedCosts = modelReportedCosts self.modelPricingCounters = modelPricingCounters self.projectName = projectName self.attributionQuality = attributionQuality @@ -465,6 +426,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { && cost >= 0 && hermesReportedCostBreakdownIsValid( reportedCost, + modelReportedCosts: modelReportedCosts, modelPricingCounters: modelPricingCounters, totalCost: cost) && hermesModelPricingCountersAreValid(modelPricingCounters, within: counters) @@ -478,6 +440,7 @@ struct HermesUsageLedgerBaseline: Codable, Equatable { || modelCounters != previous.modelCounters || cost != previous.cost || reportedCost != previous.reportedCost + || modelReportedCosts != previous.modelReportedCosts || modelPricingCounters != previous.modelPricingCounters || attributionQuality != previous.attributionQuality } diff --git a/Sources/TokiUsageReaders/HermesUsageObservation.swift b/Sources/TokiUsageReaders/HermesUsageObservation.swift new file mode 100644 index 0000000..399a679 --- /dev/null +++ b/Sources/TokiUsageReaders/HermesUsageObservation.swift @@ -0,0 +1,88 @@ +import Foundation +import TokiUsageCore + +struct HermesSessionObservation { + let sessionID: String + let startedAt: Date + let earliestActivityAt: Date? + let latestActivityAt: Date? + let model: String? + let counters: HermesTokenCounters + let modelCounters: [String: HermesTokenCounters]? + let cost: Double + let costIsDerivedFromModelPricing: Bool + let reportedCost: Double? + let modelReportedCosts: [String: Double]? + let modelPricingCounters: [String: HermesTokenCounters]? + let modelPricingTimestamp: Date? + let projectName: String? + let attributionQuality: AttributionQuality + + init( + sessionID: String, + startedAt: Date, + earliestActivityAt: Date?, + latestActivityAt: Date?, + model: String?, + counters: HermesTokenCounters, + modelCounters: [String: HermesTokenCounters]? = nil, + cost: Double, + costIsDerivedFromModelPricing: Bool = false, + reportedCost: Double? = nil, + modelReportedCosts: [String: Double]? = nil, + modelPricingCounters: [String: HermesTokenCounters]? = nil, + modelPricingTimestamp: Date? = nil, + projectName: String?, + attributionQuality: AttributionQuality) { + self.sessionID = sessionID + self.startedAt = startedAt + self.earliestActivityAt = earliestActivityAt + self.latestActivityAt = latestActivityAt + self.model = model + self.counters = counters + self.modelCounters = modelCounters + self.cost = cost + self.costIsDerivedFromModelPricing = costIsDerivedFromModelPricing + self.reportedCost = reportedCost + self.modelReportedCosts = modelReportedCosts + self.modelPricingCounters = modelPricingCounters + self.modelPricingTimestamp = modelPricingTimestamp + self.projectName = projectName + self.attributionQuality = attributionQuality + } +} + +func validateHermesUsageObservation( + _ observation: HermesSessionObservation, + observedAt: Date) throws { + guard !observation.sessionID.isEmpty, + observation.sessionID.utf8.count <= 4096, + hermesDateIsValid(observation.startedAt), + observation.startedAt <= observedAt, + observation.earliestActivityAt.map(hermesDateIsValid) ?? true, + observation.latestActivityAt.map(hermesDateIsValid) ?? true, + observation.counters.isValid(), + observation.cost.isFinite, + observation.cost >= 0, + (observation.modelPricingTimestamp.map { + hermesDateIsValid($0) && $0 <= observedAt + } ?? true), + hermesReportedCostBreakdownIsValid( + observation.reportedCost, + modelReportedCosts: observation.modelReportedCosts, + modelPricingCounters: observation.modelPricingCounters, + totalCost: observation.cost), + observation.reportedCost != nil + || observation.modelPricingCounters == nil + || observation.costIsDerivedFromModelPricing, + hermesModelPricingCountersAreValid( + observation.modelPricingCounters, + within: observation.counters), + hermesModelPricingCountersAreValid( + observation.modelCounters, + within: observation.counters), + observation.model?.utf8.count ?? 0 <= 512, + observation.projectName?.utf8.count ?? 0 <= 512 else { + throw HermesUsageLedgerError.invalidObservation + } +} diff --git a/Sources/TokiUsageReaders/HermesUsageResolver.swift b/Sources/TokiUsageReaders/HermesUsageResolver.swift index ef2362a..05906fc 100644 --- a/Sources/TokiUsageReaders/HermesUsageResolver.swift +++ b/Sources/TokiUsageReaders/HermesUsageResolver.swift @@ -14,68 +14,118 @@ struct HermesSessionModelUsage { let counters: HermesTokenCounters let cost: Double let costIsDerivedFromModelPricing: Bool + let modelPricingTimestamp: Date? } struct HermesResolvedUsageCost { let value: Double let isDerivedFromModelPricing: Bool + let modelPricingTimestamp: Date? } private struct HermesResolvedSessionCost { let value: Double let isDerivedFromModelPricing: Bool let reportedValue: Double + let modelReportedCosts: [String: Double]? let modelPricingCounters: [String: HermesTokenCounters] + let modelPricingTimestamp: Date? +} + +private struct HermesSessionModelUsageAccumulator { + var counters = HermesTokenCounters.zero + var cost = 0.0 + var reportedCost = 0.0 + var costIsDerivedFromModelPricing = true + var countersByModel: [String: HermesTokenCounters] = [:] + var reportedCostsByModel: [String: Double] = [:] + var pricingCountersByModel: [String: HermesTokenCounters] = [:] + var pricingTimestamp: Date? + var models = Set() + var hasUnattributedReportedCost = false + + var resolvedReportedCostsByModel: [String: Double]? { + hasUnattributedReportedCost ? nil : reportedCostsByModel + } + + var hasResolvedReportedCostBreakdown: Bool { + resolvedReportedCostsByModel?.isEmpty == false + } + + mutating func add(_ usage: HermesSessionModelUsage) throws { + guard usage.counters.isValid(), + counters.canAdd(usage.counters, maximum: hermesLedgerMaximumCumulativeTokens), + usage.cost.isFinite, + usage.cost >= 0, + (cost + usage.cost).isFinite, + usage.modelPricingTimestamp.map(hermesDateIsValid) ?? true else { + throw HermesUsageLedgerError.invalidObservation + } + counters = counters.adding(usage.counters) + cost += usage.cost + costIsDerivedFromModelPricing = + costIsDerivedFromModelPricing && usage.costIsDerivedFromModelPricing + if !usage.costIsDerivedFromModelPricing { + try addReportedCost(usage) + } + + guard usage.counters.totalTokens > 0, let model = usage.model else { return } + countersByModel[model] = try addingHermesCounters( + usage.counters, + to: countersByModel[model] ?? .zero) + if usage.costIsDerivedFromModelPricing { + pricingCountersByModel[model] = try addingHermesCounters( + usage.counters, + to: pricingCountersByModel[model] ?? .zero) + try recordPricingTimestamp(usage.modelPricingTimestamp) + } + models.insert(model) + } + + private mutating func addReportedCost(_ usage: HermesSessionModelUsage) throws { + reportedCost += usage.cost + guard usage.cost > 0 else { return } + guard let model = usage.model, usage.counters.totalTokens > 0 else { + hasUnattributedReportedCost = true + return + } + let existingCost = reportedCostsByModel[model] ?? 0 + guard (existingCost + usage.cost).isFinite else { + throw HermesUsageLedgerError.invalidObservation + } + reportedCostsByModel[model] = existingCost + usage.cost + } + + private mutating func recordPricingTimestamp(_ timestamp: Date?) throws { + guard let timestamp else { return } + guard pricingTimestamp == nil || pricingTimestamp == timestamp else { + throw HermesUsageLedgerError.invalidObservation + } + pricingTimestamp = timestamp + } +} + +private func addingHermesCounters( + _ counters: HermesTokenCounters, + to existing: HermesTokenCounters) throws -> HermesTokenCounters { + guard existing.canAdd( + counters, + maximum: hermesLedgerMaximumCumulativeTokens) else { + throw HermesUsageLedgerError.invalidObservation + } + return existing.adding(counters) } enum HermesUsageResolver { static func resolve( session: HermesSessionObservation, modelUsage: [HermesSessionModelUsage]) throws -> HermesSessionObservation { - var modelCounters = HermesTokenCounters.zero - var modelCost = 0.0 - var modelReportedCost = 0.0 - var modelCostIsDerivedFromModelPricing = true - var attributedModelCounters: [String: HermesTokenCounters] = [:] - var modelPricingCounters: [String: HermesTokenCounters] = [:] - var models: Set = [] - + var accumulatedModelUsage = HermesSessionModelUsageAccumulator() for usage in modelUsage { - guard usage.counters.isValid(), - modelCounters.canAdd(usage.counters, maximum: hermesLedgerMaximumCumulativeTokens), - usage.cost.isFinite, - usage.cost >= 0, - (modelCost + usage.cost).isFinite else { - throw HermesUsageLedgerError.invalidObservation - } - modelCounters = modelCounters.adding(usage.counters) - modelCost += usage.cost - modelCostIsDerivedFromModelPricing = - modelCostIsDerivedFromModelPricing && usage.costIsDerivedFromModelPricing - if !usage.costIsDerivedFromModelPricing { - modelReportedCost += usage.cost - } - if usage.counters.totalTokens > 0, let model = usage.model { - let existingCounters = attributedModelCounters[model] ?? .zero - guard existingCounters.canAdd( - usage.counters, - maximum: hermesLedgerMaximumCumulativeTokens) else { - throw HermesUsageLedgerError.invalidObservation - } - attributedModelCounters[model] = existingCounters.adding(usage.counters) - if usage.costIsDerivedFromModelPricing { - let existingPricingCounters = modelPricingCounters[model] ?? .zero - guard existingPricingCounters.canAdd( - usage.counters, - maximum: hermesLedgerMaximumCumulativeTokens) else { - throw HermesUsageLedgerError.invalidObservation - } - modelPricingCounters[model] = existingPricingCounters.adding(usage.counters) - } - models.insert(model) - } + try accumulatedModelUsage.add(usage) } + var models = accumulatedModelUsage.models if session.counters.totalTokens > 0, let model = session.model { models.insert(model) } @@ -83,10 +133,7 @@ enum HermesUsageResolver { let resolvedCost = resolveCost( session: session, hasModelUsage: !modelUsage.isEmpty, - modelCost: modelCost, - modelReportedCost: modelReportedCost, - modelCostIsDerivedFromModelPricing: modelCostIsDerivedFromModelPricing, - modelPricingCounters: modelPricingCounters) + modelUsage: accumulatedModelUsage) return HermesSessionObservation( sessionID: session.sessionID, @@ -94,12 +141,14 @@ enum HermesUsageResolver { earliestActivityAt: session.earliestActivityAt, latestActivityAt: session.latestActivityAt, model: resolvedModel, - counters: session.counters.maximum(modelCounters), - modelCounters: modelUsage.isEmpty ? nil : attributedModelCounters, + counters: session.counters.maximum(accumulatedModelUsage.counters), + modelCounters: modelUsage.isEmpty ? nil : accumulatedModelUsage.countersByModel, cost: resolvedCost.value, costIsDerivedFromModelPricing: resolvedCost.isDerivedFromModelPricing, reportedCost: resolvedCost.reportedValue, + modelReportedCosts: resolvedCost.modelReportedCosts, modelPricingCounters: resolvedCost.modelPricingCounters, + modelPricingTimestamp: resolvedCost.modelPricingTimestamp, projectName: session.projectName, attributionQuality: session.attributionQuality) } @@ -107,41 +156,47 @@ enum HermesUsageResolver { private static func resolveCost( session: HermesSessionObservation, hasModelUsage: Bool, - modelCost: Double, - modelReportedCost: Double, - modelCostIsDerivedFromModelPricing: Bool, - modelPricingCounters: [String: HermesTokenCounters]) -> HermesResolvedSessionCost { + modelUsage: HermesSessionModelUsageAccumulator) -> HermesResolvedSessionCost { let sessionPricingCounters = session.costIsDerivedFromModelPricing ? pricingCounters(model: session.model, counters: session.counters) : [:] let sessionReportedCost = session.costIsDerivedFromModelPricing ? 0 : session.cost - if !hasModelUsage || session.cost > modelCost { + if !hasModelUsage || session.cost > modelUsage.cost { return HermesResolvedSessionCost( value: session.cost, isDerivedFromModelPricing: session.costIsDerivedFromModelPricing, reportedValue: sessionReportedCost, - modelPricingCounters: sessionPricingCounters) + modelReportedCosts: nil, + modelPricingCounters: sessionPricingCounters, + modelPricingTimestamp: session.modelPricingTimestamp) } - if modelCost > session.cost { + if modelUsage.cost > session.cost { return HermesResolvedSessionCost( - value: modelCost, - isDerivedFromModelPricing: modelCostIsDerivedFromModelPricing, - reportedValue: modelReportedCost, - modelPricingCounters: modelPricingCounters) + value: modelUsage.cost, + isDerivedFromModelPricing: modelUsage.costIsDerivedFromModelPricing, + reportedValue: modelUsage.reportedCost, + modelReportedCosts: modelUsage.resolvedReportedCostsByModel, + modelPricingCounters: modelUsage.pricingCountersByModel, + modelPricingTimestamp: modelUsage.pricingTimestamp) } - if !session.costIsDerivedFromModelPricing { + if !session.costIsDerivedFromModelPricing, + !modelUsage.hasResolvedReportedCostBreakdown { return HermesResolvedSessionCost( value: session.cost, isDerivedFromModelPricing: false, reportedValue: sessionReportedCost, - modelPricingCounters: sessionPricingCounters) + modelReportedCosts: nil, + modelPricingCounters: sessionPricingCounters, + modelPricingTimestamp: session.modelPricingTimestamp) } return HermesResolvedSessionCost( - value: modelCost, - isDerivedFromModelPricing: modelCostIsDerivedFromModelPricing, - reportedValue: modelReportedCost, - modelPricingCounters: modelPricingCounters) + value: modelUsage.cost, + isDerivedFromModelPricing: modelUsage.costIsDerivedFromModelPricing, + reportedValue: modelUsage.reportedCost, + modelReportedCosts: modelUsage.resolvedReportedCostsByModel, + modelPricingCounters: modelUsage.pricingCountersByModel, + modelPricingTimestamp: modelUsage.pricingTimestamp) } private static func pricingCounters( @@ -161,16 +216,19 @@ func hermesUsageCost( if actualCost > 0 { return HermesResolvedUsageCost( value: actualCost, - isDerivedFromModelPricing: false) + isDerivedFromModelPricing: false, + modelPricingTimestamp: nil) } if estimatedCost > 0 { return HermesResolvedUsageCost( value: estimatedCost, - isDerivedFromModelPricing: false) + isDerivedFromModelPricing: false, + modelPricingTimestamp: nil) } + let modelPricingTimestamp = timestamp ?? Date() let value = model - .flatMap { modelPrice(for: $0, at: timestamp ?? Date()) }? + .flatMap { modelPrice(for: $0, at: modelPricingTimestamp) }? .cost( input: counters.inputTokens, output: counters.outputTokens + counters.reasoningTokens, @@ -178,5 +236,6 @@ func hermesUsageCost( cacheWrite: counters.cacheWriteTokens) ?? 0 return HermesResolvedUsageCost( value: value, - isDerivedFromModelPricing: true) + isDerivedFromModelPricing: true, + modelPricingTimestamp: modelPricingTimestamp) } diff --git a/TokiTests/HermesPricingRefreshTests.swift b/TokiTests/HermesPricingRefreshTests.swift index 1bf7fbc..e87b7b8 100644 --- a/TokiTests/HermesPricingRefreshTests.swift +++ b/TokiTests/HermesPricingRefreshTests.swift @@ -432,6 +432,211 @@ extension HermesPricingRefreshTests { } } +extension HermesPricingRefreshTests { + // swiftlint:disable:next function_body_length + func test_hermesReader_allocatesInitialDerivedCostsAtModelUsagePricingInstant() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { + ModelPricingSupplement.install([:]) + try? FileManager.default.removeItem(at: tempDir) + } + let firstModel = "dated-hermes-model-a" + let secondModel = "dated-hermes-model-b" + let initialFirstPrice = uniformHermesModelPrice(perMillion: 1) + let updatedFirstPrice = uniformHermesModelPrice(perMillion: 10) + let secondPrice = uniformHermesModelPrice(perMillion: 2) + ModelPricingSupplement.install(priceHistories: [ + firstModel: [ + ModelPricingSupplement.PriceVersion( + effectiveFrom: .distantPast, + price: initialFirstPrice), + ModelPricingSupplement.PriceVersion( + effectiveFrom: tokiTestISODate("2026-04-10T10:00:00Z"), + price: updatedFirstPrice), + ], + secondModel: [ + ModelPricingSupplement.PriceVersion( + effectiveFrom: .distantPast, + price: secondPrice), + ], + ]) + + let dbURL = tempDir.appendingPathComponent("state.db") + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "dated-mixed-session", + startedAt: "2026-04-10T08:30:00Z", + model: nil, + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: nil)]) + try insertHermesMessage( + databaseURL: dbURL, + sessionID: "dated-mixed-session", + timestamp: tokiTestISODate("2026-04-10T09:00:00Z")) + try insertHermesModelUsage( + databaseURL: dbURL, + rows: [ + HermesModelUsageFixture( + sessionID: "dated-mixed-session", + model: firstModel, + task: "first", + apiCallCount: 1, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 0), + HermesModelUsageFixture( + sessionID: "dated-mixed-session", + model: secondModel, + task: "second", + apiCallCount: 1, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 0), + ]) + let ledger = HermesUsageLedger(fileURL: ledgerURL) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T08:00:00Z")) + + let usage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T11:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual( + usage.perModel[firstModel]?.cost ?? -1, + updatedFirstPrice.cost(input: 100, output: 0, cacheRead: 0, cacheWrite: 0), + accuracy: 0.000001) + XCTAssertEqual( + usage.perModel[secondModel]?.cost ?? -1, + secondPrice.cost(input: 100, output: 0, cacheRead: 0, cacheWrite: 0), + accuracy: 0.000001) + XCTAssertEqual(Set(usage.tokenEvents.map(\.timestamp)), [ + tokiTestISODate("2026-04-10T09:00:00Z"), + ]) + XCTAssertEqual(usage.activityEvents.count, 1) + XCTAssertNil(usage.activityEvents.first?.key) + XCTAssertEqual(usage.activeSeconds, 30, accuracy: 0.001) + } + + // swiftlint:disable:next function_body_length + func test_hermesReader_preservesReportedCostByModelAcrossRestart() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let expensiveModel = "reported-expensive-model" + let cheapModel = "reported-cheap-model" + let dbURL = tempDir.appendingPathComponent("state.db") + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "reported-mixed-session", + startedAt: "2026-04-10T09:00:00Z", + model: nil, + inputTokens: 200, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: 11)]) + try insertHermesModelUsage( + databaseURL: dbURL, + rows: [ + HermesModelUsageFixture( + sessionID: "reported-mixed-session", + model: expensiveModel, + task: "expensive", + apiCallCount: 1, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 10), + HermesModelUsageFixture( + sessionID: "reported-mixed-session", + model: cheapModel, + task: "cheap", + apiCallCount: 1, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 1), + ]) + let ledger = HermesUsageLedger(fileURL: ledgerURL) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T08:00:00Z")) + let initialUsage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T10:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-10T10:30:00Z")) + XCTAssertEqual(initialUsage.perModel[expensiveModel]?.cost ?? -1, 10, accuracy: 0.000001) + XCTAssertEqual(initialUsage.perModel[cheapModel]?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual(initialUsage.activityEvents.count, 1) + XCTAssertNil(initialUsage.activityEvents.first?.key) + + try updateHermesModelUsage( + databaseURL: dbURL, + sessionID: "reported-mixed-session", + task: "expensive", + inputTokens: 200, + actualCost: 20) + try updateHermesModelUsage( + databaseURL: dbURL, + sessionID: "reported-mixed-session", + task: "cheap", + inputTokens: 200, + actualCost: 2) + let increment = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: HermesUsageLedger(fileURL: ledgerURL), + now: { tokiTestISODate("2026-04-10T11:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T10:30:00Z"), + to: tokiTestISODate("2026-04-10T12:00:00Z")) + + XCTAssertEqual(increment.inputTokens, 200) + XCTAssertEqual(increment.cost, 11, accuracy: 0.000001) + XCTAssertEqual(increment.perModel[expensiveModel]?.totalTokens, 100) + XCTAssertEqual(increment.perModel[cheapModel]?.totalTokens, 100) + XCTAssertEqual(increment.perModel[expensiveModel]?.cost ?? -1, 10, accuracy: 0.000001) + XCTAssertEqual(increment.perModel[cheapModel]?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual(increment.activityEvents.count, 1) + XCTAssertNil(increment.activityEvents.first?.key) + } +} + private func uniformHermesModelPrice(perMillion: Double) -> ModelPrice { ModelPrice( inputPerMillion: perMillion, @@ -600,6 +805,7 @@ func removeHermesCostBreakdownFromBaselines(at ledgerURL: URL) throws { throw NSError(domain: "HermesPricingRefreshTests", code: 2) } baseline.removeValue(forKey: "reportedCost") + baseline.removeValue(forKey: "modelReportedCosts") baseline.removeValue(forKey: "modelPricingCounters") baselines[identifier] = baseline } diff --git a/TokiTests/HermesReaderTests.swift b/TokiTests/HermesReaderTests.swift index 4bff706..9dd6fb8 100644 --- a/TokiTests/HermesReaderTests.swift +++ b/TokiTests/HermesReaderTests.swift @@ -441,6 +441,15 @@ final class HermesReaderTests: XCTestCase { XCTAssertEqual(usage.inputTokens, 123) } + func test_hermesSQLiteConnection_retriesImmutableFallbackAfterReadOnlyProbeFailure() { + XCTAssertTrue(hermesSQLiteShouldRetryImmutableFallback(after: SQLITE_CANTOPEN)) + XCTAssertTrue(hermesSQLiteShouldRetryImmutableFallback(after: SQLITE_READONLY)) + XCTAssertTrue( + hermesSQLiteShouldRetryImmutableFallback( + after: SQLITE_READONLY | Int32(2 << 8))) + XCTAssertFalse(hermesSQLiteShouldRetryImmutableFallback(after: SQLITE_BUSY)) + } + func test_hermesReader_readsUncheckpointedUsageFromLiveWAL() async throws { let fileManager = FileManager.default let tempDir = try makeHermesTemporaryDirectory() From 2bf7dfb0a02071ef9b664c84310ebb7e8df39e7c Mon Sep 17 00:00:00 2001 From: choi138 Date: Fri, 31 Jul 2026 11:42:35 +0900 Subject: [PATCH 03/10] fix(hermes): preserve residual usage attribution --- Sources/TokiUsageCore/RawTokenUsage.swift | 5 ++ Sources/TokiUsageReaders/HermesReader.swift | 20 ++--- .../HermesSQLiteConnection.swift | 1 + .../HermesUsageAttribution.swift | 46 ++++++++--- .../HermesUsageCostBreakdown.swift | 33 +++++++- .../TokiUsageReaders/HermesUsageLedger.swift | 6 ++ .../HermesUsageLedgerSupport.swift | 2 +- .../HermesUsageResolver.swift | 19 ++++- .../AgentSnapshotModelAttributionTests.swift | 7 +- Toki/Domain/Usage/UsageReportModelStats.swift | 10 ++- TokiTests/HermesPricingRefreshTests.swift | 48 +++++++++-- TokiTests/HermesReaderTests.swift | 79 +++++++++++++++++++ TokiTests/UsageServiceBehaviorTests.swift | 40 +++++++++- 13 files changed, 278 insertions(+), 38 deletions(-) diff --git a/Sources/TokiUsageCore/RawTokenUsage.swift b/Sources/TokiUsageCore/RawTokenUsage.swift index 73ace88..558b1cb 100644 --- a/Sources/TokiUsageCore/RawTokenUsage.swift +++ b/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 diff --git a/Sources/TokiUsageReaders/HermesReader.swift b/Sources/TokiUsageReaders/HermesReader.swift index e5a56cd..cd54df2 100644 --- a/Sources/TokiUsageReaders/HermesReader.swift +++ b/Sources/TokiUsageReaders/HermesReader.swift @@ -32,12 +32,13 @@ public struct HermesReader: TokenReader { } public func readUsage(from startDate: Date, to endDate: Date) async throws -> RawTokenUsage { - let observedAt = now() + let modelPricingTimestamp = now() if let observations = try readDatabaseSnapshot({ database in try readSessionObservations( from: database, - modelPricingTimestamp: observedAt) + modelPricingTimestamp: modelPricingTimestamp) }) { + let observedAt = max(modelPricingTimestamp, now()) try await usageLedger.refresh( observations: observations, observedAt: observedAt) @@ -323,11 +324,10 @@ public struct HermesReader: TokenReader { result.reasoningTokens += counters.reasoningTokens result.cost += event.cost - if let model = event.model { - result.perModel[model, default: PerModelUsage()].totalTokens += counters.totalTokens - result.perModel[model, default: PerModelUsage()].cost += event.cost - result.perModel[model, default: PerModelUsage()].sources.insert(name) - } + let modelGroupingKey = event.model ?? UsageModelGrouping.mixedOrUnattributedKey + result.perModel[modelGroupingKey, default: PerModelUsage()].totalTokens += counters.totalTokens + result.perModel[modelGroupingKey, default: PerModelUsage()].cost += event.cost + result.perModel[modelGroupingKey, default: PerModelUsage()].sources.insert(name) result.recordTokenEvent( timestamp: event.timestamp, @@ -352,14 +352,16 @@ public struct HermesReader: TokenReader { private static func activityEvents( from events: [HermesUsageLedgerEvent]) -> [ActivityTimeEvent] { - Dictionary(grouping: events) { event in + Dictionary(grouping: events.filter { $0.counters.totalTokens > 0 }) { event in HermesActivityEventIdentity( streamID: event.sessionIdentifier, timestamp: event.timestamp) } .map { identity, groupedEvents in let models = Set(groupedEvents.map(\.model)) - let model = models.count == 1 ? models.first ?? nil : nil + let model = models.count == 1 + ? models.first ?? nil + : UsageModelGrouping.mixedOrUnattributedKey return ActivityTimeEvent( streamID: identity.streamID, timestamp: identity.timestamp, diff --git a/Sources/TokiUsageReaders/HermesSQLiteConnection.swift b/Sources/TokiUsageReaders/HermesSQLiteConnection.swift index 53d861f..a311e47 100644 --- a/Sources/TokiUsageReaders/HermesSQLiteConnection.swift +++ b/Sources/TokiUsageReaders/HermesSQLiteConnection.swift @@ -150,6 +150,7 @@ struct HermesDatabaseSourceSnapshot: Equatable { let path = databaseURL.path return fileManager.fileExists(atPath: "\(path)-wal") || fileManager.fileExists(atPath: "\(path)-shm") + || fileManager.fileExists(atPath: "\(path)-journal") } } diff --git a/Sources/TokiUsageReaders/HermesUsageAttribution.swift b/Sources/TokiUsageReaders/HermesUsageAttribution.swift index c3a5972..4e392f7 100644 --- a/Sources/TokiUsageReaders/HermesUsageAttribution.swift +++ b/Sources/TokiUsageReaders/HermesUsageAttribution.swift @@ -15,30 +15,40 @@ func hermesUsageEvents( timestamp: Date, observation: HermesSessionObservation, previousModelCounters: [String: HermesTokenCounters]?, + previousReportedCost: Double?, previousModelReportedCosts: [String: Double]?, previousModelPricingCounters: [String: HermesTokenCounters]?, counters: HermesTokenCounters, cost: Double, pricingTimestamp: Date) -> [HermesUsageLedgerEvent] { + let modelReportedCostDeltas = hermesModelReportedCostDeltas( + current: observation.modelReportedCosts, + previous: previousModelReportedCosts) + let unattributedReportedCostDelta = hermesUnattributedReportedCostDelta( + currentReportedCost: observation.reportedCost, + previousReportedCost: previousReportedCost, + modelReportedCostDeltas: modelReportedCostDeltas) let parts = hermesUsageEventParts( observation: observation, previousModelCounters: previousModelCounters, - counters: counters) + counters: counters, + includeCostOnlyResidual: unattributedReportedCostDelta.map { $0 > 0 } ?? false) let modelPricingDeltas = hermesModelCounterDeltas( current: observation.modelPricingCounters, previous: previousModelPricingCounters, maximumDelta: counters) - let modelReportedCostDeltas = hermesModelReportedCostDeltas( - current: observation.modelReportedCosts, - previous: previousModelReportedCosts) let costs = hermesAllocatedUsageCosts( totalCost: cost, parts: parts, modelPricingDeltas: modelPricingDeltas, modelReportedCostDeltas: modelReportedCostDeltas, + unattributedReportedCostDelta: unattributedReportedCostDelta, pricingTimestamp: pricingTimestamp) - return zip(parts, costs).map { part, allocatedCost in - HermesUsageLedgerEvent( + return zip(parts, costs).compactMap { part, allocatedCost in + guard part.counters.totalTokens > 0 || allocatedCost > 0 else { + return nil + } + return HermesUsageLedgerEvent( sessionIdentifier: identifier, timestamp: timestamp, model: part.model, @@ -52,7 +62,8 @@ func hermesUsageEvents( private func hermesUsageEventParts( observation: HermesSessionObservation, previousModelCounters: [String: HermesTokenCounters]?, - counters: HermesTokenCounters) -> [HermesUsageEventPart] { + counters: HermesTokenCounters, + includeCostOnlyResidual: Bool) -> [HermesUsageEventPart] { guard let modelDeltas = hermesModelCounterDeltas( current: observation.modelCounters, previous: previousModelCounters, @@ -71,6 +82,8 @@ private func hermesUsageEventParts( let residual = counters.subtracting(combinedCounters) if residual.totalTokens > 0 { parts.append(HermesUsageEventPart(model: nil, counters: residual)) + } else if includeCostOnlyResidual { + parts.append(HermesUsageEventPart(model: nil, counters: .zero)) } return parts } @@ -80,6 +93,7 @@ private func hermesAllocatedUsageCosts( parts: [HermesUsageEventPart], modelPricingDeltas: [String: HermesTokenCounters]?, modelReportedCostDeltas: [String: Double]?, + unattributedReportedCostDelta: Double?, pricingTimestamp: Date) -> [Double] { guard totalCost > 0, !parts.isEmpty else { return Array(repeating: 0, count: parts.count) @@ -92,6 +106,7 @@ private func hermesAllocatedUsageCosts( parts: parts, modelPricingDeltas: modelPricingDeltas, modelReportedCostDeltas: modelReportedCostDeltas, + unattributedReportedCostDelta: unattributedReportedCostDelta, pricingTimestamp: pricingTimestamp) ?? fallbackBasis let assignedTotal = basis.assignedCosts.reduce(0, +) if assignedTotal > totalCost, assignedTotal > 0 { @@ -133,6 +148,7 @@ private func hermesUsageCostAllocationBasis( parts: [HermesUsageEventPart], modelPricingDeltas: [String: HermesTokenCounters]?, modelReportedCostDeltas: [String: Double]?, + unattributedReportedCostDelta: Double?, pricingTimestamp: Date) -> HermesUsageCostAllocationBasis? { guard modelPricingDeltas != nil || modelReportedCostDeltas != nil else { return nil @@ -142,10 +158,19 @@ private func hermesUsageCostAllocationBasis( var unmatchedReportedModels = Set(modelReportedCostDeltas?.keys.map { $0 } ?? []) var assignedCosts: [Double] = [] var remainderTokenCounts: [Int] = [] + var unattributedCostWasAssigned = unattributedReportedCostDelta.map { $0 <= 0 } ?? true for part in parts { guard let model = part.model else { - assignedCosts.append(0) - remainderTokenCounts.append(part.counters.totalTokens) + let unattributedCost = unattributedReportedCostDelta ?? 0 + guard unattributedCost.isFinite, unattributedCost >= 0 else { + return nil + } + assignedCosts.append(unattributedCost) + remainderTokenCounts.append( + modelReportedCostDeltas == nil + ? part.counters.totalTokens + : max(1, part.counters.totalTokens)) + unattributedCostWasAssigned = true continue } @@ -178,7 +203,8 @@ private func hermesUsageCostAllocationBasis( : 0) } guard unmatchedPricingModels.isEmpty, - unmatchedReportedModels.isEmpty else { + unmatchedReportedModels.isEmpty, + unattributedCostWasAssigned else { return nil } diff --git a/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift b/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift index 79beb0d..a118f6c 100644 --- a/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift +++ b/Sources/TokiUsageReaders/HermesUsageCostBreakdown.swift @@ -38,6 +38,37 @@ func hermesModelReportedCostDeltas( return deltas } +func hermesUnattributedReportedCostDelta( + currentReportedCost: Double?, + previousReportedCost: Double?, + modelReportedCostDeltas: [String: Double]?) -> Double? { + guard let currentReportedCost, + let previousReportedCost, + let modelReportedCostDeltas, + currentReportedCost.isFinite, + previousReportedCost.isFinite, + currentReportedCost >= previousReportedCost else { + return nil + } + + let reportedCostDelta = currentReportedCost - previousReportedCost + var attributedCostDelta = 0.0 + for cost in modelReportedCostDeltas.values { + guard cost.isFinite, + cost >= 0, + (attributedCostDelta + cost).isFinite else { + return nil + } + attributedCostDelta += cost + } + let tolerance = 0.000_000_001 * max(1, reportedCostDelta) + guard attributedCostDelta <= reportedCostDelta + tolerance else { + return nil + } + let residual = max(0, reportedCostDelta - attributedCostDelta) + return residual <= tolerance ? 0 : residual +} + private func hermesModelReportedCostsAreValid( _ modelReportedCosts: [String: Double]?, totalReportedCost: Double) -> Bool { @@ -55,7 +86,7 @@ private func hermesModelReportedCostsAreValid( combinedCost += cost } let tolerance = 0.000_000_001 * max(1, totalReportedCost) - return abs(combinedCost - totalReportedCost) <= tolerance + return combinedCost <= totalReportedCost + tolerance } func hermesModelPricedCost( diff --git a/Sources/TokiUsageReaders/HermesUsageLedger.swift b/Sources/TokiUsageReaders/HermesUsageLedger.swift index b686814..5705f4c 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedger.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedger.swift @@ -204,6 +204,7 @@ private extension HermesUsageLedger { timestamp: timestamp, observation: observation, previousModelCounters: previous.modelCounters, + previousReportedCost: previous.reportedCost, previousModelReportedCosts: previous.modelReportedCosts, previousModelPricingCounters: previous.modelPricingCounters, counters: delta, @@ -238,6 +239,7 @@ private extension HermesUsageLedger { timestamp: timestamp, observation: observation, previousModelCounters: observation.modelCounters.map { _ in [:] }, + previousReportedCost: observation.reportedCost.map { _ in 0 }, previousModelReportedCosts: observation.modelReportedCosts.map { _ in [:] }, previousModelPricingCounters: observation.modelPricingCounters.map { _ in [:] }, counters: observation.counters, @@ -495,6 +497,10 @@ private extension HermesUsageLedger { private func append( _ event: HermesUsageLedgerEvent, to events: inout [HermesUsageLedgerEvent]) { + guard event.counters.totalTokens > 0 else { + appendSingle(event, to: &events) + return + } let chunks = event.counters.chunks(maximum: hermesUsageLedgerMaximumEventTokenCount) var remainingCost = event.cost for (index, counters) in chunks.enumerated() { diff --git a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift index 55801a8..9526ca6 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift @@ -191,7 +191,7 @@ struct HermesUsageLedgerEvent: Codable, Equatable { hermesIdentifierIsValid(sessionIdentifier) && hermesDateIsValid(timestamp) && counters.isValid(maximum: hermesUsageLedgerMaximumEventTokenCount) - && counters.totalTokens > 0 + && (counters.totalTokens > 0 || cost > 0) && cost.isFinite && cost >= 0 && model?.utf8.count ?? 0 <= 512 diff --git a/Sources/TokiUsageReaders/HermesUsageResolver.swift b/Sources/TokiUsageReaders/HermesUsageResolver.swift index 05906fc..acdd889 100644 --- a/Sources/TokiUsageReaders/HermesUsageResolver.swift +++ b/Sources/TokiUsageReaders/HermesUsageResolver.swift @@ -42,10 +42,9 @@ private struct HermesSessionModelUsageAccumulator { var pricingCountersByModel: [String: HermesTokenCounters] = [:] var pricingTimestamp: Date? var models = Set() - var hasUnattributedReportedCost = false var resolvedReportedCostsByModel: [String: Double]? { - hasUnattributedReportedCost ? nil : reportedCostsByModel + reportedCostsByModel } var hasResolvedReportedCostBreakdown: Bool { @@ -86,7 +85,6 @@ private struct HermesSessionModelUsageAccumulator { reportedCost += usage.cost guard usage.cost > 0 else { return } guard let model = usage.model, usage.counters.totalTokens > 0 else { - hasUnattributedReportedCost = true return } let existingCost = reportedCostsByModel[model] ?? 0 @@ -161,7 +159,7 @@ enum HermesUsageResolver { ? pricingCounters(model: session.model, counters: session.counters) : [:] let sessionReportedCost = session.costIsDerivedFromModelPricing ? 0 : session.cost - if !hasModelUsage || session.cost > modelUsage.cost { + if !hasModelUsage { return HermesResolvedSessionCost( value: session.cost, isDerivedFromModelPricing: session.costIsDerivedFromModelPricing, @@ -170,6 +168,19 @@ enum HermesUsageResolver { modelPricingCounters: sessionPricingCounters, modelPricingTimestamp: session.modelPricingTimestamp) } + if session.cost > modelUsage.cost { + let modelReportedCosts = !session.costIsDerivedFromModelPricing + && modelUsage.hasResolvedReportedCostBreakdown + ? modelUsage.resolvedReportedCostsByModel + : nil + return HermesResolvedSessionCost( + value: session.cost, + isDerivedFromModelPricing: session.costIsDerivedFromModelPricing, + reportedValue: sessionReportedCost, + modelReportedCosts: modelReportedCosts, + modelPricingCounters: sessionPricingCounters, + modelPricingTimestamp: session.modelPricingTimestamp) + } if modelUsage.cost > session.cost { return HermesResolvedSessionCost( value: modelUsage.cost, diff --git a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift index bf7d5b9..cf7c3af 100644 --- a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift +++ b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift @@ -38,13 +38,16 @@ final class AgentSnapshotModelAttributionTests: XCTestCase { let snapshot = try await builder.build(configuration: fixture.configuration, now: now) let tokensByModel = snapshot.tokenEvents.reduce(into: [String: Int]()) { result, event in - result[event.model ?? "Mixed / Unattributed", default: 0] += event.totalTokens + result[ + event.model ?? UsageModelGrouping.mixedOrUnattributedLabel, + default: 0 + ] += event.totalTokens } XCTAssertEqual(tokensByModel, [ "gpt-5.6-sol": 100, "kr/claude-opus-5": 30, - "Mixed / Unattributed": 20, + UsageModelGrouping.mixedOrUnattributedLabel: 20, ]) } } diff --git a/Toki/Domain/Usage/UsageReportModelStats.swift b/Toki/Domain/Usage/UsageReportModelStats.swift index 86feec3..d4ba57a 100644 --- a/Toki/Domain/Usage/UsageReportModelStats.swift +++ b/Toki/Domain/Usage/UsageReportModelStats.swift @@ -2,8 +2,6 @@ import Foundation import TokiUsageCore import TokiUsageReaders -private let unattributedModelID = "Mixed / Unattributed" - private struct ModelSourceStatAggregate { var totalTokens = 0 var cost: Double = 0 @@ -75,9 +73,12 @@ extension UsageReportBuilder { let rowID = sourceCountByModel[key.modelID, default: 0] > 1 ? "\(key.modelID)|\(key.source)" : key.modelID + let displayModelID = key.modelID == UsageModelGrouping.mixedOrUnattributedKey + ? UsageModelGrouping.mixedOrUnattributedLabel + : key.modelID return ModelStat( id: rowID, - modelID: key.modelID, + modelID: displayModelID, totalTokens: aggregate.totalTokens, cost: aggregate.cost, activeSeconds: aggregate.activeSeconds, @@ -129,7 +130,8 @@ extension UsageReportBuilder { guard let source = event.source.trimmedNonEmpty else { continue } - let modelID = event.model?.trimmedNonEmpty ?? unattributedModelID + let modelID = event.model?.trimmedNonEmpty + ?? UsageModelGrouping.mixedOrUnattributedKey let key = ModelSourceUsageKey(modelID: modelID, source: source) aggregates[key, default: ModelSourceStatAggregate()].accumulate(event, modelID: modelID) activityEventsBySource[source, default: []].append( diff --git a/TokiTests/HermesPricingRefreshTests.swift b/TokiTests/HermesPricingRefreshTests.swift index e87b7b8..fdfefc5 100644 --- a/TokiTests/HermesPricingRefreshTests.swift +++ b/TokiTests/HermesPricingRefreshTests.swift @@ -535,7 +535,13 @@ extension HermesPricingRefreshTests { tokiTestISODate("2026-04-10T09:00:00Z"), ]) XCTAssertEqual(usage.activityEvents.count, 1) - XCTAssertNil(usage.activityEvents.first?.key) + XCTAssertEqual( + usage.activityEvents.first?.key, + UsageModelGrouping.mixedOrUnattributedKey) + XCTAssertEqual( + usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.activeSeconds ?? -1, + 30, + accuracy: 0.001) XCTAssertEqual(usage.activeSeconds, 30, accuracy: 0.001) } @@ -553,7 +559,7 @@ extension HermesPricingRefreshTests { id: "reported-mixed-session", startedAt: "2026-04-10T09:00:00Z", model: nil, - inputTokens: 200, + inputTokens: 250, outputTokens: 0, cacheReadTokens: 0, cacheWriteTokens: 0, @@ -561,7 +567,7 @@ extension HermesPricingRefreshTests { cwd: nil, gitRepoRoot: nil, estimatedCost: 0, - actualCost: 11)]) + actualCost: 12)]) try insertHermesModelUsage( databaseURL: dbURL, rows: [ @@ -601,10 +607,21 @@ extension HermesPricingRefreshTests { .readUsage( from: tokiTestISODate("2026-04-10T00:00:00Z"), to: tokiTestISODate("2026-04-10T10:30:00Z")) + XCTAssertEqual(initialUsage.inputTokens, 250) + XCTAssertEqual(initialUsage.cost, 12, accuracy: 0.000001) XCTAssertEqual(initialUsage.perModel[expensiveModel]?.cost ?? -1, 10, accuracy: 0.000001) XCTAssertEqual(initialUsage.perModel[cheapModel]?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual( + initialUsage.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.totalTokens, + 50) + XCTAssertEqual( + initialUsage.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.cost ?? -1, + 1, + accuracy: 0.000001) XCTAssertEqual(initialUsage.activityEvents.count, 1) - XCTAssertNil(initialUsage.activityEvents.first?.key) + XCTAssertEqual( + initialUsage.activityEvents.first?.key, + UsageModelGrouping.mixedOrUnattributedKey) try updateHermesModelUsage( databaseURL: dbURL, @@ -618,6 +635,12 @@ extension HermesPricingRefreshTests { task: "cheap", inputTokens: 200, actualCost: 2) + try updateHermesSessionReportedUsage( + databaseURL: dbURL, + id: "reported-mixed-session", + inputTokens: 450, + estimatedCost: 0, + actualCost: 24) let increment = try await HermesReader( dbPathOverride: dbURL.path, usageLedger: HermesUsageLedger(fileURL: ledgerURL), @@ -627,13 +650,26 @@ extension HermesPricingRefreshTests { to: tokiTestISODate("2026-04-10T12:00:00Z")) XCTAssertEqual(increment.inputTokens, 200) - XCTAssertEqual(increment.cost, 11, accuracy: 0.000001) + XCTAssertEqual(increment.cost, 12, accuracy: 0.000001) XCTAssertEqual(increment.perModel[expensiveModel]?.totalTokens, 100) XCTAssertEqual(increment.perModel[cheapModel]?.totalTokens, 100) XCTAssertEqual(increment.perModel[expensiveModel]?.cost ?? -1, 10, accuracy: 0.000001) XCTAssertEqual(increment.perModel[cheapModel]?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual( + increment.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.totalTokens, + 0) + XCTAssertEqual( + increment.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.cost ?? -1, + 1, + accuracy: 0.000001) + XCTAssertEqual( + increment.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.activeSeconds ?? -1, + 30, + accuracy: 0.001) XCTAssertEqual(increment.activityEvents.count, 1) - XCTAssertNil(increment.activityEvents.first?.key) + XCTAssertEqual( + increment.activityEvents.first?.key, + UsageModelGrouping.mixedOrUnattributedKey) } } diff --git a/TokiTests/HermesReaderTests.swift b/TokiTests/HermesReaderTests.swift index 9dd6fb8..5ec5782 100644 --- a/TokiTests/HermesReaderTests.swift +++ b/TokiTests/HermesReaderTests.swift @@ -450,6 +450,52 @@ final class HermesReaderTests: XCTestCase { XCTAssertFalse(hermesSQLiteShouldRetryImmutableFallback(after: SQLITE_BUSY)) } + func test_hermesReader_capturesObservationTimeAfterReadingSnapshot() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + + let startedAt = tokiTestISODate("2026-04-10T10:00:01Z") + let dbURL = tempDir.appendingPathComponent("state.db") + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "snapshot-race", + startedAt: "2026-04-10T10:00:01Z", + model: "gpt-5.5", + inputTokens: 123, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: nil)]) + try insertHermesMessage( + databaseURL: dbURL, + sessionID: "snapshot-race", + timestamp: startedAt) + let ledger = HermesUsageLedger( + fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T09:00:00Z")) + let nowSequence = HermesNowSequence([ + tokiTestISODate("2026-04-10T10:00:00Z"), + tokiTestISODate("2026-04-10T10:00:02Z"), + ]) + + let usage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { nowSequence.next() }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(usage.inputTokens, 123) + } + func test_hermesReader_readsUncheckpointedUsageFromLiveWAL() async throws { let fileManager = FileManager.default let tempDir = try makeHermesTemporaryDirectory() @@ -518,6 +564,20 @@ final class HermesReaderTests: XCTestCase { XCTAssertFalse(snapshot.isCurrent()) } + + func test_hermesDatabaseSourceSnapshotRejectsRollbackJournal() throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + + let dbURL = tempDir.appendingPathComponent("state.db") + try createHermesStateDB( + at: dbURL, + rows: [hermesSingleCounterFixture(id: "rollback-journal", inputTokens: 1)]) + try Data().write(to: URL(fileURLWithPath: "\(dbURL.path)-journal")) + + XCTAssertNil( + HermesDatabaseSourceSnapshot.captureForImmutableFallback(databaseURL: dbURL)) + } } extension HermesReaderTests { @@ -1339,3 +1399,22 @@ private func bindHermesDouble(_ value: Double?, at index: Int32, in statement: O } return sqlite3_bind_double(statement, index, value) == SQLITE_OK } + +private final class HermesNowSequence: @unchecked Sendable { + private let lock = NSLock() + private let values: [Date] + private var index = 0 + + init(_ values: [Date]) { + precondition(!values.isEmpty) + self.values = values + } + + func next() -> Date { + lock.lock() + defer { lock.unlock() } + let value = values[min(index, values.index(before: values.endIndex))] + index += 1 + return value + } +} diff --git a/TokiTests/UsageServiceBehaviorTests.swift b/TokiTests/UsageServiceBehaviorTests.swift index 9a8f8aa..ac76c41 100644 --- a/TokiTests/UsageServiceBehaviorTests.swift +++ b/TokiTests/UsageServiceBehaviorTests.swift @@ -248,12 +248,50 @@ extension UsageServiceBehaviorTests { await service.refresh() let usageData = await MainActor.run { service.usageData } - let unattributed = usageData.perModel.first { $0.modelID == "Mixed / Unattributed" } + let unattributed = usageData.perModel.first { + $0.modelID == UsageModelGrouping.mixedOrUnattributedLabel + } XCTAssertEqual(unattributed?.totalTokens, 120) XCTAssertEqual(unattributed?.sources, ["Hermes"]) XCTAssertEqual(unattributed?.isPriceKnown, false) } + + func test_usageService_surfacesMixedModelActivityInUnattributedBreakdown() async { + let recorder = MockReaderRecorder() + let timestamp = behaviorTestISODate("2026-04-10T09:00:00Z") + let reader = MockReader(name: "Hermes", recorder: recorder) { _, _ in + RawTokenUsage( + inputTokens: 200, + activeSeconds: 30, + perModel: [ + "model-a": PerModelUsage(totalTokens: 100, sources: ["Hermes"]), + "model-b": PerModelUsage(totalTokens: 100, sources: ["Hermes"]), + UsageModelGrouping.mixedOrUnattributedKey: PerModelUsage( + activeSeconds: 30, + sources: ["Hermes"]), + ], + activityEvents: [ + ActivityTimeEvent( + streamID: "mixed-session", + timestamp: timestamp, + key: UsageModelGrouping.mixedOrUnattributedKey), + ]) + } + + let service = await MainActor.run { UsageService(readers: [reader]) } + await MainActor.run { service.selectDay(behaviorTestISODate("2026-04-10T12:00:00Z")) } + await service.refresh() + + let usageData = await MainActor.run { service.usageData } + let unattributed = usageData.perModel.first { + $0.modelID == UsageModelGrouping.mixedOrUnattributedLabel + } + + XCTAssertEqual(unattributed?.totalTokens, 0) + XCTAssertEqual(unattributed?.activeSeconds, 30) + XCTAssertEqual(unattributed?.sources, ["Hermes"]) + } } final class UsageServicePeriodTotalsConcurrencyTests: XCTestCase { From 1522601472176e40909c38ed4abb3f5e222339c6 Mon Sep 17 00:00:00 2001 From: choi138 Date: Fri, 31 Jul 2026 12:33:09 +0900 Subject: [PATCH 04/10] fix(usage): preserve residual attribution details --- Sources/TokiUsageCore/RawTokenUsage.swift | 2 +- .../HermesUsageResolver.swift | 23 +++-- Toki/Domain/Usage/UsageData.swift | 6 ++ Toki/Domain/Usage/UsageReportBuilder.swift | 4 +- Toki/Domain/Usage/UsageReportExport.swift | 6 +- Toki/Domain/Usage/UsageReportModelStats.swift | 9 +- .../UsagePanel/PanelStatComponents.swift | 3 +- TokiTests/HermesPricingRefreshTests.swift | 95 +++++++++++++++++++ TokiTests/HermesReaderTests.swift | 55 +++++++++++ TokiTests/UsageExportTests.swift | 35 +++++++ TokiTests/UsageOriginAggregationTests.swift | 50 ++++++++++ .../UsageProjectTimelineReportTests.swift | 48 ++++++++++ TokiTests/UsageServiceBehaviorTests.swift | 6 +- 13 files changed, 320 insertions(+), 22 deletions(-) diff --git a/Sources/TokiUsageCore/RawTokenUsage.swift b/Sources/TokiUsageCore/RawTokenUsage.swift index 558b1cb..c1e1dfd 100644 --- a/Sources/TokiUsageCore/RawTokenUsage.swift +++ b/Sources/TokiUsageCore/RawTokenUsage.swift @@ -338,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) } } diff --git a/Sources/TokiUsageReaders/HermesUsageResolver.swift b/Sources/TokiUsageReaders/HermesUsageResolver.swift index acdd889..b444e47 100644 --- a/Sources/TokiUsageReaders/HermesUsageResolver.swift +++ b/Sources/TokiUsageReaders/HermesUsageResolver.swift @@ -127,10 +127,11 @@ enum HermesUsageResolver { if session.counters.totalTokens > 0, let model = session.model { models.insert(model) } + let hasTokenModelUsage = accumulatedModelUsage.counters.totalTokens > 0 let resolvedModel = models.count == 1 ? models.first : (models.isEmpty ? session.model : nil) let resolvedCost = resolveCost( session: session, - hasModelUsage: !modelUsage.isEmpty, + hasModelUsage: hasTokenModelUsage, modelUsage: accumulatedModelUsage) return HermesSessionObservation( @@ -140,7 +141,7 @@ enum HermesUsageResolver { latestActivityAt: session.latestActivityAt, model: resolvedModel, counters: session.counters.maximum(accumulatedModelUsage.counters), - modelCounters: modelUsage.isEmpty ? nil : accumulatedModelUsage.countersByModel, + modelCounters: hasTokenModelUsage ? accumulatedModelUsage.countersByModel : nil, cost: resolvedCost.value, costIsDerivedFromModelPricing: resolvedCost.isDerivedFromModelPricing, reportedCost: resolvedCost.reportedValue, @@ -169,15 +170,23 @@ enum HermesUsageResolver { modelPricingTimestamp: session.modelPricingTimestamp) } if session.cost > modelUsage.cost { - let modelReportedCosts = !session.costIsDerivedFromModelPricing - && modelUsage.hasResolvedReportedCostBreakdown - ? modelUsage.resolvedReportedCostsByModel - : nil + if !session.costIsDerivedFromModelPricing { + let retainedModelPricingCost = max( + 0, + modelUsage.cost - modelUsage.reportedCost) + return HermesResolvedSessionCost( + value: session.cost, + isDerivedFromModelPricing: false, + reportedValue: session.cost - retainedModelPricingCost, + modelReportedCosts: modelUsage.resolvedReportedCostsByModel, + modelPricingCounters: modelUsage.pricingCountersByModel, + modelPricingTimestamp: modelUsage.pricingTimestamp) + } return HermesResolvedSessionCost( value: session.cost, isDerivedFromModelPricing: session.costIsDerivedFromModelPricing, reportedValue: sessionReportedCost, - modelReportedCosts: modelReportedCosts, + modelReportedCosts: nil, modelPricingCounters: sessionPricingCounters, modelPricingTimestamp: session.modelPricingTimestamp) } diff --git a/Toki/Domain/Usage/UsageData.swift b/Toki/Domain/Usage/UsageData.swift index 9e900e7..bcc8468 100644 --- a/Toki/Domain/Usage/UsageData.swift +++ b/Toki/Domain/Usage/UsageData.swift @@ -87,6 +87,12 @@ struct ModelStat: Equatable { self.sources = sources self.isPriceKnown = isPriceKnown } + + var displayModelID: String { + modelID == UsageModelGrouping.mixedOrUnattributedKey + ? UsageModelGrouping.mixedOrUnattributedLabel + : modelID + } } struct SourceStat: Equatable { diff --git a/Toki/Domain/Usage/UsageReportBuilder.swift b/Toki/Domain/Usage/UsageReportBuilder.swift index e88e607..780e6b8 100644 --- a/Toki/Domain/Usage/UsageReportBuilder.swift +++ b/Toki/Domain/Usage/UsageReportBuilder.swift @@ -222,7 +222,7 @@ private extension UsageReportBuilder { let calendar = Calendar.autoupdatingCurrent let attributionBySession = bestAttributionsBySession(from: events, calendar: calendar) - for event in events where event.totalTokens > 0 { + for event in events where event.totalTokens > 0 || event.cost > 0 { let eventSessionKey = sessionGroupingKey(for: event, calendar: calendar) let sessionKey = SessionAggregateKey(source: event.source, sessionKey: eventSessionKey) let attribution = attributionBySession[sessionKey] ?? event.attribution @@ -259,7 +259,7 @@ private extension UsageReportBuilder { calendar: Calendar) -> [SessionUsageStat] { var aggregates: [SessionAggregateKey: SessionStatAggregate] = [:] - for event in events where event.totalTokens > 0 { + for event in events where event.totalTokens > 0 || event.cost > 0 { let attribution = event.attribution let sessionKey = sessionGroupingKey(for: event, calendar: calendar) let key = SessionAggregateKey(source: event.source, sessionKey: sessionKey) diff --git a/Toki/Domain/Usage/UsageReportExport.swift b/Toki/Domain/Usage/UsageReportExport.swift index ba1c903..93544d4 100644 --- a/Toki/Domain/Usage/UsageReportExport.swift +++ b/Toki/Domain/Usage/UsageReportExport.swift @@ -115,7 +115,7 @@ private struct UsageExportModel: Encodable { let isPriceKnown: Bool init(model: ModelStat) { - self.model = model.modelID + self.model = model.displayModelID totalTokens = model.totalTokens cost = model.cost activeSeconds = model.activeSeconds @@ -269,9 +269,9 @@ private extension UsageExport { usage.perModel.map { model in [ "model", - model.modelID, + model.displayModelID, model.sources.joined(separator: ";"), - model.modelID, + model.displayModelID, "", "", "", diff --git a/Toki/Domain/Usage/UsageReportModelStats.swift b/Toki/Domain/Usage/UsageReportModelStats.swift index d4ba57a..b479901 100644 --- a/Toki/Domain/Usage/UsageReportModelStats.swift +++ b/Toki/Domain/Usage/UsageReportModelStats.swift @@ -73,12 +73,9 @@ extension UsageReportBuilder { let rowID = sourceCountByModel[key.modelID, default: 0] > 1 ? "\(key.modelID)|\(key.source)" : key.modelID - let displayModelID = key.modelID == UsageModelGrouping.mixedOrUnattributedKey - ? UsageModelGrouping.mixedOrUnattributedLabel - : key.modelID return ModelStat( id: rowID, - modelID: displayModelID, + modelID: key.modelID, totalTokens: aggregate.totalTokens, cost: aggregate.cost, activeSeconds: aggregate.activeSeconds, @@ -204,8 +201,8 @@ extension UsageReportBuilder { if lhs.cost != rhs.cost { return lhs.cost > rhs.cost } - if lhs.modelID != rhs.modelID { - return lhs.modelID < rhs.modelID + if lhs.displayModelID != rhs.displayModelID { + return lhs.displayModelID < rhs.displayModelID } return lhs.sources.joined(separator: ",") < rhs.sources.joined(separator: ",") } diff --git a/Toki/Features/UsagePanel/PanelStatComponents.swift b/Toki/Features/UsagePanel/PanelStatComponents.swift index c86eca9..efd9896 100644 --- a/Toki/Features/UsagePanel/PanelStatComponents.swift +++ b/Toki/Features/UsagePanel/PanelStatComponents.swift @@ -76,7 +76,8 @@ struct ModelStatRowView: View, Equatable { } private var displayName: String { - let baseName = stat.modelID.hasPrefix("claude-") ? String(stat.modelID.dropFirst(7)) : stat.modelID + let modelID = stat.displayModelID + let baseName = modelID.hasPrefix("claude-") ? String(modelID.dropFirst(7)) : modelID guard !stat.sources.isEmpty else { return baseName } return "\(baseName) · \(sourceLabel)" } diff --git a/TokiTests/HermesPricingRefreshTests.swift b/TokiTests/HermesPricingRefreshTests.swift index fdfefc5..60a7620 100644 --- a/TokiTests/HermesPricingRefreshTests.swift +++ b/TokiTests/HermesPricingRefreshTests.swift @@ -256,6 +256,101 @@ extension HermesPricingRefreshTests { accuracy: 0.000001) } + // swiftlint:disable:next function_body_length + func test_hermesReader_preservesDerivedModelCostWhenReportedSessionTotalWins() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { + ModelPricingSupplement.install([:]) + try? FileManager.default.removeItem(at: tempDir) + } + let reportedModel = "reported-session-winner-model" + let derivedModel = "derived-session-winner-model" + let derivedPrice = uniformHermesModelPrice(perMillion: 10000) + ModelPricingSupplement.install([derivedModel: derivedPrice]) + + let dbURL = tempDir.appendingPathComponent("state.db") + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "session-total-wins", + startedAt: "2026-04-10T09:00:00Z", + model: nil, + inputTokens: 200, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: 12)]) + try insertHermesModelUsage( + databaseURL: dbURL, + rows: [ + HermesModelUsageFixture( + sessionID: "session-total-wins", + model: reportedModel, + task: "reported", + apiCallCount: 1, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 10), + HermesModelUsageFixture( + sessionID: "session-total-wins", + model: derivedModel, + task: "derived", + apiCallCount: 1, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 0), + ]) + let ledger = HermesUsageLedger(fileURL: ledgerURL) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T08:00:00Z")) + + let usage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T10:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + let expectedDerivedCost = derivedPrice.cost( + input: 100, + output: 0, + cacheRead: 0, + cacheWrite: 0) + let costOnlyEvents = usage.tokenEvents.filter { $0.totalTokens == 0 } + XCTAssertEqual(usage.inputTokens, 200) + XCTAssertEqual(usage.cost, 12, accuracy: 0.000001) + XCTAssertEqual(usage.perModel[reportedModel]?.cost ?? -1, 10, accuracy: 0.000001) + XCTAssertEqual( + usage.perModel[derivedModel]?.cost ?? -1, + expectedDerivedCost, + accuracy: 0.000001) + XCTAssertEqual( + usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.totalTokens, + 0) + XCTAssertEqual( + usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]?.cost ?? -1, + 1, + accuracy: 0.000001) + XCTAssertEqual(costOnlyEvents.count, 1) + XCTAssertEqual(costOnlyEvents.first?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual(usage.tokenEvents.reduce(0) { $0 + $1.cost }, 12, accuracy: 0.000001) + } + func test_hermesReader_preservesTokenDeltaWhenActualCostDropsBelowEstimate() async throws { let tempDir = try makeHermesTemporaryDirectory() defer { try? FileManager.default.removeItem(at: tempDir) } diff --git a/TokiTests/HermesReaderTests.swift b/TokiTests/HermesReaderTests.swift index 5ec5782..292c6ab 100644 --- a/TokiTests/HermesReaderTests.swift +++ b/TokiTests/HermesReaderTests.swift @@ -307,6 +307,61 @@ final class HermesReaderTests: XCTestCase { XCTAssertEqual(coverage.unmeteredMainAPICallCount, 3) } + func test_hermesReader_usesSessionModelWhenDetailRowsHaveNoTokens() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + + let model = "gpt-5.5" + let dbURL = tempDir.appendingPathComponent("state.db") + let ledger = HermesUsageLedger( + fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "zero-token-detail", + startedAt: "2026-04-10T09:00:00Z", + model: model, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: 1)]) + try insertHermesModelUsage( + databaseURL: dbURL, + rows: [HermesModelUsageFixture( + sessionID: "zero-token-detail", + model: model, + task: "", + apiCallCount: 1, + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 0)]) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T08:00:00Z")) + + let usage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T10:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(usage.inputTokens, 100) + XCTAssertEqual(usage.perModel[model]?.totalTokens, 100) + XCTAssertEqual(usage.perModel[model]?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertNil(usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]) + } + func test_hermesReader_tokenHelpersUseSameRangeAndZeroTokenFilters() async throws { let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) diff --git a/TokiTests/UsageExportTests.swift b/TokiTests/UsageExportTests.swift index 64bb69f..52a296a 100644 --- a/TokiTests/UsageExportTests.swift +++ b/TokiTests/UsageExportTests.swift @@ -1,3 +1,4 @@ +import TokiUsageCore import XCTest @testable import Toki @@ -119,4 +120,38 @@ final class UsageExportTests: XCTestCase { XCTAssertTrue(projects.isEmpty) XCTAssertTrue(sessions.isEmpty) } + + func test_exportsUseDisplayLabelForCanonicalUnattributedModelKey() throws { + let usage = UsageData( + date: tokiTestISODate("2026-04-10T00:00:00Z"), + endDate: tokiTestISODate("2026-04-11T00:00:00Z"), + inputTokens: 1, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cost: 0, + activeSeconds: 0, + perModel: [ + ModelStat( + id: UsageModelGrouping.mixedOrUnattributedKey, + totalTokens: 1, + cost: 0, + activeSeconds: 0, + sources: ["Hermes"], + isPriceKnown: false), + ]) + + let csv = UsageExport.csvString(for: usage) + let data = try XCTUnwrap(UsageExport.jsonString(for: usage).data(using: .utf8)) + let object = try XCTUnwrap( + JSONSerialization.jsonObject(with: data) as? [String: Any]) + let models = try XCTUnwrap(object["models"] as? [[String: Any]]) + + XCTAssertTrue(csv.contains(UsageModelGrouping.mixedOrUnattributedLabel)) + XCTAssertFalse(csv.contains(UsageModelGrouping.mixedOrUnattributedKey)) + XCTAssertEqual( + models.first?["model"] as? String, + UsageModelGrouping.mixedOrUnattributedLabel) + } } diff --git a/TokiTests/UsageOriginAggregationTests.swift b/TokiTests/UsageOriginAggregationTests.swift index 32eed32..770d3bf 100644 --- a/TokiTests/UsageOriginAggregationTests.swift +++ b/TokiTests/UsageOriginAggregationTests.swift @@ -108,6 +108,56 @@ final class UsageOriginAggregationTests: XCTestCase { XCTAssertEqual(Set(remoteRows.map(\.sources)), [["Codex"], ["Claude Code"]]) XCTAssertEqual(Set(remoteRows.map(\.totalTokens)), [10, 20]) } + + func test_unattributedModelKeyStaysCanonicalAcrossLocalAndRemoteAggregation() async throws { + let interval = testInterval + let groupingKey = UsageModelGrouping.mixedOrUnattributedKey + let localUsage = RawTokenUsage( + inputTokens: 10, + perModel: [ + groupingKey: PerModelUsage( + totalTokens: 10, + sources: ["Hermes"]), + ]) + let remoteUsage = RawTokenUsage( + inputTokens: 20, + perModel: [ + groupingKey: PerModelUsage( + totalTokens: 20, + sources: ["Hermes"]), + ]) + let remoteSlice = UsageOriginSlice( + origin: .remote( + deviceID: "remote-a", + name: "worker", + platform: "linux", + lastUpdatedAt: interval.start), + usage: remoteUsage, + sourceStats: [ + SourceStat( + source: "Hermes", + inputTokens: 20, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cost: 0, + activeSeconds: 0), + ]) + let aggregator = UsageAggregator(readers: [ + FixedUsageReader(name: "Hermes", usage: localUsage), + FixedOriginReader(name: "Remote Devices", slices: [remoteSlice]), + ]) + + let result = await aggregator.aggregateUsage(for: makeRequest(interval: interval)) + let rows = result.usageData.perModel.filter { $0.modelID == groupingKey } + let row = try XCTUnwrap(rows.first) + + XCTAssertEqual(rows.count, 1) + XCTAssertEqual(row.totalTokens, 30) + XCTAssertEqual(row.displayModelID, UsageModelGrouping.mixedOrUnattributedLabel) + XCTAssertEqual(row.sources, ["Hermes"]) + } } extension UsageOriginAggregationTests { diff --git a/TokiTests/UsageProjectTimelineReportTests.swift b/TokiTests/UsageProjectTimelineReportTests.swift index 111da7b..701b0a7 100644 --- a/TokiTests/UsageProjectTimelineReportTests.swift +++ b/TokiTests/UsageProjectTimelineReportTests.swift @@ -48,6 +48,54 @@ final class UsageProjectTimelineReportTests: XCTestCase { XCTAssertEqual(report.attributedCost, 0.30, accuracy: 0.000001) } + func test_usageReportIncludesCostOnlyEventsInEveryAttributedBreakdown() { + let startDate = tokiTestISODate("2026-04-10T00:00:00Z") + let endDate = tokiTestISODate("2026-04-11T00:00:00Z") + let attribution = UsageAttribution( + projectPath: "/Users/example/Toki", + sessionID: "cost-only-session", + quality: .exact) + var rawUsage = RawTokenUsage( + inputTokens: 200, + cost: 12) + rawUsage.recordTokenEvent( + timestamp: tokiTestISODate("2026-04-10T10:00:00Z"), + source: "Hermes", + model: "reported-model", + inputTokens: 100, + outputTokens: 0, + cost: 10, + attribution: attribution) + rawUsage.recordTokenEvent( + timestamp: tokiTestISODate("2026-04-10T10:00:00Z"), + source: "Hermes", + model: "derived-model", + inputTokens: 100, + outputTokens: 0, + cost: 1, + attribution: attribution) + rawUsage.recordTokenEvent( + timestamp: tokiTestISODate("2026-04-10T10:00:00Z"), + source: "Hermes", + model: nil, + inputTokens: 0, + outputTokens: 0, + cost: 1, + attribution: attribution) + + let report = UsageReportBuilder.report( + from: rawUsage, + date: startDate, + endDate: endDate, + sourceStats: []) + + XCTAssertEqual(rawUsage.tokenEvents.count, 3) + XCTAssertEqual(report.timeBuckets.reduce(0) { $0 + $1.cost }, 12, accuracy: 0.000001) + XCTAssertEqual(report.projectStats.first?.cost ?? -1, 12, accuracy: 0.000001) + XCTAssertEqual(report.sessionStats.first?.cost ?? -1, 12, accuracy: 0.000001) + XCTAssertEqual(report.attributedCost, 12, accuracy: 0.000001) + } + func test_usageReportUsesUniqueIDsForPathlessProjectStats() { let startDate = tokiTestISODate("2026-04-10T00:00:00Z") let endDate = tokiTestISODate("2026-04-11T00:00:00Z") diff --git a/TokiTests/UsageServiceBehaviorTests.swift b/TokiTests/UsageServiceBehaviorTests.swift index ac76c41..b87f6de 100644 --- a/TokiTests/UsageServiceBehaviorTests.swift +++ b/TokiTests/UsageServiceBehaviorTests.swift @@ -249,10 +249,11 @@ extension UsageServiceBehaviorTests { let usageData = await MainActor.run { service.usageData } let unattributed = usageData.perModel.first { - $0.modelID == UsageModelGrouping.mixedOrUnattributedLabel + $0.modelID == UsageModelGrouping.mixedOrUnattributedKey } XCTAssertEqual(unattributed?.totalTokens, 120) + XCTAssertEqual(unattributed?.displayModelID, UsageModelGrouping.mixedOrUnattributedLabel) XCTAssertEqual(unattributed?.sources, ["Hermes"]) XCTAssertEqual(unattributed?.isPriceKnown, false) } @@ -285,10 +286,11 @@ extension UsageServiceBehaviorTests { let usageData = await MainActor.run { service.usageData } let unattributed = usageData.perModel.first { - $0.modelID == UsageModelGrouping.mixedOrUnattributedLabel + $0.modelID == UsageModelGrouping.mixedOrUnattributedKey } XCTAssertEqual(unattributed?.totalTokens, 0) + XCTAssertEqual(unattributed?.displayModelID, UsageModelGrouping.mixedOrUnattributedLabel) XCTAssertEqual(unattributed?.activeSeconds, 30) XCTAssertEqual(unattributed?.sources, ["Hermes"]) } From 26e9966307c7bfb537f1244ceb03de9f3429f14b Mon Sep 17 00:00:00 2001 From: choi138 Date: Fri, 31 Jul 2026 13:36:37 +0900 Subject: [PATCH 05/10] fix(hermes): preserve residual usage details --- .../TokiAgentCore/AgentSnapshotBuilder.swift | 1 + .../HermesUsageAttribution.swift | 31 ++- .../TokiUsageReaders/HermesUsageLedger.swift | 23 +- .../HermesUsageLedgerSupport.swift | 16 +- .../AgentSnapshotModelAttributionTests.swift | 24 ++ Toki/Domain/Usage/UsageReportModelStats.swift | 19 +- .../RemoteSync/RemoteUsageMapper.swift | 12 +- TokiTests/HermesUsageLedgerTests.swift | 245 ++++++++++++++++++ TokiTests/RemoteUsageReaderTests.swift | 32 +++ TokiTests/UsageServiceBehaviorTests.swift | 40 +++ 10 files changed, 421 insertions(+), 22 deletions(-) diff --git a/Sources/TokiAgentCore/AgentSnapshotBuilder.swift b/Sources/TokiAgentCore/AgentSnapshotBuilder.swift index f2e3e24..cc9847f 100644 --- a/Sources/TokiAgentCore/AgentSnapshotBuilder.swift +++ b/Sources/TokiAgentCore/AgentSnapshotBuilder.swift @@ -424,6 +424,7 @@ private extension AgentSnapshotBuilder { private func remoteModel(_ model: String?) -> String? { guard let model, + model != UsageModelGrouping.mixedOrUnattributedKey, TokiSyncValidation.isSafeDisplayText( model, maximumLength: RemoteUsageSnapshotValidator.maximumModelLength) else { diff --git a/Sources/TokiUsageReaders/HermesUsageAttribution.swift b/Sources/TokiUsageReaders/HermesUsageAttribution.swift index 4e392f7..eeece9b 100644 --- a/Sources/TokiUsageReaders/HermesUsageAttribution.swift +++ b/Sources/TokiUsageReaders/HermesUsageAttribution.swift @@ -32,7 +32,9 @@ func hermesUsageEvents( observation: observation, previousModelCounters: previousModelCounters, counters: counters, - includeCostOnlyResidual: unattributedReportedCostDelta.map { $0 > 0 } ?? false) + costOnlyModels: Set(modelReportedCostDeltas?.keys.map { $0 } ?? []), + includeCostOnlyResidual: unattributedReportedCostDelta.map { $0 > 0 } ?? false, + includeCostOnlyFallback: cost > 0) let modelPricingDeltas = hermesModelCounterDeltas( current: observation.modelPricingCounters, previous: previousModelPricingCounters, @@ -63,11 +65,22 @@ private func hermesUsageEventParts( observation: HermesSessionObservation, previousModelCounters: [String: HermesTokenCounters]?, counters: HermesTokenCounters, - includeCostOnlyResidual: Bool) -> [HermesUsageEventPart] { + costOnlyModels: Set, + includeCostOnlyResidual: Bool, + includeCostOnlyFallback: Bool) -> [HermesUsageEventPart] { guard let modelDeltas = hermesModelCounterDeltas( current: observation.modelCounters, previous: previousModelCounters, maximumDelta: counters) else { + if counters.totalTokens == 0, !costOnlyModels.isEmpty { + var parts = costOnlyModels.sorted().map { + HermesUsageEventPart(model: $0, counters: .zero) + } + if includeCostOnlyResidual { + parts.append(HermesUsageEventPart(model: nil, counters: .zero)) + } + return parts + } return [HermesUsageEventPart(model: observation.model, counters: counters)] } @@ -78,6 +91,10 @@ private func hermesUsageEventParts( combinedCounters = combinedCounters.adding(modelCounters) parts.append(HermesUsageEventPart(model: model, counters: modelCounters)) } + let representedModels = Set(parts.compactMap(\.model)) + for model in costOnlyModels.subtracting(representedModels).sorted() { + parts.append(HermesUsageEventPart(model: model, counters: .zero)) + } let residual = counters.subtracting(combinedCounters) if residual.totalTokens > 0 { @@ -85,6 +102,9 @@ private func hermesUsageEventParts( } else if includeCostOnlyResidual { parts.append(HermesUsageEventPart(model: nil, counters: .zero)) } + if parts.isEmpty, includeCostOnlyFallback { + parts.append(HermesUsageEventPart(model: observation.model, counters: .zero)) + } return parts } @@ -121,6 +141,13 @@ private func hermesAllocatedUsageCosts( var remainderTokenCounts = basis.remainderTokenCounts if !remainderTokenCounts.contains(where: { $0 > 0 }) { remainderTokenCounts = parts.map(\.counters.totalTokens) + if !remainderTokenCounts.contains(where: { $0 > 0 }) { + let fallbackRecipient = parts.lastIndex(where: { $0.model == nil }) + ?? (parts.count == 1 ? parts.startIndex : nil) + if let fallbackRecipient { + remainderTokenCounts[fallbackRecipient] = 1 + } + } } let recipientIndices = remainderTokenCounts.indices.filter { remainderTokenCounts[$0] > 0 diff --git a/Sources/TokiUsageReaders/HermesUsageLedger.swift b/Sources/TokiUsageReaders/HermesUsageLedger.swift index 5705f4c..18cac65 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedger.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedger.swift @@ -175,11 +175,6 @@ private extension HermesUsageLedger { } let delta = observation.counters.subtracting(previous.counters) - guard delta.totalTokens > 0 else { - guard currentBaseline.metadataDiffers(from: previous) else { return false } - candidate.baselines[identifier] = currentBaseline - return true - } let timestamp = incrementalTimestamp( observation: observation, previous: previous, @@ -190,6 +185,9 @@ private extension HermesUsageLedger { previous: previous, delta: delta, pricingTimestamp: pricingTimestamp) else { + guard delta.totalTokens > 0 || currentBaseline.metadataDiffers(from: previous) else { + return false + } candidate.baselines[identifier] = currentBaseline try addUnattributed( identifier: identifier, @@ -199,11 +197,20 @@ private extension HermesUsageLedger { to: &candidate.unattributed) return true } + guard delta.totalTokens > 0 || cost > 0 else { + guard currentBaseline.metadataDiffers(from: previous) else { return false } + candidate.baselines[identifier] = currentBaseline + return true + } + let previousModelCounters = previous.modelCounters + ?? hermesCompleteModelCounters( + previous.modelPricingCounters, + matching: previous.counters) for event in hermesUsageEvents( identifier: identifier, timestamp: timestamp, observation: observation, - previousModelCounters: previous.modelCounters, + previousModelCounters: previousModelCounters, previousReportedCost: previous.reportedCost, previousModelReportedCosts: previous.modelReportedCosts, previousModelPricingCounters: previous.modelPricingCounters, @@ -225,7 +232,7 @@ private extension HermesUsageLedger { previousSuccessfulObservationAt: Date?, observedAt: Date, to candidate: inout HermesUsageLedgerDocument) throws -> Bool { - guard observation.counters.totalTokens > 0 else { return false } + guard observation.counters.totalTokens > 0 || observation.cost > 0 else { return false } candidate.baselines[identifier] = currentBaseline if initialUsageIsDated( observation, @@ -538,7 +545,7 @@ private extension HermesUsageLedger { cost: Double, observedAt: Date, to carryovers: inout [String: HermesUsageLedgerCarryover]) throws { - guard counters.totalTokens > 0 else { return } + guard counters.totalTokens > 0 || cost > 0 else { return } guard let existing = carryovers[identifier] else { carryovers[identifier] = HermesUsageLedgerCarryover( counters: counters, diff --git a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift index 9526ca6..0041cb3 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedgerSupport.swift @@ -129,6 +129,20 @@ func hermesModelPricingCountersAreValid( return !totalCounters.hasDecrease(comparedTo: combinedCounters) } +func hermesCompleteModelCounters( + _ modelCounters: [String: HermesTokenCounters]?, + matching totalCounters: HermesTokenCounters) -> [String: HermesTokenCounters]? { + guard let modelCounters, + hermesModelPricingCountersAreValid(modelCounters, within: totalCounters) else { + return nil + } + + let combinedCounters = modelCounters.values.reduce(HermesTokenCounters.zero) { + $0.adding($1) + } + return combinedCounters == totalCounters ? modelCounters : nil +} + func hermesModelCounterDeltas( current: [String: HermesTokenCounters]?, previous: [String: HermesTokenCounters]?, @@ -453,7 +467,7 @@ struct HermesUsageLedgerCarryover: Codable, Equatable { var isValid: Bool { counters.isValid() - && counters.totalTokens > 0 + && (counters.totalTokens > 0 || cost > 0) && cost.isFinite && cost >= 0 && hermesDateIsValid(firstObservedAt) diff --git a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift index cf7c3af..8a61b7d 100644 --- a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift +++ b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift @@ -5,6 +5,30 @@ import XCTest @testable import TokiUsageReaders final class AgentSnapshotModelAttributionTests: XCTestCase { + func test_snapshotEncodesCanonicalMixedActivityAsUnattributedRemoteModel() async throws { + let now = Date(timeIntervalSince1970: 1_784_200_000) + let eventDate = now.addingTimeInterval(-60) + let usage = RawTokenUsage(activityEvents: [ + ActivityTimeEvent( + streamID: "mixed-model-session", + timestamp: eventDate, + key: UsageModelGrouping.mixedOrUnattributedKey), + ]) + let descriptor = LocalUsageReaderDescriptor( + reader: FixedTokenReader(name: "Hermes", usage: usage), + sourceLocations: []) + let fixture = try AgentSnapshotFixture() + defer { fixture.remove() } + let builder = AgentSnapshotBuilder( + home: fixture.root, + readerDescriptors: [descriptor]) + + let snapshot = try await builder.build(configuration: fixture.configuration, now: now) + + XCTAssertEqual(snapshot.activityEvents.count, 1) + XCTAssertNil(snapshot.activityEvents.first?.model) + } + func test_snapshotPreservesMixedHermesModelAttributionAndResidual() async throws { let now = Date(timeIntervalSince1970: 1_784_200_000) let eventDate = now.addingTimeInterval(-60) diff --git a/Toki/Domain/Usage/UsageReportModelStats.swift b/Toki/Domain/Usage/UsageReportModelStats.swift index b479901..8f9b256 100644 --- a/Toki/Domain/Usage/UsageReportModelStats.swift +++ b/Toki/Domain/Usage/UsageReportModelStats.swift @@ -25,7 +25,10 @@ private struct ModelSourceStatAggregate { totalTokens += event.totalTokens cost += event.cost sources.insert(event.source) - isPriceKnown = isPriceKnown && modelPriceLookup(for: modelID, at: event.timestamp).isPriced + if event.totalTokens > 0 { + isPriceKnown = isPriceKnown + && (event.cost > 0 || modelPriceLookup(for: modelID, at: event.timestamp).isPriced) + } } var hasReportableData: Bool { @@ -123,7 +126,7 @@ extension UsageReportBuilder { var aggregates: [ModelSourceUsageKey: ModelSourceStatAggregate] = [:] var activityEventsBySource: [String: [ActivityTimeEvent]] = [:] - for event in events where event.totalTokens > 0 { + for event in events where event.totalTokens > 0 || event.cost > 0 { guard let source = event.source.trimmedNonEmpty else { continue } @@ -131,11 +134,13 @@ extension UsageReportBuilder { ?? UsageModelGrouping.mixedOrUnattributedKey let key = ModelSourceUsageKey(modelID: modelID, source: source) aggregates[key, default: ModelSourceStatAggregate()].accumulate(event, modelID: modelID) - activityEventsBySource[source, default: []].append( - ActivityTimeEvent( - streamID: modelSourceStreamID(for: event, calendar: calendar), - timestamp: event.timestamp, - key: modelID)) + if event.totalTokens > 0 { + activityEventsBySource[source, default: []].append( + ActivityTimeEvent( + streamID: modelSourceStreamID(for: event, calendar: calendar), + timestamp: event.timestamp, + key: modelID)) + } } for (source, activityEvents) in activityEventsBySource { diff --git a/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift b/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift index e7741f4..bff5962 100644 --- a/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift +++ b/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift @@ -131,7 +131,7 @@ struct RemoteUsageMapper { eventsBySource[event.source, default: []].append(ActivityTimeEvent( streamID: "\(snapshot.device.id):\(event.streamID)", timestamp: event.timestamp, - key: normalizedModelID(event.model), + key: activityModelID(event.model), agentKind: event.agentKind == .subagent ? .subagent : .main)) } return eventsBySource @@ -146,7 +146,7 @@ struct RemoteUsageMapper { return ActivityTimeEvent( streamID: "\(snapshot.device.id):\(event.streamID)", timestamp: event.timestamp, - key: normalizedModelID(event.model), + key: activityModelID(event.model), agentKind: event.agentKind == .subagent ? .subagent : .main) } } @@ -157,8 +157,8 @@ struct RemoteUsageMapper { endDate: Date, usage: inout RawTokenUsage) { for event in snapshot.activityEvents where event.timestamp >= startDate && event.timestamp < endDate { - guard let model = normalizedModelID(event.model), - usage.perModel[model] != nil else { + let model = activityModelID(event.model) + guard usage.perModel[model] != nil else { continue } usage.perModel[model]?.sources.insert( @@ -166,6 +166,10 @@ struct RemoteUsageMapper { } } + private func activityModelID(_ model: String?) -> String { + normalizedModelID(model) ?? UsageModelGrouping.mixedOrUnattributedKey + } + private func deviceSource(_ source: String, deviceName: String) -> String { "\(source) · \(deviceName)" } diff --git a/TokiTests/HermesUsageLedgerTests.swift b/TokiTests/HermesUsageLedgerTests.swift index fe3d65d..e6eb5d4 100644 --- a/TokiTests/HermesUsageLedgerTests.swift +++ b/TokiTests/HermesUsageLedgerTests.swift @@ -570,6 +570,164 @@ extension HermesUsageLedgerTests { } } +extension HermesUsageLedgerTests { + func test_hermesUsageLedger_usesCompleteLegacyV3PricingCountersAsModelBaseline() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + let startedAt = tokiTestISODate("2026-04-09T08:00:00Z") + let initialModelCounters = [ + "legacy-model-a": hermesInputCounters(60), + "legacy-model-b": hermesInputCounters(40), + ] + let ledger = HermesUsageLedger(fileURL: ledgerURL) + try await ledger.refresh( + observations: [hermesLegacyV3Observation( + startedAt: startedAt, + modelCounters: initialModelCounters, + modelPricingCounters: initialModelCounters)], + observedAt: tokiTestISODate("2026-04-09T09:00:00Z")) + try removeHermesModelCountersFromBaselines(at: ledgerURL) + + let currentModelCounters = [ + "legacy-model-a": hermesInputCounters(80), + "legacy-model-b": hermesInputCounters(50), + ] + let restartedLedger = HermesUsageLedger(fileURL: ledgerURL) + try await restartedLedger.refresh( + observations: [hermesLegacyV3Observation( + startedAt: startedAt, + modelCounters: currentModelCounters, + modelPricingCounters: currentModelCounters)], + observedAt: tokiTestISODate("2026-04-10T10:00:00Z")) + let events = try await restartedLedger.events( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual( + Dictionary(uniqueKeysWithValues: events.map { ($0.model ?? "Mixed", $0.counters.inputTokens) }), + [ + "legacy-model-a": 20, + "legacy-model-b": 10, + ]) + } + + func test_hermesUsageLedger_keepsLegacyV3DeltaMixedWhenPricingCountersArePartial() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + let startedAt = tokiTestISODate("2026-04-09T08:00:00Z") + let ledger = HermesUsageLedger(fileURL: ledgerURL) + try await ledger.refresh( + observations: [hermesLegacyV3Observation( + startedAt: startedAt, + modelCounters: [ + "legacy-model-a": hermesInputCounters(60), + "legacy-model-b": hermesInputCounters(40), + ], + modelPricingCounters: [ + "legacy-model-a": hermesInputCounters(60), + ])], + observedAt: tokiTestISODate("2026-04-09T09:00:00Z")) + try removeHermesModelCountersFromBaselines(at: ledgerURL) + + let restartedLedger = HermesUsageLedger(fileURL: ledgerURL) + try await restartedLedger.refresh( + observations: [hermesLegacyV3Observation( + startedAt: startedAt, + modelCounters: [ + "legacy-model-a": hermesInputCounters(80), + "legacy-model-b": hermesInputCounters(50), + ], + modelPricingCounters: [ + "legacy-model-a": hermesInputCounters(80), + ])], + observedAt: tokiTestISODate("2026-04-10T10:00:00Z")) + let events = try await restartedLedger.events( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(events.count, 1) + XCTAssertNil(events.first?.model) + XCTAssertEqual(events.first?.counters.inputTokens, 30) + } + + func test_hermesUsageLedger_recordsDatedInitialCostOnlyObservation() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledger = HermesUsageLedger( + fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + let baselineAt = tokiTestISODate("2026-04-10T08:00:00Z") + let startedAt = tokiTestISODate("2026-04-10T09:00:00Z") + try await ledger.refresh(observations: [], observedAt: baselineAt) + + try await ledger.refresh( + observations: [hermesCostOnlyObservation( + sessionID: "dated-cost-only", + startedAt: startedAt, + activityAt: startedAt, + cost: 1.25)], + observedAt: tokiTestISODate("2026-04-10T10:00:00Z")) + let events = try await ledger.events( + from: baselineAt, + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(events.count, 1) + XCTAssertEqual(events.first?.counters, .zero) + XCTAssertEqual(events.first?.cost ?? -1, 1.25, accuracy: 0.000_001) + } + + func test_hermesUsageLedger_persistsUndatedInitialCostOnlyCarryover() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + let ledger = HermesUsageLedger(fileURL: ledgerURL) + + try await ledger.refresh( + observations: [hermesCostOnlyObservation( + sessionID: "undated-cost-only", + startedAt: tokiTestISODate("2026-04-09T08:00:00Z"), + cost: 0.75)], + observedAt: tokiTestISODate("2026-04-10T10:00:00Z")) + let restartedLedger = HermesUsageLedger(fileURL: ledgerURL) + let status = try await restartedLedger.status() + let events = try await restartedLedger.events( + from: tokiTestISODate("2026-04-09T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(status.unattributedSessionCount, 1) + XCTAssertEqual(status.unattributedTokens, 0) + XCTAssertTrue(events.isEmpty) + } + + func test_hermesUsageLedger_recordsIncrementalCostOnlyModelUpdate() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledger = HermesUsageLedger( + fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + let startedAt = tokiTestISODate("2026-04-09T08:00:00Z") + try await ledger.refresh( + observations: [hermesReportedModelObservation( + startedAt: startedAt, + cost: 1)], + observedAt: tokiTestISODate("2026-04-09T09:00:00Z")) + + try await ledger.refresh( + observations: [hermesReportedModelObservation( + startedAt: startedAt, + cost: 2)], + observedAt: tokiTestISODate("2026-04-10T10:00:00Z")) + let events = try await ledger.events( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(events.count, 1) + XCTAssertEqual(events.first?.model, "reported-cost-model") + XCTAssertEqual(events.first?.counters, .zero) + XCTAssertEqual(events.first?.cost ?? -1, 1, accuracy: 0.000_001) + } +} + private struct HermesUsageLedgerV1Fixture: Encodable { let schemaVersion: Int let identifierKey: String @@ -663,3 +821,90 @@ private func assertHermesLedgerPrivacy(at ledgerURL: URL) throws { FileManager.default.attributesOfItem(atPath: ledgerURL.path)[.posixPermissions] as? NSNumber) XCTAssertEqual(ledgerPermissions.intValue & 0o077, 0) } + +private func hermesInputCounters(_ inputTokens: Int) -> HermesTokenCounters { + HermesTokenCounters( + inputTokens: inputTokens, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0) +} + +private func hermesLegacyV3Observation( + startedAt: Date, + modelCounters: [String: HermesTokenCounters], + modelPricingCounters: [String: HermesTokenCounters]) -> HermesSessionObservation { + let counters = modelCounters.values.reduce(HermesTokenCounters.zero) { + $0.adding($1) + } + return HermesSessionObservation( + sessionID: "legacy-v3-model-session", + startedAt: startedAt, + earliestActivityAt: nil, + latestActivityAt: nil, + model: nil, + counters: counters, + modelCounters: modelCounters, + cost: 0, + costIsDerivedFromModelPricing: true, + modelPricingCounters: modelPricingCounters, + projectName: nil, + attributionQuality: .unknown) +} + +private func removeHermesModelCountersFromBaselines(at ledgerURL: URL) throws { + guard var document = try JSONSerialization.jsonObject( + with: Data(contentsOf: ledgerURL)) as? [String: Any], + var baselines = document["baselines"] as? [String: Any], + !baselines.isEmpty else { + throw NSError(domain: "HermesUsageLedgerTests", code: 1) + } + for (identifier, value) in baselines { + guard var baseline = value as? [String: Any] else { + throw NSError(domain: "HermesUsageLedgerTests", code: 2) + } + baseline.removeValue(forKey: "modelCounters") + baselines[identifier] = baseline + } + document["baselines"] = baselines + let data = try JSONSerialization.data(withJSONObject: document, options: [.sortedKeys]) + try writePrivateHermesTestData(data, to: ledgerURL) +} + +private func hermesCostOnlyObservation( + sessionID: String, + startedAt: Date, + activityAt: Date? = nil, + cost: Double) -> HermesSessionObservation { + HermesSessionObservation( + sessionID: sessionID, + startedAt: startedAt, + earliestActivityAt: activityAt, + latestActivityAt: activityAt, + model: "cost-only-model", + counters: .zero, + cost: cost, + projectName: nil, + attributionQuality: .unknown) +} + +private func hermesReportedModelObservation( + startedAt: Date, + cost: Double) -> HermesSessionObservation { + let counters = hermesInputCounters(100) + return HermesSessionObservation( + sessionID: "reported-cost-session", + startedAt: startedAt, + earliestActivityAt: nil, + latestActivityAt: nil, + model: "reported-cost-model", + counters: counters, + modelCounters: ["reported-cost-model": counters], + cost: cost, + reportedCost: cost, + modelReportedCosts: ["reported-cost-model": cost], + modelPricingCounters: [:], + projectName: nil, + attributionQuality: .unknown) +} diff --git a/TokiTests/RemoteUsageReaderTests.swift b/TokiTests/RemoteUsageReaderTests.swift index 0e891ec..594b1a6 100644 --- a/TokiTests/RemoteUsageReaderTests.swift +++ b/TokiTests/RemoteUsageReaderTests.swift @@ -131,6 +131,38 @@ final class RemoteUsageReaderTests: XCTestCase { XCTAssertEqual(grouped["Claude"]?.count, 1) XCTAssertEqual(grouped.values.flatMap { $0 }.count, 3) } + + func test_unattributedRemoteActivityRestoresCanonicalMixedModelTime() throws { + let fixture = try makeFixture() + let original = try SnapshotCipher.open(fixture.envelope, key: fixture.encryptionKey) + let snapshot = RemoteUsageSnapshot( + device: original.device, + generatedAt: original.generatedAt, + coveredFrom: original.coveredFrom, + coveredTo: original.coveredTo, + tokenEvents: [], + activityEvents: [ + RemoteActivityEvent( + timestamp: fixture.start.addingTimeInterval(60), + source: "Hermes", + model: nil, + streamID: "mixed-model-session", + agentKind: .main), + ]) + + let slice = try XCTUnwrap(RemoteUsageMapper().usageSlice( + from: snapshot, + startDate: fixture.start, + endDate: fixture.end)) + let mixedUsage = try XCTUnwrap( + slice.usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]) + + XCTAssertEqual( + slice.usage.activityEvents.first?.key, + UsageModelGrouping.mixedOrUnattributedKey) + XCTAssertEqual(mixedUsage.activeSeconds, 30) + XCTAssertEqual(mixedUsage.sources, ["Hermes · build-server"]) + } } extension RemoteUsageReaderTests { diff --git a/TokiTests/UsageServiceBehaviorTests.swift b/TokiTests/UsageServiceBehaviorTests.swift index b87f6de..d1002b1 100644 --- a/TokiTests/UsageServiceBehaviorTests.swift +++ b/TokiTests/UsageServiceBehaviorTests.swift @@ -258,6 +258,46 @@ extension UsageServiceBehaviorTests { XCTAssertEqual(unattributed?.isPriceKnown, false) } + func test_usageService_exportsReportedCostOnlyUnattributedEventAsKnownPrice() async throws { + let recorder = MockReaderRecorder() + let reader = MockReader(name: "Hermes", recorder: recorder) { _, _ in + var usage = RawTokenUsage( + cost: 1, + perModel: [ + UsageModelGrouping.mixedOrUnattributedKey: PerModelUsage( + cost: 1, + sources: ["Hermes"]), + ]) + usage.recordTokenEvent( + timestamp: behaviorTestISODate("2026-04-10T09:00:00Z"), + source: "Hermes", + model: nil, + inputTokens: 0, + outputTokens: 0, + cost: 1) + return usage + } + + let service = await MainActor.run { UsageService(readers: [reader]) } + await MainActor.run { service.selectDay(behaviorTestISODate("2026-04-10T12:00:00Z")) } + await service.refresh() + + let usageData = await MainActor.run { service.usageData } + let unattributed = try XCTUnwrap(usageData.perModel.first { + $0.modelID == UsageModelGrouping.mixedOrUnattributedKey + }) + let modelCSVRow = try XCTUnwrap( + UsageExport.csvString(for: usageData) + .split(separator: "\n") + .first { $0.hasPrefix("model,") }) + + XCTAssertEqual(unattributed.totalTokens, 0) + XCTAssertEqual(unattributed.cost, 1, accuracy: 0.000_001) + XCTAssertEqual(unattributed.activeSeconds, 0) + XCTAssertTrue(unattributed.isPriceKnown) + XCTAssertTrue(modelCSVRow.contains(",1.000000,")) + } + func test_usageService_surfacesMixedModelActivityInUnattributedBreakdown() async { let recorder = MockReaderRecorder() let timestamp = behaviorTestISODate("2026-04-10T09:00:00Z") From 0f82e2ecc5b99b716cfe7791789ee51f2fc1a1f4 Mon Sep 17 00:00:00 2001 From: choi138 Date: Fri, 31 Jul 2026 14:46:04 +0900 Subject: [PATCH 06/10] fix(hermes): preserve cost-only usage and current pricing --- .../TokiAgentCore/AgentSnapshotBuilder.swift | 11 ++- .../TokiSyncProtocol/SnapshotValidation.swift | 9 +- Sources/TokiSyncProtocol/UsageSnapshot.swift | 5 +- .../HermesUsageAttribution.swift | 26 ++++-- .../TokiUsageReaders/HermesUsageLedger.swift | 2 +- .../HermesUsageObservation.swift | 6 ++ .../HermesUsageResolver.swift | 27 +++--- .../AgentSnapshotModelAttributionTests.swift | 31 +++++++ .../SnapshotCipherTests.swift | 3 +- .../SyncSecurityTests.swift | 30 +++++++ .../RemoteSync/RemoteUsageMapper.swift | 12 +-- TokiTests/HermesPricingRefreshTests.swift | 83 +++++++++++++++++++ TokiTests/HermesReaderTests.swift | 63 ++++++++++++++ TokiTests/RemoteUsageReaderTests.swift | 39 +++++++++ 14 files changed, 316 insertions(+), 31 deletions(-) diff --git a/Sources/TokiAgentCore/AgentSnapshotBuilder.swift b/Sources/TokiAgentCore/AgentSnapshotBuilder.swift index cc9847f..ebac330 100644 --- a/Sources/TokiAgentCore/AgentSnapshotBuilder.swift +++ b/Sources/TokiAgentCore/AgentSnapshotBuilder.swift @@ -411,7 +411,8 @@ 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 activityEventSort(_ lhs: RemoteActivityEvent, _ rhs: RemoteActivityEvent) -> Bool { @@ -442,8 +443,11 @@ private extension AgentSnapshotBuilder { event.reasoningTokens, ] let validRange = 0...RemoteUsageSnapshotValidator.maximumTokenCountPerBucket + let validCostRange = 0...RemoteUsageSnapshotValidator.maximumCostPerEvent guard counts.allSatisfy(validRange.contains), - counts.contains(where: { $0 > 0 }) else { + event.cost.isFinite, + validCostRange.contains(event.cost), + counts.contains(where: { $0 > 0 }) || event.cost > 0 else { return nil } return RemoteTokenEvent( @@ -454,7 +458,8 @@ 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 var platformName: String { diff --git a/Sources/TokiSyncProtocol/SnapshotValidation.swift b/Sources/TokiSyncProtocol/SnapshotValidation.swift index aed6093..8fd7ce5 100644 --- a/Sources/TokiSyncProtocol/SnapshotValidation.swift +++ b/Sources/TokiSyncProtocol/SnapshotValidation.swift @@ -4,6 +4,7 @@ public enum RemoteUsageSnapshotValidator { public static let maximumTokenEventCount = 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 { @@ -41,7 +42,9 @@ 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 || (event.cost ?? 0) > 0 else { throw RemoteUsageSnapshotValidationError.invalidTokenEvent } } @@ -61,6 +64,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 } diff --git a/Sources/TokiSyncProtocol/UsageSnapshot.swift b/Sources/TokiSyncProtocol/UsageSnapshot.swift index 87109df..4e3c7b2 100644 --- a/Sources/TokiSyncProtocol/UsageSnapshot.swift +++ b/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 { diff --git a/Sources/TokiUsageReaders/HermesUsageAttribution.swift b/Sources/TokiUsageReaders/HermesUsageAttribution.swift index eeece9b..bc0e6c5 100644 --- a/Sources/TokiUsageReaders/HermesUsageAttribution.swift +++ b/Sources/TokiUsageReaders/HermesUsageAttribution.swift @@ -72,16 +72,24 @@ private func hermesUsageEventParts( current: observation.modelCounters, previous: previousModelCounters, maximumDelta: counters) else { - if counters.totalTokens == 0, !costOnlyModels.isEmpty { - var parts = costOnlyModels.sorted().map { - HermesUsageEventPart(model: $0, counters: .zero) - } - if includeCostOnlyResidual { - parts.append(HermesUsageEventPart(model: nil, counters: .zero)) - } - return parts + var parts: [HermesUsageEventPart] = [] + if counters.totalTokens > 0 { + parts.append(HermesUsageEventPart( + model: observation.model, + counters: counters)) + } + let representedModels = Set(parts.compactMap(\.model)) + for model in costOnlyModels.subtracting(representedModels).sorted() { + parts.append(HermesUsageEventPart(model: model, counters: .zero)) + } + if includeCostOnlyResidual, + !parts.contains(where: { $0.model == nil }) { + parts.append(HermesUsageEventPart(model: nil, counters: .zero)) + } + if parts.isEmpty, includeCostOnlyFallback { + parts.append(HermesUsageEventPart(model: observation.model, counters: .zero)) } - return [HermesUsageEventPart(model: observation.model, counters: counters)] + return parts } var combinedCounters = HermesTokenCounters.zero diff --git a/Sources/TokiUsageReaders/HermesUsageLedger.swift b/Sources/TokiUsageReaders/HermesUsageLedger.swift index 18cac65..5809ee4 100644 --- a/Sources/TokiUsageReaders/HermesUsageLedger.swift +++ b/Sources/TokiUsageReaders/HermesUsageLedger.swift @@ -179,7 +179,7 @@ private extension HermesUsageLedger { observation: observation, previous: previous, observedAt: observedAt) - let pricingTimestamp = observation.modelPricingTimestamp ?? timestamp + let pricingTimestamp = observation.incrementalModelPricingTimestamp ?? timestamp guard let cost = hermesIncrementalCost( observation: observation, previous: previous, diff --git a/Sources/TokiUsageReaders/HermesUsageObservation.swift b/Sources/TokiUsageReaders/HermesUsageObservation.swift index 399a679..7ebda04 100644 --- a/Sources/TokiUsageReaders/HermesUsageObservation.swift +++ b/Sources/TokiUsageReaders/HermesUsageObservation.swift @@ -15,6 +15,7 @@ struct HermesSessionObservation { let modelReportedCosts: [String: Double]? let modelPricingCounters: [String: HermesTokenCounters]? let modelPricingTimestamp: Date? + let incrementalModelPricingTimestamp: Date? let projectName: String? let attributionQuality: AttributionQuality @@ -32,6 +33,7 @@ struct HermesSessionObservation { modelReportedCosts: [String: Double]? = nil, modelPricingCounters: [String: HermesTokenCounters]? = nil, modelPricingTimestamp: Date? = nil, + incrementalModelPricingTimestamp: Date? = nil, projectName: String?, attributionQuality: AttributionQuality) { self.sessionID = sessionID @@ -47,6 +49,7 @@ struct HermesSessionObservation { self.modelReportedCosts = modelReportedCosts self.modelPricingCounters = modelPricingCounters self.modelPricingTimestamp = modelPricingTimestamp + self.incrementalModelPricingTimestamp = incrementalModelPricingTimestamp self.projectName = projectName self.attributionQuality = attributionQuality } @@ -67,6 +70,9 @@ func validateHermesUsageObservation( (observation.modelPricingTimestamp.map { hermesDateIsValid($0) && $0 <= observedAt } ?? true), + (observation.incrementalModelPricingTimestamp.map { + hermesDateIsValid($0) && $0 <= observedAt + } ?? true), hermesReportedCostBreakdownIsValid( observation.reportedCost, modelReportedCosts: observation.modelReportedCosts, diff --git a/Sources/TokiUsageReaders/HermesUsageResolver.swift b/Sources/TokiUsageReaders/HermesUsageResolver.swift index b444e47..f2fd6d9 100644 --- a/Sources/TokiUsageReaders/HermesUsageResolver.swift +++ b/Sources/TokiUsageReaders/HermesUsageResolver.swift @@ -30,6 +30,7 @@ private struct HermesResolvedSessionCost { let modelReportedCosts: [String: Double]? let modelPricingCounters: [String: HermesTokenCounters] let modelPricingTimestamp: Date? + let incrementalModelPricingTimestamp: Date? } private struct HermesSessionModelUsageAccumulator { @@ -84,9 +85,7 @@ private struct HermesSessionModelUsageAccumulator { private mutating func addReportedCost(_ usage: HermesSessionModelUsage) throws { reportedCost += usage.cost guard usage.cost > 0 else { return } - guard let model = usage.model, usage.counters.totalTokens > 0 else { - return - } + guard let model = usage.model else { return } let existingCost = reportedCostsByModel[model] ?? 0 guard (existingCost + usage.cost).isFinite else { throw HermesUsageLedgerError.invalidObservation @@ -128,10 +127,11 @@ enum HermesUsageResolver { models.insert(model) } let hasTokenModelUsage = accumulatedModelUsage.counters.totalTokens > 0 + let hasCostModelUsage = accumulatedModelUsage.cost > 0 let resolvedModel = models.count == 1 ? models.first : (models.isEmpty ? session.model : nil) let resolvedCost = resolveCost( session: session, - hasModelUsage: hasTokenModelUsage, + hasModelUsage: hasTokenModelUsage || hasCostModelUsage, modelUsage: accumulatedModelUsage) return HermesSessionObservation( @@ -148,6 +148,7 @@ enum HermesUsageResolver { modelReportedCosts: resolvedCost.modelReportedCosts, modelPricingCounters: resolvedCost.modelPricingCounters, modelPricingTimestamp: resolvedCost.modelPricingTimestamp, + incrementalModelPricingTimestamp: resolvedCost.incrementalModelPricingTimestamp, projectName: session.projectName, attributionQuality: session.attributionQuality) } @@ -167,7 +168,8 @@ enum HermesUsageResolver { reportedValue: sessionReportedCost, modelReportedCosts: nil, modelPricingCounters: sessionPricingCounters, - modelPricingTimestamp: session.modelPricingTimestamp) + modelPricingTimestamp: session.modelPricingTimestamp, + incrementalModelPricingTimestamp: nil) } if session.cost > modelUsage.cost { if !session.costIsDerivedFromModelPricing { @@ -180,7 +182,8 @@ enum HermesUsageResolver { reportedValue: session.cost - retainedModelPricingCost, modelReportedCosts: modelUsage.resolvedReportedCostsByModel, modelPricingCounters: modelUsage.pricingCountersByModel, - modelPricingTimestamp: modelUsage.pricingTimestamp) + modelPricingTimestamp: modelUsage.pricingTimestamp, + incrementalModelPricingTimestamp: modelUsage.pricingTimestamp) } return HermesResolvedSessionCost( value: session.cost, @@ -188,7 +191,8 @@ enum HermesUsageResolver { reportedValue: sessionReportedCost, modelReportedCosts: nil, modelPricingCounters: sessionPricingCounters, - modelPricingTimestamp: session.modelPricingTimestamp) + modelPricingTimestamp: session.modelPricingTimestamp, + incrementalModelPricingTimestamp: nil) } if modelUsage.cost > session.cost { return HermesResolvedSessionCost( @@ -197,7 +201,8 @@ enum HermesUsageResolver { reportedValue: modelUsage.reportedCost, modelReportedCosts: modelUsage.resolvedReportedCostsByModel, modelPricingCounters: modelUsage.pricingCountersByModel, - modelPricingTimestamp: modelUsage.pricingTimestamp) + modelPricingTimestamp: modelUsage.pricingTimestamp, + incrementalModelPricingTimestamp: modelUsage.pricingTimestamp) } if !session.costIsDerivedFromModelPricing, @@ -208,7 +213,8 @@ enum HermesUsageResolver { reportedValue: sessionReportedCost, modelReportedCosts: nil, modelPricingCounters: sessionPricingCounters, - modelPricingTimestamp: session.modelPricingTimestamp) + modelPricingTimestamp: session.modelPricingTimestamp, + incrementalModelPricingTimestamp: nil) } return HermesResolvedSessionCost( value: modelUsage.cost, @@ -216,7 +222,8 @@ enum HermesUsageResolver { reportedValue: modelUsage.reportedCost, modelReportedCosts: modelUsage.resolvedReportedCostsByModel, modelPricingCounters: modelUsage.pricingCountersByModel, - modelPricingTimestamp: modelUsage.pricingTimestamp) + modelPricingTimestamp: modelUsage.pricingTimestamp, + incrementalModelPricingTimestamp: modelUsage.pricingTimestamp) } private static func pricingCounters( diff --git a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift index 8a61b7d..6aa03b9 100644 --- a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift +++ b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift @@ -1,4 +1,5 @@ import Foundation +import TokiSyncProtocol import TokiUsageCore import XCTest @testable import TokiAgentCore @@ -29,6 +30,36 @@ final class AgentSnapshotModelAttributionTests: XCTestCase { XCTAssertNil(snapshot.activityEvents.first?.model) } + func test_snapshotPreservesCostOnlyEventForRemoteMapping() async throws { + let now = Date(timeIntervalSince1970: 1_784_200_000) + let eventDate = now.addingTimeInterval(-60) + var usage = RawTokenUsage(cost: 1.25) + usage.recordTokenEvent( + timestamp: eventDate, + source: "Hermes", + model: nil, + inputTokens: 0, + outputTokens: 0, + cost: 1.25) + let descriptor = LocalUsageReaderDescriptor( + reader: FixedTokenReader(name: "Hermes", usage: usage), + sourceLocations: []) + let fixture = try AgentSnapshotFixture() + defer { fixture.remove() } + let builder = AgentSnapshotBuilder( + home: fixture.root, + readerDescriptors: [descriptor]) + + let snapshot = try await builder.build(configuration: fixture.configuration, now: now) + let event = try XCTUnwrap(snapshot.tokenEvents.first) + + XCTAssertEqual(snapshot.tokenEvents.count, 1) + XCTAssertEqual(event.totalTokens, 0) + XCTAssertNil(event.model) + XCTAssertEqual(event.cost ?? -1, 1.25, accuracy: 0.000001) + XCTAssertNoThrow(try RemoteUsageSnapshotValidator.validate(snapshot, now: now)) + } + func test_snapshotPreservesMixedHermesModelAttributionAndResidual() async throws { let now = Date(timeIntervalSince1970: 1_784_200_000) let eventDate = now.addingTimeInterval(-60) diff --git a/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift b/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift index 3762ec8..0e05ed0 100644 --- a/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift +++ b/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift @@ -104,7 +104,8 @@ final class SnapshotCipherTests: XCTestCase { outputTokens: 3, cacheReadTokens: 2, cacheWriteTokens: 0, - reasoningTokens: 1), + reasoningTokens: 1, + cost: 0.25), ], activityEvents: [ RemoteActivityEvent( diff --git a/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift b/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift index 654164f..b5e3b61 100644 --- a/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift +++ b/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift @@ -104,6 +104,25 @@ final class SyncSecurityTests: XCTestCase { XCTAssertEqual(summary.syncIntervalSeconds, TokiSyncLimits.defaultSyncIntervalSeconds) } + func test_v1TokenEventsDefaultMissingCost() throws { + let timestamp = Date(timeIntervalSince1970: 1_750_000_000) + let data = try TokiSyncCoding.makeEncoder().encode(LegacyRemoteTokenEvent( + timestamp: timestamp, + source: "Codex", + model: "gpt-5", + inputTokens: 10, + outputTokens: 3, + cacheReadTokens: 2, + cacheWriteTokens: 0, + reasoningTokens: 1)) + + let event = try TokiSyncCoding.makeDecoder().decode(RemoteTokenEvent.self, from: data) + + XCTAssertEqual(event.timestamp, timestamp) + XCTAssertEqual(event.totalTokens, 16) + XCTAssertNil(event.cost) + } + func test_v1PairingBundleDefaultsMissingRetentionAndSyncInterval() throws { let hubURL = try XCTUnwrap(URL(string: "https://hub.example.test")) let data = try TokiSyncCoding.makeEncoder().encode(LegacyAgentPairingBundle( @@ -138,6 +157,17 @@ private struct LegacyRemoteDeviceSummary: Encodable { let latestSequence: UInt64? } +private struct LegacyRemoteTokenEvent: Encodable { + let timestamp: Date + let source: String + let model: String? + let inputTokens: Int + let outputTokens: Int + let cacheReadTokens: Int + let cacheWriteTokens: Int + let reasoningTokens: Int +} + private struct LegacyAgentPairingBundle: Encodable { let schemaVersion: Int let hubURL: URL diff --git a/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift b/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift index bff5962..0202543 100644 --- a/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift +++ b/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift @@ -95,11 +95,10 @@ struct RemoteUsageMapper { result.cacheWriteTokens += event.cacheWriteTokens result.reasoningTokens += event.reasoningTokens result.cost += cost - if let model { - result.perModel[model, default: PerModelUsage()].totalTokens += event.totalTokens - result.perModel[model, default: PerModelUsage()].cost += cost - result.perModel[model, default: PerModelUsage()].sources.insert(source) - } + let modelGroupingKey = model ?? UsageModelGrouping.mixedOrUnattributedKey + result.perModel[modelGroupingKey, default: PerModelUsage()].totalTokens += event.totalTokens + result.perModel[modelGroupingKey, default: PerModelUsage()].cost += cost + result.perModel[modelGroupingKey, default: PerModelUsage()].sources.insert(source) result.recordTokenEvent( timestamp: event.timestamp, @@ -114,6 +113,9 @@ struct RemoteUsageMapper { } private func tokenCost(for event: RemoteTokenEvent, model: String?) -> Double { + if let cost = event.cost, cost.isFinite, cost >= 0 { + return cost + } guard let model, let price = modelPrice(for: model, at: event.timestamp) else { return 0 } return price.cost( input: event.inputTokens, diff --git a/TokiTests/HermesPricingRefreshTests.swift b/TokiTests/HermesPricingRefreshTests.swift index 60a7620..9bdc208 100644 --- a/TokiTests/HermesPricingRefreshTests.swift +++ b/TokiTests/HermesPricingRefreshTests.swift @@ -71,6 +71,89 @@ final class HermesPricingRefreshTests: XCTestCase { updatedPrice.cost(input: 60, output: 0, cacheRead: 0, cacheWrite: 0), accuracy: 0.000001) } +} + +extension HermesPricingRefreshTests { + func test_hermesReader_pricesSessionOnlyIncrementAtLatestActivityAfterRateChange() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { + ModelPricingSupplement.install([:]) + try? FileManager.default.removeItem(at: tempDir) + } + let model = "session-only-pricing-model" + let initialPrice = uniformHermesModelPrice(perMillion: 1) + let updatedPrice = uniformHermesModelPrice(perMillion: 10) + ModelPricingSupplement.install(priceHistories: [ + model: [ + ModelPricingSupplement.PriceVersion( + effectiveFrom: .distantPast, + price: initialPrice), + ModelPricingSupplement.PriceVersion( + effectiveFrom: tokiTestISODate("2026-04-10T10:00:00Z"), + price: updatedPrice), + ], + ]) + + let dbURL = tempDir.appendingPathComponent("state.db") + let ledgerURL = tempDir.appendingPathComponent("hermes-usage-ledger.json") + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "session-only-pricing", + startedAt: "2026-04-10T09:00:00Z", + model: model, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: 0)]) + let ledger = HermesUsageLedger(fileURL: ledgerURL) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T08:00:00Z")) + + let initialUsage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T09:30:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-10T10:00:00Z")) + XCTAssertEqual( + initialUsage.cost, + initialPrice.cost(input: 100, output: 0, cacheRead: 0, cacheWrite: 0), + accuracy: 0.000001) + + try updateHermesSession( + databaseURL: dbURL, + id: "session-only-pricing", + model: model, + inputTokens: 160) + let latestActivity = tokiTestISODate("2026-04-10T10:30:00Z") + try insertHermesMessage( + databaseURL: dbURL, + sessionID: "session-only-pricing", + timestamp: latestActivity) + + let increment = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: HermesUsageLedger(fileURL: ledgerURL), + now: { tokiTestISODate("2026-04-10T11:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T10:00:00Z"), + to: tokiTestISODate("2026-04-10T12:00:00Z")) + + XCTAssertEqual(increment.inputTokens, 60) + XCTAssertEqual( + increment.cost, + updatedPrice.cost(input: 60, output: 0, cacheRead: 0, cacheWrite: 0), + accuracy: 0.000001) + XCTAssertEqual(increment.tokenEvents.map(\.timestamp), [latestActivity]) + } func test_hermesReader_pricesMultiModelTokenDeltasAfterRateChanges() async throws { let tempDir = try makeHermesTemporaryDirectory() diff --git a/TokiTests/HermesReaderTests.swift b/TokiTests/HermesReaderTests.swift index 292c6ab..82fdcd3 100644 --- a/TokiTests/HermesReaderTests.swift +++ b/TokiTests/HermesReaderTests.swift @@ -362,6 +362,69 @@ final class HermesReaderTests: XCTestCase { XCTAssertNil(usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]) } + func test_hermesReader_preservesReportedCostFromZeroTokenDetailRow() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + + let sessionModel = "session-model-without-catalog-price" + let costModel = "cost-only-detail-model" + let dbURL = tempDir.appendingPathComponent("state.db") + let ledger = HermesUsageLedger( + fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + try createHermesStateDB( + at: dbURL, + rows: [HermesSessionFixture( + id: "zero-token-cost-detail", + startedAt: "2026-04-10T09:00:00Z", + model: sessionModel, + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cwd: nil, + gitRepoRoot: nil, + estimatedCost: 0, + actualCost: 0)]) + try insertHermesModelUsage( + databaseURL: dbURL, + rows: [HermesModelUsageFixture( + sessionID: "zero-token-cost-detail", + model: costModel, + task: "cost-only", + apiCallCount: 1, + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + estimatedCost: 0, + actualCost: 2)]) + try await ledger.refresh( + observations: [], + observedAt: tokiTestISODate("2026-04-10T08:00:00Z")) + + let usage = try await HermesReader( + dbPathOverride: dbURL.path, + usageLedger: ledger, + now: { tokiTestISODate("2026-04-10T10:00:00Z") }) + .readUsage( + from: tokiTestISODate("2026-04-10T00:00:00Z"), + to: tokiTestISODate("2026-04-11T00:00:00Z")) + + XCTAssertEqual(usage.inputTokens, 100) + XCTAssertEqual(usage.cost, 2, accuracy: 0.000001) + XCTAssertEqual(usage.perModel[sessionModel]?.totalTokens, 100) + XCTAssertEqual(usage.perModel[sessionModel]?.cost ?? -1, 0, accuracy: 0.000001) + XCTAssertEqual(usage.perModel[costModel]?.totalTokens, 0) + XCTAssertEqual(usage.perModel[costModel]?.cost ?? -1, 2, accuracy: 0.000001) + XCTAssertEqual( + usage.tokenEvents.first { $0.model == costModel }?.cost ?? -1, + 2, + accuracy: 0.000001) + XCTAssertNil(usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]) + } + func test_hermesReader_tokenHelpersUseSameRangeAndZeroTokenFilters() async throws { let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString) try FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true) diff --git a/TokiTests/RemoteUsageReaderTests.swift b/TokiTests/RemoteUsageReaderTests.swift index 594b1a6..6f36c93 100644 --- a/TokiTests/RemoteUsageReaderTests.swift +++ b/TokiTests/RemoteUsageReaderTests.swift @@ -163,6 +163,45 @@ final class RemoteUsageReaderTests: XCTestCase { XCTAssertEqual(mixedUsage.activeSeconds, 30) XCTAssertEqual(mixedUsage.sources, ["Hermes · build-server"]) } + + func test_costOnlyRemoteEventPreservesOverallSourceAndMixedModelCost() throws { + let fixture = try makeFixture() + let original = try SnapshotCipher.open(fixture.envelope, key: fixture.encryptionKey) + let snapshot = RemoteUsageSnapshot( + device: original.device, + generatedAt: original.generatedAt, + coveredFrom: original.coveredFrom, + coveredTo: original.coveredTo, + tokenEvents: [ + RemoteTokenEvent( + timestamp: fixture.start.addingTimeInterval(60), + source: "Hermes", + model: nil, + inputTokens: 0, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0, + cost: 1.25), + ], + activityEvents: []) + + let slice = try XCTUnwrap(RemoteUsageMapper().usageSlice( + from: snapshot, + startDate: fixture.start, + endDate: fixture.end)) + let mixedUsage = try XCTUnwrap( + slice.usage.perModel[UsageModelGrouping.mixedOrUnattributedKey]) + + XCTAssertEqual(slice.usage.totalTokens, 0) + XCTAssertEqual(slice.usage.cost, 1.25, accuracy: 0.000001) + XCTAssertEqual(slice.usage.tokenEvents.count, 1) + XCTAssertEqual(slice.usage.tokenEvents.first?.cost ?? -1, 1.25, accuracy: 0.000001) + XCTAssertEqual(mixedUsage.totalTokens, 0) + XCTAssertEqual(mixedUsage.cost, 1.25, accuracy: 0.000001) + XCTAssertEqual(mixedUsage.sources, ["Hermes · build-server"]) + XCTAssertEqual(slice.sourceStats.first?.cost ?? -1, 1.25, accuracy: 0.000001) + } } extension RemoteUsageReaderTests { From 06631362fb2acb2fd1c1745c77d97a443ebd7794 Mon Sep 17 00:00:00 2001 From: choi138 Date: Fri, 31 Jul 2026 15:17:31 +0900 Subject: [PATCH 07/10] fix(hermes): preserve detailed cost attribution --- .../TokiAgentCore/AgentSnapshotBuilder.swift | 40 ++++- .../TokiSyncProtocol/SnapshotValidation.swift | 20 ++- Sources/TokiSyncProtocol/UsageSnapshot.swift | 21 +++ .../HermesUsageResolver.swift | 43 ++---- .../AgentSnapshotModelAttributionTests.swift | 8 +- .../SnapshotCipherTests.swift | 7 + .../SyncSecurityTests.swift | 56 ++++++- .../RemoteSync/RemoteUsageMapper.swift | 51 +++++++ TokiTests/HermesPricingRefreshTests.swift | 137 ++++++++++++++++++ TokiTests/RemoteUsageReaderTests.swift | 10 +- 10 files changed, 344 insertions(+), 49 deletions(-) diff --git a/Sources/TokiAgentCore/AgentSnapshotBuilder.swift b/Sources/TokiAgentCore/AgentSnapshotBuilder.swift index ebac330..e11e727 100644 --- a/Sources/TokiAgentCore/AgentSnapshotBuilder.swift +++ b/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)) } @@ -415,6 +421,13 @@ private extension AgentSnapshotBuilder { 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 { if lhs.timestamp != rhs.timestamp { return lhs.timestamp < rhs.timestamp } if lhs.source != rhs.source { return lhs.source < rhs.source } @@ -447,7 +460,7 @@ private extension AgentSnapshotBuilder { guard counts.allSatisfy(validRange.contains), event.cost.isFinite, validCostRange.contains(event.cost), - counts.contains(where: { $0 > 0 }) || event.cost > 0 else { + counts.contains(where: { $0 > 0 }) else { return nil } return RemoteTokenEvent( @@ -462,6 +475,28 @@ private extension AgentSnapshotBuilder { 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 { #if os(Linux) "linux" @@ -484,6 +519,7 @@ private struct AgentSnapshotContent: Encodable { let coveredFrom: Date let coveredTo: Date let tokenEvents: [RemoteTokenEvent] + let costEvents: [RemoteCostEvent]? let activityEvents: [RemoteActivityEvent] } diff --git a/Sources/TokiSyncProtocol/SnapshotValidation.swift b/Sources/TokiSyncProtocol/SnapshotValidation.swift index 8fd7ce5..a77a54b 100644 --- a/Sources/TokiSyncProtocol/SnapshotValidation.swift +++ b/Sources/TokiSyncProtocol/SnapshotValidation.swift @@ -2,6 +2,7 @@ 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 @@ -29,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 } @@ -44,11 +48,22 @@ public enum RemoteUsageSnapshotValidator { validTokenCount(event.cacheWriteTokens), validTokenCount(event.reasoningTokens), event.cost.map(validCost) ?? true, - event.totalTokens > 0 || (event.cost ?? 0) > 0 else { + 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, @@ -84,6 +99,7 @@ public enum RemoteUsageSnapshotValidationError: LocalizedError { case invalidDateRange case tooManyEvents case invalidTokenEvent + case invalidCostEvent case invalidActivityEvent public var errorDescription: String? { @@ -98,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." } diff --git a/Sources/TokiSyncProtocol/UsageSnapshot.swift b/Sources/TokiSyncProtocol/UsageSnapshot.swift index 4e3c7b2..c3f5fd3 100644 --- a/Sources/TokiSyncProtocol/UsageSnapshot.swift +++ b/Sources/TokiSyncProtocol/UsageSnapshot.swift @@ -59,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 } @@ -95,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( @@ -104,6 +123,7 @@ public struct RemoteUsageSnapshot: Codable, Equatable, Sendable { coveredFrom: Date, coveredTo: Date, tokenEvents: [RemoteTokenEvent], + costEvents: [RemoteCostEvent]? = nil, activityEvents: [RemoteActivityEvent]) { self.schemaVersion = schemaVersion self.device = device @@ -111,6 +131,7 @@ public struct RemoteUsageSnapshot: Codable, Equatable, Sendable { self.coveredFrom = coveredFrom self.coveredTo = coveredTo self.tokenEvents = tokenEvents + self.costEvents = costEvents self.activityEvents = activityEvents } } diff --git a/Sources/TokiUsageReaders/HermesUsageResolver.swift b/Sources/TokiUsageReaders/HermesUsageResolver.swift index f2fd6d9..1f3d47b 100644 --- a/Sources/TokiUsageReaders/HermesUsageResolver.swift +++ b/Sources/TokiUsageReaders/HermesUsageResolver.swift @@ -48,10 +48,6 @@ private struct HermesSessionModelUsageAccumulator { reportedCostsByModel } - var hasResolvedReportedCostBreakdown: Bool { - resolvedReportedCostsByModel?.isEmpty == false - } - mutating func add(_ usage: HermesSessionModelUsage) throws { guard usage.counters.isValid(), counters.canAdd(usage.counters, maximum: hermesLedgerMaximumCumulativeTokens), @@ -172,27 +168,17 @@ enum HermesUsageResolver { incrementalModelPricingTimestamp: nil) } if session.cost > modelUsage.cost { - if !session.costIsDerivedFromModelPricing { - let retainedModelPricingCost = max( - 0, - modelUsage.cost - modelUsage.reportedCost) - return HermesResolvedSessionCost( - value: session.cost, - isDerivedFromModelPricing: false, - reportedValue: session.cost - retainedModelPricingCost, - modelReportedCosts: modelUsage.resolvedReportedCostsByModel, - modelPricingCounters: modelUsage.pricingCountersByModel, - modelPricingTimestamp: modelUsage.pricingTimestamp, - incrementalModelPricingTimestamp: modelUsage.pricingTimestamp) - } + let retainedModelPricingCost = max( + 0, + modelUsage.cost - modelUsage.reportedCost) return HermesResolvedSessionCost( value: session.cost, - isDerivedFromModelPricing: session.costIsDerivedFromModelPricing, - reportedValue: sessionReportedCost, - modelReportedCosts: nil, - modelPricingCounters: sessionPricingCounters, - modelPricingTimestamp: session.modelPricingTimestamp, - incrementalModelPricingTimestamp: nil) + isDerivedFromModelPricing: false, + reportedValue: session.cost - retainedModelPricingCost, + modelReportedCosts: modelUsage.resolvedReportedCostsByModel, + modelPricingCounters: modelUsage.pricingCountersByModel, + modelPricingTimestamp: modelUsage.pricingTimestamp, + incrementalModelPricingTimestamp: modelUsage.pricingTimestamp) } if modelUsage.cost > session.cost { return HermesResolvedSessionCost( @@ -205,17 +191,6 @@ enum HermesUsageResolver { incrementalModelPricingTimestamp: modelUsage.pricingTimestamp) } - if !session.costIsDerivedFromModelPricing, - !modelUsage.hasResolvedReportedCostBreakdown { - return HermesResolvedSessionCost( - value: session.cost, - isDerivedFromModelPricing: false, - reportedValue: sessionReportedCost, - modelReportedCosts: nil, - modelPricingCounters: sessionPricingCounters, - modelPricingTimestamp: session.modelPricingTimestamp, - incrementalModelPricingTimestamp: nil) - } return HermesResolvedSessionCost( value: modelUsage.cost, isDerivedFromModelPricing: modelUsage.costIsDerivedFromModelPricing, diff --git a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift index 6aa03b9..1241b57 100644 --- a/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift +++ b/Tests/TokiAgentTests/AgentSnapshotModelAttributionTests.swift @@ -51,12 +51,12 @@ final class AgentSnapshotModelAttributionTests: XCTestCase { readerDescriptors: [descriptor]) let snapshot = try await builder.build(configuration: fixture.configuration, now: now) - let event = try XCTUnwrap(snapshot.tokenEvents.first) + let event = try XCTUnwrap(snapshot.costEvents?.first) - XCTAssertEqual(snapshot.tokenEvents.count, 1) - XCTAssertEqual(event.totalTokens, 0) + XCTAssertTrue(snapshot.tokenEvents.isEmpty) + XCTAssertEqual(snapshot.costEvents?.count, 1) XCTAssertNil(event.model) - XCTAssertEqual(event.cost ?? -1, 1.25, accuracy: 0.000001) + XCTAssertEqual(event.cost, 1.25, accuracy: 0.000001) XCTAssertNoThrow(try RemoteUsageSnapshotValidator.validate(snapshot, now: now)) } diff --git a/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift b/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift index 0e05ed0..a43de38 100644 --- a/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift +++ b/Tests/TokiSyncProtocolTests/SnapshotCipherTests.swift @@ -107,6 +107,13 @@ final class SnapshotCipherTests: XCTestCase { reasoningTokens: 1, cost: 0.25), ], + costEvents: [ + RemoteCostEvent( + timestamp: generatedAt.addingTimeInterval(-30), + source: "Hermes", + model: nil, + cost: 0.5), + ], activityEvents: [ RemoteActivityEvent( timestamp: generatedAt.addingTimeInterval(-60), diff --git a/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift b/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift index b5e3b61..324f240 100644 --- a/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift +++ b/Tests/TokiSyncProtocolTests/SyncSecurityTests.swift @@ -141,6 +141,50 @@ final class SyncSecurityTests: XCTestCase { } } +extension SyncSecurityTests { + func test_v1SnapshotsKeepCostEventsBackwardCompatible() throws { + let generatedAt = Date(timeIntervalSince1970: 1_750_000_000) + let device = RemoteDeviceDescriptor(id: "device-1", name: "ubuntu", platform: "linux") + let legacySnapshot = LegacyRemoteUsageSnapshot( + schemaVersion: TokiSyncProtocolVersion.current, + device: device, + generatedAt: generatedAt, + coveredFrom: generatedAt.addingTimeInterval(-3600), + coveredTo: generatedAt.addingTimeInterval(1), + tokenEvents: [], + activityEvents: []) + let encoder = TokiSyncCoding.makeEncoder() + let decoder = TokiSyncCoding.makeDecoder() + + let decodedCurrentSnapshot = try decoder.decode( + RemoteUsageSnapshot.self, + from: encoder.encode(legacySnapshot)) + XCTAssertNil(decodedCurrentSnapshot.costEvents) + + let currentSnapshot = RemoteUsageSnapshot( + device: device, + generatedAt: generatedAt, + coveredFrom: generatedAt.addingTimeInterval(-3600), + coveredTo: generatedAt.addingTimeInterval(1), + tokenEvents: [], + costEvents: [ + RemoteCostEvent( + timestamp: generatedAt.addingTimeInterval(-60), + source: "Hermes", + model: nil, + cost: 1.25), + ], + activityEvents: []) + let decodedLegacySnapshot = try decoder.decode( + LegacyRemoteUsageSnapshot.self, + from: encoder.encode(currentSnapshot)) + + XCTAssertEqual(decodedLegacySnapshot.schemaVersion, TokiSyncProtocolVersion.current) + XCTAssertTrue(decodedLegacySnapshot.tokenEvents.isEmpty) + XCTAssertNoThrow(try RemoteUsageSnapshotValidator.validate(currentSnapshot, now: generatedAt)) + } +} + private struct OversizedEncodable: Encodable { let value: String } @@ -157,7 +201,7 @@ private struct LegacyRemoteDeviceSummary: Encodable { let latestSequence: UInt64? } -private struct LegacyRemoteTokenEvent: Encodable { +private struct LegacyRemoteTokenEvent: Codable { let timestamp: Date let source: String let model: String? @@ -168,6 +212,16 @@ private struct LegacyRemoteTokenEvent: Encodable { let reasoningTokens: Int } +private struct LegacyRemoteUsageSnapshot: Codable { + let schemaVersion: Int + let device: RemoteDeviceDescriptor + let generatedAt: Date + let coveredFrom: Date + let coveredTo: Date + let tokenEvents: [LegacyRemoteTokenEvent] + let activityEvents: [RemoteActivityEvent] +} + private struct LegacyAgentPairingBundle: Encodable { let schemaVersion: Int let hubURL: URL diff --git a/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift b/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift index 0202543..5aa016f 100644 --- a/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift +++ b/Toki/Infrastructure/RemoteSync/RemoteUsageMapper.swift @@ -31,6 +31,12 @@ struct RemoteUsageMapper { to: &sourceUsage) usageBySource[event.source] = sourceUsage } + appendCostEvents( + from: snapshot, + startDate: startDate, + endDate: endDate, + usage: &usage, + usageBySource: &usageBySource) usage.activityEvents.append(contentsOf: mappedActivityEvents( from: snapshot, @@ -124,6 +130,25 @@ struct RemoteUsageMapper { cacheWrite: event.cacheWriteTokens) } + private func appendCostEvent( + _ event: RemoteCostEvent, + model: String?, + source: String, + to result: inout RawTokenUsage) { + result.cost += event.cost + let modelGroupingKey = model ?? UsageModelGrouping.mixedOrUnattributedKey + result.perModel[modelGroupingKey, default: PerModelUsage()].cost += event.cost + result.perModel[modelGroupingKey, default: PerModelUsage()].sources.insert(source) + + result.recordTokenEvent( + timestamp: event.timestamp, + source: source, + model: model, + inputTokens: 0, + outputTokens: 0, + cost: event.cost) + } + func mappedActivityEventsBySource( from snapshot: RemoteUsageSnapshot, startDate: Date, @@ -176,3 +201,29 @@ struct RemoteUsageMapper { "\(source) · \(deviceName)" } } + +private extension RemoteUsageMapper { + func appendCostEvents( + from snapshot: RemoteUsageSnapshot, + startDate: Date, + endDate: Date, + usage: inout RawTokenUsage, + usageBySource: inout [String: RawTokenUsage]) { + for event in snapshot.costEvents ?? [] where event.timestamp >= startDate && event.timestamp < endDate { + let model = normalizedModelID(event.model) + appendCostEvent( + event, + model: model, + source: deviceSource(event.source, deviceName: snapshot.device.name), + to: &usage) + + var sourceUsage = usageBySource[event.source] ?? RawTokenUsage() + appendCostEvent( + event, + model: model, + source: event.source, + to: &sourceUsage) + usageBySource[event.source] = sourceUsage + } + } +} diff --git a/TokiTests/HermesPricingRefreshTests.swift b/TokiTests/HermesPricingRefreshTests.swift index 9bdc208..4cec0cc 100644 --- a/TokiTests/HermesPricingRefreshTests.swift +++ b/TokiTests/HermesPricingRefreshTests.swift @@ -851,6 +851,143 @@ extension HermesPricingRefreshTests { } } +extension HermesPricingRefreshTests { + func test_resolverRetainsReportedModelCostsWhenDerivedSessionTotalWins() throws { + let timestamp = tokiTestISODate("2026-04-10T10:00:00Z") + let expensiveModel = "reported-detail-expensive" + let cheapModel = "reported-detail-cheap" + let expensiveCounters = HermesTokenCounters( + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0) + let cheapCounters = HermesTokenCounters( + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0) + let sessionCounters = expensiveCounters.adding(cheapCounters) + let session = HermesSessionObservation( + sessionID: "derived-session-total-wins", + startedAt: timestamp, + earliestActivityAt: timestamp, + latestActivityAt: timestamp, + model: nil, + counters: sessionCounters, + cost: 20, + costIsDerivedFromModelPricing: true, + modelPricingTimestamp: timestamp, + projectName: nil, + attributionQuality: .exact) + let resolved = try HermesUsageResolver.resolve( + session: session, + modelUsage: [ + HermesSessionModelUsage( + model: expensiveModel, + counters: expensiveCounters, + cost: 10, + costIsDerivedFromModelPricing: false, + modelPricingTimestamp: nil), + HermesSessionModelUsage( + model: cheapModel, + counters: cheapCounters, + cost: 1, + costIsDerivedFromModelPricing: false, + modelPricingTimestamp: nil), + ]) + + let events = hermesUsageEvents( + identifier: "derived-session-total-wins", + timestamp: timestamp, + observation: resolved, + previousModelCounters: [:], + previousReportedCost: 0, + previousModelReportedCosts: [:], + previousModelPricingCounters: [:], + counters: resolved.counters, + cost: resolved.cost, + pricingTimestamp: timestamp) + + XCTAssertEqual(resolved.cost, 20, accuracy: 0.000001) + XCTAssertEqual(resolved.reportedCost ?? -1, 20, accuracy: 0.000001) + XCTAssertEqual(resolved.modelReportedCosts?[expensiveModel] ?? -1, 10, accuracy: 0.000001) + XCTAssertEqual(resolved.modelReportedCosts?[cheapModel] ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual(events.first { $0.model == expensiveModel }?.cost ?? -1, 10, accuracy: 0.000001) + XCTAssertEqual(events.first { $0.model == cheapModel }?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertEqual(events.first { $0.model == nil }?.cost ?? -1, 9, accuracy: 0.000001) + } + + func test_resolverRetainsDetailedPricingWhenReportedSessionTotalTies() throws { + let timestamp = tokiTestISODate("2026-04-10T10:00:00Z") + let expensiveModel = "derived-detail-expensive" + let cheapModel = "derived-detail-cheap" + defer { ModelPricingSupplement.install([:]) } + ModelPricingSupplement.install([ + expensiveModel: uniformHermesModelPrice(perMillion: 10), + cheapModel: uniformHermesModelPrice(perMillion: 1), + ]) + let expensiveCounters = HermesTokenCounters( + inputTokens: 1_000_000, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0) + let cheapCounters = HermesTokenCounters( + inputTokens: 1_000_000, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0) + let session = HermesSessionObservation( + sessionID: "reported-session-total-ties", + startedAt: timestamp, + earliestActivityAt: timestamp, + latestActivityAt: timestamp, + model: nil, + counters: expensiveCounters.adding(cheapCounters), + cost: 11, + projectName: nil, + attributionQuality: .exact) + let resolved = try HermesUsageResolver.resolve( + session: session, + modelUsage: [ + HermesSessionModelUsage( + model: expensiveModel, + counters: expensiveCounters, + cost: 10, + costIsDerivedFromModelPricing: true, + modelPricingTimestamp: timestamp), + HermesSessionModelUsage( + model: cheapModel, + counters: cheapCounters, + cost: 1, + costIsDerivedFromModelPricing: true, + modelPricingTimestamp: timestamp), + ]) + + let events = hermesUsageEvents( + identifier: "reported-session-total-ties", + timestamp: timestamp, + observation: resolved, + previousModelCounters: [:], + previousReportedCost: 0, + previousModelReportedCosts: [:], + previousModelPricingCounters: [:], + counters: resolved.counters, + cost: resolved.cost, + pricingTimestamp: timestamp) + + XCTAssertEqual(resolved.cost, 11, accuracy: 0.000001) + XCTAssertEqual(resolved.reportedCost ?? -1, 0, accuracy: 0.000001) + XCTAssertEqual(resolved.modelPricingCounters?.count, 2) + XCTAssertEqual(events.first { $0.model == expensiveModel }?.cost ?? -1, 10, accuracy: 0.000001) + XCTAssertEqual(events.first { $0.model == cheapModel }?.cost ?? -1, 1, accuracy: 0.000001) + XCTAssertNil(events.first { $0.model == nil }) + } +} + private func uniformHermesModelPrice(perMillion: Double) -> ModelPrice { ModelPrice( inputPerMillion: perMillion, diff --git a/TokiTests/RemoteUsageReaderTests.swift b/TokiTests/RemoteUsageReaderTests.swift index 6f36c93..aa64e21 100644 --- a/TokiTests/RemoteUsageReaderTests.swift +++ b/TokiTests/RemoteUsageReaderTests.swift @@ -172,16 +172,12 @@ final class RemoteUsageReaderTests: XCTestCase { generatedAt: original.generatedAt, coveredFrom: original.coveredFrom, coveredTo: original.coveredTo, - tokenEvents: [ - RemoteTokenEvent( + tokenEvents: [], + costEvents: [ + RemoteCostEvent( timestamp: fixture.start.addingTimeInterval(60), source: "Hermes", model: nil, - inputTokens: 0, - outputTokens: 0, - cacheReadTokens: 0, - cacheWriteTokens: 0, - reasoningTokens: 0, cost: 1.25), ], activityEvents: []) From e25ead8a7dd236f3851a218f2ee21c3798845dbc Mon Sep 17 00:00:00 2001 From: choi138 Date: Tue, 4 Aug 2026 16:07:13 +0900 Subject: [PATCH 08/10] feat(pricing): add Korea-routed Claude Opus 5 rates The custom billing provider reports kr/claude-opus-5 as its own catalog entry rather than as a provider prefix on claude-opus-5, so usage under that ID was left unpriced. Add an explicit exact key at the claude-opus-5 rates. Keep the entry exact-only so it cannot lend its rates to future tiers, and cover that slash-bearing keys stay independent of prefix stripping. --- Sources/TokiUsageReaders/ModelPricing.swift | 5 +++ .../Claude5ModelPricingBehaviorTests.swift | 42 +++++++++++++++++++ 2 files changed, 47 insertions(+) diff --git a/Sources/TokiUsageReaders/ModelPricing.swift b/Sources/TokiUsageReaders/ModelPricing.swift index ea6ddcc..597fe86 100644 --- a/Sources/TokiUsageReaders/ModelPricing.swift +++ b/Sources/TokiUsageReaders/ModelPricing.swift @@ -74,6 +74,10 @@ private let exactPricingTable: [String: ModelPrice] = [ // fast-mode usage is under-estimated by this table. "claude-fable-5": price(10.0, 50.0, 1.00, 12.5), "claude-opus-5": price(5.0, 25.0, 0.50, 6.25), + // Reported by the custom billing provider as its own catalog entry rather + // than as a provider prefix on claude-opus-5, so it needs an explicit key. + // Rates match claude-opus-5; it is exact-only for the same reason. + "kr/claude-opus-5": price(5.0, 25.0, 0.50, 6.25), // Introductory pricing through 2026-08-31 (UTC); the standard rate from // 2026-09-01 is applied per usage timestamp via scheduledPriceChanges. "claude-sonnet-5": price(2.0, 10.0, 0.20, 2.50), @@ -144,6 +148,7 @@ private let exactPricingTable: [String: ModelPrice] = [ private let exactOnlyPricingKeys: Set = [ "claude-fable-5", "claude-opus-5", + "kr/claude-opus-5", "claude-sonnet-5", "claude-opus-4", "gpt-5", diff --git a/TokiTests/Claude5ModelPricingBehaviorTests.swift b/TokiTests/Claude5ModelPricingBehaviorTests.swift index eaadf56..7fbdca5 100644 --- a/TokiTests/Claude5ModelPricingBehaviorTests.swift +++ b/TokiTests/Claude5ModelPricingBehaviorTests.swift @@ -99,6 +99,48 @@ final class Claude5ModelPricingBehaviorTests: XCTestCase { XCTAssertNil(modelPrice(for: "claude-sonnet-5-1")) } + func test_modelPrice_matchesKoreaRoutedClaudeOpus5() throws { + // kr/claude-opus-5 is the model ID reported by the custom billing + // provider and bills at the same rates as claude-opus-5. It is a + // distinct catalog entry rather than a provider prefix, so it carries + // its own exact key instead of inheriting through a stripping rule. + let lookup = modelPriceLookup(for: "kr/claude-opus-5", at: Self.claude5IntroductoryDate) + let price = try XCTUnwrap(lookup.price) + let base = try XCTUnwrap(modelPrice(for: "claude-opus-5", at: Self.claude5IntroductoryDate)) + + XCTAssertEqual(lookup.match, .exact(modelId: "kr/claude-opus-5")) + XCTAssertEqual(price.inputPerMillion, base.inputPerMillion, accuracy: 0.0001) + XCTAssertEqual(price.outputPerMillion, base.outputPerMillion, accuracy: 0.0001) + XCTAssertEqual(price.cacheReadPerMillion, base.cacheReadPerMillion, accuracy: 0.0001) + XCTAssertEqual(price.cacheWritePerMillion, base.cacheWritePerMillion, accuracy: 0.0001) + } + + func test_modelPriceIsKnown_reportsKoreaRoutedClaudeOpus5AsPriced() { + let interval = DateInterval(start: Self.claude5IntroductoryDate, duration: 86400) + + XCTAssertTrue(modelPriceIsKnown(for: "kr/claude-opus-5", throughout: interval)) + } + + func test_modelPrice_treatsKoreaRoutedClaudeOpus5AsExactOnly() { + // The kr/ entry must not become a prefix that lends its rates to + // future tiers, matching how claude-opus-5 itself is exact-only. + XCTAssertNil(modelPrice(for: "kr/claude-opus-5-1")) + XCTAssertNil(modelPrice(for: "kr/claude-opus-5-mini")) + // Other kr/-reported models stay unpriced until their own rates are known. + XCTAssertNil(modelPrice(for: "kr/claude-fable-5")) + XCTAssertNil(modelPrice(for: "kr/gpt-5.6-sol")) + } + + func test_modelPrice_keepsSlashBearingModelKeysIndependent() throws { + // zai-org/ is part of the model ID, not a provider prefix. Adding a + // slash-bearing key must not turn the segment before the slash into + // something strippable. + XCTAssertNotNil(try XCTUnwrap(modelPrice(for: "zai-org/GLM-5.2"))) + XCTAssertNotNil(try XCTUnwrap(modelPrice(for: "zai-org/GLM-5.2-Batch"))) + XCTAssertNil(modelPrice(for: "GLM-5.2")) + XCTAssertNil(modelPrice(for: "claude-opus-5-1")) + } + func test_modelPrice_matchesClaudeOpus47And48() throws { for modelID in ["claude-opus-4-7", "claude-opus-4-8"] { let price = try XCTUnwrap(modelPrice(for: modelID)) From f3df13dab48940682fb47473b29a502c51d11071 Mon Sep 17 00:00:00 2001 From: choi138 Date: Tue, 4 Aug 2026 16:08:24 +0900 Subject: [PATCH 09/10] chore(git): ignore local agent scratch directories Hermes, omo, and senpi write task logs and plans into the working tree, which left dozens of untracked files in every status check and risked committing local agent state. --- .gitignore | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.gitignore b/.gitignore index 31aa62c..0ecd495 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,9 @@ rabbit_preview.png .codegraph/ .agents/hooks/state/ .agents/settings.local.json + +# Local agent scratch state +.hermes/ +.omo/ +.senpi/ +.debug-journal.md From 30f7e040db8ca7109efa15f683653310b0c8f593 Mon Sep 17 00:00:00 2001 From: choi138 Date: Tue, 4 Aug 2026 16:17:39 +0900 Subject: [PATCH 10/10] fix(hermes): attribute single-model residuals to the resolved model When session counters exceeded a partial session_model_usage breakdown, the residual was always emitted with a nil model. A 150-token session with a 100-token detail row on the same model was therefore reported as 100 tokens for that model and 50 as Mixed / Unattributed, and the residual cost followed the same wrong split. The resolver leaves the observation model non-nil only when the session and every detail row agree on one model, so use it for the residual and reserve nil for genuinely ambiguous sessions. Merge into the existing part for that model so its pricing counters are not charged twice, and keep emitting an unattributed carrier when reported cost belongs to no model. --- .../HermesUsageAttribution.swift | 19 ++++++- TokiTests/HermesUsageLedgerTests.swift | 52 +++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) diff --git a/Sources/TokiUsageReaders/HermesUsageAttribution.swift b/Sources/TokiUsageReaders/HermesUsageAttribution.swift index bc0e6c5..e45aea5 100644 --- a/Sources/TokiUsageReaders/HermesUsageAttribution.swift +++ b/Sources/TokiUsageReaders/HermesUsageAttribution.swift @@ -106,8 +106,23 @@ private func hermesUsageEventParts( let residual = counters.subtracting(combinedCounters) if residual.totalTokens > 0 { - parts.append(HermesUsageEventPart(model: nil, counters: residual)) - } else if includeCostOnlyResidual { + // The resolver leaves observation.model non-nil only when the session and + // every detail row agree on one model, so a residual there belongs to that + // model rather than to the unattributed bucket. Merge into the existing + // part instead of appending a second one for the same model, otherwise its + // pricing counters would be charged twice during cost allocation. + if let residualModel = observation.model, + let existingIndex = parts.firstIndex(where: { $0.model == residualModel }) { + parts[existingIndex] = HermesUsageEventPart( + model: residualModel, + counters: parts[existingIndex].counters.adding(residual)) + } else { + parts.append(HermesUsageEventPart(model: observation.model, counters: residual)) + } + } + // Reported cost that no model claims still needs an unattributed carrier, even + // when the residual tokens themselves resolved to a model. + if includeCostOnlyResidual, !parts.contains(where: { $0.model == nil }) { parts.append(HermesUsageEventPart(model: nil, counters: .zero)) } if parts.isEmpty, includeCostOnlyFallback { diff --git a/TokiTests/HermesUsageLedgerTests.swift b/TokiTests/HermesUsageLedgerTests.swift index e6eb5d4..59ebae9 100644 --- a/TokiTests/HermesUsageLedgerTests.swift +++ b/TokiTests/HermesUsageLedgerTests.swift @@ -118,6 +118,58 @@ final class HermesUsageLedgerTests: XCTestCase { XCTAssertEqual(events.reduce(0) { $0 + $1.counters.totalTokens }, 150) } + func test_hermesUsageLedger_attributesPartialDetailResidualToResolvedModel() async throws { + let tempDir = try makeHermesTemporaryDirectory() + defer { try? FileManager.default.removeItem(at: tempDir) } + let ledger = HermesUsageLedger(fileURL: tempDir.appendingPathComponent("hermes-usage-ledger.json")) + let baselineAt = tokiTestISODate("2026-04-10T08:00:00Z") + let startedAt = tokiTestISODate("2026-04-10T09:00:00Z") + let observedAt = tokiTestISODate("2026-04-10T10:00:00Z") + try await ledger.refresh(observations: [], observedAt: baselineAt) + + // The session counters exceed the partial session_model_usage breakdown, but + // the session and the detail row agree on one model, so the resolver keeps + // that model on the observation. The 50-token residual and its share of the + // cost belong to that model rather than to Mixed / Unattributed. + try await ledger.refresh( + observations: [HermesSessionObservation( + sessionID: "single-model-session", + startedAt: startedAt, + earliestActivityAt: startedAt, + latestActivityAt: observedAt, + model: "kr/claude-opus-5", + counters: HermesTokenCounters( + inputTokens: 150, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + modelCounters: [ + "kr/claude-opus-5": HermesTokenCounters( + inputTokens: 100, + outputTokens: 0, + cacheReadTokens: 0, + cacheWriteTokens: 0, + reasoningTokens: 0), + ], + cost: 3, + projectName: nil, + attributionQuality: .unknown)], + observedAt: observedAt) + + let events = try await ledger.events( + from: baselineAt, + to: observedAt.addingTimeInterval(1)) + let tokensByModel = events.reduce(into: [String: Int]()) { result, event in + result[event.model ?? "Mixed / Unattributed", default: 0] += event.counters.totalTokens + } + + XCTAssertEqual(tokensByModel, ["kr/claude-opus-5": 150]) + XCTAssertFalse(events.contains { $0.model == nil }) + XCTAssertEqual(events.reduce(0) { $0 + $1.counters.totalTokens }, 150) + XCTAssertEqual(events.reduce(0) { $0 + $1.cost }, 3, accuracy: 0.0001) + } + func test_hermesUsageLedger_establishesModelBaselineBeforeSplittingLegacySessionDeltas() async throws { let tempDir = try makeHermesTemporaryDirectory() defer { try? FileManager.default.removeItem(at: tempDir) }