Issue #5529: Restore engine state of tabs in tab collections.

components/browser/session/src/main/java/mozilla/components/browser/session/LegacySessionManager.kt

@@ -137,9 +137,10 @@ class LegacySessionManager(
         session: Session,
         selected: Boolean = false,
         engineSession: EngineSession? = null,
+        engineSessionState: EngineSessionState? = null,
         parent: Session? = null
     ) = synchronized(values) {
-        addInternal(session, selected, engineSession, parent = parent, viaRestore = false)
+        addInternal(session, selected, engineSession, engineSessionState, parent = parent, viaRestore = false)
     }
 
     @Suppress("LongParameterList", "ComplexMethod")

components/browser/session/src/main/java/mozilla/components/browser/session/SessionManager.kt

@@ -147,10 +147,12 @@ class SessionManager(
         session: Session,
         selected: Boolean = false,
         engineSession: EngineSession? = null,
+        engineSessionState: EngineSessionState? = null,
         parent: Session? = null
     ) {
         // Add store to Session so that it can dispatch actions whenever it changes.
         session.store = store
+
         if (parent != null) {
             require(all.contains(parent)) { "The parent does not exist" }
             session.parentId = parent.id
@@ -171,7 +173,17 @@ class SessionManager(
             )
         }
 
-        delegate.add(session, selected, engineSession, parent)
+        if (engineSessionState != null && engineSession == null) {
+            // If the caller passed us an engine session state then also notify the store. We only
+            // do this if there is no engine session, because this mirrors the behavior in
+            // LegacySessionManager, which will either link an engine session or keep its state.
+            store?.syncDispatch(UpdateEngineSessionStateAction(
+                session.id,
+                engineSessionState
+            ))
+        }
+
+        delegate.add(session, selected, engineSession, engineSessionState, parent)
     }
 
     /**

components/browser/session/src/test/java/mozilla/components/browser/session/SessionManagerMigrationTest.kt

@@ -1281,6 +1281,58 @@ class SessionManagerMigrationTest {
         verify(engineSession3, never()).close()
         verify(engineSession4, never()).close()
     }
+
+    @Test
+    fun `Adding session with engine session state`() {
+        val store = BrowserStore()
+        val manager = SessionManager(engine = mock(), store = store)
+
+        val state: EngineSessionState = mock()
+        val session = Session("https://www.mozilla.org")
+
+        manager.add(session, engineSessionState = state)
+
+        assertEquals(state, session.engineSessionHolder.engineSessionState)
+        assertNull(session.engineSessionHolder.engineSession)
+
+        assertEquals(state, store.state.tabs[0].engineState.engineSessionState)
+        assertNull(store.state.tabs[0].engineState.engineSession)
+    }
+
+    @Test
+    fun `Adding session with engine session and engine session state`() {
+        val store = BrowserStore()
+        val manager = SessionManager(engine = mock(), store = store)
+
+        val state: EngineSessionState = mock()
+        val engineSession: EngineSession = mock()
+        val session = Session("https://www.mozilla.org")
+
+        manager.add(session, engineSession = engineSession, engineSessionState = state)
+
+        assertNull(session.engineSessionHolder.engineSessionState)
+        assertEquals(engineSession, session.engineSessionHolder.engineSession)
+
+        assertNull(store.state.tabs[0].engineState.engineSessionState)
+        assertEquals(engineSession, session.engineSessionHolder.engineSession)
+    }
+
+    @Test
+    fun `Adding session with engine session`() {
+        val store = BrowserStore()
+        val manager = SessionManager(engine = mock(), store = store)
+
+        val engineSession: EngineSession = mock()
+        val session = Session("https://www.mozilla.org")
+
+        manager.add(session, engineSession = engineSession)
+
+        assertNull(session.engineSessionHolder.engineSessionState)
+        assertEquals(engineSession, session.engineSessionHolder.engineSession)
+
+        assertNull(store.state.tabs[0].engineState.engineSessionState)
+        assertEquals(engineSession, session.engineSessionHolder.engineSession)
+    }
 }
 
 private fun createMockEngineSessionWithState(): EngineSession {

components/feature/customtabs/src/test/java/mozilla/components/feature/customtabs/CustomTabIntentProcessorTest.kt

@@ -63,7 +63,7 @@ class CustomTabIntentProcessorTest {
         whenever(intent.putExtra(any<String>(), any<String>())).thenReturn(intent)
 
         handler.process(intent)
-        verify(sessionManager).add(anySession(), eq(false), eq(null), eq(null))
+        verify(sessionManager).add(anySession(), eq(false), eq(null), eq(null), eq(null))
         verify(engineSession).loadUrl("http://mozilla.org", flags = LoadUrlFlags.external())
         verify(intent).putExtra(eq(EXTRA_SESSION_ID), any<String>())
 
@@ -92,7 +92,7 @@ class CustomTabIntentProcessorTest {
         whenever(intent.putExtra(any<String>(), any<String>())).thenReturn(intent)
 
         handler.process(intent)
-        verify(sessionManager).add(anySession(), eq(false), eq(null), eq(null))
+        verify(sessionManager).add(anySession(), eq(false), eq(null), eq(null), eq(null))
         verify(engineSession).loadUrl("http://mozilla.org", flags = LoadUrlFlags.external())
         verify(intent).putExtra(eq(EXTRA_SESSION_ID), any<String>())
 

components/feature/search/src/test/java/mozilla/components/feature/search/SearchUseCasesTest.kt

@@ -107,7 +107,7 @@ class SearchUseCasesTest {
         useCases.newPrivateTabSearch.invoke(searchTerms)
 
         val captor = argumentCaptor<Session>()
-        verify(sessionManager).add(captor.capture(), eq(true), eq(null), eq(null))
+        verify(sessionManager).add(captor.capture(), eq(true), eq(null), eq(null), eq(null))
         assertTrue(captor.value.private)
         verify(engineSession).loadUrl(searchUrl)
     }

components/feature/tab-collections/src/androidTest/java/mozilla/components/feature/tab/collections/TabCollectionStorageTest.kt

@@ -191,8 +191,8 @@ class TabCollectionStorageTest {
             assertEquals(session1, sessions[0])
             assertEquals(session2, sessions[1])
 
-            assertEquals(session1.id, sessions[0].id)
-            assertEquals(session2.id, sessions[1].id)
+            assertEquals(session1.id, sessions[0].session.id)
+            assertEquals(session2.id, sessions[1].session.id)
         }
 
         getAllCollections().let { collections ->
@@ -207,14 +207,14 @@ class TabCollectionStorageTest {
             assertNotEquals(session1, sessions[0])
             assertNotEquals(session2, sessions[1])
 
-            assertNotEquals(session1.id, sessions[0].id)
-            assertNotEquals(session2.id, sessions[1].id)
+            assertNotEquals(session1.id, sessions[0].session.id)
+            assertNotEquals(session2.id, sessions[1].session.id)
 
-            assertEquals(session1.url, sessions[0].url)
-            assertEquals(session2.url, sessions[1].url)
+            assertEquals(session1.url, sessions[0].session.url)
+            assertEquals(session2.url, sessions[1].session.url)
 
-            assertEquals(session1.title, sessions[0].title)
-            assertEquals(session2.title, sessions[1].title)
+            assertEquals(session1.title, sessions[0].session.title)
+            assertEquals(session2.title, sessions[1].session.title)
         }
     }
 

components/feature/tab-collections/src/main/java/mozilla/components/feature/tab/collections/Tab.kt

@@ -6,6 +6,7 @@ package mozilla.components.feature.tab.collections
 
 import android.content.Context
 import mozilla.components.browser.session.Session
+import mozilla.components.browser.session.SessionManager
 import mozilla.components.concept.engine.Engine
 
 /**
@@ -37,7 +38,6 @@ interface Tab {
     fun restore(
         context: Context,
         engine: Engine,
-        tab: Tab,
         restoreSessionId: Boolean = false
-    ): Session?
+    ): SessionManager.Snapshot.Item?
 }

components/feature/tab-collections/src/main/java/mozilla/components/feature/tab/collections/TabCollection.kt

@@ -6,6 +6,7 @@ package mozilla.components.feature.tab.collections
 
 import android.content.Context
 import mozilla.components.browser.session.Session
+import mozilla.components.browser.session.SessionManager
 import mozilla.components.concept.engine.Engine
 
 /**
@@ -28,7 +29,8 @@ interface TabCollection {
     val tabs: List<Tab>
 
     /**
-     * Restores all tabs in this collection and returns a matching list of [Session] objects.
+     * Restores all tabs in this collection and returns a matching list of
+     * [SessionManager.Snapshot.Item] objects.
      *
      * @param restoreSessionId If true the original [Session.id] of [Session]s will be restored. Otherwise a new ID
      * will be generated. An app may prefer to use a new ID if it expects sessions to get restored multiple times -
@@ -38,10 +40,11 @@ interface TabCollection {
         context: Context,
         engine: Engine,
         restoreSessionId: Boolean = false
-    ): List<Session>
+    ): List<SessionManager.Snapshot.Item>
 
     /**
-     * Restores a subset of the tabs in this collection and returns a matching list of [Session] objects.
+     * Restores a subset of the tabs in this collection and returns a matching list of
+     * [SessionManager.Snapshot.Item] objects.
      *
      * @param restoreSessionId If true the original [Session.id] of [Session]s will be restored. Otherwise a new ID
      * will be generated. An app may prefer to use a new ID if it expects sessions to get restored multiple times -
@@ -52,5 +55,5 @@ interface TabCollection {
         engine: Engine,
         tabs: List<Tab>,
         restoreSessionId: Boolean = false
-    ): List<Session>
+    ): List<SessionManager.Snapshot.Item>
 }

components/feature/tab-collections/src/main/java/mozilla/components/feature/tab/collections/adapter/TabAdapter.kt

@@ -5,7 +5,6 @@
 package mozilla.components.feature.tab.collections.adapter
 
 import android.content.Context
-import mozilla.components.browser.session.Session
 import mozilla.components.browser.session.SessionManager
 import mozilla.components.browser.session.ext.readSnapshotItem
 import mozilla.components.concept.engine.Engine
@@ -24,18 +23,13 @@ internal class TabAdapter(
     override val url: String
         get() = entity.url
 
-    /**
-     * Restores a single tab from this collection and returns a matching [SessionManager.Snapshot].
-     */
     override fun restore(
         context: Context,
         engine: Engine,
-        tab: Tab,
         restoreSessionId: Boolean
-    ): Session? {
+    ): SessionManager.Snapshot.Item? {
         return entity.getStateFile(context.filesDir)
             .readSnapshotItem(engine, restoreSessionId, false)
-            ?.session
     }
 
     override fun equals(other: Any?): Boolean {

components/feature/tab-collections/src/main/java/mozilla/components/feature/tab/collections/adapter/TabCollectionAdapter.kt

@@ -5,7 +5,7 @@
 package mozilla.components.feature.tab.collections.adapter
 
 import android.content.Context
-import mozilla.components.browser.session.Session
+import mozilla.components.browser.session.SessionManager
 import mozilla.components.browser.session.ext.readSnapshotItem
 import mozilla.components.concept.engine.Engine
 import mozilla.components.feature.tab.collections.Tab
@@ -29,7 +29,11 @@ internal class TabCollectionAdapter(
     override val id: Long
         get() = entity.collection.id!!
 
-    override fun restore(context: Context, engine: Engine, restoreSessionId: Boolean): List<Session> {
+    override fun restore(
+        context: Context,
+        engine: Engine,
+        restoreSessionId: Boolean
+    ): List<SessionManager.Snapshot.Item> {
         return restore(context, engine, entity.tabs, restoreSessionId)
     }
 
@@ -38,7 +42,7 @@ internal class TabCollectionAdapter(
         engine: Engine,
         tabs: List<Tab>,
         restoreSessionId: Boolean
-    ): List<Session> {
+    ): List<SessionManager.Snapshot.Item> {
         val entities = entity.tabs.filter {
             candidate -> tabs.find { tab -> tab.id == candidate.id } != null
         }
@@ -50,9 +54,9 @@ internal class TabCollectionAdapter(
         engine: Engine,
         tabs: List<TabEntity>,
         restoreSessionId: Boolean
-    ): List<Session> {
+    ): List<SessionManager.Snapshot.Item> {
         return tabs.mapNotNull { tab ->
-            tab.getStateFile(context.filesDir).readSnapshotItem(engine, restoreSessionId)?.session
+            tab.getStateFile(context.filesDir).readSnapshotItem(engine, restoreSessionId)
         }
     }
 

components/feature/tab-collections/src/main/java/mozilla/components/feature/tab/collections/ext/SessionManager.kt

+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+package mozilla.components.feature.tab.collections.ext
+
+import android.content.Context
+import mozilla.components.browser.session.SessionManager
+import mozilla.components.concept.engine.Engine
+import mozilla.components.feature.tab.collections.Tab
+import mozilla.components.feature.tab.collections.TabCollection
+
+/**
+ * Restores the given [Tab] from a [TabCollection]. Will invoke [onTabRestored] on successful restore
+ * and [onFailure] otherwise.
+ */
+fun SessionManager.restore(
+    context: Context,
+    engine: Engine,
+    tab: Tab,
+    onTabRestored: () -> Unit,
+    onFailure: () -> Unit
+) {
+    val item = tab.restore(
+        context = context,
+        engine = engine,
+        restoreSessionId = false
+    )
+
+    if (item == null) {
+        // We were unable to restore the tab. Let the app know so that it can workaround that
+        onFailure()
+    } else {
+        add(
+            session = item.session,
+            selected = true,
+            engineSessionState = item.engineSessionState
+        )
+        onTabRestored()
+    }
+}
+
+/**
+ * Restores the given [TabCollection]. Will invoke [onFailure] if restoring a single [Tab] of the
+ * collection failed. The URL of the tab will be passed to [onFailure].
+ */
+fun SessionManager.restore(
+    context: Context,
+    engine: Engine,
+    collection: TabCollection,
+    onFailure: (String) -> Unit
+) {
+    collection.tabs.reversed().forEach { tab ->
+        val item = tab.restore(
+            context = context,
+            engine = engine,
+            restoreSessionId = false
+        )
+
+        if (item == null) {
+            // We were unable to restore the tab. Let the app know so that it can workaround that
+            onFailure(tab.url)
+        } else {
+            add(
+                session = item.session,
+                selected = selectedSession == null,
+                engineSessionState = item.engineSessionState
+            )
+        }
+    }
+}