wehub-resource-sync
492 lines
20 KiB
Swift
492 lines
20 KiB
Swift
import Foundation
|
||
|
||
#if os(macOS)
|
||
import AppKit
|
||
import UserNotifications
|
||
|
||
/// Manages automatic session keepalive for Augment to prevent cookie expiration.
|
||
///
|
||
/// This actor monitors cookie expiration and proactively refreshes the session
|
||
/// before cookies expire, ensuring uninterrupted access to Augment APIs.
|
||
@MainActor
|
||
public final class AugmentSessionKeepalive {
|
||
// MARK: - Configuration
|
||
|
||
/// How often to check if session needs refresh (default: 1 minute for faster recovery)
|
||
private let checkInterval: TimeInterval = 60
|
||
|
||
/// Refresh session this many seconds before cookie expiration (default: 5 minutes)
|
||
private let refreshBufferSeconds: TimeInterval = 300
|
||
|
||
/// Minimum time between refresh attempts (default: 1 minute for faster recovery)
|
||
private let minRefreshInterval: TimeInterval = 60
|
||
|
||
/// Maximum time to wait for session refresh (default: 30 seconds)
|
||
private let refreshTimeout: TimeInterval = 30
|
||
|
||
// MARK: - State
|
||
|
||
private var timerTask: Task<Void, Never>?
|
||
private var lastRefreshAttempt: Date?
|
||
private var lastSuccessfulRefresh: Date?
|
||
private var isRefreshing = false
|
||
private let logger: ((String) -> Void)?
|
||
private var onSessionRecovered: (() async -> Void)?
|
||
|
||
/// Track consecutive failures to stop retrying after too many failures
|
||
private var consecutiveFailures = 0
|
||
private let maxConsecutiveFailures = 3 // Stop after 3 failures
|
||
private var hasGivenUp = false
|
||
|
||
// MARK: - Initialization
|
||
|
||
public init(logger: ((String) -> Void)? = nil, onSessionRecovered: (() async -> Void)? = nil) {
|
||
self.logger = logger
|
||
self.onSessionRecovered = onSessionRecovered
|
||
}
|
||
|
||
deinit {
|
||
self.timerTask?.cancel()
|
||
}
|
||
|
||
// MARK: - Public API
|
||
|
||
/// Start the automatic session keepalive timer
|
||
public func start() {
|
||
guard self.timerTask == nil else {
|
||
self.log("Keepalive already running")
|
||
return
|
||
}
|
||
|
||
self.log("🚀 Starting Augment session keepalive")
|
||
self.log(
|
||
" - Check interval: \(Int(self.checkInterval))s "
|
||
+ "(\(Self.durationDescription(seconds: self.checkInterval)))")
|
||
self.log(
|
||
" - Refresh buffer: \(Int(self.refreshBufferSeconds))s "
|
||
+ "(\(Self.durationDescription(seconds: self.refreshBufferSeconds)) before expiry)")
|
||
self.log(
|
||
" - Min refresh interval: \(Int(self.minRefreshInterval))s "
|
||
+ "(\(Self.durationDescription(seconds: self.minRefreshInterval)))")
|
||
|
||
self.timerTask = Task.detached(priority: .utility) { [weak self] in
|
||
while !Task.isCancelled {
|
||
try? await Task.sleep(for: .seconds(self?.checkInterval ?? 300))
|
||
await self?.checkAndRefreshIfNeeded()
|
||
}
|
||
}
|
||
|
||
self.log("✅ Keepalive timer started successfully")
|
||
}
|
||
|
||
/// Stop the automatic session keepalive timer
|
||
public func stop() {
|
||
self.log("Stopping Augment session keepalive")
|
||
self.timerTask?.cancel()
|
||
self.timerTask = nil
|
||
}
|
||
|
||
/// Manually trigger a session refresh (bypasses rate limiting)
|
||
public func forceRefresh() async {
|
||
self.log("Force refresh requested")
|
||
await self.performRefresh(forced: true)
|
||
}
|
||
|
||
// MARK: - Private Implementation
|
||
|
||
private func checkAndRefreshIfNeeded() async {
|
||
guard !self.isRefreshing else {
|
||
self.log("Refresh already in progress, skipping check")
|
||
return
|
||
}
|
||
|
||
// Stop trying if we've given up
|
||
if self.hasGivenUp {
|
||
self.log("⏸️ Keepalive has given up after \(self.maxConsecutiveFailures) consecutive failures")
|
||
self.log(" User must manually log in to Augment and click 'Refresh Session'")
|
||
return
|
||
}
|
||
|
||
// Rate limit: don't refresh too frequently
|
||
if let lastAttempt = self.lastRefreshAttempt {
|
||
let timeSinceLastAttempt = Date().timeIntervalSince(lastAttempt)
|
||
if timeSinceLastAttempt < self.minRefreshInterval {
|
||
self.log(
|
||
"Skipping refresh (last attempt \(Int(timeSinceLastAttempt))s ago, " +
|
||
"min interval: \(Int(self.minRefreshInterval))s)")
|
||
return
|
||
}
|
||
}
|
||
|
||
// Check if cookies are about to expire
|
||
let shouldRefresh = await self.shouldRefreshSession()
|
||
if shouldRefresh {
|
||
await self.performRefresh(forced: false)
|
||
}
|
||
}
|
||
|
||
private func shouldRefreshSession() async -> Bool {
|
||
do {
|
||
let session = try AugmentCookieImporter.importSession(logger: self.logger)
|
||
|
||
self.log("📊 Cookie Status Check:")
|
||
self.log(" Total cookies: \(session.cookies.count)")
|
||
self.log(" Source: \(session.sourceLabel)")
|
||
|
||
// Log each cookie's expiration status
|
||
for cookie in session.cookies {
|
||
if let expiry = cookie.expiresDate {
|
||
let timeUntil = expiry.timeIntervalSinceNow
|
||
let status = timeUntil > 0 ? "expires in \(Int(timeUntil))s" : "EXPIRED \(Int(-timeUntil))s ago"
|
||
self.log(" - \(cookie.name): \(status)")
|
||
} else {
|
||
self.log(" - \(cookie.name): session cookie (no expiry)")
|
||
}
|
||
}
|
||
|
||
// Find the earliest expiration date among session cookies
|
||
let expirationDates = session.cookies.compactMap(\.expiresDate)
|
||
|
||
guard !expirationDates.isEmpty else {
|
||
// Session cookies (no expiration) - refresh periodically
|
||
self.log(" All cookies are session cookies (no expiration dates)")
|
||
if let lastRefresh = self.lastSuccessfulRefresh {
|
||
let timeSinceRefresh = Date().timeIntervalSince(lastRefresh)
|
||
// Refresh every 30 minutes for session cookies
|
||
if timeSinceRefresh > 1800 {
|
||
self.log(" ⚠️ Need periodic refresh (\(Int(timeSinceRefresh))s since last refresh)")
|
||
return true
|
||
} else {
|
||
self.log(" ✅ Recently refreshed (\(Int(timeSinceRefresh))s ago)")
|
||
return false
|
||
}
|
||
} else {
|
||
// Never refreshed - do it now
|
||
self.log(" ⚠️ Never refreshed - doing initial refresh")
|
||
return true
|
||
}
|
||
}
|
||
|
||
let earliestExpiration = expirationDates.min()!
|
||
let timeUntilExpiration = earliestExpiration.timeIntervalSinceNow
|
||
let expiringCookie = session.cookies.first { $0.expiresDate == earliestExpiration }
|
||
|
||
if timeUntilExpiration < self.refreshBufferSeconds {
|
||
self.log(" ⚠️ REFRESH NEEDED:")
|
||
self.log(" Earliest expiring cookie: \(expiringCookie?.name ?? "unknown")")
|
||
self.log(" Time until expiration: \(Int(timeUntilExpiration))s")
|
||
self.log(" Refresh threshold: \(Int(self.refreshBufferSeconds))s")
|
||
return true
|
||
} else {
|
||
self.log(" ✅ Session healthy:")
|
||
self.log(" Earliest expiring cookie: \(expiringCookie?.name ?? "unknown")")
|
||
self.log(" Time until expiration: \(Int(timeUntilExpiration))s")
|
||
return false
|
||
}
|
||
} catch {
|
||
self.log("✗ Failed to check session: \(error.localizedDescription)")
|
||
return false
|
||
}
|
||
}
|
||
|
||
private func performRefresh(forced: Bool) async {
|
||
self.isRefreshing = true
|
||
self.lastRefreshAttempt = Date()
|
||
defer { self.isRefreshing = false }
|
||
|
||
self.log(forced ? "Performing forced session refresh..." : "Performing automatic session refresh...")
|
||
|
||
// If this is a forced refresh (user clicked "Refresh Session"), reset failure tracking
|
||
if forced {
|
||
self.consecutiveFailures = 0
|
||
self.hasGivenUp = false
|
||
self.log("🔄 Manual refresh - resetting failure tracking")
|
||
}
|
||
|
||
do {
|
||
// Step 1: Ping the session endpoint to trigger cookie refresh
|
||
let refreshed = try await self.pingSessionEndpoint()
|
||
|
||
if refreshed {
|
||
// Step 2: Re-import cookies from browser
|
||
try await Task.sleep(for: .seconds(1)) // Brief delay for browser to update cookies
|
||
let newSession = try AugmentCookieImporter.importSession(logger: self.logger)
|
||
|
||
await AugmentSessionStore.shared.setCookies(newSession.cookies)
|
||
CookieHeaderCache.store(
|
||
provider: .augment,
|
||
cookieHeader: newSession.cookieHeader,
|
||
sourceLabel: newSession.sourceLabel)
|
||
|
||
self.log(
|
||
"✅ Session refresh successful - imported \(newSession.cookies.count) cookies " +
|
||
"from \(newSession.sourceLabel)")
|
||
self.lastSuccessfulRefresh = Date()
|
||
|
||
// Reset failure tracking on success
|
||
self.consecutiveFailures = 0
|
||
self.hasGivenUp = false
|
||
|
||
if let callback = self.onSessionRecovered {
|
||
self.log("🔄 Triggering usage refresh after session refresh")
|
||
await callback()
|
||
}
|
||
} else {
|
||
self.log("⚠️ Session refresh returned no new cookies")
|
||
self.consecutiveFailures += 1
|
||
self.checkIfShouldGiveUp()
|
||
}
|
||
} catch AugmentSessionKeepaliveError.sessionExpired {
|
||
self.log("🔐 Session expired - attempting automatic recovery...")
|
||
self.consecutiveFailures += 1
|
||
|
||
if self.consecutiveFailures >= self.maxConsecutiveFailures {
|
||
self.log("❌ Too many consecutive failures (\(self.consecutiveFailures)) - giving up")
|
||
self.log(" User must manually log in to Augment and click 'Refresh Session'")
|
||
self.hasGivenUp = true
|
||
self.notifyUserLoginRequired()
|
||
} else {
|
||
await self.attemptSessionRecovery()
|
||
}
|
||
} catch {
|
||
self.log("✗ Session refresh failed: \(error.localizedDescription)")
|
||
self.consecutiveFailures += 1
|
||
self.checkIfShouldGiveUp()
|
||
}
|
||
}
|
||
|
||
private func checkIfShouldGiveUp() {
|
||
if self.consecutiveFailures >= self.maxConsecutiveFailures {
|
||
self.log("❌ Too many consecutive failures (\(self.consecutiveFailures)) - giving up")
|
||
self.log(" User must manually log in to Augment and click 'Refresh Session'")
|
||
self.hasGivenUp = true
|
||
self.notifyUserLoginRequired()
|
||
}
|
||
}
|
||
|
||
/// Attempt to recover from an expired session by triggering browser re-authentication
|
||
private func attemptSessionRecovery() async {
|
||
self.log("🔄 Attempting automatic session recovery...")
|
||
self.log(" Strategy: Open Augment dashboard to trigger browser re-auth")
|
||
|
||
#if os(macOS)
|
||
// Open the Augment dashboard in the default browser
|
||
// This will trigger the browser to re-authenticate if the user is still logged in
|
||
if let url = URL(string: "https://app.augmentcode.com") {
|
||
_ = await MainActor.run {
|
||
NSWorkspace.shared.open(url)
|
||
}
|
||
self.log(" ✅ Opened Augment dashboard in browser")
|
||
self.log(" ⏳ Waiting 5 seconds for browser to re-authenticate...")
|
||
|
||
// Wait for browser to potentially re-authenticate
|
||
try? await Task.sleep(for: .seconds(5))
|
||
|
||
// Try to import cookies again
|
||
do {
|
||
let newSession = try AugmentCookieImporter.importSession(logger: self.logger)
|
||
self.log(" ✅ Session recovery successful - imported \(newSession.cookies.count) cookies")
|
||
self.lastSuccessfulRefresh = Date()
|
||
|
||
// Verify the session is actually valid by pinging the API
|
||
let isValid = try await self.pingSessionEndpoint()
|
||
if isValid {
|
||
self.log(" ✅ Session verified - recovery complete!")
|
||
// Notify UsageStore to refresh Augment usage
|
||
if let callback = self.onSessionRecovered {
|
||
self.log(" 🔄 Triggering usage refresh after successful recovery")
|
||
await callback()
|
||
}
|
||
} else {
|
||
self.log(" ⚠️ Session imported but not yet valid - may need manual login")
|
||
self.notifyUserLoginRequired()
|
||
}
|
||
} catch {
|
||
self.log(" ✗ Session recovery failed: \(error.localizedDescription)")
|
||
self.log(" ℹ️ User needs to manually log in to Augment")
|
||
self.notifyUserLoginRequired()
|
||
}
|
||
}
|
||
#else
|
||
self.log(" ✗ Automatic recovery not supported on this platform")
|
||
#endif
|
||
}
|
||
|
||
/// Notify the user that they need to log in to Augment
|
||
private func notifyUserLoginRequired() {
|
||
#if os(macOS)
|
||
self.log("📢 Sending notification: Augment session expired")
|
||
|
||
Task {
|
||
let center = UNUserNotificationCenter.current()
|
||
|
||
// Request authorization if needed
|
||
do {
|
||
let granted = try await center.requestAuthorization(options: [.alert, .sound])
|
||
guard granted else {
|
||
self.log("⚠️ Notification permission denied")
|
||
return
|
||
}
|
||
} catch {
|
||
self.log("✗ Failed to request notification permission: \(error)")
|
||
return
|
||
}
|
||
|
||
// Create notification content
|
||
let content = UNMutableNotificationContent()
|
||
content.title = "Augment Session Expired"
|
||
content.body = "Please log in to app.augmentcode.com to restore your session."
|
||
content.sound = .default
|
||
|
||
// Create trigger (deliver immediately)
|
||
let request = UNNotificationRequest(
|
||
identifier: "augment-session-expired-\(UUID().uuidString)",
|
||
content: content,
|
||
trigger: nil)
|
||
|
||
// Deliver notification
|
||
do {
|
||
try await center.add(request)
|
||
self.log("✅ Notification delivered successfully")
|
||
} catch {
|
||
self.log("✗ Failed to deliver notification: \(error)")
|
||
}
|
||
}
|
||
#endif
|
||
}
|
||
|
||
/// Ping Augment's session endpoint to trigger cookie refresh
|
||
private func pingSessionEndpoint() async throws -> Bool {
|
||
// Try to get current cookies first
|
||
let currentSession = try? AugmentCookieImporter.importSession(logger: self.logger)
|
||
guard let cookieHeader = currentSession?.cookieHeader else {
|
||
self.log("No cookies available for session ping")
|
||
return false
|
||
}
|
||
|
||
self.log("🔄 Attempting session refresh...")
|
||
self.log(" Cookies being sent: \(cookieHeader.prefix(100))...")
|
||
|
||
// Try multiple endpoints - Augment might use different auth patterns
|
||
let endpoints = [
|
||
"https://app.augmentcode.com/api/auth/session", // NextAuth pattern
|
||
"https://app.augmentcode.com/api/session", // Alternative
|
||
"https://app.augmentcode.com/api/user", // User endpoint
|
||
]
|
||
|
||
var receivedUnauthorized = false
|
||
|
||
for (index, urlString) in endpoints.enumerated() {
|
||
self.log(" Trying endpoint \(index + 1)/\(endpoints.count): \(urlString)")
|
||
|
||
guard let sessionURL = URL(string: urlString) else { continue }
|
||
var request = URLRequest(url: sessionURL)
|
||
request.timeoutInterval = self.refreshTimeout
|
||
request.setValue(cookieHeader, forHTTPHeaderField: "Cookie")
|
||
request.setValue("application/json", forHTTPHeaderField: "Accept")
|
||
request.setValue("https://app.augmentcode.com", forHTTPHeaderField: "Origin")
|
||
request.setValue("https://app.augmentcode.com", forHTTPHeaderField: "Referer")
|
||
|
||
do {
|
||
let (data, response) = try await ProviderHTTPClient.shared.data(for: request)
|
||
|
||
guard let httpResponse = response as? HTTPURLResponse else {
|
||
self.log(" ✗ Invalid response type")
|
||
continue
|
||
}
|
||
|
||
self.log(" Response: HTTP \(httpResponse.statusCode)")
|
||
|
||
// Log Set-Cookie headers if present
|
||
if let setCookies = httpResponse.allHeaderFields["Set-Cookie"] as? String {
|
||
self.log(" Set-Cookie headers received: \(setCookies.prefix(100))...")
|
||
}
|
||
|
||
if httpResponse.statusCode == 200 {
|
||
// Check if we got a valid session response
|
||
if let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] {
|
||
self.log(" JSON response keys: \(json.keys.joined(separator: ", "))")
|
||
|
||
if json["user"] != nil || json["email"] != nil || json["session"] != nil {
|
||
self.log(" ✅ Valid session data found!")
|
||
return true
|
||
} else {
|
||
self.log(" ⚠️ 200 OK but no session data in response")
|
||
// Try next endpoint
|
||
continue
|
||
}
|
||
} else {
|
||
self.log(" ⚠️ 200 OK but response is not JSON")
|
||
if let responseText = String(data: data, encoding: .utf8) {
|
||
self.log(" Response text: \(responseText.prefix(200))...")
|
||
}
|
||
continue
|
||
}
|
||
} else if httpResponse.statusCode == 401 {
|
||
self.log(" ✗ 401 Unauthorized - session expired")
|
||
receivedUnauthorized = true
|
||
// Don't throw immediately - try all endpoints first
|
||
continue
|
||
} else if httpResponse.statusCode == 404 {
|
||
self.log(" ✗ 404 Not Found - trying next endpoint")
|
||
continue
|
||
} else {
|
||
self.log(" ✗ HTTP \(httpResponse.statusCode) - trying next endpoint")
|
||
continue
|
||
}
|
||
} catch {
|
||
self.log(" ✗ Request failed: \(error.localizedDescription)")
|
||
continue
|
||
}
|
||
}
|
||
|
||
// If we got 401 from all endpoints, the session is definitely expired
|
||
if receivedUnauthorized {
|
||
self.log("⚠️ All endpoints returned 401 - session is expired")
|
||
throw AugmentSessionKeepaliveError.sessionExpired
|
||
}
|
||
|
||
self.log("⚠️ All session endpoints failed or returned no valid data")
|
||
return false
|
||
}
|
||
|
||
private static let log = CodexBarLog.logger(LogCategories.augmentKeepalive)
|
||
|
||
private static func durationDescription(seconds: TimeInterval) -> String {
|
||
let totalSeconds = max(0, Int(seconds.rounded()))
|
||
if totalSeconds >= 60, totalSeconds % 60 == 0 {
|
||
let minutes = totalSeconds / 60
|
||
return "\(minutes) minute\(minutes == 1 ? "" : "s")"
|
||
}
|
||
return "\(totalSeconds) second\(totalSeconds == 1 ? "" : "s")"
|
||
}
|
||
|
||
private func log(_ message: String) {
|
||
let timestamp = Date().formatted(date: .omitted, time: .standard)
|
||
let fullMessage = "[\(timestamp)] [AugmentKeepalive] \(message)"
|
||
self.logger?(fullMessage)
|
||
Self.log.debug(fullMessage)
|
||
}
|
||
}
|
||
|
||
// MARK: - Errors
|
||
|
||
public enum AugmentSessionKeepaliveError: LocalizedError, Sendable {
|
||
case invalidResponse
|
||
case sessionExpired
|
||
case networkError(String)
|
||
|
||
public var errorDescription: String? {
|
||
switch self {
|
||
case .invalidResponse:
|
||
"Invalid response from session endpoint"
|
||
case .sessionExpired:
|
||
"Session has expired"
|
||
case let .networkError(message):
|
||
"Network error: \(message)"
|
||
}
|
||
}
|
||
}
|
||
|
||
#endif
|