feat(auth): retry auth requests, and schedule next refresh retry in background (#395)

* feat(auth): auto refresh token

* Schedule next token refresh

* move auto refresh token logic to actor

* dont expose StoredSession

* generic RetryableError

* Use request interceptor for retrying requests

* revert

* Revert retryable error

* schedule refresh when accessing session for the first time

* fix tests

Author: Guilherme Souza

Sources/Auth/AuthClient.swift

@@ -17,6 +17,8 @@ public final class AuthClient: Sendable {
   private var date: @Sendable () -> Date { Current.date }
   private var sessionManager: SessionManager { Current.sessionManager }
   private var eventEmitter: AuthStateChangeEventEmitter { Current.eventEmitter }
+  private var logger: (any SupabaseLogger)? { Current.logger }
+  private var storage: any AuthLocalStorage { Current.configuration.localStorage }
 
   /// Returns the session, refreshing it if necessary.
   ///
@@ -27,6 +29,20 @@ public final class AuthClient: Sendable {
     }
   }
 
+  /// Returns the current session, if any.
+  ///
+  /// The session returned by this property may be expired. Use ``session`` for a session that is guaranteed to be valid.
+  public var currentSession: Session? {
+    try? storage.getSession()
+  }
+
+  /// Returns the current user, if any.
+  ///
+  /// The user returned by this property may be outdated. Use ``user(jwt:)`` method to get an up-to-date user instance.
+  public var currentUser: User? {
+    try? storage.getSession()?.user
+  }
+
   /// Namespace for accessing multi-factor authentication API.
   public let mfa = AuthMFA()
   /// Namespace for the GoTrue admin methods.
@@ -41,9 +57,6 @@ public final class AuthClient: Sendable {
   public init(configuration: Configuration) {
     Current = Dependencies(
       configuration: configuration,
-      sessionRefresher: SessionRefresher { [weak self] in
-        try await self?.refreshSession(refreshToken: $0) ?? .empty
-      },
       http: HTTPClient(configuration: configuration)
     )
   }
@@ -158,13 +171,14 @@ public final class AuthClient: Sendable {
 
   private func _signUp(request: HTTPRequest) async throws -> AuthResponse {
     await sessionManager.remove()
+
     let response = try await api.execute(request).decoded(
       as: AuthResponse.self,
       decoder: configuration.decoder
     )
 
     if let session = response.session {
-      try await sessionManager.update(session)
+      await sessionManager.update(session)
       eventEmitter.emit(.signedIn, session: session)
     }
 
@@ -264,7 +278,7 @@ public final class AuthClient: Sendable {
       decoder: configuration.decoder
     )
 
-    try await sessionManager.update(session)
+    await sessionManager.update(session)
     eventEmitter.emit(.signedIn, session: session)
 
     return session
@@ -445,7 +459,7 @@ public final class AuthClient: Sendable {
 
     codeVerifierStorage.set(nil)
 
-    try await sessionManager.update(session)
+    await sessionManager.update(session)
     eventEmitter.emit(.signedIn, session: session)
 
     return session
@@ -640,7 +654,7 @@ public final class AuthClient: Sendable {
       user: user
     )
 
-    try await sessionManager.update(session)
+    await sessionManager.update(session)
     eventEmitter.emit(.signedIn, session: session)
 
     if let type = params["type"], type == "recovery" {
@@ -688,7 +702,7 @@ public final class AuthClient: Sendable {
       )
     }
 
-    try await sessionManager.update(session)
+    await sessionManager.update(session)
     eventEmitter.emit(.signedIn, session: session)
     return session
   }
@@ -805,7 +819,7 @@ public final class AuthClient: Sendable {
     )
 
     if let session = response.session {
-      try await sessionManager.update(session)
+      await sessionManager.update(session)
       eventEmitter.emit(.signedIn, session: session)
     }
 
@@ -889,20 +903,6 @@ public final class AuthClient: Sendable {
     )
   }
 
-  /// Returns the current session, if any.
-  ///
-  /// The session returned by this property may be expired. Use ``session`` for a session that is guaranteed to be valid.
-  public var currentSession: Session? {
-    try? configuration.localStorage.getSession()?.session
-  }
-
-  /// Returns the current user, if any.
-  ///
-  /// The user returned by this property may be outdated. Use ``user(jwt:)`` method to get an up-to-date user instance.
-  public var currentUser: User? {
-    try? configuration.localStorage.getSession()?.session.user
-  }
-
   /// Gets the current user details if there is an existing session.
   /// - Parameter jwt: Takes in an optional access token jwt. If no jwt is provided, user() will
   /// attempt to get the jwt from the current session.
@@ -945,7 +945,7 @@ public final class AuthClient: Sendable {
       )
     ).decoded(as: User.self, decoder: configuration.decoder)
     session.user = updatedUser
-    try await sessionManager.update(session)
+    await sessionManager.update(session)
     eventEmitter.emit(.userUpdated, session: session)
     return updatedUser
   }
@@ -1094,30 +1094,11 @@ public final class AuthClient: Sendable {
   /// - Returns: A new session.
   @discardableResult
   public func refreshSession(refreshToken: String? = nil) async throws -> Session {
-    var credentials = UserCredentials(refreshToken: refreshToken)
-    if credentials.refreshToken == nil {
-      credentials.refreshToken = try await sessionManager.session(shouldValidateExpiration: false)
-        .refreshToken
+    guard let refreshToken = refreshToken ?? currentSession?.refreshToken else {
+      throw AuthError.sessionNotFound
     }
 
-    let session = try await api.execute(
-      .init(
-        url: configuration.url.appendingPathComponent("token"),
-        method: .post,
-        query: [URLQueryItem(name: "grant_type", value: "refresh_token")],
-        body: configuration.encoder.encode(credentials)
-      )
-    ).decoded(as: Session.self, decoder: configuration.decoder)
-
-    if session.user.phoneConfirmedAt != nil || session.user.emailConfirmedAt != nil
-      || session
-      .user.confirmedAt != nil
-    {
-      try await sessionManager.update(session)
-      eventEmitter.emit(.tokenRefreshed, session: session)
-    }
-
-    return session
+    return try await sessionManager.refreshSession(refreshToken)
   }
 
   private func emitInitialSession(forToken token: ObservationToken) async {

Sources/Auth/AuthMFA.swift

@@ -60,7 +60,7 @@ public struct AuthMFA: Sendable {
       )
     ).decoded(decoder: decoder)
 
-    try await sessionManager.update(response)
+    await sessionManager.update(response)
 
     eventEmitter.emit(.mfaChallengeVerified, session: response, token: nil)
 
@@ -116,9 +116,7 @@ public struct AuthMFA: Sendable {
   /// Returns the Authenticator Assurance Level (AAL) for the active session.
   ///
   /// - Returns: An authentication response with the Authenticator Assurance Level.
-  public func getAuthenticatorAssuranceLevel() async throws
-    -> AuthMFAGetAuthenticatorAssuranceLevelResponse
-  {
+  public func getAuthenticatorAssuranceLevel() async throws -> AuthMFAGetAuthenticatorAssuranceLevelResponse {
     do {
       let session = try await sessionManager.session()
       let payload = try decode(jwt: session.accessToken)

Sources/Auth/Internal/APIClient.swift

@@ -8,6 +8,14 @@ extension HTTPClient {
       interceptors.append(LoggerInterceptor(logger: logger))
     }
 
+    interceptors.append(
+      RetryRequestInterceptor(
+        retryableHTTPMethods: RetryRequestInterceptor.defaultRetryableHTTPMethods.union(
+          [.post] // Add POST method so refresh token are also retried.
+        )
+      )
+    )
+
     self.init(fetch: configuration.fetch, interceptors: interceptors)
   }
 }

Sources/Auth/Internal/Dependencies.swift

@@ -4,9 +4,8 @@ import Foundation
 
 struct Dependencies: Sendable {
   var configuration: AuthClient.Configuration
-  var sessionRefresher: SessionRefresher
   var http: any HTTPClientType
-  var sessionManager = SessionManager()
+  var sessionManager = SessionManager.live
   var api = APIClient()
 
   var eventEmitter: AuthStateChangeEventEmitter = .shared

Sources/Auth/Internal/Helpers.swift

@@ -1,3 +1,4 @@
+import _Helpers
 import Foundation
 
 /// Extracts parameters encoded in the URL both in the query and fragment.

Sources/Auth/Internal/SessionManager.swift

@@ -1,50 +1,132 @@
 import _Helpers
 import Foundation
 
-struct SessionRefresher: Sendable {
+struct SessionManager: Sendable {
+  var session: @Sendable () async throws -> Session
   var refreshSession: @Sendable (_ refreshToken: String) async throws -> Session
-}
 
-actor SessionManager {
-  private var task: Task<Session, any Error>?
+  var update: @Sendable (_ session: Session) async -> Void
+  var remove: @Sendable () async -> Void
+}
 
-  private var storage: any AuthLocalStorage {
-    Current.configuration.localStorage
+extension SessionManager {
+  static var live: Self {
+    let instance = LiveSessionManager()
+    return Self(
+      session: { try await instance.session() },
+      refreshSession: { try await instance.refreshSession($0) },
+      update: { await instance.update($0) },
+      remove: { await instance.remove() }
+    )
   }
+}
+
+private actor LiveSessionManager {
+  private var configuration: AuthClient.Configuration { Current.configuration }
+  private var storage: any AuthLocalStorage { Current.configuration.localStorage }
+  private var eventEmitter: AuthStateChangeEventEmitter { Current.eventEmitter }
+  private var logger: (any SupabaseLogger)? { Current.logger }
+  private var api: APIClient { Current.api }
+
+  private var inFlightRefreshTask: Task<Session, any Error>?
+  private var scheduledNextRefreshTask: Task<Void, Never>?
+
+  func session() async throws -> Session {
+    guard let currentSession = try storage.getSession() else {
+      throw AuthError.sessionNotFound
+    }
+
+    if currentSession.isValid {
+      scheduleNextTokenRefresh(currentSession)
 
-  private var sessionRefresher: SessionRefresher {
-    Current.sessionRefresher
+      return currentSession
+    }
+
+    return try await refreshSession(currentSession.refreshToken)
   }
 
-  func session(shouldValidateExpiration: Bool = true) async throws -> Session {
-    if let task {
-      return try await task.value
+  func refreshSession(_ refreshToken: String) async throws -> Session {
+    logger?.debug("begin")
+    defer { logger?.debug("end") }
+
+    if let inFlightRefreshTask {
+      logger?.debug("refresh already in flight")
+      return try await inFlightRefreshTask.value
     }
 
-    task = Task {
-      defer { task = nil }
+    inFlightRefreshTask = Task {
+      logger?.debug("refresh task started")
 
-      guard let currentSession = try storage.getSession() else {
-        throw AuthError.sessionNotFound
+      defer {
+        inFlightRefreshTask = nil
+        logger?.debug("refresh task ended")
       }
 
-      if currentSession.isValid || !shouldValidateExpiration {
-        return currentSession.session
-      }
+      let session = try await api.execute(
+        HTTPRequest(
+          url: configuration.url.appendingPathComponent("token"),
+          method: .post,
+          query: [
+            URLQueryItem(name: "grant_type", value: "refresh_token"),
+          ],
+          body: configuration.encoder.encode(UserCredentials(refreshToken: refreshToken))
+        )
+      )
+      .decoded(as: Session.self, decoder: configuration.decoder)
+
+      update(session)
+      eventEmitter.emit(.tokenRefreshed, session: session)
+
+      scheduleNextTokenRefresh(session)
 
-      let session = try await sessionRefresher.refreshSession(currentSession.session.refreshToken)
-      try update(session)
       return session
     }
 
-    return try await task!.value
+    return try await inFlightRefreshTask!.value
   }
 
-  func update(_ session: Session) throws {
-    try storage.storeSession(StoredSession(session: session))
+  func update(_ session: Session) {
+    do {
+      try storage.storeSession(session)
+    } catch {
+      logger?.error("Failed to store session: \(error)")
+    }
   }
 
   func remove() {
-    try? storage.deleteSession()
+    do {
+      try storage.deleteSession()
+    } catch {
+      logger?.error("Failed to remove session: \(error)")
+    }
+  }
+
+  private func scheduleNextTokenRefresh(_ refreshedSession: Session, source: StaticString = #function) {
+    logger?.debug("source: \(source)")
+
+    guard scheduledNextRefreshTask == nil else {
+      logger?.debug("source: \(source) refresh task already scheduled")
+      return
+    }
+
+    scheduledNextRefreshTask = Task {
+      defer { scheduledNextRefreshTask = nil }
+
+      let expiresAt = Date(timeIntervalSince1970: refreshedSession.expiresAt)
+      let expiresIn = expiresAt.timeIntervalSinceNow
+
+      // if expiresIn < 0, it will refresh right away.
+      let timeToRefresh = max(expiresIn * 0.9, 0)
+
+      logger?.debug("source: \(source) scheduled next token refresh in: \(timeToRefresh)s")
+
+      try? await Task.sleep(nanoseconds: NSEC_PER_SEC * UInt64(timeToRefresh))
+
+      if Task.isCancelled {
+        return
+      }
+
+      _ = try? await refreshSession(refreshedSession.refreshToken)
+    }
   }
 }