Fixes for back button navigation and black screen

app/src/main/java/chat/zekochat/activities/MainActivity.kt

@@ -134,6 +134,7 @@ import dagger.hilt.android.lifecycle.HiltViewModel
 import dagger.hilt.android.qualifiers.ApplicationContext
 import io.ktor.client.request.get
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -147,6 +148,9 @@ class MainActivityViewModel @Inject constructor(
     val nextDestination = MutableStateFlow<String?>(null)
     var isConnected = MutableStateFlow(false)
     val isReady = MutableStateFlow(false)
+    // Controls when we allow the first UI draw.
+    // Authentication/navigation readiness is represented by `isReady`.
+    val isUiReady = MutableStateFlow(false)
     val couldNotLogIn = MutableStateFlow(false)
 
     private fun hasInternetConnection(): Boolean {
@@ -184,11 +188,9 @@ class MainActivityViewModel @Inject constructor(
         isReady.emit(true)
     }
 
-    private fun doPreStartupTasks() {
+    private fun doPostStartupTasks() {
-        Log.d("MainActivity", "Performing pre-startup tasks")
+        Log.d("MainActivity", "Performing post-startup tasks")
         viewModelScope.launch {
-            Log.d("MainActivity", "Hydrating Experiments from KV")
-            Experiments.hydrateWithKv()
             Log.d("MainActivity", "Performing health check")
             doHealthCheck()
             Log.d("MainActivity", "Performing update geo state")
@@ -225,6 +227,8 @@ class MainActivityViewModel @Inject constructor(
     fun checkLoggedInState() {
         viewModelScope.launch {
             Log.d("MainActivity", "Checking logged in state")
+            Log.d("MainActivity", "Hydrating Experiments from KV")
+            Experiments.hydrateWithKv()
 
             isConnected.emit(hasInternetConnection())
 
@@ -339,8 +343,13 @@ class MainActivityViewModel @Inject constructor(
 
     init {
         Log.d("MainActivity", "Starting up")
-        doPreStartupTasks()
+        // Do not block initial rendering on slow network/login checks.
+        isUiReady.value = true
         checkLoggedInState()
+        viewModelScope.launch {
+            isReady.first { it }
+            doPostStartupTasks()
+        }
     }
 }
 
@@ -419,7 +428,7 @@ class MainActivity : AppCompatActivity() {
             object : ViewTreeObserver.OnPreDrawListener {
                 override fun onPreDraw(): Boolean {
                     // Check whether the initial data is ready.
-                    return if (viewModel.isReady.value) {
+                    return if (viewModel.isUiReady.value) {
                         // The content is ready. Start drawing.
                         content.viewTreeObserver.removeOnPreDrawListener(this)
                         true

app/src/main/java/chat/zekochat/screens/DefaultDestinationScreen.kt

@@ -3,9 +3,11 @@ package chat.zekochat.screens
 import android.app.Activity
 import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
+import androidx.compose.material3.CircularProgressIndicator
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.Alignment
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.toArgb
 import androidx.compose.ui.platform.LocalContext
@@ -47,5 +49,12 @@ fun DefaultDestinationScreen(
         }
     }
 
+    if (nextDestination == null) {
+        Box(Modifier.fillMaxSize(), contentAlignment = Alignment.Center) {
+            CircularProgressIndicator()
+        }
+        return
+    }
+
     Box(Modifier.fillMaxSize())
 }

app/src/main/java/chat/zekochat/screens/chooseplatform/ChoosePlatformScreen.kt

@@ -45,6 +45,22 @@ fun ChoosePlatformScreen(navController: NavController) {
     val apiUrlValue = remember {
         mutableStateOf("")
     }
+    fun continueWithSelectedPlatform() {
+        if (apiUrlValue.value.isNotEmpty()) {
+            when (apiUrlValue.value) {
+                ApplicationPlatform.PEP.baseUrl -> PeptideAPI.setPlatform(
+                    ApplicationPlatform.PEP
+                )
+
+                ApplicationPlatform.REVOLT.baseUrl -> PeptideAPI.setPlatform(
+                    ApplicationPlatform.REVOLT
+                )
+
+                else -> return
+            }
+        }
+        navController.navigate("login/greeting")
+    }
 
     Column(
         modifier = Modifier
@@ -118,6 +134,7 @@ fun ChoosePlatformScreen(navController: NavController) {
                         .size(72.dp)
                         .clickable {
                             apiUrlValue.value = ApplicationPlatform.PEP.baseUrl
+                            continueWithSelectedPlatform()
                         },
                 ) {
                     Box(
@@ -152,6 +169,7 @@ fun ChoosePlatformScreen(navController: NavController) {
                         .size(72.dp)
                         .clickable {
                             apiUrlValue.value = ApplicationPlatform.REVOLT.baseUrl
+                            continueWithSelectedPlatform()
                         },
                 ) {
                     Box(
@@ -199,20 +217,7 @@ fun ChoosePlatformScreen(navController: NavController) {
                 .padding(vertical = 8.dp)
         )
         SquareButton(
-            onClick = {
+            onClick = ::continueWithSelectedPlatform,
-                if(apiUrlValue.value.isNotEmpty()){
-                    when(apiUrlValue.value){
-                        ApplicationPlatform.PEP.baseUrl -> PeptideAPI.setPlatform(
-                            ApplicationPlatform.PEP
-                        )
-                        ApplicationPlatform.REVOLT.baseUrl -> PeptideAPI.setPlatform(
-                            ApplicationPlatform.REVOLT
-                        )
-                        else -> return@SquareButton
-                    }
-                }
-                navController.navigate("login/greeting")
-            },
             modifier = Modifier
                 .padding(top = 32.dp, bottom = 16.dp)
                 .fillMaxWidth()

app/src/main/java/chat/zekochat/screens/login/LoginScreen.kt

@@ -198,7 +198,8 @@ fun LoginScreen(navController: NavController, viewModel: LoginViewModel = hiltVi
 
             "home" -> {
                 navController.navigate("chat") {
-                    popUpTo("login/greeting") { inclusive = true }
+                    // Clear auth flow history while keeping graph/state restoration stable.
+                    popUpTo("choose-platform") { inclusive = true }
                 }
             }
 

app/src/main/java/chat/zekochat/screens/login/MfaScreen.kt

@@ -149,7 +149,8 @@ fun MfaScreen(
     LaunchedEffect(viewModel.navigateToHome) {
         if (viewModel.navigateToHome) {
             navController.navigate("chat") {
-                popUpTo("login/greeting") { inclusive = true }
+                // Clear auth flow history while keeping graph/state restoration stable.
+                popUpTo("choose-platform") { inclusive = true }
             }
         }
     }