1582 lines · 73.5 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.codexAccountSnapshots        _ = self.kiloScopeSnapshots        _ = self.tokenSnapshots        _ = self.tokenErrors        _ = self.tokenRefreshInFlight        _ = self.credits        _ = self.lastCreditsError        _ = self.openAIDashboard        _ = self.lastOpenAIDashboardError        _ = self.openAIDashboardRequiresLogin        _ = self.openAIDashboardAttachmentRevision        _ = self.versions        _ = self.isRefreshing        _ = self.refreshingProviders        _ = self.pathDebugInfo        _ = self.statuses        _ = self.probeLogs        _ = self.historicalPaceRevision        _ = self.planUtilizationHistoryRevision        _ = self.providerStorageFootprints        return 0    }     var iconObservationToken: Int {        _ = self.snapshots        _ = self.errors        _ = self.credits        _ = self.lastCreditsError        _ = self.openAIDashboard        _ = self.lastOpenAIDashboardError        _ = self.openAIDashboardRequiresLogin        _ = self.refreshingProviders        _ = self.statuses        _ = self.historicalPaceRevision        return 0    }     func observeSettingsChanges() {        withObservationTracking {            _ = self.backgroundWorkSettingsObservationToken        } onChange: { [weak self] in            Task { @MainActor [weak self] in                guard let self else { return }                self.observeSettingsChanges()                self.invalidateProviderAvailabilityCache()                self.probeLogs = [:]                guard self.startupBehavior.automaticallyStartsBackgroundWork else { return }                self.startTimer()                self.updateProviderRuntimes()                await self.refreshHistoricalDatasetIfNeeded()                await self.refreshForSettingsChange()            }        }    }     var backgroundWorkSettingsObservationToken: Int {        _ = self.settings.refreshFrequency        _ = self.settings.statusChecksEnabled        _ = self.settings.sessionQuotaNotificationsEnabled        _ = self.settings.quotaWarningNotificationsEnabled        _ = self.settings.quotaWarningThresholds        _ = self.settings.quotaWarningThresholds(.session)        _ = self.settings.quotaWarningThresholds(.weekly)        _ = self.settings.quotaWarningSoundEnabled        _ = self.settings.usageBarsShowUsed        _ = self.settings.costUsageEnabled        _ = self.settings.costUsageHistoryDays        _ = self.settings.randomBlinkEnabled        _ = self.settings.configRevision        for implementation in ProviderCatalog.all {            implementation.observeSettings(self.settings)        }        _ = self.settings.multiAccountMenuLayout        _ = self.settings.tokenAccountsByProvider        _ = self.settings.mergeIcons        _ = self.settings.debugLoadingPattern        _ = self.settings.debugKeepCLISessionsAlive        _ = self.settings.historicalTrackingEnabled        _ = self.settings.providerStorageFootprintsEnabled        return 0    }     var attachedOpenAIDashboardSnapshot: OpenAIDashboardSnapshot? {        guard self.openAIDashboardAttachmentAuthorized else { return nil }        return self.openAIDashboard    }} @MainActor@Observablefinal class UsageStore {    private struct ProviderAvailabilityCacheEntry {        let available: Bool        let configRevision: Int        let expiresAt: Date         func isValid(now: Date, configRevision: Int) -> Bool {            self.configRevision == configRevision && self.expiresAt > now        }    }     struct AccountInfoCacheEntry {        let account: AccountInfo        let configRevision: Int        let expiresAt: Date         func isValid(now: Date, configRevision: Int) -> Bool {            self.configRevision == configRevision && self.expiresAt > now        }    }     enum CodexCreditsSource {        case none        case api        case dashboardWeb    }     var snapshots: [UsageProvider: UsageSnapshot] = [:]    var errors: [UsageProvider: String] = [:]    var lastSourceLabels: [UsageProvider: String] = [:]    var lastFetchAttempts: [UsageProvider: [ProviderFetchAttempt]] = [:]    var accountSnapshots: [UsageProvider: [TokenAccountUsageSnapshot]] = [:]    var codexAccountSnapshots: [CodexAccountUsageSnapshot] = []    var kiloScopeSnapshots: [KiloScopeSnapshot] = []    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    var planUtilizationHistoryRevision: Int = 0    var providerStorageFootprints: [UsageProvider: ProviderStorageFootprint] = [:]    @ObservationIgnored var lastCreditsSnapshot: CreditsSnapshot?    @ObservationIgnored var lastCreditsSnapshotAccountKey: String?    @ObservationIgnored var lastCreditsSource: CodexCreditsSource = .none    @ObservationIgnored var creditsFailureStreak: Int = 0    @ObservationIgnored var openAIDashboardAttachmentAuthorized: Bool = false {        didSet {            guard self.openAIDashboardAttachmentAuthorized != oldValue else { return }            self.openAIDashboardAttachmentRevision &+= 1        }    }     var openAIDashboardAttachmentRevision = 0    @ObservationIgnored var lastOpenAIDashboardSnapshot: OpenAIDashboardSnapshot?    @ObservationIgnored var lastOpenAIDashboardAttachmentAuthorized: Bool = false    @ObservationIgnored var lastOpenAIDashboardTargetEmail: String?    @ObservationIgnored var lastOpenAIDashboardTargetIsolationKey: 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 creditsRefreshTask: Task<Void, Never>?    @ObservationIgnored var creditsRefreshTaskKey: String?    @ObservationIgnored var openAIDashboardBackgroundRefreshTask: Task<Void, Never>?    @ObservationIgnored var openAIDashboardBackgroundRefreshTaskKey: String?    @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,        Bool,        TimeInterval) async throws -> OpenAIDashboardSnapshot)?    @ObservationIgnored var _test_codexCreditsLoaderOverride: (@MainActor () async throws -> CreditsSnapshot)?    @ObservationIgnored var _test_widgetSnapshotSaveOverride: (@MainActor (WidgetSnapshot) async -> Void)?    @ObservationIgnored var _test_providerRefreshOverride: (@MainActor (UsageProvider) async -> Void)?    @ObservationIgnored var _test_tokenUsageRefreshOverride: (@MainActor (UsageProvider, Bool) async -> Void)?    @ObservationIgnored var _test_providerStatusFetchOverride: (@MainActor (        UsageProvider) async throws -> ProviderStatus)?    @ObservationIgnored var _test_startupConnectivityRetryScheduled: (@MainActor (Int, TimeInterval) -> Void)?    @ObservationIgnored var _test_startupConnectivityRetrySleepOverride: (@MainActor (        TimeInterval) async throws -> Void)?    @ObservationIgnored var widgetSnapshotPersistTask: Task<Void, Never>?     @ObservationIgnored let codexFetcher: UsageFetcher    @ObservationIgnored let claudeFetcher: any ClaudeUsageFetching    @ObservationIgnored let costUsageFetcher: CostUsageFetcher    @ObservationIgnored let browserDetection: BrowserDetection    @ObservationIgnored private let registry: ProviderRegistry    @ObservationIgnored let settings: SettingsStore    @ObservationIgnored let environmentBase: [String: String]    @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 var providerRefreshCoordinator = ProviderRefreshCoordinator<UsageProvider>()    @ObservationIgnored private var providerAvailabilityCache: [UsageProvider: ProviderAvailabilityCacheEntry] = [:]    @ObservationIgnored var accountInfoCache: [UsageProvider: AccountInfoCacheEntry] = [:]    @ObservationIgnored private var timerTask: Task<Void, Never>?    @ObservationIgnored private var tokenTimerTask: Task<Void, Never>?    @ObservationIgnored private var tokenRefreshSequenceTask: Task<Void, Never>?    @ObservationIgnored var memoryPressureReliefTask: Task<Void, Never>?    @ObservationIgnored var startupConnectivityRetryTask: Task<Void, Never>?    @ObservationIgnored var startupConnectivityRetryNeeded = false    @ObservationIgnored var startupConnectivityRetryRefreshActive = false    @ObservationIgnored var storageRefreshTask: Task<Void, Never>?    @ObservationIgnored var storageRefreshGeneration: UInt64 = 0    @ObservationIgnored var storageRefreshInFlightSignature: String?    @ObservationIgnored var storageRefreshInFlightRequestKey: String?    @ObservationIgnored var lastStorageRefreshSignature: String?    @ObservationIgnored var lastStorageRefreshRequestKey: String?    @ObservationIgnored var lastStorageRefreshAt: Date?    @ObservationIgnored var managedCodexAccountsForStorageOverride: [ManagedCodexAccount]?    @ObservationIgnored private var pathDebugRefreshTask: Task<Void, Never>?    @ObservationIgnored var codexPlanHistoryBackfillTask: Task<Void, Never>?    @ObservationIgnored let historicalUsageHistoryStore: HistoricalUsageHistoryStore    @ObservationIgnored let planUtilizationHistoryStore: PlanUtilizationHistoryStore    @ObservationIgnored let codexAccountUsageSnapshotStore: (any CodexAccountUsageSnapshotStoring)?    @ObservationIgnored var codexHistoricalDataset: CodexHistoricalDataset?    @ObservationIgnored var codexHistoricalDatasetAccountKey: String?    @ObservationIgnored var lastKnownResetSnapshots: [UsageProvider: UsageSnapshot] = [:]    @ObservationIgnored var lastKnownSessionRemaining: [UsageProvider: Double] = [:]    @ObservationIgnored var lastKnownSessionWindowSource: [UsageProvider: SessionQuotaWindowSource] = [:]    @ObservationIgnored var quotaWarningState: [QuotaWarningStateKey: QuotaWarningState] = [:]    @ObservationIgnored var lastPermissionPromptNotificationAt: [UsageProvider: Date] = [:]    @ObservationIgnored var lastTokenFetchAt: [UsageProvider: Date] = [:]    @ObservationIgnored var lastTokenFetchScope: [UsageProvider: String] = [:]    @ObservationIgnored var planUtilizationHistory: [UsageProvider: PlanUtilizationHistoryBuckets] = [:]    @ObservationIgnored var weeklyLimitResetDetectorStates: [String: WeeklyLimitResetDetectorState] = [:]    @ObservationIgnored private var hasCompletedInitialRefresh: Bool = false    @ObservationIgnored private let providerAvailabilityCacheTTL: TimeInterval = 1    @ObservationIgnored let accountInfoCacheTTL: TimeInterval = 30    @ObservationIgnored private let tokenFetchTTL: TimeInterval = 60 * 60    @ObservationIgnored private let tokenFetchTimeout: TimeInterval = 10 * 60    @ObservationIgnored 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(),        codexAccountUsageSnapshotStore: (any CodexAccountUsageSnapshotStoring)? = nil,        sessionQuotaNotifier: any SessionQuotaNotifying = SessionQuotaNotifier(),        startupBehavior: StartupBehavior = .automatic,        environmentBase: [String: String] = ProcessInfo.processInfo.environment)    {        self.codexFetcher = fetcher        self.browserDetection = browserDetection        self.claudeFetcher = claudeFetcher ?? ClaudeUsageFetcher(browserDetection: browserDetection)        self.costUsageFetcher = costUsageFetcher        self.settings = settings        self.registry = registry        self.environmentBase = environmentBase        self.historicalUsageHistoryStore = historicalUsageHistoryStore        self.planUtilizationHistoryStore = planUtilizationHistoryStore        self.sessionQuotaNotifier = sessionQuotaNotifier        self.startupBehavior = startupBehavior.resolved(isRunningTests: Self.isRunningTestsProcess())        self.codexAccountUsageSnapshotStore = codexAccountUsageSnapshotStore ??            (self.startupBehavior.automaticallyStartsBackgroundWork ? FileCodexAccountUsageSnapshotStore() : nil)        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,            environmentBase: environmentBase)        self.providerRuntimes = Dictionary(uniqueKeysWithValues: ProviderCatalog.all.compactMap { implementation in            implementation.makeRuntime().map { (implementation.id, $0) }        })        self.planUtilizationHistory = planUtilizationHistoryStore.load()        self.weeklyLimitResetDetectorStates = Self.loadWeeklyLimitResetDetectorStates(from: settings.userDefaults)        if let codexAccountUsageSnapshotStore = self.codexAccountUsageSnapshotStore {            self.codexAccountSnapshots = codexAccountUsageSnapshotStore.load(                for: self.freshCodexVisibleAccountsForSnapshotHydration())        }        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.hydrateCachedTokenSnapshots()        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)        let now = Date()        return enabled.filter { self.isProviderAvailable($0, now: now) }    }     /// 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 {        self.isProviderAvailable(provider, now: Date())    }     private func isProviderAvailable(_ provider: UsageProvider, now: Date) -> Bool {        guard provider != .codex else { return true }         let configRevision = self.settings.configRevision        if let cached = self.providerAvailabilityCache[provider],           cached.isValid(now: now, configRevision: configRevision)        {

Lines 1–500 of 1582.