1027 lines · 38.1 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif public enum OpenCodeGoUsageError: LocalizedError {    case invalidCredentials    case networkError(String)    case apiError(String)    case parseFailed(String)     public var errorDescription: String? {        switch self {        case .invalidCredentials:            "OpenCode Go session cookie is invalid or expired."        case let .networkError(message):            "OpenCode Go network error: \(message)"        case let .apiError(message):            "OpenCode Go API error: \(message)"        case let .parseFailed(message):            "OpenCode Go parse error: \(message)"        }    }} public struct OpenCodeGoUsageFetcher: Sendable {    private static let log = CodexBarLog.logger(LogCategories.opencodeGoUsage)    private static let baseURL = URL(string: "https://opencode.ai")!    private static let serverURL = URL(string: "https://opencode.ai/_server")!    private static let workspacesServerID = "def39973159c7f0483d8793a822b8dbb10d067e12c65455fcb4608459ba0234f"    private static let billingServerID = "c83b78a614689c38ebee981f9b39a8b377716db85c1fd7dbab604adc02d3313d"     private static let userAgent =        "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) " +        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36"     private final class RedirectGuardDelegate: NSObject, URLSessionTaskDelegate {        func urlSession(            _ session: URLSession,            task: URLSessionTask,            willPerformHTTPRedirection response: HTTPURLResponse,            newRequest request: URLRequest,            completionHandler: @escaping (URLRequest?) -> Void)        {            guard OpenCodeGoUsageFetcher.allowsRedirect(                from: task.originalRequest?.url,                to: request.url)            else {                completionHandler(nil)                return            }            completionHandler(request)        }    }     private struct ServerRequest {        let serverID: String        let args: String?        let method: String        let referer: URL    }     // swiftformat:disable:next redundantSendable    struct ZenBalanceRequest: Sendable {        let workspaceID: String        let cookieHeader: String        let timeout: TimeInterval        let session: URLSession    }     private static let percentKeys = [        "usagePercent",        "usedPercent",        "percentUsed",        "percent",        "usage_percent",        "used_percent",        "utilization",        "utilizationPercent",        "utilization_percent",        "usage",    ]    private static let resetInKeys = [        "resetInSec",        "resetInSeconds",        "resetSeconds",        "reset_sec",        "reset_in_sec",        "resetsInSec",        "resetsInSeconds",        "resetIn",        "resetSec",    ]    private static let resetAtKeys = [        "resetAt",        "resetsAt",        "reset_at",        "resets_at",        "nextReset",        "next_reset",        "renewAt",        "renew_at",    ]    private static let renewAtKeys = [        "renewAt",        "renew_at",    ]    private static let redirectGuardDelegate = RedirectGuardDelegate()    private static let redirectGuardSession: URLSession = {        let configuration = URLSessionConfiguration.ephemeral        configuration.httpCookieStorage = nil        return URLSession(            configuration: configuration,            delegate: OpenCodeGoUsageFetcher.redirectGuardDelegate,            delegateQueue: nil)    }()     public static func fetchUsage(        cookieHeader: String,        timeout: TimeInterval,        now: Date = Date(),        workspaceIDOverride: String? = nil,        includeZenBalance: Bool = true,        session: URLSession? = nil) async throws -> OpenCodeGoUsageSnapshot    {        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)        }        let zenBalanceRequest = ZenBalanceRequest(            workspaceID: workspaceID,            cookieHeader: requestCookieHeader,            timeout: timeout,            session: session)        let subscriptionTask = Task {            try await self.fetchUsagePage(                workspaceID: workspaceID,                cookieHeader: requestCookieHeader,                timeout: timeout,                session: session)        }        let zenBalanceTask = includeZenBalance ? Task {            try await Task.sleep(for: self.optionalZenBalanceStartDelay)            return try await self.fetchZenBalance(                workspaceID: workspaceID,                cookieHeader: requestCookieHeader,                timeout: timeout,                session: session)        } : nil        defer {            subscriptionTask.cancel()            zenBalanceTask?.cancel()        }        let subscriptionText: String        do {            subscriptionText = try await withTaskCancellationHandler {                try await subscriptionTask.value            } onCancel: {                subscriptionTask.cancel()                zenBalanceTask?.cancel()            }        } catch is CancellationError {            throw CancellationError()        } catch let error as URLError where error.code == .cancelled {            throw CancellationError()        } catch let error as OpenCodeGoUsageError {            return try await self.requiredZenBalanceFallback(                from: zenBalanceTask,                for: error,                request: zenBalanceRequest,                now: now)        } catch {            throw error        }        let snapshot: OpenCodeGoUsageSnapshot        do {            snapshot = try self.parseSubscription(text: subscriptionText, now: now)        } catch let error as OpenCodeGoUsageError {            return try await self.requiredZenBalanceFallback(                from: zenBalanceTask,                for: error,                request: zenBalanceRequest,                now: now)        }        guard let zenBalanceTask else {            return snapshot        }        let zenBalance = try await self.completedOptionalZenBalance(from: zenBalanceTask)        return snapshot.withZenBalanceUSD(zenBalance)    }     static func requiredZenBalanceFallback(        from task: Task<Double?, Error>?,        for error: OpenCodeGoUsageError,        request: ZenBalanceRequest,        now: Date) async throws -> OpenCodeGoUsageSnapshot    {        guard case let .parseFailed(message) = error,              message.contains("Missing usage fields")        else {            throw error        }        let task = task ?? Task {            try await self.fetchZenBalance(                workspaceID: request.workspaceID,                cookieHeader: request.cookieHeader,                timeout: request.timeout,                session: request.session)        }        defer {            task.cancel()        }        let zenBalance = try await self.completedRequiredZenBalance(from: task)        guard let zenBalance else {            throw error        }        return OpenCodeGoUsageSnapshot.zenBalanceOnly(balanceUSD: zenBalance, updatedAt: now)    }     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(),              sourceHost == destinationHost,              destinationURL?.scheme?.lowercased() == "https"        else { return false }        return true    }     public static func dashboardURL(workspaceID raw: String?) -> URL {        guard let workspaceID = self.normalizeWorkspaceID(raw),              let url = URL(string: "\(self.baseURL.absoluteString)/workspace/\(workspaceID)/go")        else {            return self.baseURL        }        return url    }} extension OpenCodeGoUsageFetcher {    static func fetchZenBillingText(        workspaceID: String,        cookieHeader: String,        timeout: TimeInterval,        session: URLSession) async throws -> String    {        let argsData = try JSONSerialization.data(withJSONObject: [workspaceID])        guard let args = String(data: argsData, encoding: .utf8) else {            throw OpenCodeGoUsageError.parseFailed("Could not encode billing request.")        }        return try await self.fetchServerText(            request: ServerRequest(                serverID: self.billingServerID,                args: args,                method: "GET",                referer: self.zenDashboardURL(workspaceID: workspaceID)),            cookieHeader: cookieHeader,            timeout: timeout,            session: session)    }     private static func fetchWorkspaceID(        cookieHeader: String,        timeout: TimeInterval,        session: URLSession) async throws -> String    {        let text = try await self.fetchServerText(            request: ServerRequest(                serverID: self.workspacesServerID,                args: nil,                method: "GET",                referer: self.baseURL),            cookieHeader: cookieHeader,            timeout: timeout,            session: session)        if self.looksSignedOut(text: text) {            throw OpenCodeGoUsageError.invalidCredentials        }        var ids = self.parseWorkspaceIDs(text: text)        if ids.isEmpty {            ids = self.parseWorkspaceIDsFromJSON(text: text)        }        if ids.isEmpty {            Self.log.error("OpenCode Go workspace ids missing after GET; retrying with POST.")            let fallback = try await self.fetchServerText(                request: ServerRequest(                    serverID: self.workspacesServerID,                    args: "[]",                    method: "POST",                    referer: self.baseURL),                cookieHeader: cookieHeader,                timeout: timeout,                session: session)            if self.looksSignedOut(text: fallback) {                throw OpenCodeGoUsageError.invalidCredentials            }            ids = self.parseWorkspaceIDs(text: fallback)            if ids.isEmpty {                ids = self.parseWorkspaceIDsFromJSON(text: fallback)            }            if ids.isEmpty {                throw OpenCodeGoUsageError.parseFailed("Missing workspace id.")            }            return ids[0]        }        return ids[0]    }     static func normalizeWorkspaceID(_ raw: String?) -> String? {        guard let raw else { return nil }        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)        if trimmed.hasPrefix("wrk_"), trimmed.count > 4 {            return trimmed        }        if let url = URL(string: trimmed) {            let parts = url.pathComponents            if let index = parts.firstIndex(of: "workspace"),               parts.count > index + 1            {                let candidate = parts[index + 1]                if candidate.hasPrefix("wrk_"), candidate.count > 4 {                    return candidate                }            }        }        if let match = trimmed.range(of: #"wrk_[A-Za-z0-9]+"#, options: .regularExpression) {            return String(trimmed[match])        }        return nil    }     static func parseWorkspaceIDs(text: String) -> [String] {        let pattern = #"id\s*:\s*\"(wrk_[^\"]+)\""#        guard let regex = try? NSRegularExpression(pattern: pattern, options: []) else { return [] }        let nsrange = NSRange(text.startIndex..<text.endIndex, in: text)        return regex.matches(in: text, options: [], range: nsrange).compactMap { match in            guard let range = Range(match.range(at: 1), in: text) else { return nil }            return String(text[range])        }    }     private static func parseWorkspaceIDsFromJSON(text: String) -> [String] {        guard let data = text.data(using: .utf8),              let object = try? JSONSerialization.jsonObject(with: data, options: [])        else {            return []        }        var results: [String] = []        self.collectWorkspaceIDs(object: object, out: &results)        return results    }     private static func collectWorkspaceIDs(object: Any, out: inout [String]) {        if let dict = object as? [String: Any] {            for (_, value) in dict {                self.collectWorkspaceIDs(object: value, out: &out)            }            return        }        if let array = object as? [Any] {            for value in array {                self.collectWorkspaceIDs(object: value, out: &out)            }            return        }        if let string = object as? String,           string.hasPrefix("wrk_"),           !out.contains(string)        {            out.append(string)        }    }     private static func fetchUsagePage(        workspaceID: String,        cookieHeader: String,        timeout: TimeInterval,        session: URLSession) async throws -> String    {        let url = URL(string: "https://opencode.ai/workspace/\(workspaceID)/go") ?? self.baseURL        let text = try await self.fetchPageText(            url: url,            cookieHeader: cookieHeader,            timeout: timeout,            session: session)        if self.looksSignedOut(text: text) {            throw OpenCodeGoUsageError.invalidCredentials        }        guard self.parseSubscriptionJSON(text: text, now: Date()) != nil ||            self.extractDouble(                pattern: #"rollingUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#,                text: text) != nil        else {            Self.log.error("OpenCode Go usage page payload missing usage fields.")            throw OpenCodeGoUsageError.parseFailed("Missing usage fields.")        }        return text    }     static func parseSubscription(text: String, now: Date) throws -> OpenCodeGoUsageSnapshot {        if let snapshot = self.parseSubscriptionJSON(text: text, now: now) {            return snapshot        }         guard let rollingPercent = self.extractDouble(            pattern: #"rollingUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#,            text: text),            let rollingReset = self.extractInt(                pattern: #"rollingUsage[^}]*?resetInSec\s*:\s*([0-9]+)"#,                text: text)        else {            throw OpenCodeGoUsageError.parseFailed("Missing usage fields.")        }         let weeklyPercent = self.extractDouble(            pattern: #"weeklyUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#,            text: text)        let weeklyReset = self.extractInt(            pattern: #"weeklyUsage[^}]*?resetInSec\s*:\s*([0-9]+)"#,            text: text)        let hasWeeklyUsage = weeklyPercent != nil && weeklyReset != nil         let monthlyPercent = self.extractDouble(            pattern: #"monthlyUsage[^}]*?usagePercent\s*:\s*([0-9]+(?:\.[0-9]+)?)"#,            text: text)        let monthlyReset = self.extractInt(            pattern: #"monthlyUsage[^}]*?resetInSec\s*:\s*([0-9]+)"#,            text: text)         return OpenCodeGoUsageSnapshot(            hasWeeklyUsage: hasWeeklyUsage,            hasMonthlyUsage: monthlyPercent != nil || monthlyReset != nil,            rollingUsagePercent: rollingPercent,            weeklyUsagePercent: weeklyPercent ?? 0,            monthlyUsagePercent: monthlyPercent ?? 0,            rollingResetInSec: rollingReset,            weeklyResetInSec: weeklyReset ?? 0,            monthlyResetInSec: monthlyReset ?? 0,            updatedAt: now)    }     private static func parseSubscriptionJSON(text: String, now: Date) -> OpenCodeGoUsageSnapshot? {        guard let data = text.data(using: .utf8),              let object = try? JSONSerialization.jsonObject(with: data, options: []),              let dict = object as? [String: Any]        else {            return nil        }         let renewsAt = self.dateValue(from: self.value(from: dict, keys: self.renewAtKeys))        if let snapshot = self.parseUsageDictionary(dict, now: now, inheritedRenewsAt: renewsAt) {            return snapshot        }        for key in ["data", "result", "usage", "billing", "payload"] {            if let nested = dict[key] as? [String: Any],               let snapshot = self.parseUsageDictionary(nested, now: now, inheritedRenewsAt: renewsAt)            {                return snapshot            }        }        if let snapshot = self.parseUsageNested(dict, now: now, depth: 0, inheritedRenewsAt: renewsAt) {            return snapshot        }        return self.parseUsageFromCandidates(object: object, now: now, inheritedRenewsAt: renewsAt)    }     private static func parseUsageDictionary(        _ dict: [String: Any],

Lines 1–500 of 1027.