How to navigate to different screens based on the state in viewmodel with jetpack compose

Question

I have an authentication module which contains login, otp confirmation among other screens. I have an AuthViewModel which handles UI events as below:

@HiltViewModel
class AuthViewModel @Inject constructor(private val repo: AuthRepository) : ViewModel() {
    private val _authState = MutableStateFlow(AuthState.UNAUTHENTICATED)
    val authState: StateFlow<AuthState> = _authState.asStateFlow()

    fun login(credential: Login) {
        viewModelScope.launch {
            //.. call repo method here
            _authState.value = AuthState.OTP
         }
     }

     fun confirm(otp: String) {
        viewModelScope.launch {
            //.. call repo method here
            _authState.value = AuthState.AUTHENTICATED
         }
     }
}

As you can see I update the authState to navigate to different screens when login() and confirm() are successful. However, I'm unable to observe the authState from a common composable function and trigger the navigation accordingly.

Here is my navigation graph looks like:

const val authRoute = "auth"
fun NavGraphBuilder.authNavigation(navController: NavController) {
    navigation(startDestination = loginRoute, route = authRoute) {
        composable(loginRoute) {
            val parentEntry = remember(it) { navController.getBackStackEntry(authRoute) }
            val vm: AuthViewModel = hiltViewModel(parentEntry)
            LoginScreen(vm, onOtpSent = {
                navController.navigate(otpRoute)
            })
        }

        composable(otpRoute) {
            val parentEntry = remember(it) { navController.getBackStackEntry(authRoute) }
            val vm: AuthViewModel = hiltViewModel(parentEntry)
            OtpScreen(vm, onAuthenticated = {
                navController.navigate(homeRoute)
            })
        }
    }
}

Currently, I observe the vm.authState from each screen and then trigger callback lamdbda functions as below:

@Composable
fun LoginScreen(vm: AuthViewModel, onOtpSent: () -> Unit) {
    val authState = vm.authState.collectAsStateWithLifecycle().value
    val currentOnOtpSent = rememberUpdatedState(onOtpSent)

    LaunchedEffect(authState) {
        if (authState == AuthState.OTP) {
            currentOnOtpSent.value()
        }
    }

    val state = vm.loginUiState.collectAsStateWithLifecycle().value
    LoginUi(state, onLogin = vm::login)
}

However, the problem with this approach is that, when user presses the back button from the OtpScreen, the LaunchedEffect coroutine in the LoginScreen gets called again navigating back to the OtpScreen again instead of staying in the LoginScreen as the authState in viewmodel still contains navigate to OTP state. This is a likely issue if I were to observe the authState from all the composable routes instead of from just one place.

In the View world I used to solve this issue by only observing the authState in ViewModel from the Activity as a common place for all the Auth fragments and navigate to appropriate fragments based on the authState.

However, in the compose world, I can't find a common place to navigate between the auth screens based on the ViewModel state. How do I do that?

Any help would be highly appreciated!

Answer 1

After hours of experiments, I ended up doing as below as per the documentation provided here for Navigation events

@Composable
fun LoginScreen(vm: AuthViewModel, onOtpSent: () -> Unit, onRegister: () -> Unit) {
    val state = vm.loginUiState.collectAsStateWithLifecycle().value
    var isLoggingIn by rememberSaveable { mutableStateOf(false) }

    LoginUi(
        state,
        onLogin = {
            isLoggingIn = true
            vm.login(it)
        },
        onRegister = onRegister
    )

    val lifecycle = LocalLifecycleOwner.current.lifecycle
    val currentOnOtpSent = rememberUpdatedState(onOtpSent).value

    if (isLoggingIn) {
        LaunchedEffect(vm, lifecycle) {
            snapshotFlow { vm.authState }
                .filter { it == AuthState.OTP }
                .flowWithLifecycle(lifecycle)
                .collect {
                    currentOnOtpSent()
                    isLoggingIn = false
                }
        }
    }
}

The issue here is you would still have to replicate this code throughout the authentication screens whereever you expect to have a navigation event from ViewModel.

The isLoggingIn state keeps track of the event that's triggered to potentially update the navigation state and the LaunchedEffect would only trigger when there is event triggered from UI (login, register etc). This prevents the issue with the back navigation.