1330 lines · 55.6 KB · Swift
RawDownload
import AppKitimport CodexBarCoreimport Foundationimport Observation enum IconStyle {    case codex    case claude    case zai    case gemini    case antigravity    case cursor    case factory    case combined} // MARK: - Observation helpers @MainActorextension UsageStore {    var menuObservationToken: Int {        _ = self.snapshots        _ = self.errors        _ = 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.claudeAccountEmail        _ = self.claudeAccountOrganization        _ = 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.ccusageCostUsageEnabled            _ = self.settings.randomBlinkEnabled            _ = self.settings.claudeWebExtrasEnabled            _ = self.settings.claudeUsageDataSource            _ = self.settings.mergeIcons            _ = self.settings.selectedMenuProvider            _ = self.settings.debugLoadingPattern        } onChange: { [weak self] in            Task { @MainActor [weak self] in                guard let self else { return }                self.observeSettingsChanges()                self.startTimer()                await self.refresh()            }        }    }} 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    }     func _setTokenSnapshotForTesting(_ snapshot: CCUsageTokenSnapshot?, 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 tokenSnapshots: [UsageProvider: CCUsageTokenSnapshot] = [:]    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 claudeAccountEmail: String?    var claudeAccountOrganization: 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 ccusageFetcher: CCUsageFetcher    @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     init(        fetcher: UsageFetcher,        claudeFetcher: any ClaudeUsageFetching = ClaudeUsageFetcher(),        ccusageFetcher: CCUsageFetcher = CCUsageFetcher(),        settings: SettingsStore,        registry: ProviderRegistry = .shared,        sessionQuotaNotifier: SessionQuotaNotifier = SessionQuotaNotifier())    {        self.codexFetcher = fetcher        self.claudeFetcher = claudeFetcher        self.ccusageFetcher = ccusageFetcher        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: claudeFetcher)        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()    }     /// Returns the login method (plan type) for the specified provider, if available.    private func loginMethod(for provider: UsageProvider) -> String? {        self.snapshots[provider]?.loginMethod    }     /// 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        }    }     var preferredSnapshot: UsageSnapshot? {        if self.isEnabled(.codex), let codexSnapshot {            return codexSnapshot        }        if self.isEnabled(.claude), let claudeSnapshot {            return claudeSnapshot        }        if self.isEnabled(.zai), let snap = self.snapshots[.zai] {            return snap        }        if self.isEnabled(.gemini), let snap = self.snapshots[.gemini] {            return snap        }        if self.isEnabled(.antigravity), let snap = self.snapshots[.antigravity] {            return snap        }        if self.isEnabled(.cursor), let snap = self.snapshots[.cursor] {            return snap        }        return nil    }     var iconStyle: IconStyle {        let enabled = self.enabledProviders()        if enabled.count > 1 { return .combined }        if self.isEnabled(.cursor) { return .cursor }        if self.isEnabled(.antigravity) { return .antigravity }        if self.isEnabled(.gemini) { return .gemini }        if self.isEnabled(.zai) { return .zai }        if self.isEnabled(.claude) { return .claude }        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)    }     func enabledProviders() -> [UsageProvider] {        self.settings.orderedProviders().filter { self.isEnabled($0) }    }     var statusChecksEnabled: Bool {        self.settings.statusChecksEnabled    }     func metadata(for provider: UsageProvider) -> ProviderMetadata {        self.providerMetadata[provider]!    }     func snapshot(for provider: UsageProvider) -> UsageSnapshot? {        self.snapshots[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 {        self.settings.isProviderEnabled(provider: provider, metadata: self.metadata(for: provider))    }     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 {                if self.isEnabled(.codex) {                    self.snapshots[.codex] = current                } else if self.isEnabled(.claude) {                    self.snapshots[.claude] = current                }            }            self.debugForceAnimation = false        }    }     // MARK: - Private     private func bindSettings() {        self.observeSettingsChanges()    }     private func startTimer() {        self.timerTask?.cancel()        guard let wait = self.settings.refreshFrequency.seconds else { return }         // Background poller so the menu stays responsive; canceled when settings change or store deallocates.        self.timerTask = Task.detached(priority: .utility) { [weak self] in            while !Task.isCancelled {                try? await Task.sleep(for: .seconds(wait))                await self?.refresh()            }        }    }     private func startTokenTimer() {        self.tokenTimerTask?.cancel()        let wait = self.tokenFetchTTL        self.tokenTimerTask = Task.detached(priority: .utility) { [weak self] in            while !Task.isCancelled {                try? await Task.sleep(for: .seconds(wait))                await self?.scheduleTokenRefresh(force: false)            }        }    }     private func scheduleTokenRefresh(force: Bool) {        if force {            self.tokenRefreshSequenceTask?.cancel()            self.tokenRefreshSequenceTask = nil        } else if self.tokenRefreshSequenceTask != nil {            return        }         self.tokenRefreshSequenceTask = Task(priority: .utility) { [weak self] in            guard let self else { return }            defer {                Task { @MainActor [weak self] in                    self?.tokenRefreshSequenceTask = nil                }            }            for provider in UsageProvider.allCases {                if Task.isCancelled { break }                await self.refreshTokenUsage(provider, force: force)            }        }    }     deinit {        self.timerTask?.cancel()        self.tokenTimerTask?.cancel()        self.tokenRefreshSequenceTask?.cancel()    }     private func refreshProvider(_ provider: UsageProvider) async {        guard let spec = self.providerSpecs[provider] else { return }         if !spec.isEnabled() {            self.refreshingProviders.remove(provider)            await MainActor.run {                self.snapshots.removeValue(forKey: provider)                self.errors[provider] = nil                self.tokenSnapshots.removeValue(forKey: provider)                self.tokenErrors[provider] = nil                self.failureGates[provider]?.reset()                self.tokenFailureGates[provider]?.reset()                self.statuses.removeValue(forKey: provider)                self.lastKnownSessionRemaining.removeValue(forKey: provider)                self.lastTokenFetchAt.removeValue(forKey: provider)            }            return        }         self.refreshingProviders.insert(provider)        defer { self.refreshingProviders.remove(provider) }         do {            let fetchClosure = spec.fetch            let task = Task(priority: .utility) { () -> UsageSnapshot in                try await fetchClosure()            }            let snapshot = try await task.value            await MainActor.run {                self.handleSessionQuotaTransition(provider: provider, snapshot: snapshot)                self.snapshots[provider] = snapshot                self.errors[provider] = nil                self.failureGates[provider]?.recordSuccess()                if provider == .claude {                    self.claudeAccountEmail = snapshot.accountEmail                    self.claudeAccountOrganization = snapshot.accountOrganization                }            }        } catch {            await MainActor.run {                let hadPriorData = self.snapshots[provider] != nil                let shouldSurface = self.failureGates[provider]?                    .shouldSurfaceError(onFailureWithPriorData: hadPriorData) ?? true                if shouldSurface {                    self.errors[provider] = error.localizedDescription                    self.snapshots.removeValue(forKey: provider)                } else {                    self.errors[provider] = nil                }            }        }    } 

Lines 1–500 of 1330.