Make CoreATProtocol auth proxy support session-aware with direct fallback

Sources/CoreATProtocol/OAuth/ATProtoOAuth.swift

···
     public let storePrivateKey: @Sendable (Data) async throws -> Void
     public let retrieveAuthProxyKeyID: (@Sendable () async throws -> String?)?
     public let storeAuthProxyKeyID: (@Sendable (String) async throws -> Void)?
+    public let clearAuthProxyKeyID: (@Sendable () async throws -> Void)?
+    public let retrieveAuthProxyKeyIDForLogin: (@Sendable (Login) async throws -> String?)?
+    public let storeAuthProxyKeyIDForLogin: (@Sendable (Login, String) async throws -> Void)?
+    public let clearAuthProxyKeyIDForLogin: (@Sendable (Login) async throws -> Void)?
 
     public init(
         retrieveLogin: @escaping @Sendable () async throws -> Login?,
···
         retrievePrivateKey: @escaping @Sendable () async throws -> Data?,
         storePrivateKey: @escaping @Sendable (Data) async throws -> Void,
         retrieveAuthProxyKeyID: (@Sendable () async throws -> String?)? = nil,
-        storeAuthProxyKeyID: (@Sendable (String) async throws -> Void)? = nil
+        storeAuthProxyKeyID: (@Sendable (String) async throws -> Void)? = nil,
+        clearAuthProxyKeyID: (@Sendable () async throws -> Void)? = nil,
+        retrieveAuthProxyKeyIDForLogin: (@Sendable (Login) async throws -> String?)? = nil,
+        storeAuthProxyKeyIDForLogin: (@Sendable (Login, String) async throws -> Void)? = nil,
+        clearAuthProxyKeyIDForLogin: (@Sendable (Login) async throws -> Void)? = nil
     ) {
         self.retrieveLogin = retrieveLogin
         self.storeLogin = storeLogin
···
         self.storePrivateKey = storePrivateKey
         self.retrieveAuthProxyKeyID = retrieveAuthProxyKeyID
         self.storeAuthProxyKeyID = storeAuthProxyKeyID
+        self.clearAuthProxyKeyID = clearAuthProxyKeyID
+        self.retrieveAuthProxyKeyIDForLogin = retrieveAuthProxyKeyIDForLogin
+        self.storeAuthProxyKeyIDForLogin = storeAuthProxyKeyIDForLogin
+        self.clearAuthProxyKeyIDForLogin = clearAuthProxyKeyIDForLogin
     }
 }
 
···
     // JWT signing keys (pattern from AtProtocol)
     private var keys: JWTKeyCollection
     private var privateKey: ES256PrivateKey
+
+    private struct AuthenticationAttemptResult {
+        let login: Login
+        let usedAuthProxy: Bool
+    }
 
     public init(config: ATProtoOAuthConfig, storage: ATProtoAuthStorage) async {
         self.config = config
···
             try await self.generateJWT(params: params)
         }
 
-        // Step 7: Create proxy URL loader if auth proxy is configured
+        // Step 7: Create authenticator configuration
+        let tokenHandling = buildTokenHandling(
+            accountHint: identifier,
+            server: serverConfig,
+            jwtGenerator: jwtGenerator,
+            expectedSubjectDID: identity.did,
+            expectedAuthorizationServer: identity.authorizationServer
+        )
+
+        let makeAuthenticatorConfiguration = {
+            Authenticator.Configuration(
+                appCredentials: clientConfig.credentials,
+                loginStorage: loginStorage,
+                tokenHandling: tokenHandling,
+                mode: .manualOnly,
+                userAuthenticator: userAuthenticator
+            )
+        }
+
+        // Step 8: Create authenticator
         let proxyKeyIDStorage: AuthProxyKeyIDStorage?
-        let proxyURLLoader: URLResponseProvider?
+        let directAuthenticator = Authenticator(config: makeAuthenticatorConfiguration())
+        let authenticator: Authenticator
 
         if let proxyBaseURL = config.authProxyBaseURL {
-            let initialKeyID = try? await storage.retrieveAuthProxyKeyID?()
-            let keyIDStorage = AuthProxyKeyIDStorage(keyID: initialKeyID)
+            let keyIDStorage = AuthProxyKeyIDStorage()
             proxyKeyIDStorage = keyIDStorage
-            proxyURLLoader = makeProxyURLLoader(
+            let proxyURLLoader = makeProxyURLLoader(
                 proxyBaseURL: proxyBaseURL,
                 tokenEndpoint: serverConfig.tokenEndpoint,
                 parEndpoint: serverConfig.pushedAuthorizationRequestEndpoint,
                 issuer: serverConfig.issuer,
-                keyIDStorage: keyIDStorage
+                keyIDStorage: keyIDStorage,
+                fallbackProvider: provider
             )
+            authenticator = Authenticator(config: makeAuthenticatorConfiguration(), urlLoader: proxyURLLoader)
         } else {
             proxyKeyIDStorage = nil
-            proxyURLLoader = nil
+            authenticator = directAuthenticator
         }
 
-        // Step 8: Create authenticator
-        let tokenHandling = buildTokenHandling(
-            accountHint: identifier,
-            server: serverConfig,
-            jwtGenerator: jwtGenerator,
-            expectedSubjectDID: identity.did,
-            expectedAuthorizationServer: identity.authorizationServer
-        )
-
-        let authenticatorConfig = Authenticator.Configuration(
-            appCredentials: clientConfig.credentials,
-            loginStorage: loginStorage,
-            tokenHandling: tokenHandling,
-            mode: .manualOnly,
-            userAuthenticator: userAuthenticator
-        )
-
-        let authenticator = Authenticator(config: authenticatorConfig, urlLoader: proxyURLLoader)
-
         // Step 9: Trigger authentication with user interaction
-        let login: Login
+        let authenticationResult: AuthenticationAttemptResult
         do {
-            login = try await authenticator.authenticate()
+            authenticationResult = try await performAuthentication(
+                preferredAuthenticator: authenticator,
+                fallbackAuthenticator: config.authProxyBaseURL == nil ? nil : directAuthenticator,
+                prefersAuthProxy: config.authProxyBaseURL != nil
+            )
         } catch {
-            if shouldRecoverFromRefreshFailure(error) {
-                try await storage.storeLogin(Login(token: "invalid", validUntilDate: .distantPast))
-                try await resetDPoPKey()
-                do {
-                    login = try await authenticator.authenticate()
-                } catch {
-                    throw ATProtoOAuthError.authenticationFailed("OAuth flow failed: \(error.localizedDescription)")
-                }
-            } else {
-                throw ATProtoOAuthError.authenticationFailed("OAuth flow failed: \(error.localizedDescription)")
-            }
+            throw ATProtoOAuthError.authenticationFailed("OAuth flow failed: \(error.localizedDescription)")
         }
+        let login = authenticationResult.login
 
         // Step 10: Persist auth proxy key ID
-        if let proxyKeyIDStorage, let keyID = await proxyKeyIDStorage.keyID {
-            try? await storage.storeAuthProxyKeyID?(keyID)
+        if config.authProxyBaseURL != nil {
+            if authenticationResult.usedAuthProxy,
+               let proxyKeyIDStorage,
+               let keyID = await proxyKeyIDStorage.keyID {
+                await persistAuthProxyKeyID(keyID, for: login)
+            } else {
+                await clearAuthProxyKeyID(for: login)
+            }
         }
 
         // Step 11: Setup CoreATProtocol environment
···
         let proxyKeyIDStorage: AuthProxyKeyIDStorage?
 
         if let proxyBaseURL = config.authProxyBaseURL {
-            let initialKeyID = try? await storage.retrieveAuthProxyKeyID?()
+            let initialKeyID = await retrieveAuthProxyKeyID(for: login)
             let keyIDStorage = AuthProxyKeyIDStorage(keyID: initialKeyID)
             proxyKeyIDStorage = keyIDStorage
             baseProvider = makeProxyURLLoader(
···
                 tokenEndpoint: serverConfig.tokenEndpoint,
                 parEndpoint: serverConfig.pushedAuthorizationRequestEndpoint,
                 issuer: serverConfig.issuer,
-                keyIDStorage: keyIDStorage
+                keyIDStorage: keyIDStorage,
+                fallbackProvider: provider
             )
         } else {
             baseProvider = provider
             proxyKeyIDStorage = nil
         }
 
-        let responseProvider: URLResponseProvider = { request in
+        let directResponseProvider: URLResponseProvider = { request in
+            try await self.dpopRequestActor.response(
+                request: request,
+                jwtGenerator: jwtGenerator,
+                provider: provider,
+                issuingServer: issuer
+            )
+        }
+
+        let proxyResponseProvider: URLResponseProvider = { request in
             try await self.dpopRequestActor.response(
                 request: request,
                 jwtGenerator: jwtGenerator,
···
             )
         }
 
-        let refreshedLogin = try await refreshProvider(login, clientConfig.credentials, responseProvider)
+        let refreshedLogin: Login
+        let usedAuthProxy: Bool
+
+        if config.authProxyBaseURL != nil {
+            do {
+                refreshedLogin = try await refreshProvider(login, clientConfig.credentials, proxyResponseProvider)
+                usedAuthProxy = true
+            } catch {
+                guard shouldRetryWithoutAuthProxy(after: error) else {
+                    throw error
+                }
+
+                refreshedLogin = try await refreshProvider(login, clientConfig.credentials, directResponseProvider)
+                usedAuthProxy = false
+            }
+        } else {
+            refreshedLogin = try await refreshProvider(login, clientConfig.credentials, directResponseProvider)
+            usedAuthProxy = false
+        }
+
         try await storage.storeLogin(refreshedLogin)
 
         // Persist updated auth proxy key ID
-        if let proxyKeyIDStorage, let keyID = await proxyKeyIDStorage.keyID {
-            try? await storage.storeAuthProxyKeyID?(keyID)
+        if config.authProxyBaseURL != nil {
+            if usedAuthProxy, let proxyKeyIDStorage, let keyID = await proxyKeyIDStorage.keyID {
+                await persistAuthProxyKeyID(keyID, for: refreshedLogin)
+            } else {
+                await clearAuthProxyKeyID(for: refreshedLogin)
+            }
         }
 
         return refreshedLogin
···
         try await persistPrivateKey()
     }
 
+    private func performAuthentication(
+        preferredAuthenticator: Authenticator,
+        fallbackAuthenticator: Authenticator?,
+        prefersAuthProxy: Bool
+    ) async throws -> AuthenticationAttemptResult {
+        do {
+            return AuthenticationAttemptResult(
+                login: try await preferredAuthenticator.authenticate(),
+                usedAuthProxy: prefersAuthProxy
+            )
+        } catch {
+            if shouldRecoverFromRefreshFailure(error) {
+                let previousLogin = try? await storage.retrieveLogin()
+                try await storage.storeLogin(Login(token: "invalid", validUntilDate: .distantPast))
+                await clearAuthProxyKeyID(for: previousLogin)
+                try await resetDPoPKey()
+
+                do {
+                    return AuthenticationAttemptResult(
+                        login: try await preferredAuthenticator.authenticate(),
+                        usedAuthProxy: prefersAuthProxy
+                    )
+                } catch {
+                    if let fallbackAuthenticator, shouldRetryWithoutAuthProxy(after: error) {
+                        return AuthenticationAttemptResult(
+                            login: try await fallbackAuthenticator.authenticate(),
+                            usedAuthProxy: false
+                        )
+                    }
+
+                    throw error
+                }
+            }
+
+            if let fallbackAuthenticator, shouldRetryWithoutAuthProxy(after: error) {
+                return AuthenticationAttemptResult(
+                    login: try await fallbackAuthenticator.authenticate(),
+                    usedAuthProxy: false
+                )
+            }
+
+            throw error
+        }
+    }
+
     private func shouldRecoverFromRefreshFailure(_ error: Error) -> Bool {
         guard let authError = error as? AuthenticatorError else {
             return false
···
         default:
             return false
         }
+    }
+
+    private func shouldRetryWithoutAuthProxy(after error: Error) -> Bool {
+        guard let proxyError = error as? AuthProxyError else {
+            return false
+        }
+
+        switch proxyError {
+        case .unavailable:
+            return true
+        }
+    }
+
+    private func retrieveAuthProxyKeyID(for login: Login) async -> String? {
+        if let retrieveForLogin = storage.retrieveAuthProxyKeyIDForLogin {
+            return try? await retrieveForLogin(login)
+        }
+
+        if let retrieve = storage.retrieveAuthProxyKeyID {
+            return try? await retrieve()
+        }
+
+        return nil
+    }
+
+    private func persistAuthProxyKeyID(_ keyID: String, for login: Login) async {
+        if let storeForLogin = storage.storeAuthProxyKeyIDForLogin {
+            try? await storeForLogin(login, keyID)
+            return
+        }
+
+        try? await storage.storeAuthProxyKeyID?(keyID)
+    }
+
+    private func clearAuthProxyKeyID(for login: Login?) async {
+        if let login, let clearForLogin = storage.clearAuthProxyKeyIDForLogin {
+            try? await clearForLogin(login)
+            return
+        }
+
+        try? await storage.clearAuthProxyKeyID?()
     }
 
     private func buildTokenHandling(

Sources/CoreATProtocol/OAuth/AuthProxy.swift

···
     }
 }
 
+// MARK: - Auth Proxy Errors
+
+enum AuthProxyRequestKind: String, Sendable {
+    case pushedAuthorizationRequest
+    case tokenExchange
+    case refresh
+
+    var description: String {
+        switch self {
+        case .pushedAuthorizationRequest:
+            "PAR request"
+        case .tokenExchange:
+            "token exchange"
+        case .refresh:
+            "refresh request"
+        }
+    }
+}
+
+enum AuthProxyError: LocalizedError, Sendable {
+    case unavailable(requestKind: AuthProxyRequestKind, detail: String)
+
+    var errorDescription: String? {
+        switch self {
+        case .unavailable(let requestKind, let detail):
+            "Auth proxy unavailable during \(requestKind.description): \(detail)"
+        }
+    }
+}
+
 // MARK: - Key ID Storage
 
 /// Thread-safe storage for the auth proxy key ID, used during a single auth flow.
 actor AuthProxyKeyIDStorage {
     var keyID: String?
+    var prefersDirectRequests = false
 
     init(keyID: String? = nil) {
         self.keyID = keyID
···
     func update(_ newKeyID: String) {
         keyID = newKeyID
     }
+
+    func bypassProxy() {
+        prefersDirectRequests = true
+        keyID = nil
+    }
 }
 
 // MARK: - Proxy URL Loader
···
     tokenEndpoint: String,
     parEndpoint: String,
     issuer: String,
-    keyIDStorage: AuthProxyKeyIDStorage
+    keyIDStorage: AuthProxyKeyIDStorage,
+    fallbackProvider: @escaping URLResponseProvider = URLSession.defaultProvider
 ) -> URLResponseProvider {
     { request in
-        guard let requestURL = request.url?.absoluteString else {
-            return try await URLSession.shared.data(for: request)
+        if await keyIDStorage.prefersDirectRequests {
+            return try await fallbackProvider(request)
         }
 
-        let isTokenRequest = requestURL == tokenEndpoint
-        let isPARRequest = requestURL == parEndpoint
+        guard let requestURL = request.url?.absoluteString else {
+            return try await fallbackProvider(request)
+        }
 
-        guard isTokenRequest || isPARRequest else {
-            return try await URLSession.shared.data(for: request)
+        guard let requestKind = authProxyRequestKind(
+            requestURL: requestURL,
+            tokenEndpoint: tokenEndpoint,
+            parEndpoint: parEndpoint,
+            requestBody: request.httpBody
+        ) else {
+            return try await fallbackProvider(request)
         }
 
-        let proxyPath = isTokenRequest ? "/oauth/token" : "/oauth/par"
+        let proxyPath = requestKind == .pushedAuthorizationRequest ? "/oauth/par" : "/oauth/token"
         guard let proxyURL = URL(string: proxyBaseURL + proxyPath) else {
-            return try await URLSession.shared.data(for: request)
+            if requestKind == .pushedAuthorizationRequest {
+                await keyIDStorage.bypassProxy()
+                return try await fallbackProvider(request)
+            }
+
+            throw AuthProxyError.unavailable(
+                requestKind: requestKind,
+                detail: "Invalid auth proxy base URL: \(proxyBaseURL)"
+            )
         }
 
         var proxyRequest = URLRequest(url: proxyURL)
···
 
         let currentKeyID = await keyIDStorage.keyID
 
-        if isPARRequest {
+        if requestKind == .pushedAuthorizationRequest {
             // Convert form-encoded PAR body to JSON for the proxy
             let formParams = parseFormEncoded(request.httpBody)
             let parBody = AuthProxyPARRequest(
···
             proxyRequest.httpBody = try JSONSerialization.data(withJSONObject: bodyDict)
         }
 
-        let (data, response) = try await URLSession.shared.data(for: proxyRequest)
+        do {
+            let (data, response) = try await fallbackProvider(proxyRequest)
 
-        // Capture Auth-Proxy-Key-ID from every proxy response
-        if let httpResponse = response as? HTTPURLResponse,
-           let newKeyID = httpResponse.value(forHTTPHeaderField: "Auth-Proxy-Key-ID") {
-            await keyIDStorage.update(newKeyID)
-        }
+            if let httpResponse = response as? HTTPURLResponse {
+                if let newKeyID = httpResponse.value(forHTTPHeaderField: "Auth-Proxy-Key-ID") {
+                    await keyIDStorage.update(newKeyID)
+                }
 
-        return (data, response)
+                if shouldTreatProxyResponseAsUnavailable(httpResponse) {
+                    if requestKind == .pushedAuthorizationRequest {
+                        await keyIDStorage.bypassProxy()
+                        return try await fallbackProvider(request)
+                    }
+
+                    throw AuthProxyError.unavailable(
+                        requestKind: requestKind,
+                        detail: "Proxy returned HTTP \(httpResponse.statusCode)"
+                    )
+                }
+            }
+
+            return (data, response)
+        } catch let error as AuthProxyError {
+            throw error
+        } catch {
+            guard shouldBypassUnavailableAuthProxy(error) else {
+                throw error
+            }
+
+            if requestKind == .pushedAuthorizationRequest {
+                await keyIDStorage.bypassProxy()
+                return try await fallbackProvider(request)
+            }
+
+            throw AuthProxyError.unavailable(
+                requestKind: requestKind,
+                detail: error.localizedDescription
+            )
+        }
     }
 }
 
 // MARK: - Helpers
+
+func authProxyRequestKind(
+    requestURL: String,
+    tokenEndpoint: String,
+    parEndpoint: String,
+    requestBody: Data?
+) -> AuthProxyRequestKind? {
+    if requestURL == parEndpoint {
+        return .pushedAuthorizationRequest
+    }
+
+    guard requestURL == tokenEndpoint else {
+        return nil
+    }
+
+    guard let requestBody,
+          let json = (try? JSONSerialization.jsonObject(with: requestBody)) as? [String: Any],
+          let grantType = json["grant_type"] as? String else {
+        return .tokenExchange
+    }
+
+    return grantType == "refresh_token" ? .refresh : .tokenExchange
+}
+
+func shouldTreatProxyResponseAsUnavailable(_ response: HTTPURLResponse) -> Bool {
+    return switch response.statusCode {
+    case 404, 405, 502, 503, 504:
+        true
+    default:
+        false
+    }
+}
+
+func shouldBypassUnavailableAuthProxy(_ error: Error) -> Bool {
+    guard let urlError = error as? URLError else {
+        return false
+    }
+
+    return switch urlError.code {
+    case .cannotFindHost, .cannotConnectToHost, .dnsLookupFailed, .notConnectedToInternet, .resourceUnavailable:
+        true
+    default:
+        false
+    }
+}
 
 /// Parses a `application/x-www-form-urlencoded` body into a dictionary.
 func parseFormEncoded(_ data: Data?) -> [String: String] {

Tests/CoreATProtocolTests/OAuthTests.swift

···
         #expect(parseFormEncoded(Data()).isEmpty)
     }
 
+    @Test("Auth proxy request kind detects PAR, token exchange, and refresh")
+    func testAuthProxyRequestKind() throws {
+        let tokenExchangeBody = try JSONSerialization.data(withJSONObject: [
+            "grant_type": "authorization_code",
+            "code": "abc123"
+        ])
+        let refreshBody = try JSONSerialization.data(withJSONObject: [
+            "grant_type": "refresh_token",
+            "refresh_token": "refresh"
+        ])
+
+        #expect(
+            authProxyRequestKind(
+                requestURL: "https://bsky.social/oauth/par",
+                tokenEndpoint: "https://bsky.social/oauth/token",
+                parEndpoint: "https://bsky.social/oauth/par",
+                requestBody: nil
+            ) == .pushedAuthorizationRequest
+        )
+        #expect(
+            authProxyRequestKind(
+                requestURL: "https://bsky.social/oauth/token",
+                tokenEndpoint: "https://bsky.social/oauth/token",
+                parEndpoint: "https://bsky.social/oauth/par",
+                requestBody: tokenExchangeBody
+            ) == .tokenExchange
+        )
+        #expect(
+            authProxyRequestKind(
+                requestURL: "https://bsky.social/oauth/token",
+                tokenEndpoint: "https://bsky.social/oauth/token",
+                parEndpoint: "https://bsky.social/oauth/par",
+                requestBody: refreshBody
+            ) == .refresh
+        )
+    }
+
+    @Test("Auth proxy availability checks only retry for recoverable transport errors")
+    func testRecoverableProxyTransportErrors() {
+        #expect(shouldBypassUnavailableAuthProxy(URLError(.cannotConnectToHost)))
+        #expect(shouldBypassUnavailableAuthProxy(URLError(.dnsLookupFailed)))
+        #expect(!shouldBypassUnavailableAuthProxy(URLError(.timedOut)))
+        #expect(!shouldBypassUnavailableAuthProxy(URLError(.badServerResponse)))
+    }
+
+    @Test("Auth proxy availability checks only retry for recoverable HTTP statuses")
+    func testRecoverableProxyStatuses() throws {
+        let unavailableResponse = try #require(
+            HTTPURLResponse(
+                url: URL(string: "https://auth.example.com/oauth/token")!,
+                statusCode: 503,
+                httpVersion: nil,
+                headerFields: nil
+            )
+        )
+        let invalidRequestResponse = try #require(
+            HTTPURLResponse(
+                url: URL(string: "https://auth.example.com/oauth/token")!,
+                statusCode: 400,
+                httpVersion: nil,
+                headerFields: nil
+            )
+        )
+
+        #expect(shouldTreatProxyResponseAsUnavailable(unavailableResponse))
+        #expect(!shouldTreatProxyResponseAsUnavailable(invalidRequestResponse))
+    }
+
     @Test("AuthProxyKeyIDStorage stores and retrieves key ID")
     func testKeyIDStorage() async {
         let storage = AuthProxyKeyIDStorage()
···
         #expect(rotated == "key-2")
     }
 
+    @Test("AuthProxyKeyIDStorage can bypass proxy for remaining flow")
+    func testKeyIDStorageBypassesProxy() async {
+        let storage = AuthProxyKeyIDStorage(keyID: "existing-key")
+
+        #expect(await storage.prefersDirectRequests == false)
+        #expect(await storage.keyID == "existing-key")
+
+        await storage.bypassProxy()
+
+        #expect(await storage.prefersDirectRequests)
+        #expect(await storage.keyID == nil)
+    }
+
     @Test("AuthProxyKeyIDStorage initializes with existing key ID")
     func testKeyIDStorageWithInitialValue() async {
         let storage = AuthProxyKeyIDStorage(keyID: "existing-key")
···
         #expect(retrieved == "test-key")
     }
 
+    @Test("ATProtoAuthStorage accepts session-scoped auth proxy key ID closures")
+    func testStorageWithSessionScopedProxyKeyID() async throws {
+        actor SessionKeyIDStore {
+            var records: [String: String] = [:]
+
+            func retrieve(for login: Login) -> String? {
+                records[login.accessToken.value]
+            }
+
+            func store(login: Login, keyID: String) {
+                records[login.accessToken.value] = keyID
+            }
+
+            func clear(for login: Login) {
+                records.removeValue(forKey: login.accessToken.value)
+            }
+        }
+
+        let sessionStore = SessionKeyIDStore()
+        let login = Login(
+            accessToken: Token(value: "access-token"),
+            refreshToken: Token(value: "refresh-token")
+        )
+
+        let storage = ATProtoAuthStorage(
+            retrieveLogin: { nil },
+            storeLogin: { _ in },
+            retrievePrivateKey: { nil },
+            storePrivateKey: { _ in },
+            retrieveAuthProxyKeyIDForLogin: { login in
+                await sessionStore.retrieve(for: login)
+            },
+            storeAuthProxyKeyIDForLogin: { login, keyID in
+                await sessionStore.store(login: login, keyID: keyID)
+            },
+            clearAuthProxyKeyIDForLogin: { login in
+                await sessionStore.clear(for: login)
+            }
+        )
+
+        #expect(storage.retrieveAuthProxyKeyIDForLogin != nil)
+        #expect(storage.storeAuthProxyKeyIDForLogin != nil)
+        #expect(storage.clearAuthProxyKeyIDForLogin != nil)
+
+        let initial = try await storage.retrieveAuthProxyKeyIDForLogin?(login)
+        #expect(initial == nil)
+
+        try await storage.storeAuthProxyKeyIDForLogin?(login, "proxy-key")
+        let stored = try await storage.retrieveAuthProxyKeyIDForLogin?(login)
+        #expect(stored == "proxy-key")
+
+        try await storage.clearAuthProxyKeyIDForLogin?(login)
+        let cleared = try await storage.retrieveAuthProxyKeyIDForLogin?(login)
+        #expect(cleared == nil)
+    }
+
     @Test("ATProtoAuthStorage defaults proxy key ID closures to nil")
     func testStorageWithoutProxyKeyID() {
         let storage = ATProtoAuthStorage(
···
 
         #expect(storage.retrieveAuthProxyKeyID == nil)
         #expect(storage.storeAuthProxyKeyID == nil)
+        #expect(storage.clearAuthProxyKeyID == nil)
+        #expect(storage.retrieveAuthProxyKeyIDForLogin == nil)
+        #expect(storage.storeAuthProxyKeyIDForLogin == nil)
+        #expect(storage.clearAuthProxyKeyIDForLogin == nil)
     }
 }