Implement phone friendly link device registration flow.

2026-08-04 20:34:14 +01:00 · 2026-07-02 15:25:16 -04:00
parent a849af6880
commit 512df1ec1b
11 changed files with 266 additions and 61 deletions
@@ -10,12 +10,19 @@ import org.signal.core.models.MasterKey
 import org.signal.core.util.censor

 sealed interface RegistrationFlowEvent {
-  /** Navigate to a specific screen. */
-  data class NavigateToScreen(val route: RegistrationRoute) : RegistrationFlowEvent
+  /**
+   * Navigate to a specific screen.
+   *
+   * @param popCurrent Remove the current screen from the backstack
+   */
+  data class NavigateToScreen(val route: RegistrationRoute, val popCurrent: Boolean = false) : RegistrationFlowEvent

  /** Navigate back one screen. */
  data object NavigateBack : RegistrationFlowEvent

+  /** Pop the back stack back to an existing screen, removing everything above it. Replaces the current screen if the route isn't on the stack. */
+  data class NavigateBackToScreen(val route: RegistrationRoute) : RegistrationFlowEvent
+
  /** We've encountered some irrecoverable state where the best course of action is to completely reset registration. */
  data object ResetState : RegistrationFlowEvent

@@ -73,6 +73,7 @@ class RegistrationViewModel(private val repository: RegistrationRepository, save
      is RegistrationFlowEvent.Registered -> state.copy(accountEntropyPool = event.accountEntropyPool, storageCapable = event.storageCapable)
      is RegistrationFlowEvent.MasterKeyRestoredFromSvr -> state.copy(temporaryMasterKey = event.masterKey)
      is RegistrationFlowEvent.NavigateToScreen -> applyNavigationToScreenEvent(state, event)
+      is RegistrationFlowEvent.NavigateBackToScreen -> applyNavigateBackToScreenEvent(state, event)
      is RegistrationFlowEvent.NavigateBack -> {
        if (state.backStack.size > 1) {
          state.copy(backStack = state.backStack.dropLast(1))
@@ -98,10 +99,10 @@ class RegistrationViewModel(private val repository: RegistrationRepository, save
  }

  private fun applyNavigationToScreenEvent(inputState: RegistrationFlowState, event: RegistrationFlowEvent.NavigateToScreen): RegistrationFlowState {
-    val backStack = if (event.route.clearsBackStack()) {
-      listOf(event.route)
-    } else {
-      inputState.backStack + event.route
+    val backStack = when {
+      event.route.clearsBackStack() -> listOf(event.route)
+      event.popCurrent -> inputState.backStack.dropLast(1) + event.route
+      else -> inputState.backStack + event.route
    }
    return inputState.copy(backStack = backStack)
  }
@@ -122,6 +123,15 @@ class RegistrationViewModel(private val repository: RegistrationRepository, save
    }
  }

+  private fun applyNavigateBackToScreenEvent(inputState: RegistrationFlowState, event: RegistrationFlowEvent.NavigateBackToScreen): RegistrationFlowState {
+    val index = inputState.backStack.indexOfLast { it == event.route }
+    return if (index >= 0) {
+      inputState.copy(backStack = inputState.backStack.take(index + 1))
+    } else {
+      inputState.copy(backStack = inputState.backStack.dropLast(1) + event.route)
+    }
+  }
+
  /**
   * Validates a restored flow state by checking if the session is still valid.
   *
@@ -177,6 +187,7 @@ class RegistrationViewModel(private val repository: RegistrationRepository, save
        }
      }
      is RegistrationFlowEvent.NavigateBack,
+      is RegistrationFlowEvent.NavigateBackToScreen,
      is RegistrationFlowEvent.SessionUpdated,
      is RegistrationFlowEvent.E164Chosen,
      is RegistrationFlowEvent.RecoveryPasswordInvalid,
@@ -8,20 +8,26 @@ package org.signal.registration.screens.linkaccount
 import androidx.compose.animation.AnimatedContent
 import androidx.compose.animation.AnimatedVisibility
 import androidx.compose.animation.AnimatedVisibilityScope
+import androidx.compose.animation.BoundsTransform
 import androidx.compose.animation.SharedTransitionLayout
 import androidx.compose.animation.SharedTransitionScope
+import androidx.compose.animation.core.tween
 import androidx.compose.animation.fadeIn
 import androidx.compose.animation.fadeOut
+import androidx.compose.animation.togetherWith
 import androidx.compose.foundation.background
 import androidx.compose.foundation.layout.Arrangement
 import androidx.compose.foundation.layout.Arrangement.spacedBy
 import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.BoxWithConstraints
 import androidx.compose.foundation.layout.Column
 import androidx.compose.foundation.layout.Row
 import androidx.compose.foundation.layout.Spacer
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.foundation.layout.fillMaxWidth
+import androidx.compose.foundation.layout.offset
 import androidx.compose.foundation.layout.padding
+import androidx.compose.foundation.layout.safeDrawingPadding
 import androidx.compose.foundation.layout.size
 import androidx.compose.foundation.rememberScrollState
 import androidx.compose.foundation.shape.RoundedCornerShape
@@ -35,6 +41,7 @@ import androidx.compose.material3.Text
 import androidx.compose.material3.TextButton
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
@@ -55,6 +62,7 @@ import androidx.compose.ui.text.style.TextDecoration
 import androidx.compose.ui.text.withLink
 import androidx.compose.ui.unit.Dp
 import androidx.compose.ui.unit.dp
+import kotlinx.coroutines.delay
 import org.signal.core.ui.WindowBreakpoint
 import org.signal.core.ui.compose.AllDevicePreviews
 import org.signal.core.ui.compose.Buttons
@@ -75,6 +83,10 @@ import org.signal.registration.screens.attachDebugLogHelper
 import org.signal.registration.screens.quickrestore.QrState
 import org.signal.registration.test.TestTags

+private val OVERLAY_HORIZONTAL_PADDING = 24.dp
+private const val EXPAND_BUTTON_FADE_DURATION_MS = 100
+private const val QR_MORPH_DURATION_MS = 300
+
 /**
 * Screen which will display a QR code for linking this device as a secondary.
 */
@@ -85,23 +97,38 @@ fun LinkAccountScreen(
  modifier: Modifier = Modifier
 ) {
  val layoutParams = RegistrationScaffold.rememberLayoutParams()
+  val isPhone = rememberWindowBreakpoint() is WindowBreakpoint.Small
+
+  // Sequence the expand button animation with the QR morph
+  var expandButtonVisible by remember { mutableStateOf(!state.displayQrOverlay) }
+  LaunchedEffect(state.displayQrOverlay) {
+    if (state.displayQrOverlay) {
+      expandButtonVisible = false
+    } else {
+      delay(QR_MORPH_DURATION_MS.toLong())
+      expandButtonVisible = true
+    }
+  }

  Surface(modifier = modifier.testTag(TestTags.LINK_ACCOUNT_SCREEN)) {
    SharedTransitionLayout {
      AnimatedContent(
        targetState = state.displayQrOverlay,
-        label = "qr_code_fullscreen_transition"
+        label = "qr_code_fullscreen_transition",
+        transitionSpec = {
+          fadeIn(animationSpec = tween(durationMillis = QR_MORPH_DURATION_MS, delayMillis = if (targetState) EXPAND_BUTTON_FADE_DURATION_MS else 0)) togetherWith fadeOut(tween(EXPAND_BUTTON_FADE_DURATION_MS))
+        }
      ) { target ->
        CompositionLocalProvider(
          LocalSharedTransitionScope provides this@SharedTransitionLayout,
          LocalAnimateVisibilityScope provides this
        ) {
          if (target) {
-            QrCodeOverlay(state, onEvent)
+            QrCodeOverlay(state, onEvent, isPhone)
          } else {
            when (layoutParams) {
-              is RegistrationScaffold.Params.OnePane -> OnePane(layoutParams, state, onEvent)
-              is RegistrationScaffold.Params.TwoPane -> TwoPane(layoutParams, state, onEvent)
+              is RegistrationScaffold.Params.OnePane -> OnePane(layoutParams, isPhone, expandButtonVisible, state, onEvent)
+              is RegistrationScaffold.Params.TwoPane -> TwoPane(layoutParams, expandButtonVisible, state, onEvent)
            }
          }
        }
@@ -139,6 +166,8 @@ private fun StateDialogs(
@Composable
 private fun OnePane(
  params: RegistrationScaffold.Params.OnePane,
+  isPhone: Boolean,
+  expandButtonVisible: Boolean,
  state: LinkAccountScreenState,
  onEvent: (LinkAccountScreenEvent) -> Unit
 ) {
@@ -149,16 +178,20 @@ private fun OnePane(
    content = { paddingValues ->
      Column(
        horizontalAlignment = Alignment.CenterHorizontally,
-        verticalArrangement = spacedBy(64.dp),
+        verticalArrangement = spacedBy(if (isPhone) 32.dp else 64.dp),
        modifier = Modifier
          .verticalScroll(scrollState)
          .padding(paddingValues)
      ) {
        Title()

-        QrCodeContent(state = state, onEvent = onEvent)
+        QrCodeContent(state = state, onEvent = onEvent, isPhone = isPhone, expandButtonVisible = expandButtonVisible)

-        Steps(verticalArrangement = spacedBy(32.dp), onEvent)
+        Steps(
+          verticalArrangement = spacedBy(if (isPhone) 20.dp else 32.dp),
+          centerGetHelp = isPhone,
+          onEvent = onEvent
+        )
      }
    },
    footer = {
@@ -174,6 +207,7 @@ private fun OnePane(
@Composable
 private fun TwoPane(
  params: RegistrationScaffold.Params.TwoPane,
+  expandButtonVisible: Boolean,
  state: LinkAccountScreenState,
  onEvent: (LinkAccountScreenEvent) -> Unit
 ) {
@@ -193,7 +227,8 @@ private fun TwoPane(
        onEvent = onEvent,
        modifier = Modifier
          .weight(1f)
-          .padding(paddingValues)
+          .padding(paddingValues),
+        expandButtonVisible = expandButtonVisible
      )
    },
    footer = {
@@ -217,7 +252,7 @@ private fun FirstPaneContent(
  ) {
    Title()

-    Steps(verticalArrangement = spacedBy(32.dp), onEvent = onEvent)
+    Steps(verticalArrangement = spacedBy(32.dp), centerGetHelp = false, onEvent = onEvent)
  }
 }

@@ -235,6 +270,7 @@ private fun Title() {
@Composable
 private fun Steps(
  verticalArrangement: Arrangement.Vertical,
+  centerGetHelp: Boolean,
  onEvent: (LinkAccountScreenEvent) -> Unit
 ) {
  Column(verticalArrangement = verticalArrangement) {
@@ -253,17 +289,21 @@ private fun Steps(
      text = stringResource(R.string.LinkAccountScreen__tap_linked_devices_and_link_new_device)
    )

-    GetHelp(onEvent)
+    GetHelp(
+      onEvent = onEvent,
+      modifier = if (centerGetHelp) Modifier.align(Alignment.CenterHorizontally) else Modifier
+    )
  }
 }

@Composable
 private fun GetHelp(
-  onEvent: (LinkAccountScreenEvent) -> Unit
+  onEvent: (LinkAccountScreenEvent) -> Unit,
+  modifier: Modifier = Modifier
 ) {
  TextButton(
    onClick = { onEvent(LinkAccountScreenEvent.GetHelpClick) },
-    modifier = Modifier.testTag(TestTags.LINK_ACCOUNT_GET_HELP_BUTTON)
+    modifier = modifier.testTag(TestTags.LINK_ACCOUNT_GET_HELP_BUTTON)
  ) {
    Text(text = stringResource(R.string.LinkAccountScreen__get_help_with_these_steps))
  }
@@ -285,15 +325,25 @@ private fun Step(icon: ImageVector, text: String) {
 private fun QrCodeContent(
  state: LinkAccountScreenState,
  onEvent: (LinkAccountScreenEvent) -> Unit,
-  modifier: Modifier = Modifier
+  modifier: Modifier = Modifier,
+  isPhone: Boolean = false,
+  isInOverlay: Boolean = false,
+  expandButtonVisible: Boolean = true,
+  overlayMaxWidth: Dp? = null
 ) {
  val sharedTransitionScope = LocalSharedTransitionScope.current!!
  val animatedVisibilityScope = LocalAnimateVisibilityScope.current!!

+  // Delay the morph on expand (so the button can fade out first) but not on collapse
+  val expanding = state.displayQrOverlay
+  val qrBoundsTransform = remember(expanding) {
+    BoundsTransform { _, _ ->
+      tween(durationMillis = QR_MORPH_DURATION_MS, delayMillis = if (expanding) EXPAND_BUTTON_FADE_DURATION_MS else 0)
+    }
+  }
+
  Box(
-    contentAlignment = if (!state.displayQrOverlay) {
-      Alignment.CenterEnd
-    } else Alignment.Center,
+    contentAlignment = if (isInOverlay) Alignment.Center else Alignment.CenterEnd,
    modifier = modifier
  ) {
    with(sharedTransitionScope) {
@@ -302,20 +352,22 @@ private fun QrCodeContent(
        modifier = Modifier
          .sharedElement(
            sharedContentState = rememberSharedContentState("qr_code_outer_border"),
-            animatedVisibilityScope = animatedVisibilityScope
+            animatedVisibilityScope = animatedVisibilityScope,
+            boundsTransform = qrBoundsTransform
          )
-          .size(getQrOuterBorderSize(state.displayQrOverlay))
-          .background(color = colorResource(org.signal.core.ui.R.color.signal_light_colorPrimary), shape = RoundedCornerShape(64.dp))
+          .size(getQrOuterBorderSize(isInOverlay, overlayMaxWidth))
+          .background(color = colorResource(org.signal.core.ui.R.color.signal_light_colorPrimary), shape = RoundedCornerShape(if (isPhone) 48.dp else 64.dp))
      ) {
        AnimatedContent(
          targetState = state.qrCodeState,
          modifier = Modifier
            .sharedElement(
              sharedContentState = rememberSharedContentState("qr_code_inner_border"),
-              animatedVisibilityScope = animatedVisibilityScope
+              animatedVisibilityScope = animatedVisibilityScope,
+              boundsTransform = qrBoundsTransform
            )
-            .size(getQrInnerBorderSize(state.displayQrOverlay))
-            .background(color = Color.White, shape = RoundedCornerShape(24.dp))
+            .size(getQrInnerBorderSize(isInOverlay, overlayMaxWidth))
+            .background(color = Color.White, shape = RoundedCornerShape(if (isPhone) 26.dp else 24.dp))
        ) { target ->
          Box(
            contentAlignment = Alignment.Center,
@@ -323,7 +375,7 @@ private fun QrCodeContent(
          ) {
            when (target) {
              QrState.Failed -> QrCodeFailed(onEvent)
-              is QrState.Loaded -> QrCodeDisplay(target.qrCodeData, state.displayQrOverlay, sharedTransitionScope, animatedVisibilityScope)
+              is QrState.Loaded -> QrCodeDisplay(target.qrCodeData, isInOverlay, overlayMaxWidth, qrBoundsTransform, sharedTransitionScope, animatedVisibilityScope)
              QrState.Loading -> QrCodeLoading()
              QrState.Scanned -> QrCodeScanned()
            }
@@ -333,14 +385,17 @@ private fun QrCodeContent(
    }

    AnimatedVisibility(
-      visible = state.qrCodeState is QrState.Loaded && !state.displayQrOverlay,
-      modifier = Modifier.align(Alignment.TopEnd),
-      enter = fadeIn(),
-      exit = fadeOut()
+      visible = state.qrCodeState is QrState.Loaded && !isInOverlay && expandButtonVisible,
+      modifier = Modifier
+        .align(Alignment.TopEnd)
+        .then(if (isPhone) Modifier.offset(x = 6.dp, y = (-6).dp) else Modifier)
+        .then(with(sharedTransitionScope) { Modifier.renderInSharedTransitionScopeOverlay(zIndexInOverlay = 1f) }),
+      enter = fadeIn(tween(EXPAND_BUTTON_FADE_DURATION_MS)),
+      exit = fadeOut(tween(EXPAND_BUTTON_FADE_DURATION_MS))
    ) {
      IconButtons.IconButton(
        onClick = { onEvent(LinkAccountScreenEvent.DisplayOverlayClick) },
-        size = 53.dp,
+        size = if (isPhone) 40.dp else 53.dp,
        colors = IconButtons.iconButtonColors(
          containerColor = Color(0xFF506DCD),
          contentColor = colorResource(org.signal.core.ui.R.color.signal_light_colorOnPrimary)
@@ -360,6 +415,8 @@ private fun QrCodeContent(
 private fun QrCodeDisplay(
  qrCodeData: QrCodeData,
  isInOverlay: Boolean,
+  overlayMaxWidth: Dp?,
+  boundsTransform: BoundsTransform,
  sharedTransitionScope: SharedTransitionScope,
  animatedVisibilityScope: AnimatedVisibilityScope
 ) {
@@ -370,9 +427,10 @@ private fun QrCodeDisplay(
      modifier = Modifier
        .sharedElement(
          sharedContentState = rememberSharedContentState("qr_code_display"),
-          animatedVisibilityScope = animatedVisibilityScope
+          animatedVisibilityScope = animatedVisibilityScope,
+          boundsTransform = boundsTransform
        )
-        .size(getQrCodeSize(isInOverlay))
+        .size(getQrCodeSize(isInOverlay, overlayMaxWidth))
    )
  }
 }
@@ -431,18 +489,26 @@ private fun QrCodeFailed(
 }

@Composable
-fun QrCodeOverlay(
+private fun QrCodeOverlay(
  state: LinkAccountScreenState,
-  onEvent: (LinkAccountScreenEvent) -> Unit
+  onEvent: (LinkAccountScreenEvent) -> Unit,
+  isPhone: Boolean
 ) {
  Surface(
    modifier = Modifier.fillMaxSize()
  ) {
-    Box {
+    BoxWithConstraints(
+      modifier = Modifier
+        .fillMaxSize()
+        .safeDrawingPadding()
+    ) {
      QrCodeContent(
        state = state,
        onEvent = onEvent,
-        modifier = Modifier.align(Alignment.Center)
+        modifier = Modifier.align(Alignment.Center),
+        isPhone = isPhone,
+        isInOverlay = true,
+        overlayMaxWidth = maxWidth
      )

      IconButtons.IconButton(
@@ -459,47 +525,52 @@ fun QrCodeOverlay(
 }

@Composable
-fun getQrOuterBorderSize(isInOverlay: Boolean): Dp {
+private fun getQrOuterBorderSize(isInOverlay: Boolean, overlayMaxWidth: Dp? = null): Dp {
  if (isInOverlay) {
-    return 456.dp
+    return overlayOuterBorderSize(overlayMaxWidth)
  }

  val breakpoint = rememberWindowBreakpoint()
  return when (breakpoint) {
-    is WindowBreakpoint.Small -> 296.dp
+    is WindowBreakpoint.Small -> 272.dp
    is WindowBreakpoint.Medium -> 296.dp
    is WindowBreakpoint.Large -> 364.dp
  }
 }

@Composable
-fun getQrInnerBorderSize(isInOverlay: Boolean): Dp {
+private fun getQrInnerBorderSize(isInOverlay: Boolean, overlayMaxWidth: Dp? = null): Dp {
  if (isInOverlay) {
-    return 360.dp
+    return overlayOuterBorderSize(overlayMaxWidth) * (360f / 456f)
  }

  val breakpoint = rememberWindowBreakpoint()
  return when (breakpoint) {
-    is WindowBreakpoint.Small -> 232.dp
+    is WindowBreakpoint.Small -> 208.dp
    is WindowBreakpoint.Medium -> 232.dp
    is WindowBreakpoint.Large -> 284.dp
  }
 }

@Composable
-fun getQrCodeSize(isInOverlay: Boolean): Dp {
+private fun getQrCodeSize(isInOverlay: Boolean, overlayMaxWidth: Dp? = null): Dp {
  if (isInOverlay) {
-    return 297.dp
+    return overlayOuterBorderSize(overlayMaxWidth) * (297f / 456f)
  }

  val breakpoint = rememberWindowBreakpoint()
  return when (breakpoint) {
-    is WindowBreakpoint.Small -> 208.dp
+    is WindowBreakpoint.Small -> 176.dp
    is WindowBreakpoint.Medium -> 208.dp
    is WindowBreakpoint.Large -> 256.dp
  }
 }

+private fun overlayOuterBorderSize(overlayMaxWidth: Dp?): Dp {
+  overlayMaxWidth ?: return 456.dp
+  return (overlayMaxWidth - OVERLAY_HORIZONTAL_PADDING * 2).coerceAtMost(456.dp)
+}
+
@Composable
 private fun OnePaneFooterContent(
  params: RegistrationScaffold.Params.OnePane,
@@ -65,7 +65,12 @@ class LinkAccountViewModel(
    val result = when (event) {
      LinkAccountScreenEvent.GetHelpClick -> error("This event is handled in the nav-entry.")
      LinkAccountScreenEvent.CreateAccountClick -> {
-        parentEventEmitter.navigateTo(RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry))
+        // Revisit permission screen if necessary
+        if (parentState.value.backStack.any { it == RegistrationRoute.PhoneNumberEntry }) {
+          parentEventEmitter(RegistrationFlowEvent.NavigateBackToScreen(RegistrationRoute.PhoneNumberEntry))
+        } else {
+          parentEventEmitter.navigateTo(RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry), popCurrent = true)
+        }
        state
      }
      LinkAccountScreenEvent.DisplayOverlayClick -> state.copy(displayQrOverlay = true)
@@ -228,7 +228,7 @@ private fun OnePaneLayout(

  OnePaneRegistrationScaffold(
    params = params,
-    topBar = { TopAppBar(scrollBehavior = topBarScrollBehavior) },
+    topBar = { TopAppBar(scrollBehavior = topBarScrollBehavior, onEvent = onEvent) },
    content = { paddingValues ->
      Column(
        modifier = Modifier
@@ -289,7 +289,7 @@ private fun TwoPaneLayout(

  TwoPaneRegistrationScaffold(
    params = params,
-    topBar = { TopAppBar(scrollBehavior = topBarScrollBehavior) },
+    topBar = { TopAppBar(scrollBehavior = topBarScrollBehavior, onEvent = onEvent) },
    firstPane = { paddingValues ->
      Column(
        modifier = Modifier
@@ -344,7 +344,8 @@ private fun TwoPaneLayout(
@OptIn(ExperimentalMaterial3Api::class)
@Composable
 fun TopAppBar(
-  scrollBehavior: TopAppBarScrollBehavior
+  scrollBehavior: TopAppBarScrollBehavior,
+  onEvent: (PhoneNumberEntryScreenEvents) -> Unit
 ) {
  val context = LocalContext.current

@@ -382,7 +383,7 @@ fun TopAppBar(
        DropdownMenus.Item(
          text = { Text(text = stringResource(R.string.RegistrationActivity_link_device)) },
          onClick = {
-            TODO("Handle link device")
+            onEvent(PhoneNumberEntryScreenEvents.LinkDevice)
            menuController.hide()
          }
        )
@@ -36,6 +36,10 @@ sealed class PhoneNumberEntryScreenEvents {

  /** The user requested to open the country picker.  */
  data object CountryPicker : PhoneNumberEntryScreenEvents()
+
+  /** The user chose to link this device to an existing account instead of registering a new number. */
+  data object LinkDevice : PhoneNumberEntryScreenEvents()
+
  data class CaptchaCompleted(val token: String) : PhoneNumberEntryScreenEvents() {
    override fun toString(): String = "CaptchaCompleted(token=${token.censor()})"
  }
@@ -132,6 +132,9 @@ class PhoneNumberEntryViewModel(
          )
        }
      }
+      is PhoneNumberEntryScreenEvents.LinkDevice -> {
+        parentEventEmitter.navigateTo(RegistrationRoute.LinkAccount)
+      }
      is PhoneNumberEntryScreenEvents.CaptchaCompleted -> {
        stateEmitter(applyCaptchaCompleted(state, event.token, parentEventEmitter))
      }
@@ -11,8 +11,8 @@ import org.signal.registration.RegistrationRoute
 /**
 * Convenience function to emit a navigation event to a parentEmitter.
 */
-fun ((RegistrationFlowEvent) -> Unit).navigateTo(route: RegistrationRoute) {
-  this(RegistrationFlowEvent.NavigateToScreen(route))
+fun ((RegistrationFlowEvent) -> Unit).navigateTo(route: RegistrationRoute, popCurrent: Boolean = false) {
+  this(RegistrationFlowEvent.NavigateToScreen(route, popCurrent))
 }

 /**
@@ -459,6 +459,68 @@ class RegistrationViewModelTest {
    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, preRegisterSelection))
  }

+  @Test
+  fun `applyEvent NavigateToScreen with popCurrent replaces the current screen`() = runTest(testDispatcher) {
+    coEvery { mockRepository.restoreFlowState() } returns null
+    coEvery { mockRepository.getPreExistingRegistrationData() } returns null
+
+    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
+    advanceUntilIdle()
+
+    val initialState = RegistrationFlowState(
+      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry, RegistrationRoute.LinkAccount)
+    )
+
+    val result = viewModel.applyEvent(
+      initialState,
+      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry), popCurrent = true)
+    )
+
+    assertThat(result.backStack).isEqualTo(
+      listOf(
+        RegistrationRoute.Welcome,
+        RegistrationRoute.PhoneNumberEntry,
+        RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry)
+      )
+    )
+  }
+
+  @Test
+  fun `applyEvent NavigateBackToScreen pops back to the target screen`() = runTest(testDispatcher) {
+    coEvery { mockRepository.restoreFlowState() } returns null
+    coEvery { mockRepository.getPreExistingRegistrationData() } returns null
+
+    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
+    advanceUntilIdle()
+
+    val initialState = RegistrationFlowState(
+      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry, RegistrationRoute.LinkAccount)
+    )
+
+    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBackToScreen(RegistrationRoute.PhoneNumberEntry))
+
+    assertThat(result.backStack).isEqualTo(listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry))
+  }
+
+  @Test
+  fun `applyEvent NavigateBackToScreen replaces the current screen when the route is not on the back stack`() = runTest(testDispatcher) {
+    coEvery { mockRepository.restoreFlowState() } returns null
+    coEvery { mockRepository.getPreExistingRegistrationData() } returns null
+
+    val viewModel = RegistrationViewModel(mockRepository, SavedStateHandle())
+    advanceUntilIdle()
+
+    val initialState = RegistrationFlowState(
+      backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.LinkAccount)
+    )
+
+    val result = viewModel.applyEvent(initialState, RegistrationFlowEvent.NavigateBackToScreen(RegistrationRoute.PhoneNumberEntry))
+
+    assertThat(result.backStack).isEqualTo(
+      listOf(RegistrationRoute.Welcome, RegistrationRoute.PhoneNumberEntry)
+    )
+  }
+
  @Test
  fun `applyEvent ResetState returns default state`() = runTest(testDispatcher) {
    coEvery { mockRepository.restoreFlowState() } returns null
@@ -66,10 +66,10 @@ class LinkAccountViewModelTest {
    Dispatchers.resetMain()
  }

-  private fun TestScope.createViewModel(): LinkAccountViewModel {
+  private fun TestScope.createViewModel(parentState: RegistrationFlowState = RegistrationFlowState()): LinkAccountViewModel {
    val viewModel = LinkAccountViewModel(
      repository = mockRepository,
-      parentState = MutableStateFlow(RegistrationFlowState()),
+      parentState = MutableStateFlow(parentState),
      parentEventEmitter = parentEventEmitter
    )
    // Keep the WhileSubscribed state flow hot so state.value reflects updates during the test.
@@ -211,13 +211,38 @@ class LinkAccountViewModelTest {
  }

  @Test
-  fun `applyEvent CreateAccountClick navigates to Permissions`() = runTest(testDispatcher) {
-    val viewModel = createViewModel()
+  fun `applyEvent CreateAccountClick from link-device-first flow routes through Permissions`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(
+      RegistrationFlowState(backStack = listOf(RegistrationRoute.Welcome, RegistrationRoute.LinkAccount))
+    )

    viewModel.applyEvent(LinkAccountScreenState(), LinkAccountScreenEvent.CreateAccountClick, stateEmitter)

    assertThat(emittedParentEvents).contains(
-      RegistrationFlowEvent.NavigateToScreen(RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry))
+      RegistrationFlowEvent.NavigateToScreen(
+        route = RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
+        popCurrent = true
+      )
+    )
+  }
+
+  @Test
+  fun `applyEvent CreateAccountClick from phone number screen returns to it without Permissions`() = runTest(testDispatcher) {
+    val viewModel = createViewModel(
+      RegistrationFlowState(
+        backStack = listOf(
+          RegistrationRoute.Welcome,
+          RegistrationRoute.Permissions(nextRoute = RegistrationRoute.PhoneNumberEntry),
+          RegistrationRoute.PhoneNumberEntry,
+          RegistrationRoute.LinkAccount
+        )
+      )
+    )
+
+    viewModel.applyEvent(LinkAccountScreenState(), LinkAccountScreenEvent.CreateAccountClick, stateEmitter)
+
+    assertThat(emittedParentEvents).contains(
+      RegistrationFlowEvent.NavigateBackToScreen(RegistrationRoute.PhoneNumberEntry)
    )
  }

@@ -249,6 +249,22 @@ class PhoneNumberEntryViewModelTest {
    )
  }

+  @Test
+  fun `LinkDevice navigates to the link account flow`() = runTest {
+    viewModel.applyEvent(
+      PhoneNumberEntryState(),
+      PhoneNumberEntryScreenEvents.LinkDevice,
+      parentEventEmitter,
+      stateEmitter
+    )
+
+    assertThat(emittedEvents).hasSize(1)
+    assertThat(emittedEvents.first())
+      .isInstanceOf<RegistrationFlowEvent.NavigateToScreen>()
+      .prop(RegistrationFlowEvent.NavigateToScreen::route)
+      .isInstanceOf<RegistrationRoute.LinkAccount>()
+  }
+
  @Test
  fun `ConsumeInnerOneTimeEvent clears inner event`() = runTest {
    val initialState = PhoneNumberEntryState(