import CodexBarCoreimport Foundation extension UsageStore {    func restartAugmentKeepaliveIfNeeded() {        #if os(macOS)        let shouldRun = self.isEnabled(.augment)        let isRunning = self.augmentKeepalive != nil         if shouldRun, !isRunning {            self.startAugmentKeepalive()        } else if !shouldRun, isRunning {            Task { @MainActor in                self.augmentKeepalive?.stop()                self.augmentKeepalive = nil                self.augmentLogger.info("Augment keepalive stopped (provider disabled)")            }        }        #endif    }     func startAugmentKeepalive() {        #if os(macOS)        self.augmentLogger.info(            "Augment keepalive check",            metadata: [                "enabled": self.isEnabled(.augment) ? "1" : "0",                "available": self.isProviderAvailable(.augment) ? "1" : "0",            ])         // Only start keepalive if Augment is enabled        guard self.isEnabled(.augment) else {            self.augmentLogger.warning("Augment keepalive not started (provider disabled)")            return        }         let logger: (String) -> Void = { [augmentLogger = self.augmentLogger] message in            augmentLogger.verbose(message)        }         // Callback to refresh Augment usage after successful session recovery        let onSessionRecovered: () async -> Void = { [weak self] in            guard let self else { return }            self.augmentLogger.info("Augment session recovered; refreshing usage")            await self.refreshProvider(.augment)        }         self.augmentKeepalive = AugmentSessionKeepalive(logger: logger, onSessionRecovered: onSessionRecovered)        self.augmentKeepalive?.start()        self.augmentLogger.info("Augment keepalive started")        #endif    }     /// Force refresh Augment session (called from UI button)    func forceRefreshAugmentSession() async {        #if os(macOS)        self.augmentLogger.info("Augment force refresh requested")        guard let keepalive = self.augmentKeepalive else {            self.augmentLogger.warning("Augment keepalive not running; starting")            self.startAugmentKeepalive()            // Give it a moment to start            try? await Task.sleep(for: .seconds(1))            guard let keepalive = self.augmentKeepalive else {                self.augmentLogger.error("Augment keepalive failed to start")                return            }            await keepalive.forceRefresh()            return        }         await keepalive.forceRefresh()         // Refresh usage after forcing session refresh        self.augmentLogger.info("Refreshing Augment usage after session refresh")        await self.refreshProvider(.augment)        #endif    }     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()            }        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                }                 // Trigger immediate session recovery for Augment when session expires                if provider == .augment, error.localizedDescription.contains("session expired") {                    self.augmentLogger.warning("Augment session expired; triggering recovery")                    Task {                        await self.forceRefreshAugmentSession()                    }                }            }        }    }}