1164 lines · 46.7 KB · Swift
RawDownload
import CodexBarCoreimport Foundation extension UsageStore {    private nonisolated static let weeklyLimitResetThreshold = 1.0    private nonisolated static let weeklyLimitResetDetectorDefaultsKey = "weeklyLimitResetDetectorStates"    private nonisolated static let weeklyWindowMinutes = 7 * 24 * 60     struct WeeklyLimitResetDetectorState: Codable, Equatable {        let wasAboveThreshold: Bool        let lastObservedAt: Date    }     func supportsPlanUtilizationHistory(for provider: UsageProvider) -> Bool {        switch provider {        case .codex, .claude:            true        default:            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    }     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()        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,        shouldUpdatePreferredAccountKey: Bool = true,        shouldAdoptUnscopedHistory: Bool = true,        now: Date = Date())        async    {        let samples = self.planUtilizationSeriesSamples(provider: provider, snapshot: snapshot, capturedAt: now)        guard !samples.isEmpty else { return }         let detectorAccountKey = self.planUtilizationAccountKey(            for: provider,            snapshot: snapshot,            preferredAccount: account)        await MainActor.run {            self.postWeeklyLimitResetCelebrationIfNeeded(                provider: provider,                account: account,                snapshot: snapshot,                accountKey: detectorAccountKey,                samples: samples)        }         guard self.supportsPlanUtilizationHistory(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 preferredAccount = account ?? self.settings.selectedTokenAccount(for: provider)            let accountKey = self.resolvePlanUtilizationAccountKey(                provider: provider,                snapshot: snapshot,                preferredAccount: preferredAccount,                shouldUpdatePreferredAccountKey: shouldUpdatePreferredAccountKey,                shouldAdoptUnscopedHistory: shouldAdoptUnscopedHistory,                providerBuckets: &providerBuckets)            let histories = providerBuckets.histories(for: accountKey)             guard let updatedHistories = Self.updatedPlanUtilizationHistories(                existingHistories: histories,                samples: samples)            else {                return            }             providerBuckets.setHistories(updatedHistories, for: accountKey)            self.planUtilizationHistory[provider] = providerBuckets            self.planUtilizationHistoryRevision &+= 1            snapshotToPersist = self.planUtilizationHistory        }         guard let snapshotToPersist else { return }        await self.planUtilizationPersistenceCoordinator.enqueue(snapshotToPersist)    }     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 postWeeklyLimitResetCelebrationIfNeeded(        provider: UsageProvider,        account: ProviderTokenAccount?,        snapshot: UsageSnapshot,        accountKey: String?,        samples: [PlanUtilizationSeriesSample])    {        guard let weeklySample = samples.last(where: { $0.name == .weekly }) else { return }         let accountIdentifier = self.weeklyLimitResetAccountIdentifier(            provider: provider,            account: account,            snapshot: snapshot,            accountKey: accountKey)        let detectorKey = Self.weeklyLimitResetDetectorStateKey(            provider: provider,            accountIdentifier: accountIdentifier)        let currentUsed = weeklySample.entry.usedPercent        let currentObservedAt = weeklySample.entry.capturedAt        let wasAboveThreshold = currentUsed > Self.weeklyLimitResetThreshold        if let existingState = self.weeklyLimitResetDetectorStates[detectorKey],           currentObservedAt <= existingState.lastObservedAt        {            return        }         let shouldPost = self.weeklyLimitResetDetectorStates[detectorKey]?.wasAboveThreshold == true            && !wasAboveThreshold        self.weeklyLimitResetDetectorStates[detectorKey] = WeeklyLimitResetDetectorState(            wasAboveThreshold: wasAboveThreshold,            lastObservedAt: currentObservedAt)        self.persistWeeklyLimitResetDetectorStates()         guard shouldPost else { return }        let accountLabel = self.weeklyLimitResetAccountLabel(            provider: provider,            account: account,            snapshot: snapshot)        let event = WeeklyLimitResetEvent(            provider: provider,            accountIdentifier: accountIdentifier,            accountLabel: accountLabel,            usedPercent: currentUsed)         CodexBarLog.logger(LogCategories.confetti).info(            "Weekly limit reset",            metadata: [                "provider": provider.rawValue,                "accountIdentifier": accountIdentifier,                "accountLabel": accountLabel ?? "",                "usedPercent": String(format: "%.2f", currentUsed),                "observedAt": String(format: "%.0f", currentObservedAt.timeIntervalSince1970),            ])        NotificationCenter.default.post(name: .codexbarWeeklyLimitReset, object: event)    }     private func planUtilizationSeriesSamples(        provider: UsageProvider,        snapshot: UsageSnapshot,        capturedAt: Date) -> [PlanUtilizationSeriesSample]    {        var samplesByKey: [PlanUtilizationSeriesKey: PlanUtilizationSeriesSample] = [:]         func appendWindow(_ window: RateWindow?, name: PlanUtilizationSeriesName?) {            guard let name,                  let window,                  let windowMinutes = window.windowMinutes,                  windowMinutes > 0,                  let usedPercent = Self.clampedPercent(window.usedPercent)            else {                return            }             let canonicalWindowMinutes = name.canonicalWindowMinutes(windowMinutes)            let key = PlanUtilizationSeriesKey(name: name, windowMinutes: canonicalWindowMinutes)            samplesByKey[key] = PlanUtilizationSeriesSample(                name: name,                windowMinutes: canonicalWindowMinutes,                entry: PlanUtilizationHistoryEntry(                    capturedAt: capturedAt,                    usedPercent: usedPercent,                    resetsAt: window.resetsAt))        }         switch provider {        case .codex:            let projection = self.codexConsumerProjection(                surface: .liveCard,                snapshotOverride: snapshot,                now: capturedAt)            for lane in projection.planUtilizationLanes {                appendWindow(lane.window, name: lane.role)            }        case .claude:            appendWindow(snapshot.primary, name: .session)            appendWindow(snapshot.secondary, name: .weekly)            appendWindow(snapshot.tertiary, name: .opus)        default:            for window in [snapshot.primary, snapshot.secondary, snapshot.tertiary] {                guard let window, window.windowMinutes == Self.weeklyWindowMinutes else { continue }                appendWindow(window, name: .weekly)            }        }         return samplesByKey.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 planUtilizationHourBucket(for date: Date) -> Int64 {        Int64(floor(date.timeIntervalSince1970 / self.planUtilizationMinSampleIntervalSeconds))    }     private nonisolated static func planUtilizationHourRange(        entries: [PlanUtilizationHistoryEntry],        insertionIndex: Int,        hourBucket: Int64) -> Range<Int>    {        var lowerBound = insertionIndex        while lowerBound > entries.startIndex {            let previousIndex = lowerBound - 1            let previousHourBucket = self.planUtilizationHourBucket(for: entries[previousIndex].capturedAt)            guard previousHourBucket == hourBucket else { break }            lowerBound = previousIndex        }         var upperBound = insertionIndex        while upperBound < entries.endIndex {            let currentHourBucket = self.planUtilizationHourBucket(for: entries[upperBound].capturedAt)            guard currentHourBucket == hourBucket else { break }            upperBound += 1        }         return lowerBound..<upperBound    }     private nonisolated static func canonicalPlanUtilizationHourEntries(        existingHourEntries: [PlanUtilizationHistoryEntry],        incomingEntry: PlanUtilizationHistoryEntry) -> [PlanUtilizationHistoryEntry]    {        let hourlyObservations = (existingHourEntries + [incomingEntry]).sorted { lhs, rhs in            if lhs.capturedAt != rhs.capturedAt {                return lhs.capturedAt < rhs.capturedAt            }            if lhs.usedPercent != rhs.usedPercent {                return lhs.usedPercent < rhs.usedPercent            }            let lhsReset = lhs.resetsAt?.timeIntervalSince1970 ?? Date.distantPast.timeIntervalSince1970            let rhsReset = rhs.resetsAt?.timeIntervalSince1970 ?? Date.distantPast.timeIntervalSince1970            return lhsReset < rhsReset        }        guard var activeSegmentPeak = hourlyObservations.first else { return [] }         var peakBeforeLatestReset: PlanUtilizationHistoryEntry?         for observation in hourlyObservations.dropFirst() {            if self.startsNewPlanUtilizationResetSegment(                activeSegmentPeak: activeSegmentPeak,                observation: observation)            {                if peakBeforeLatestReset == nil {                    peakBeforeLatestReset = activeSegmentPeak                }                activeSegmentPeak = observation                continue            }             activeSegmentPeak = self.segmentPeakEntry(                existingPeak: activeSegmentPeak,                observation: observation)        }         if let peakBeforeLatestReset {            return [peakBeforeLatestReset, activeSegmentPeak]        }        return [activeSegmentPeak]    }     private nonisolated static func startsNewPlanUtilizationResetSegment(        activeSegmentPeak: PlanUtilizationHistoryEntry,        observation: PlanUtilizationHistoryEntry) -> Bool    {        self.haveMeaningfullyDifferentResetBoundaries(            activeSegmentPeak.resetsAt,            observation.resetsAt)    }     private nonisolated static func segmentPeakEntry(        existingPeak: PlanUtilizationHistoryEntry,        observation: PlanUtilizationHistoryEntry) -> PlanUtilizationHistoryEntry    {        if existingPeak.resetsAt == nil, observation.resetsAt != nil {            return observation        }         let hasHigherUsage = observation.usedPercent > existingPeak.usedPercent        let tiesUsageAndIsMoreRecent = observation.usedPercent == existingPeak.usedPercent            && observation.capturedAt >= existingPeak.capturedAt        let observationShouldReplacePeak = hasHigherUsage || tiesUsageAndIsMoreRecent        let peakSource = observationShouldReplacePeak ? observation : existingPeak        let preferObservationMetadata = observation.capturedAt >= existingPeak.capturedAt         return PlanUtilizationHistoryEntry(            capturedAt: peakSource.capturedAt,            usedPercent: peakSource.usedPercent,            resetsAt: self.preferredResetBoundary(                existing: existingPeak.resetsAt,                incoming: observation.resetsAt,                preferIncoming: preferObservationMetadata))    }

Lines 1–500 of 1164.