chore(auth): Create use cases for device management APIs (#2979)

Co-authored-by: Tyler Roach <[email protected]>

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/AWSCognitoAuthPlugin.kt

@@ -297,16 +297,16 @@ class AWSCognitoAuthPlugin : AuthPlugin<AWSCognitoAuthService>() {
         enqueue(onSuccess, onError) { queueFacade.fetchAuthSession() }
 
     override fun rememberDevice(onSuccess: Action, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.rememberDevice() }
+        enqueue(onSuccess, onError) { useCaseFactory.rememberDevice().execute() }
 
     override fun forgetDevice(onSuccess: Action, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.forgetDevice() }
+        enqueue(onSuccess, onError) { useCaseFactory.forgetDevice().execute() }
 
     override fun forgetDevice(device: AuthDevice, onSuccess: Action, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.forgetDevice(device) }
+        enqueue(onSuccess, onError) { useCaseFactory.forgetDevice().execute(device) }
 
     override fun fetchDevices(onSuccess: Consumer<List<AuthDevice>>, onError: Consumer<AuthException>) =
-        enqueue(onSuccess, onError) { queueFacade.fetchDevices() }
+        enqueue(onSuccess, onError) { useCaseFactory.fetchDevices().execute() }
 
     override fun resetPassword(
         username: String,

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/AuthStateMachine.kt

@@ -33,6 +33,7 @@ import com.amplifyframework.auth.cognito.actions.SignUpCognitoActions
 import com.amplifyframework.auth.cognito.actions.UserAuthSignInCognitoActions
 import com.amplifyframework.auth.cognito.actions.WebAuthnSignInCognitoActions
 import com.amplifyframework.auth.exceptions.InvalidStateException
+import com.amplifyframework.auth.exceptions.SignedOutException
 import com.amplifyframework.statemachine.Environment
 import com.amplifyframework.statemachine.StateMachine
 import com.amplifyframework.statemachine.StateMachineResolver
@@ -137,3 +138,12 @@ internal suspend inline fun <reified T : AuthenticationState> AuthStateMachine.r
         )
     }
 }
+
+// Returns the SignedInState or throws SignedOutException or InvalidStateException
+internal suspend fun AuthStateMachine.requireSignedInState(): AuthenticationState.SignedIn {
+    return when (val state = getCurrentState().authNState) {
+        is AuthenticationState.SignedIn -> state
+        is AuthenticationState.SignedOut -> throw SignedOutException()
+        else -> throw InvalidStateException()
+    }
+}

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/KotlinAuthFacadeInternal.kt

@@ -18,7 +18,6 @@ package com.amplifyframework.auth.cognito
 import android.app.Activity
 import android.content.Intent
 import com.amplifyframework.auth.AuthCodeDeliveryDetails
-import com.amplifyframework.auth.AuthDevice
 import com.amplifyframework.auth.AuthProvider
 import com.amplifyframework.auth.AuthSession
 import com.amplifyframework.auth.AuthUser
@@ -259,43 +258,6 @@ internal class KotlinAuthFacadeInternal(private val delegate: RealAWSCognitoAuth
         }
     }
 
-    suspend fun rememberDevice() {
-        return suspendCoroutine { continuation ->
-            delegate.rememberDevice(
-                { continuation.resume(Unit) },
-                { continuation.resumeWithException(it) }
-            )
-        }
-    }
-
-    suspend fun forgetDevice() {
-        return suspendCoroutine { continuation ->
-            delegate.forgetDevice(
-                { continuation.resume(Unit) },
-                { continuation.resumeWithException(it) }
-            )
-        }
-    }
-
-    suspend fun forgetDevice(device: AuthDevice) {
-        return suspendCoroutine { continuation ->
-            delegate.forgetDevice(
-                device,
-                { continuation.resume(Unit) },
-                { continuation.resumeWithException(it) }
-            )
-        }
-    }
-
-    suspend fun fetchDevices(): List<AuthDevice> {
-        return suspendCoroutine { continuation ->
-            delegate.fetchDevices(
-                { continuation.resume(it) },
-                { continuation.resumeWithException(it) }
-            )
-        }
-    }
-
     suspend fun resetPassword(
         username: String
     ): AuthResetPasswordResult {

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/RealAWSCognitoAuthPlugin.kt

@@ -25,15 +25,11 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.model.AnalyticsMetadataTy
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.AttributeType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChallengeNameType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.ChangePasswordRequest
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.DeviceRememberedStatusType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.EmailMfaSettingsType
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.ForgetDeviceRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.GetUserAttributeVerificationCodeRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.GetUserRequest
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.ListDevicesRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.SmsMfaSettingsType
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.SoftwareTokenMfaSettingsType
-import aws.sdk.kotlin.services.cognitoidentityprovider.model.UpdateDeviceStatusRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.UpdateUserAttributesRequest
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.UpdateUserAttributesResponse
 import aws.sdk.kotlin.services.cognitoidentityprovider.model.VerifySoftwareTokenResponseType
@@ -48,7 +44,6 @@ import com.amplifyframework.auth.AWSCredentials
 import com.amplifyframework.auth.AWSTemporaryCredentials
 import com.amplifyframework.auth.AuthChannelEventName
 import com.amplifyframework.auth.AuthCodeDeliveryDetails
-import com.amplifyframework.auth.AuthDevice
 import com.amplifyframework.auth.AuthException
 import com.amplifyframework.auth.AuthFactorType
 import com.amplifyframework.auth.AuthProvider
@@ -1497,135 +1492,6 @@ internal class RealAWSCognitoAuthPlugin(
         )
     }
 
-    fun rememberDevice(onSuccess: Action, onError: Consumer<AuthException>) {
-        authStateMachine.getCurrentState { authState ->
-            when (val state = authState.authNState) {
-                is AuthenticationState.SignedIn -> {
-                    GlobalScope.launch {
-                        updateDevice(
-                            authEnvironment.getDeviceMetadata(state.signedInData.username)?.deviceKey,
-                            DeviceRememberedStatusType.Remembered,
-                            onSuccess,
-                            onError
-                        )
-                    }
-                }
-                is AuthenticationState.SignedOut -> {
-                    onError.accept(SignedOutException())
-                }
-                else -> {
-                    onError.accept(InvalidStateException())
-                }
-            }
-        }
-    }
-
-    fun forgetDevice(onSuccess: Action, onError: Consumer<AuthException>) {
-        forgetDevice(AuthDevice.fromId(""), onSuccess, onError)
-    }
-
-    private fun updateDevice(
-        alternateDeviceId: String?,
-        rememberedStatusType: DeviceRememberedStatusType,
-        onSuccess: Action,
-        onError: Consumer<AuthException>
-    ) {
-        GlobalScope.async {
-            try {
-                val tokens = getSession().userPoolTokensResult
-                authEnvironment.cognitoAuthService.cognitoIdentityProviderClient?.updateDeviceStatus(
-                    UpdateDeviceStatusRequest.invoke {
-                        accessToken = tokens.value?.accessToken
-                        deviceKey = alternateDeviceId
-                        deviceRememberedStatus = rememberedStatusType
-                    }
-                )
-                onSuccess.call()
-            } catch (e: Exception) {
-                onError.accept(CognitoAuthExceptionConverter.lookup(e, "Update device ID failed."))
-            }
-        }
-    }
-
-    fun forgetDevice(device: AuthDevice, onSuccess: Action, onError: Consumer<AuthException>) {
-        authStateMachine.getCurrentState { authState ->
-            when (val authNState = authState.authNState) {
-                is AuthenticationState.SignedIn -> {
-                    GlobalScope.launch {
-                        try {
-                            if (device.id.isEmpty()) {
-                                val deviceKey = authEnvironment.getDeviceMetadata(authNState.signedInData.username)
-                                    ?.deviceKey
-                                forgetDevice(deviceKey)
-                            } else {
-                                forgetDevice(device.id)
-                            }
-                            onSuccess.call()
-                        } catch (e: Exception) {
-                            onError.accept(CognitoAuthExceptionConverter.lookup(e, "Failed to forget device."))
-                        }
-                    }
-                }
-                is AuthenticationState.SignedOut -> {
-                    onError.accept(SignedOutException())
-                }
-                else -> {
-                    onError.accept(InvalidStateException())
-                }
-            }
-        }
-    }
-
-    private suspend fun forgetDevice(alternateDeviceId: String?) {
-        val tokens = getSession().userPoolTokensResult
-        authEnvironment.cognitoAuthService.cognitoIdentityProviderClient?.forgetDevice(
-            ForgetDeviceRequest.invoke {
-                accessToken = tokens.value?.accessToken
-                deviceKey = alternateDeviceId
-            }
-        )
-    }
-
-    fun fetchDevices(onSuccess: Consumer<List<AuthDevice>>, onError: Consumer<AuthException>) {
-        authStateMachine.getCurrentState { authState ->
-            when (authState.authNState) {
-                is AuthenticationState.SignedIn -> {
-                    _fetchDevices(onSuccess, onError)
-                }
-                is AuthenticationState.SignedOut -> {
-                    onError.accept(SignedOutException())
-                }
-                else -> {
-                    onError.accept(InvalidStateException())
-                }
-            }
-        }
-    }
-
-    private fun _fetchDevices(onSuccess: Consumer<List<AuthDevice>>, onError: Consumer<AuthException>) {
-        GlobalScope.async {
-            try {
-                val tokens = getSession().userPoolTokensResult
-                val response =
-                    authEnvironment.cognitoAuthService.cognitoIdentityProviderClient?.listDevices(
-                        ListDevicesRequest.invoke {
-                            accessToken = tokens.value?.accessToken
-                        }
-                    )
-
-                val devices = response?.devices?.map { device ->
-                    val id = device.deviceKey ?: ""
-                    val name = device.deviceAttributes?.find { it.name == "device_name" }?.value
-                    AuthDevice.fromId(id, name)
-                } ?: emptyList()
-
-                onSuccess.accept(devices)
-            } catch (e: Exception) {
-                onError.accept(CognitoAuthExceptionConverter.lookup(e, "Fetch devices failed."))
-            }
-        }
-    }
-
     @OptIn(DelicateCoroutinesApi::class)
     fun resetPassword(
         username: String,

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/AssociateWebAuthnCredentialUseCase.kt

@@ -22,9 +22,8 @@ import aws.sdk.kotlin.services.cognitoidentityprovider.startWebAuthnRegistration
 import com.amplifyframework.auth.cognito.AuthStateMachine
 import com.amplifyframework.auth.cognito.helpers.WebAuthnHelper
 import com.amplifyframework.auth.cognito.helpers.authLogger
-import com.amplifyframework.auth.cognito.requireAuthenticationState
+import com.amplifyframework.auth.cognito.requireSignedInState
 import com.amplifyframework.auth.options.AuthAssociateWebAuthnCredentialsOptions
-import com.amplifyframework.statemachine.codegen.states.AuthenticationState.SignedIn
 import com.amplifyframework.statemachine.util.mask
 import com.amplifyframework.util.JsonDocument
 import com.amplifyframework.util.toJsonString
@@ -40,7 +39,7 @@ internal class AssociateWebAuthnCredentialUseCase(
     @Suppress("UNUSED_PARAMETER")
     suspend fun execute(callingActivity: Activity, options: AuthAssociateWebAuthnCredentialsOptions) {
         // User must be signed in to call this API
-        stateMachine.requireAuthenticationState<SignedIn>()
+        stateMachine.requireSignedInState()
 
         val accessToken = fetchAuthSession.execute().accessToken
 

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/AuthUseCaseFactory.kt

@@ -47,4 +47,24 @@ internal class AuthUseCaseFactory(
         fetchAuthSession = fetchAuthSession(),
         stateMachine = stateMachine
     )
+
+    fun rememberDevice() = RememberDeviceUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        stateMachine = stateMachine,
+        environment = authEnvironment
+    )
+
+    fun forgetDevice() = ForgetDeviceUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        stateMachine = stateMachine,
+        environment = authEnvironment
+    )
+
+    fun fetchDevices() = FetchDevicesUseCase(
+        client = authEnvironment.requireIdentityProviderClient(),
+        fetchAuthSession = fetchAuthSession(),
+        stateMachine = stateMachine
+    )
 }

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/DeleteWebAuthnCredentialUseCase.kt

@@ -18,9 +18,8 @@ package com.amplifyframework.auth.cognito.usecases
 import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
 import aws.sdk.kotlin.services.cognitoidentityprovider.deleteWebAuthnCredential
 import com.amplifyframework.auth.cognito.AuthStateMachine
-import com.amplifyframework.auth.cognito.requireAuthenticationState
+import com.amplifyframework.auth.cognito.requireSignedInState
 import com.amplifyframework.auth.options.AuthDeleteWebAuthnCredentialOptions
-import com.amplifyframework.statemachine.codegen.states.AuthenticationState.SignedIn
 
 internal class DeleteWebAuthnCredentialUseCase(
     private val client: CognitoIdentityProviderClient,
@@ -30,7 +29,7 @@ internal class DeleteWebAuthnCredentialUseCase(
     @Suppress("UNUSED_PARAMETER")
     suspend fun execute(credentialId: String, options: AuthDeleteWebAuthnCredentialOptions) {
         // User must be signed in to call this API
-        stateMachine.requireAuthenticationState<SignedIn>()
+        stateMachine.requireSignedInState()
 
         val accessToken = fetchAuthSession.execute().accessToken
         client.deleteWebAuthnCredential {

aws-auth-cognito/src/main/java/com/amplifyframework/auth/cognito/usecases/FetchDevicesUseCase.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package com.amplifyframework.auth.cognito.usecases
+
+import aws.sdk.kotlin.services.cognitoidentityprovider.CognitoIdentityProviderClient
+import aws.sdk.kotlin.services.cognitoidentityprovider.listDevices
+import com.amplifyframework.auth.AuthDevice
+import com.amplifyframework.auth.cognito.AuthStateMachine
+import com.amplifyframework.auth.cognito.requireSignedInState
+
+internal class FetchDevicesUseCase(
+    private val client: CognitoIdentityProviderClient,
+    private val fetchAuthSession: FetchAuthSessionUseCase,
+    private val stateMachine: AuthStateMachine
+) {
+    suspend fun execute(): List<AuthDevice> {
+        stateMachine.requireSignedInState()
+        val token = fetchAuthSession.execute().accessToken
+        val response = client.listDevices { accessToken = token }
+        return response.devices?.map { device ->
+            val id = device.deviceKey ?: ""
+            val name = device.deviceAttributes?.find { it.name == "device_name" }?.value
+            AuthDevice.fromId(id, name)
+        } ?: emptyList()
+    }
+}