619 lines · 26.7 KB · Swift
RawDownload
import CodexBarCoreimport Foundation extension UsageStore {    private struct ProviderRefreshOutcomeContext {        let generation: UInt64        let codexExpectedGuard: CodexAccountScopedRefreshGuard?        let claudeCredentialsChanged: Bool        let shouldConsumeClaudeKeychainFingerprint: Bool    }     func refreshForSettingsChange() async {        await self.runRefresh(            startupConnectivityRetryAttempt: nil,            coalesceProviderRefreshesOverride: false)    }     func prepareRefreshState(for provider: UsageProvider? = nil) {        guard provider == nil || provider == .codex else { return }        _ = self.settings.persistResolvedCodexActiveSourceCorrectionIfNeeded()    }     /// Force refresh Augment session (called from UI button)    func forceRefreshAugmentSession() async {        await self.performRuntimeAction(.forceSessionRefresh, for: .augment)    }     private func providerRefreshSpec(_ provider: UsageProvider) async -> ProviderSpec? {        if let override = self._test_providerRefreshOverride {            await override(provider)            return nil        }        return self.providerSpecs[provider]    }     func refreshProvider(        _ provider: UsageProvider,        allowDisabled: Bool = false,        coalesceIfRefreshing: Bool = false) async    {        while coalesceIfRefreshing,              let states = self.providerRefreshTasks[provider],              let latestGeneration = self.latestProviderRefreshGenerations[provider],              let existingState = states.last(where: { $0.generation == latestGeneration })        {            await self.waitForProviderRefresh(provider, state: existingState)            if Task.isCancelled { return }            if existingState.shouldRetry {                self.removeProviderRefreshTask(provider, state: existingState)                continue            }            return        }         self.providerRefreshTaskGeneration &+= 1        let generation = self.providerRefreshTaskGeneration        let predecessorStates = self.providerRefreshTasks[provider] ?? []        for predecessorState in predecessorStates {            predecessorState.cancelTask()        }        self.latestProviderRefreshGenerations[provider] = generation        let state = ProviderRefreshTaskState(generation: generation)        let task = Task { @MainActor [weak self] in            guard let self else { return }            var snapshotUpdatedAtBeforeRefresh: Date?            var didStartRefresh = false            for predecessorState in predecessorStates {                await predecessorState.waitForTaskCompletion()            }            if !Task.isCancelled, self.isCurrentProviderRefreshGeneration(provider, generation: generation) {                snapshotUpdatedAtBeforeRefresh = self.snapshot(for: provider)?.updatedAt                didStartRefresh = true                await self.refreshProviderTracked(                    provider,                    allowDisabled: allowDisabled,                    generation: generation)            }            let publishedNewSnapshot = didStartRefresh &&                self.snapshot(for: provider)?.updatedAt != snapshotUpdatedAtBeforeRefresh            let retryRequired = Task.isCancelled && !publishedNewSnapshot            self.providerRefreshDidComplete(provider, state: state, retryRequired: retryRequired)        }        state.install(task: task)        self.providerRefreshTasks[provider, default: []].append(state)        await self.waitForProviderRefresh(provider, state: state)    }     private func waitForProviderRefresh(_ provider: UsageProvider, state: ProviderRefreshTaskState) async {        self.providerRefreshWaiterGeneration &+= 1        let waiterID = self.providerRefreshWaiterGeneration        guard let task = state.addWaiter(waiterID) else { return }        await withTaskCancellationHandler {            await task.value        } onCancel: {            state.cancelWaiter(waiterID)        }        state.finishWaiter(waiterID)        if state.canRemove {            self.scheduleProviderRefreshTaskRemoval(provider, state: state)        }    }     private func providerRefreshDidComplete(        _ provider: UsageProvider,        state: ProviderRefreshTaskState,        retryRequired: Bool)    {        state.markCompleted(retryRequired: retryRequired)        self.scheduleProviderRefreshTaskRemoval(provider, state: state)    }     private func removeProviderRefreshTask(_ provider: UsageProvider, state: ProviderRefreshTaskState) {        guard var states = self.providerRefreshTasks[provider] else { return }        states.removeAll { $0 === state }        if states.isEmpty {            self.providerRefreshTasks.removeValue(forKey: provider)        } else {            self.providerRefreshTasks[provider] = states        }    }     private func scheduleProviderRefreshTaskRemoval(_ provider: UsageProvider, state: ProviderRefreshTaskState) {        Task { @MainActor [weak self] in            await Task.yield()            guard let self,                  self.providerRefreshTasks[provider]?.contains(where: { $0 === state }) == true,                  state.canRemove            else {                return            }            self.removeProviderRefreshTask(provider, state: state)        }    }     func isCurrentProviderRefreshGeneration(_ provider: UsageProvider, generation: UInt64?) -> Bool {        guard let generation else { return true }        return self.latestProviderRefreshGenerations[provider] == generation    }     private func refreshProviderTracked(        _ provider: UsageProvider,        allowDisabled: Bool,        generation: UInt64) async    {        self.providerRefreshCounts[provider, default: 0] += 1        self.refreshingProviders.insert(provider)        defer {            let remaining = max(0, self.providerRefreshCounts[provider, default: 1] - 1)            if remaining == 0 {                self.providerRefreshCounts.removeValue(forKey: provider)                self.refreshingProviders.remove(provider)            } else {                self.providerRefreshCounts[provider] = remaining            }        }        await self.refreshProviderNow(            provider,            allowDisabled: allowDisabled,            generation: generation)    }     private func refreshProviderNow(        _ provider: UsageProvider,        allowDisabled: Bool,        generation: UInt64) async    {        self.prepareRefreshState(for: provider)        guard let spec = await self.providerRefreshSpec(provider) else { return }        guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }        let codexExpectedGuard = provider == .codex ? self.freshCodexAccountScopedRefreshGuard() : nil         if !spec.isEnabled(), !allowDisabled {            await self.clearDisabledProviderRefreshState(provider)            return        }         if provider == .codex, self.shouldFetchAllCodexVisibleAccounts() {            await self.refreshCodexVisibleAccountsForMenu(generation: generation)            return        } else if provider == .codex {            self.codexAccountSnapshots = []        }         if provider == .kilo, self.shouldFanOutKiloScopes() {            await self.refreshKiloScopes(generation: generation)            guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }            // Continue to also fetch the personal snapshot through the regular path            // so the existing single-card render keeps working when only personal is shown.            // The presence of multi-element kiloScopeSnapshots triggers stacked rendering.        } else if provider == .kilo {            await MainActor.run { self.kiloScopeSnapshots = [] }        }         let tokenAccounts = self.tokenAccounts(for: provider)        if self.shouldFetchAllTokenAccounts(provider: provider, accounts: tokenAccounts) {            await self.refreshTokenAccounts(                provider: provider,                accounts: tokenAccounts,                generation: generation)            return        } else {            _ = await MainActor.run {                self.accountSnapshots.removeValue(forKey: provider)            }        }         let claudeAuthStateBeforeFetch = provider == .claude            ? await Self.captureClaudeRefreshAuthState(invalidateCredentialsFile: true)            : nil        let fetchContext = self.makeFetchContext(provider: provider, override: nil)        let descriptor = spec.descriptor        // Keep provider fetch work off MainActor so slow keychain/process reads don't stall menu/UI responsiveness.        let outcome = await withTaskGroup(            of: ProviderFetchOutcome.self,            returning: ProviderFetchOutcome.self)        { group in            group.addTask {                await descriptor.fetchOutcome(context: fetchContext)            }            return await group.next()!        }        guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }        let claudeAuthFingerprintAfterFetch = provider == .claude            ? await Self.captureClaudeAuthFingerprintToken()            : nil        let claudeAuthChangedDuringFetch = Self.claudeAuthChangedDuringFetch(            provider: provider,            beforeFetch: claudeAuthStateBeforeFetch,            afterFetchFingerprintToken: claudeAuthFingerprintAfterFetch)        await Self.invalidateClaudeCredentialsFileCacheIfNeeded(changedDuringFetch: claudeAuthChangedDuringFetch)        let claudeCredentialsChanged = Self.claudeCredentialsChanged(            beforeFetch: claudeAuthStateBeforeFetch,            changedDuringFetch: claudeAuthChangedDuringFetch)        let shouldConsumeClaudeKeychainFingerprint = Self.shouldConsumeClaudeKeychainFingerprintChange(            beforeFetch: claudeAuthStateBeforeFetch,            changedDuringFetch: claudeAuthChangedDuringFetch)        guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }        await self.applyProviderRefreshOutcome(            provider: provider,            outcome: outcome,            context: ProviderRefreshOutcomeContext(                generation: generation,                codexExpectedGuard: codexExpectedGuard,                claudeCredentialsChanged: claudeCredentialsChanged,                shouldConsumeClaudeKeychainFingerprint: shouldConsumeClaudeKeychainFingerprint))    }     private func applyProviderRefreshOutcome(        provider: UsageProvider,        outcome: ProviderFetchOutcome,        context: ProviderRefreshOutcomeContext) async    {        await MainActor.run {            self.lastFetchAttempts[provider] = outcome.attempts        }         switch outcome.result {        case let .success(result):            let scoped = result.usage.scoped(to: provider)            if provider == .codex,               let codexExpectedGuard = context.codexExpectedGuard,               !self.shouldApplyCodexUsageResult(expectedGuard: codexExpectedGuard, usage: scoped)            {                return            }            let backfilled = await MainActor.run { () -> UsageSnapshot? in                guard self.isCurrentProviderRefreshGeneration(provider, generation: context.generation) else {                    return nil                }                if context.claudeCredentialsChanged {                    self.clearClaudeCredentialDerivedStateForCredentialSwapNow()                }                let resetBackfillSource = provider == .codex                    ? self.codexLastKnownResetSnapshot(matching: context.codexExpectedGuard)                    : self.lastKnownResetSnapshots[provider]                let backfilled = scoped.backfillingResetTimes(from: resetBackfillSource)                self.handleQuotaWarningTransitions(provider: provider, snapshot: backfilled)                self.handleSessionQuotaTransition(provider: provider, snapshot: backfilled)                self.lastKnownResetSnapshots[provider] = backfilled                self.snapshots[provider] = backfilled                if let tokenSnapshot = self.tokenSnapshot(fromProviderSnapshot: backfilled, provider: provider) {                    self.tokenSnapshots[provider] = tokenSnapshot                    self.tokenErrors[provider] = nil                    self.tokenFailureGates[provider]?.recordSuccess()                } else if Self.tokenCostRequiresProviderSnapshot(provider) {                    self.tokenSnapshots.removeValue(forKey: provider)                    self.tokenErrors[provider] = nil                }                self.lastSourceLabels[provider] = result.sourceLabel                self.errors[provider] = nil                self.failureGates[provider]?.recordSuccess()                if provider == .codex {                    self.rememberLiveSystemCodexEmailIfNeeded(scoped.accountEmail(for: .codex))                    self.seedCodexAccountScopedRefreshGuard(accountEmail: scoped.accountEmail(for: .codex))                }                return backfilled            }            guard let backfilled else { return }            if context.shouldConsumeClaudeKeychainFingerprint {                _ = await Self.consumeClaudeKeychainFingerprintChangeWithoutPrompt()            }            await self.recordPlanUtilizationHistorySample(                provider: provider,                snapshot: backfilled)            guard self.isCurrentProviderRefreshGeneration(provider, generation: context.generation) else { return }            if let runtime = self.providerRuntimes[provider] {                let context = ProviderRuntimeContext(                    provider: provider, settings: self.settings, store: self)                runtime.providerDidRefresh(context: context, provider: provider)            }            if provider == .codex {                self.recordCodexHistoricalSampleIfNeeded(snapshot: backfilled)            }        case let .failure(error):            if provider == .codex,               let codexExpectedGuard = context.codexExpectedGuard,               !self.shouldApplyCodexScopedFailure(expectedGuard: codexExpectedGuard)            {                return            }            guard self.isCurrentProviderRefreshGeneration(provider, generation: context.generation) else { return }            self.recordStartupConnectivityRetryableFailure(error)            if context.claudeCredentialsChanged {                await self.clearClaudeCredentialDerivedStateForCredentialSwap()            }            if context.shouldConsumeClaudeKeychainFingerprint {                _ = await Self.consumeClaudeKeychainFingerprintChangeWithoutPrompt()            }            await self.handleProviderFetchFailure(                provider: provider,                error: error,                generation: context.generation)        }    }     private func clearDisabledProviderRefreshState(_ provider: UsageProvider) async {        self.refreshingProviders.remove(provider)        await MainActor.run {            self.snapshots.removeValue(forKey: provider)            self.lastKnownResetSnapshots.removeValue(forKey: provider)            self.errors[provider] = nil            self.lastSourceLabels.removeValue(forKey: provider)            self.lastFetchAttempts.removeValue(forKey: provider)            self.accountSnapshots.removeValue(forKey: provider)            if provider == .codex {                self.codexAccountSnapshots = []            }            if provider == .kilo {                self.kiloScopeSnapshots = []            }            self.tokenSnapshots.removeValue(forKey: provider)            self.tokenErrors[provider] = nil            self.failureGates[provider]?.reset()            self.tokenFailureGates[provider]?.reset()            self.statuses.removeValue(forKey: provider)            self.lastKnownSessionRemaining.removeValue(forKey: provider)            self.lastKnownSessionWindowSource.removeValue(forKey: provider)            self.quotaWarningState = self.quotaWarningState.filter { $0.key.provider != provider }            self.lastTokenFetchAt.removeValue(forKey: provider)        }    }     private struct ClaudeRefreshAuthState {        let fingerprintToken: String        let credentialsFileChanged: Bool        let keychainFingerprintChanged: Bool    }     private nonisolated static func claudeCredentialsChanged(        beforeFetch: ClaudeRefreshAuthState?,        changedDuringFetch: Bool) -> Bool    {        beforeFetch?.credentialsFileChanged == true ||            beforeFetch?.keychainFingerprintChanged == true ||            changedDuringFetch    }     private nonisolated static func shouldConsumeClaudeKeychainFingerprintChange(        beforeFetch: ClaudeRefreshAuthState?,        changedDuringFetch: Bool) -> Bool    {        beforeFetch?.keychainFingerprintChanged == true || changedDuringFetch    }     private nonisolated static func claudeAuthChangedDuringFetch(        provider: UsageProvider,        beforeFetch: ClaudeRefreshAuthState?,        afterFetchFingerprintToken: String?) -> Bool    {        provider == .claude && afterFetchFingerprintToken != beforeFetch?.fingerprintToken    }     private nonisolated static func captureClaudeRefreshAuthState(        invalidateCredentialsFile: Bool) async -> ClaudeRefreshAuthState    {        await withTaskGroup(of: ClaudeRefreshAuthState.self, returning: ClaudeRefreshAuthState.self) { group in            group.addTask {                let fingerprintToken = ClaudeOAuthCredentialsStore.authFingerprintToken()                let credentialsFileChanged = invalidateCredentialsFile                    ? ClaudeOAuthCredentialsStore.invalidateCacheIfCredentialsFileChanged()                    : false                let keychainFingerprintChanged = ClaudeOAuthCredentialsStore                    .claudeKeychainFingerprintChangedWithoutConsuming()                return ClaudeRefreshAuthState(                    fingerprintToken: fingerprintToken,                    credentialsFileChanged: credentialsFileChanged,                    keychainFingerprintChanged: keychainFingerprintChanged)            }            return await group.next()!        }    }     private nonisolated static func captureClaudeAuthFingerprintToken() async -> String {        await withTaskGroup(of: String.self, returning: String.self) { group in            group.addTask {                ClaudeOAuthCredentialsStore.authFingerprintToken()            }            return await group.next()!        }    }     private nonisolated static func invalidateClaudeCredentialsFileCacheIfChanged() async -> Bool {        await withTaskGroup(of: Bool.self, returning: Bool.self) { group in            group.addTask {                ClaudeOAuthCredentialsStore.invalidateCacheIfCredentialsFileChanged()            }            return await group.next()!        }    }     private nonisolated static func invalidateClaudeCredentialsFileCacheIfNeeded(changedDuringFetch: Bool) async {        guard changedDuringFetch else { return }        _ = await self.invalidateClaudeCredentialsFileCacheIfChanged()    }     private nonisolated static func consumeClaudeKeychainFingerprintChangeWithoutPrompt() async -> Bool {        await withTaskGroup(of: Bool.self, returning: Bool.self) { group in            group.addTask {                ClaudeOAuthCredentialsStore.consumeClaudeKeychainFingerprintChangeWithoutPrompt()            }            return await group.next()!        }    }     private func clearClaudeCredentialDerivedStateForCredentialSwap() async {        await MainActor.run {            self.clearClaudeCredentialDerivedStateForCredentialSwapNow()        }    }     private func clearClaudeCredentialDerivedStateForCredentialSwapNow() {        self.snapshots.removeValue(forKey: .claude)        self.lastKnownResetSnapshots.removeValue(forKey: .claude)        self.errors[.claude] = nil        self.lastSourceLabels.removeValue(forKey: .claude)        self.accountSnapshots.removeValue(forKey: .claude)        self.tokenSnapshots.removeValue(forKey: .claude)        self.tokenErrors[.claude] = nil        self.failureGates[.claude]?.reset()        self.tokenFailureGates[.claude]?.reset()        self.lastKnownSessionRemaining.removeValue(forKey: .claude)        self.lastKnownSessionWindowSource.removeValue(forKey: .claude)        self.quotaWarningState = self.quotaWarningState.filter { $0.key.provider != .claude }        self.lastTokenFetchAt.removeValue(forKey: .claude)    }     private func handleProviderFetchFailure(        provider: UsageProvider,        error: Error,        generation: UInt64) async    {        let shouldNotifyPermissionPrompt = Self.isPermissionPromptWaiting(error)        await MainActor.run {            guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }            let hadPriorData = self.snapshots[provider] != nil            let preservesPriorData = Self.shouldPreservePriorSnapshot(                after: error,                hadPriorData: hadPriorData)            let shouldSurface =                self.failureGates[provider]?                    .shouldSurfaceError(onFailureWithPriorData: hadPriorData) ?? true            let preservesClaudeWebSessionFailure =                provider == .claude &&                hadPriorData &&                Self.isClaudeWebSessionRefreshFailure(error)            if preservesClaudeWebSessionFailure,               !shouldSurface            {                self.errors[provider] = nil                return            }            if provider == .claude,               preservesPriorData,               Self.isClaudeUsageProbeTimeout(error)            {                self.errors[provider] = nil                return            }            if preservesPriorData, !shouldSurface {                self.errors[provider] = nil

Lines 1–500 of 619.