1276 lines · 57.6 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.versions        _ = self.isRefreshing        _ = self.refreshingProviders        _ = self.pathDebugInfo        _ = self.statuses        _ = self.probeLogs        _ = self.historicalPaceRevision        return 0    }     var iconObservationToken: Int {        _ = self.snapshots        _ = self.errors        _ = self.credits        _ = self.lastCreditsError        _ = self.openAIDashboard        _ = self.lastOpenAIDashboardError        _ = self.openAIDashboardRequiresLogin        _ = self.isRefreshing        _ = self.refreshingProviders        _ = self.statuses        _ = self.historicalPaceRevision        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            _ = self.settings.historicalTrackingEnabled        } onChange: { [weak self] in            Task { @MainActor [weak self] in                guard let self else { return }                self.observeSettingsChanges()                self.probeLogs = [:]                guard self.startupBehavior.automaticallyStartsBackgroundWork else { return }                self.startTimer()                self.updateProviderRuntimes()                await self.refreshHistoricalDatasetIfNeeded()                await self.refresh()            }        }    }     var attachedOpenAIDashboardSnapshot: OpenAIDashboardSnapshot? {        guard self.openAIDashboardAttachmentAuthorized else { return nil }        return self.openAIDashboard    }} @MainActor@Observablefinal class UsageStore {    enum CodexCreditsSource {        case none        case api        case dashboardWeb    }     enum StartupBehavior {        case automatic        case full        case testing         var automaticallyStartsBackgroundWork: Bool {            switch self {            case .automatic, .full:                true            case .testing:                false            }        }         func resolved(isRunningTests: Bool) -> StartupBehavior {            switch self {            case .automatic:                isRunningTests ? .testing : .full            case .full, .testing:                self            }        }    }     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] = [:]    var historicalPaceRevision: Int = 0    @ObservationIgnored var lastCreditsSnapshot: CreditsSnapshot?    @ObservationIgnored var lastCreditsSnapshotAccountKey: String?    @ObservationIgnored var lastCreditsSource: CodexCreditsSource = .none    @ObservationIgnored var creditsFailureStreak: Int = 0    @ObservationIgnored var openAIDashboardAttachmentAuthorized: Bool = false    @ObservationIgnored var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot?    @ObservationIgnored var lastOpenAIDashboardAttachmentAuthorized: Bool = false    @ObservationIgnored var lastOpenAIDashboardTargetEmail: String?    @ObservationIgnored var lastOpenAIDashboardAttemptAt: Date?    @ObservationIgnored var lastOpenAIDashboardCookieImportAttemptAt: Date?    @ObservationIgnored var lastOpenAIDashboardCookieImportEmail: String?    @ObservationIgnored var lastCodexAccountScopedRefreshGuard: CodexAccountScopedRefreshGuard?    @ObservationIgnored var lastKnownLiveSystemCodexEmail: String?    @ObservationIgnored var openAIWebAccountDidChange: Bool = false    @ObservationIgnored var openAIDashboardRefreshTask: Task<Void, Never>?    @ObservationIgnored var openAIDashboardRefreshTaskKey: String?    @ObservationIgnored var openAIDashboardRefreshTaskToken: UUID?    @ObservationIgnored var _test_openAIDashboardCookieImportOverride: (@MainActor (        String?,        Bool,        ProviderCookieSource,        CookieHeaderCache.Scope?,        @escaping (String) -> Void) async throws -> OpenAIDashboardBrowserCookieImporter.ImportResult)?    @ObservationIgnored var _test_openAIDashboardLoaderOverride: (@MainActor (        String?,        @escaping (String) -> Void,        TimeInterval) async throws -> OpenAIDashboardSnapshot)?    @ObservationIgnored var _test_codexCreditsLoaderOverride: (@MainActor () async throws -> CreditsSnapshot)?    @ObservationIgnored var _test_widgetSnapshotSaveOverride: (@MainActor (WidgetSnapshot) async -> Void)?    @ObservationIgnored var widgetSnapshotPersistTask: Task<Void, Never>?     @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: any SessionQuotaNotifying    @ObservationIgnored private let sessionQuotaLogger = CodexBarLog.logger(LogCategories.sessionQuota)    @ObservationIgnored 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 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 codexPlanHistoryBackfillTask: Task<Void, Never>?    @ObservationIgnored let historicalUsageHistoryStore: HistoricalUsageHistoryStore    @ObservationIgnored let planUtilizationHistoryStore: PlanUtilizationHistoryStore    @ObservationIgnored var codexHistoricalDataset: CodexHistoricalDataset?    @ObservationIgnored var codexHistoricalDatasetAccountKey: String?    @ObservationIgnored var lastKnownSessionRemaining: [UsageProvider: Double] = [:]    @ObservationIgnored var lastKnownSessionWindowSource: [UsageProvider: SessionQuotaWindowSource] = [:]    @ObservationIgnored var lastTokenFetchAt: [UsageProvider: Date] = [:]    @ObservationIgnored var planUtilizationHistory: [UsageProvider: PlanUtilizationHistoryBuckets] = [:]    @ObservationIgnored private var hasCompletedInitialRefresh: Bool = false    @ObservationIgnored private let tokenFetchTTL: TimeInterval = 60 * 60    @ObservationIgnored private let tokenFetchTimeout: TimeInterval = 10 * 60    @ObservationIgnored private let startupBehavior: StartupBehavior    @ObservationIgnored let planUtilizationPersistenceCoordinator: PlanUtilizationHistoryPersistenceCoordinator     init(        fetcher: UsageFetcher,        browserDetection: BrowserDetection,        claudeFetcher: (any ClaudeUsageFetching)? = nil,        costUsageFetcher: CostUsageFetcher = CostUsageFetcher(),        settings: SettingsStore,        registry: ProviderRegistry = .shared,        historicalUsageHistoryStore: HistoricalUsageHistoryStore = HistoricalUsageHistoryStore(),        planUtilizationHistoryStore: PlanUtilizationHistoryStore = .defaultAppSupport(),        sessionQuotaNotifier: any SessionQuotaNotifying = SessionQuotaNotifier(),        startupBehavior: StartupBehavior = .automatic)    {        self.codexFetcher = fetcher        self.browserDetection = browserDetection        self.claudeFetcher = claudeFetcher ?? ClaudeUsageFetcher(browserDetection: browserDetection)        self.costUsageFetcher = costUsageFetcher        self.settings = settings        self.registry = registry        self.historicalUsageHistoryStore = historicalUsageHistoryStore        self.planUtilizationHistoryStore = planUtilizationHistoryStore        self.sessionQuotaNotifier = sessionQuotaNotifier        self.startupBehavior = startupBehavior.resolved(isRunningTests: Self.isRunningTestsProcess())        self.planUtilizationPersistenceCoordinator = PlanUtilizationHistoryPersistenceCoordinator(            store: planUtilizationHistoryStore)        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.planUtilizationHistory = planUtilizationHistoryStore.load()        self.logStartupState()        self.bindSettings()        self.pathDebugInfo = PathDebugSnapshot(            codexBinary: nil,            claudeBinary: nil,            geminiBinary: nil,            effectivePATH: PathBuilder.effectivePATH(purposes: [.rpc, .tty, .nodeTooling]),            loginShellPATH: LoginShellPathCache.shared.current?.joined(separator: ":"))        guard self.startupBehavior.automaticallyStartsBackgroundWork else { return }        self.detectVersions()        self.updateProviderRuntimes()        Task { @MainActor [weak self] in            self?.schedulePathDebugInfoRefresh()        }        LoginShellPathCache.shared.captureOnce { [weak self] _ in            Task { @MainActor [weak self] in                self?.schedulePathDebugInfoRefresh()            }        }        Task { @MainActor [weak self] in            await self?.refreshHistoricalDatasetIfNeeded()        }        Task { await self.refresh() }        self.startTimer()        self.startTokenTimer()    }     private static func isRunningTestsProcess() -> Bool {        let environment = ProcessInfo.processInfo.environment        if environment["XCTestConfigurationFilePath"] != nil { return true }        if environment["XCTestSessionIdentifier"] != nil { return true }        if environment["SWIFT_TESTING_ENABLED"] != nil { return true }        return CommandLine.arguments.contains { argument in            argument.contains("xctest") || argument.contains("swift-testing")        }    }     /// 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 {        ClaudePlan.isSubscriptionLoginMethod(loginMethod)    }     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) }    }     /// Enabled providers without availability filtering. Used for display (switcher, merge-icons).    func enabledProvidersForDisplay() -> [UsageProvider] {        self.settings.enabledProvidersOrdered(metadataByProvider: self.providerMetadata)    }     /// Providers that should actually participate in background refresh/status/token work.    func enabledProvidersForBackgroundWork() -> [UsageProvider] {        self.enabledProviders()    }     var statusChecksEnabled: Bool {        self.settings.statusChecksEnabled    }     func metadata(for provider: UsageProvider) -> ProviderMetadata {        self.providerMetadata[provider]!    }     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 }        self.prepareRefreshState()        let refreshPhase: ProviderRefreshPhase = self.hasCompletedInitialRefresh ? .regular : .startup        let displayEnabledProviders = self.enabledProvidersForDisplay()        let enabledProviderSet = Set(displayEnabledProviders)        let refreshProviders = self.enabledProvidersForBackgroundWork()        let availableRefreshProviders = Set(self.enabledProviders())        let refreshStartedAt = Date()         await ProviderRefreshContext.$current.withValue(refreshPhase) {            self.isRefreshing = true            defer {                self.isRefreshing = false                self.hasCompletedInitialRefresh = true            }             self.clearDisabledProviderState(enabledProviders: enabledProviderSet)            self.clearUnavailableProviderState(                displayEnabledProviders: enabledProviderSet,                availableProviders: availableRefreshProviders)             await withTaskGroup(of: Void.self) { group in                for provider in refreshProviders {                    group.addTask { await self.refreshProvider(provider) }                    if availableRefreshProviders.contains(provider) {                        group.addTask { await self.refreshStatus(provider) }                    }                }                group.addTask { await self.refreshCreditsIfNeeded(minimumSnapshotUpdatedAt: refreshStartedAt) }            }             // 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.            self.syncOpenAIWebState()            let refreshPolicy = OpenAIWebRefreshPolicyContext(                accessEnabled: self.isEnabled(.codex) &&                    self.settings.openAIWebAccessEnabled &&                    self.settings.codexCookieSource.isEnabled,                batterySaverEnabled: self.settings.openAIWebBatterySaverEnabled,                force: forceTokenUsage)            let shouldRefreshOpenAIWeb = Self.shouldRunOpenAIWebRefresh(refreshPolicy)            self.openAIWebLogger.debug(                "OpenAI web refresh gate",                metadata: [

Lines 1–500 of 1276.