Consume CloudKit web auth token rotation from response headers (#462).

Middleware now reads X-Apple-CloudKit-Web-Auth-Token and feeds rotated tokens back through TokenManager so long-lived sessions can extend their window; WebAuthTokenManager becomes an actor to hold mutable state on macOS 11+.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: leogdion

Sources/MistKit/Authentication/AdaptiveTokenManager+Transitions.swift

@@ -61,4 +61,27 @@ extension AdaptiveTokenManager {
 
     return authenticator
   }
+
+  /// Adopts a rotated web auth token from a CloudKit response header.
+  public func didReceiveRotatedWebAuthToken(_ token: String) async throws(TokenManagerError) {
+    guard webAuthToken != nil else {
+      return
+    }
+
+    let authenticator = try WebAuthTokenAuthenticator(
+      apiToken: apiToken,
+      webAuthToken: token
+    )
+    self.webAuthToken = token
+
+    if let storage = storage {
+      do {
+        try await storage.store(authenticator, identifier: apiToken)
+      } catch {
+        Logger(subsystem: .auth).warning(
+          "Failed to store credentials after token rotation: \(error.localizedDescription)"
+        )
+      }
+    }
+  }
 }

Sources/MistKit/Authentication/AuthenticationMiddleware.swift

@@ -29,6 +29,7 @@
 
 internal import Foundation
 internal import HTTPTypes
+internal import Logging
 internal import OpenAPIRuntime
 
 /// Authentication middleware that delegates request mutation to whichever
@@ -50,6 +51,16 @@ internal struct AuthenticationMiddleware: ClientMiddleware {
     var modifiedRequest = request
     var modifiedBody = body
     try await authenticator.authenticate(request: &modifiedRequest, body: &modifiedBody)
-    return try await next(modifiedRequest, modifiedBody, baseURL)
+    let (response, responseBody) = try await next(modifiedRequest, modifiedBody, baseURL)
+    if let rotated = response.headerFields[.cloudKitWebAuthToken] {
+      do {
+        try await tokenManager.didReceiveRotatedWebAuthToken(rotated)
+      } catch {
+        let message = "Failed to consume rotated web auth token: \(error.localizedDescription)"
+        Logger(subsystem: .auth).warning("\(message)")
+        RotatedWebAuthTokenFailureReporter.assertionHandler(message)
+      }
+    }
+    return (response, responseBody)
   }
 }

Sources/MistKit/Authentication/HTTPField.Name+CloudKit.swift

@@ -49,6 +49,11 @@ extension HTTPField.Name {
     "X-Apple-CloudKit-Request-SignatureV1"
   )
 
+  /// Rotated web authentication token returned on every CloudKit response.
+  internal static let cloudKitWebAuthToken = Self.knownFieldName(
+    "X-Apple-CloudKit-Web-Auth-Token"
+  )
+
   private static func knownFieldName(_ name: String) -> HTTPField.Name {
     guard let fieldName = HTTPField.Name(name) else {
       preconditionFailure("Invalid HTTP field name: \(name)")

Sources/MistKit/Authentication/RotatedWebAuthTokenFailureReporter.swift

@@ -0,0 +1,35 @@
+//
+//  RotatedWebAuthTokenFailureReporter.swift
+//  MistKit
+//
+//  Created by Leo Dion.
+//  Copyright © 2026 BrightDigit.
+//
+//  Permission is hereby granted, free of charge, to any person
+//  obtaining a copy of this software and associated documentation
+//  files (the "Software"), to deal in the Software without
+//  restriction, including without limitation the rights to use,
+//  copy, modify, merge, publish, distribute, sublicense, and/or
+//  sell copies of the Software, and to permit persons to whom the
+//  Software is furnished to do so, subject to the following
+//  conditions:
+//
+//  The above copyright notice and this permission notice shall be
+//  included in all copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+//  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+//  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+//  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+//  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+//  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+//  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+//  OTHER DEALINGS IN THE SOFTWARE.
+//
+
+/// Trap hook when rotated web auth token adoption fails in middleware.
+internal enum RotatedWebAuthTokenFailureReporter {
+  @TaskLocal internal static var assertionHandler: @Sendable (String) -> Void = { message in
+    assertionFailure(message)
+  }
+}

Sources/MistKit/Authentication/TokenManager+Rotation.swift

@@ -0,0 +1,37 @@
+//
+//  TokenManager+Rotation.swift
+//  MistKit
+//
+//  Created by Leo Dion.
+//  Copyright © 2026 BrightDigit.
+//
+//  Permission is hereby granted, free of charge, to any person
+//  obtaining a copy of this software and associated documentation
+//  files (the "Software"), to deal in the Software without
+//  restriction, including without limitation the rights to use,
+//  copy, modify, merge, publish, distribute, sublicense, and/or
+//  sell copies of the Software, and to permit persons to whom the
+//  Software is furnished to do so, subject to the following
+//  conditions:
+//
+//  The above copyright notice and this permission notice shall be
+//  included in all copies or substantial portions of the Software.
+//
+//  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+//  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+//  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+//  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+//  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+//  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+//  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+//  OTHER DEALINGS IN THE SOFTWARE.
+//
+
+internal import Foundation
+
+extension TokenManager {
+  /// Adopts a rotated web auth token from the `X-Apple-CloudKit-Web-Auth-Token`
+  /// response header. Default implementation is a no-op for managers that do
+  /// not use web authentication.
+  public func didReceiveRotatedWebAuthToken(_ token: String) async throws(TokenManagerError) {}
+}

Sources/MistKit/Authentication/TokenManager.swift

@@ -46,4 +46,8 @@ public protocol TokenManager: Sendable {
   /// Returns the authenticator that should be used for the next request,
   /// or `nil` if no credentials are available.
   func currentAuthenticator() async throws(TokenManagerError) -> (any Authenticator)?
+
+  /// Adopts a rotated web auth token from the `X-Apple-CloudKit-Web-Auth-Token`
+  /// response header.
+  func didReceiveRotatedWebAuthToken(_ token: String) async throws(TokenManagerError)
 }

Sources/MistKit/Authentication/WebAuthTokenManager.swift

@@ -31,9 +31,9 @@ internal import Foundation
 
 /// Token manager for web authentication with API token + web auth token.
 /// Provides user-specific access to CloudKit Web Services.
-public final class WebAuthTokenManager: TokenManager, Sendable {
+public actor WebAuthTokenManager: TokenManager {
   internal let apiToken: String
-  internal let webAuthToken: String
+  internal var webAuthToken: String
 
   // MARK: - TokenManager Protocol
 
@@ -66,4 +66,10 @@ public final class WebAuthTokenManager: TokenManager, Sendable {
   public func currentAuthenticator() async throws(TokenManagerError) -> (any Authenticator)? {
     try WebAuthTokenAuthenticator(apiToken: apiToken, webAuthToken: webAuthToken)
   }
+
+  /// Adopts a rotated web auth token from a CloudKit response header.
+  public func didReceiveRotatedWebAuthToken(_ token: String) async throws(TokenManagerError) {
+    _ = try WebAuthTokenAuthenticator(apiToken: apiToken, webAuthToken: token)
+    self.webAuthToken = token
+  }
 }

Tests/MistKitTests/Authentication/Middleware/AuthenticationMiddlewareTests+TokenRotation.swift

@@ -0,0 +1,137 @@
+internal import Foundation
+internal import HTTPTypes
+internal import OpenAPIRuntime
+internal import Testing
+
+@testable import MistKit
+
+extension AuthenticationMiddlewareTests {
+  @Suite("Token Rotation")
+  internal struct TokenRotation {
+    private static let validAPIToken = TestConstants.apiToken
+    private static let validWebAuthToken = TestConstants.webAuthToken
+    private static let rotatedWebAuthToken =
+      "rotatedwebauthtokenabcdef0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+      + "abcdefghijklmnopqrstuvwxyz0123456789AB=="
+    private static let testOperationID = TestConstants.operationID
+
+    private static func makeRequest() -> HTTPRequest {
+      HTTPRequest(
+        method: .get,
+        scheme: "https",
+        authority: "api.apple-cloudkit.com",
+        path: "/database/1/iCloud.com.example.app/private/records/query"
+      )
+    }
+
+    private static func responseWithRotatedToken(_ token: String) -> HTTPResponse {
+      var response = HTTPResponse(status: .ok)
+      response.headerFields[.cloudKitWebAuthToken] = token
+      return response
+    }
+
+    @Test("Middleware forwards rotated token to token manager")
+    internal func middlewareForwardsRotatedToken() async throws {
+      let mockManager = MockTokenManagerWithRotation()
+      let middleware = AuthenticationMiddleware(tokenManager: mockManager)
+
+      _ = try await middleware.intercept(
+        Self.makeRequest(),
+        body: nil as HTTPBody?,
+        baseURL: CloudKitService.baseURL,
+        operationID: Self.testOperationID,
+        next: { _, _, _ in
+          (Self.responseWithRotatedToken(Self.rotatedWebAuthToken), nil)
+        }
+      )
+
+      let received = await mockManager.receivedRotatedTokens
+      #expect(received == [Self.rotatedWebAuthToken])
+    }
+
+    @Test("Middleware skips rotation hook when header is absent")
+    internal func middlewareSkipsRotationWithoutHeader() async throws {
+      let mockManager = MockTokenManagerWithRotation()
+      let middleware = AuthenticationMiddleware(tokenManager: mockManager)
+
+      _ = try await middleware.intercept(
+        Self.makeRequest(),
+        body: nil as HTTPBody?,
+        baseURL: CloudKitService.baseURL,
+        operationID: Self.testOperationID,
+        next: { _, _, _ in
+          (HTTPResponse(status: .ok), nil)
+        }
+      )
+
+      let received = await mockManager.receivedRotatedTokens
+      #expect(received.isEmpty)
+    }
+
+    @Test("Middleware returns response when rotation adoption fails")
+    internal func middlewareReturnsResponseWhenRotationFails() async throws {
+      let mockManager = MockTokenManagerWithRotationFailure()
+      let middleware = AuthenticationMiddleware(tokenManager: mockManager)
+
+      let (response, _) = try await RotatedWebAuthTokenFailureReporter.$assertionHandler.withValue(
+        { _ in },
+        operation: {
+          try await middleware.intercept(
+            Self.makeRequest(),
+            body: nil as HTTPBody?,
+            baseURL: CloudKitService.baseURL,
+            operationID: Self.testOperationID,
+            next: { _, _, _ in
+              (Self.responseWithRotatedToken("short"), nil)
+            }
+          )
+        }
+      )
+
+      #expect(response.status == .ok)
+    }
+
+    @Test("WebAuthTokenManager adopts rotated token from middleware")
+    internal func webAuthTokenManagerAdoptsRotatedToken() async throws {
+      let tokenManager = WebAuthTokenManager(
+        apiToken: Self.validAPIToken,
+        webAuthToken: Self.validWebAuthToken
+      )
+      let middleware = AuthenticationMiddleware(tokenManager: tokenManager)
+
+      _ = try await middleware.intercept(
+        Self.makeRequest(),
+        body: nil as HTTPBody?,
+        baseURL: CloudKitService.baseURL,
+        operationID: Self.testOperationID,
+        next: { _, _, _ in
+          (Self.responseWithRotatedToken(Self.rotatedWebAuthToken), nil)
+        }
+      )
+
+      #expect(await tokenManager.webAuthToken == Self.rotatedWebAuthToken)
+      let authenticator = try await tokenManager.currentAuthenticator()
+      let web = try #require(authenticator as? WebAuthTokenAuthenticator)
+      #expect(web.webAuthToken == Self.rotatedWebAuthToken)
+    }
+
+    @Test("AdaptiveTokenManager adopts rotated token from middleware")
+    internal func adaptiveTokenManagerAdoptsRotatedToken() async throws {
+      let tokenManager = AdaptiveTokenManager(apiToken: Self.validAPIToken)
+      try await tokenManager.upgradeToWebAuthentication(webAuthToken: Self.validWebAuthToken)
+      let middleware = AuthenticationMiddleware(tokenManager: tokenManager)
+
+      _ = try await middleware.intercept(
+        Self.makeRequest(),
+        body: nil as HTTPBody?,
+        baseURL: CloudKitService.baseURL,
+        operationID: Self.testOperationID,
+        next: { _, _, _ in
+          (Self.responseWithRotatedToken(Self.rotatedWebAuthToken), nil)
+        }
+      )
+
+      #expect(await tokenManager.webAuthToken == Self.rotatedWebAuthToken)
+    }
+  }
+}

Tests/MistKitTests/Authentication/WebAuth/WebAuthTokenManagerTests+Basic.swift

@@ -22,26 +22,26 @@ extension WebAuthTokenManagerTests {
 
     /// Tests WebAuthTokenManager initialization with valid tokens
     @Test("WebAuthTokenManager initialization with valid tokens")
-    internal func initializationWithValidTokens() {
+    internal func initializationWithValidTokens() async {
       let manager = WebAuthTokenManager(
         apiToken: Self.validAPIToken,
         webAuthToken: Self.validWebAuthToken
       )
 
-      #expect(manager.apiToken == Self.validAPIToken)
-      #expect(manager.webAuthToken == Self.validWebAuthToken)
+      #expect(await manager.apiToken == Self.validAPIToken)
+      #expect(await manager.webAuthToken == Self.validWebAuthToken)
     }
 
     /// Tests WebAuthTokenManager initialization with storage
     @Test("WebAuthTokenManager initialization with storage")
-    internal func initializationWithStorage() {
+    internal func initializationWithStorage() async {
       let manager = WebAuthTokenManager(
         apiToken: Self.validAPIToken,
         webAuthToken: Self.validWebAuthToken
       )
 
-      #expect(manager.apiToken == Self.validAPIToken)
-      #expect(manager.webAuthToken == Self.validWebAuthToken)
+      #expect(await manager.apiToken == Self.validAPIToken)
+      #expect(await manager.webAuthToken == Self.validWebAuthToken)
     }
 
     // MARK: - TokenManager Protocol Tests

Tests/MistKitTests/Mocks/TokenManagers/MockTokenManagerWithRotation.swift

@@ -0,0 +1,47 @@
+//
+//  MockTokenManagerWithRotation.swift
+//  MistKit
+//
+
+internal import Foundation
+internal import HTTPTypes
+internal import OpenAPIRuntime
+
+@testable import MistKit
+
+/// Mock TokenManager that records rotated web auth tokens from responses.
+internal final class MockTokenManagerWithRotation: TokenManager {
+  private actor State {
+    private var receivedTokens: [String] = []
+
+    func append(_ token: String) {
+      receivedTokens.append(token)
+    }
+
+    func tokens() -> [String] {
+      receivedTokens
+    }
+  }
+
+  private let state = State()
+
+  internal var hasCredentials: Bool {
+    get async { true }
+  }
+
+  internal var receivedRotatedTokens: [String] {
+    get async { await state.tokens() }
+  }
+
+  internal func validateCredentials() async throws(TokenManagerError) -> Bool {
+    true
+  }
+
+  internal func currentAuthenticator() async throws(TokenManagerError) -> (any Authenticator)? {
+    try APITokenAuthenticator(token: TestConstants.apiToken)
+  }
+
+  internal func didReceiveRotatedWebAuthToken(_ token: String) async throws(TokenManagerError) {
+    await state.append(token)
+  }
+}