wehub-resource-sync
122 lines
4.2 KiB
Swift
122 lines
4.2 KiB
Swift
import CodexBarCore
|
|
import Foundation
|
|
|
|
@MainActor
|
|
final 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
|
|
}
|
|
}
|