diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b65a4ea8..123293ee1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - Display: add optional workday markers for weekly progress bars (#1102). Thanks @Yuxin-Qiao! ### Fixed +- OpenCode Go: read local usage history before falling back to browser-cookie dashboard fetches (#1021). Thanks @sopenlaz0! - Menu bar: show extra-usage spend as currency text for Claude and Cursor when that metric is selected (#1107). Thanks @Yuxin-Qiao! - Codex: run regular credits and OpenAI dashboard refreshes in the background while coalescing overlapping refresh work (#1078). Thanks @ptstory! diff --git a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift new file mode 100644 index 000000000..13098a445 --- /dev/null +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoLocalUsageReader.swift @@ -0,0 +1,320 @@ +import Foundation + +#if os(macOS) +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)" + } + } +} + +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 { + throw OpenCodeGoLocalUsageError.historyUnavailable("database not found") + } + throw OpenCodeGoLocalUsageError.notDetected + } + + 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) + } + + 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 = 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 { + 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 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 + } + + 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 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) + let nowComponents = calendar.dateComponents([.year, .month], from: now) + + 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 = self.anchoredMonth( + calendar: calendar, + month: self.monthComponents(after: startMonthComponents, calendar: calendar), + anchor: anchorComponents) + return (Int64(start.timeIntervalSince1970 * 1000), Int64(end.timeIntervalSince1970 * 1000)) + } + startMonthComponents = calendar.dateComponents([.year, .month], from: previous) + start = self.anchoredMonth(calendar: calendar, month: startMonthComponents, anchor: anchorComponents) + } + 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, + 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..51b2ac417 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoProviderDescriptor.swift @@ -33,11 +33,84 @@ 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 [ + OpenCodeGoUsageFetchStrategy(), + OpenCodeGoLocalUsageFetchStrategy(), + ] + } +} + +struct OpenCodeGoLocalUsageFetchStrategy: ProviderFetchStrategy { + let id: String = "opencodego.local" + let kind: ProviderFetchKind = .localProbe + + func isAvailable(_: ProviderFetchContext) async -> Bool { + true + } + + func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult { + let snapshot = try await self.snapshot(context: context) + return self.makeResult( + usage: snapshot.toUsageSnapshot(), + sourceLabel: "local") + } + + 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 { @@ -81,11 +154,19 @@ struct OpenCodeGoUsageFetchStrategy: ProviderFetchStrategy { } } - func shouldFallback(on _: Error, context _: ProviderFetchContext) -> Bool { - false + func shouldFallback(on error: Error, context: ProviderFetchContext) -> Bool { + guard context.sourceMode == .auto else { return false } + return switch error { + case OpenCodeGoSettingsError.missingCookie, + OpenCodeGoSettingsError.invalidCookie, + OpenCodeGoUsageError.invalidCredentials: + true + default: + 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 39b290f8b..70e18ec04 100644 --- a/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoUsageFetcher.swift +++ b/Sources/CodexBarCore/Providers/OpenCodeGo/OpenCodeGoUsageFetcher.swift @@ -151,6 +151,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(), @@ -168,7 +193,9 @@ public struct OpenCodeGoUsageFetcher: Sendable { } return url } +} +extension OpenCodeGoUsageFetcher { private static func fetchWorkspaceID( cookieHeader: String, timeout: TimeInterval, diff --git a/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift new file mode 100644 index 000000000..ca060b0a3 --- /dev/null +++ b/Tests/CodexBarTests/OpenCodeGoLocalUsageReaderTests.swift @@ -0,0 +1,299 @@ +#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 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) + + #expect(throws: OpenCodeGoLocalUsageError.historyUnavailable("database not found")) { + _ = try reader.fetch(now: Date(timeIntervalSince1970: 1_772_798_400)) + } + } + + @Test + func `auth with unreadable 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) + 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) + + #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 == 25) + #expect(snapshot.weeklyUsagePercent == 10) + #expect(snapshot.monthlyUsagePercent == 5) + } + + @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 ( + 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_updated INTEGER + ); + """) + } + + @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 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 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, 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 } + } + + 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 diff --git a/Tests/CodexBarTests/OpenCodeGoProviderStrategyTests.swift b/Tests/CodexBarTests/OpenCodeGoProviderStrategyTests.swift new file mode 100644 index 000000000..4c4d4d398 --- /dev/null +++ b/Tests/CodexBarTests/OpenCodeGoProviderStrategyTests.swift @@ -0,0 +1,64 @@ +import Foundation +import Testing +@testable import CodexBarCore + +struct OpenCodeGoProviderStrategyTests { + private struct StubClaudeFetcher: ClaudeUsageFetching { + func loadLatestUsage(model _: String) async throws -> ClaudeUsageSnapshot { + throw ClaudeUsageError.parseFailed("stub") + } + + func debugRawProbe(model _: String) async -> String { + "stub" + } + + func detectVersion() -> String? { + nil + } + } + + private func makeContext(sourceMode: ProviderSourceMode = .auto) -> ProviderFetchContext { + let env: [String: String] = [:] + return ProviderFetchContext( + runtime: .app, + sourceMode: sourceMode, + includeCredits: false, + webTimeout: 1, + webDebugDumpHTML: false, + verbose: false, + env: env, + settings: nil, + fetcher: UsageFetcher(environment: env), + claudeFetcher: StubClaudeFetcher(), + browserDetection: BrowserDetection(cacheTTL: 0)) + } + + @Test + func `auto source prefers web before local fallback`() async { + let descriptor = OpenCodeGoProviderDescriptor.makeDescriptor() + let strategies = await descriptor.fetchPlan.pipeline.resolveStrategies(self.makeContext()) + + #expect(strategies.map(\.id) == ["opencodego.web", "opencodego.local"]) + } + + @Test + func `web source does not include local fallback`() async { + let descriptor = OpenCodeGoProviderDescriptor.makeDescriptor() + let strategies = await descriptor.fetchPlan.pipeline.resolveStrategies(self.makeContext(sourceMode: .web)) + + #expect(strategies.map(\.id) == ["opencodego.web"]) + } + + @Test + func `web strategy falls back to local only for auth setup failures in auto mode`() { + let strategy = OpenCodeGoUsageFetchStrategy() + let autoContext = self.makeContext() + let webContext = self.makeContext(sourceMode: .web) + + #expect(strategy.shouldFallback(on: OpenCodeGoSettingsError.missingCookie, context: autoContext)) + #expect(strategy.shouldFallback(on: OpenCodeGoSettingsError.invalidCookie, context: autoContext)) + #expect(strategy.shouldFallback(on: OpenCodeGoUsageError.invalidCredentials, context: autoContext)) + #expect(!strategy.shouldFallback(on: OpenCodeGoUsageError.networkError("timeout"), context: autoContext)) + #expect(!strategy.shouldFallback(on: OpenCodeGoSettingsError.missingCookie, context: webContext)) + } +} 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 {