chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,324 @@
|
||||
import Foundation
|
||||
|
||||
#if canImport(SQLite3)
|
||||
import SQLite3
|
||||
#elseif canImport(CSQLite3)
|
||||
import CSQLite3
|
||||
#endif
|
||||
|
||||
#if canImport(SQLite3) || canImport(CSQLite3)
|
||||
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
|
||||
@@ -0,0 +1,192 @@
|
||||
import Foundation
|
||||
|
||||
public enum OpenCodeGoProviderDescriptor {
|
||||
public static let descriptor: ProviderDescriptor = Self.makeDescriptor()
|
||||
|
||||
static func makeDescriptor() -> ProviderDescriptor {
|
||||
ProviderDescriptor(
|
||||
id: .opencodego,
|
||||
metadata: ProviderMetadata(
|
||||
id: .opencodego,
|
||||
displayName: "OpenCode Go",
|
||||
sessionLabel: "5-hour",
|
||||
weeklyLabel: "Weekly",
|
||||
opusLabel: "Monthly",
|
||||
supportsOpus: true,
|
||||
supportsCredits: false,
|
||||
creditsHint: "",
|
||||
toggleTitle: "Show OpenCode Go usage",
|
||||
cliName: "opencodego",
|
||||
defaultEnabled: false,
|
||||
isPrimaryProvider: false,
|
||||
usesAccountFallback: false,
|
||||
browserCookieOrder: ProviderBrowserCookieDefaults.defaultImportOrder,
|
||||
dashboardURL: "https://opencode.ai",
|
||||
statusPageURL: nil),
|
||||
branding: ProviderBranding(
|
||||
iconStyle: .opencodego,
|
||||
iconResourceName: "ProviderIcon-opencodego",
|
||||
color: ProviderColor(red: 59 / 255, green: 130 / 255, blue: 246 / 255)),
|
||||
tokenCost: ProviderTokenCostConfig(
|
||||
supportsTokenCost: false,
|
||||
noDataMessage: { "OpenCode Go cost summary is not supported." }),
|
||||
fetchPlan: ProviderFetchPlan(
|
||||
sourceModes: [.auto, .web],
|
||||
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<Double?, Error> {
|
||||
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 {
|
||||
let id: String = "opencodego.web"
|
||||
let kind: ProviderFetchKind = .web
|
||||
|
||||
func isAvailable(_ context: ProviderFetchContext) async -> Bool {
|
||||
guard context.settings?.opencodego?.cookieSource != .off else { return false }
|
||||
return true
|
||||
}
|
||||
|
||||
func fetch(_ context: ProviderFetchContext) async throws -> ProviderFetchResult {
|
||||
let workspaceOverride = context.settings?.opencodego?.workspaceID
|
||||
?? context.env["CODEXBAR_OPENCODEGO_WORKSPACE_ID"]
|
||||
let cookieSource = context.settings?.opencodego?.cookieSource ?? .auto
|
||||
do {
|
||||
let cookieHeader = try Self.resolveCookieHeader(context: context, allowCached: true)
|
||||
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
|
||||
cookieHeader: cookieHeader,
|
||||
timeout: context.webTimeout,
|
||||
workspaceIDOverride: workspaceOverride,
|
||||
includeZenBalance: context.includeOptionalUsage)
|
||||
return self.makeResult(
|
||||
usage: snapshot.toUsageSnapshot(),
|
||||
sourceLabel: "web")
|
||||
} catch OpenCodeGoUsageError.invalidCredentials where cookieSource != .manual {
|
||||
#if os(macOS)
|
||||
CookieHeaderCache.clear(provider: .opencodego)
|
||||
let cookieHeader = try Self.resolveCookieHeader(context: context, allowCached: false)
|
||||
let snapshot = try await OpenCodeGoUsageFetcher.fetchUsage(
|
||||
cookieHeader: cookieHeader,
|
||||
timeout: context.webTimeout,
|
||||
workspaceIDOverride: workspaceOverride,
|
||||
includeZenBalance: context.includeOptionalUsage)
|
||||
return self.makeResult(
|
||||
usage: snapshot.toUsageSnapshot(),
|
||||
sourceLabel: "web")
|
||||
#else
|
||||
throw OpenCodeGoUsageError.invalidCredentials
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
static func resolveCookieHeader(context: ProviderFetchContext, allowCached: Bool) throws -> String {
|
||||
try OpenCodeWebCookieSupport.resolveCookieHeader(
|
||||
context: OpenCodeWebCookieSupport.Context(
|
||||
settings: context.settings?.opencodego,
|
||||
provider: .opencodego,
|
||||
browserDetection: context.browserDetection,
|
||||
allowCached: allowCached),
|
||||
invalidCookie: OpenCodeGoSettingsError.invalidCookie,
|
||||
missingCookie: OpenCodeGoSettingsError.missingCookie)
|
||||
}
|
||||
}
|
||||
|
||||
enum OpenCodeGoSettingsError: LocalizedError {
|
||||
case missingCookie
|
||||
case invalidCookie
|
||||
|
||||
var errorDescription: String? {
|
||||
switch self {
|
||||
case .missingCookie:
|
||||
"No OpenCode Go session cookies found in browsers."
|
||||
case .invalidCookie:
|
||||
"OpenCode Go cookie header is invalid."
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,145 @@
|
||||
import Foundation
|
||||
|
||||
public struct OpenCodeGoUsageSnapshot: Sendable {
|
||||
public let isBalanceOnly: Bool
|
||||
public let hasWeeklyUsage: Bool
|
||||
public let hasMonthlyUsage: Bool
|
||||
public let rollingUsagePercent: Double
|
||||
public let weeklyUsagePercent: Double
|
||||
public let monthlyUsagePercent: Double
|
||||
public let rollingResetInSec: Int
|
||||
public let weeklyResetInSec: Int
|
||||
public let monthlyResetInSec: Int
|
||||
public let zenBalanceUSD: Double?
|
||||
public let renewsAt: Date?
|
||||
public let updatedAt: Date
|
||||
|
||||
public init(
|
||||
isBalanceOnly: Bool = false,
|
||||
hasWeeklyUsage: Bool = true,
|
||||
hasMonthlyUsage: Bool,
|
||||
rollingUsagePercent: Double,
|
||||
weeklyUsagePercent: Double,
|
||||
monthlyUsagePercent: Double,
|
||||
rollingResetInSec: Int,
|
||||
weeklyResetInSec: Int,
|
||||
monthlyResetInSec: Int,
|
||||
zenBalanceUSD: Double? = nil,
|
||||
renewsAt: Date? = nil,
|
||||
updatedAt: Date)
|
||||
{
|
||||
self.isBalanceOnly = isBalanceOnly
|
||||
self.hasWeeklyUsage = hasWeeklyUsage
|
||||
self.hasMonthlyUsage = hasMonthlyUsage
|
||||
self.rollingUsagePercent = rollingUsagePercent
|
||||
self.weeklyUsagePercent = weeklyUsagePercent
|
||||
self.monthlyUsagePercent = monthlyUsagePercent
|
||||
self.rollingResetInSec = rollingResetInSec
|
||||
self.weeklyResetInSec = weeklyResetInSec
|
||||
self.monthlyResetInSec = monthlyResetInSec
|
||||
self.zenBalanceUSD = zenBalanceUSD
|
||||
self.renewsAt = renewsAt
|
||||
self.updatedAt = updatedAt
|
||||
}
|
||||
|
||||
public static func zenBalanceOnly(balanceUSD: Double, updatedAt: Date) -> OpenCodeGoUsageSnapshot {
|
||||
OpenCodeGoUsageSnapshot(
|
||||
isBalanceOnly: true,
|
||||
hasWeeklyUsage: false,
|
||||
hasMonthlyUsage: false,
|
||||
rollingUsagePercent: 0,
|
||||
weeklyUsagePercent: 0,
|
||||
monthlyUsagePercent: 0,
|
||||
rollingResetInSec: 0,
|
||||
weeklyResetInSec: 0,
|
||||
monthlyResetInSec: 0,
|
||||
zenBalanceUSD: balanceUSD,
|
||||
updatedAt: updatedAt)
|
||||
}
|
||||
|
||||
public func toUsageSnapshot() -> UsageSnapshot {
|
||||
if self.isBalanceOnly {
|
||||
return UsageSnapshot(
|
||||
primary: nil,
|
||||
secondary: nil,
|
||||
providerCost: self.providerCostSnapshot,
|
||||
updatedAt: self.updatedAt,
|
||||
identity: nil)
|
||||
}
|
||||
|
||||
let rollingReset = self.updatedAt.addingTimeInterval(TimeInterval(self.rollingResetInSec))
|
||||
let primary = RateWindow(
|
||||
usedPercent: self.rollingUsagePercent,
|
||||
windowMinutes: 5 * 60,
|
||||
resetsAt: rollingReset,
|
||||
resetDescription: nil)
|
||||
let secondary: RateWindow?
|
||||
if self.hasWeeklyUsage {
|
||||
let weeklyReset = self.updatedAt.addingTimeInterval(TimeInterval(self.weeklyResetInSec))
|
||||
secondary = RateWindow(
|
||||
usedPercent: self.weeklyUsagePercent,
|
||||
windowMinutes: 7 * 24 * 60,
|
||||
resetsAt: weeklyReset,
|
||||
resetDescription: nil)
|
||||
} else {
|
||||
secondary = nil
|
||||
}
|
||||
let tertiary: RateWindow?
|
||||
if self.hasMonthlyUsage {
|
||||
let monthlyReset = self.updatedAt.addingTimeInterval(TimeInterval(self.monthlyResetInSec))
|
||||
tertiary = RateWindow(
|
||||
usedPercent: self.monthlyUsagePercent,
|
||||
windowMinutes: 30 * 24 * 60,
|
||||
resetsAt: monthlyReset,
|
||||
resetDescription: nil)
|
||||
} else {
|
||||
tertiary = nil
|
||||
}
|
||||
|
||||
var extraWindows: [NamedRateWindow]?
|
||||
if let renewsAt = self.renewsAt {
|
||||
let renewalWindow = RateWindow(
|
||||
usedPercent: 0,
|
||||
windowMinutes: nil,
|
||||
resetsAt: renewsAt,
|
||||
resetDescription: nil)
|
||||
extraWindows = [NamedRateWindow(id: "renewal", title: "Renews", window: renewalWindow)]
|
||||
}
|
||||
|
||||
return UsageSnapshot(
|
||||
primary: primary,
|
||||
secondary: secondary,
|
||||
tertiary: tertiary,
|
||||
extraRateWindows: extraWindows,
|
||||
providerCost: self.providerCostSnapshot,
|
||||
updatedAt: self.updatedAt,
|
||||
identity: nil)
|
||||
}
|
||||
|
||||
private var providerCostSnapshot: ProviderCostSnapshot? {
|
||||
self.zenBalanceUSD.map {
|
||||
ProviderCostSnapshot(
|
||||
used: $0,
|
||||
limit: 0,
|
||||
currencyCode: "USD",
|
||||
period: "Zen balance",
|
||||
updatedAt: self.updatedAt)
|
||||
}
|
||||
}
|
||||
|
||||
public func withZenBalanceUSD(_ balance: Double?) -> OpenCodeGoUsageSnapshot {
|
||||
OpenCodeGoUsageSnapshot(
|
||||
isBalanceOnly: self.isBalanceOnly,
|
||||
hasWeeklyUsage: self.hasWeeklyUsage,
|
||||
hasMonthlyUsage: self.hasMonthlyUsage,
|
||||
rollingUsagePercent: self.rollingUsagePercent,
|
||||
weeklyUsagePercent: self.weeklyUsagePercent,
|
||||
monthlyUsagePercent: self.monthlyUsagePercent,
|
||||
rollingResetInSec: self.rollingResetInSec,
|
||||
weeklyResetInSec: self.weeklyResetInSec,
|
||||
monthlyResetInSec: self.monthlyResetInSec,
|
||||
zenBalanceUSD: balance,
|
||||
renewsAt: self.renewsAt,
|
||||
updatedAt: self.updatedAt)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,166 @@
|
||||
import Foundation
|
||||
#if canImport(FoundationNetworking)
|
||||
import FoundationNetworking
|
||||
#endif
|
||||
|
||||
private final class OpenCodeGoZenBalanceTaskRace: @unchecked Sendable {
|
||||
private let lock = NSLock()
|
||||
private let sourceTask: Task<Double?, Error>
|
||||
private var result: Result<Double?, Error>?
|
||||
private var continuation: CheckedContinuation<Double?, Error>?
|
||||
private var observerTask: Task<Void, Never>?
|
||||
private var timeoutTask: Task<Void, Never>?
|
||||
|
||||
init(sourceTask: Task<Double?, Error>) {
|
||||
self.sourceTask = sourceTask
|
||||
}
|
||||
|
||||
func value(timeout: Duration? = nil) async throws -> Double? {
|
||||
try await withTaskCancellationHandler {
|
||||
try await withCheckedThrowingContinuation { continuation in
|
||||
self.lock.lock()
|
||||
if let result = self.result {
|
||||
self.lock.unlock()
|
||||
continuation.resume(with: result)
|
||||
return
|
||||
}
|
||||
|
||||
self.continuation = continuation
|
||||
let sourceTask = self.sourceTask
|
||||
self.observerTask = Task { [weak self] in
|
||||
do {
|
||||
let value = try await sourceTask.value
|
||||
self?.resolve(.success(value), cancelSource: false)
|
||||
} catch {
|
||||
self?.resolve(.failure(error), cancelSource: false)
|
||||
}
|
||||
}
|
||||
if let timeout {
|
||||
self.timeoutTask = Task { [weak self] in
|
||||
do {
|
||||
try await Task.sleep(for: timeout)
|
||||
self?.resolve(.success(nil), cancelSource: true)
|
||||
} catch {
|
||||
// The source completed or the caller canceled the race.
|
||||
}
|
||||
}
|
||||
}
|
||||
self.lock.unlock()
|
||||
}
|
||||
} onCancel: {
|
||||
self.resolve(.failure(CancellationError()), cancelSource: true)
|
||||
}
|
||||
}
|
||||
|
||||
private func resolve(_ result: Result<Double?, Error>, cancelSource: Bool) {
|
||||
self.lock.lock()
|
||||
guard self.result == nil else {
|
||||
self.lock.unlock()
|
||||
return
|
||||
}
|
||||
|
||||
self.result = result
|
||||
let continuation = self.continuation
|
||||
self.continuation = nil
|
||||
let observerTask = self.observerTask
|
||||
let timeoutTask = self.timeoutTask
|
||||
self.observerTask = nil
|
||||
self.timeoutTask = nil
|
||||
self.lock.unlock()
|
||||
|
||||
if cancelSource {
|
||||
self.sourceTask.cancel()
|
||||
}
|
||||
observerTask?.cancel()
|
||||
timeoutTask?.cancel()
|
||||
continuation?.resume(with: result)
|
||||
}
|
||||
}
|
||||
|
||||
extension OpenCodeGoUsageFetcher {
|
||||
static let optionalZenBalanceTimeout: TimeInterval = 5
|
||||
static let optionalZenBalanceStartDelay: Duration = .milliseconds(25)
|
||||
static let optionalZenBalanceJoinGrace: Duration = .milliseconds(250)
|
||||
|
||||
public static func zenDashboardURL(workspaceID raw: String?) -> URL {
|
||||
guard let workspaceID = self.normalizeWorkspaceID(raw),
|
||||
let url = URL(string: "https://opencode.ai/workspace/\(workspaceID)")
|
||||
else {
|
||||
return URL(string: "https://opencode.ai")!
|
||||
}
|
||||
return url
|
||||
}
|
||||
|
||||
static func fetchOptionalZenBalance(
|
||||
workspaceID: String,
|
||||
cookieHeader: String,
|
||||
timeout: TimeInterval,
|
||||
session: URLSession) async throws -> Double?
|
||||
{
|
||||
do {
|
||||
let balance = try await self.fetchZenBalance(
|
||||
workspaceID: workspaceID,
|
||||
cookieHeader: cookieHeader,
|
||||
timeout: timeout,
|
||||
session: session)
|
||||
try Task.checkCancellation()
|
||||
return balance
|
||||
} catch is CancellationError {
|
||||
throw CancellationError()
|
||||
} catch {
|
||||
if Task.isCancelled {
|
||||
throw CancellationError()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
static func completedOptionalZenBalance(from task: Task<Double?, Error>) async throws -> Double? {
|
||||
let race = OpenCodeGoZenBalanceTaskRace(sourceTask: task)
|
||||
do {
|
||||
return try await race.value(timeout: self.optionalZenBalanceJoinGrace)
|
||||
} catch is CancellationError {
|
||||
throw CancellationError()
|
||||
} catch {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
static func completedRequiredZenBalance(from task: Task<Double?, Error>) async throws -> Double? {
|
||||
let race = OpenCodeGoZenBalanceTaskRace(sourceTask: task)
|
||||
return try await race.value()
|
||||
}
|
||||
|
||||
static func parseZenBalance(text: String) -> Double? {
|
||||
OpenCodeGoZenBalanceParser.parse(text: text)
|
||||
}
|
||||
|
||||
static func fetchZenBalance(
|
||||
workspaceID: String,
|
||||
cookieHeader: String,
|
||||
timeout: TimeInterval,
|
||||
session: URLSession) async throws -> Double?
|
||||
{
|
||||
let text = try await self.fetchPageText(
|
||||
url: self.zenDashboardURL(workspaceID: workspaceID),
|
||||
cookieHeader: cookieHeader,
|
||||
timeout: timeout,
|
||||
session: session)
|
||||
if self.looksSignedOut(text: text) {
|
||||
throw OpenCodeGoUsageError.invalidCredentials
|
||||
}
|
||||
if let balance = self.parseZenBalance(text: text) {
|
||||
return balance
|
||||
}
|
||||
|
||||
let billingText = try await self.fetchZenBillingText(
|
||||
workspaceID: workspaceID,
|
||||
cookieHeader: cookieHeader,
|
||||
timeout: timeout,
|
||||
session: session)
|
||||
if self.looksSignedOut(text: billingText) {
|
||||
throw OpenCodeGoUsageError.invalidCredentials
|
||||
}
|
||||
return OpenCodeGoZenBalanceParser.parseBillingServerResponse(text: billingText)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,155 @@
|
||||
import Foundation
|
||||
|
||||
enum OpenCodeGoZenBalanceParser {
|
||||
private static let billingScale = 100_000_000.0
|
||||
|
||||
static func parse(text: String) -> Double? {
|
||||
if let value = self.parseJSON(text: text) {
|
||||
return value
|
||||
}
|
||||
let localizedPattern = [
|
||||
#"(?i)(?:current\s+balance|zen\s+balance|現在の残高)"#,
|
||||
#"[^$]{0,80}\$\s*([0-9][0-9,]*(?:\.[0-9]+)?)"#,
|
||||
].joined()
|
||||
if let value = self.extractDollarValue(pattern: localizedPattern, text: text) {
|
||||
return value
|
||||
}
|
||||
let nearbyPattern = #"(?i)(?:balance|残高)[\s\S]{0,120}?\$\s*([0-9][0-9,]*(?:\.[0-9]+)?)"#
|
||||
return self.extractDollarValue(pattern: nearbyPattern, text: text)
|
||||
}
|
||||
|
||||
static func parseBillingServerResponse(text: String) -> Double? {
|
||||
if let data = text.data(using: .utf8),
|
||||
let object = try? JSONSerialization.jsonObject(with: data, options: []),
|
||||
let rawBalance = self.findRawBillingBalance(in: object)
|
||||
{
|
||||
return rawBalance / self.billingScale
|
||||
}
|
||||
|
||||
let customerPattern =
|
||||
#"(?:\"customerID\"|customerID)\s*:\s*(?:\$R\[\d+\]\s*=\s*)?\"[^\"]+\""#
|
||||
guard self.containsMatch(pattern: customerPattern, text: text) else {
|
||||
return nil
|
||||
}
|
||||
let pattern = #"(?:\"balance\"|balance)\s*:\s*(?:\$R\[\d+\]\s*=\s*)?(-?[0-9]+(?:\.[0-9]+)?)"#
|
||||
guard let rawBalance = self.extractNumber(pattern: pattern, text: text) else {
|
||||
return nil
|
||||
}
|
||||
return rawBalance / self.billingScale
|
||||
}
|
||||
|
||||
private static func parseJSON(text: String) -> Double? {
|
||||
guard let data = text.data(using: .utf8),
|
||||
let object = try? JSONSerialization.jsonObject(with: data, options: [])
|
||||
else {
|
||||
return nil
|
||||
}
|
||||
return self.findBalanceValue(in: object, path: [])
|
||||
}
|
||||
|
||||
private static func findBalanceValue(in object: Any, path: [String]) -> Double? {
|
||||
if let dict = object as? [String: Any] {
|
||||
for (key, value) in dict {
|
||||
let nextPath = path + [key]
|
||||
if self.isExplicitBalanceAmountKey(key),
|
||||
let number = self.doubleValue(from: value)
|
||||
{
|
||||
return number
|
||||
}
|
||||
if let found = self.findBalanceValue(in: value, path: nextPath) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if let array = object as? [Any] {
|
||||
for (index, value) in array.enumerated() {
|
||||
if let found = self.findBalanceValue(in: value, path: path + ["[\(index)]"]) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
private static func findRawBillingBalance(in object: Any) -> Double? {
|
||||
if let dict = object as? [String: Any] {
|
||||
if dict["balance"] != nil {
|
||||
guard let customerID = dict["customerID"] as? String,
|
||||
!customerID.isEmpty
|
||||
else {
|
||||
return nil
|
||||
}
|
||||
guard let rawBalance = self.doubleValue(from: dict["balance"]) else {
|
||||
return nil
|
||||
}
|
||||
return rawBalance
|
||||
}
|
||||
for value in dict.values {
|
||||
if let found = self.findRawBillingBalance(in: value) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
} else if let array = object as? [Any] {
|
||||
for value in array {
|
||||
if let found = self.findRawBillingBalance(in: value) {
|
||||
return found
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
private static func containsMatch(pattern: String, text: String) -> Bool {
|
||||
guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return false }
|
||||
let nsrange = NSRange(text.startIndex..<text.endIndex, in: text)
|
||||
return regex.firstMatch(in: text, options: [], range: nsrange) != nil
|
||||
}
|
||||
|
||||
private static func isExplicitBalanceAmountKey(_ key: String) -> Bool {
|
||||
let normalized = key
|
||||
.lowercased()
|
||||
.filter { $0.isLetter || $0.isNumber }
|
||||
return [
|
||||
"zenbalance",
|
||||
"zencurrentbalance",
|
||||
"currentbalance",
|
||||
"currentbalanceusd",
|
||||
"balanceusd",
|
||||
"usdbalance",
|
||||
].contains(normalized)
|
||||
}
|
||||
|
||||
private static func extractDollarValue(pattern: String, text: String) -> Double? {
|
||||
self.extractNumber(pattern: pattern, text: text)
|
||||
}
|
||||
|
||||
private static func extractNumber(pattern: String, text: String) -> Double? {
|
||||
guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return nil }
|
||||
let nsrange = NSRange(text.startIndex..<text.endIndex, in: text)
|
||||
guard let match = regex.firstMatch(in: text, options: [], range: nsrange),
|
||||
let range = Range(match.range(at: 1), in: text)
|
||||
else {
|
||||
return nil
|
||||
}
|
||||
return Double(text[range].replacingOccurrences(of: ",", with: ""))
|
||||
}
|
||||
|
||||
private static func doubleValue(from value: Any?) -> Double? {
|
||||
switch value {
|
||||
case is Bool:
|
||||
nil
|
||||
case let number as Double:
|
||||
number
|
||||
case let number as NSNumber:
|
||||
number.doubleValue
|
||||
case let string as String:
|
||||
Double(
|
||||
string
|
||||
.trimmingCharacters(in: .whitespacesAndNewlines)
|
||||
.replacingOccurrences(of: ",", with: ""))
|
||||
default:
|
||||
nil
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user