From 751e680c6eadcd982a75570c9740c0f02c1bb222 Mon Sep 17 00:00:00 2001 From: sopenlaz0 Date: Mon, 18 May 2026 16:22:31 +0700 Subject: [PATCH 1/2] Fix OpenCode Go local usage --- .../OpenCodeGoLocalUsageReader.swift | 259 ++++++++++++++++++ .../OpenCodeGoProviderDescriptor.swift | 32 ++- .../OpenCodeGoLocalUsageReaderTests.swift | 166 +++++++++++ 3 files changed, 456 insertions(+), 1 deletion(-) create mode 100644 Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift create mode 100644 Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift new file mode 100644 index 000000000..bf75ca1b9 --- /dev/null +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift @@ -0,0 +1,259 @@ +import Foundation + +#if os(macOS) +import SQLite3 + +public enum OpenCodeGoLocalUsageError: LocalizedError, Sendable, Equatable { + case notDetected + case sqliteFailed(String) + + public var errorDescription: String? { + switch self { + case .notDetected: + "OpenCode Go not detected. Log in with OpenCode Go or use it locally first." + case let .sqliteFailed(message): + "SQLite error reading OpenCode Go usage: \(message)" + } + } +} + +public struct OpenCodeGoLocalUsageReader: Sendable { + private static let fiveHours: TimeInterval = 5 * 60 * 60 + private static let week: TimeInterval = 7 * 24 * 60 * 60 + private static let limits = (session: 12.0, weekly: 30.0, monthly: 60.0) + + private let authURL: URL + private let databaseURL: URL + + public init(homeDirectory: URL = FileManager.default.homeDirectoryForCurrentUser) { + let openCodeDirectory = homeDirectory + .appendingPathComponent(".local", isDirectory: true) + .appendingPathComponent("share", isDirectory: true) + .appendingPathComponent("opencode", isDirectory: true) + self.authURL = openCodeDirectory.appendingPathComponent("auth.json", isDirectory: false) + self.databaseURL = openCodeDirectory.appendingPathComponent("opencode.db", isDirectory: false) + } + + public init(authURL: URL, databaseURL: URL) { + self.authURL = authURL + self.databaseURL = databaseURL + } + + public func fetch(now: Date = Date()) throws -> OpenCodeGoUsageSnapshot { + let hasAuth = Self.hasAuthKey(at: self.authURL) + guard FileManager.default.fileExists(atPath: self.databaseURL.path) else { + if hasAuth { + return Self.emptySnapshot(now: now) + } + throw OpenCodeGoLocalUsageError.notDetected + } + + let rows: [UsageRow] + do { + rows = try self.readRows() + } catch let error as OpenCodeGoLocalUsageError { + if hasAuth { + return Self.emptySnapshot(now: now) + } + throw error + } + guard hasAuth || !rows.isEmpty else { + throw OpenCodeGoLocalUsageError.notDetected + } + return Self.snapshot(rows: rows, now: now) + } + + private func readRows() throws -> [UsageRow] { + var db: OpaquePointer? + guard sqlite3_open_v2(self.databaseURL.path, &db, SQLITE_OPEN_READONLY, nil) == SQLITE_OK else { + let message = db.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "unknown error" + sqlite3_close(db) + throw OpenCodeGoLocalUsageError.sqliteFailed(message) + } + defer { sqlite3_close(db) } + sqlite3_busy_timeout(db, 250) + + let sql = """ + SELECT + CAST(COALESCE(json_extract(data, '$.time.created'), time_created) AS INTEGER) AS createdMs, + CAST(json_extract(data, '$.cost') AS REAL) AS cost + FROM message + WHERE json_valid(data) + AND json_extract(data, '$.providerID') = 'opencode-go' + AND json_extract(data, '$.role') = 'assistant' + AND json_type(data, '$.cost') IN ('integer', 'real') + """ + + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else { + let message = db.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "unknown error" + throw OpenCodeGoLocalUsageError.sqliteFailed(message) + } + defer { sqlite3_finalize(stmt) } + + var rows: [UsageRow] = [] + while true { + let step = sqlite3_step(stmt) + if step == SQLITE_DONE { break } + guard step == SQLITE_ROW else { + let message = db.flatMap { String(cString: sqlite3_errmsg($0)) } ?? "unknown error" + throw OpenCodeGoLocalUsageError.sqliteFailed(message) + } + + let createdMs = sqlite3_column_int64(stmt, 0) + let cost = sqlite3_column_double(stmt, 1) + guard createdMs > 0, cost >= 0, cost.isFinite else { continue } + rows.append(UsageRow(createdMs: createdMs, cost: cost)) + } + return rows + } + + private struct UsageRow { + let createdMs: Int64 + let cost: Double + } + + private static func hasAuthKey(at url: URL) -> Bool { + guard let data = try? Data(contentsOf: url), + let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any], + let entry = object["opencode-go"] as? [String: Any], + let key = entry["key"] as? String + else { + return false + } + return !key.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty + } + + private static func snapshot(rows: [UsageRow], now: Date) -> OpenCodeGoUsageSnapshot { + let nowMs = Int64(now.timeIntervalSince1970 * 1000) + let sessionStart = nowMs - Int64(Self.fiveHours * 1000) + let weekStart = self.startOfUTCWeek(now: now).timeIntervalSince1970 * 1000 + let weekStartMs = Int64(weekStart) + let weekEndMs = weekStartMs + Int64(Self.week * 1000) + let earliestMs = rows.map(\.createdMs).min() + let monthBounds = self.monthBounds(now: now, anchorMs: earliestMs) + + let sessionCost = self.sum(rows: rows, startMs: sessionStart, endMs: nowMs) + let weeklyCost = self.sum(rows: rows, startMs: weekStartMs, endMs: weekEndMs) + let monthlyCost = self.sum(rows: rows, startMs: monthBounds.startMs, endMs: monthBounds.endMs) + + return OpenCodeGoUsageSnapshot( + hasMonthlyUsage: true, + rollingUsagePercent: self.percent(used: sessionCost, limit: self.limits.session), + weeklyUsagePercent: self.percent(used: weeklyCost, limit: self.limits.weekly), + monthlyUsagePercent: self.percent(used: monthlyCost, limit: self.limits.monthly), + rollingResetInSec: self.rollingReset(rows: rows, nowMs: nowMs), + weeklyResetInSec: max(0, Int((weekEndMs - nowMs) / 1000)), + monthlyResetInSec: max(0, Int((monthBounds.endMs - nowMs) / 1000)), + updatedAt: now) + } + + private static func emptySnapshot(now: Date) -> OpenCodeGoUsageSnapshot { + self.snapshot(rows: [], now: now) + } + + private static func sum(rows: [UsageRow], startMs: Int64, endMs: Int64) -> Double { + rows.reduce(0) { total, row in + guard row.createdMs >= startMs, row.createdMs < endMs else { return total } + return total + row.cost + } + } + + private static func percent(used: Double, limit: Double) -> Double { + guard used.isFinite, limit > 0 else { return 0 } + let value = max(0, min(100, used / limit * 100)) + return (value * 10).rounded() / 10 + } + + private static func rollingReset(rows: [UsageRow], nowMs: Int64) -> Int { + let sessionStart = nowMs - Int64(Self.fiveHours * 1000) + let oldest = rows + .filter { $0.createdMs >= sessionStart && $0.createdMs < nowMs } + .map(\.createdMs) + .min() ?? nowMs + return max(0, Int((oldest + Int64(Self.fiveHours * 1000) - nowMs) / 1000)) + } + + private static func startOfUTCWeek(now: Date) -> Date { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(secondsFromGMT: 0) ?? TimeZone.current + calendar.firstWeekday = 2 + calendar.minimumDaysInFirstWeek = 4 + let components = calendar.dateComponents([.yearForWeekOfYear, .weekOfYear], from: now) + return calendar.date(from: components) ?? now + } + + private static func monthBounds(now: Date, anchorMs: Int64?) -> (startMs: Int64, endMs: Int64) { + var calendar = Calendar(identifier: .gregorian) + calendar.timeZone = TimeZone(secondsFromGMT: 0) ?? TimeZone.current + + guard let anchorMs else { + let start = calendar.date(from: calendar.dateComponents([.year, .month], from: now)) ?? now + let end = calendar.date(byAdding: .month, value: 1, to: start) ?? start + return (Int64(start.timeIntervalSince1970 * 1000), Int64(end.timeIntervalSince1970 * 1000)) + } + + let anchor = Date(timeIntervalSince1970: TimeInterval(anchorMs) / 1000) + let anchorComponents = calendar.dateComponents([.day, .hour, .minute, .second, .nanosecond], from: anchor) + var nowComponents = calendar.dateComponents([.year, .month], from: now) + + var start = self.anchoredMonth(calendar: calendar, month: nowComponents, anchor: anchorComponents) + if start > now { + guard let previous = calendar.date(byAdding: .month, value: -1, to: start) else { + let end = calendar.date(byAdding: .month, value: 1, to: start) ?? start + return (Int64(start.timeIntervalSince1970 * 1000), Int64(end.timeIntervalSince1970 * 1000)) + } + nowComponents = calendar.dateComponents([.year, .month], from: previous) + start = self.anchoredMonth(calendar: calendar, month: nowComponents, anchor: anchorComponents) + } + let end = calendar.date(byAdding: .month, value: 1, to: start) ?? start + return (Int64(start.timeIntervalSince1970 * 1000), Int64(end.timeIntervalSince1970 * 1000)) + } + + private static func anchoredMonth( + calendar: Calendar, + month: DateComponents, + anchor: DateComponents) -> Date + { + var components = DateComponents() + components.calendar = calendar + components.timeZone = calendar.timeZone + components.year = month.year + components.month = month.month + components.day = anchor.day + components.hour = anchor.hour + components.minute = anchor.minute + components.second = anchor.second + components.nanosecond = anchor.nanosecond + + if let date = calendar.date(from: components), + calendar.component(.month, from: date) == month.month + { + return date + } + + components.day = calendar.range(of: .day, in: .month, for: calendar.date(from: month) ?? Date())?.count + return calendar.date(from: components) ?? Date() + } +} + +#else + +public enum OpenCodeGoLocalUsageError: LocalizedError, Sendable, Equatable { + case notSupported + + public var errorDescription: String? { + "OpenCode Go local usage is only supported on macOS." + } +} + +public struct OpenCodeGoLocalUsageReader: Sendable { + public init(homeDirectory _: URL = FileManager.default.homeDirectoryForCurrentUser) {} + public init(authURL _: URL, databaseURL _: URL) {} + + public func fetch(now _: Date = Date()) throws -> OpenCodeGoUsageSnapshot { + throw OpenCodeGoLocalUsageError.notSupported + } +} + +#endif diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift index c5ed2a77f..28d60ca00 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift @@ -33,11 +33,41 @@ public enum OpenCodeGoProviderDescriptor { noDataMessage: { "OpenCode Go cost summary is not supported." }), fetchPlan: ProviderFetchPlan( sourceModes: [.auto, .web], - pipeline: ProviderFetchPipeline(resolveStrategies: { _ in [OpenCodeGoUsageFetchStrategy()] })), + pipeline: ProviderFetchPipeline(resolveStrategies: self.resolveStrategies)), cli: ProviderCLIConfig( name: "opencodego", versionDetector: nil)) } + + private static func resolveStrategies(context: ProviderFetchContext) async -> [any ProviderFetchStrategy] { + if context.sourceMode == .web { + return [OpenCodeGoUsageFetchStrategy()] + } + return [ + OpenCodeGoLocalUsageFetchStrategy(), + OpenCodeGoUsageFetchStrategy(), + ] + } +} + +struct OpenCodeGoLocalUsageFetchStrategy: ProviderFetchStrategy { + let id: String = "opencodego.local" + let kind: ProviderFetchKind = .localProbe + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try OpenCodeGoLocalUsageReader().fetch() + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "local") + } + + func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool { + error is OpenCodeGoLocalUsageError + } } struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy { diff --git a/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift new file mode 100644 index 000000000..86ba2ae1f --- /dev/null +++ b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift @@ -0,0 +1,166 @@ +#if os(macOS) + +import Foundation +import SQLite3 +import Testing +@testable import CodexBarCore + +struct OpenCodeGoLocalUsageReaderTests { + @Test + func `reads local OpenCode Go history into usage windows`() 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) + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-05T12:00:00.000Z"), + cost: 6.0) + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-02-25T07:53:16.000Z"), + cost: 2.0) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + + #expect(snapshot.rollingUsagePercent == 25) + #expect(snapshot.weeklyUsagePercent == 30) + #expect(snapshot.monthlyUsagePercent == 18.3) + #expect(snapshot.rollingResetInSec == 14400) + #expect(snapshot.weeklyResetInSec == 216_000) + #expect(snapshot.monthlyResetInSec == 1_626_796) + } + + @Test + func `auth without history returns zeroed bars`() throws { + let env = try Self.makeEnvironment() + defer { try? FileManager.default.removeItem(at: env.root) } + + try Self.writeAuth(to: env.authURL) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + + #expect(snapshot.rollingUsagePercent == 0) + #expect(snapshot.weeklyUsagePercent == 0) + #expect(snapshot.monthlyUsagePercent == 0) + } + + @Test + func `auth with unreadable history returns zeroed bars`() throws { + let env = try Self.makeEnvironment() + defer { try? FileManager.default.removeItem(at: env.root) } + + try Self.writeAuth(to: env.authURL) + var db: OpaquePointer? + guard sqlite3_open(env.databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } + sqlite3_close(db) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + + #expect(snapshot.rollingUsagePercent == 0) + #expect(snapshot.weeklyUsagePercent == 0) + #expect(snapshot.monthlyUsagePercent == 0) + } + + @Test + func `missing auth and history is not detected`() throws { + let env = try Self.makeEnvironment() + defer { try? FileManager.default.removeItem(at: env.root) } + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + + #expect(throws: OpenCodeGoLocalUsageError.notDetected) { + _ = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + } + } + + private static func makeEnvironment() throws -> (root: URL, authURL: URL, databaseURL: URL) { + let root = FileManager.default.temporaryDirectory + .appendingPathComponent("OpenCodeGoLocalUsageReaderTests-\(UUID().uuidString)", isDirectory: true) + let directory = root + .appendingPathComponent(".local", isDirectory: true) + .appendingPathComponent("share", isDirectory: true) + .appendingPathComponent("opencode", isDirectory: true) + try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true) + return ( + root, + directory.appendingPathComponent("auth.json", isDirectory: false), + directory.appendingPathComponent("opencode.db", isDirectory: false)) + } + + private static func writeAuth(to url: URL) throws { + let data = Data(#"{"opencode-go":{"type":"api-key","key":"go-key"}}"#.utf8) + try data.write(to: url) + } + + private static func createDatabase(at url: URL) throws { + var db: OpaquePointer? + guard sqlite3_open(url.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } + defer { sqlite3_close(db) } + try Self.exec( + db: db, + sql: """ + CREATE TABLE message ( + data TEXT NOT NULL, + time_created INTEGER + ); + """) + } + + private static func insertMessage(databaseURL: URL, createdMs: Int64, cost: Double) throws { + var db: OpaquePointer? + guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } + defer { sqlite3_close(db) } + + let payload: [String: Any] = [ + "providerID": "opencode-go", + "role": "assistant", + "time": ["created": createdMs], + "cost": cost, + ] + let data = try JSONSerialization.data(withJSONObject: payload) + let json = String(data: data, encoding: .utf8) ?? "{}" + + var stmt: OpaquePointer? + guard sqlite3_prepare_v2(db, "INSERT INTO message (data, time_created) VALUES (?, ?)", -1, &stmt, nil) + == SQLITE_OK + else { throw SQLiteTestError.prepare } + defer { sqlite3_finalize(stmt) } + + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + sqlite3_bind_text(stmt, 1, json, -1, transient) + sqlite3_bind_int64(stmt, 2, createdMs) + guard sqlite3_step(stmt) == SQLITE_DONE else { throw SQLiteTestError.step } + } + + private static func exec(db: OpaquePointer?, sql: String) throws { + var message: UnsafeMutablePointer? + guard sqlite3_exec(db, sql, nil, nil, &message) == SQLITE_OK else { + sqlite3_free(message) + throw SQLiteTestError.exec + } + } + + private static func ms(_ iso: String) -> Int64 { + let formatter = ISO8601DateFormatter() + formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds] + return Int64((formatter.date(from: iso)?.timeIntervalSince1970 ?? 0) * 1000) + } + + private enum SQLiteTestError: Error { + case open + case prepare + case step + case exec + } +} + +#endif From ff8057a65ab82d35961628ee5969eedd010c98f6 Mon Sep 17 00:00:00 2001 From: sopenlaz0 Date: Mon, 18 May 2026 18:26:48 +0700 Subject: [PATCH 2/2] Fix OpenCode Go monthly anchor --- .../OpenCodeGoLocalUsageReader.swift | 121 +++++++++---- .../OpenCodeGoProviderDescriptor.swift | 49 +++++- .../OpenCodeGo/OpenCodeGoUsageFetcher.swift | 25 +++ .../OpenCodeGoLocalUsageReaderTests.swift | 165 ++++++++++++++++-- .../OpenCodeGoUsageFetcherErrorTests.swift | 28 +++ 5 files changed, 339 insertions(+), 49 deletions(-) diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift index bf75ca1b9..13098a445 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift @@ -5,12 +5,15 @@ import SQLite3 public enum OpenCodeGoLocalUsageError: LocalizedError, Sendable, Equatable { case notDetected + case historyUnavailable(String) case sqliteFailed(String) public var errorDescription: String? { switch self { case .notDetected: "OpenCode Go not detected. Log in with OpenCode Go or use it locally first." + case let .historyUnavailable(message): + "OpenCode Go local usage history is unavailable: \(message)" case let .sqliteFailed(message): "SQLite error reading OpenCode Go usage: \(message)" } @@ -43,23 +46,18 @@ public struct OpenCodeGoLocalUsageReader: Sendable { let hasAuth = Self.hasAuthKey(at: self.authURL) guard FileManager.default.fileExists(atPath: self.databaseURL.path) else { if hasAuth { - return Self.emptySnapshot(now: now) + throw OpenCodeGoLocalUsageError.historyUnavailable("database not found") } throw OpenCodeGoLocalUsageError.notDetected } - let rows: [UsageRow] - do { - rows = try self.readRows() - } catch let error as OpenCodeGoLocalUsageError { - if hasAuth { - return Self.emptySnapshot(now: now) - } - throw error - } + let rows = try self.readRows() guard hasAuth || !rows.isEmpty else { throw OpenCodeGoLocalUsageError.notDetected } + guard !rows.isEmpty else { + throw OpenCodeGoLocalUsageError.historyUnavailable("no local usage rows") + } return Self.snapshot(rows: rows, now: now) } @@ -73,16 +71,7 @@ public struct OpenCodeGoLocalUsageReader: Sendable { defer { sqlite3_close(db) } sqlite3_busy_timeout(db, 250) - let sql = """ - SELECT - CAST(COALESCE(json_extract(data, '$.time.created'), time_created) AS INTEGER) AS createdMs, - CAST(json_extract(data, '$.cost') AS REAL) AS cost - FROM message - WHERE json_valid(data) - AND json_extract(data, '$.providerID') = 'opencode-go' - AND json_extract(data, '$.role') = 'assistant' - AND json_type(data, '$.cost') IN ('integer', 'real') - """ + let sql = self.hasTable(named: "part", db: db) ? Self.messageAndPartUsageSQL : Self.messageUsageSQL var stmt: OpaquePointer? guard sqlite3_prepare_v2(db, sql, -1, &stmt, nil) == SQLITE_OK else { @@ -108,6 +97,69 @@ public struct OpenCodeGoLocalUsageReader: Sendable { return rows } + private func hasTable(named name: String, db: OpaquePointer?) -> Bool { + var stmt: OpaquePointer? + guard sqlite3_prepare_v2( + db, + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ? LIMIT 1", + -1, + &stmt, + nil) == SQLITE_OK + else { + return false + } + defer { sqlite3_finalize(stmt) } + + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + sqlite3_bind_text(stmt, 1, name, -1, transient) + return sqlite3_step(stmt) == SQLITE_ROW + } + + 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 + FROM message + WHERE json_valid(data) + AND json_extract(data, '$.providerID') = 'opencode-go' + AND json_extract(data, '$.role') = 'assistant' + AND json_type(data, '$.cost') IN ('integer', 'real') + """ + + private static let messageAndPartUsageSQL = """ + WITH message_costs AS ( + SELECT + id AS messageID, + CAST(COALESCE(json_extract(data, '$.time.created'), time_created) AS INTEGER) AS createdMs, + CAST(json_extract(data, '$.cost') AS REAL) AS cost + FROM message + WHERE json_valid(data) + AND json_extract(data, '$.providerID') = 'opencode-go' + AND json_extract(data, '$.role') = 'assistant' + AND json_type(data, '$.cost') IN ('integer', 'real') + ) + SELECT createdMs, cost + FROM message_costs + UNION ALL + SELECT + CAST(COALESCE(json_extract(p.data, '$.time.created'), p.time_created, m.time_created) AS INTEGER) + AS createdMs, + CAST(json_extract(p.data, '$.cost') AS REAL) AS cost + FROM part p + JOIN message m ON m.id = p.message_id + WHERE json_valid(p.data) + AND json_valid(m.data) + AND json_extract(p.data, '$.type') = 'step-finish' + AND json_type(p.data, '$.cost') IN ('integer', 'real') + AND json_extract(m.data, '$.providerID') = 'opencode-go' + AND json_extract(m.data, '$.role') = 'assistant' + AND NOT EXISTS ( + SELECT 1 + FROM message_costs + WHERE message_costs.messageID = p.message_id + ) + """ + private struct UsageRow { let createdMs: Int64 let cost: Double @@ -148,10 +200,6 @@ public struct OpenCodeGoLocalUsageReader: Sendable { updatedAt: now) } - private static func emptySnapshot(now: Date) -> OpenCodeGoUsageSnapshot { - self.snapshot(rows: [], now: now) - } - private static func sum(rows: [UsageRow], startMs: Int64, endMs: Int64) -> Double { rows.reduce(0) { total, row in guard row.createdMs >= startMs, row.createdMs < endMs else { return total } @@ -195,21 +243,34 @@ public struct OpenCodeGoLocalUsageReader: Sendable { let anchor = Date(timeIntervalSince1970: TimeInterval(anchorMs) / 1000) let anchorComponents = calendar.dateComponents([.day, .hour, .minute, .second, .nanosecond], from: anchor) - var nowComponents = calendar.dateComponents([.year, .month], from: now) + let nowComponents = calendar.dateComponents([.year, .month], from: now) - var start = self.anchoredMonth(calendar: calendar, month: nowComponents, anchor: anchorComponents) + var startMonthComponents = nowComponents + var start = self.anchoredMonth(calendar: calendar, month: startMonthComponents, anchor: anchorComponents) if start > now { guard let previous = calendar.date(byAdding: .month, value: -1, to: start) else { - let end = calendar.date(byAdding: .month, value: 1, to: start) ?? start + let end = self.anchoredMonth( + calendar: calendar, + month: self.monthComponents(after: startMonthComponents, calendar: calendar), + anchor: anchorComponents) return (Int64(start.timeIntervalSince1970 * 1000), Int64(end.timeIntervalSince1970 * 1000)) } - nowComponents = calendar.dateComponents([.year, .month], from: previous) - start = self.anchoredMonth(calendar: calendar, month: nowComponents, anchor: anchorComponents) + startMonthComponents = calendar.dateComponents([.year, .month], from: previous) + start = self.anchoredMonth(calendar: calendar, month: startMonthComponents, anchor: anchorComponents) } - let end = calendar.date(byAdding: .month, value: 1, to: start) ?? start + let end = self.anchoredMonth( + calendar: calendar, + month: self.monthComponents(after: startMonthComponents, calendar: calendar), + anchor: anchorComponents) return (Int64(start.timeIntervalSince1970 * 1000), Int64(end.timeIntervalSince1970 * 1000)) } + private static func monthComponents(after month: DateComponents, calendar: Calendar) -> DateComponents { + let monthStart = calendar.date(from: month) ?? Date() + let nextMonth = calendar.date(byAdding: .month, value: 1, to: monthStart) ?? monthStart + return calendar.dateComponents([.year, .month], from: nextMonth) + } + private static func anchoredMonth( calendar: Calendar, month: DateComponents, diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift index 28d60ca00..e16e63431 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift @@ -58,8 +58,8 @@ struct OpenCodeGoLocalUsageFetchStrategy: ProviderFetchStrategy { true } - func fetch(_: ProviderFetchContext) async throws -> ProviderFetchResult { - let snapshot = try OpenCodeGoLocalUsageReader().fetch() + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await self.snapshot(context: context) return self.makeResult( usage: snapshot.toUsageSnapshot(), sourceLabel: "local") @@ -68,6 +68,49 @@ struct OpenCodeGoLocalUsageFetchStrategy: ProviderFetchStrategy { func shouldFallback(on error: Error, context _: ProviderFetchContext) -> Bool { error is OpenCodeGoLocalUsageError } + + private func snapshot(context: ProviderFetchContext) async throws -> OpenCodeGoUsageSnapshot { + let snapshot = try OpenCodeGoLocalUsageReader().fetch() + guard context.includeOptionalUsage, + context.settings?.opencodego?.cookieSource != .off + else { + return snapshot + } + + guard let cookieHeader = Self.cachedOrManualCookieHeader(context: context) else { + return snapshot + } + + let workspaceOverride = context.settings?.opencodego?.workspaceID + ?? context.env["CODEXBAR_OPENCODEGO_WORKSPACE_ID"] + let zenBalanceTask = Task { + do { + return try await OpenCodeGoUsageFetcher.fetchOptionalZenBalance( + cookieHeader: cookieHeader, + timeout: context.webTimeout, + workspaceIDOverride: workspaceOverride) + } catch is CancellationError { + throw CancellationError() + } catch { + return nil + } + } + let zenBalance = try await OpenCodeGoUsageFetcher.completedOptionalZenBalance(from: zenBalanceTask) + return snapshot.withZenBalanceUSD(zenBalance) + } + + private static func cachedOrManualCookieHeader(context: ProviderFetchContext) -> String? { + if let settings = context.settings?.opencodego, settings.cookieSource == .manual { + return OpenCodeWebCookieSupport.requestCookieHeader(from: settings.manualCookieHeader) + } + + #if os(macOS) + guard let cached = CookieHeaderCache.load(provider: .opencodego) else { return nil } + return OpenCodeWebCookieSupport.requestCookieHeader(from: cached.cookieHeader) + #else + return nil + #endif + } } struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy { @@ -115,7 +158,7 @@ struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy { false } - private static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String { + static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String { try OpenCodeWebCookieSupport.resolveCookieHeader( context: OpenCodeWebCookieSupport.Context( settings: context.settings?.opencodego, diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoUsageFetcher.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoUsageFetcher.swift index f3cae18f4..940d71daf 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoUsageFetcher.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoUsageFetcher.swift @@ -147,6 +147,31 @@ public struct OpenCodeGoUsageFetcher: Sendable { return snapshot.withZenBalanceUSD(zenBalance) } + static func fetchOptionalZenBalance( + cookieHeader: String, + timeout: TimeInterval, + workspaceIDOverride: String? = nil, + session: URLSession? = nil) async throws -> Double? + { + let session = session ?? self.redirectGuardSession + guard let requestCookieHeader = OpenCodeWebCookieSupport.requestCookieHeader(from: cookieHeader) else { + throw OpenCodeGoUsageError.invalidCredentials + } + let workspaceID: String = if let override = self.normalizeWorkspaceID(workspaceIDOverride) { + override + } else { + try await self.fetchWorkspaceID( + cookieHeader: requestCookieHeader, + timeout: timeout, + session: session) + } + return try await self.fetchOptionalZenBalance( + workspaceID: workspaceID, + cookieHeader: requestCookieHeader, + timeout: min(timeout, self.optionalZenBalanceTimeout), + session: session) + } + static func allowsRedirect(from sourceURL: URL?, to destinationURL: URL?) -> Bool { guard let sourceHost = sourceURL?.host?.lowercased(), let destinationHost = destinationURL?.host?.lowercased(), diff --git a/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift index 86ba2ae1f..ca060b0a3 100644 --- a/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift +++ b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift @@ -38,22 +38,21 @@ struct OpenCodeGoLocalUsageReaderTests { } @Test - func `auth without history returns zeroed bars`() throws { + func `auth without history falls through to web strategy`() throws { let env = try Self.makeEnvironment() defer { try? FileManager.default.removeItem(at: env.root) } try Self.writeAuth(to: env.authURL) let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) - let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) - #expect(snapshot.rollingUsagePercent == 0) - #expect(snapshot.weeklyUsagePercent == 0) - #expect(snapshot.monthlyUsagePercent == 0) + #expect(throws: OpenCodeGoLocalUsageError.historyUnavailable("database not found")) { + _ = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + } } @Test - func `auth with unreadable history returns zeroed bars`() throws { + func `auth with unreadable history falls through to web strategy`() throws { let env = try Self.makeEnvironment() defer { try? FileManager.default.removeItem(at: env.root) } @@ -62,12 +61,85 @@ struct OpenCodeGoLocalUsageReaderTests { guard sqlite3_open(env.databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } sqlite3_close(db) + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + + #expect(throws: OpenCodeGoLocalUsageError.self) { + _ = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + } + } + + @Test + func `monthly window keeps original anchor after shorter month clamp`() 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-01-31T00:00:00.000Z"), + cost: 1.0) + try Self.insertMessage( + databaseURL: env.databaseURL, + createdMs: Self.ms("2026-03-29T10:00:00.000Z"), + cost: 6.0) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let now = Date(timeIntervalSince1970: TimeInterval(Self.ms("2026-03-29T12:00:00.000Z")) / 1000) + let snapshot = try reader.fetch(now: now) + + #expect(snapshot.monthlyUsagePercent == 10) + #expect(snapshot.monthlyResetInSec == 129_600) + } + + @Test + func `reads step finish parts when message only stores metadata`() 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: nil) + try Self.insertStepFinishPart( + databaseURL: env.databaseURL, + messageID: messageID, + createdMs: Self.ms("2026-03-06T11:00:00.000Z"), + cost: 3.0) + + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) + let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + + #expect(snapshot.rollingUsagePercent == 25) + #expect(snapshot.weeklyUsagePercent == 10) + #expect(snapshot.monthlyUsagePercent == 5) + } + + @Test + func `does not double count step finish parts when message has cost`() 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) + try Self.insertStepFinishPart( + databaseURL: env.databaseURL, + messageID: messageID, + createdMs: Self.ms("2026-03-06T11:00:00.000Z"), + cost: 3.0) + let reader = OpenCodeGoLocalUsageReader(authURL: env.authURL, databaseURL: env.databaseURL) let snapshot = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) - #expect(snapshot.rollingUsagePercent == 0) - #expect(snapshot.weeklyUsagePercent == 0) - #expect(snapshot.monthlyUsagePercent == 0) + #expect(snapshot.rollingUsagePercent == 25) + #expect(snapshot.weeklyUsagePercent == 10) + #expect(snapshot.monthlyUsagePercent == 5) } @Test @@ -109,35 +181,96 @@ struct OpenCodeGoLocalUsageReaderTests { db: db, sql: """ CREATE TABLE message ( + id TEXT PRIMARY KEY, + session_id TEXT NOT NULL, + data TEXT NOT NULL, + time_created INTEGER, + time_updated INTEGER + ); + CREATE TABLE part ( + id TEXT PRIMARY KEY, + message_id TEXT NOT NULL, + session_id TEXT NOT NULL, data TEXT NOT NULL, - time_created INTEGER + time_created INTEGER, + time_updated INTEGER ); """) } - private static func insertMessage(databaseURL: URL, createdMs: Int64, cost: Double) throws { + @discardableResult + private static func insertMessage(databaseURL: URL, createdMs: Int64, cost: Double?) throws -> String { var db: OpaquePointer? guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } defer { sqlite3_close(db) } - let payload: [String: Any] = [ + let messageID = UUID().uuidString + var payload: [String: Any] = [ "providerID": "opencode-go", "role": "assistant", "time": ["created": createdMs], + ] + if let cost { + payload["cost"] = cost + } + let data = try JSONSerialization.data(withJSONObject: payload) + let json = String(data: data, encoding: .utf8) ?? "{}" + + var stmt: OpaquePointer? + guard sqlite3_prepare_v2( + db, + "INSERT INTO message (id, session_id, data, time_created, time_updated) VALUES (?, ?, ?, ?, ?)", + -1, + &stmt, + nil) == SQLITE_OK + else { throw SQLiteTestError.prepare } + defer { sqlite3_finalize(stmt) } + + let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) + sqlite3_bind_text(stmt, 1, messageID, -1, transient) + sqlite3_bind_text(stmt, 2, "session-1", -1, transient) + sqlite3_bind_text(stmt, 3, json, -1, transient) + sqlite3_bind_int64(stmt, 4, createdMs) + sqlite3_bind_int64(stmt, 5, createdMs) + guard sqlite3_step(stmt) == SQLITE_DONE else { throw SQLiteTestError.step } + return messageID + } + + private static func insertStepFinishPart( + databaseURL: URL, + messageID: String, + createdMs: Int64, + cost: Double) throws + { + var db: OpaquePointer? + guard sqlite3_open(databaseURL.path, &db) == SQLITE_OK else { throw SQLiteTestError.open } + defer { sqlite3_close(db) } + + let payload: [String: Any] = [ + "type": "step-finish", "cost": cost, + "tokens": ["input": 1, "output": 1, "total": 2], ] let data = try JSONSerialization.data(withJSONObject: payload) let json = String(data: data, encoding: .utf8) ?? "{}" var stmt: OpaquePointer? - guard sqlite3_prepare_v2(db, "INSERT INTO message (data, time_created) VALUES (?, ?)", -1, &stmt, nil) - == SQLITE_OK + guard sqlite3_prepare_v2( + db, + "INSERT INTO part (id, message_id, session_id, data, time_created, time_updated) VALUES (?, ?, ?, ?, ?, ?)", + -1, + &stmt, + nil) == SQLITE_OK else { throw SQLiteTestError.prepare } defer { sqlite3_finalize(stmt) } let transient = unsafeBitCast(-1, to: sqlite3_destructor_type.self) - sqlite3_bind_text(stmt, 1, json, -1, transient) - sqlite3_bind_int64(stmt, 2, createdMs) + sqlite3_bind_text(stmt, 1, UUID().uuidString, -1, transient) + sqlite3_bind_text(stmt, 2, messageID, -1, transient) + sqlite3_bind_text(stmt, 3, "session-1", -1, transient) + sqlite3_bind_text(stmt, 4, json, -1, transient) + sqlite3_bind_int64(stmt, 5, createdMs) + sqlite3_bind_int64(stmt, 6, createdMs) guard sqlite3_step(stmt) == SQLITE_DONE else { throw SQLiteTestError.step } } diff --git a/Tests/CodexBarTests/OpenCodeGoUsageFetcherErrorTests.swift b/Tests/CodexBarTests/OpenCodeGoUsageFetcherErrorTests.swift index 7929ae3bf..131bfddd1 100644 --- a/Tests/CodexBarTests/OpenCodeGoUsageFetcherErrorTests.swift +++ b/Tests/CodexBarTests/OpenCodeGoUsageFetcherErrorTests.swift @@ -271,6 +271,34 @@ struct OpenCodeGoUsageFetcherErrorTests { #expect(snapshot.toUsageSnapshot().providerCost?.period == "Zen balance") } + @Test + func `optional zen balance helper uses normalized cookie and workspace override`() async throws { + defer { + OpenCodeGoStubURLProtocol.handler = nil + } + + var observedCookie: String? + OpenCodeGoStubURLProtocol.handler = { request in + guard let url = request.url else { throw URLError(.badURL) } + observedCookie = request.value(forHTTPHeaderField: "Cookie") + #expect(url.path == "/workspace/wrk_TEST123") + return Self.makeResponse( + url: url, + body: #"

現在の残高 $98.76

"#, + statusCode: 200, + contentType: "text/html") + } + + let balance = try await OpenCodeGoUsageFetcher.fetchOptionalZenBalance( + cookieHeader: "provider=google; auth=test", + timeout: 2, + workspaceIDOverride: "https://opencode.ai/workspace/wrk_TEST123/go", + session: self.makeSession()) + + #expect(balance == 98.76) + #expect(observedCookie == "auth=test") + } + @Test func `optional zen balance failure does not fail subscription usage`() async throws { defer {