import CodexBarCoreimport Foundation struct TokenAccountUsageSnapshot: Identifiable {    let id: UUID    let account: ProviderTokenAccount    let snapshot: UsageSnapshot?    let error: String?    let sourceLabel: String?     init(account: ProviderTokenAccount, snapshot: UsageSnapshot?, error: String?, sourceLabel: String?) {        self.id = account.id        self.account = account        self.snapshot = snapshot        self.error = error        self.sourceLabel = sourceLabel    }} struct CodexAccountUsageSnapshot: Identifiable {    let id: String    let account: CodexVisibleAccount    let snapshot: UsageSnapshot?    let error: String?    let sourceLabel: String?     init(account: CodexVisibleAccount, snapshot: UsageSnapshot?, error: String?, sourceLabel: String?) {        self.id = account.id        self.account = account        self.snapshot = snapshot        self.error = error        self.sourceLabel = sourceLabel    }} extension UsageStore {    func activateCachedTokenAccountSnapshot(provider: UsageProvider, accountID: UUID) {        guard let cached = self.accountSnapshots[provider]?.first(where: { $0.account.id == accountID }) else {            self.snapshots.removeValue(forKey: provider)            self.errors.removeValue(forKey: provider)            self.lastSourceLabels.removeValue(forKey: provider)            self.lastKnownResetSnapshots.removeValue(forKey: provider)            return        }         if let snapshot = cached.snapshot {            self.snapshots[provider] = snapshot            self.lastKnownResetSnapshots[provider] = snapshot        } else {            self.snapshots.removeValue(forKey: provider)            self.lastKnownResetSnapshots.removeValue(forKey: provider)        }        self.errors[provider] = cached.error        if let sourceLabel = cached.sourceLabel {            self.lastSourceLabels[provider] = sourceLabel        } else {            self.lastSourceLabels.removeValue(forKey: provider)        }    }} private struct TokenAccountFetchResult {    let index: Int    let account: ProviderTokenAccount    let outcome: ProviderFetchOutcome} private struct CodexAccountFetchResult {    let index: Int    let account: CodexVisibleAccount    let outcome: ProviderFetchOutcome} private struct CodexManagedVisibleAccountRuntimeState {    let authFingerprint: String?    let workspaceAccountID: String?} extension UsageStore {    static let tokenAccountMenuSnapshotLimit = 6    private static let codexSessionWindowMinutes = 5 * 60    private static let codexWeeklyWindowMinutes = 7 * 24 * 60     func freshCodexVisibleAccountsForSnapshotHydration() -> [CodexVisibleAccount] {        self.freshCodexVisibleAccountProjectionForAccountRefresh().visibleAccounts    }     func tokenAccounts(for provider: UsageProvider) -> [ProviderTokenAccount] {        guard TokenAccountSupportCatalog.support(for: provider) != nil else { return [] }        return self.settings.tokenAccounts(for: provider)    }     func shouldFetchAllTokenAccounts(provider: UsageProvider, accounts: [ProviderTokenAccount]) -> Bool {        guard TokenAccountSupportCatalog.support(for: provider) != nil else { return false }        return self.settings.multiAccountMenuLayout == .stacked && accounts.count > 1    }     func shouldFetchAllCodexVisibleAccounts() -> Bool {        let projection = self.freshCodexVisibleAccountProjectionForAccountRefresh()        return self.settings.multiAccountMenuLayout == .stacked &&            projection.visibleAccounts.count > 1    }     func refreshCodexVisibleAccountsForMenu(generation: UInt64? = nil) async {        let projection = self.freshCodexVisibleAccountProjectionForAccountRefresh()        let accounts = self.limitedCodexVisibleAccounts(            projection.visibleAccounts,            snapshots: self.codexAccountSnapshots,            activeVisibleAccountID: projection.activeVisibleAccountID)        guard accounts.count > 1 else {            self.codexAccountSnapshots = []            return        }        let managedAccountIDsWithReadableAuthAtStart = self.codexManagedAccountIDsWithReadableAuth()         let originalVisibleAccountID = projection.activeVisibleAccountID        let originalSelectionSource = originalVisibleAccountID.flatMap {            projection.source(forVisibleAccountID: $0)        }        let originalVisibleAccount = originalVisibleAccountID.flatMap { id in            accounts.first { $0.id == id }        }        let priorByAccountID = Dictionary(uniqueKeysWithValues: self.codexAccountSnapshots.map { ($0.id, $0) })        var snapshots: [CodexAccountUsageSnapshot] = []        var selectedOutcome: ProviderFetchOutcome?        var selectedAccount: CodexVisibleAccount?        var selectedSnapshot: UsageSnapshot?        var selectedSourceLabel: String?        var sawAnyNonCancellationOutcome = false         let results = await self.fetchCodexVisibleAccountOutcomes(accounts)        for result in results {            let account = result.account            let outcome = result.outcome            let isCancellation = Self.outcomeIsCancellation(outcome)            if !isCancellation {                sawAnyNonCancellationOutcome = true            }            let resolved = self.resolveCodexAccountOutcome(                outcome,                account: account,                priorSnapshot: priorByAccountID[account.id],                resetBackfillSnapshots: self.codexResetBackfillSnapshots(                    for: account,                    priorSnapshot: priorByAccountID[account.id],                    activeVisibleAccountID: originalVisibleAccountID))            if let snapshot = resolved.snapshot {                snapshots.append(snapshot)            }            if account.id == originalVisibleAccountID {                selectedOutcome = outcome                selectedAccount = account                selectedSnapshot = resolved.usage                selectedSourceLabel = resolved.sourceLabel            }        }         let currentProjection = self.freshCodexVisibleAccountProjectionForAccountRefresh(            requireLiveManagedAuthFor: managedAccountIDsWithReadableAuthAtStart)        guard self.isCurrentProviderRefreshGeneration(.codex, generation: generation) else { return }        let currentSnapshots = snapshots.compactMap { snapshot -> CodexAccountUsageSnapshot? in            guard let currentAccount = Self.currentCodexVisibleAccount(                matching: snapshot.account,                projection: currentProjection,                allowProviderAccountAuthFingerprintMismatch: snapshot.error == nil)            else {                return nil            }            guard currentAccount != snapshot.account else { return snapshot }            return CodexAccountUsageSnapshot(                account: currentAccount,                snapshot: Self.codexVisibleAccountSnapshotRelabeledForCurrentProjection(                    snapshot.snapshot,                    account: currentAccount),                error: snapshot.error,                sourceLabel: snapshot.sourceLabel)        }        let shouldPreservePriorState = !sawAnyNonCancellationOutcome &&            currentSnapshots.allSatisfy { $0.snapshot == nil }        if !shouldPreservePriorState {            self.codexAccountSnapshots = currentSnapshots            self.codexAccountUsageSnapshotStore?.store(currentSnapshots)        }         let selectionStillMatches = self.codexVisibleSelectionStillMatches(            originalVisibleAccountID: originalVisibleAccountID,            originalSelectionSource: originalSelectionSource,            originalAccount: originalVisibleAccount,            currentProjection: currentProjection)        guard let selectedOutcome, let selectedAccount else { return }        guard selectionStillMatches else {            _ = self.prepareCodexAccountScopedRefreshIfNeeded()            return        }         let allowSelectedAuthFingerprintMismatch = switch selectedOutcome.result {        case .success:            true        case .failure:            false        }        let currentSelectedAccount = Self.currentCodexVisibleAccount(            matching: selectedAccount,            projection: currentProjection,            allowProviderAccountAuthFingerprintMismatch: allowSelectedAuthFingerprintMismatch)        if let currentSelectedAccount {            let currentSelectedSnapshot = Self.codexVisibleAccountSnapshotRelabeledForCurrentProjection(                selectedSnapshot,                account: currentSelectedAccount)            if self.shouldApplySelectedCodexVisibleAccountOutcome(                selectedOutcome,                snapshot: currentSelectedSnapshot)            {                await self.applySelectedCodexVisibleAccountOutcome(                    selectedOutcome,                    account: currentSelectedAccount,                    snapshot: currentSelectedSnapshot,                    sourceLabel: selectedSourceLabel,                    generation: generation)            }        } else {            _ = self.prepareCodexAccountScopedRefreshIfNeeded()        }    }     func codexVisibleSelectionStillMatches(        originalVisibleAccountID: String?,        originalSelectionSource: CodexActiveSource?,        originalAccount: CodexVisibleAccount? = nil,        currentProjection: CodexVisibleAccountProjection? = nil) -> Bool    {        let currentProjection = currentProjection ?? self.settings.codexVisibleAccountProjection        let currentActiveAccount = currentProjection.activeVisibleAccountID.flatMap { id in            currentProjection.visibleAccounts.first { $0.id == id }        }        let currentSelectionSource = currentActiveAccount?.selectionSource        if currentProjection.activeVisibleAccountID == originalVisibleAccountID,           currentSelectionSource == originalSelectionSource        {            return true        }        guard let originalAccount, let currentActiveAccount, currentSelectionSource == originalSelectionSource else {            return false        }        return Self.codexVisibleAccountMatchesCurrentProjection(originalAccount, account: currentActiveAccount)    }     private func freshCodexVisibleAccountProjectionForAccountRefresh(        requireLiveManagedAuthFor accountIDs: Set<UUID> = []) -> CodexVisibleAccountProjection    {        // Auth files can change while account fetches are in flight, so account refreshes bypass the        // short-lived reconciliation cache used for normal menu rendering and stale-result guards.        self.settings.invalidateCodexAccountReconciliationSnapshotCache()        let snapshot = self.settings.codexAccountReconciliationSnapshot        return Self.codexVisibleAccountProjectionWithFreshManagedAuthFingerprints(            CodexVisibleAccountProjection.make(from: snapshot),            snapshot: snapshot,            requireLiveManagedAuthFor: accountIDs)    }     private func codexManagedAccountIDsWithReadableAuth() -> Set<UUID> {        Set(self.settings.codexAccountReconciliationSnapshot.storedAccounts.compactMap { account in            CodexAuthFingerprint.fingerprint(homePath: account.managedHomePath) == nil ? nil : account.id        })    }     private nonisolated static func codexVisibleAccountProjectionWithFreshManagedAuthFingerprints(        _ projection: CodexVisibleAccountProjection,        snapshot: CodexAccountReconciliationSnapshot,        requireLiveManagedAuthFor accountIDs: Set<UUID> = []) -> CodexVisibleAccountProjection    {        let managedRuntimeStates = Dictionary(            uniqueKeysWithValues: snapshot.storedAccounts.map { account in                let workspaceAccountID: String? = switch snapshot.runtimeIdentity(for: account) {                case let .providerAccount(id):                    id                case .emailOnly, .unresolved:                    nil                }                let authFingerprint = CodexAuthFingerprint.fingerprint(homePath: account.managedHomePath)                let requiresLiveAuth = accountIDs.contains(account.id)                return (account.id, CodexManagedVisibleAccountRuntimeState(                    authFingerprint: authFingerprint ?? (requiresLiveAuth ? nil : account.authFingerprint),                    workspaceAccountID: authFingerprint == nil && requiresLiveAuth                        ? nil                        : (workspaceAccountID ?? account.workspaceAccountID)))            })        let visibleAccounts = projection.visibleAccounts.map { account in            guard case let .managedAccount(id) = account.selectionSource else { return account }            let accountWorkspaceAccountID = account.workspaceAccountID                .map(CodexOpenAIWorkspaceIdentity.normalizeWorkspaceAccountID)            let runtimeWorkspaceAccountID = managedRuntimeStates[id]?.workspaceAccountID                .map(CodexOpenAIWorkspaceIdentity.normalizeWorkspaceAccountID)            guard let runtimeState = managedRuntimeStates[id],                  runtimeState.authFingerprint != account.authFingerprint ||                  runtimeWorkspaceAccountID != accountWorkspaceAccountID            else {                return account            }            return CodexVisibleAccount(                id: account.id,                email: account.email,                workspaceLabel: account.workspaceLabel,                workspaceAccountID: runtimeState.workspaceAccountID,                authFingerprint: runtimeState.authFingerprint,                storedAccountID: account.storedAccountID,                selectionSource: account.selectionSource,                isActive: account.isActive,                isLive: account.isLive,                canReauthenticate: account.canReauthenticate,                canRemove: account.canRemove)        }        return CodexVisibleAccountProjection(            visibleAccounts: visibleAccounts,            activeVisibleAccountID: projection.activeVisibleAccountID,            liveVisibleAccountID: projection.liveVisibleAccountID,            hasUnreadableAddedAccountStore: projection.hasUnreadableAddedAccountStore)    }     private static func currentCodexVisibleAccount(        matching account: CodexVisibleAccount,        projection: CodexVisibleAccountProjection,        allowProviderAccountAuthFingerprintMismatch: Bool = true) -> CodexVisibleAccount?    {        if let currentAccount = projection.visibleAccounts.first(where: { $0.id == account.id }),           self.codexVisibleAccountMatchesCurrentProjection(               account,               account: currentAccount,               allowProviderAccountAuthFingerprintMismatch: allowProviderAccountAuthFingerprintMismatch)        {            return currentAccount        }        return projection.visibleAccounts.first {            self.codexVisibleAccountMatchesCurrentProjection(                account,                account: $0,                allowProviderAccountAuthFingerprintMismatch: allowProviderAccountAuthFingerprintMismatch)        }    }     private static func codexVisibleAccountSnapshotRelabeledForCurrentProjection(        _ snapshot: UsageSnapshot?,        account: CodexVisibleAccount) -> UsageSnapshot?    {        guard let snapshot else { return nil }        let existing = snapshot.identity(for: .codex)        return snapshot.withIdentity(ProviderIdentitySnapshot(            providerID: .codex,            accountEmail: account.email,            accountOrganization: existing?.accountOrganization,            loginMethod: existing?.loginMethod ?? account.workspaceLabel))    }     private static func codexVisibleAccountMatchesCurrentProjection(        _ prior: CodexVisibleAccount,        account: CodexVisibleAccount,        allowProviderAccountAuthFingerprintMismatch: Bool = true) -> Bool    {        guard prior.selectionSource == account.selectionSource else { return false }         let priorEmail = CodexIdentityResolver.normalizeEmail(prior.email)        let accountEmail = CodexIdentityResolver.normalizeEmail(account.email)         let priorWorkspaceID = self.normalizedCodexVisibleAccountText(prior.workspaceAccountID)            .map(CodexOpenAIWorkspaceIdentity.normalizeWorkspaceAccountID)        let accountWorkspaceID = self.normalizedCodexVisibleAccountText(account.workspaceAccountID)            .map(CodexOpenAIWorkspaceIdentity.normalizeWorkspaceAccountID)        if priorWorkspaceID != nil || accountWorkspaceID != nil {            guard priorWorkspaceID == accountWorkspaceID else { return false }            if !allowProviderAccountAuthFingerprintMismatch {                guard self.codexVisibleAccountAuthFingerprintMatches(prior, account: account) else { return false }            }            switch account.selectionSource {            case .managedAccount:                if !self.codexVisibleAccountAuthFingerprintMatches(prior, account: account) {                    return priorEmail != nil && priorEmail == accountEmail                }                return true            case .liveSystem:                return priorEmail != nil && priorEmail == accountEmail            case .profileHome:                if !allowProviderAccountAuthFingerprintMismatch {                    guard self.codexVisibleAccountAuthFingerprintMatches(prior, account: account) else { return false }                }                return priorEmail != nil && priorEmail == accountEmail            }        }         let priorAuthFingerprint = CodexAuthFingerprint.normalize(prior.authFingerprint)        let accountAuthFingerprint = CodexAuthFingerprint.normalize(account.authFingerprint)        if priorAuthFingerprint != nil || accountAuthFingerprint != nil {            guard priorAuthFingerprint == accountAuthFingerprint else { return false }        }         return priorEmail != nil && priorEmail == accountEmail    }     private static func codexVisibleAccountAuthFingerprintMatches(        _ prior: CodexVisibleAccount,        account: CodexVisibleAccount) -> Bool    {        let priorAuthFingerprint = CodexAuthFingerprint.normalize(prior.authFingerprint)        let accountAuthFingerprint = CodexAuthFingerprint.normalize(account.authFingerprint)        if priorAuthFingerprint != nil || accountAuthFingerprint != nil {            return priorAuthFingerprint == accountAuthFingerprint        }        return true    }     func shouldApplySelectedCodexVisibleAccountOutcome(        _ outcome: ProviderFetchOutcome,        snapshot: UsageSnapshot?) -> Bool    {        switch outcome.result {        case .success:            snapshot != nil        case .failure:            true        }    }     func refreshTokenAccounts(        provider: UsageProvider,        accounts: [ProviderTokenAccount],        generation: UInt64? = nil) async    {        let selectedAccount = self.settings.selectedTokenAccount(for: provider)        let limitedAccounts = self.limitedTokenAccounts(accounts, selected: selectedAccount)        let effectiveSelected = selectedAccount ?? limitedAccounts.first         // Capture the prior per-account snapshot state so we can preserve last-good        // data when an in-flight refresh is cancelled (e.g. menu tab switches). Without        // this, cancellation produces empty/error snapshots and the menu briefly shows        // misleading cards for accounts that previously had valid data.        let priorSnapshots = await MainActor.run { self.accountSnapshots[provider] ?? [] }        let priorByAccountID = Dictionary(uniqueKeysWithValues: priorSnapshots.map { ($0.account.id, $0) })         var snapshots: [TokenAccountUsageSnapshot] = []        var historySamples: [(account: ProviderTokenAccount, snapshot: UsageSnapshot)] = []        var selectedOutcome: ProviderFetchOutcome?        var selectedSnapshot: UsageSnapshot?        var sawAnyNonCancellationOutcome = false         let results = await self.fetchTokenAccountOutcomes(provider: provider, accounts: limitedAccounts)        guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }        for result in results {            let account = result.account            let outcome = result.outcome            let isCancellation = Self.outcomeIsCancellation(outcome)            if !isCancellation {                sawAnyNonCancellationOutcome = true            }            let resolved = self.resolveAccountOutcome(                outcome,                provider: provider,                account: account,                priorSnapshot: priorByAccountID[account.id])            if let snapshot = resolved.snapshot {                snapshots.append(snapshot)            }            if let usage = resolved.usage {                historySamples.append((account: account, snapshot: usage))            }            if account.id == effectiveSelected?.id {                selectedOutcome = outcome                selectedSnapshot = resolved.usage            }        }         // If every fetch was cancelled (e.g. the user closed/reopened the menu mid-flight)        // and we have no usable snapshots, leave the prior per-account state alone.        // Wiping it would produce a menu of useless "cancelled" placeholders.        let shouldPreservePriorState = !sawAnyNonCancellationOutcome &&            snapshots.allSatisfy { $0.snapshot == nil }        if !shouldPreservePriorState {            await MainActor.run {                self.accountSnapshots[provider] = snapshots            }        }         if let selectedOutcome {            await self.applySelectedOutcome(                selectedOutcome,                provider: provider,                account: effectiveSelected,                fallbackSnapshot: selectedSnapshot,                generation: generation)        }         guard self.isCurrentProviderRefreshGeneration(provider, generation: generation) else { return }        await self.recordFetchedTokenAccountPlanUtilizationHistory(            provider: provider,            samples: historySamples,            selectedAccount: effectiveSelected)    }     private static func outcomeIsCancellation(_ outcome: ProviderFetchOutcome) -> Bool {        if case let .failure(error) = outcome.result, error is CancellationError {            return true        }

Lines 1–500 of 1308. Load more