300 lines · 12.2 KB · Swift
RawDownload
import CodexBarCoreimport Foundation @MainActorextension UsageStore {    private struct StorageRefreshRequest {        let providers: [UsageProvider]        let candidatePathsByProvider: [UsageProvider: [String]]        let signature: String    }     private static let automaticStorageRefreshInterval: TimeInterval = 5 * 60     var isStorageRefreshInFlight: Bool {        self.storageRefreshTask != nil    }     func storageFootprint(for provider: UsageProvider) -> ProviderStorageFootprint? {        guard self.settings.providerStorageFootprintsEnabled else { return nil }        return self.providerStorageFootprints[provider]    }     func storageFootprintText(for provider: UsageProvider) -> String? {        guard let footprint = self.storageFootprint(for: provider) else { return nil }        if footprint.hasLocalData {            return UsageFormatter.byteCountString(footprint.totalBytes)        }        return "No local data found"    }     func refreshStorageFootprintsForOverview() {        self.scheduleStorageFootprintRefresh(for: self.enabledProvidersForDisplay())    }     func refreshStorageFootprintsForOverviewNow() async {        await self.refreshStorageFootprintsNow(for: self.enabledProvidersForDisplay())    }     func scheduleStorageFootprintRefreshForOverview(force: Bool = false) {        self.scheduleStorageFootprintRefresh(for: self.enabledProvidersForDisplay(), force: force)    }     func refreshStorageFootprintsNow(for providers: [UsageProvider]) async {        guard self.settings.providerStorageFootprintsEnabled else {            self.clearStorageFootprints()            return        }        let environment = self.environmentBase        let managedAccountsOverride = self.managedCodexAccountsForStorageOverride        let request = await Task.detached(priority: .utility) {            let managedAccounts = Self.loadManagedCodexAccountsForStorage(override: managedAccountsOverride)            return Self.makeStorageRefreshRequest(                for: providers,                environment: environment,                managedAccounts: managedAccounts)        }.value        guard let request else {            self.clearStorageFootprints()            return        }         self.storageRefreshTask?.cancel()        self.storageRefreshGeneration &+= 1        let generation = self.storageRefreshGeneration        self.storageRefreshInFlightSignature = request.signature         let footprints = await Task.detached(priority: .utility) {            Self.scanStorageFootprints(candidatePathsByProvider: request.candidatePathsByProvider)        }.value         guard generation == self.storageRefreshGeneration else { return }        self.applyStorageFootprints(            footprints,            providers: request.providers,            signature: request.signature,            updatedAt: Date())        self.storageRefreshTask = nil        self.storageRefreshInFlightSignature = nil        self.storageRefreshInFlightRequestKey = nil    }     func scheduleStorageFootprintRefresh(for providers: [UsageProvider], force: Bool = false) {        guard self.settings.providerStorageFootprintsEnabled else {            self.clearStorageFootprints()            return        }        let managedAccountsOverride = self.managedCodexAccountsForStorageOverride        let requestKey = Self.storageRefreshRequestKey(            for: providers,            managedAccountsOverride: managedAccountsOverride)        guard !requestKey.isEmpty else {            self.clearStorageFootprints()            return        }         let now = Date()        if self.storageRefreshTask != nil,           self.storageRefreshInFlightRequestKey == nil || self.storageRefreshInFlightRequestKey == requestKey        {            return        }        if !force {            if self.lastStorageRefreshRequestKey == requestKey,               let lastStorageRefreshAt,               now.timeIntervalSince(lastStorageRefreshAt) < Self.automaticStorageRefreshInterval            {                return            }        }         self.storageRefreshTask?.cancel()        self.storageRefreshGeneration &+= 1        let generation = self.storageRefreshGeneration        self.storageRefreshInFlightSignature = nil        self.storageRefreshInFlightRequestKey = requestKey        let environment = self.environmentBase         self.storageRefreshTask = Task.detached(priority: .utility) { [weak self] in            let managedAccounts = Self.loadManagedCodexAccountsForStorage(override: managedAccountsOverride)            guard let request = Self.makeStorageRefreshRequest(                for: providers,                environment: environment,                managedAccounts: managedAccounts)            else {                await MainActor.run { [weak self] in                    guard let self,                          !Task.isCancelled,                          generation == self.storageRefreshGeneration                    else { return }                    self.providerStorageFootprints.removeAll()                    self.storageRefreshTask = nil                    self.storageRefreshInFlightSignature = nil                    self.storageRefreshInFlightRequestKey = nil                    self.lastStorageRefreshSignature = nil                    self.lastStorageRefreshRequestKey = requestKey                    self.lastStorageRefreshAt = Date()                }                return            }            let footprints = Self.scanStorageFootprints(candidatePathsByProvider: request.candidatePathsByProvider)             await MainActor.run { [weak self] in                guard let self,                      !Task.isCancelled,                      generation == self.storageRefreshGeneration                else { return }                 self.applyStorageFootprints(                    footprints,                    providers: request.providers,                    signature: request.signature,                    requestKey: requestKey,                    updatedAt: Date())                self.storageRefreshTask = nil                self.storageRefreshInFlightSignature = nil                self.storageRefreshInFlightRequestKey = nil            }        }    }     private func clearStorageFootprints() {        self.storageRefreshTask?.cancel()        self.storageRefreshTask = nil        self.storageRefreshInFlightSignature = nil        self.storageRefreshInFlightRequestKey = nil        self.lastStorageRefreshSignature = nil        self.lastStorageRefreshRequestKey = nil        self.lastStorageRefreshAt = nil        self.providerStorageFootprints.removeAll()    }     private func applyStorageFootprints(        _ footprints: [UsageProvider: ProviderStorageFootprint],        providers: [UsageProvider],        signature: String,        requestKey: String? = nil,        updatedAt: Date)    {        let providerSet = Set(providers)        var updated = self.providerStorageFootprints.filter { !providerSet.contains($0.key) }        for provider in providers {            // Reuse the existing footprint when only its scan timestamp would change, so the equality            // guard below treats an unchanged scan as a no-op.            if let incoming = footprints[provider],               let existing = self.providerStorageFootprints[provider],               existing.hasSameContents(as: incoming)            {                updated[provider] = existing            } else {                updated[provider] = footprints[provider]            }        }        // Only republish the observable footprints when a value actually changed. Storage scans run        // on every menu open and roughly every 5 minutes; an unconditional re-assignment wakes        // `menuObservationToken` -> `invalidateMenus` churn (clearing menu caches) even when the        // scanned bytes are identical.        if updated != self.providerStorageFootprints {            self.providerStorageFootprints = updated        }        self.lastStorageRefreshSignature = signature        self.lastStorageRefreshRequestKey = requestKey ?? signature        self.lastStorageRefreshAt = updatedAt    }     private nonisolated static func makeStorageRefreshRequest(        for providers: [UsageProvider],        environment: [String: String],        managedAccounts: [ManagedCodexAccount])        -> StorageRefreshRequest?    {        let uniqueProviders = Array(Set(providers)).sorted { $0.rawValue < $1.rawValue }        guard !uniqueProviders.isEmpty else { return nil }         var candidatePathsByProvider: [UsageProvider: [String]] = [:]         for provider in uniqueProviders {            let candidatePaths = ProviderStoragePathCatalog.candidatePaths(                for: provider,                environment: environment,                managedCodexAccounts: managedAccounts)            guard !candidatePaths.isEmpty else { continue }            candidatePathsByProvider[provider] = candidatePaths        }         let providersWithPaths = uniqueProviders.filter { candidatePathsByProvider[$0] != nil }        guard !providersWithPaths.isEmpty else { return nil }         let signature = providersWithPaths            .map { provider in                let paths = candidatePathsByProvider[provider]?.joined(separator: "\u{1f}") ?? ""                return "\(provider.rawValue)=\(paths)"            }            .joined(separator: "\u{1e}")        return StorageRefreshRequest(            providers: providersWithPaths,            candidatePathsByProvider: candidatePathsByProvider,            signature: signature)    }     private nonisolated static func storageRefreshRequestKey(        for providers: [UsageProvider],        managedAccountsOverride: [ManagedCodexAccount]?)        -> String    {        let uniqueProviders = Array(Set(providers))            .sorted { $0.rawValue < $1.rawValue }        let providerKey = uniqueProviders.map(\.rawValue).joined(separator: ",")        guard uniqueProviders.contains(.codex) else { return providerKey }         let managedAccountsRevision: String        if let managedAccountsOverride {            managedAccountsRevision = Array(Set(managedAccountsOverride.map(\.managedHomePath)))                .sorted()                .joined(separator: "\u{1f}")        } else {            let fileURL = FileManagedCodexAccountStore.defaultURL()            let attributes = try? FileManager.default.attributesOfItem(atPath: fileURL.path)            let modificationDate = (attributes?[.modificationDate] as? Date)?                .timeIntervalSinceReferenceDate.bitPattern ?? 0            let fileNumber = (attributes?[.systemFileNumber] as? NSNumber)?.uint64Value ?? 0            let fileSize = (attributes?[.size] as? NSNumber)?.uint64Value ?? 0            managedAccountsRevision = "\(fileNumber):\(modificationDate):\(fileSize)"        }        return "\(providerKey)\u{1e}\(managedAccountsRevision)"    }     private nonisolated static func loadManagedCodexAccountsForStorage(        override: [ManagedCodexAccount]?)        -> [ManagedCodexAccount]    {        if let override {            return override        }        return (try? FileManagedCodexAccountStore().loadAccounts().accounts) ?? []    }     private nonisolated static func scanStorageFootprints(        candidatePathsByProvider: [UsageProvider: [String]])        -> [UsageProvider: ProviderStorageFootprint]    {        let scanner = ProviderStorageScanner()        var footprints: [UsageProvider: ProviderStorageFootprint] = [:]        var pathCache: [String: ProviderStorageFootprint] = [:]         for provider in candidatePathsByProvider.keys.sorted(by: { $0.rawValue < $1.rawValue }) {            if Task.isCancelled { return footprints }            guard let candidatePaths = candidatePathsByProvider[provider] else { continue }            let pathKey = candidatePaths.joined(separator: "\u{1f}")            if let cached = pathCache[pathKey] {                footprints[provider] = cached.replacingProvider(provider)                continue            }            let footprint = scanner.scan(provider: provider, candidatePaths: candidatePaths)            pathCache[pathKey] = footprint            footprints[provider] = footprint        }         return footprints    }}