Remove runBlocking call to restore sessions when the app starts (#6193)

* Remove `runBlocking` call to restore sessions when the app starts

Sadly, to do this we need to manually handle restoring the state from Appyx using internal values. At least it doesn't seem like they're going to change any time soon (or ever).

This should take care of a few ANRs, although it may make loading the initial state a bit slower

* Add `ReplaceAllOperation` for state restoration

* Add warning comment for Appyx dependency

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"