Skip to content

Read OpenCode Go local token counts #3995

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Yuxin-Qiao wants to merge 1 commit into from
Closed

Read OpenCode Go local token counts #3995

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
154 changes: 137 additions & 17 deletions Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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'
Expand All @@ -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'
Expand All @@ -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 (
Expand All @@ -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 {
Expand Down Expand Up @@ -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)
}
Expand Down
121 changes: 118 additions & 3 deletions Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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 }
Expand All @@ -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) ?? "{}"

Expand Down Expand Up @@ -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 }
Expand All @@ -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) ?? "{}"
Expand Down
Loading
Loading