Peter Steinbergerfix: cache provider cookies680ce0910 months agoHistory
1360 lines · 51.0 KB · Swift
RawDownload
import Foundationimport SweetCookieKit #if os(macOS) private let factoryCookieImportOrder: BrowserCookieImportOrder =    ProviderDefaults.metadata[.factory]?.browserCookieOrder ?? Browser.defaultImportOrder // MARK: - Factory Cookie Importer /// Imports Factory session cookies from browser cookies.public enum FactoryCookieImporter {    private static let cookieClient = BrowserCookieClient()    private static let sessionCookieNames: Set<String> = [        "wos-session",        "__Secure-next-auth.session-token",        "next-auth.session-token",        "__Secure-authjs.session-token",        "__Host-authjs.csrf-token",        "authjs.session-token",        "session",        "access-token",    ]     private static let authSessionCookieNames: Set<String> = [        "__Secure-next-auth.session-token",        "next-auth.session-token",        "__Secure-authjs.session-token",        "authjs.session-token",    ]    private static let appBaseURL = URL(string: "https://app.factory.ai")!    private static let authBaseURL = URL(string: "https://auth.factory.ai")!    private static let apiBaseURL = URL(string: "https://api.factory.ai")!     public struct SessionInfo: Sendable {        public let cookies: [HTTPCookie]        public let sourceLabel: String         public init(cookies: [HTTPCookie], sourceLabel: String) {            self.cookies = cookies            self.sourceLabel = sourceLabel        }         public var cookieHeader: String {            self.cookies.map { "\($0.name)=\($0.value)" }.joined(separator: "; ")        }    }     /// Returns all Factory sessions across supported browsers.    public static func importSessions(        browserDetection: BrowserDetection,        logger: ((String) -> Void)? = nil) throws -> [SessionInfo]    {        let log: (String) -> Void = { msg in logger?("[factory-cookie] \(msg)") }        var sessions: [SessionInfo] = []         // Filter to cookie-eligible browsers to avoid unnecessary keychain prompts        let installedBrowsers = factoryCookieImportOrder.cookieImportCandidates(using: browserDetection)        for browserSource in installedBrowsers {            do {                let perSource = try self.importSessions(from: browserSource, logger: logger)                sessions.append(contentsOf: perSource)            } catch {                log("\(browserSource.displayName) cookie import failed: \(error.localizedDescription)")            }        }         guard !sessions.isEmpty else {            throw FactoryStatusProbeError.noSessionCookie        }        return sessions    }     public static func importSessions(        from browserSource: Browser,        logger: ((String) -> Void)? = nil) throws -> [SessionInfo]    {        let log: (String) -> Void = { msg in logger?("[factory-cookie] \(msg)") }        let cookieDomains = ["factory.ai", "app.factory.ai", "auth.factory.ai"]        let query = BrowserCookieQuery(domains: cookieDomains)        let sources = try Self.cookieClient.records(            matching: query,            in: browserSource,            logger: log)         var sessions: [SessionInfo] = []        for source in sources where !source.records.isEmpty {            let httpCookies = BrowserCookieClient.makeHTTPCookies(source.records, origin: query.origin)            if httpCookies.contains(where: { Self.sessionCookieNames.contains($0.name) }) {                log("Found \(httpCookies.count) Factory cookies in \(source.label)")                log("\(source.label) cookie names: \(self.cookieNames(from: httpCookies))")                if let token = httpCookies.first(where: { $0.name == "access-token" })?.value {                    let hint = token.contains(".") ? "jwt" : "opaque"                    log("\(source.label) access-token cookie: \(token.count) chars (\(hint))")                }                if let token = httpCookies.first(where: { self.authSessionCookieNames.contains($0.name) })?.value {                    let hint = token.contains(".") ? "jwt" : "opaque"                    log("\(source.label) session cookie: \(token.count) chars (\(hint))")                }                sessions.append(SessionInfo(cookies: httpCookies, sourceLabel: source.label))            } else {                log("\(source.label) cookies found, but no Factory session cookie present")            }        }        return sessions    }     /// Attempts to import Factory cookies using the standard browser import order.    public static func importSession(        browserDetection: BrowserDetection,        logger: ((String) -> Void)? = nil) throws -> SessionInfo    {        let sessions = try self.importSessions(browserDetection: browserDetection, logger: logger)        guard let first = sessions.first else {            throw FactoryStatusProbeError.noSessionCookie        }        return first    }     /// Check if Factory session cookies are available    public static func hasSession(browserDetection: BrowserDetection, logger: ((String) -> Void)? = nil) -> Bool {        do {            return try !(self.importSessions(browserDetection: browserDetection, logger: logger)).isEmpty        } catch {            return false        }    }     private static func cookieNames(from cookies: [HTTPCookie]) -> String {        let names = Set(cookies.map { "\($0.name)@\($0.domain)" }).sorted()        return names.joined(separator: ", ")    }} // MARK: - Factory API Models public struct FactoryAuthResponse: Codable, Sendable {    public let featureFlags: FactoryFeatureFlags?    public let organization: FactoryOrganization?} public struct FactoryFeatureFlags: Codable, Sendable {    public let flags: [String: Bool]?    public let configs: [String: AnyCodable]?} public struct FactoryOrganization: Codable, Sendable {    public let id: String?    public let name: String?    public let subscription: FactorySubscription?     enum CodingKeys: String, CodingKey {        case id        case name        case subscription    }} public struct FactorySubscription: Codable, Sendable {    public let factoryTier: String?    public let orbSubscription: FactoryOrbSubscription?} public struct FactoryOrbSubscription: Codable, Sendable {    public let plan: FactoryPlan?    public let status: String?} public struct FactoryPlan: Codable, Sendable {    public let name: String?    public let id: String?} public struct FactoryUsageResponse: Codable, Sendable {    public let usage: FactoryUsageData?    public let source: String?    public let userId: String?} public struct FactoryUsageData: Codable, Sendable {    public let startDate: Int64?    public let endDate: Int64?    public let standard: FactoryTokenUsage?    public let premium: FactoryTokenUsage?} public struct FactoryTokenUsage: Codable, Sendable {    public let userTokens: Int64?    public let orgTotalTokensUsed: Int64?    public let totalAllowance: Int64?    public let usedRatio: Double?    public let orgOverageUsed: Int64?    public let basicAllowance: Int64?    public let orgOverageLimit: Int64?} /// Helper for encoding arbitrary JSONpublic struct AnyCodable: Codable, Sendable {    public init(from decoder: Decoder) throws {        let container = try decoder.singleValueContainer()        if container.decodeNil() {            return        }        _ = try? container.decode([String: AnyCodable].self)    }     public func encode(to encoder: Encoder) throws {        var container = encoder.singleValueContainer()        try container.encodeNil()    }} // MARK: - Factory Status Snapshot public struct FactoryStatusSnapshot: Sendable {    /// Standard token usage (user)    public let standardUserTokens: Int64    /// Standard token usage (org total)    public let standardOrgTokens: Int64    /// Standard token allowance    public let standardAllowance: Int64    /// Premium token usage (user)    public let premiumUserTokens: Int64    /// Premium token usage (org total)    public let premiumOrgTokens: Int64    /// Premium token allowance    public let premiumAllowance: Int64    /// Billing period start    public let periodStart: Date?    /// Billing period end    public let periodEnd: Date?    /// Plan name    public let planName: String?    /// Factory tier (enterprise, team, etc.)    public let tier: String?    /// Organization name    public let organizationName: String?    /// User email    public let accountEmail: String?    /// User ID    public let userId: String?    /// Raw JSON for debugging    public let rawJSON: String?     public init(        standardUserTokens: Int64,        standardOrgTokens: Int64,        standardAllowance: Int64,        premiumUserTokens: Int64,        premiumOrgTokens: Int64,        premiumAllowance: Int64,        periodStart: Date?,        periodEnd: Date?,        planName: String?,        tier: String?,        organizationName: String?,        accountEmail: String?,        userId: String?,        rawJSON: String?)    {        self.standardUserTokens = standardUserTokens        self.standardOrgTokens = standardOrgTokens        self.standardAllowance = standardAllowance        self.premiumUserTokens = premiumUserTokens        self.premiumOrgTokens = premiumOrgTokens        self.premiumAllowance = premiumAllowance        self.periodStart = periodStart        self.periodEnd = periodEnd        self.planName = planName        self.tier = tier        self.organizationName = organizationName        self.accountEmail = accountEmail        self.userId = userId        self.rawJSON = rawJSON    }     /// Convert to UsageSnapshot for the common provider interface    public func toUsageSnapshot() -> UsageSnapshot {        // Primary: Standard tokens used (as percentage of allowance, capped reasonably)        let standardPercent = self.calculateUsagePercent(            used: self.standardUserTokens,            allowance: self.standardAllowance)         let primary = RateWindow(            usedPercent: standardPercent,            windowMinutes: nil,            resetsAt: self.periodEnd,            resetDescription: self.periodEnd.map { Self.formatResetDate($0) })         // Secondary: Premium tokens used        let premiumPercent = self.calculateUsagePercent(            used: self.premiumUserTokens,            allowance: self.premiumAllowance)         let secondary = RateWindow(            usedPercent: premiumPercent,            windowMinutes: nil,            resetsAt: self.periodEnd,            resetDescription: self.periodEnd.map { Self.formatResetDate($0) })         // Format login method as tier + plan        let loginMethod: String? = {            var parts: [String] = []            if let tier = self.tier, !tier.isEmpty {                parts.append("Factory \(tier.capitalized)")            }            if let plan = self.planName, !plan.isEmpty, !plan.lowercased().contains("factory") {                parts.append(plan)            }            return parts.isEmpty ? nil : parts.joined(separator: " - ")        }()         let identity = ProviderIdentitySnapshot(            providerID: .factory,            accountEmail: self.accountEmail,            accountOrganization: self.organizationName,            loginMethod: loginMethod)        return UsageSnapshot(            primary: primary,            secondary: secondary,            tertiary: nil,            providerCost: nil,            updatedAt: Date(),            identity: identity)    }     private func calculateUsagePercent(used: Int64, allowance: Int64) -> Double {        // Treat very large allowances (> 1 trillion) as unlimited        let unlimitedThreshold: Int64 = 1_000_000_000_000        if allowance > unlimitedThreshold {            // For unlimited, show a token count-based pseudo-percentage (capped at 100%)            // Use 100M tokens as a reference point for "100%"            let referenceTokens: Double = 100_000_000            return min(100, Double(used) / referenceTokens * 100)        }        guard allowance > 0 else { return 0 }        return min(100, Double(used) / Double(allowance) * 100)    }     private static func formatResetDate(_ date: Date) -> String {        let formatter = DateFormatter()        formatter.dateFormat = "MMM d 'at' h:mma"        formatter.locale = Locale(identifier: "en_US_POSIX")        return "Resets " + formatter.string(from: date)    }} // MARK: - Factory Status Probe Error public enum FactoryStatusProbeError: LocalizedError, Sendable {    case notLoggedIn    case networkError(String)    case parseFailed(String)    case noSessionCookie     public var errorDescription: String? {        switch self {        case .notLoggedIn:            "Not logged in to Factory. Please log in via the CodexBar menu."        case let .networkError(msg):            "Factory API error: \(msg)"        case let .parseFailed(msg):            "Could not parse Factory usage: \(msg)"        case .noSessionCookie:            "No Factory session found. Please log in to app.factory.ai in \(factoryCookieImportOrder.loginHint)."        }    }} // MARK: - Factory Session Store public actor FactorySessionStore {    public static let shared = FactorySessionStore()     private var sessionCookies: [HTTPCookie] = []    private var bearerToken: String?    private var refreshToken: String?    private let fileURL: URL    private var didLoadFromDisk = false     private init() {        let fm = FileManager.default        let appSupport = fm.urls(for: .applicationSupportDirectory, in: .userDomainMask).first            ?? fm.temporaryDirectory        let dir = appSupport.appendingPathComponent("CodexBar", isDirectory: true)        try? fm.createDirectory(at: dir, withIntermediateDirectories: true)        self.fileURL = dir.appendingPathComponent("factory-session.json")    }     public func setCookies(_ cookies: [HTTPCookie]) {        self.didLoadFromDisk = true        self.sessionCookies = cookies        self.saveToDisk()    }     public func getCookies() -> [HTTPCookie] {        self.loadFromDiskIfNeeded()        return self.sessionCookies    }     public func setBearerToken(_ token: String?) {        self.didLoadFromDisk = true        self.bearerToken = token        self.saveToDisk()    }     public func getBearerToken() -> String? {        self.loadFromDiskIfNeeded()        return self.bearerToken    }     public func setRefreshToken(_ token: String?) {        self.didLoadFromDisk = true        self.refreshToken = token        self.saveToDisk()    }     public func getRefreshToken() -> String? {        self.loadFromDiskIfNeeded()        return self.refreshToken    }     public func clearSession() {        self.didLoadFromDisk = true        self.sessionCookies = []        self.bearerToken = nil        self.refreshToken = nil        try? FileManager.default.removeItem(at: self.fileURL)    }     public func hasValidSession() -> Bool {        self.loadFromDiskIfNeeded()        return !self.sessionCookies.isEmpty || self.bearerToken != nil || self.refreshToken != nil    }     private func saveToDisk() {        let cookieData = self.sessionCookies.compactMap { cookie -> [String: Any]? in            guard let props = cookie.properties else { return nil }            var serializable: [String: Any] = [:]            for (key, value) in props {                let keyString = key.rawValue                if let date = value as? Date {                    serializable[keyString] = date.timeIntervalSince1970                    serializable[keyString + "_isDate"] = true                } else if let url = value as? URL {                    serializable[keyString] = url.absoluteString                    serializable[keyString + "_isURL"] = true                } else if JSONSerialization.isValidJSONObject([value]) ||                    value is String ||                    value is Bool ||                    value is NSNumber                {                    serializable[keyString] = value                }            }            return serializable        }         var payload: [String: Any] = [:]        if !cookieData.isEmpty {            payload["cookies"] = cookieData        }        if let bearerToken {            payload["bearerToken"] = bearerToken        }        if let refreshToken {            payload["refreshToken"] = refreshToken        }         guard !payload.isEmpty,              let data = try? JSONSerialization.data(withJSONObject: payload, options: [.prettyPrinted])        else {            return        }        try? data.write(to: self.fileURL)    }     private func loadFromDisk() {        guard let data = try? Data(contentsOf: self.fileURL),              let json = try? JSONSerialization.jsonObject(with: data)        else { return }         var cookieArray: [[String: Any]] = []        if let dict = json as? [String: Any] {            if let stored = dict["cookies"] as? [[String: Any]] {                cookieArray = stored            }            self.bearerToken = dict["bearerToken"] as? String            self.refreshToken = dict["refreshToken"] as? String        } else if let stored = json as? [[String: Any]] {            cookieArray = stored        }         self.sessionCookies = cookieArray.compactMap { props in            var cookieProps: [HTTPCookiePropertyKey: Any] = [:]            for (key, value) in props {                if key.hasSuffix("_isDate") || key.hasSuffix("_isURL") { continue }                 let propKey = HTTPCookiePropertyKey(key) 

Lines 1–500 of 1360.