1678 lines · 69.9 KB · Swift
RawDownload
import CodexBarCoreimport Foundation extension UsageStore {    private nonisolated static let limitResetThreshold = 1.0    nonisolated static let sessionLimitResetDetectorDefaultsKey = "sessionLimitResetDetectorStates"    private nonisolated static let weeklyLimitResetDetectorDefaultsKey = "weeklyLimitResetDetectorStates"    private nonisolated static let claudeOAuthAccountUuidMapDefaultsKey = "ClaudeOAuthHistoryOwnerAccountUuidMapV1"    private nonisolated static let claudeOAuthAccountCandidateMapDefaultsKey =        "ClaudeOAuthHistoryOwnerAccountCandidateMapV1"    private nonisolated static let weeklyWindowMinutes = 7 * 24 * 60    private nonisolated static let planUtilizationUnscopedPreferredKey = "__unscoped__"    private nonisolated static let claudeOAuthPlanUtilizationAccountKeyPrefix = "__claude_oauth__:"     enum ClaudeOAuthActiveAccountObservation: Equatable, Sendable {        case stable(identity: String?)        case changed    }     struct ClaudeOAuthAccountBindingCandidate: Codable, Equatable {        let identity: String        let observedAt: Date    }     private struct ClaudeOAuthHistoryEvidence {        let owner: String        let persistentRefHash: String?        let keychainCredentialMismatch: Bool        let keychainCredentialAbsent: Bool        let keychainCredentialUnavailable: Bool        let activeAccountObservation: ClaudeOAuthActiveAccountObservation        let observedAt: Date    }     struct LimitResetDetectorState: Codable, Equatable {        let wasAboveThreshold: Bool        let lastObservedAt: Date        let sourceRawValue: String?    }     func supportsPlanUtilizationHistory(for provider: UsageProvider) -> Bool {        switch provider {        case .codex, .claude:            true        default:            if self.planUtilizationHistory[provider]?.isEmpty == false {                true            } else if self.settings.historicalTrackingEnabled, let snapshot = self.snapshots[provider] {                !self.planUtilizationSeriesSamples(                    provider: provider,                    snapshot: snapshot,                    capturedAt: snapshot.updatedAt).isEmpty            } else {                false            }        }    }     private nonisolated static let planUtilizationMinSampleIntervalSeconds: TimeInterval = 60 * 60    private nonisolated static let planUtilizationResetEquivalenceToleranceSeconds: TimeInterval = 2 * 60    private nonisolated static let planUtilizationMaxSamples: Int = 24 * 730     private struct PlanUtilizationSeriesKey: Hashable {        let name: PlanUtilizationSeriesName        let windowMinutes: Int    }     private struct PlanUtilizationSeriesSample {        let name: PlanUtilizationSeriesName        let windowMinutes: Int        let entry: PlanUtilizationHistoryEntry    }     private struct LimitResetDetectionContext {        let provider: UsageProvider        let account: ProviderTokenAccount?        let snapshot: UsageSnapshot        let accountKey: String?        let capturedAt: Date    }     private struct LimitResetObservation {        let usedPercent: Double        let observedAt: Date        let source: SessionQuotaWindowSource?    }     private struct LimitResetDetectionDescriptor {        let seriesName: PlanUtilizationSeriesName        let defaultsKey: String        let resetKind: String    }     func planUtilizationHistory(for provider: UsageProvider) -> [PlanUtilizationSeriesHistory] {        self.planUtilizationHistorySelection(for: provider).histories    }     func planUtilizationHistorySelection(for provider: UsageProvider)        -> (accountKey: String?, histories: [PlanUtilizationSeriesHistory])    {        var providerBuckets = self.planUtilizationHistory[provider] ?? PlanUtilizationHistoryBuckets()        if provider == .claude,           providerBuckets.preferredAccountKey == Self.planUtilizationUnscopedPreferredKey           || Self.isClaudeOAuthPlanUtilizationAccountKey(providerBuckets.preferredAccountKey)        {            // Persisted OAuth provenance outranks an unrelated configured token account. The unscoped            // sentinel intentionally resolves to nil, including after the history store is reloaded.            let accountKey = self.stickyPlanUtilizationAccountKey(providerBuckets: providerBuckets)            return (accountKey, providerBuckets.histories(for: accountKey))        }        let originalProviderBuckets = providerBuckets        let accountKey = self.resolvePlanUtilizationAccountKey(            provider: provider,            snapshot: self.snapshots[provider],            preferredAccount: nil,            providerBuckets: &providerBuckets)        self.planUtilizationHistory[provider] = providerBuckets        if providerBuckets != originalProviderBuckets {            let snapshotToPersist = self.planUtilizationHistory            Task {                await self.planUtilizationPersistenceCoordinator.enqueue(snapshotToPersist)            }        }        return (accountKey, providerBuckets.histories(for: accountKey))    }     func codexPlanUtilizationHistories(forVisibleAccount account: CodexVisibleAccount)        -> [PlanUtilizationSeriesHistory]    {        var providerBuckets = self.planUtilizationHistory[.codex] ?? PlanUtilizationHistoryBuckets()        let originalProviderBuckets = providerBuckets        let ownership = self.codexOwnershipContext(forVisibleAccount: account)        guard let canonicalKey = ownership.canonicalKey else { return [] }         if ownership.hasAdjacentEmailScopeAmbiguity {            guard canonicalKey != ownership.canonicalEmailHashKey else { return [] }            return providerBuckets.histories(for: canonicalKey)        }         let accountKey = self.materializeCodexPlanUtilizationHistoryIfNeeded(            into: canonicalKey,            ownership: ownership,            shouldAdoptUnscopedHistory: true,            providerBuckets: &providerBuckets)        self.planUtilizationHistory[.codex] = providerBuckets        if providerBuckets != originalProviderBuckets {            let snapshotToPersist = self.planUtilizationHistory            Task {                await self.planUtilizationPersistenceCoordinator.enqueue(snapshotToPersist)            }        }        return providerBuckets.histories(for: accountKey)    }     func shouldShowRefreshingMenuCard(for provider: UsageProvider) -> Bool {        let isRefreshing = self.isRefreshing || self.refreshingProviders.contains(provider)        return isRefreshing            && self.snapshots[provider] == nil            && self.error(for: provider) == nil    }     func shouldShowRefreshingMenuCardIndicator(for provider: UsageProvider) -> Bool {        let isRefreshing = self.isRefreshing || self.refreshingProviders.contains(provider)        return isRefreshing && self.error(for: provider) == nil    }     func shouldHidePlanUtilizationMenuItem(for provider: UsageProvider) -> Bool {        guard self.supportsPlanUtilizationHistory(for: provider) else { return true }        return self.shouldShowRefreshingMenuCard(for: provider)    }     func recordPlanUtilizationHistorySample(        provider: UsageProvider,        snapshot: UsageSnapshot,        account: ProviderTokenAccount? = nil,        claudeOAuthPersistentRefHash: String? = nil,        claudeOAuthHistoryOwnerIdentifier: String? = nil,        claudeOAuthKeychainCredentialMismatch: Bool = false,        claudeOAuthKeychainCredentialAbsent: Bool = false,        claudeOAuthKeychainCredentialUnavailable: Bool = false,        claudeOAuthActiveAccountObservation: ClaudeOAuthActiveAccountObservation = .stable(identity: nil),        isClaudeOAuthSample: Bool = false,        shouldUpdatePreferredAccountKey: Bool = true,        shouldAdoptUnscopedHistory: Bool = true,        now: Date = Date())        async    {        let samples = self.planUtilizationSeriesSamples(provider: provider, snapshot: snapshot, capturedAt: now)        var effectiveOwner = claudeOAuthHistoryOwnerIdentifier        if provider == .claude, isClaudeOAuthSample, let owner = claudeOAuthHistoryOwnerIdentifier {            effectiveOwner = self.resolvedClaudeOAuthHistoryOwner(evidence: ClaudeOAuthHistoryEvidence(                owner: owner,                persistentRefHash: claudeOAuthPersistentRefHash,                keychainCredentialMismatch: claudeOAuthKeychainCredentialMismatch,                keychainCredentialAbsent: claudeOAuthKeychainCredentialAbsent,                keychainCredentialUnavailable: claudeOAuthKeychainCredentialUnavailable,                activeAccountObservation: claudeOAuthActiveAccountObservation,                observedAt: now))        }        let detectorAccountKey = if provider == .claude, isClaudeOAuthSample {            Self.claudeOAuthPlanUtilizationAccountKey(                historyOwnerIdentifier: effectiveOwner,                corroboratingPersistentRefHash: claudeOAuthPersistentRefHash)        } else {            self.planUtilizationAccountKey(                for: provider,                snapshot: snapshot,                preferredAccount: account)        }        if provider == .claude, isClaudeOAuthSample, detectorAccountKey == nil {            // Persisting without a high-entropy owner would merge unrelated OAuth accounts into `unscoped`.            return        }        let detectorContext = LimitResetDetectionContext(            provider: provider,            account: account,            snapshot: snapshot,            accountKey: detectorAccountKey,            capturedAt: now)        await MainActor.run {            self.postLimitResetCelebrationsIfNeeded(                context: detectorContext,                samples: samples)        }         guard !samples.isEmpty else { return }        guard self.shouldRecordPlanUtilizationHistory(for: provider) else { return }        guard !self.shouldDeferClaudePlanUtilizationHistory(provider: provider) else { return }         var snapshotToPersist: [UsageProvider: PlanUtilizationHistoryBuckets]?        await MainActor.run {            var providerBuckets = self.planUtilizationHistory[provider] ?? PlanUtilizationHistoryBuckets()            let originalProviderBuckets = providerBuckets            let preferredAccount = account ?? self.settings.selectedTokenAccount(for: provider)            let accountKey = self.resolvePlanUtilizationAccountKey(                provider: provider,                snapshot: snapshot,                preferredAccount: preferredAccount,                claudeOAuthPersistentRefHash: claudeOAuthPersistentRefHash,                claudeOAuthHistoryOwnerIdentifier: effectiveOwner,                isClaudeOAuthSample: isClaudeOAuthSample,                shouldUpdatePreferredAccountKey: shouldUpdatePreferredAccountKey,                shouldAdoptUnscopedHistory: shouldAdoptUnscopedHistory,                providerBuckets: &providerBuckets)            let histories = providerBuckets.histories(for: accountKey)             if let updatedHistories = Self.updatedPlanUtilizationHistories(                existingHistories: histories,                samples: samples)            {                providerBuckets.setHistories(updatedHistories, for: accountKey)            }             guard providerBuckets != originalProviderBuckets else { return }            self.planUtilizationHistory[provider] = providerBuckets            self.planUtilizationHistoryRevision &+= 1            snapshotToPersist = self.planUtilizationHistory        }         guard let snapshotToPersist else { return }        await self.planUtilizationPersistenceCoordinator.enqueue(snapshotToPersist)    }     private func shouldRecordPlanUtilizationHistory(for provider: UsageProvider) -> Bool {        switch provider {        case .codex, .claude:            true        default:            self.settings.historicalTrackingEnabled        }    }     private nonisolated static func updatedPlanUtilizationHistories(        existingHistories: [PlanUtilizationSeriesHistory],        samples: [PlanUtilizationSeriesSample]) -> [PlanUtilizationSeriesHistory]?    {        guard !samples.isEmpty else { return nil }         var historiesByKey: [PlanUtilizationSeriesKey: PlanUtilizationSeriesHistory] = [:]        var didChange = false        for history in existingHistories {            let canonicalWindowMinutes = history.name.canonicalWindowMinutes(history.windowMinutes)            let key = PlanUtilizationSeriesKey(name: history.name, windowMinutes: canonicalWindowMinutes)            let canonicalHistory = PlanUtilizationSeriesHistory(                name: history.name,                windowMinutes: canonicalWindowMinutes,                entries: history.entries)            if let existingHistory = historiesByKey[key] {                historiesByKey[key] = PlanUtilizationSeriesHistory(                    name: history.name,                    windowMinutes: canonicalWindowMinutes,                    entries: self.mergedPlanUtilizationEntries(existingHistory.entries + canonicalHistory.entries))                didChange = true            } else {                historiesByKey[key] = canonicalHistory                didChange = didChange || canonicalWindowMinutes != history.windowMinutes            }        }         for sample in samples {            let canonicalWindowMinutes = sample.name.canonicalWindowMinutes(sample.windowMinutes)            let key = PlanUtilizationSeriesKey(name: sample.name, windowMinutes: canonicalWindowMinutes)            if let existingHistory = historiesByKey[key] {                guard let updatedEntries = self.updatedPlanUtilizationEntries(                    existingEntries: existingHistory.entries,                    entry: sample.entry)                else {                    continue                }                historiesByKey[key] = PlanUtilizationSeriesHistory(                    name: sample.name,                    windowMinutes: canonicalWindowMinutes,                    entries: updatedEntries)            } else {                historiesByKey[key] = PlanUtilizationSeriesHistory(                    name: sample.name,                    windowMinutes: canonicalWindowMinutes,                    entries: [sample.entry])            }            didChange = true        }         guard didChange else { return nil }        return historiesByKey.values.sorted { lhs, rhs in            if lhs.windowMinutes != rhs.windowMinutes {                return lhs.windowMinutes < rhs.windowMinutes            }            return lhs.name.rawValue < rhs.name.rawValue        }    }     private nonisolated static func mergedPlanUtilizationEntries(        _ entries: [PlanUtilizationHistoryEntry]) -> [PlanUtilizationHistoryEntry]    {        entries.reduce(into: []) { result, entry in            guard !result.contains(entry) else { return }            result.append(entry)        }    }     private nonisolated static func updatedPlanUtilizationEntries(        existingEntries: [PlanUtilizationHistoryEntry],        entry: PlanUtilizationHistoryEntry) -> [PlanUtilizationHistoryEntry]?    {        var entries = existingEntries        let insertionIndex = entries.firstIndex(where: { $0.capturedAt > entry.capturedAt }) ?? entries.endIndex        let sampleHourBucket = self.planUtilizationHourBucket(for: entry.capturedAt)        let sameHourRange = self.planUtilizationHourRange(            entries: entries,            insertionIndex: insertionIndex,            hourBucket: sampleHourBucket)        let existingHourEntries = Array(entries[sameHourRange])        let canonicalHourEntries = self.canonicalPlanUtilizationHourEntries(            existingHourEntries: existingHourEntries,            incomingEntry: entry)         guard canonicalHourEntries != existingHourEntries else { return nil }        entries.replaceSubrange(sameHourRange, with: canonicalHourEntries)         if entries.count > self.planUtilizationMaxSamples {            entries.removeFirst(entries.count - self.planUtilizationMaxSamples)        }        return entries    }     #if DEBUG    nonisolated static func _updatedPlanUtilizationEntriesForTesting(        existingEntries: [PlanUtilizationHistoryEntry],        entry: PlanUtilizationHistoryEntry) -> [PlanUtilizationHistoryEntry]?    {        self.updatedPlanUtilizationEntries(existingEntries: existingEntries, entry: entry)    }     nonisolated static func _updatedPlanUtilizationHistoriesForTesting(        existingHistories: [PlanUtilizationSeriesHistory],        samples: [PlanUtilizationSeriesHistory]) -> [PlanUtilizationSeriesHistory]?    {        let normalized = samples.flatMap { history in            history.entries.map { entry in                PlanUtilizationSeriesSample(name: history.name, windowMinutes: history.windowMinutes, entry: entry)            }        }        return self.updatedPlanUtilizationHistories(existingHistories: existingHistories, samples: normalized)    }     nonisolated static var _planUtilizationMaxSamplesForTesting: Int {        self.planUtilizationMaxSamples    }    #endif     private nonisolated static func clampedPercent(_ value: Double?) -> Double? {        guard let value else { return nil }        return max(0, min(100, value))    }     private func postLimitResetCelebrationsIfNeeded(        context: LimitResetDetectionContext,        samples: [PlanUtilizationSeriesSample])    {        let shouldIgnoreCommandCode = context.provider == .commandcode            && context.snapshot.commandCodeSubscriptionEnrichmentUnavailable        let sessionObservation: LimitResetObservation? = if shouldIgnoreCommandCode {            nil        } else if context.provider == .codex {            samples.last(where: { $0.name == .session }).map {                LimitResetObservation(                    usedPercent: $0.entry.usedPercent,                    observedAt: $0.entry.capturedAt,                    source: nil)            }        } else {            self.sessionQuotaWindow(provider: context.provider, snapshot: context.snapshot).flatMap { resolved in                guard Self.isSemanticSessionResetWindow(resolved) else { return nil }                return Self.clampedPercent(resolved.window.usedPercent).map {                    LimitResetObservation(                        usedPercent: $0,                        observedAt: context.capturedAt,                        source: resolved.source)                }            }        }        self.postLimitResetCelebrationIfNeeded(            states: &self.sessionLimitResetDetectorStates,            context: context,            descriptor: LimitResetDetectionDescriptor(                seriesName: .session,                defaultsKey: Self.sessionLimitResetDetectorDefaultsKey,                resetKind: "session"),            observation: sessionObservation)        let weeklyObservation = samples.last(where: { $0.name == .weekly }).map {            LimitResetObservation(                usedPercent: $0.entry.usedPercent,                observedAt: $0.entry.capturedAt,                source: nil)        }        self.postLimitResetCelebrationIfNeeded(            states: &self.weeklyLimitResetDetectorStates,            context: context,            descriptor: LimitResetDetectionDescriptor(                seriesName: .weekly,                defaultsKey: Self.weeklyLimitResetDetectorDefaultsKey,                resetKind: "weekly"),            observation: weeklyObservation)    }     private static func isSemanticSessionResetWindow(        _ resolved: (window: RateWindow, source: SessionQuotaWindowSource)) -> Bool    {        switch resolved.source {        case .primary:            guard let minutes = resolved.window.windowMinutes else { return false }            return minutes > 0 && minutes <= 6 * 60        case .copilotSecondaryFallback, .zaiTertiary, .antigravityQuotaSummary, .antigravityLegacy:            return true        }    }     private func postLimitResetCelebrationIfNeeded(        states: inout [String: LimitResetDetectorState],        context: LimitResetDetectionContext,        descriptor: LimitResetDetectionDescriptor,        observation: LimitResetObservation?)    {        guard let observation else { return }         let accountIdentifier = self.limitResetAccountIdentifier(            provider: context.provider,            account: context.account,            snapshot: context.snapshot,            accountKey: context.accountKey)        let detectorKey = Self.limitResetDetectorStateKey(            provider: context.provider,            accountIdentifier: accountIdentifier)        let currentUsed = observation.usedPercent        let currentObservedAt = observation.observedAt        let wasAboveThreshold = currentUsed > Self.limitResetThreshold        if let existingState = states[detectorKey],           currentObservedAt <= existingState.lastObservedAt        {            return        }         let previousState = states[detectorKey]        let sourceRawValue = observation.source?.rawValue        let sourceChanged = descriptor.seriesName == .session            && previousState?.sourceRawValue != nil            && previousState?.sourceRawValue != sourceRawValue        let shouldPost = !sourceChanged            && previousState?.wasAboveThreshold == true            && !wasAboveThreshold        states[detectorKey] = LimitResetDetectorState(            wasAboveThreshold: wasAboveThreshold,            lastObservedAt: currentObservedAt,            sourceRawValue: sourceRawValue)        self.persistLimitResetDetectorStates(            states,            defaultsKey: descriptor.defaultsKey,            logName: descriptor.resetKind)         guard shouldPost else { return }

Lines 1–500 of 1678.

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t