Remove runBlocking call to restore sessions when the app starts

appnav/src/main/kotlin/io/element/android/appnav/ReplaceAllOperation.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) 2026 Element Creations Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial.
+ * Please see LICENSE files in the repository root for full details.
+ */
+
+package io.element.android.appnav
+
+import com.bumble.appyx.core.navigation.NavElements
+import com.bumble.appyx.core.navigation.Operation
+import com.bumble.appyx.navmodel.backstack.BackStack
+import kotlinx.parcelize.Parcelize
+
+/**
+ * Replaces all the current elements with the provided [navElements], keeping their [BackStack.State] too.
+ */
+@Parcelize
+class ReplaceAllOperation<NavTarget : Any>(
+    private val navElements: NavElements<NavTarget, BackStack.State>
+) : Operation<NavTarget, BackStack.State> {
+    override fun isApplicable(elements: NavElements<NavTarget, BackStack.State>): Boolean {
+        return true
+    }
+
+    override fun invoke(existing: NavElements<NavTarget, BackStack.State>): NavElements<NavTarget, BackStack.State> {
+        return navElements
+    }
+}

appnav/src/main/kotlin/io/element/android/appnav/RootFlowNode.kt

@@ -16,9 +16,12 @@ import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
 import androidx.lifecycle.lifecycleScope
 import com.bumble.appyx.core.modality.BuildContext
+import com.bumble.appyx.core.navigation.NavElements
+import com.bumble.appyx.core.navigation.NavKey
 import com.bumble.appyx.core.node.Node
 import com.bumble.appyx.core.plugin.Plugin
 import com.bumble.appyx.core.state.MutableSavedStateMap
+import com.bumble.appyx.core.state.SavedStateMap
 import com.bumble.appyx.navmodel.backstack.BackStack
 import com.bumble.appyx.navmodel.backstack.operation.pop
 import com.bumble.appyx.navmodel.backstack.operation.push
@@ -50,6 +53,7 @@ import io.element.android.libraries.architecture.createNode
 import io.element.android.libraries.architecture.waitForChildAttached
 import io.element.android.libraries.core.uri.ensureProtocol
 import io.element.android.libraries.deeplink.api.DeeplinkData
+import io.element.android.libraries.di.annotations.AppCoroutineScope
 import io.element.android.libraries.featureflag.api.FeatureFlagService
 import io.element.android.libraries.featureflag.api.FeatureFlags
 import io.element.android.libraries.matrix.api.core.EventId
@@ -67,7 +71,9 @@ import io.element.android.services.analytics.api.AnalyticsLongRunningTransaction
 import io.element.android.services.analytics.api.AnalyticsService
 import io.element.android.services.analytics.api.watchers.AnalyticsColdStartWatcher
 import io.element.android.services.appnavstate.api.ROOM_OPENED_FROM_NOTIFICATION
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.drop
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.launch
@@ -93,19 +99,27 @@ class RootFlowNode(
     private val announcementService: AnnouncementService,
     private val analyticsService: AnalyticsService,
     private val analyticsColdStartWatcher: AnalyticsColdStartWatcher,
+    @AppCoroutineScope private val appCoroutineScope: CoroutineScope,
 ) : BaseFlowNode<RootFlowNode.NavTarget>(
     backstack = BackStack(
         initialElement = NavTarget.SplashScreen,
-        savedStateMap = buildContext.savedStateMap,
+        savedStateMap = null,
     ),
     buildContext = buildContext,
     plugins = plugins
 ) {
     override fun onBuilt() {
         analyticsColdStartWatcher.start()
-        matrixSessionCache.restoreWithSavedState(buildContext.savedStateMap)
+        appCoroutineScope.launch {
+            matrixSessionCache.restoreWithSavedState(buildContext.savedStateMap)
+            if (buildContext.savedStateMap != null) {
+                restoreSavedState(buildContext.savedStateMap)
+                observeNavState(true)
+            } else {
+                observeNavState(false)
+            }
+        }
         super.onBuilt()
-        observeNavState()
     }
 
     override fun onSaveInstanceState(state: MutableSavedStateMap) {
@@ -114,25 +128,68 @@ class RootFlowNode(
         navStateFlowFactory.saveIntoSavedState(state)
     }
 
-    private fun observeNavState() {
-        navStateFlowFactory.create(buildContext.savedStateMap).distinctUntilChanged().onEach { navState ->
-            Timber.v("navState=$navState")
-            when (navState.loggedInState) {
-                is LoggedInState.LoggedIn -> {
-                    if (navState.loggedInState.isTokenValid) {
-                        tryToRestoreLatestSession(
-                            onSuccess = { sessionId -> switchToLoggedInFlow(sessionId, navState.cacheIndex) },
-                            onFailure = { switchToNotLoggedInFlow(null) }
-                        )
-                    } else {
-                        switchToSignedOutFlow(SessionId(navState.loggedInState.sessionId))
+    private fun observeNavState(skipFirst: Boolean) {
+        navStateFlowFactory.create(buildContext.savedStateMap)
+            .distinctUntilChanged()
+            .drop(if (skipFirst) 1 else 0)
+            .onEach { navState ->
+                Timber.v("navState=$navState")
+                when (navState.loggedInState) {
+                    is LoggedInState.LoggedIn -> {
+                        if (navState.loggedInState.isTokenValid) {
+                            val sessionId = SessionId(navState.loggedInState.sessionId)
+                            if (matrixSessionCache.getOrNull(sessionId) != null) {
+                                switchToLoggedInFlow(sessionId, navState.cacheIndex)
+                            } else {
+                                tryToRestoreLatestSession(
+                                    onSuccess = { sessionId -> switchToLoggedInFlow(sessionId, navState.cacheIndex) },
+                                    onFailure = { switchToNotLoggedInFlow(null) }
+                                )
+                            }
+                        } else {
+                            switchToSignedOutFlow(SessionId(navState.loggedInState.sessionId))
+                        }
+                    }
+                    LoggedInState.NotLoggedIn -> {
+                        switchToNotLoggedInFlow(null)
                     }
-                }
-                LoggedInState.NotLoggedIn -> {
-                    switchToNotLoggedInFlow(null)
                 }
             }
-        }.launchIn(lifecycleScope)
+            .launchIn(lifecycleScope)
+    }
+
+    /**
+     * Restore the saved state for navigation in the current backstack.
+     *
+     * **WARNING:** this is an unsafe operation abusing the internals of the Appyx library, but it's the only way allow async state
+     * restoration and not having to block the main thread when the app starts.
+     *
+     * Modify with utmost care and double check any possible Appyx updates that might break this.
+     */
+    @Suppress("UNCHECKED_CAST")
+    private fun restoreSavedState(savedStateMap: SavedStateMap?) {
+        if (savedStateMap == null) return
+
+        // 'NavModel' is the key used for storing the nav model state data in the map in Appyx
+        val savedElements = buildContext.savedStateMap?.get("NavModel") as? NavElements<NavTarget, BackStack.State>
+        if (savedElements != null) {
+            backstack.accept(ReplaceAllOperation(savedElements))
+        }
+    }
+
+    /**
+     * Extract the saved state for navigation in the [navTarget].
+     *
+     * **WARNING:** this is an unsafe operation abusing the internals of the Appyx library, but it's the only way allow async state
+     * restoration and not having to block the main thread when the app starts.
+     *
+     * Modify with utmost care and double check any possible Appyx updates that might break this.
+     */
+    @Suppress("UNCHECKED_CAST")
+    private fun extractSavedStateForNavTarget(navTarget: NavTarget, savedStateMap: SavedStateMap?): SavedStateMap? {
+        // 'ChildrenState' is the key used for storing the children state data in the map in Appyx
+        val childrenState = savedStateMap?.get("ChildrenState") as? Map<NavKey<NavTarget>, SavedStateMap> ?: return null
+        return childrenState.entries.find { (key, _) -> key.navTarget == navTarget }?.value
     }
 
     private fun switchToLoggedInFlow(sessionId: SessionId, navId: Int) {
@@ -244,6 +301,13 @@ class RootFlowNode(
                         backstack.push(NavTarget.NotLoggedInFlow(null))
                     }
                 }
+                val savedNavState = extractSavedStateForNavTarget(navTarget, this.buildContext.savedStateMap)
+                val buildContext = if (savedNavState != null) {
+                    Timber.d("Creating a $navTarget with restored saved state")
+                    buildContext.copy(savedStateMap = savedNavState)
+                } else {
+                    buildContext.copy(savedStateMap = savedNavState)
+                }
                 createNode<LoggedInAppScopeFlowNode>(buildContext, plugins = listOf(inputs, callback))
             }
             is NavTarget.NotLoggedInFlow -> {

appnav/src/main/kotlin/io/element/android/appnav/di/MatrixSessionCache.kt

@@ -21,7 +21,6 @@ import io.element.android.libraries.matrix.api.auth.MatrixAuthenticationService
 import io.element.android.libraries.matrix.api.core.SessionId
 import io.element.android.services.analytics.api.AnalyticsService
 import io.element.android.services.analyticsproviders.api.AnalyticsUserData
-import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
 import timber.log.Timber
@@ -77,20 +76,18 @@ class MatrixSessionCache(
     }
 
     @Suppress("UNCHECKED_CAST")
-    fun restoreWithSavedState(state: SavedStateMap?) {
+    suspend fun restoreWithSavedState(state: SavedStateMap?) {
         Timber.d("Restore state")
         if (state == null || sessionIdsToMatrixSession.isNotEmpty()) {
-            Timber.w("Restore with non-empty map")
+            Timber.w("No need to restore saved state")
             return
         }
         val sessionIds = state[SAVE_INSTANCE_KEY] as? Array<SessionId>
         Timber.d("Restore matrix session keys = ${sessionIds?.map { it.value }}")
         if (sessionIds.isNullOrEmpty()) return
         // Not ideal but should only happens in case of process recreation. This ensure we restore all the active sessions before restoring the node graphs.
-        runBlocking {
-            sessionIds.forEach { sessionId ->
-                getOrRestore(sessionId)
-            }
+        sessionIds.forEach { sessionId ->
+            getOrRestore(sessionId)
         }
     }
 

gradle/libs.versions.toml

@@ -42,6 +42,8 @@ serialization_json = "1.10.0"
 coil = "3.3.0"
 # Rollback to 1.0.4, 1.0.5 has this issue: https://github.com/airbnb/Showkase/issues/420
 showkase = "1.0.5"
+# There is some custom logic in `RootFlowNode` that may break because it reuses some Appyx internal APIs.
+# When upgrading this version, check state restoration still works fine.
 appyx = "1.7.1"
 sqldelight = "2.2.1"
 wysiwyg = "2.41.1"