909 lines · 33.8 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 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    }     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 subscriptionText: String        do {            subscriptionText = try await self.fetchUsagePage(                workspaceID: workspaceID,                cookieHeader: requestCookieHeader,                timeout: timeout,                session: session)        } catch {            throw error        }        let snapshot = try self.parseSubscription(text: subscriptionText, now: now)        let zenBalanceTask = includeZenBalance ? Task {            try await self.fetchOptionalZenBalance(                workspaceID: workspaceID,                cookieHeader: requestCookieHeader,                timeout: min(timeout, self.optionalZenBalanceTimeout),                session: session)        } : nil        guard let zenBalanceTask else {            return snapshot        }        let zenBalance = try await self.completedOptionalZenBalance(from: zenBalanceTask)        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(),              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 {    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),            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)        else {            throw OpenCodeGoUsageError.parseFailed("Missing usage fields.")        }         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(            hasMonthlyUsage: monthlyPercent != nil || monthlyReset != nil,            rollingUsagePercent: rollingPercent,            weeklyUsagePercent: weeklyPercent,            monthlyUsagePercent: monthlyPercent ?? 0,            rollingResetInSec: rollingReset,            weeklyResetInSec: weeklyReset,            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],        now: Date,        inheritedRenewsAt: Date?) -> OpenCodeGoUsageSnapshot?    {        let renewsAt = self.dateValue(from: self.value(from: dict, keys: self.renewAtKeys)) ?? inheritedRenewsAt        if let usage = dict["usage"] as? [String: Any],           let snapshot = self.parseUsageDictionary(usage, now: now, inheritedRenewsAt: renewsAt)        {            return snapshot        }         let rollingKeys = ["rollingUsage", "rolling", "rolling_usage", "rollingWindow", "rolling_window"]        let weeklyKeys = ["weeklyUsage", "weekly", "weekly_usage", "weeklyWindow", "weekly_window"]        let monthlyKeys = ["monthlyUsage", "monthly", "monthly_usage", "monthlyWindow", "monthly_window"]         let rolling = self.firstDict(from: dict, keys: rollingKeys)        let weekly = self.firstDict(from: dict, keys: weeklyKeys)        let monthly = self.firstDict(from: dict, keys: monthlyKeys)         guard let rolling, let weekly else { return nil }         return self.buildSnapshot(rolling: rolling, weekly: weekly, monthly: monthly, now: now, renewsAt: renewsAt)    }     private static func parseUsageNested(        _ dict: [String: Any],        now: Date,        depth: Int,        inheritedRenewsAt: Date?) -> OpenCodeGoUsageSnapshot?    {        if depth > 3 { return nil }        let renewsAt = self.dateValue(from: self.value(from: dict, keys: self.renewAtKeys)) ?? inheritedRenewsAt        var rolling: [String: Any]?        var weekly: [String: Any]?        var monthly: [String: Any]?         for (key, value) in dict {            guard let sub = value as? [String: Any] else { continue }            let lower = key.lowercased()            if lower.contains("rolling") || lower.contains("hour") || lower.contains("5h") || lower.contains("5-hour") {                rolling = sub            } else if lower.contains("weekly") || lower.contains("week") {                weekly = sub            } else if lower.contains("monthly") || lower.contains("month") {                monthly = sub            }        }         if let rolling, let weekly {            let snapshot = self.buildSnapshot(                rolling: rolling,                weekly: weekly,                monthly: monthly,                now: now,                renewsAt: renewsAt)            if let snapshot { return snapshot }        }         for value in dict.values {            if let sub = value as? [String: Any],               let snapshot = self.parseUsageNested(                   sub,                   now: now,                   depth: depth + 1,                   inheritedRenewsAt: renewsAt)            {                return snapshot            }        }         return nil    }     private static func parseUsageFromCandidates(        object: Any,        now: Date,        inheritedRenewsAt: Date? = nil) -> OpenCodeGoUsageSnapshot?    {        let candidates = self.collectWindowCandidates(object: object, now: now)        guard !candidates.isEmpty else { return nil }         let rollingCandidates = candidates.filter { candidate in            candidate.pathLower.contains("rolling") ||                candidate.pathLower.contains("hour") ||                candidate.pathLower.contains("5h") ||                candidate.pathLower.contains("5-hour")        }        let weeklyCandidates = candidates.filter { candidate in            candidate.pathLower.contains("weekly") ||                candidate.pathLower.contains("week")        }        let monthlyCandidates = candidates.filter { candidate in            candidate.pathLower.contains("monthly") ||                candidate.pathLower.contains("month")        }         let rolling = self.pickCandidate(            preferred: rollingCandidates,            fallback: candidates,

Lines 1–500 of 909.