1655 lines · 69.8 KB · Swift
RawDownload
import AppKitimport CodexBarCoreimport Foundationimport Observationimport SweetCookieKit // MARK: - Observation helpers @MainActorextension UsageStore {    var menuObservationToken: Int {        _ = self.snapshots        _ = self.errors        _ = self.lastSourceLabels        _ = self.lastFetchAttempts        _ = self.tokenSnapshots        _ = self.tokenErrors        _ = self.tokenRefreshInFlight        _ = self.credits        _ = self.lastCreditsError        _ = self.openAIDashboard        _ = self.lastOpenAIDashboardError        _ = self.openAIDashboardRequiresLogin        _ = self.openAIDashboardCookieImportStatus        _ = self.openAIDashboardCookieImportDebugLog        _ = self.codexVersion        _ = self.claudeVersion        _ = self.geminiVersion        _ = self.zaiVersion        _ = self.antigravityVersion        _ = self.kiroVersion        _ = self.isRefreshing        _ = self.refreshingProviders        _ = self.pathDebugInfo        _ = self.statuses        _ = self.probeLogs        return 0    }     func observeSettingsChanges() {        withObservationTracking {            _ = self.settings.refreshFrequency            _ = self.settings.statusChecksEnabled            _ = self.settings.sessionQuotaNotificationsEnabled            _ = self.settings.usageBarsShowUsed            _ = self.settings.costUsageEnabled            _ = self.settings.randomBlinkEnabled            _ = self.settings.claudeWebExtrasEnabled            _ = self.settings.codexUsageDataSource            _ = self.settings.claudeUsageDataSource            _ = self.settings.codexCookieSource            _ = self.settings.claudeCookieSource            _ = self.settings.cursorCookieSource            _ = self.settings.factoryCookieSource            _ = self.settings.minimaxCookieSource            _ = self.settings.codexCookieHeader            _ = self.settings.claudeCookieHeader            _ = self.settings.cursorCookieHeader            _ = self.settings.factoryCookieHeader            _ = self.settings.minimaxCookieHeader            _ = self.settings.mergeIcons            _ = self.settings.selectedMenuProvider            _ = self.settings.debugLoadingPattern            _ = self.settings.augmentCookieSource        } onChange: { [weak self] in            Task { @MainActor [weak self] in                guard let self else { return }                self.observeSettingsChanges()                self.startTimer()                self.restartAugmentKeepaliveIfNeeded()                await self.refresh()            }        }    }     private func restartAugmentKeepaliveIfNeeded() {        #if os(macOS)        let shouldRun = self.isEnabled(.augment)        let isRunning = self.augmentKeepalive != nil         if shouldRun, !isRunning {            self.startAugmentKeepalive()        } else if !shouldRun, isRunning {            Task { @MainActor in                self.augmentKeepalive?.stop()                self.augmentKeepalive = nil                print("[CodexBar] Augment session keepalive stopped (provider disabled)")            }        }        #endif    }} // MARK: - OpenAI web error messaging extension UsageStore {    private func openAIDashboardFriendlyError(        body: String,        targetEmail: String?,        cookieImportStatus: String?) -> String?    {        let trimmed = body.trimmingCharacters(in: .whitespacesAndNewlines)        let status = cookieImportStatus?.trimmingCharacters(in: .whitespacesAndNewlines)        if trimmed.isEmpty {            return [                "OpenAI web dashboard returned an empty page.",                "Sign in to chatgpt.com and update OpenAI cookies in Providers → Codex.",            ].joined(separator: " ")        }         let lower = trimmed.lowercased()        let looksLikePublicLanding = lower.contains("skip to content")            && (lower.contains("about") || lower.contains("openai") || lower.contains("chatgpt"))        let looksLoggedOut = lower.contains("sign in")            || lower.contains("log in")            || lower.contains("create account")            || lower.contains("continue with google")            || lower.contains("continue with apple")            || lower.contains("continue with microsoft")         guard looksLikePublicLanding || looksLoggedOut else { return nil }        let emailLabel = targetEmail?.trimmingCharacters(in: .whitespacesAndNewlines)        let targetLabel = (emailLabel?.isEmpty == false) ? emailLabel! : "your OpenAI account"        if let status, !status.isEmpty {            if status.contains("cookies do not match Codex account")                || status.localizedCaseInsensitiveContains("cookie import failed")            {                return [                    status,                    "Sign in to chatgpt.com as \(targetLabel), then update OpenAI cookies in Providers → Codex.",                ].joined(separator: " ")            }        }        return [            "OpenAI web dashboard returned a public page (not signed in).",            "Sign in to chatgpt.com as \(targetLabel), then update OpenAI cookies in Providers → Codex.",        ].joined(separator: " ")    }} enum ProviderStatusIndicator: String {    case none    case minor    case major    case critical    case maintenance    case unknown     var hasIssue: Bool {        switch self {        case .none: false        default: true        }    }     var label: String {        switch self {        case .none: "Operational"        case .minor: "Partial outage"        case .major: "Major outage"        case .critical: "Critical issue"        case .maintenance: "Maintenance"        case .unknown: "Status unknown"        }    }} #if DEBUGextension UsageStore {    func _setSnapshotForTesting(_ snapshot: UsageSnapshot?, provider: UsageProvider) {        self.snapshots[provider] = snapshot?.scoped(to: provider)    }     func _setTokenSnapshotForTesting(_ snapshot: CostUsageTokenSnapshot?, provider: UsageProvider) {        self.tokenSnapshots[provider] = snapshot    }     func _setTokenErrorForTesting(_ error: String?, provider: UsageProvider) {        self.tokenErrors[provider] = error    }     func _setErrorForTesting(_ error: String?, provider: UsageProvider) {        self.errors[provider] = error    }}#endif struct ProviderStatus {    let indicator: ProviderStatusIndicator    let description: String?    let updatedAt: Date?} /// Tracks consecutive failures so we can ignore a single flake when we previously had fresh data.struct ConsecutiveFailureGate {    private(set) var streak: Int = 0     mutating func recordSuccess() {        self.streak = 0    }     mutating func reset() {        self.streak = 0    }     /// Returns true when the caller should surface the error to the UI.    mutating func shouldSurfaceError(onFailureWithPriorData hadPriorData: Bool) -> Bool {        self.streak += 1        if hadPriorData, self.streak == 1 { return false }        return true    }} @MainActor@Observablefinal class UsageStore {    private(set) var snapshots: [UsageProvider: UsageSnapshot] = [:]    private(set) var errors: [UsageProvider: String] = [:]    private(set) var lastSourceLabels: [UsageProvider: String] = [:]    private(set) var lastFetchAttempts: [UsageProvider: [ProviderFetchAttempt]] = [:]    private(set) var tokenSnapshots: [UsageProvider: CostUsageTokenSnapshot] = [:]    private(set) var tokenErrors: [UsageProvider: String] = [:]    private(set) var tokenRefreshInFlight: Set<UsageProvider> = []    var credits: CreditsSnapshot?    var lastCreditsError: String?    var openAIDashboard: OpenAIDashboardSnapshot?    var lastOpenAIDashboardError: String?    private(set) var openAIDashboardRequiresLogin: Bool = false    var openAIDashboardCookieImportStatus: String?    var openAIDashboardCookieImportDebugLog: String?    var codexVersion: String?    var claudeVersion: String?    var geminiVersion: String?    var zaiVersion: String?    var antigravityVersion: String?    var cursorVersion: String?    var kiroVersion: String?    var isRefreshing = false    private(set) var refreshingProviders: Set<UsageProvider> = []    var debugForceAnimation = false    var pathDebugInfo: PathDebugSnapshot = .empty    private(set) var statuses: [UsageProvider: ProviderStatus] = [:]    private(set) var probeLogs: [UsageProvider: String] = [:]    @ObservationIgnored private var lastCreditsSnapshot: CreditsSnapshot?    @ObservationIgnored private var creditsFailureStreak: Int = 0    @ObservationIgnored private var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot?    @ObservationIgnored private var lastOpenAIDashboardTargetEmail: String?    @ObservationIgnored private var lastOpenAIDashboardCookieImportAttemptAt: Date?    @ObservationIgnored private var lastOpenAIDashboardCookieImportEmail: String?    @ObservationIgnored private var openAIWebAccountDidChange: Bool = false     @ObservationIgnored private let codexFetcher: UsageFetcher    @ObservationIgnored private let claudeFetcher: any ClaudeUsageFetching    @ObservationIgnored private let costUsageFetcher: CostUsageFetcher    @ObservationIgnored let browserDetection: BrowserDetection    @ObservationIgnored private let registry: ProviderRegistry    @ObservationIgnored private let settings: SettingsStore    @ObservationIgnored private let sessionQuotaNotifier: SessionQuotaNotifier    @ObservationIgnored private let sessionQuotaLogger = CodexBarLog.logger("sessionQuota")    @ObservationIgnored private let openAIWebLogger = CodexBarLog.logger("openai-web")    @ObservationIgnored private let tokenCostLogger = CodexBarLog.logger("token-cost")    @ObservationIgnored private var openAIWebDebugLines: [String] = []    @ObservationIgnored private var failureGates: [UsageProvider: ConsecutiveFailureGate] = [:]    @ObservationIgnored private var tokenFailureGates: [UsageProvider: ConsecutiveFailureGate] = [:]    @ObservationIgnored private var providerSpecs: [UsageProvider: ProviderSpec] = [:]    @ObservationIgnored private let providerMetadata: [UsageProvider: ProviderMetadata]    @ObservationIgnored private var timerTask: Task<Void, Never>?    @ObservationIgnored private var tokenTimerTask: Task<Void, Never>?    @ObservationIgnored private var tokenRefreshSequenceTask: Task<Void, Never>?    @ObservationIgnored private var lastKnownSessionRemaining: [UsageProvider: Double] = [:]    @ObservationIgnored private(set) var lastTokenFetchAt: [UsageProvider: Date] = [:]    @ObservationIgnored private let tokenFetchTTL: TimeInterval = 60 * 60    @ObservationIgnored private let tokenFetchTimeout: TimeInterval = 10 * 60    @ObservationIgnored private var augmentKeepalive: AugmentSessionKeepalive?     init(        fetcher: UsageFetcher,        browserDetection: BrowserDetection,        claudeFetcher: (any ClaudeUsageFetching)? = nil,        costUsageFetcher: CostUsageFetcher = CostUsageFetcher(),        settings: SettingsStore,        registry: ProviderRegistry = .shared,        sessionQuotaNotifier: SessionQuotaNotifier = SessionQuotaNotifier())    {        self.codexFetcher = fetcher        self.browserDetection = browserDetection        self.claudeFetcher = claudeFetcher ?? ClaudeUsageFetcher(browserDetection: browserDetection)        self.costUsageFetcher = costUsageFetcher        self.settings = settings        self.registry = registry        self.sessionQuotaNotifier = sessionQuotaNotifier        self.providerMetadata = registry.metadata        self            .failureGates = Dictionary(                uniqueKeysWithValues: UsageProvider.allCases                    .map { ($0, ConsecutiveFailureGate()) })        self.tokenFailureGates = Dictionary(            uniqueKeysWithValues: UsageProvider.allCases                .map { ($0, ConsecutiveFailureGate()) })        self.providerSpecs = registry.specs(            settings: settings,            metadata: self.providerMetadata,            codexFetcher: fetcher,            claudeFetcher: self.claudeFetcher,            browserDetection: browserDetection)        self.bindSettings()        self.detectVersions()        self.refreshPathDebugInfo()        LoginShellPathCache.shared.captureOnce { [weak self] _ in            Task { @MainActor in self?.refreshPathDebugInfo() }        }        Task { await self.refresh() }        self.startTimer()        self.startTokenTimer()        self.startAugmentKeepalive()    }     /// Returns the login method (plan type) for the specified provider, if available.    private func loginMethod(for provider: UsageProvider) -> String? {        self.snapshots[provider]?.loginMethod(for: provider)    }     /// Returns true if the Claude account appears to be a subscription (Max, Pro, Ultra, Team).    /// Returns false for API users or when plan cannot be determined.    func isClaudeSubscription() -> Bool {        Self.isSubscriptionPlan(self.loginMethod(for: .claude))    }     /// Determines if a login method string indicates a Claude subscription plan.    /// Known subscription indicators: Max, Pro, Ultra, Team (case-insensitive).    nonisolated static func isSubscriptionPlan(_ loginMethod: String?) -> Bool {        guard let method = loginMethod?.lowercased(), !method.isEmpty else {            return false        }        let subscriptionIndicators = ["max", "pro", "ultra", "team"]        return subscriptionIndicators.contains { method.contains($0) }    }     func version(for provider: UsageProvider) -> String? {        switch provider {        case .codex: self.codexVersion        case .claude: self.claudeVersion        case .zai: self.zaiVersion        case .gemini: self.geminiVersion        case .antigravity: self.antigravityVersion        case .cursor: self.cursorVersion        case .factory: nil        case .copilot: nil        case .minimax: nil        case .vertexai: nil        case .kiro: self.kiroVersion        case .augment: nil        }    }     var preferredSnapshot: UsageSnapshot? {        for provider in self.enabledProviders() {            if let snap = self.snapshots[provider] { return snap }        }        return nil    }     var iconStyle: IconStyle {        let enabled = self.enabledProviders()        if enabled.count > 1 { return .combined }        if let provider = enabled.first {            return self.style(for: provider)        }        return .codex    }     var isStale: Bool {        (self.isEnabled(.codex) && self.lastCodexError != nil) ||            (self.isEnabled(.claude) && self.lastClaudeError != nil) ||            (self.isEnabled(.zai) && self.errors[.zai] != nil) ||            (self.isEnabled(.gemini) && self.errors[.gemini] != nil) ||            (self.isEnabled(.antigravity) && self.errors[.antigravity] != nil) ||            (self.isEnabled(.cursor) && self.errors[.cursor] != nil) ||            (self.isEnabled(.factory) && self.errors[.factory] != nil) ||            (self.isEnabled(.copilot) && self.errors[.copilot] != nil) ||            (self.isEnabled(.minimax) && self.errors[.minimax] != nil)    }     func enabledProviders() -> [UsageProvider] {        // Use cached enablement to avoid repeated UserDefaults lookups in animation ticks.        let enabled = self.settings.enabledProvidersOrdered(metadataByProvider: self.providerMetadata)        return enabled.filter { self.isProviderAvailable($0) }    }     var statusChecksEnabled: Bool {        self.settings.statusChecksEnabled    }     func metadata(for provider: UsageProvider) -> ProviderMetadata {        self.providerMetadata[provider]!    }     private var codexBrowserCookieOrder: BrowserCookieImportOrder {        self.metadata(for: .codex).browserCookieOrder ?? Browser.defaultImportOrder    }     func snapshot(for provider: UsageProvider) -> UsageSnapshot? {        self.snapshots[provider]    }     func sourceLabel(for provider: UsageProvider) -> String {        var label = self.lastSourceLabels[provider] ?? ""        if label.isEmpty {            let descriptor = ProviderDescriptorRegistry.descriptor(for: provider)            let modes = descriptor.fetchPlan.sourceModes            if modes.count == 1, let mode = modes.first {                label = mode.rawValue            } else if provider == .codex {                label = self.settings.codexUsageDataSource.rawValue            } else if provider == .claude {                label = self.settings.claudeUsageDataSource.rawValue            } else {                label = "auto"            }        }         // When OpenAI web extras are active, show a blended label like `oauth + openai-web`.        if provider == .codex,           self.settings.codexCookieSource.isEnabled,           self.openAIDashboard != nil,           !self.openAIDashboardRequiresLogin,           !label.contains("openai-web")        {            return "\(label) + openai-web"        }        return label    }     func fetchAttempts(for provider: UsageProvider) -> [ProviderFetchAttempt] {        self.lastFetchAttempts[provider] ?? []    }     func style(for provider: UsageProvider) -> IconStyle {        self.providerSpecs[provider]?.style ?? .codex    }     func isStale(provider: UsageProvider) -> Bool {        self.errors[provider] != nil    }     func isEnabled(_ provider: UsageProvider) -> Bool {        let enabled = self.settings.isProviderEnabledCached(            provider: provider,            metadataByProvider: self.providerMetadata)        guard enabled else { return false }        return self.isProviderAvailable(provider)    }     private func isProviderAvailable(_ provider: UsageProvider) -> Bool {        if provider == .zai {            if ZaiSettingsReader.apiToken(environment: ProcessInfo.processInfo.environment) != nil {                return true            }            self.settings.ensureZaiAPITokenLoaded()            return !self.settings.zaiAPIToken.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty        }        return true    }     func refresh(forceTokenUsage: Bool = false) async {        guard !self.isRefreshing else { return }        self.isRefreshing = true        defer { self.isRefreshing = false }         await withTaskGroup(of: Void.self) { group in            for provider in UsageProvider.allCases {                group.addTask { await self.refreshProvider(provider) }                group.addTask { await self.refreshStatus(provider) }            }            group.addTask { await self.refreshCreditsIfNeeded() }        }         // Token-cost usage can be slow; run it outside the refresh group so we don't block menu updates.        self.scheduleTokenRefresh(force: forceTokenUsage)         // OpenAI web scrape depends on the current Codex account email (which can change after login/account switch).        // Run this after Codex usage refresh so we don't accidentally scrape with stale credentials.        await self.refreshOpenAIDashboardIfNeeded(force: forceTokenUsage)         if self.openAIDashboardRequiresLogin {            await self.refreshProvider(.codex)            await self.refreshCreditsIfNeeded()        }         self.persistWidgetSnapshot(reason: "refresh")    }     /// For demo/testing: drop the snapshot so the loading animation plays, then restore the last snapshot.    func replayLoadingAnimation(duration: TimeInterval = 3) {        let current = self.preferredSnapshot        self.snapshots.removeAll()        self.debugForceAnimation = true        Task { @MainActor in            try? await Task.sleep(for: .seconds(duration))            if let current, let provider = self.enabledProviders().first {

Lines 1–500 of 1655.