915 lines · 37.6 KB · Swift
RawDownload
import AppKitimport CodexBarCoreimport Combineimport Foundationimport OSLog enum IconStyle {    case codex    case claude    case gemini    case combined} 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"        }    }} 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    }} @MainActorfinal class UsageStore: ObservableObject {    @Published private var snapshots: [UsageProvider: UsageSnapshot] = [:]    @Published private var errors: [UsageProvider: String] = [:]    @Published var credits: CreditsSnapshot?    @Published var lastCreditsError: String?    @Published var openAIDashboard: OpenAIDashboardSnapshot?    @Published var lastOpenAIDashboardError: String?    @Published private(set) var openAIDashboardRequiresLogin: Bool = false    @Published var openAIDashboardCookieImportStatus: String?    @Published var openAIDashboardCookieImportDebugLog: String?    @Published var codexVersion: String?    @Published var claudeVersion: String?    @Published var geminiVersion: String?    @Published var claudeAccountEmail: String?    @Published var claudeAccountOrganization: String?    @Published var isRefreshing = false    @Published var debugForceAnimation = false    @Published var pathDebugInfo: PathDebugSnapshot = .empty    @Published private var statuses: [UsageProvider: ProviderStatus] = [:]    @Published private(set) var probeLogs: [UsageProvider: String] = [:]    private var lastCreditsSnapshot: CreditsSnapshot?    private var creditsFailureStreak: Int = 0    private var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot?    private var lastOpenAIDashboardTargetEmail: String?    private var lastOpenAIDashboardCookieImportAttemptAt: Date?    private var lastOpenAIDashboardCookieImportEmail: String?    private var openAIWebAccountDidChange: Bool = false     private let codexFetcher: UsageFetcher    private let claudeFetcher: any ClaudeUsageFetching    private let registry: ProviderRegistry    private let settings: SettingsStore    private let sessionQuotaNotifier: SessionQuotaNotifier    private let sessionQuotaLogger = Logger(subsystem: "com.steipete.codexbar", category: "sessionQuota")    private let openAIWebLogger = Logger(subsystem: "com.steipete.codexbar", category: "openai-web")    private var openAIWebDebugLines: [String] = []    private var failureGates: [UsageProvider: ConsecutiveFailureGate] = [:]    private var providerSpecs: [UsageProvider: ProviderSpec] = [:]    private let providerMetadata: [UsageProvider: ProviderMetadata]    private var timerTask: Task<Void, Never>?    private var cancellables = Set<AnyCancellable>()    private var lastKnownSessionRemaining: [UsageProvider: Double] = [:]     init(        fetcher: UsageFetcher,        claudeFetcher: any ClaudeUsageFetching = ClaudeUsageFetcher(),        settings: SettingsStore,        registry: ProviderRegistry = .shared,        sessionQuotaNotifier: SessionQuotaNotifier = SessionQuotaNotifier())    {        self.codexFetcher = fetcher        self.claudeFetcher = claudeFetcher        self.settings = settings        self.registry = registry        self.sessionQuotaNotifier = sessionQuotaNotifier        self.providerMetadata = registry.metadata        self            .failureGates = 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()    }     var codexSnapshot: UsageSnapshot? { self.snapshots[.codex] }    var claudeSnapshot: UsageSnapshot? { self.snapshots[.claude] }    var lastCodexError: String? { self.errors[.codex] }    var lastClaudeError: String? { self.errors[.claude] }    func error(for provider: UsageProvider) -> String? { self.errors[provider] }    func metadata(for provider: UsageProvider) -> ProviderMetadata { self.providerMetadata[provider]! }    func status(for provider: UsageProvider) -> ProviderStatus? {        guard self.settings.statusChecksEnabled else { return nil }        return self.statuses[provider]    }     func statusIndicator(for provider: UsageProvider) -> ProviderStatusIndicator {        self.status(for: provider)?.indicator ?? .none    }     /// 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 .gemini: self.geminiVersion        }    }     var preferredSnapshot: UsageSnapshot? {        if self.isEnabled(.codex), let codexSnapshot {            return codexSnapshot        }        if self.isEnabled(.claude), let claudeSnapshot {            return claudeSnapshot        }        if self.isEnabled(.gemini), let snap = self.snapshots[.gemini] {            return snap        }        return nil    }     var iconStyle: IconStyle {        if self.isEnabled(.claude) { return .claude }        if self.isEnabled(.gemini) { return .gemini }        return .codex    }     var isStale: Bool {        (self.isEnabled(.codex) && self.lastCodexError != nil) ||            (self.isEnabled(.claude) && self.lastClaudeError != nil) ||            (self.isEnabled(.gemini) && self.errors[.gemini] != nil)    }     func enabledProviders() -> [UsageProvider] {        UsageProvider.allCases.filter { self.isEnabled($0) }    }     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() 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() }        }         // 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()    }     /// 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.settings.$refreshFrequency            .sink { [weak self] _ in                self?.startTimer()            }            .store(in: &self.cancellables)         self.settings.objectWillChange            .sink { [weak self] _ in                Task { await self?.refresh() }            }            .store(in: &self.cancellables)    }     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()            }        }    }     deinit {        self.timerTask?.cancel()    }     private func refreshProvider(_ provider: UsageProvider) async {        guard let spec = self.providerSpecs[provider] else { return }         if !spec.isEnabled() {            await MainActor.run {                self.snapshots.removeValue(forKey: provider)                self.errors[provider] = nil                self.failureGates[provider]?.reset()                self.statuses.removeValue(forKey: provider)                self.lastKnownSessionRemaining.removeValue(forKey: provider)            }            return        }         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                }            }        }    }     private func handleSessionQuotaTransition(provider: UsageProvider, snapshot: UsageSnapshot) {        let currentRemaining = snapshot.primary.remainingPercent        let previousRemaining = self.lastKnownSessionRemaining[provider]         defer { self.lastKnownSessionRemaining[provider] = currentRemaining }         guard self.settings.sessionQuotaNotificationsEnabled else {            if SessionQuotaNotificationLogic.isDepleted(currentRemaining) ||                SessionQuotaNotificationLogic.isDepleted(previousRemaining)            {                let providerText = provider.rawValue                let message =                    "notifications disabled: provider=\(providerText) " +                    "prev=\(previousRemaining ?? -1) curr=\(currentRemaining)"                self.sessionQuotaLogger.debug("\(message, privacy: .public)")            }            return        }         guard previousRemaining != nil else {            if SessionQuotaNotificationLogic.isDepleted(currentRemaining) {                let providerText = provider.rawValue                let message = "startup depleted: provider=\(providerText) curr=\(currentRemaining)"                self.sessionQuotaLogger.info("\(message, privacy: .public)")                self.sessionQuotaNotifier.post(transition: .depleted, provider: provider)            }            return        }         let transition = SessionQuotaNotificationLogic.transition(            previousRemaining: previousRemaining,            currentRemaining: currentRemaining)        guard transition != .none else {            if SessionQuotaNotificationLogic.isDepleted(currentRemaining) ||                SessionQuotaNotificationLogic.isDepleted(previousRemaining)            {                let providerText = provider.rawValue                let message =                    "no transition: provider=\(providerText) " +                    "prev=\(previousRemaining ?? -1) curr=\(currentRemaining)"                self.sessionQuotaLogger.debug("\(message, privacy: .public)")            }            return        }         let providerText = provider.rawValue        let transitionText = String(describing: transition)        let message =            "transition \(transitionText): provider=\(providerText) " +            "prev=\(previousRemaining ?? -1) curr=\(currentRemaining)"        self.sessionQuotaLogger.info("\(message, privacy: .public)")         self.sessionQuotaNotifier.post(transition: transition, provider: provider)    }     private func refreshStatus(_ provider: UsageProvider) async {        guard self.settings.statusChecksEnabled else { return }        guard let urlString = self.providerMetadata[provider]?.statusPageURL,              let baseURL = URL(string: urlString) else { return }         do {            let status = try await Self.fetchStatus(from: baseURL)            await MainActor.run { self.statuses[provider] = status }        } catch {            // Keep the previous status to avoid flapping when the API hiccups.            await MainActor.run {                if self.statuses[provider] == nil {                    self.statuses[provider] = ProviderStatus(                        indicator: .unknown,                        description: error.localizedDescription,                        updatedAt: nil)                }            }        }    }     private func refreshCreditsIfNeeded() async {        guard self.isEnabled(.codex) else { return }        do {            let credits = try await self.codexFetcher.loadLatestCredits()            await MainActor.run {                self.credits = credits                self.lastCreditsError = nil                self.lastCreditsSnapshot = credits                self.creditsFailureStreak = 0            }        } catch {            let message = error.localizedDescription            if message.localizedCaseInsensitiveContains("data not available yet") {                await MainActor.run {                    if let cached = self.lastCreditsSnapshot {                        self.credits = cached                        self.lastCreditsError = nil                    } else {                        self.credits = nil                        self.lastCreditsError = "Codex credits are still loading; will retry shortly."                    }                }                return            }             await MainActor.run {                self.creditsFailureStreak += 1                if let cached = self.lastCreditsSnapshot {                    self.credits = cached                    let stamp = cached.updatedAt.formatted(date: .abbreviated, time: .shortened)                    self.lastCreditsError =                        "Last Codex credits refresh failed: \(message). Cached values from \(stamp)."                } else {                    self.lastCreditsError = message                    self.credits = nil                }            }        }    }     private func applyOpenAIDashboard(_ dash: OpenAIDashboardSnapshot, targetEmail: String?) async {        await MainActor.run {            self.openAIDashboard = dash            self.lastOpenAIDashboardError = nil            self.lastOpenAIDashboardSnapshot = dash            self.openAIDashboardRequiresLogin = false        }         if let email = targetEmail, !email.isEmpty {            OpenAIDashboardCacheStore.save(OpenAIDashboardCache(accountEmail: email, snapshot: dash))        }    }     private func applyOpenAIDashboardFailure(message: String) async {        await MainActor.run {            if let cached = self.lastOpenAIDashboardSnapshot {                self.openAIDashboard = cached                let stamp = cached.updatedAt.formatted(date: .abbreviated, time: .shortened)                self.lastOpenAIDashboardError =                    "Last OpenAI dashboard refresh failed: \(message). Cached values from \(stamp)."            } else {                self.lastOpenAIDashboardError = message                self.openAIDashboard = nil            }        }    }     private func refreshOpenAIDashboardIfNeeded() async {        guard self.isEnabled(.codex) else { return }        guard self.settings.openAIDashboardEnabled else {            await MainActor.run {                self.openAIDashboard = nil                self.lastOpenAIDashboardError = nil                self.lastOpenAIDashboardSnapshot = nil                self.lastOpenAIDashboardTargetEmail = nil                self.openAIDashboardRequiresLogin = false                self.openAIDashboardCookieImportStatus = nil                self.openAIDashboardCookieImportDebugLog = nil                self.lastOpenAIDashboardCookieImportAttemptAt = nil                self.lastOpenAIDashboardCookieImportEmail = nil

Lines 1–500 of 915.