1045 lines · 34.8 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif public enum ChutesUsageError: LocalizedError, Sendable {    case missingCredentials    case invalidCredentials    case invalidURL    case apiError(String)    case parseFailed(String)     public var errorDescription: String? {        switch self {        case .missingCredentials:            "Missing Chutes API key. Set apiKey in ~/.codexbar/config.json or CHUTES_API_KEY."        case .invalidCredentials:            "Chutes API key was rejected. Check the API key in Settings."        case .invalidURL:            "Chutes usage URL is invalid."        case let .apiError(message):            "Chutes usage API error: \(message)"        case let .parseFailed(message):            "Chutes usage parse error: \(message)"        }    }} public enum ChutesSubscriptionState: String, Sendable, Codable {    case active    case inactive    case unknown} public struct ChutesQuotaWindow: Sendable, Equatable {    public let label: String?    public let used: Double?    public let limit: Double?    public let remaining: Double?    public let usedPercent: Double?    public let windowMinutes: Int?    public let resetsAt: Date?    public let unit: String?     public init(        label: String?,        used: Double?,        limit: Double?,        remaining: Double?,        usedPercent: Double?,        windowMinutes: Int?,        resetsAt: Date?,        unit: String? = nil)    {        self.label = label        self.used = used        self.limit = limit        self.remaining = remaining        self.usedPercent = usedPercent        self.windowMinutes = windowMinutes        self.resetsAt = resetsAt        self.unit = unit    }     public func rateWindow(defaultWindowMinutes: Int?) -> RateWindow? {        guard let percent = self.usagePercent else { return nil }        return RateWindow(            usedPercent: max(0, min(percent, 100)),            windowMinutes: self.windowMinutes ?? defaultWindowMinutes,            resetsAt: self.resetsAt,            resetDescription: self.usageDescription)    }     public var usagePercent: Double? {        if let usedPercent { return max(0, min(usedPercent, 100)) }         var used = self.used        var limit = self.limit        var remaining = self.remaining         if limit == nil, let used, let remaining {            limit = used + remaining        }        if used == nil, let limit, let remaining {            used = limit - remaining        }        if remaining == nil, let limit, let used {            remaining = max(0, limit - used)        }         guard let used, let limit, limit > 0 else { return nil }        return (used / limit) * 100    }     private var usageDescription: String? {        guard let limit, limit > 0 else { return nil }         let used: Double        if let explicitUsed = self.used {            used = explicitUsed        } else if let remaining {            used = max(0, limit - remaining)        } else {            return nil        }         let unit = self.unit?.trimmingCharacters(in: .whitespacesAndNewlines)        let suffix = if let unit, !unit.isEmpty {            " \(unit)"        } else {            ""        }        return "\(Self.formatAmount(used))/\(Self.formatAmount(limit))\(suffix)"    }     private static func formatAmount(_ value: Double) -> String {        let rounded = value.rounded()        if abs(value - rounded) < 0.0001 {            return String(Int(rounded))        }         var text = String(format: "%.2f", value)        while text.contains("."), text.hasSuffix("0") {            text.removeLast()        }        if text.hasSuffix(".") {            text.removeLast()        }        return text    }} public struct ChutesUsageSnapshot: Sendable, Equatable {    public let rollingWindow: ChutesQuotaWindow?    public let monthlyWindow: ChutesQuotaWindow?    public let fallbackWindows: [ChutesQuotaWindow]    public let subscriptionState: ChutesSubscriptionState    public let planName: String?    public let subscriptionRenewsAt: Date?    public let updatedAt: Date     public init(        rollingWindow: ChutesQuotaWindow?,        monthlyWindow: ChutesQuotaWindow?,        fallbackWindows: [ChutesQuotaWindow] = [],        subscriptionState: ChutesSubscriptionState,        planName: String?,        subscriptionRenewsAt: Date?,        updatedAt: Date)    {        self.rollingWindow = rollingWindow        self.monthlyWindow = monthlyWindow        self.fallbackWindows = fallbackWindows        self.subscriptionState = subscriptionState        self.planName = planName        self.subscriptionRenewsAt = subscriptionRenewsAt        self.updatedAt = updatedAt    }     public var hasUsageData: Bool {        self.toUsageSnapshot().hasRateLimitWindows    }     public func toUsageSnapshot() -> UsageSnapshot {        let fallbackRateWindows = self.fallbackWindows.compactMap {            $0.rateWindow(defaultWindowMinutes: $0.windowMinutes)        }        let monthly = self.monthlyWindow?.rateWindow(defaultWindowMinutes: Self.monthlyWindowMinutes)        let rolling = self.rollingWindow?.rateWindow(defaultWindowMinutes: Self.rollingWindowMinutes)         let primary = rolling ?? (monthly == nil ? fallbackRateWindows.first : nil)        let secondary: RateWindow? = if let monthly {            monthly        } else if rolling != nil {            fallbackRateWindows.first        } else if primary != nil {            fallbackRateWindows.dropFirst().first        } else {            nil        }         let identity = ProviderIdentitySnapshot(            providerID: .chutes,            accountEmail: nil,            accountOrganization: nil,            loginMethod: self.loginMethod(hasWindows: primary != nil || secondary != nil))         return UsageSnapshot(            primary: primary,            secondary: secondary,            tertiary: nil,            providerCost: nil,            subscriptionRenewsAt: self.subscriptionRenewsAt ?? monthly?.resetsAt,            updatedAt: self.updatedAt,            identity: identity)    }     func preservingSubscriptionContext(from fallback: ChutesUsageSnapshot) -> ChutesUsageSnapshot {        ChutesUsageSnapshot(            rollingWindow: fallback.rollingWindow ?? self.rollingWindow,            monthlyWindow: fallback.monthlyWindow ?? self.monthlyWindow,            fallbackWindows: fallback.fallbackWindows + self.fallbackWindows,            subscriptionState: fallback.subscriptionState,            planName: fallback.planName,            subscriptionRenewsAt: fallback.subscriptionRenewsAt,            updatedAt: fallback.updatedAt)    }     private func loginMethod(hasWindows: Bool) -> String? {        if let planName = self.planName?.trimmingCharacters(in: .whitespacesAndNewlines), !planName.isEmpty {            return planName        }        switch self.subscriptionState {        case .active:            return nil        case .inactive:            return "No active subscription"        case .unknown:            return hasWindows ? nil : "No usage data"        }    }     static let rollingWindowMinutes = 4 * 60    static let monthlyWindowMinutes = 30 * 24 * 60} public struct ChutesUsageFetcher: Sendable {    private static let requestTimeoutSeconds: TimeInterval = 15     public static func fetchUsage(        apiKey: String,        environment: [String: String] = ProcessInfo.processInfo.environment,        transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,        now: Date = Date()) async throws -> ChutesUsageSnapshot    {        guard let token = ChutesSettingsReader.cleaned(apiKey) else {            throw ChutesUsageError.missingCredentials        }        try ChutesSettingsReader.validateEndpointOverrides(environment: environment)         let baseURL = ChutesSettingsReader.apiURL(environment: environment)        let subscription = try await self.fetchSnapshot(            pathComponents: ["users", "me", "subscription_usage"],            apiKey: token,            baseURL: baseURL,            transport: transport,            now: now)         guard subscription.rollingWindow == nil || subscription.monthlyWindow == nil else {            return subscription        }         do {            let quotas = try await self.fetchQuotaSnapshot(                apiKey: token,                baseURL: baseURL,                transport: transport,                now: now)            return quotas.hasUsageData ? quotas.preservingSubscriptionContext(from: subscription) : subscription        } catch ChutesUsageError.invalidCredentials {            throw ChutesUsageError.invalidCredentials        } catch {            return subscription        }    }     private static func fetchSnapshot(        pathComponents: [String],        apiKey: String,        baseURL: URL,        transport: any ProviderHTTPTransport,        now: Date) async throws -> ChutesUsageSnapshot    {        let data = try await self.fetchData(            pathComponents: pathComponents,            apiKey: apiKey,            baseURL: baseURL,            transport: transport)        do {            return try ChutesUsageParser.parse(data: data, now: now)        } catch let error as ChutesUsageError {            throw error        } catch {            throw ChutesUsageError.parseFailed(error.localizedDescription)        }    }     private static func fetchQuotaSnapshot(        apiKey: String,        baseURL: URL,        transport: any ProviderHTTPTransport,        now: Date) async throws -> ChutesUsageSnapshot    {        let data = try await self.fetchData(            pathComponents: ["users", "me", "quotas"],            apiKey: apiKey,            baseURL: baseURL,            transport: transport)        let fallback = try ChutesUsageParser.parse(data: data, now: now)        let definitions = try self.quotaDefinitions(from: data)        guard !definitions.isEmpty else { return fallback }         var enrichedDefinitions: [[String: Any]] = []        for definition in definitions {            guard let identifier = self.quotaIdentifier(from: definition) else {                enrichedDefinitions.append(definition)                continue            }             do {                let usageData = try await self.fetchData(                    pathComponents: ["users", "me", "quota_usage", identifier],                    apiKey: apiKey,                    baseURL: baseURL,                    transport: transport)                guard let usage = try self.responseDictionary(from: usageData) else {                    enrichedDefinitions.append(definition)                    continue                }                enrichedDefinitions.append(definition.merging(usage) { _, usageValue in usageValue })            } catch ChutesUsageError.invalidCredentials {                throw ChutesUsageError.invalidCredentials            } catch {                enrichedDefinitions.append(definition)            }        }         let enrichedData = try JSONSerialization.data(            withJSONObject: ["quotas": enrichedDefinitions],            options: [])        let enriched = try ChutesUsageParser.parse(data: enrichedData, now: now)        return enriched.hasUsageData ? enriched : fallback    }     private static func fetchData(        pathComponents: [String],        apiKey: String,        baseURL: URL,        transport: any ProviderHTTPTransport) async throws -> Data    {        let url = pathComponents.reduce(baseURL) { partial, component in            partial.appendingPathComponent(component)        }        guard url.scheme?.lowercased() == "https" else {            throw ChutesUsageError.invalidURL        }         var request = URLRequest(url: url)        request.httpMethod = "GET"        request.setValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")        request.setValue("application/json", forHTTPHeaderField: "Accept")        request.timeoutInterval = Self.requestTimeoutSeconds         let response = try await transport.response(for: request)        guard (200..<300).contains(response.statusCode) else {            if response.statusCode == 401 || response.statusCode == 403 {                throw ChutesUsageError.invalidCredentials            }            throw ChutesUsageError.apiError("HTTP \(response.statusCode): \(Self.responseSummary(response.data))")        }        return response.data    }     private static func quotaDefinitions(from data: Data) throws -> [[String: Any]] {        let object = try JSONSerialization.jsonObject(with: data, options: [])        if let array = object as? [Any] {            return array.compactMap { $0 as? [String: Any] }        }        guard let dictionary = object as? [String: Any] else { return [] }        if let quotas = dictionary["quotas"] as? [Any] {            return quotas.compactMap { $0 as? [String: Any] }        }        if let data = dictionary["data"] as? [Any] {            return data.compactMap { $0 as? [String: Any] }        }        if let data = dictionary["data"] as? [String: Any],           let quotas = data["quotas"] as? [Any]        {            return quotas.compactMap { $0 as? [String: Any] }        }        return []    }     private static func quotaIdentifier(from definition: [String: Any]) -> String? {        for key in ["chute_id", "chuteId", "id"] {            if let value = definition[key] as? String {                let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)                if !trimmed.isEmpty { return trimmed }            }            if let value = definition[key] as? NSNumber {                return value.stringValue            }        }        return nil    }     private static func responseDictionary(from data: Data) throws -> [String: Any]? {        guard let dictionary = try JSONSerialization.jsonObject(with: data, options: []) as? [String: Any] else {            return nil        }        return dictionary["data"] as? [String: Any]            ?? dictionary["result"] as? [String: Any]            ?? dictionary    }     private static func responseSummary(_ data: Data) -> String {        String(bytes: data.prefix(240), encoding: .utf8)?            .trimmingCharacters(in: .whitespacesAndNewlines)            ?? ""    }} private final class ChutesISO8601FormatterBox: @unchecked Sendable {    let lock = NSLock()    let withFractional: ISO8601DateFormatter = {        let formatter = ISO8601DateFormatter()        formatter.formatOptions = [.withInternetDateTime, .withFractionalSeconds]        return formatter    }()     let plain: ISO8601DateFormatter = {        let formatter = ISO8601DateFormatter()        formatter.formatOptions = [.withInternetDateTime]        return formatter    }()} private enum ChutesTimestampParser {    static let box = ChutesISO8601FormatterBox()     static func parse(_ value: Any?) -> Date? {        switch value {        case let date as Date:            date        case let number as NSNumber:            self.date(fromEpochValue: number.doubleValue)        case let text as String:            self.parseString(text)        default:            nil        }    }     private static func parseString(_ raw: String) -> Date? {        let text = raw.trimmingCharacters(in: .whitespacesAndNewlines)        guard !text.isEmpty else { return nil }        if let number = Double(text) {            return self.date(fromEpochValue: number)        }         self.box.lock.lock()        defer { self.box.lock.unlock() }        return self.box.withFractional.date(from: text) ?? self.box.plain.date(from: text)    }     private static func date(fromEpochValue value: Double) -> Date? {        guard value.isFinite, value > 0 else { return nil }        let seconds = value > 10_000_000_000 ? value / 1000 : value        return Date(timeIntervalSince1970: seconds)    }} enum ChutesUsageParser {    static func parse(data: Data, now: Date = Date()) throws -> ChutesUsageSnapshot {        let object = try JSONSerialization.jsonObject(with: data, options: [])        return self.parse(object: object, now: now)    }     private static func parse(object: Any, now: Date) -> ChutesUsageSnapshot {        let root: [String: Any] = {            if let dict = object as? [String: Any] { return dict }            if let array = object as? [Any] { return ["quotas": array] }            return [:]        }()        let dataRoot = self.dictionaryValue(self.value(in: root, keys: ["data", "result"])) ?? root        let subscription = self.subscriptionPayload(root: root, dataRoot: dataRoot)         let explicitRolling = self.firstDictionary(            in: root,            dataRoot: dataRoot,            keys: self.rollingPayloadKeys)            .flatMap {                self.parseQuota(                    $0,                    defaultLabel: "4-hour quota",                    defaultWindowMinutes: ChutesUsageSnapshot.rollingWindowMinutes)            }        let explicitMonthly = self.firstDictionary(            in: root,            dataRoot: dataRoot,            keys: self.monthlyPayloadKeys)            .flatMap {                self.parseQuota(                    $0,                    defaultLabel: "Monthly quota",                    defaultWindowMinutes: ChutesUsageSnapshot.monthlyWindowMinutes)            }         let quotaWindows = self.fallbackQuotaObjects(from: root, dataRoot: dataRoot).compactMap {            self.parseQuota($0, defaultLabel: nil, defaultWindowMinutes: nil)

Lines 1–500 of 1045.