fix(navigation): reset graph state after client removal [WPB-27228] [WPB-27362]

app/src/androidTest/kotlin/com/wire/android/ui/WireActivityNavigatorTest.kt

@@ -21,7 +21,11 @@ package com.wire.android.ui
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.ui.test.junit4.createComposeRule
 import androidx.navigation.NavHostController
+import androidx.navigation.compose.composable
+import androidx.navigation.createGraph
+import org.junit.Assert.assertEquals
 import org.junit.Assert.assertNotSame
+import org.junit.Assert.assertNull
 import org.junit.Rule
 import org.junit.Test
 
@@ -63,4 +67,56 @@ class WireActivityNavigatorTest {
             assertNotSame(sessionController, currentController)
         }
     }
+
+    @Test
+    fun givenNavigatorIsRecreatedWithoutActivityRecreation_whenReadingBackStack_thenStateTracksNewGraph() {
+        val isUserUiBlocked = mutableStateOf(false)
+        lateinit var currentController: NavHostController
+        var currentRoute: String? = null
+
+        composeTestRule.setContent {
+            val navigator = rememberWireActivityNavigator(
+                isUserUiBlocked = isUserUiBlocked.value,
+                finish = {},
+                isAllowedToNavigate = { true },
+            )
+            currentController = navigator.navController
+            currentRoute = wireActivityCurrentBackStackEntryAsState(navigator)
+                .value
+                ?.destination
+                ?.route
+        }
+
+        composeTestRule.runOnIdle {
+            currentController.graph = currentController.createGraph(startDestination = SESSION_ROUTE) {
+                composable(SESSION_ROUTE) {}
+            }
+        }
+        composeTestRule.waitForIdle()
+        composeTestRule.runOnIdle {
+            assertEquals(SESSION_ROUTE, currentRoute)
+            isUserUiBlocked.value = true
+        }
+        composeTestRule.waitForIdle()
+        composeTestRule.runOnIdle {
+            assertNull(currentRoute)
+            isUserUiBlocked.value = false
+        }
+        composeTestRule.waitForIdle()
+        composeTestRule.runOnIdle {
+            assertNull(currentRoute)
+            currentController.graph = currentController.createGraph(startDestination = AUTH_ROUTE) {
+                composable(AUTH_ROUTE) {}
+            }
+        }
+        composeTestRule.waitForIdle()
+        composeTestRule.runOnIdle {
+            assertEquals(AUTH_ROUTE, currentRoute)
+        }
+    }
+
+    private companion object {
+        const val AUTH_ROUTE = "authentication"
+        const val SESSION_ROUTE = "session"
+    }
 }

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

@@ -39,7 +39,8 @@
 import androidx.compose.runtime.CompositionLocalProvider
 import androidx.compose.runtime.DisposableEffect
 import androidx.compose.runtime.LaunchedEffect
-import androidx.compose.runtime.derivedStateOf
+import androidx.compose.runtime.Stable
+import androidx.compose.runtime.State
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.key
 import androidx.compose.runtime.mutableStateOf
@@ -57,12 +58,15 @@
 import androidx.core.splashscreen.SplashScreen.Companion.installSplashScreen
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.ViewModel
+import androidx.lifecycle.ViewModelStore
+import androidx.lifecycle.ViewModelStoreOwner
 import androidx.lifecycle.flowWithLifecycle
 import androidx.lifecycle.lifecycleScope
 import androidx.lifecycle.repeatOnLifecycle
 import androidx.lifecycle.viewmodel.compose.viewModel
 import androidx.lifecycle.viewmodel.initializer
 import androidx.lifecycle.viewmodel.viewModelFactory
+import androidx.navigation.NavBackStackEntry
 import androidx.navigation.NavController
 import androidx.navigation.compose.currentBackStackEntryAsState
 import com.ramcosta.composedestinations.generated.app.destinations.CreateAccountCodeScreenDestination
@@ -346,7 +350,7 @@
         sessionGraphStore: SessionGraphStoreViewModel = viewModel(
             factory = viewModelFactory {
                 initializer {
-                    SessionGraphStoreViewModel(appGraph)
+                    SessionGraphStoreViewModel(appGraph::createSessionViewModelGraph)
                 }
             }
         ),
@@ -393,7 +397,7 @@
             finish = this@WireActivity::finish,
             isAllowedToNavigate = ::isNavigationAllowed
         )
-        val currentBackStackEntryState = navigator.navController.currentBackStackEntryAsState()
+        val currentBackStackEntryState = wireActivityCurrentBackStackEntryAsState(navigator)
         val currentBaseRoute = currentBackStackEntryState.value
             ?.destination
             ?.route
@@ -404,6 +408,7 @@
             ?.sessionBackedAuthenticationUserId()
         val isAuthenticationRoute = currentBaseRoute in authenticationGraphRoutes
         val isSessionTransitionInProgress = viewModel.globalAppState.isSessionTransitionInProgress
+        val sessionTransitionReason = viewModel.globalAppState.sessionTransitionReason
         var noSessionAuthenticationStartedWithoutSession by remember { mutableStateOf(false) }
         LaunchedEffect(currentBaseRoute, currentUserId) {
             when {
@@ -437,6 +442,17 @@
         if (graphContext?.sessionGraph != null) {
             lastSessionGraphContext.value = graphContext
         }
+        val shouldInvalidateRetainedSessionGraph = shouldInvalidateWireActivitySessionGraph(
+            isUserUiBlocked = isUserUiBlocked,
+            sessionTransitionReason = sessionTransitionReason,
+        )
+        LaunchedEffect(shouldInvalidateRetainedSessionGraph) {
+            lastSessionGraphContext.value = invalidateRetainedSessionGraphIfNeeded(
+                shouldInvalidate = shouldInvalidateRetainedSessionGraph,
+                lastSessionGraphContext = lastSessionGraphContext.value,
+                sessionGraphStore = sessionGraphStore,
+            )
+        }
         val graphContextWithRetainedImageLoader = graphContext?.copy(
             imageLoaderSessionGraph = resolveWireActivityImageLoaderSessionGraph(
                 activeSessionGraph = graphContext.sessionGraph,
@@ -449,12 +465,8 @@
             currentBaseRoute != null && canRetainSessionGraphForContent -> lastSessionGraphContext.value
             else -> null
         }
-        val backgroundType by remember {
-            derivedStateOf {
-                currentBackStackEntryState.value?.safeDestination()?.style.let {
-                    (it as? BackgroundStyle)?.backgroundType() ?: BackgroundType.Default
-                }
-            }
+        val backgroundType = currentBackStackEntryState.value?.safeDestination()?.style.let {
+            (it as? BackgroundStyle)?.backgroundType() ?: BackgroundType.Default
         }
 
         HandleSessionGraphEffects(
@@ -489,6 +501,17 @@
         noSessionAuthenticationStartedWithoutSession = noSessionAuthenticationStartedWithoutSession,
     )
 
+    private fun invalidateRetainedSessionGraphIfNeeded(
+        shouldInvalidate: Boolean,
+        lastSessionGraphContext: WireActivityGraphContext?,
+        sessionGraphStore: SessionGraphStoreViewModel,
+    ): WireActivityGraphContext? {
+        if (!shouldInvalidate) return lastSessionGraphContext
+
+        sessionGraphStore.invalidateActive()
+        return null
+    }
+
     private fun isNavigationAllowed(navigationCommand: NavigationCommand): Boolean {
         if (navigationCommand.destination.baseRoute != NewLoginScreenDestination.baseRoute) return true
 
@@ -614,7 +637,7 @@
         val usesAuthenticationGraph = effectiveBaseRoute in authenticationGraphRoutes
         val usesSessionBackedAuthenticationGraph = effectiveBaseRoute in sessionBackedAuthenticationGraphRoutes
         val usesInvalidSessionBackedAuthenticationGraph = usesSessionBackedAuthenticationGraph && currentUserId == null
-        val sessionGraph = remember(
+        val retainedSessionGraph = remember(
             appGraph,
             currentUserId,
             sessionBackedAuthenticationUserId,
@@ -630,6 +653,7 @@
                 isSessionTransitionInProgress = isSessionTransitionInProgress,
             )
         }
+        val sessionGraph = retainedSessionGraph?.graph
         val graph = resolveActiveGraph(
             WireActivityActiveGraphRequest(
                 authenticationViewModelGraph = authenticationViewModelGraph,
@@ -642,7 +666,7 @@
                 isSessionTransitionInProgress = isSessionTransitionInProgress,
             )
         )
-        val activityViewModels = sessionGraph?.let {
+        val activityViewModels = retainedSessionGraph?.let {
             wireActivityScopedViewModels(it)
         }
         return graph?.let {
@@ -662,12 +686,12 @@
         usesNoSessionAuthenticationGraph: Boolean,
         usesInvalidSessionBackedAuthenticationGraph: Boolean,
         isSessionTransitionInProgress: Boolean,
-    ): AppSessionViewModelGraph? = when {
+    ): RetainedSessionGraph? = when {
         usesNoSessionAuthenticationGraph -> null
         usesInvalidSessionBackedAuthenticationGraph -> null
         isSessionTransitionInProgress -> null
-        sessionBackedAuthenticationUserId != null -> graphFor(sessionBackedAuthenticationUserId)
-        currentUserId != null -> graphFor(currentUserId)
+        sessionBackedAuthenticationUserId != null -> retainedGraphFor(sessionBackedAuthenticationUserId)
+        currentUserId != null -> retainedGraphFor(currentUserId)
         else -> null
     }
 
@@ -746,10 +770,12 @@
     }
 
     @Composable
-    private fun wireActivityScopedViewModels(graph: AppSessionViewModelGraph): WireActivityScopedViewModels {
+    private fun wireActivityScopedViewModels(retainedSessionGraph: RetainedSessionGraph): WireActivityScopedViewModels {
+        val graph = retainedSessionGraph.graph
         val scopeKey = graph.viewModelScopeKey
         return WireActivityScopedViewModels(
             callFeedbackViewModel = viewModel(
+                viewModelStoreOwner = retainedSessionGraph,
                 key = "CallFeedbackViewModel:$scopeKey",
                 factory = viewModelFactory {
                     initializer {
@@ -758,6 +784,7 @@
                 }
             ),
             featureFlagNotificationViewModel = viewModel(
+                viewModelStoreOwner = retainedSessionGraph,
                 key = "FeatureFlagNotificationViewModel:$scopeKey",
                 factory = viewModelFactory {
                     initializer {
@@ -766,6 +793,7 @@
                 }
             ),
             commonTopAppBarViewModel = viewModel(
+                viewModelStoreOwner = retainedSessionGraph,
                 key = "CommonTopAppBarViewModel:$scopeKey",
                 factory = viewModelFactory {
                     initializer {
@@ -776,6 +804,7 @@
                 }
             ),
             legalHoldRequestedViewModel = viewModel(
+                viewModelStoreOwner = retainedSessionGraph,
                 key = "LegalHoldRequestedViewModel:$scopeKey",
                 factory = viewModelFactory {
                     initializer {
@@ -784,6 +813,7 @@
                 }
             ),
             legalHoldDeactivatedViewModel = viewModel(
+                viewModelStoreOwner = retainedSessionGraph,
                 key = "LegalHoldDeactivatedViewModel:$scopeKey",
                 factory = viewModelFactory {
                     initializer {
@@ -1276,6 +1306,15 @@
     if (isLocked) userId else null
 }.filterNotNull()
 
+@Composable
+internal fun wireActivityCurrentBackStackEntryAsState(
+    navigator: Navigator,
+): State<NavBackStackEntry?> = key(navigator.navController) {
+    // collectAsState retains its previous value while changing flows, so reset its composition
+    // identity with the controller to avoid exposing a route from the discarded navigation graph.
+    navigator.navController.currentBackStackEntryAsState()
+}
+
 private data class WireActivityScopedViewModels(
     val callFeedbackViewModel: CallFeedbackViewModel,
     val featureFlagNotificationViewModel: FeatureFlagNotificationViewModel,
@@ -1292,16 +1331,51 @@
     val activityViewModels: WireActivityScopedViewModels?,
 )
 
-private class SessionGraphStoreViewModel(
-    private val appGraph: WireApplicationGraph,
+internal class SessionGraphStoreViewModel(
+    private val createSessionGraph: (UserId) -> AppSessionViewModelGraph,
 ) : ViewModel() {
-    private val sessionGraphs = mutableMapOf<UserId, AppSessionViewModelGraph>()
+    private val sessionGraphs = mutableMapOf<UserId, RetainedSessionGraph>()
+    private var activeUserId: UserId? = null
 
-    fun graphFor(userId: UserId): AppSessionViewModelGraph =
-        sessionGraphs.getOrPut(userId) {
+    fun retainedGraphFor(userId: UserId): RetainedSessionGraph {
+        activeUserId = userId
+        return sessionGraphs.getOrPut(userId) {
             appLogger.i("WireActivity creating lifecycle-retained session graph for $userId")
-            appGraph.createSessionViewModelGraph(userId)
+            RetainedSessionGraph(createSessionGraph(userId))
         }
+    }
+
+    fun invalidateActive() {
+        activeUserId?.let(::invalidate)
+    }
+
+    fun invalidate(userId: UserId) {
+        sessionGraphs.remove(userId)?.also {
+            appLogger.i("WireActivity clearing lifecycle-retained session graph for $userId")
+            it.clear()
+        }
+        if (activeUserId == userId) {
+            activeUserId = null
+        }
+    }
+
+    override fun onCleared() {
+        sessionGraphs.values.forEach(RetainedSessionGraph::clear)
+        sessionGraphs.clear()
+        activeUserId = null
+    }
+}
+
+@Stable
+internal class RetainedSessionGraph(
+    val graph: AppSessionViewModelGraph,
+    override val viewModelStore: ViewModelStore = ViewModelStore(),
+) : ViewModelStoreOwner {
+
+    fun clear() {
+        viewModelStore.clear()
+        graph.wireSessionImageLoader.shutdown()
+    }
 }
 
 internal data class WireActivityActiveGraphRequest(
@@ -1414,6 +1488,11 @@
     retainedSessionGraph: AppSessionViewModelGraph?,
 ): AppSessionViewModelGraph? = activeSessionGraph ?: retainedSessionGraph
 
+internal fun shouldInvalidateWireActivitySessionGraph(
+    isUserUiBlocked: Boolean,
+    sessionTransitionReason: SessionTransitionReason?,
+): Boolean = isUserUiBlocked || sessionTransitionReason != null
+
 private fun Bundle.sessionBackedAuthenticationUserId(): UserId? {
     val value = getString(SessionBackedAuthenticationNavArgs.USER_ID_VALUE_KEY)
     val domain = getString(SessionBackedAuthenticationNavArgs.USER_ID_DOMAIN_KEY)

app/src/test/kotlin/com/wire/android/feature/AccountSwitchUseCaseTest.kt

@@ -37,6 +37,9 @@ import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertEquals
 import org.junit.jupiter.api.Test
+import org.junit.jupiter.params.ParameterizedTest
+import org.junit.jupiter.params.provider.MethodSource
+import java.util.stream.Stream
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class AccountSwitchUseCaseTest {
@@ -86,9 +89,12 @@ class AccountSwitchUseCaseTest {
             }
         }
 
-    @Test
-    fun givenCurrentSessionIsInvalid_whenSwitchingToAccount_thenUpdateCurrentSessionAndDeleteTheOldOne() = testScope.runTest {
-        val currentAccount = ACCOUNT_INVALID_3
+    @ParameterizedTest
+    @MethodSource("terminalLogoutReasons")
+    fun givenCurrentSessionIsInvalid_whenSwitchingToAccount_thenUpdateCurrentSessionAndDeleteTheOldOne(
+        logoutReason: LogoutReason
+    ) = testScope.runTest {
+        val currentAccount = AccountInfo.Invalid(ACCOUNT_INVALID_3.userId, logoutReason)
         val switchTo = ACCOUNT_VALID_2
 
         val expectedResult = SwitchAccountResult.SwitchedToAnotherAccount
@@ -151,6 +157,16 @@ class AccountSwitchUseCaseTest {
         val ACCOUNT_VALID_2 = AccountInfo.Valid(UserId("userId_valid_2", "domain_valid_2"))
         val ACCOUNT_INVALID_3 =
             AccountInfo.Invalid(UserId("userId_invalid_3", "domain_invalid_3"), LogoutReason.SELF_SOFT_LOGOUT)
+
+        @JvmStatic
+        fun terminalLogoutReasons(): Stream<LogoutReason> = Stream.of(
+            LogoutReason.SELF_SOFT_LOGOUT,
+            LogoutReason.SELF_HARD_LOGOUT,
+            LogoutReason.MIGRATION_TO_CC_FAILED,
+            LogoutReason.DELETED_ACCOUNT,
+            LogoutReason.REMOVED_CLIENT,
+            LogoutReason.SESSION_EXPIRED,
+        )
     }
 
     private class Arrangement(val testScope: TestScope) {