Ensure that all regV5 viewmodels follow the right patterns.

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModel.kt

@@ -7,13 +7,17 @@ package org.signal.registration.screens.aepentry
 
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.update
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.core.util.logging.Log
 import org.signal.registration.RegistrationFlowEvent
+import org.signal.registration.screens.EventDrivenViewModel
 import org.signal.registration.screens.util.navigateBack
 
 class EnterAepForLocalBackupViewModel(
@@ -21,7 +25,7 @@ class EnterAepForLocalBackupViewModel(
   private val resultBus: ResultEventBus,
   private val resultKey: String,
   isPasswordManagerAvailable: Boolean = false
-) : ViewModel() {
+) : EventDrivenViewModel<EnterAepEvents>(TAG) {
 
   companion object {
     private val TAG = Log.tag(EnterAepForLocalBackupViewModel::class)
@@ -30,8 +34,13 @@ class EnterAepForLocalBackupViewModel(
   private val _state = MutableStateFlow(EnterAepState(isPasswordManagerAvailable = isPasswordManagerAvailable))
   val state: StateFlow<EnterAepState> = _state.asStateFlow()
 
-  fun onEvent(event: EnterAepEvents) {
-    Log.d(TAG, "[Event] $event")
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
+
+  override suspend fun processEvent(event: EnterAepEvents) {
     when (event) {
       is EnterAepEvents.BackupKeyChanged -> {
         _state.update { EnterAepScreenEventHandler.applyEvent(it, event) }

feature/registration/src/main/java/org/signal/registration/screens/aepentry/EnterAepForRemoteBackupPreRegistrationViewModel.kt

@@ -8,9 +8,12 @@ package org.signal.registration.screens.aepentry
 import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import org.signal.core.models.AccountEntropyPool
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
@@ -33,9 +36,14 @@ class EnterAepForRemoteBackupPreRegistrationViewModel(
   }
 
   private val _state = MutableStateFlow(EnterAepState(isPasswordManagerAvailable = isPasswordManagerAvailable))
-
   val state: StateFlow<EnterAepState> = _state.asStateFlow()
 
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
+
   override suspend fun processEvent(event: EnterAepEvents) {
     applyEvent(_state.value, event) { _state.value = it }
   }

feature/registration/src/main/java/org/signal/registration/screens/countrycode/CountryCodePickerViewModel.kt

@@ -11,11 +11,14 @@ import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.core.util.logging.Log
 import org.signal.registration.RegistrationFlowEvent
+import org.signal.registration.screens.EventDrivenViewModel
 import org.signal.registration.screens.util.navigateBack
 
 /**
@@ -28,7 +31,7 @@ class CountryCodePickerViewModel(
   private val resultBus: ResultEventBus,
   private val resultKey: String,
   initialCountry: Country? = null
-) : ViewModel() {
+) : EventDrivenViewModel<CountryCodePickerScreenEvents>(TAG) {
 
   companion object {
     private val TAG = Log.tag(CountryCodePickerViewModel::class)
@@ -38,11 +41,14 @@ class CountryCodePickerViewModel(
   val state: StateFlow<CountryCodeState> = _state.asStateFlow()
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     loadCountries(initialCountry)
   }
 
-  fun onEvent(event: CountryCodePickerScreenEvents) {
-    Log.d(TAG, "[Event] $event")
+  override suspend fun processEvent(event: CountryCodePickerScreenEvents) {
     when (event) {
       is CountryCodePickerScreenEvents.Search -> applySearchEvent(event.query)
       is CountryCodePickerScreenEvents.CountrySelected -> {

feature/registration/src/main/java/org/signal/registration/screens/devicetransfer/complete/DeviceTransferCompleteViewModel.kt

@@ -8,8 +8,11 @@ package org.signal.registration.screens.devicetransfer.complete
 import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import org.signal.core.util.logging.Log
 import org.signal.registration.RegistrationFlowEvent
 import org.signal.registration.RegistrationRepository
@@ -28,6 +31,12 @@ class DeviceTransferCompleteViewModel(
   private val _state = MutableStateFlow(DeviceTransferCompleteState())
   val state: StateFlow<DeviceTransferCompleteState> = _state
 
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
+
   override suspend fun processEvent(event: DeviceTransferCompleteScreenEvents) {
     applyEvent(state.value, event, parentEventEmitter, repository) { _state.value = it }
   }

feature/registration/src/main/java/org/signal/registration/screens/devicetransfer/instructions/DeviceTransferInstructionsViewModel.kt

@@ -8,8 +8,11 @@ package org.signal.registration.screens.devicetransfer.instructions
 import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import org.signal.core.util.logging.Log
 import org.signal.registration.RegistrationFlowEvent
 import org.signal.registration.RegistrationRoute
@@ -28,6 +31,12 @@ class DeviceTransferInstructionsViewModel(
   private val _state = MutableStateFlow(DeviceTransferInstructionsState())
   val state: StateFlow<DeviceTransferInstructionsState> = _state
 
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
+
   override suspend fun processEvent(event: DeviceTransferInstructionsScreenEvents) {
     applyEvent(state.value, event, parentEventEmitter) { _state.value = it }
   }

feature/registration/src/main/java/org/signal/registration/screens/devicetransfer/progress/DeviceTransferProgressViewModel.kt

@@ -17,6 +17,8 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import org.greenrobot.eventbus.EventBus
 import org.greenrobot.eventbus.Subscribe
@@ -60,6 +62,10 @@ class DeviceTransferProgressViewModel(
   val showCancelDialog: StateFlow<Boolean> = _showCancelDialog
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     viewModelScope.launch {
       progressEvents.collect { handleProgressEvent(it) }
     }

feature/registration/src/main/java/org/signal/registration/screens/devicetransfer/setup/DeviceTransferSetupViewModel.kt

@@ -24,6 +24,8 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.flowOn
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import org.greenrobot.eventbus.EventBus
 import org.greenrobot.eventbus.Subscribe
@@ -76,6 +78,10 @@ class DeviceTransferSetupViewModel(
   private var shutdown: Boolean = false
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     subscribeToSetupEvents()
     onEvent(DeviceTransferSetupScreenEvents.CheckPermissions)
   }

feature/registration/src/main/java/org/signal/registration/screens/discoverability/PhoneNumberDiscoverabilityViewModel.kt

@@ -8,8 +8,11 @@ package org.signal.registration.screens.discoverability
 import androidx.annotation.VisibleForTesting
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
+import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.core.util.logging.Log
 import org.signal.registration.RegistrationFlowEvent
@@ -30,6 +33,12 @@ class PhoneNumberDiscoverabilityViewModel(
   private val _state = MutableStateFlow(PhoneNumberDiscoverabilityState(discoverable = initialDiscoverable))
   val state: StateFlow<PhoneNumberDiscoverabilityState> = _state
 
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
+
   override suspend fun processEvent(event: PhoneNumberDiscoverabilityScreenEvents) {
     applyEvent(state.value, event, parentEventEmitter, resultBus, resultKey) { _state.value = it }
   }

feature/registration/src/main/java/org/signal/registration/screens/linkaccount/LinkAccountViewModel.kt

@@ -11,10 +11,10 @@ import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import org.signal.core.ui.compose.QrCodeData
@@ -47,13 +47,15 @@ class LinkAccountViewModel(
   }
 
   private val _state = MutableStateFlow(LinkAccountScreenState(showCreateAccount = showCreateAccount))
-  val state: StateFlow<LinkAccountScreenState> = _state
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), _state.value)
+  val state: StateFlow<LinkAccountScreenState> = _state.asStateFlow()
 
   private var provisioningJob: Job? = null
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     startProvisioning()
   }
 

feature/registration/src/main/java/org/signal/registration/screens/localbackuprestore/LocalBackupRestoreViewModel.kt

@@ -12,9 +12,9 @@ import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import org.signal.archive.LocalBackupRestoreProgress
 import org.signal.core.models.AccountEntropyPool
@@ -40,15 +40,19 @@ class LocalBackupRestoreViewModel(
     private val TAG = Log.tag(LocalBackupRestoreViewModel::class)
   }
 
-  private val _localState = MutableStateFlow(LocalBackupRestoreState())
-  val state = _localState
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), LocalBackupRestoreState())
+  private val _state = MutableStateFlow(LocalBackupRestoreState())
+  val state = _state.asStateFlow()
 
   private var restoreJob: Job? = null
 
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
+
   override suspend fun processEvent(event: LocalBackupRestoreEvents) {
-    applyEvent(state.value, event) { _localState.value = it }
+    applyEvent(state.value, event) { _state.value = it }
   }
 
   @VisibleForTesting
@@ -139,21 +143,21 @@ class LocalBackupRestoreViewModel(
         val backups = repository.scanLocalBackupFolder(uri)
         val mostRecent = backups.firstOrNull()
         if (mostRecent != null) {
-          _localState.value = LocalBackupRestoreState(
+          _state.value = LocalBackupRestoreState(
             restorePhase = LocalBackupRestoreState.RestorePhase.BackupFound,
             backupInfo = mostRecent,
             allBackups = backups,
             selectedFolderUri = uri
           )
         } else {
-          _localState.value = LocalBackupRestoreState(
+          _state.value = LocalBackupRestoreState(
             restorePhase = LocalBackupRestoreState.RestorePhase.NoBackupFound,
             selectedFolderUri = uri
           )
         }
       } catch (e: Exception) {
         Log.w(TAG, "Error scanning backup folder", e)
-        _localState.value = LocalBackupRestoreState(
+        _state.value = LocalBackupRestoreState(
           restorePhase = LocalBackupRestoreState.RestorePhase.Error,
           errorMessage = e.message
         )
@@ -164,13 +168,13 @@ class LocalBackupRestoreViewModel(
   private fun startRestore(backup: LocalBackupInfo, rootUri: Uri?, credential: String, aep: AccountEntropyPool?) {
     restoreJob?.cancel()
     restoreJob = viewModelScope.launch {
-      val currentState = _localState.value
+      val currentState = _state.value
       val restoreFlow = when (backup.type) {
         LocalBackupInfo.BackupType.V1 -> repository.restoreV1Backup(backup.uri, passphrase = credential)
         LocalBackupInfo.BackupType.V2 -> repository.restoreV2Backup(rootUri = rootUri!!, backupUri = backup.uri, aep = aep!!)
       }
       restoreFlow.collect { progress ->
-        _localState.value = when (progress) {
+        _state.value = when (progress) {
           is LocalBackupRestoreProgress.Preparing -> LocalBackupRestoreState(
             restorePhase = LocalBackupRestoreState.RestorePhase.Preparing,
             aep = currentState.aep,
@@ -183,8 +187,8 @@ class LocalBackupRestoreViewModel(
             v1Passphrase = currentState.v1Passphrase
           )
           is LocalBackupRestoreProgress.Complete -> {
-            onRestoreComplete(_localState.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase), progress)
-            _localState.value
+            onRestoreComplete(_state.value.copy(aep = currentState.aep, v1Passphrase = currentState.v1Passphrase), progress)
+            _state.value
           }
           is LocalBackupRestoreProgress.IncorrectCredential -> {
             Log.w(TAG, "Restore failed: incorrect passphrase/recovery key")

feature/registration/src/main/java/org/signal/registration/screens/messagesync/MessageSyncViewModel.kt

@@ -12,10 +12,10 @@ import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import org.signal.core.util.bytes
@@ -42,14 +42,16 @@ class MessageSyncViewModel(
   }
 
   private val _state = MutableStateFlow(MessageSyncScreenState())
-  val state: StateFlow<MessageSyncScreenState> = _state
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), MessageSyncScreenState())
+  val state: StateFlow<MessageSyncScreenState> = _state.asStateFlow()
 
   private var restoreJob: Job? = null
   private var finishJob: Job? = null
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     startRestore()
   }
 

feature/registration/src/main/java/org/signal/registration/screens/pincreation/PinCreationScreenEvents.kt

@@ -5,7 +5,12 @@
 
 package org.signal.registration.screens.pincreation
 
+import org.signal.registration.RegistrationFlowState
+
 sealed class PinCreationScreenEvents {
+  /** The parent registration flow state changed and needs to be merged into this screen's state. */
+  data class ParentStateChanged(val parentState: RegistrationFlowState) : PinCreationScreenEvents()
+
   data class PinSubmitted(val pin: String) : PinCreationScreenEvents() {
     override fun toString(): String = "PinSubmitted(pin=${pin.length} chars)"
   }

feature/registration/src/main/java/org/signal/registration/screens/pincreation/PinCreationViewModel.kt

@@ -10,11 +10,10 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
 import org.signal.registration.NetworkController
@@ -32,7 +31,7 @@ import kotlin.time.toKotlinDuration
  */
 class PinCreationViewModel(
   private val repository: RegistrationRepository,
-  private val parentState: StateFlow<RegistrationFlowState>,
+  parentState: StateFlow<RegistrationFlowState>,
   private val parentEventEmitter: (RegistrationFlowEvent) -> Unit
 ) : EventDrivenViewModel<PinCreationScreenEvents>(TAG) {
 
@@ -41,19 +40,28 @@ class PinCreationViewModel(
   }
 
   private val _state = MutableStateFlow(PinCreationState())
+  val state: StateFlow<PinCreationState> = _state.asStateFlow()
 
-  val state: StateFlow<PinCreationState> = _state
-    .combine(parentState) { state, parentState -> applyParentState(state, parentState) }
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), PinCreationState())
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
+    parentState
+      .onEach { onEvent(PinCreationScreenEvents.ParentStateChanged(it)) }
+      .launchIn(viewModelScope)
+  }
 
   override suspend fun processEvent(event: PinCreationScreenEvents) {
-    applyEvent(state.value, event)
+    applyEvent(_state.value, event)
   }
 
   @VisibleForTesting
   suspend fun applyEvent(state: PinCreationState, event: PinCreationScreenEvents) {
     when (event) {
+      is PinCreationScreenEvents.ParentStateChanged -> {
+        _state.value = applyParentState(state, event.parentState)
+      }
       is PinCreationScreenEvents.PinSubmitted -> {
         when {
           !state.isConfirmEnabled -> {
@@ -105,8 +113,7 @@ class PinCreationViewModel(
     parentEventEmitter(RegistrationFlowEvent.RegistrationComplete)
   }
 
-  @VisibleForTesting
-  fun applyParentState(state: PinCreationState, parentState: RegistrationFlowState): PinCreationState {
+  private fun applyParentState(state: PinCreationState, parentState: RegistrationFlowState): PinCreationState {
     return state.copy(accountEntropyPool = parentState.accountEntropyPool)
   }
 

feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForRegistrationLockViewModel.kt

@@ -10,10 +10,10 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import org.signal.core.models.MasterKey
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
@@ -50,9 +50,13 @@ class PinEntryForRegistrationLockViewModel(
     )
   )
 
-  val state: StateFlow<PinEntryState> = _state
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), PinEntryState(showNeedHelp = true))
+  val state: StateFlow<PinEntryState> = _state.asStateFlow()
+
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
 
   override suspend fun processEvent(event: PinEntryScreenEvents) {
     applyEvent(state.value, event, parentEventEmitter) { _state.value = it }
@@ -70,7 +74,8 @@ class PinEntryForRegistrationLockViewModel(
         throw NotImplementedError("Skip is not a valid action during registration lock PIN entry")
       }
       is PinEntryScreenEvents.CreateNewPin,
-      is PinEntryScreenEvents.ContactSupport -> Unit
+      is PinEntryScreenEvents.ContactSupport,
+      is PinEntryScreenEvents.ParentStateChanged -> Unit
       is PinEntryScreenEvents.ToggleKeyboard -> {
         stateEmitter(PinEntryScreenEventHandler.applyEvent(state, event))
       }

feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForSmsBypassViewModel.kt

@@ -10,11 +10,10 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import org.signal.core.models.MasterKey
 import org.signal.core.util.Hex
 import org.signal.core.util.logging.Log
@@ -50,13 +49,20 @@ class PinEntryForSmsBypassViewModel(
     )
   )
 
-  val state: StateFlow<PinEntryState> = _state
-    .combine(parentState) { state, parentState -> applyParentState(state, parentState) }
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), PinEntryState(showNeedHelp = true))
+  val state: StateFlow<PinEntryState> = _state.asStateFlow()
+
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
+    parentState
+      .onEach { onEvent(PinEntryScreenEvents.ParentStateChanged(it)) }
+      .launchIn(viewModelScope)
+  }
 
   override suspend fun processEvent(event: PinEntryScreenEvents) {
-    applyEvent(state.value, event, parentEventEmitter) { _state.value = it }
+    applyEvent(_state.value, event, parentEventEmitter) { _state.value = it }
   }
 
   @VisibleForTesting
@@ -67,6 +73,9 @@ class PinEntryForSmsBypassViewModel(
     stateEmitter: (PinEntryState) -> Unit
   ) {
     when (event) {
+      is PinEntryScreenEvents.ParentStateChanged -> {
+        stateEmitter(applyParentState(state, event.parentState))
+      }
       is PinEntryScreenEvents.PinEntered -> {
         val localState = state.copy(loading = true)
         stateEmitter(localState)
@@ -83,7 +92,7 @@ class PinEntryForSmsBypassViewModel(
     }
   }
 
-  fun applyParentState(state: PinEntryState, parentState: RegistrationFlowState): PinEntryState {
+  private fun applyParentState(state: PinEntryState, parentState: RegistrationFlowState): PinEntryState {
     return state.copy(e164 = parentState.sessionE164)
   }
 

feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryForSvrRestoreViewModel.kt

@@ -10,10 +10,10 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
 import org.signal.registration.NetworkController
@@ -47,9 +47,13 @@ class PinEntryForSvrRestoreViewModel(
     )
   )
 
-  val state: StateFlow<PinEntryState> = _state
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), PinEntryState(showNeedHelp = true))
+  val state: StateFlow<PinEntryState> = _state.asStateFlow()
+
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+  }
 
   override suspend fun processEvent(event: PinEntryScreenEvents) {
     applyEvent(state.value, event, parentEventEmitter) { _state.value = it }
@@ -83,6 +87,7 @@ class PinEntryForSvrRestoreViewModel(
       is PinEntryScreenEvents.ToggleKeyboard -> {
         stateEmitter(PinEntryScreenEventHandler.applyEvent(state, event))
       }
+      is PinEntryScreenEvents.ParentStateChanged -> Unit
     }
   }
 

feature/registration/src/main/java/org/signal/registration/screens/pinentry/PinEntryScreenEvents.kt

@@ -5,7 +5,12 @@
 
 package org.signal.registration.screens.pinentry
 
+import org.signal.registration.RegistrationFlowState
+
 sealed class PinEntryScreenEvents {
+  /** The parent registration flow state changed and needs to be merged into this screen's state. */
+  data class ParentStateChanged(val parentState: RegistrationFlowState) : PinEntryScreenEvents()
+
   data class PinEntered(val pin: String) : PinEntryScreenEvents() {
     override fun toString(): String = "PinEntered(pin=${pin.length} chars)"
   }

feature/registration/src/main/java/org/signal/registration/screens/quickrestore/QuickRestoreQrViewModel.kt

@@ -13,6 +13,8 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import org.signal.core.ui.compose.QrCodeData
 import org.signal.core.util.logging.Log
@@ -34,17 +36,21 @@ class QuickRestoreQrViewModel(
     private val TAG = Log.tag(QuickRestoreQrViewModel::class)
   }
 
-  private val _localState = MutableStateFlow(QuickRestoreQrState())
-  val state: StateFlow<QuickRestoreQrState> = _localState.asStateFlow()
+  private val _state = MutableStateFlow(QuickRestoreQrState())
+  val state: StateFlow<QuickRestoreQrState> = _state.asStateFlow()
 
   private var provisioningJob: Job? = null
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     startProvisioning()
   }
 
   override suspend fun processEvent(event: QuickRestoreQrEvents) {
-    applyEvent(state.value, event) { _localState.value = it }
+    applyEvent(state.value, event) { _state.value = it }
   }
 
   @VisibleForTesting
@@ -73,7 +79,7 @@ class QuickRestoreQrViewModel(
         when (event) {
           is NetworkController.ProvisioningEvent.QrCodeReady -> {
             Log.d(TAG, "[Provisioning] QR code ready")
-            _localState.value = _localState.value.copy(
+            _state.value = _state.value.copy(
               qrState = QrState.Loaded(
                 qrCodeData = QrCodeData.forData(data = event.url, supportIconOverlay = false)
               )
@@ -85,7 +91,7 @@ class QuickRestoreQrViewModel(
           }
           is NetworkController.ProvisioningEvent.Error -> {
             Log.w(TAG, "[Provisioning] Error", event.cause)
-            _localState.value = _localState.value.copy(qrState = QrState.Failed)
+            _state.value = _state.value.copy(qrState = QrState.Failed)
           }
         }
       }
@@ -101,7 +107,7 @@ class QuickRestoreQrViewModel(
       return
     }
 
-    _localState.value = _localState.value.copy(isRegistering = true, qrState = QrState.Scanned)
+    _state.value = _state.value.copy(isRegistering = true, qrState = QrState.Scanned)
 
     val registerResult = repository.registerAccountWithProvisioningData(message)
 
@@ -116,7 +122,7 @@ class QuickRestoreQrViewModel(
         when (val error = registerResult.error) {
           is NetworkController.RegisterAccountError.RateLimited -> {
             Log.w(TAG, "[Register] Rate limited (retryAfter: ${error.retryAfter}).")
-            _localState.value = _localState.value.copy(
+            _state.value = _state.value.copy(
               isRegistering = false,
               showRegistrationError = true,
               errorMessage = null
@@ -124,7 +130,7 @@ class QuickRestoreQrViewModel(
           }
           is NetworkController.RegisterAccountError.RegistrationRecoveryPasswordIncorrect -> {
             Log.w(TAG, "[Register] Recovery password incorrect: ${error.message}")
-            _localState.value = _localState.value.copy(
+            _state.value = _state.value.copy(
               isRegistering = false,
               showRegistrationError = true,
               errorMessage = null
@@ -141,7 +147,7 @@ class QuickRestoreQrViewModel(
           }
           is NetworkController.RegisterAccountError.SessionNotFoundOrNotVerified -> {
             Log.w(TAG, "[Register] Session not found or not verified: ${error.message}")
-            _localState.value = _localState.value.copy(
+            _state.value = _state.value.copy(
               isRegistering = false,
               showRegistrationError = true,
               errorMessage = null
@@ -153,7 +159,7 @@ class QuickRestoreQrViewModel(
           }
           is NetworkController.RegisterAccountError.InvalidRequest -> {
             Log.w(TAG, "[Register] Invalid request: ${error.message}")
-            _localState.value = _localState.value.copy(
+            _state.value = _state.value.copy(
               isRegistering = false,
               showRegistrationError = true,
               errorMessage = null
@@ -163,7 +169,7 @@ class QuickRestoreQrViewModel(
       }
       is RequestResult.RetryableNetworkError -> {
         Log.w(TAG, "[Register] Network error.", registerResult.networkError)
-        _localState.value = _localState.value.copy(
+        _state.value = _state.value.copy(
           isRegistering = false,
           showRegistrationError = true,
           errorMessage = null
@@ -171,7 +177,7 @@ class QuickRestoreQrViewModel(
       }
       is RequestResult.ApplicationError -> {
         Log.w(TAG, "[Register] Application error.", registerResult.cause)
-        _localState.value = _localState.value.copy(
+        _state.value = _state.value.copy(
           isRegistering = false,
           showRegistrationError = true,
           errorMessage = null

feature/registration/src/main/java/org/signal/registration/screens/remotebackuprestore/RemoteBackupRestoreViewModel.kt

@@ -11,10 +11,10 @@ import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import org.signal.core.models.AccountEntropyPool
@@ -40,12 +40,13 @@ class RemoteBackupRestoreViewModel(
   }
 
   private val _state = MutableStateFlow(RemoteBackupRestoreState(aep))
-
-  val state: StateFlow<RemoteBackupRestoreState> = _state
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), RemoteBackupRestoreState(aep))
+  val state: StateFlow<RemoteBackupRestoreState> = _state.asStateFlow()
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
     loadBackupInfo()
   }
 

feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionScreenEvents.kt

@@ -5,7 +5,12 @@
 
 package org.signal.registration.screens.restoreselection
 
+import org.signal.registration.RegistrationFlowState
+
 sealed class ArchiveRestoreSelectionScreenEvents {
+  /** The parent registration flow state changed and needs to be merged into this screen's state. */
+  data class ParentStateChanged(val parentState: RegistrationFlowState) : ArchiveRestoreSelectionScreenEvents()
+
   data class RestoreOptionSelected(val option: ArchiveRestoreOption) : ArchiveRestoreSelectionScreenEvents()
 
   data object ConfirmSkip : ArchiveRestoreSelectionScreenEvents()

feature/registration/src/main/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModel.kt

@@ -10,10 +10,10 @@ import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
@@ -36,7 +36,7 @@ class ArchiveRestoreSelectionViewModel(
   private val restoreOptions: List<ArchiveRestoreOption>,
   private val registeredState: RegisteredState,
   private val repository: RegistrationRepository,
-  private val parentState: StateFlow<RegistrationFlowState>,
+  parentState: StateFlow<RegistrationFlowState>,
   private val parentEventEmitter: (RegistrationFlowEvent) -> Unit
 ) : EventDrivenViewModel<ArchiveRestoreSelectionScreenEvents>(TAG) {
 
@@ -44,32 +44,37 @@ class ArchiveRestoreSelectionViewModel(
     private val TAG = Log.tag(ArchiveRestoreSelectionViewModel::class)
   }
 
-  private val _localState = MutableStateFlow(
+  private val _state = MutableStateFlow(
     ArchiveRestoreSelectionState(
       restoreOptions = restoreOptions
     )
   )
+  val state: StateFlow<ArchiveRestoreSelectionState> = _state.asStateFlow()
 
-  val state: StateFlow<ArchiveRestoreSelectionState> = _localState
-    .combine(parentState) { state, parentState -> applyParentState(state, parentState) }
-    .stateIn(
-      viewModelScope,
-      SharingStarted.WhileSubscribed(5000),
-      ArchiveRestoreSelectionState(restoreOptions = restoreOptions)
-    )
+  init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
 
-  override suspend fun processEvent(event: ArchiveRestoreSelectionScreenEvents) {
-    applyEvent(state.value, event) { _localState.value = it }
+    parentState
+      .onEach { onEvent(ArchiveRestoreSelectionScreenEvents.ParentStateChanged(it)) }
+      .launchIn(viewModelScope)
   }
 
-  @VisibleForTesting
-  fun applyParentState(state: ArchiveRestoreSelectionState, parentState: RegistrationFlowState): ArchiveRestoreSelectionState {
+  override suspend fun processEvent(event: ArchiveRestoreSelectionScreenEvents) {
+    applyEvent(_state.value, event) { _state.value = it }
+  }
+
+  private fun applyParentState(state: ArchiveRestoreSelectionState, parentState: RegistrationFlowState): ArchiveRestoreSelectionState {
     return state.copy(restoreMethodToken = parentState.restoreMethodToken, storageCapable = parentState.storageCapable)
   }
 
   @VisibleForTesting
   suspend fun applyEvent(state: ArchiveRestoreSelectionState, event: ArchiveRestoreSelectionScreenEvents, stateEmitter: (ArchiveRestoreSelectionState) -> Unit) {
     val result = when (event) {
+      is ArchiveRestoreSelectionScreenEvents.ParentStateChanged -> {
+        applyParentState(state, event.parentState)
+      }
       is ArchiveRestoreSelectionScreenEvents.RestoreOptionSelected -> {
         when (event.option) {
           ArchiveRestoreOption.SignalSecureBackup -> {

feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeScreenEvents.kt

@@ -6,8 +6,12 @@
 package org.signal.registration.screens.verificationcode
 
 import org.signal.core.util.censor
+import org.signal.registration.RegistrationFlowState
 
 sealed class VerificationCodeScreenEvents {
+  /** The parent registration flow state changed and needs to be merged into this screen's state. */
+  data class ParentStateChanged(val parentState: RegistrationFlowState) : VerificationCodeScreenEvents()
+
   data class CodeEntered(val code: String) : VerificationCodeScreenEvents() {
     override fun toString(): String = "CodeEntered(code=${code.censor()})"
   }

feature/registration/src/main/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModel.kt

@@ -20,13 +20,12 @@ import com.google.android.gms.common.api.Status
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.callbackFlow
-import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.emptyFlow
+import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.flow.stateIn
 import kotlinx.coroutines.launch
 import org.signal.core.util.logging.Log
 import org.signal.libsignal.net.RequestResult
@@ -97,15 +96,21 @@ class VerificationCodeViewModel(
     }
   }
 
-  private val _localState = MutableStateFlow(VerificationCodeState())
-  val state = combine(_localState, parentState) { state, parentState -> applyParentState(state, parentState) }
-    .onEach { Log.d(TAG, "[State] $it") }
-    .stateIn(viewModelScope, SharingStarted.WhileSubscribed(5000), VerificationCodeState())
+  private val _state = MutableStateFlow(VerificationCodeState())
+  val state: StateFlow<VerificationCodeState> = _state.asStateFlow()
 
   private var nextSmsAvailableAt: Duration = 0.seconds
   private var nextCallAvailableAt: Duration = 0.seconds
 
   init {
+    _state
+      .onEach { Log.d(TAG, "[State] $it") }
+      .launchIn(viewModelScope)
+
+    parentState
+      .onEach { onEvent(VerificationCodeScreenEvents.ParentStateChanged(it)) }
+      .launchIn(viewModelScope)
+
     viewModelScope.launch {
       smsCodeEvents.collect { code ->
         onEvent(VerificationCodeScreenEvents.CodeAutoFilled(code))
@@ -114,12 +119,13 @@ class VerificationCodeViewModel(
   }
 
   override suspend fun processEvent(event: VerificationCodeScreenEvents) {
-    applyEvent(state.value, event) { _localState.value = it }
+    applyEvent(_state.value, event) { _state.value = it }
   }
 
   @VisibleForTesting
   suspend fun applyEvent(state: VerificationCodeState, event: VerificationCodeScreenEvents, stateEmitter: (VerificationCodeState) -> Unit) {
     val result = when (event) {
+      is VerificationCodeScreenEvents.ParentStateChanged -> applyParentState(state, event.parentState)
       is VerificationCodeScreenEvents.CodeEntered -> submitCode(state, event.code, stateEmitter)
       is VerificationCodeScreenEvents.DigitChanged -> applyDigitChanged(state, event.index, event.value, stateEmitter)
       is VerificationCodeScreenEvents.CodeAutoFilled -> state.copy(autoFillCode = event.code)
@@ -152,8 +158,7 @@ class VerificationCodeViewModel(
     return state
   }
 
-  @VisibleForTesting
-  fun applyParentState(state: VerificationCodeState, parentState: RegistrationFlowState): VerificationCodeState {
+  private fun applyParentState(state: VerificationCodeState, parentState: RegistrationFlowState): VerificationCodeState {
     if (parentState.sessionMetadata == null || parentState.sessionE164 == null) {
       Log.w(TAG, "Parent state is missing session metadata or e164! Resetting.")
       parentEventEmitter(RegistrationFlowEvent.ResetState)

feature/registration/src/test/java/org/signal/registration/screens/aepentry/EnterAepForLocalBackupViewModelTest.kt

@@ -10,11 +10,20 @@ import assertk.assertions.hasSize
 import assertk.assertions.isEmpty
 import assertk.assertions.isEqualTo
 import assertk.assertions.isNull
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.advanceUntilIdle
+import kotlinx.coroutines.test.resetMain
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.signal.core.ui.navigation.ResultEventBus
 import org.signal.registration.RegistrationFlowEvent
 
+@OptIn(ExperimentalCoroutinesApi::class)
 class EnterAepForLocalBackupViewModelTest {
 
   private lateinit var viewModel: EnterAepForLocalBackupViewModel
@@ -24,8 +33,11 @@ class EnterAepForLocalBackupViewModelTest {
 
   private val resultKey = "test-result-key"
 
+  private val testDispatcher = StandardTestDispatcher()
+
   @Before
   fun setup() {
+    Dispatchers.setMain(testDispatcher)
     resultBus = ResultEventBus()
     emittedParentEvents = mutableListOf()
     parentEventEmitter = { event -> emittedParentEvents.add(event) }
@@ -36,13 +48,19 @@ class EnterAepForLocalBackupViewModelTest {
     )
   }
 
+  @After
+  fun tearDown() {
+    Dispatchers.resetMain()
+  }
+
   // ==================== BackupKeyChanged Tests ====================
 
   @Test
-  fun `BackupKeyChanged updates backup key in state`() {
+  fun `BackupKeyChanged updates backup key in state`() = runTest {
     val testKey = VALID_AEP
 
     viewModel.onEvent(EnterAepEvents.BackupKeyChanged(testKey))
+    advanceUntilIdle()
 
     assertThat(viewModel.state.value.backupKey).isEqualTo(testKey)
   }
@@ -50,9 +68,10 @@ class EnterAepForLocalBackupViewModelTest {
   // ==================== Submit Tests ====================
 
   @Test
-  fun `Submit with valid key sends result via resultBus and emits NavigateBack`() {
+  fun `Submit with valid key sends result via resultBus and emits NavigateBack`() = runTest {
     viewModel.onEvent(EnterAepEvents.BackupKeyChanged(VALID_AEP))
     viewModel.onEvent(EnterAepEvents.Submit)
+    advanceUntilIdle()
 
     val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull()
     assertThat(result).isEqualTo(VALID_AEP)
@@ -61,9 +80,10 @@ class EnterAepForLocalBackupViewModelTest {
   }
 
   @Test
-  fun `Submit with invalid key does not send result or navigate`() {
+  fun `Submit with invalid key does not send result or navigate`() = runTest {
     viewModel.onEvent(EnterAepEvents.BackupKeyChanged("too-short"))
     viewModel.onEvent(EnterAepEvents.Submit)
+    advanceUntilIdle()
 
     assertThat(resultBus.channelMap[resultKey]).isNull()
     assertThat(emittedParentEvents).isEmpty()
@@ -72,8 +92,9 @@ class EnterAepForLocalBackupViewModelTest {
   // ==================== Cancel Tests ====================
 
   @Test
-  fun `Cancel emits NavigateBack`() {
+  fun `Cancel emits NavigateBack`() = runTest {
     viewModel.onEvent(EnterAepEvents.Cancel)
+    advanceUntilIdle()
 
     assertThat(emittedParentEvents).hasSize(1)
     assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack)
@@ -82,8 +103,9 @@ class EnterAepForLocalBackupViewModelTest {
   // ==================== DismissError Tests ====================
 
   @Test
-  fun `DismissError clears registrationError from state`() {
+  fun `DismissError clears registrationError from state`() = runTest {
     viewModel.onEvent(EnterAepEvents.DismissError)
+    advanceUntilIdle()
 
     assertThat(viewModel.state.value.registrationError).isNull()
   }

feature/registration/src/test/java/org/signal/registration/screens/countrycode/CountryCodePickerViewModelTest.kt

@@ -160,6 +160,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("United"))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.query).isEqualTo("United")
@@ -173,6 +174,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("49"))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.filteredList).isNotEmpty()
@@ -185,6 +187,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("+49"))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.filteredList).isNotEmpty()
@@ -197,6 +200,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("germany"))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.filteredList).isNotEmpty()
@@ -209,6 +213,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("usa"))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.filteredList).isNotEmpty()
@@ -222,10 +227,12 @@ class CountryCodePickerViewModelTest {
 
     // First, search for something
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("United"))
+    advanceUntilIdle()
     assertThat(viewModel.state.value.filteredList).isNotEmpty()
 
     // Then clear it
     viewModel.onEvent(CountryCodePickerScreenEvents.Search(""))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.query).isEqualTo("")
@@ -238,6 +245,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Search("xyznonexistent"))
+    advanceUntilIdle()
 
     val state = viewModel.state.value
     assertThat(state.filteredList).isEqualTo(emptyList())
@@ -252,6 +260,7 @@ class CountryCodePickerViewModelTest {
 
     val country = Country("\uD83C\uDDFA\uD83C\uDDF8", "United States", 1, "US")
     viewModel.onEvent(CountryCodePickerScreenEvents.CountrySelected(country))
+    advanceUntilIdle()
 
     val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull()
     assertThat(result).isEqualTo(country)
@@ -266,6 +275,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Dismissed)
+    advanceUntilIdle()
 
     val result = resultBus.channelMap[resultKey]?.tryReceive()?.getOrNull()
     assertThat(result).isNull()
@@ -279,6 +289,7 @@ class CountryCodePickerViewModelTest {
     advanceUntilIdle()
 
     viewModel.onEvent(CountryCodePickerScreenEvents.Dismissed)
+    advanceUntilIdle()
 
     assertThat(emittedEvents).hasSize(1)
     assertThat(emittedEvents.first()).isEqualTo(RegistrationFlowEvent.NavigateBack)

feature/registration/src/test/java/org/signal/registration/screens/pincreation/PinCreationViewModelTest.kt

@@ -251,26 +251,36 @@ class PinCreationViewModelTest {
     assertThat(emittedParentEvents.first()).isEqualTo(RegistrationFlowEvent.RegistrationComplete)
   }
 
-  // ==================== applyParentState Tests ====================
+  // ==================== ParentStateChanged Tests ====================
 
   @Test
-  fun `applyParentState copies accountEntropyPool from parent`() {
+  fun `ParentStateChanged copies accountEntropyPool from parent`() = runTest(testDispatcher) {
+    val states = collectStates()
     val aep = AccountEntropyPool.generate()
     val parentFlowState = RegistrationFlowState(accountEntropyPool = aep)
-    val initialState = PinCreationState()
 
-    val result = viewModel.applyParentState(initialState, parentFlowState)
+    viewModel.applyEvent(PinCreationState(), PinCreationScreenEvents.ParentStateChanged(parentFlowState))
 
-    assertThat(result.accountEntropyPool).isEqualTo(aep)
+    assertThat(states.last().accountEntropyPool).isEqualTo(aep)
   }
 
   @Test
-  fun `applyParentState with null accountEntropyPool keeps null`() {
+  fun `ParentStateChanged with null accountEntropyPool keeps null`() = runTest(testDispatcher) {
+    val states = collectStates()
     val parentFlowState = RegistrationFlowState(accountEntropyPool = null)
-    val initialState = PinCreationState()
 
-    val result = viewModel.applyParentState(initialState, parentFlowState)
+    viewModel.applyEvent(PinCreationState(), PinCreationScreenEvents.ParentStateChanged(parentFlowState))
 
-    assertThat(result.accountEntropyPool).isNull()
+    assertThat(states.last().accountEntropyPool).isNull()
+  }
+
+  @Test
+  fun `parent state changes are merged into state through the event stream`() = runTest(testDispatcher) {
+    val states = collectStates()
+    val aep = AccountEntropyPool.generate()
+
+    parentState.value = RegistrationFlowState(accountEntropyPool = aep)
+
+    assertThat(states.last().accountEntropyPool).isEqualTo(aep)
   }
 }

feature/registration/src/test/java/org/signal/registration/screens/pinentry/PinEntryForSmsBypassViewModelTest.kt

@@ -360,26 +360,26 @@ class PinEntryForSmsBypassViewModelTest {
     assertThat(emittedParentEvents[1]).isEqualTo(RegistrationFlowEvent.ResetState)
   }
 
-  // ==================== applyParentState Tests ====================
+  // ==================== ParentStateChanged Tests ====================
 
   @Test
-  fun `applyParentState copies e164 from parent state`() {
+  fun `ParentStateChanged copies e164 from parent state`() = runTest {
     val state = PinEntryState(mode = PinEntryState.Mode.SmsBypass)
     val parentFlowState = RegistrationFlowState(sessionE164 = "+15559876543")
 
-    val result = viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, PinEntryScreenEvents.ParentStateChanged(parentFlowState), parentEventEmitter, stateEmitter)
 
-    assertThat(result.e164).isEqualTo("+15559876543")
+    assertThat(emittedStates.last().e164).isEqualTo("+15559876543")
   }
 
   @Test
-  fun `applyParentState with null e164 in parent state sets null e164`() {
+  fun `ParentStateChanged with null e164 in parent state sets null e164`() = runTest {
     val state = PinEntryState(mode = PinEntryState.Mode.SmsBypass, e164 = "+15551234567")
     val parentFlowState = RegistrationFlowState(sessionE164 = null)
 
-    val result = viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, PinEntryScreenEvents.ParentStateChanged(parentFlowState), parentEventEmitter, stateEmitter)
 
-    assertThat(result.e164).isEqualTo(null)
+    assertThat(emittedStates.last().e164).isEqualTo(null)
   }
 
   // ==================== ToggleKeyboard Tests ====================

feature/registration/src/test/java/org/signal/registration/screens/restoreselection/ArchiveRestoreSelectionViewModelTest.kt

@@ -264,12 +264,14 @@ class ArchiveRestoreSelectionViewModelTest {
     assertThat(viewModel.state.value.restoreOptions).isEqualTo(options)
   }
 
+  // ==================== ParentStateChanged Tests ====================
+
   @Test
-  fun `applyParentState copies storageCapable from parent`() = runTest {
+  fun `ParentStateChanged copies storageCapable from parent`() = runTest {
     val viewModel = createViewModel()
 
-    val result = viewModel.applyParentState(ArchiveRestoreSelectionState(), RegistrationFlowState(storageCapable = true))
+    viewModel.applyEvent(ArchiveRestoreSelectionState(), ArchiveRestoreSelectionScreenEvents.ParentStateChanged(RegistrationFlowState(storageCapable = true)), stateEmitter)
 
-    assertThat(result.storageCapable).isTrue()
+    assertThat(emittedStates.last().storageCapable).isTrue()
   }
 }

feature/registration/src/test/java/org/signal/registration/screens/verificationcode/VerificationCodeViewModelTest.kt

@@ -76,52 +76,52 @@ class VerificationCodeViewModelTest {
     Dispatchers.resetMain()
   }
 
-  // ==================== applyParentState Tests ====================
+  // ==================== ParentStateChanged Tests ====================
 
   @Test
-  fun `applyParentState with null sessionMetadata emits ResetState`() {
+  fun `ParentStateChanged with null sessionMetadata emits ResetState`() = runTest {
     val state = VerificationCodeState()
     val parentFlowState = RegistrationFlowState(
       sessionMetadata = null,
       sessionE164 = "+15551234567"
     )
 
-    viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, VerificationCodeScreenEvents.ParentStateChanged(parentFlowState), stateEmitter)
 
     assertThat(emittedEvents).hasSize(1)
     assertThat(emittedEvents.first()).isEqualTo(RegistrationFlowEvent.ResetState)
   }
 
   @Test
-  fun `applyParentState with null sessionE164 emits ResetState`() {
+  fun `ParentStateChanged with null sessionE164 emits ResetState`() = runTest {
     val state = VerificationCodeState()
     val parentFlowState = RegistrationFlowState(
       sessionMetadata = createSessionMetadata(),
       sessionE164 = null
     )
 
-    viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, VerificationCodeScreenEvents.ParentStateChanged(parentFlowState), stateEmitter)
 
     assertThat(emittedEvents).hasSize(1)
     assertThat(emittedEvents.first()).isEqualTo(RegistrationFlowEvent.ResetState)
   }
 
   @Test
-  fun `applyParentState with both null values emits ResetState`() {
+  fun `ParentStateChanged with both null values emits ResetState`() = runTest {
     val state = VerificationCodeState()
     val parentFlowState = RegistrationFlowState(
       sessionMetadata = null,
       sessionE164 = null
     )
 
-    viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, VerificationCodeScreenEvents.ParentStateChanged(parentFlowState), stateEmitter)
 
     assertThat(emittedEvents).hasSize(1)
     assertThat(emittedEvents.first()).isEqualTo(RegistrationFlowEvent.ResetState)
   }
 
   @Test
-  fun `applyParentState with valid session copies metadata and e164`() {
+  fun `ParentStateChanged with valid session copies metadata and e164`() = runTest {
     val state = VerificationCodeState()
     val sessionMetadata = createSessionMetadata(id = "test-session")
     val e164 = "+15551234567"
@@ -130,15 +130,15 @@ class VerificationCodeViewModelTest {
       sessionE164 = e164
     )
 
-    val result = viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, VerificationCodeScreenEvents.ParentStateChanged(parentFlowState), stateEmitter)
 
     assertThat(emittedEvents).hasSize(0)
-    assertThat(result.sessionMetadata).isEqualTo(sessionMetadata)
-    assertThat(result.e164).isEqualTo(e164)
+    assertThat(emittedStates.last().sessionMetadata).isEqualTo(sessionMetadata)
+    assertThat(emittedStates.last().e164).isEqualTo(e164)
   }
 
   @Test
-  fun `applyParentState preserves existing oneTimeEvent`() {
+  fun `ParentStateChanged preserves existing oneTimeEvent`() = runTest {
     val state = VerificationCodeState(oneTimeEvent = VerificationCodeState.OneTimeEvent.NetworkError)
     val sessionMetadata = createSessionMetadata()
     val parentFlowState = RegistrationFlowState(
@@ -146,9 +146,9 @@ class VerificationCodeViewModelTest {
       sessionE164 = "+15551234567"
     )
 
-    val result = viewModel.applyParentState(state, parentFlowState)
+    viewModel.applyEvent(state, VerificationCodeScreenEvents.ParentStateChanged(parentFlowState), stateEmitter)
 
-    assertThat(result.oneTimeEvent).isEqualTo(VerificationCodeState.OneTimeEvent.NetworkError)
+    assertThat(emittedStates.last().oneTimeEvent).isEqualTo(VerificationCodeState.OneTimeEvent.NetworkError)
   }
 
   // ==================== applyEvent: ConsumeInnerOneTimeEvent Tests ====================