82 lines · 3.5 KB · Swift
RawDownload
import CodexBarCoreimport Foundation extension UsageStore {    /// Force refresh Augment session (called from UI button)    func forceRefreshAugmentSession() async {        await self.performRuntimeAction(.forceSessionRefresh, for: .augment)    }     func refreshProvider(_ provider: UsageProvider, allowDisabled: Bool = false) async {        guard let spec = self.providerSpecs[provider] else { return }         if !spec.isEnabled(), !allowDisabled {            self.refreshingProviders.remove(provider)            await MainActor.run {                self.snapshots.removeValue(forKey: provider)                self.errors[provider] = nil                self.lastSourceLabels.removeValue(forKey: provider)                self.lastFetchAttempts.removeValue(forKey: provider)                self.accountSnapshots.removeValue(forKey: provider)                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.lastTokenFetchAt.removeValue(forKey: provider)            }            return        }         self.refreshingProviders.insert(provider)        defer { self.refreshingProviders.remove(provider) }         let tokenAccounts = self.tokenAccounts(for: provider)        if self.shouldFetchAllTokenAccounts(provider: provider, accounts: tokenAccounts) {            await self.refreshTokenAccounts(provider: provider, accounts: tokenAccounts)            return        } else {            _ = await MainActor.run {                self.accountSnapshots.removeValue(forKey: provider)            }        }         let outcome = await spec.fetch()        await MainActor.run {            self.lastFetchAttempts[provider] = outcome.attempts        }         switch outcome.result {        case let .success(result):            let scoped = result.usage.scoped(to: provider)            await MainActor.run {                self.handleSessionQuotaTransition(provider: provider, snapshot: scoped)                self.snapshots[provider] = scoped                self.lastSourceLabels[provider] = result.sourceLabel                self.errors[provider] = nil                self.failureGates[provider]?.recordSuccess()            }            if let runtime = self.providerRuntimes[provider] {                let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self)                runtime.providerDidRefresh(context: context, provider: provider)            }        case let .failure(error):            await MainActor.run {                let hadPriorData = self.snapshots[provider] != nil                let shouldSurface = self.failureGates[provider]?                    .shouldSurfaceError(onFailureWithPriorData: hadPriorData) ?? true                if shouldSurface {                    self.errors[provider] = error.localizedDescription                    self.snapshots.removeValue(forKey: provider)                } else {                    self.errors[provider] = nil                }            }            if let runtime = self.providerRuntimes[provider] {                let context = ProviderRuntimeContext(provider: provider, settings: self.settings, store: self)                runtime.providerDidFail(context: context, provider: provider, error: error)            }        }    }}