fix: session graph handling during auth transitions on 4.29 [WPB-26563] [WPB-26562]

app/src/main/kotlin/com/wire/android/ui/WireActivity.kt

@@ -43,6 +43,7 @@
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberUpdatedState
+import androidx.compose.runtime.setValue
 import androidx.compose.runtime.staticCompositionLocalOf
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
@@ -354,6 +355,21 @@
         val isUserUiBlocked = viewModel.globalAppState.blockUserUI != null
         val isAuthenticationRoute = currentBaseRoute in authenticationGraphRoutes
         val isSessionTransitionInProgress = viewModel.globalAppState.isSessionTransitionInProgress
+        var noSessionAuthenticationStartedWithoutSession by remember { mutableStateOf(false) }
+        LaunchedEffect(currentBaseRoute, currentUserId) {
+            when {
+                currentUserId == null && currentBaseRoute in noSessionLoginAuthenticationRoutes ->
+                    noSessionAuthenticationStartedWithoutSession = true
+                currentBaseRoute != null && currentBaseRoute !in authenticationGraphRoutes ->
+                    noSessionAuthenticationStartedWithoutSession = false
+            }
+        }
+        val navHostStartDestination = resolveNavHostStartDestination(
+            initialStartDestination = startDestination,
+            currentUserId = currentUserId,
+            currentBaseRoute = currentBaseRoute,
+            noSessionAuthenticationStartedWithoutSession = noSessionAuthenticationStartedWithoutSession,
+        )
         LaunchedEffect(currentUserId) {
             currentUserId?.let(lastKnownCurrentAccount::update)
         }
@@ -364,21 +380,26 @@
             currentUserId = currentUserId,
             sessionBackedAuthenticationUserId = sessionBackedAuthenticationUserId,
             currentBaseRoute = currentBaseRoute,
-            startDestinationBaseRoute = startDestination.baseRoute,
+            startDestinationBaseRoute = navHostStartDestination.baseRoute,
             isUserUiBlocked = isUserUiBlocked,
             isSessionTransitionInProgress = isSessionTransitionInProgress,
         )
         val lastSessionGraphContext = remember { mutableStateOf<WireActivityGraphContext?>(null) }
         if (graphContext?.sessionGraph != null) {
             lastSessionGraphContext.value = graphContext
         }
+        val graphContextWithRetainedImageLoader = graphContext?.copy(
+            imageLoaderSessionGraph = resolveWireActivityImageLoaderSessionGraph(
+                activeSessionGraph = graphContext.sessionGraph,
+                retainedSessionGraph = lastSessionGraphContext.value?.sessionGraph,
+            )
+        )
         val canRenderSessionBackedRoute = currentUserId != null || sessionBackedAuthenticationUserId != null
         val canRetainSessionGraphForContent = !isAuthenticationRoute && canRenderSessionBackedRoute
-        val contentGraphContext = graphContext ?: when {
+        val contentGraphContext = graphContextWithRetainedImageLoader ?: when {
             currentBaseRoute != null && canRetainSessionGraphForContent -> lastSessionGraphContext.value
             else -> null
         }
-        val navHostStartDestination = resolveNavHostStartDestination(startDestination, currentUserId, currentBaseRoute)
         val backgroundType by remember {
             derivedStateOf {
                 currentBackStackEntryState.value?.safeDestination()?.style.let {
@@ -410,25 +431,14 @@
         initialStartDestination: Direction,
         currentUserId: UserId?,
         currentBaseRoute: String?,
-    ): Direction {
-        if (currentUserId != null || initialStartDestination.baseRoute in authenticationGraphRoutes) {
-            return initialStartDestination
-        }
-        return when (currentBaseRoute) {
-            NewLoginScreenDestination.baseRoute -> NewLoginScreenDestination()
-            LoginScreenDestination.baseRoute -> WelcomeScreenDestination()
-            WelcomeScreenDestination.baseRoute -> WelcomeScreenDestination()
-            NewWelcomeEmptyStartScreenDestination.baseRoute,
-            null -> loggedOutStartDestination()
-            else -> initialStartDestination
-        }
-    }
-
-    private fun loggedOutStartDestination(): Direction =
-        when (loginTypeSelector.canUseNewLogin()) {
-            true -> NewWelcomeEmptyStartScreenDestination()
-            false -> WelcomeScreenDestination()
-        }
+        noSessionAuthenticationStartedWithoutSession: Boolean,
+    ): Direction = resolveWireActivityNavHostStartDestination(
+        initialStartDestination = initialStartDestination,
+        currentUserId = currentUserId,
+        currentBaseRoute = currentBaseRoute,
+        canUseNewLogin = loginTypeSelector.canUseNewLogin(),
+        noSessionAuthenticationStartedWithoutSession = noSessionAuthenticationStartedWithoutSession,
+    )
 
     private fun isNavigationAllowed(navigationCommand: NavigationCommand): Boolean {
         if (navigationCommand.destination.baseRoute != NewLoginScreenDestination.baseRoute) return true
@@ -591,6 +601,7 @@
                 graph = it,
                 viewModelFactory = (it as? ViewModelGraph)?.metroViewModelFactory ?: appGraph.metroViewModelFactory,
                 sessionGraph = sessionGraph,
+                imageLoaderSessionGraph = sessionGraph,
                 activityViewModels = activityViewModels,
             )
         }
@@ -673,7 +684,7 @@
         logoutAction: (wipeData: Boolean) -> Unit,
         content: @Composable () -> Unit,
     ) {
-        val imageLoader = sessionGraph?.wireSessionImageLoader
+        val imageLoader = imageLoaderSessionGraph?.wireSessionImageLoader
         CompositionLocalProvider(
             LocalMetroViewModelFactory provides viewModelFactory,
             LocalWireViewModelScopeKey provides graph.viewModelScopeKey,
@@ -1210,6 +1221,7 @@
     val graph: MetroViewModelGraph,
     val viewModelFactory: MetroViewModelFactory,
     val sessionGraph: AppSessionViewModelGraph?,
+    val imageLoaderSessionGraph: AppSessionViewModelGraph?,
     val activityViewModels: WireActivityScopedViewModels?,
 )
 
@@ -1265,6 +1277,11 @@
 
 private val noSessionAuthenticationGraphRoutes = loginContinuationAuthenticationRoutes + accountCreationAuthenticationRoutes
 
+private val noSessionLoginAuthenticationRoutes = setOf(
+    LoginScreenDestination.baseRoute,
+    NewLoginScreenDestination.baseRoute,
+) + loginContinuationAuthenticationRoutes
+
 private val sessionBackedAuthenticationGraphRoutes = setOf(
     RegisterDeviceScreenDestination.baseRoute,
     RemoveDeviceScreenDestination.baseRoute,
@@ -1285,10 +1302,51 @@
     request.effectiveBaseRoute in authenticationGraphRoutes &&
             request.effectiveBaseRoute !in sessionBackedAuthenticationGraphRoutes -> request.authenticationViewModelGraph
     request.sessionGraph != null -> request.sessionGraph
-    request.currentBaseRoute == null -> request.authenticationViewModelGraph
+    request.currentBaseRoute == null && request.effectiveBaseRoute in authenticationGraphRoutes -> request.authenticationViewModelGraph
     else -> null
 }
 
+internal fun resolveWireActivityNavHostStartDestination(
+    initialStartDestination: Direction,
+    currentUserId: UserId?,
+    currentBaseRoute: String?,
+    canUseNewLogin: Boolean,
+    noSessionAuthenticationStartedWithoutSession: Boolean = false,
+): Direction = when {
+    currentUserId != null &&
+            noSessionAuthenticationStartedWithoutSession &&
+            currentBaseRoute in noSessionLoginAuthenticationRoutes ->
+        resolveWireActivityLoggedOutStartDestination(canUseNewLogin)
+
+    currentUserId != null || initialStartDestination.baseRoute in authenticationGraphRoutes ->
+        initialStartDestination
+
+    currentBaseRoute == NewLoginScreenDestination.baseRoute ||
+            currentBaseRoute == NewWelcomeEmptyStartScreenDestination.baseRoute ||
+            currentBaseRoute in noSessionAuthenticationGraphRoutes ->
+        resolveWireActivityLoggedOutStartDestination(canUseNewLogin)
+
+    currentBaseRoute == null ->
+        resolveWireActivityLoggedOutStartDestination(canUseNewLogin)
+
+    currentBaseRoute == LoginScreenDestination.baseRoute ||
+            currentBaseRoute == WelcomeScreenDestination.baseRoute ->
+        WelcomeScreenDestination()
+
+    else -> initialStartDestination
+}
+
+internal fun resolveWireActivityLoggedOutStartDestination(canUseNewLogin: Boolean): Direction =
+    when (canUseNewLogin) {
+        true -> NewWelcomeEmptyStartScreenDestination()
+        false -> WelcomeScreenDestination()
+    }
+
+internal fun resolveWireActivityImageLoaderSessionGraph(
+    activeSessionGraph: AppSessionViewModelGraph?,
+    retainedSessionGraph: AppSessionViewModelGraph?,
+): AppSessionViewModelGraph? = activeSessionGraph ?: retainedSessionGraph
+
 private fun Bundle.sessionBackedAuthenticationUserId(): UserId? {
     val value = getString(SessionBackedAuthenticationNavArgs.USER_ID_VALUE_KEY)
     val domain = getString(SessionBackedAuthenticationNavArgs.USER_ID_DOMAIN_KEY)