121 lines · 4.2 KB · Swift
RawDownload
import CodexBarCoreimport Foundation @MainActorfinal class AugmentProviderRuntime: ProviderRuntime {    let id: UsageProvider = .augment    private var keepalive: AugmentSessionKeepalive?    #if DEBUG    private(set) var _test_keepaliveStopCount = 0    var _test_isKeepaliveRunning: Bool {        self.keepalive != nil    }    #endif     func start(context: ProviderRuntimeContext) {        self.updateKeepalive(context: context)    }     func stop(context: ProviderRuntimeContext) {        self.stopKeepalive(context: context, reason: "provider disabled")    }     func settingsDidChange(context: ProviderRuntimeContext) {        self.updateKeepalive(context: context)    }     func providerDidFail(context: ProviderRuntimeContext, provider: UsageProvider, error: Error) {        guard provider == .augment else { return }        let message = error.localizedDescription        guard message.contains("session expired") else { return }        context.store.augmentLogger.warning("Augment session expired; triggering recovery")        Task { [weak self] in            guard let self else { return }            await self.forceRefresh(context: context)        }    }     func perform(action: ProviderRuntimeAction, context: ProviderRuntimeContext) async {        switch action {        case .forceSessionRefresh:            await self.forceRefresh(context: context)        case .openAIWebAccessToggled:            break        }    }     private func updateKeepalive(context: ProviderRuntimeContext) {        #if os(macOS)        let shouldRun = context.store.isEnabled(.augment)        let isRunning = self.keepalive != nil         if shouldRun, !isRunning {            self.startKeepalive(context: context)        } else if !shouldRun, isRunning {            self.stopKeepalive(context: context, reason: "provider disabled")        }        #endif    }     private func startKeepalive(context: ProviderRuntimeContext) {        #if os(macOS)        context.store.augmentLogger.info(            "Augment keepalive check",            metadata: [                "enabled": context.store.isEnabled(.augment) ? "1" : "0",                "available": context.store.isProviderAvailable(.augment) ? "1" : "0",            ])         guard context.store.isEnabled(.augment) else {            context.store.augmentLogger.warning("Augment keepalive not started (provider disabled)")            return        }         let logger: (String) -> Void = { [augmentLogger = context.store.augmentLogger] message in            augmentLogger.verbose(message)        }         let onSessionRecovered: () async -> Void = { [weak store = context.store] in            guard let store else { return }            store.augmentLogger.info("Augment session recovered; refreshing usage")            await store.refreshProvider(.augment)        }         self.keepalive = AugmentSessionKeepalive(logger: logger, onSessionRecovered: onSessionRecovered)        self.keepalive?.start()        context.store.augmentLogger.info("Augment keepalive started")        #endif    }     private func stopKeepalive(context: ProviderRuntimeContext, reason: String) {        #if os(macOS)        guard let keepalive = self.keepalive else { return }        keepalive.stop()        self.keepalive = nil        #if DEBUG        self._test_keepaliveStopCount += 1        #endif        context.store.augmentLogger.info("Augment keepalive stopped (\(reason))")        #endif    }     private func forceRefresh(context: ProviderRuntimeContext) async {        #if os(macOS)        context.store.augmentLogger.info("Augment force refresh requested")        CookieHeaderCache.clear(provider: .augment)        guard let keepalive = self.keepalive else {            context.store.augmentLogger.warning("Augment keepalive not running; starting")            self.startKeepalive(context: context)            try? await Task.sleep(for: .seconds(1))            guard let keepalive = self.keepalive else {                context.store.augmentLogger.error("Augment keepalive failed to start")                return            }            await keepalive.forceRefresh()            return        }         await keepalive.forceRefresh()        #endif    }}