From 003bc4187c537d083ab04a9ae2453c1495104e91 Mon Sep 17 00:00:00 2001 From: Yuxin Qiao <104957188+Yuxin-Qiao@users.noreply.github.com> Date: Sat, 26 Sep 2026 00:44:02 +0800 Subject: [PATCH] Read OpenCode Go local token counts Co-authored-by: Cursor --- .../OpenCodeGoLocalUsageReader.swift | 154 ++++++++++++++++-- .../OpenCodeGoLocalUsageReaderTests.swift | 121 +++++++++++++- .../SpendActivityHeatmapTests.swift | 35 ++++ 3 files changed, 290 insertions(+), 20 deletions(-) diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift index f3e29c4e73..3185d72993 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift @@ -123,11 +123,47 @@ public struct OpenCodeGoLocalUsageReader: Sendable { guard createdMs > 0, cost >= 0, cost.isFinite else { continue } let requestCount = max(1, Int(sqlite3_column_int64(stmt, 2))) let model = sqlite3_column_text(stmt, 3).map { String(cString: $0) } ?? "" - rows.append(UsageRow(createdMs: createdMs, cost: cost, requestCount: requestCount, model: model)) + rows.append(UsageRow( + createdMs: createdMs, + cost: cost, + requestCount: requestCount, + model: model, + tokens: Self.tokenCounts(stmt: stmt, firstColumn: 4))) } return rows } + /// Reads the `hasTokens, input, output, reasoning, cacheRead, cacheWrite, total` column run. + /// OpenCode's `total` already includes reasoning and cache tokens; older rows omit it. + private static func tokenCounts(stmt: OpaquePointer?, firstColumn: Int32) -> TokenCounts? { + guard sqlite3_column_int(stmt, firstColumn) != 0 else { return nil } + func count(_ offset: Int32) -> Int? { + let column = firstColumn + offset + guard sqlite3_column_type(stmt, column) != SQLITE_NULL else { return nil } + return Int(sqlite3_column_int64(stmt, column)) + } + let components = [count(1), count(2), count(3), count(4), count(5)] + guard components.allSatisfy({ ($0 ?? 0) >= 0 }) else { return nil } + let input = components[0] ?? 0 + let output = components[1] ?? 0 + let reasoning = components[2] ?? 0 + let cacheRead = components[3] ?? 0 + let cacheWrite = components[4] ?? 0 + let total: Int? = if let explicitTotal = count(6) { + explicitTotal >= 0 ? explicitTotal : nil + } else { + TokenCounts.sum([input, output, reasoning, cacheRead, cacheWrite]) + } + guard let total else { return nil } + return TokenCounts( + input: input, + output: output, + reasoning: reasoning, + cacheRead: cacheRead, + cacheWrite: cacheWrite, + total: total) + } + private var walSidecarsAreMissing: Bool { !FileManager.default.fileExists(atPath: self.databaseURL.path + "-wal") && !FileManager.default.fileExists(atPath: self.databaseURL.path + "-shm") @@ -162,12 +198,28 @@ public struct OpenCodeGoLocalUsageReader: Sendable { message: db.map { String(cString: sqlite3_errmsg($0)) } ?? "unknown error") } + private static func tokenColumnsSQL(_ data: String) -> String { + """ + json_type(\(data), '$.tokens') = 'object' AS hasTokens, + CAST(json_extract(\(data), '$.tokens.input') AS INTEGER) AS inputTokens, + CAST(json_extract(\(data), '$.tokens.output') AS INTEGER) AS outputTokens, + CAST(json_extract(\(data), '$.tokens.reasoning') AS INTEGER) AS reasoningTokens, + CAST(json_extract(\(data), '$.tokens.cache.read') AS INTEGER) AS cacheReadTokens, + CAST(json_extract(\(data), '$.tokens.cache.write') AS INTEGER) AS cacheWriteTokens, + CAST(json_extract(\(data), '$.tokens.total') AS INTEGER) AS totalTokens + """ + } + + private static let tokenColumnNames = + "hasTokens, inputTokens, outputTokens, reasoningTokens, cacheReadTokens, cacheWriteTokens, totalTokens" + private static let messageUsageSQL = """ SELECT CAST(COALESCE(json_extract(data, '$.time.created'), time_created) AS INTEGER) AS createdMs, CAST(json_extract(data, '$.cost') AS REAL) AS cost, 1 AS requestCount, - COALESCE(json_extract(data, '$.modelID'), '') AS modelID + COALESCE(json_extract(data, '$.modelID'), '') AS modelID, + \(Self.tokenColumnsSQL("data")) FROM message WHERE json_valid(data) AND json_extract(data, '$.providerID') = 'opencode-go' @@ -182,7 +234,8 @@ public struct OpenCodeGoLocalUsageReader: Sendable { CAST(COALESCE(json_extract(data, '$.time.created'), time_created) AS INTEGER) AS createdMs, CAST(json_extract(data, '$.cost') AS REAL) AS cost, json_type(data, '$.cost') IN ('integer', 'real') AS hasCost, - COALESCE(json_extract(data, '$.modelID'), '') AS modelID + COALESCE(json_extract(data, '$.modelID'), '') AS modelID, + \(Self.tokenColumnsSQL("data")) FROM message WHERE json_valid(data) AND json_extract(data, '$.providerID') = 'opencode-go' @@ -193,14 +246,15 @@ public struct OpenCodeGoLocalUsageReader: Sendable { AS createdMs, CAST(json_extract(p.data, '$.cost') AS REAL) AS cost, 1 AS requestCount, - m.modelID AS modelID + m.modelID AS modelID, + \(Self.tokenColumnsSQL("p.data")) FROM part p JOIN provider_messages m ON m.messageID = p.message_id WHERE json_valid(p.data) AND json_extract(p.data, '$.type') = 'step-finish' AND json_type(p.data, '$.cost') IN ('integer', 'real') UNION ALL - SELECT createdMs, cost, 1 AS requestCount, modelID + SELECT createdMs, cost, 1 AS requestCount, modelID, \(Self.tokenColumnNames) FROM provider_messages m WHERE hasCost AND NOT EXISTS ( @@ -220,6 +274,60 @@ public struct OpenCodeGoLocalUsageReader: Sendable { let requestCount: Int /// The underlying model behind the `opencode-go` Zen proxy; empty when unattributed. let model: String + /// Nil when the row carries no readable `tokens` object. + let tokens: TokenCounts? + } + + private struct TokenCounts { + static let zero = TokenCounts(input: 0, output: 0, reasoning: 0, cacheRead: 0, cacheWrite: 0, total: 0) + + let input: Int + let output: Int + let reasoning: Int + let cacheRead: Int + let cacheWrite: Int + let total: Int + + static func sum(_ values: [Int]) -> Int? { + var result = 0 + for value in values { + let addition = result.addingReportingOverflow(value) + guard !addition.overflow else { return nil } + result = addition.partialValue + } + return result + } + + func adding(_ other: TokenCounts) -> TokenCounts? { + guard let input = Self.sum([self.input, other.input]), + let output = Self.sum([self.output, other.output]), + let reasoning = Self.sum([self.reasoning, other.reasoning]), + let cacheRead = Self.sum([self.cacheRead, other.cacheRead]), + let cacheWrite = Self.sum([self.cacheWrite, other.cacheWrite]), + let total = Self.sum([self.total, other.total]) + else { return nil } + return TokenCounts( + input: input, + output: output, + reasoning: reasoning, + cacheRead: cacheRead, + cacheWrite: cacheWrite, + total: total) + } + } + + /// Per-model day bucket. `tokens` turns nil once any contributing row lacks token data, so a + /// partially known day is never reported as a complete total. + private struct DailyBucket { + var cost: Double = 0 + var requestCount = 0 + var tokens: TokenCounts? = .zero + + mutating func add(_ row: UsageRow) { + self.cost += row.cost + self.requestCount += row.requestCount + self.tokens = self.tokens.flatMap { current in row.tokens.flatMap(current.adding) } + } } private struct SQLiteReadFailure: Error { @@ -322,35 +430,47 @@ public struct OpenCodeGoLocalUsageReader: Sendable { } let sinceStartOfDay = calendar.startOfDay(for: since) - var totalsByModel: [String: [String: (cost: Double, requestCount: Int)]] = [:] + var totalsByModel: [String: [String: DailyBucket]] = [:] for row in rows { let date = Date(timeIntervalSince1970: TimeInterval(row.createdMs) / 1000) guard date >= sinceStartOfDay, date <= now else { continue } let key = CostUsageScanner.CostUsageDayRange.dayKey(from: date) let trimmedModel = row.model.trimmingCharacters(in: .whitespacesAndNewlines) let model = trimmedModel.isEmpty ? Self.unknownModelName : trimmedModel - totalsByModel[key, default: [:]][model, default: (0, 0)].cost += row.cost - totalsByModel[key, default: [:]][model, default: (0, 0)].requestCount += row.requestCount + totalsByModel[key, default: [:]][model, default: DailyBucket()].add(row) } return totalsByModel.keys.sorted().compactMap { key in guard let dayTotals = totalsByModel[key] else { return nil } let modelBreakdowns = dayTotals.keys.sorted().map { model in - let bucket = dayTotals[model] ?? (cost: 0, requestCount: 0) + let bucket = dayTotals[model] ?? DailyBucket() return CostUsageDailyReport.ModelBreakdown( modelName: model, costUSD: bucket.cost, - requestCount: bucket.requestCount) + totalTokens: bucket.tokens?.total, + requestCount: bucket.requestCount, + inputTokens: bucket.tokens?.input, + outputTokens: bucket.tokens?.output, + cacheReadTokens: bucket.tokens?.cacheRead, + cacheCreationTokens: bucket.tokens?.cacheWrite, + reasoningTokens: bucket.tokens?.reasoning) }.sorted { ($0.costUSD ?? 0) > ($1.costUSD ?? 0) } - let totalCost = dayTotals.values.reduce(0) { $0 + $1.cost } - let totalRequests = dayTotals.values.reduce(0) { $0 + $1.requestCount } + var day = DailyBucket() + for bucket in dayTotals.values { + day.cost += bucket.cost + day.requestCount += bucket.requestCount + day.tokens = day.tokens.flatMap { current in bucket.tokens.flatMap(current.adding) } + } return CostUsageDailyReport.Entry( date: key, - inputTokens: nil, - outputTokens: nil, - totalTokens: nil, - requestCount: totalRequests, - costUSD: totalCost, + inputTokens: day.tokens?.input, + outputTokens: day.tokens?.output, + cacheReadTokens: day.tokens?.cacheRead, + cacheCreationTokens: day.tokens?.cacheWrite, + reasoningTokens: day.tokens?.reasoning, + totalTokens: day.tokens?.total, + requestCount: day.requestCount, + costUSD: day.cost, modelsUsed: dayTotals.keys.sorted(), modelBreakdowns: modelBreakdowns) } diff --git a/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift index 64f4bc0e46..086a50e2a7 100644 --- a/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift +++ b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift @@ -401,6 +401,116 @@ struct OpenCodeGoLocalUsageReaderTests { #expect(breakdowns.first?.requestCount == 2) } + @Test + func `daily entries carry message token counts per day and model`() throws { + let env = try Self.makeEnvironment() + defer { try? FileManager.default.removeItem(at: env.root) } + + try Self.writeAuth(to: env.authURL) + try Self.createDatabase(at: env.databaseURL) + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-06T11:00:00.000Z"), + cost: 3.0, + model: "claude-sonnet-4-5", + tokens: [ + "total": 1600, + "input": 100, + "output": 20, + "reasoning": 30, + "cache": ["read": 1400, "write": 50], + ]) + // Older OpenCode rows omit `total`; it is the sum of every component. + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-06T12:00:00.000Z"), + cost: 2.0, + model: "gpt-5.1-codex", + tokens: ["input": 10, "output": 5, "reasoning": 1, "cache": ["read": 4, "write": 0]]) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let now = Date(timeIntervalSince1970: TimeInterval(Self.ms("2026-03-06T15:00:00.000Z")) / 1000) + let snapshot = try reader.fetch(now: now, historyDays: 30) + + let entry = try #require(snapshot.daily.first) + #expect(entry.totalTokens == 1620) + #expect(entry.inputTokens == 110) + #expect(entry.outputTokens == 25) + #expect(entry.reasoningTokens == 31) + #expect(entry.cacheReadTokens == 1404) + #expect(entry.cacheCreationTokens == 50) + let breakdowns = try #require(entry.modelBreakdowns) + #expect(breakdowns.first { $0.modelName == "claude-sonnet-4-5" }?.totalTokens == 1600) + #expect(breakdowns.first { $0.modelName == "gpt-5.1-codex" }?.totalTokens == 20) + + let tokenSnapshot = snapshot.toCostUsageTokenSnapshot(historyDays: 30) + #expect(tokenSnapshot.last30DaysTokens == 1620) + } + + @Test + func `step finish tokens replace message tokens when parts exist`() throws { + let env = try Self.makeEnvironment() + defer { try? FileManager.default.removeItem(at: env.root) } + + try Self.writeAuth(to: env.authURL) + try Self.createDatabase(at: env.databaseURL) + let messageID = try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-06T11:00:00.000Z"), + cost: 3.0, + tokens: ["total": 999, "input": 999]) + try Self.insertStepFinishPart( + databaseURL: env.databaseURL, + messageID: messageID, + createdMs: Self.ms("2026-03-06T11:00:00.000Z"), + cost: 1.0, + tokens: ["total": 40, "input": 30, "output": 10]) + try Self.insertStepFinishPart( + databaseURL: env.databaseURL, + messageID: messageID, + createdMs: Self.ms("2026-03-06T11:05:00.000Z"), + cost: 2.0, + tokens: ["total": 60, "input": 50, "output": 10]) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + + let entry = try #require(snapshot.daily.first) + #expect(entry.totalTokens == 100) + #expect(entry.inputTokens == 80) + #expect(entry.outputTokens == 20) + } + + @Test + func `a day with a tokenless row keeps its token total unknown`() throws { + let env = try Self.makeEnvironment() + defer { try? FileManager.default.removeItem(at: env.root) } + + try Self.writeAuth(to: env.authURL) + try Self.createDatabase(at: env.databaseURL) + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-06T11:00:00.000Z"), + cost: 3.0, + model: "claude-sonnet-4-5", + tokens: ["total": 50, "input": 50]) + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-06T12:00:00.000Z"), + cost: 1.0, + model: "claude-sonnet-4-5") + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let now = Date(timeIntervalSince1970: TimeInterval(Self.ms("2026-03-06T15:00:00.000Z")) / 1000) + let snapshot = try reader.fetch(now: now, historyDays: 30) + + let entry = try #require(snapshot.daily.first) + #expect(entry.costUSD == 4.0) + #expect(entry.totalTokens == nil) + #expect(entry.modelBreakdowns?.first?.totalTokens == nil) + #expect(snapshot.toCostUsageTokenSnapshot(historyDays: 30).last30DaysTokens == nil) + } + @Test func `missing auth and history is not detected`() throws { let env = try Self.makeEnvironment() @@ -481,7 +591,8 @@ struct OpenCodeGoLocalUsageReaderTests { databaseURL: URL, createdMs: Int64, cost: Double?, - model: String? = nil) throws -> String + model: String? = nil, + tokens: [String: Any]? = nil) throws -> String { var db: OpaquePointer? guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } @@ -499,6 +610,9 @@ struct OpenCodeGoLocalUsageReaderTests { if let model { payload["modelID"] = model } + if let tokens { + payload["tokens"] = tokens + } let data = try JSONSerialization.data(withJSONObject: payload) let json = String(data: data, encoding: .utf8) ?? "{}" @@ -526,7 +640,8 @@ struct OpenCodeGoLocalUsageReaderTests { databaseURL: URL, messageID: String, createdMs: Int64, - cost: Double) throws + cost: Double, + tokens: [String: Any] = ["input": 1, "output": 1, "total": 2]) throws { var db: OpaquePointer? guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } @@ -535,7 +650,7 @@ struct OpenCodeGoLocalUsageReaderTests { let payload: [String: Any] = [ "type": "step-finish", "cost": cost, - "tokens": ["input": 1, "output": 1, "total": 2], + "tokens": tokens, ] let data = try JSONSerialization.data(withJSONObject: payload) let json = String(data: data, encoding: .utf8) ?? "{}" diff --git a/Tests/CodexBarTests/SpendActivityHeatmapTests.swift b/Tests/CodexBarTests/SpendActivityHeatmapTests.swift index 0e8736c599..8f5e52c229 100644 --- a/Tests/CodexBarTests/SpendActivityHeatmapTests.swift +++ b/Tests/CodexBarTests/SpendActivityHeatmapTests.swift @@ -217,6 +217,41 @@ struct SpendActivityHeatmapTests { #expect(point?.isScanned == true) } + @Test + func `OpenCode Go token history keeps shared activity days covered`() throws { + let now = try #require(Self.calendar.date(from: DateComponents(year: 2026, month: 7, day: 16, hour: 12))) + let today = Self.calendar.startOfDay(for: now) + let yesterday = try #require(Self.calendar.date(byAdding: .day, value: -1, to: today)) + let openCodeGo = OpenCodeGoUsageSnapshot( + hasMonthlyUsage: true, + rollingUsagePercent: 0, + weeklyUsagePercent: 0, + monthlyUsagePercent: 0, + rollingResetInSec: 0, + weeklyResetInSec: 0, + monthlyResetInSec: 0, + daily: [Self.entry(day: "2026-07-16", cost: 2, tokens: 60)], + updatedAt: now).toCostUsageTokenSnapshot(historyDays: 30) + let claude = Self.snapshot( + entries: [ + Self.entry(day: "2026-07-15", cost: 1, tokens: 15), + Self.entry(day: "2026-07-16", cost: 1, tokens: 40), + ], + historyDays: 30, + last30DaysTokens: 55) + let model = SpendDashboardModel.build( + inputs: [ + .init(id: "claude", provider: .claude, displayName: "Claude", snapshot: claude), + .init(id: "opencodego", provider: .opencodego, displayName: "OpenCode Go", snapshot: openCodeGo), + ], + requestedDays: 30, + now: now, + calendar: Self.calendar) + + #expect(model.tokenActivity.first { $0.day == today }?.totalTokens == 100) + #expect(model.tokenActivity.first { $0.day == yesterday }?.totalTokens == 15) + } + @Test func `covered empty history is zero while unestablished history remains unavailable`() throws { let now = try #require(Self.calendar.date(from: DateComponents(year: 2026, month: 7, day: 16)))