1653 lines · 71.7 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.accountSnapshots        _ = self.tokenSnapshots        _ = self.tokenErrors        _ = self.tokenRefreshInFlight        _ = self.credits        _ = self.lastCreditsError        _ = self.openAIDashboard        _ = self.lastOpenAIDashboardError        _ = self.openAIDashboardRequiresLogin        _ = self.openAIDashboardCookieImportStatus        _ = self.openAIDashboardCookieImportDebugLog        _ = self.versions        _ = 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.configRevision            for implementation in ProviderCatalog.all {                implementation.observeSettings(self.settings)            }            _ = self.settings.showAllTokenAccountsInMenu            _ = self.settings.tokenAccountsByProvider            _ = self.settings.mergeIcons            _ = self.settings.selectedMenuProvider            _ = self.settings.debugLoadingPattern            _ = self.settings.debugKeepCLISessionsAlive        } onChange: { [weak self] in            Task { @MainActor [weak self] in                guard let self else { return }                self.observeSettingsChanges()                self.startTimer()                self.updateProviderRuntimes()                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?.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 {    var snapshots: [UsageProvider: UsageSnapshot] = [:]    var errors: [UsageProvider: String] = [:]    var lastSourceLabels: [UsageProvider: String] = [:]    var lastFetchAttempts: [UsageProvider: [ProviderFetchAttempt]] = [:]    var accountSnapshots: [UsageProvider: [TokenAccountUsageSnapshot]] = [:]    var tokenSnapshots: [UsageProvider: CostUsageTokenSnapshot] = [:]    var tokenErrors: [UsageProvider: String] = [:]    var tokenRefreshInFlight: Set<UsageProvider> = []    var credits: CreditsSnapshot?    var lastCreditsError: String?    var openAIDashboard: OpenAIDashboardSnapshot?    var lastOpenAIDashboardError: String?    var openAIDashboardRequiresLogin: Bool = false    var openAIDashboardCookieImportStatus: String?    var openAIDashboardCookieImportDebugLog: String?    var versions: [UsageProvider: String] = [:]    var isRefreshing = false    var refreshingProviders: Set<UsageProvider> = []    var debugForceAnimation = false    var pathDebugInfo: PathDebugSnapshot = .empty    var statuses: [UsageProvider: ProviderStatus] = [:]    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 let codexFetcher: UsageFetcher    @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching    @ObservationIgnored private let costUsageFetcher: CostUsageFetcher    @ObservationIgnored let browserDetection: BrowserDetection    @ObservationIgnored private let registry: ProviderRegistry    @ObservationIgnored let settings: SettingsStore    @ObservationIgnored private let sessionQuotaNotifier: SessionQuotaNotifier    @ObservationIgnored private let sessionQuotaLogger = CodexBarLog.logger(LogCategories.sessionQuota)    @ObservationIgnored private let openAIWebLogger = CodexBarLog.logger(LogCategories.openAIWeb)    @ObservationIgnored private let tokenCostLogger = CodexBarLog.logger(LogCategories.tokenCost)    @ObservationIgnored let augmentLogger = CodexBarLog.logger(LogCategories.augment)    @ObservationIgnored let providerLogger = CodexBarLog.logger(LogCategories.providers)    @ObservationIgnored private var openAIWebDebugLines: [String] = []    @ObservationIgnored var failureGates: [UsageProvider: ConsecutiveFailureGate] = [:]    @ObservationIgnored var tokenFailureGates: [UsageProvider: ConsecutiveFailureGate] = [:]    @ObservationIgnored var providerSpecs: [UsageProvider: ProviderSpec] = [:]    @ObservationIgnored let providerMetadata: [UsageProvider: ProviderMetadata]    @ObservationIgnored var providerRuntimes: [UsageProvider: any ProviderRuntime] = [:]    @ObservationIgnored private var timerTask: Task<Void, Never>?    @ObservationIgnored private var tokenTimerTask: Task<Void, Never>?    @ObservationIgnored private var tokenRefreshSequenceTask: Task<Void, Never>?    @ObservationIgnored private var pathDebugRefreshTask: Task<Void, Never>?    @ObservationIgnored var lastKnownSessionRemaining: [UsageProvider: Double] = [:]    @ObservationIgnored var lastTokenFetchAt: [UsageProvider: Date] = [:]    @ObservationIgnored private var hasCompletedInitialRefresh: Bool = false    @ObservationIgnored private let tokenFetchTTL: TimeInterval = 60 * 60    @ObservationIgnored private let tokenFetchTimeout: TimeInterval = 10 * 60     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.providerRuntimes = Dictionary(uniqueKeysWithValues: ProviderCatalog.all.compactMap { implementation in            implementation.makeRuntime().map { (implementation.id, $0) }        })        self.logStartupState()        self.bindSettings()        self.detectVersions()        self.updateProviderRuntimes()        self.pathDebugInfo = PathDebugSnapshot(            codexBinary: nil,            claudeBinary: nil,            geminiBinary: nil,            effectivePATH: PathBuilder.effectivePATH(purposes: [.rpc, .tty, .nodeTooling]),            loginShellPATH: LoginShellPathCache.shared.current?.joined(separator: ":"))        Task { @MainActor [weak self] in            self?.schedulePathDebugInfoRefresh()        }        LoginShellPathCache.shared.captureOnce { [weak self] _ in            Task { @MainActor [weak self] in                self?.schedulePathDebugInfoRefresh()            }        }        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(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? {        self.versions[provider]    }     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 {        for provider in self.enabledProviders() where self.errors[provider] != nil {            return true        }        return false    }     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) }    }     /// Returns the enabled provider with the highest usage percentage (closest to rate limit).    /// Excludes providers already at 100% since they're fully rate-limited.    func providerWithHighestUsage() -> (provider: UsageProvider, usedPercent: Double)? {        var highest: (provider: UsageProvider, usedPercent: Double)?        for provider in self.enabledProviders() {            guard let snapshot = self.snapshots[provider] else { continue }            // Use the same window selection logic as menuBarPercentWindow:            // Factory uses secondary (premium) first, others use primary (session) first.            let window: RateWindow? = if provider == .factory {                snapshot.secondary ?? snapshot.primary            } else {                snapshot.primary ?? snapshot.secondary            }            let percent = window?.usedPercent ?? 0            // Skip providers already at 100% - they're fully rate-limited            guard percent < 100 else { continue }            if highest == nil || percent > highest!.usedPercent {                highest = (provider, percent)            }        }        return highest    }     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 {                let context = ProviderSourceLabelContext(                    provider: provider,                    settings: self.settings,                    store: self,                    descriptor: descriptor)                label = ProviderCatalog.implementation(for: provider)?                    .defaultSourceLabel(context: context)                    ?? "auto"            }        }         let context = ProviderSourceLabelContext(            provider: provider,            settings: self.settings,            store: self,            descriptor: ProviderDescriptorRegistry.descriptor(for: provider))        return ProviderCatalog.implementation(for: provider)?            .decorateSourceLabel(context: context, baseLabel: label)            ?? 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)    }     func isProviderAvailable(_ provider: UsageProvider) -> Bool {        // Availability should mirror the effective fetch environment, including token-account overrides.        // Otherwise providers (notably token-account-backed API providers) can fetch successfully but be        // hidden from the menu because their credentials are not in ProcessInfo's environment.        let environment = ProviderRegistry.makeEnvironment(            base: ProcessInfo.processInfo.environment,            provider: provider,            settings: self.settings,            tokenOverride: nil)        let context = ProviderAvailabilityContext(            provider: provider,            settings: self.settings,            environment: environment)        return ProviderCatalog.implementation(for: provider)?            .isAvailable(context: context)            ?? true    }     func performRuntimeAction(_ action: ProviderRuntimeAction, for provider: UsageProvider) async {        guard let runtime = self.providerRuntimes[provider] else { return }        let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self)        await runtime.perform(action: action, context: context)    }     private func updateProviderRuntimes() {        for (provider, runtime) in self.providerRuntimes {            let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self)            if self.isEnabled(provider) {                runtime.start(context: context)            } else {                runtime.stop(context: context)            }            runtime.settingsDidChange(context: context)        }    }     func refresh(forceTokenUsage: Bool = false) async {        guard !self.isRefreshing else { return }        let refreshPhase: ProviderRefreshPhase = self.hasCompletedInitialRefresh ? .regular : .startup         await ProviderRefreshContext.$current.withValue(refreshPhase) {            self.isRefreshing = true            defer {                self.isRefreshing = false                self.hasCompletedInitialRefresh = true            }             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 {                self.snapshots[provider] = 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()            }        }    }

Lines 1–500 of 1653.