Peter Steinbergerstyle: apply release formattingb1d3e405 months agoHistory
801 lines · 31.6 KB · Swift
RawDownload
import Foundation#if canImport(FoundationNetworking)import FoundationNetworking#endif#if os(macOS)import SweetCookieKit#endif public struct CopilotBudgetWebFetcher: Sendable {    public enum Error: Swift.Error, LocalizedError, Equatable {        case noSessionCookie        case notLoggedIn        case accountMismatch(expected: String, actual: String?)        case badStatus(Int)        case invalidResponse         public var errorDescription: String? {            switch self {            case .noSessionCookie:                "No GitHub browser session cookie found."            case .notLoggedIn:                "GitHub browser session is not logged in."            case let .accountMismatch(expected, actual):                "GitHub browser session belongs to \(actual ?? "an unknown account"), expected \(expected)."            case let .badStatus(status):                "GitHub budgets request failed with HTTP \(status)."            case .invalidResponse:                "GitHub budgets response could not be decoded."            }        }    }     struct BudgetResponse: Decodable {        let budgets: [Budget]        let hasNextPage: Bool?         private enum CodingKeys: String, CodingKey {            case budgets            case payload            case hasNextPage            case hasNextPageSnake = "has_next_page"        }         init(budgets: [Budget], hasNextPage: Bool? = nil) {            self.budgets = budgets            self.hasNextPage = hasNextPage        }         init(from decoder: Decoder) throws {            let container = try decoder.container(keyedBy: CodingKeys.self)            if let payload = try container.decodeIfPresent(BudgetResponse.self, forKey: .payload) {                self = payload                return            }            self.budgets = try container.decodeIfPresent([Budget].self, forKey: .budgets) ?? []            self.hasNextPage = try container.decodeIfPresent(Bool.self, forKey: .hasNextPage)                ?? container.decodeIfPresent(Bool.self, forKey: .hasNextPageSnake)        }    }     struct Budget: Decodable, Equatable {        let id: String?        let name: String?        let budgetType: String?        let budgetProductSkus: [String]        let budgetScope: String?        let budgetEntityName: String?        let budgetAmount: Double        let currentAmount: Double         init(            id: String? = nil,            name: String? = nil,            budgetType: String? = nil,            budgetProductSkus: [String] = [],            budgetScope: String? = nil,            budgetEntityName: String? = nil,            budgetAmount: Double,            currentAmount: Double = 0)        {            self.id = id            self.name = name            self.budgetType = budgetType            self.budgetProductSkus = budgetProductSkus            self.budgetScope = budgetScope            self.budgetEntityName = budgetEntityName            self.budgetAmount = budgetAmount            self.currentAmount = currentAmount        }         init(from decoder: Decoder) throws {            let container = try decoder.container(keyedBy: DynamicCodingKey.self)            self.id = Self.decodeString(container: container, keys: ["id", "uuid", "budget_id", "budgetId"])            self.name = Self.decodeString(container: container, keys: ["name", "display_name", "displayName", "title"])            self.budgetType = Self.decodeString(                container: container,                keys: ["budget_type", "budgetType", "type", "pricing_target_type", "pricingTargetType"])            self.budgetProductSkus = Self.decodeStringArray(                container: container,                keys: [                    "budget_product_skus",                    "budgetProductSkus",                    "budget_product_sku",                    "budgetProductSku",                    "product_skus",                    "productSkus",                    "skus",                    "sku",                    "product",                    "product_name",                    "productName",                    "pricing_target_id",                    "pricingTargetId",                ])            self.budgetScope = Self.decodeString(container: container, keys: ["budget_scope", "budgetScope", "scope"])            self.budgetEntityName = Self.decodeString(                container: container,                keys: [                    "budget_entity_name",                    "budgetEntityName",                    "entity_name",                    "entityName",                    "target_name",                    "targetName",                ])            self.budgetAmount = Self.decodeDouble(                container: container,                keys: [                    "budget_amount",                    "budgetAmount",                    "target_amount",                    "targetAmount",                    "spending_limit",                    "spendingLimit",                    "limit",                    "amount",                    "max",                ]) ?? 0            self.currentAmount = Self.decodeDouble(                container: container,                keys: [                    "current_usage",                    "currentUsage",                    "current_amount",                    "currentAmount",                    "usage_amount",                    "usageAmount",                    "usage",                    "spent",                    "amount_used",                    "amountUsed",                ]) ?? 0        }         var normalizedSelectors: Set<String> {            let values = self.budgetProductSkus + [                self.budgetType,                self.budgetEntityName,                self.name,            ].compactMap(\.self)            return Set(values.compactMap(CopilotBudgetWebFetcher.normalizedBillingIdentifier))        }         private static func decodeString(            container: KeyedDecodingContainer<DynamicCodingKey>,            keys: [String]) -> String?        {            for key in keys {                guard let codingKey = DynamicCodingKey(key) else { continue }                if let value = try? container.decodeIfPresent(String.self, forKey: codingKey), !value.isEmpty {                    return value                }                if let value = try? container.decodeIfPresent(Int.self, forKey: codingKey) {                    return String(value)                }            }            return nil        }         private static func decodeStringArray(            container: KeyedDecodingContainer<DynamicCodingKey>,            keys: [String]) -> [String]        {            for key in keys {                guard let codingKey = DynamicCodingKey(key) else { continue }                if let values = try? container.decodeIfPresent([String].self, forKey: codingKey), !values.isEmpty {                    return values                }                if let value = try? container.decodeIfPresent(String.self, forKey: codingKey), !value.isEmpty {                    return [value]                }                if let values = try? container.decodeIfPresent([ProductSKU].self, forKey: codingKey),                   !values.isEmpty                {                    return values.flatMap(\.selectors)                }            }            return []        }         private static func decodeDouble(            container: KeyedDecodingContainer<DynamicCodingKey>,            keys: [String]) -> Double?        {            for key in keys {                guard let codingKey = DynamicCodingKey(key) else { continue }                if let value = try? container.decodeIfPresent(Double.self, forKey: codingKey) {                    return value                }                if let value = try? container.decodeIfPresent(Int.self, forKey: codingKey) {                    return Double(value)                }                if let value = try? container.decodeIfPresent(String.self, forKey: codingKey),                   let parsed = Self.parseAmount(value)                {                    return parsed                }                if let value = try? container.decodeIfPresent(AmountValue.self, forKey: codingKey) {                    return value.amount                }            }            return nil        }         fileprivate static func parseAmount(_ value: String) -> Double? {            let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)            let isNegative = trimmed.first == "-"            guard !trimmed.dropFirst(isNegative ? 1 : 0).contains("-") else { return nil }            let unsigned = trimmed.filter { $0.isNumber || $0 == "." }            guard !unsigned.isEmpty else { return nil }            return Double(isNegative ? "-\(unsigned)" : unsigned)        }    }     struct GitHubWebIdentity: Equatable {        let id: String?        let login: String?         init(id: String?, login: String?) {            let trimmedID = id?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""            let trimmedLogin = login?.trimmingCharacters(in: .whitespacesAndNewlines) ?? ""            self.id = trimmedID.isEmpty ? nil : trimmedID            self.login = trimmedLogin.isEmpty ? nil : trimmedLogin        }         var displayName: String? {            self.login ?? self.id.map { "github:user:\($0)" }        }    }     private struct ProductSKU: Decodable, Equatable {        let selectors: [String]         init(from decoder: Decoder) throws {            let container = try decoder.container(keyedBy: DynamicCodingKey.self)            self.selectors = [                "sku",                "name",                "display_name",                "displayName",                "product",                "product_name",                "productName",            ].compactMap { key in                guard let codingKey = DynamicCodingKey(key) else { return nil }                return try? container.decodeIfPresent(String.self, forKey: codingKey)            }        }    }     private struct AmountValue: Decodable {        let amount: Double?         init(from decoder: Decoder) throws {            let container = try decoder.container(keyedBy: DynamicCodingKey.self)            self.amount = [                "amount",                "value",                "total",                "cents",                "formatted",            ].lazy.compactMap { key -> Double? in                guard let codingKey = DynamicCodingKey(key) else { return nil }                if let value = try? container.decodeIfPresent(Double.self, forKey: codingKey) {                    return key == "cents" ? value / 100 : value                }                if let value = try? container.decodeIfPresent(Int.self, forKey: codingKey) {                    return key == "cents" ? Double(value) / 100 : Double(value)                }                if let value = try? container.decodeIfPresent(String.self, forKey: codingKey) {                    return Budget.parseAmount(value)                }                return nil            }.first        }    }     private struct DynamicCodingKey: CodingKey {        let stringValue: String        let intValue: Int?         init?(_ stringValue: String) {            self.stringValue = stringValue            self.intValue = nil        }         init?(stringValue: String) {            self.init(stringValue)        }         init?(intValue: Int) {            self.stringValue = String(intValue)            self.intValue = intValue        }    }     private static let copilotProductID = "copilot"    private static let copilotPremiumRequestSKU = "copilot_premium_request"    private static let copilotAgentPremiumRequestSKU = "copilot_agent_premium_request"    private static let sparkPremiumRequestSKU = "spark_premium_request"    private static let copilotBudgetSelectors: Set<String> = [        copilotProductID,        copilotPremiumRequestSKU,        copilotAgentPremiumRequestSKU,        sparkPremiumRequestSKU,    ]     private let cookieHeaderOverride: String?    private let expectedGitHubAccountIdentifier: String?    private let browserDetection: BrowserDetection    private let transport: any ProviderHTTPTransport    private let now: @Sendable () -> Date     public init(        cookieHeaderOverride: String? = nil,        expectedGitHubAccountIdentifier: String? = nil,        browserDetection: BrowserDetection = BrowserDetection(),        transport: any ProviderHTTPTransport = ProviderHTTPClient.shared,        now: @escaping @Sendable () -> Date = { Date() })    {        self.cookieHeaderOverride = CookieHeaderNormalizer.normalize(cookieHeaderOverride ?? "")        self.expectedGitHubAccountIdentifier = Self.normalizedExpectedAccountIdentifier(            expectedGitHubAccountIdentifier)        self.browserDetection = browserDetection        self.transport = transport        self.now = now    }     public func fetchBudgetWindows() async throws -> [NamedRateWindow] {        if let cookieHeaderOverride, !cookieHeaderOverride.isEmpty {            return try await self.fetchBudgetWindows(cookieHeader: cookieHeaderOverride)        }         if let cached = CookieHeaderCache.load(provider: .copilot),           !cached.cookieHeader.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty        {            do {                return try await self.fetchBudgetWindows(cookieHeader: cached.cookieHeader)            } catch {                if case Error.notLoggedIn = error {                    CookieHeaderCache.clear(provider: .copilot)                } else if case Error.accountMismatch = error {                    CookieHeaderCache.clear(provider: .copilot)                } else {                    throw error                }            }        }         #if os(macOS)        for session in CopilotGitHubCookieImporter.importSessions(browserDetection: self.browserDetection) {            do {                let windows = try await self.fetchBudgetWindows(cookieHeader: session.cookieHeader)                CookieHeaderCache.store(                    provider: .copilot,                    cookieHeader: session.cookieHeader,                    sourceLabel: session.sourceLabel)                return windows            } catch {                if case Error.notLoggedIn = error {                    continue                }                if case Error.accountMismatch = error {                    continue                }                throw error            }        }        #endif         throw Error.noSessionCookie    }     func fetchBudgetWindows(cookieHeader: String) async throws -> [NamedRateWindow] {        let nonce = try await self.boundFetchNonce(cookieHeader: cookieHeader)        var allBudgets: [Budget] = []        var page = 1        let maxPages = 20        var shouldContinue = true        while shouldContinue, page <= maxPages {            let response = try await self.fetchBudgetPage(                cookieHeader: cookieHeader,                nonce: nonce,                page: page)            allBudgets.append(contentsOf: response.budgets)            shouldContinue = response.hasNextPage == true            page += 1        }        if shouldContinue {            CodexBarLog.logger(LogCategories.providers).warning(                "Copilot budget pagination reached page cap",                metadata: ["pageCap": "\(maxPages)"])        }        return Self.extraRateWindows(from: allBudgets, now: self.now())    }     private func boundFetchNonce(cookieHeader: String) async throws -> String? {        guard self.expectedGitHubAccountIdentifier != nil else {            return await self.bestEffortFetchNonce(cookieHeader: cookieHeader)        }        let metadata = try await self.fetchBudgetPageMetadata(cookieHeader: cookieHeader)        try self.verifyExpectedGitHubAccount(metadata.identity)        return metadata.nonce    }     private func bestEffortFetchNonce(cookieHeader: String) async -> String? {        do {            return try await self.fetchBudgetPageMetadata(cookieHeader: cookieHeader).nonce        } catch {            // GitHub accepts some budget requests without a nonce. Keep auth failures and            // page-shape changes non-fatal so a valid Cookie header can still be tried.            return nil        }    }     private struct BudgetPageMetadata {        let nonce: String?        let identity: GitHubWebIdentity?    }     private func fetchBudgetPageMetadata(cookieHeader: String) async throws -> BudgetPageMetadata {        guard let url = URL(string: "https://github.com/settings/billing/budgets") else {            throw URLError(.badURL)        }        var request = URLRequest(url: url)        request.timeoutInterval = 15        request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")        request.setValue("text/html,application/xhtml+xml", forHTTPHeaderField: "Accept")        request.setValue("CodexBar", forHTTPHeaderField: "User-Agent")         let response = try await self.transport.response(for: request)        switch response.statusCode {        case 200:            guard let html = String(data: response.data, encoding: .utf8) else { throw Error.invalidResponse }            return BudgetPageMetadata(                nonce: Self.extractFetchNonce(from: html),                identity: Self.extractGitHubWebIdentity(from: html))        case 401, 403:            throw Error.notLoggedIn        default:            throw Error.badStatus(response.statusCode)        }    }     private func verifyExpectedGitHubAccount(_ actual: GitHubWebIdentity?) throws {        guard let expected = self.expectedGitHubAccountIdentifier else { return }        guard let actual else { throw Error.accountMismatch(expected: expected, actual: nil) }        guard Self.webIdentity(actual, matches: expected) else {            throw Error.accountMismatch(expected: expected, actual: actual.displayName)        }    }     private func fetchBudgetPage(cookieHeader: String, nonce: String?, page: Int) async throws -> BudgetResponse {        guard var components = URLComponents(string: "https://github.com/settings/billing/budgets") else {            throw URLError(.badURL)        }        components.queryItems = [            URLQueryItem(name: "page", value: "\(page)"),            URLQueryItem(name: "page_size", value: "10"),            URLQueryItem(name: "scope", value: "customer"),        ]        guard let url = components.url else { throw URLError(.badURL) }        var request = URLRequest(url: url)        request.timeoutInterval = 15        request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")        request.setValue("application/json", forHTTPHeaderField: "Accept")        request.setValue("https://github.com/settings/billing/budgets", forHTTPHeaderField: "Referer")        request.setValue("XMLHttpRequest", forHTTPHeaderField: "X-Requested-With")        request.setValue("true", forHTTPHeaderField: "GitHub-Verified-Fetch")        request.setValue("CodexBar", forHTTPHeaderField: "User-Agent")        if let nonce, !nonce.isEmpty {            request.setValue(nonce, forHTTPHeaderField: "X-Fetch-Nonce")        }         let response = try await self.transport.response(for: request)        switch response.statusCode {        case 200:            do {                return try JSONDecoder().decode(BudgetResponse.self, from: response.data)            } catch {

Lines 1–500 of 801.