[PM-29309] [BWA-209] Fix TOTP countdown freeze when returning to Authenticator app (change Flow to StateFlow)

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/manager/TotpCodeManager.kt

@@ -3,19 +3,21 @@ package com.bitwarden.authenticator.data.authenticator.manager
 import com.bitwarden.authenticator.data.authenticator.datasource.disk.entity.AuthenticatorItemAlgorithm
 import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
 import com.bitwarden.authenticator.data.authenticator.repository.model.AuthenticatorItem
-import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.StateFlow
 
 /**
  * Manages the flows for getting verification codes.
  */
 interface TotpCodeManager {
 
     /**
-     * Flow for getting a DataState with multiple verification code items.
+     * StateFlow for getting multiple verification code items. Returns a StateFlow that emits
+     * updated verification codes every second. The StateFlow is cached per-item to prevent
+     * recreation on each subscribe, ensuring smooth UI updates when returning from background.
      */
     fun getTotpCodesFlow(
         itemList: List<AuthenticatorItem>,
-    ): Flow<List<VerificationCodeItem>>
+    ): StateFlow<List<VerificationCodeItem>>
 
     @Suppress("UndocumentedPublicClass")
     companion object {

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/manager/TotpCodeManagerImpl.kt

@@ -4,13 +4,18 @@ import com.bitwarden.authenticator.data.authenticator.datasource.sdk.Authenticat
 import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
 import com.bitwarden.authenticator.data.authenticator.repository.model.AuthenticatorItem
 import com.bitwarden.core.DateTime
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.currentCoroutineContext
 import kotlinx.coroutines.delay
-import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.onCompletion
+import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.isActive
 import java.time.Clock
 import java.util.UUID
@@ -20,68 +25,124 @@ private const val ONE_SECOND_MILLISECOND = 1000L
 
 /**
  * Primary implementation of [TotpCodeManager].
+ *
+ * This implementation uses per-item [StateFlow] caching to prevent flow recreation on each
+ * subscribe, ensuring smooth UI updates when returning from background. The pattern mirrors
+ * the Password Manager's [com.x8bit.bitwarden.data.vault.manager.TotpCodeManagerImpl].
  */
 class TotpCodeManagerImpl @Inject constructor(
     private val authenticatorSdkSource: AuthenticatorSdkSource,
     private val clock: Clock,
+    private val dispatcherManager: DispatcherManager,
 ) : TotpCodeManager {
 
+    private val unconfinedScope = CoroutineScope(dispatcherManager.unconfined)
+
+    /**
+     * Cache of per-item StateFlows to prevent recreation on each subscribe.
+     * Key is the [AuthenticatorItem], value is the cached [StateFlow] for that item.
+     */
+    private val mutableItemVerificationCodeStateFlowMap =
+        mutableMapOf<AuthenticatorItem, StateFlow<VerificationCodeItem?>>()
+
     override fun getTotpCodesFlow(
         itemList: List<AuthenticatorItem>,
-    ): Flow<List<VerificationCodeItem>> {
+    ): StateFlow<List<VerificationCodeItem>> {
         if (itemList.isEmpty()) {
-            return flowOf(emptyList())
+            return MutableStateFlow(emptyList())
         }
-        val flows = itemList.map { it.toFlowOfVerificationCodes() }
-        return combine(flows) { it.toList() }
+
+        val stateFlows = itemList.map { getOrCreateItemStateFlow(it) }
+
+        return combine(stateFlows) { results ->
+            results.filterNotNull()
+        }
+            .stateIn(
+                scope = unconfinedScope,
+                started = SharingStarted.WhileSubscribed(),
+                initialValue = emptyList(),
+            )
     }
 
-    private fun AuthenticatorItem.toFlowOfVerificationCodes(): Flow<VerificationCodeItem> {
-        val otpUri = this.otpUri
-        return flow {
-            var item: VerificationCodeItem? = null
-            while (currentCoroutineContext().isActive) {
-                val time = (clock.millis() / ONE_SECOND_MILLISECOND).toInt()
-                if (item == null || item.isExpired(clock)) {
-                    // If the item is expired or we haven't generated our first item,
-                    // generate a new code using the SDK:
-                    item = authenticatorSdkSource
-                        .generateTotp(otpUri, DateTime.now())
-                        .getOrNull()
-                        ?.let { response ->
-                            VerificationCodeItem(
-                                code = response.code,
-                                periodSeconds = response.period.toInt(),
-                                timeLeftSeconds = response.period.toInt() -
-                                    time % response.period.toInt(),
-                                issueTime = clock.millis(),
-                                id = when (source) {
-                                    is AuthenticatorItem.Source.Local -> source.cipherId
-                                    is AuthenticatorItem.Source.Shared -> UUID.randomUUID()
-                                        .toString()
-                                },
-                                issuer = issuer,
-                                label = label,
-                                source = source,
-                            )
-                        }
-                        ?: run {
-                            // We are assuming that our otp URIs can generate a valid code.
-                            // If they can't, we'll just silently omit that code from the list.
-                            currentCoroutineContext().cancel()
-                            return@flow
-                        }
-                } else {
-                    // Item is not expired, just update time left:
-                    item = item.copy(
-                        timeLeftSeconds = item.periodSeconds - (time % item.periodSeconds),
-                    )
+    /**
+     * Gets an existing [StateFlow] for the given [item] or creates a new one if it doesn't exist.
+     * Each item gets its own [CoroutineScope] to manage its lifecycle independently.
+     */
+    private fun getOrCreateItemStateFlow(
+        item: AuthenticatorItem,
+    ): StateFlow<VerificationCodeItem?> {
+        return mutableItemVerificationCodeStateFlowMap.getOrPut(item) {
+            // Define a per-item scope so that we can clear the Flow from the map when it is
+            // no longer needed.
+            val itemScope = CoroutineScope(dispatcherManager.unconfined)
+
+            createVerificationCodeFlow(item)
+                .onCompletion {
+                    mutableItemVerificationCodeStateFlowMap.remove(item)
+                    itemScope.cancel()
                 }
-                // Emit item
-                emit(item)
-                // Wait one second before heading to the top of the loop:
-                delay(ONE_SECOND_MILLISECOND)
+                .stateIn(
+                    scope = itemScope,
+                    started = SharingStarted.Eagerly,
+                    initialValue = null,
+                )
+        }
+    }
+
+    /**
+     * Creates a flow that emits [VerificationCodeItem] updates every second for the given [item].
+     */
+    @Suppress("LongMethod")
+    private fun createVerificationCodeFlow(
+        item: AuthenticatorItem,
+    ) = flow<VerificationCodeItem?> {
+        val otpUri = item.otpUri
+        var verificationCodeItem: VerificationCodeItem? = null
+
+        while (currentCoroutineContext().isActive) {
+            val time = (clock.millis() / ONE_SECOND_MILLISECOND).toInt()
+
+            if (verificationCodeItem == null || verificationCodeItem.isExpired(clock)) {
+                // If the item is expired or we haven't generated our first item,
+                // generate a new code using the SDK:
+                verificationCodeItem = authenticatorSdkSource
+                    .generateTotp(otpUri, DateTime.now())
+                    .getOrNull()
+                    ?.let { response ->
+                        VerificationCodeItem(
+                            code = response.code,
+                            periodSeconds = response.period.toInt(),
+                            timeLeftSeconds = response.period.toInt() -
+                                time % response.period.toInt(),
+                            issueTime = clock.millis(),
+                            id = when (item.source) {
+                                is AuthenticatorItem.Source.Local -> item.source.cipherId
+                                is AuthenticatorItem.Source.Shared -> UUID.randomUUID().toString()
+                            },
+                            issuer = item.issuer,
+                            label = item.label,
+                            source = item.source,
+                        )
+                    }
+                    ?: run {
+                        // We are assuming that our otp URIs can generate a valid code.
+                        // If they can't, we'll just silently omit that code from the list.
+                        emit(null)
+                        return@flow
+                    }
+            } else {
+                // Item is not expired, just update time left:
+                verificationCodeItem = verificationCodeItem.copy(
+                    timeLeftSeconds = verificationCodeItem.periodSeconds -
+                        (time % verificationCodeItem.periodSeconds),
+                )
             }
+
+            // Emit item
+            emit(verificationCodeItem)
+
+            // Wait one second before heading to the top of the loop:
+            delay(ONE_SECOND_MILLISECOND)
         }
     }
 }

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/manager/di/AuthenticatorManagerModule.kt

@@ -41,5 +41,6 @@ object AuthenticatorManagerModule {
     ): TotpCodeManager = TotpCodeManagerImpl(
         authenticatorSdkSource = authenticatorSdkSource,
         clock = clock,
+        dispatcherManager = dispatcherManager,
     )
 }

authenticator/src/main/kotlin/com/bitwarden/authenticator/data/authenticator/repository/AuthenticatorRepositoryImpl.kt

@@ -157,7 +157,7 @@ class AuthenticatorRepositoryImpl @Inject constructor(
             .flatMapLatest { it.toSharedVerificationCodesStateFlow() }
             .stateIn(
                 scope = unconfinedScope,
-                started = SharingStarted.WhileSubscribed(STOP_TIMEOUT_DELAY_MS),
+                started = SharingStarted.WhileSubscribed(),
                 initialValue = SharedVerificationCodesState.Loading,
             )
     }
@@ -197,7 +197,7 @@ class AuthenticatorRepositoryImpl @Inject constructor(
             }
             .stateIn(
                 scope = unconfinedScope,
-                started = SharingStarted.WhileSubscribed(STOP_TIMEOUT_DELAY_MS),
+                started = SharingStarted.WhileSubscribed(),
                 initialValue = DataState.Loading,
             )
     }

authenticator/src/test/kotlin/com/bitwarden/authenticator/data/authenticator/manager/util/TotpCodeManagerTest.kt

@@ -4,6 +4,7 @@ import app.cash.turbine.test
 import com.bitwarden.authenticator.data.authenticator.datasource.sdk.AuthenticatorSdkSource
 import com.bitwarden.authenticator.data.authenticator.manager.TotpCodeManagerImpl
 import com.bitwarden.authenticator.data.authenticator.manager.model.VerificationCodeItem
+import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
 import io.mockk.mockk
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertEquals
@@ -19,18 +20,19 @@ class TotpCodeManagerTest {
         ZoneOffset.UTC,
     )
     private val authenticatorSdkSource: AuthenticatorSdkSource = mockk()
+    private val dispatcherManager = FakeDispatcherManager()
 
     private val manager = TotpCodeManagerImpl(
         authenticatorSdkSource = authenticatorSdkSource,
         clock = clock,
+        dispatcherManager = dispatcherManager,
     )
 
     @Test
-    fun `getTotpCodesFlow should return flow that emits empty list when input list is empty`() =
+    fun `getTotpCodesFlow should return StateFlow that emits empty list when input list is empty`() =
         runTest {
             manager.getTotpCodesFlow(emptyList()).test {
                 assertEquals(emptyList<VerificationCodeItem>(), awaitItem())
-                awaitComplete()
             }
         }
 }

authenticator/src/test/kotlin/com/bitwarden/authenticator/data/authenticator/repository/AuthenticatorRepositoryTest.kt

@@ -24,7 +24,6 @@ import io.mockk.runs
 import io.mockk.unmockkStatic
 import io.mockk.verify
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.AfterEach
 import org.junit.jupiter.api.Assertions.assertEquals
@@ -145,7 +144,7 @@ class AuthenticatorRepositoryTest {
             every { sharedAccounts.toAuthenticatorItems() } returns authenticatorItems
             every {
                 mockTotpCodeManager.getTotpCodesFlow(authenticatorItems)
-            } returns flowOf(verificationCodes)
+            } returns MutableStateFlow(verificationCodes)
             authenticatorRepository.sharedCodesStateFlow.test {
                 assertEquals(SharedVerificationCodesState.Loading, awaitItem())
                 mutableAccountSyncStateFlow.value = AccountSyncState.Success(sharedAccounts)