diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index b213203..358d703 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -148,7 +148,7 @@
         });
     }
 
-    private void processLoadedTasks(ArrayList<GroupTask> tasks) {
+    private void processLoadedTasks(List<GroupTask> tasks) {
         // Only store MAX_TASK tasks, from most to least recent
         Collections.reverse(tasks);
         mTasks = tasks.stream()
@@ -157,7 +157,7 @@
         mNumHiddenTasks = Math.max(0, tasks.size() - MAX_TASKS);
     }
 
-    private void processLoadedTasksOnDesktop(ArrayList<GroupTask> tasks) {
+    private void processLoadedTasksOnDesktop(List<GroupTask> tasks) {
         // Find the single desktop task that contains a grouping of desktop tasks
         DesktopTask desktopTask = findDesktopTask(tasks);
 
@@ -173,7 +173,7 @@
     }
 
     @Nullable
-    private DesktopTask findDesktopTask(ArrayList<GroupTask> tasks) {
+    private DesktopTask findDesktopTask(List<GroupTask> tasks) {
         return (DesktopTask) tasks.stream()
                 .filter(t -> t instanceof DesktopTask)
                 .findFirst()
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 711882c..37b4dca 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -31,6 +31,7 @@
 import android.os.RemoteException;
 import android.util.SparseBooleanArray;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.util.LooperExecutor;
@@ -44,6 +45,7 @@
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.List;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -137,7 +139,7 @@
      * @return The change id of the current task list
      */
     public synchronized int getTasks(boolean loadKeysOnly,
-            Consumer<ArrayList<GroupTask>> callback, Predicate<GroupTask> filter) {
+            @Nullable Consumer<List<GroupTask>> callback, Predicate<GroupTask> filter) {
         final int requestLoadId = mChangeId;
         if (mResultsUi.isValidForRequest(requestLoadId, loadKeysOnly)) {
             // The list is up to date, send the callback on the next frame,
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index 89351aa..98c1eb4 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -33,6 +33,7 @@
 import android.os.Process;
 import android.os.UserHandle;
 
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.icons.IconProvider;
@@ -40,6 +41,7 @@
 import com.android.launcher3.util.Executors.SimpleThreadFactory;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.SafeCloseable;
+import com.android.quickstep.recents.data.RecentTasksDataSource;
 import com.android.quickstep.util.GroupTask;
 import com.android.quickstep.util.TaskVisualsChangeListener;
 import com.android.systemui.shared.recents.model.Task;
@@ -60,8 +62,8 @@
  * Singleton class to load and manage recents model.
  */
 @TargetApi(Build.VERSION_CODES.O)
-public class RecentsModel implements IconChangeListener, TaskStackChangeListener,
-        TaskVisualsChangeListener, SafeCloseable {
+public class RecentsModel implements RecentTasksDataSource, IconChangeListener,
+        TaskStackChangeListener, TaskVisualsChangeListener, SafeCloseable {
 
     // We do not need any synchronization for this variable as its only written on UI thread.
     public static final MainThreadInitializedObject<RecentsModel> INSTANCE =
@@ -141,7 +143,8 @@
      *                always called on the UI thread.
      * @return the request id associated with this call.
      */
-    public int getTasks(Consumer<ArrayList<GroupTask>> callback) {
+    @Override
+    public int getTasks(@Nullable Consumer<List<GroupTask>> callback) {
         return mTaskList.getTasks(false /* loadKeysOnly */, callback,
                 RecentsFilterState.DEFAULT_FILTER);
     }
@@ -155,7 +158,7 @@
      *                callback.
      * @return the request id associated with this call.
      */
-    public int getTasks(Consumer<ArrayList<GroupTask>> callback, Predicate<GroupTask> filter) {
+    public int getTasks(@Nullable Consumer<List<GroupTask>> callback, Predicate<GroupTask> filter) {
         return mTaskList.getTasks(false /* loadKeysOnly */, callback, filter);
     }
 
diff --git a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
index 7ebb767..38e927f 100644
--- a/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
+++ b/quickstep/src/com/android/quickstep/TaskThumbnailCache.java
@@ -21,11 +21,13 @@
 import android.content.Context;
 import android.content.res.Resources;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.R;
 import com.android.launcher3.util.CancellableTask;
 import com.android.launcher3.util.Preconditions;
+import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource;
 import com.android.quickstep.util.TaskKeyByLastActiveTimeCache;
 import com.android.quickstep.util.TaskKeyCache;
 import com.android.quickstep.util.TaskKeyLruCache;
@@ -38,7 +40,7 @@
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
 
-public class TaskThumbnailCache {
+public class TaskThumbnailCache implements TaskThumbnailDataSource {
 
     private final Executor mBgExecutor;
     private final TaskKeyCache<ThumbnailData> mCache;
@@ -148,8 +150,9 @@
      * @param callback The callback to receive the task after its data has been populated.
      * @return A cancelable handle to the request
      */
+    @Override
     public CancellableTask<ThumbnailData> updateThumbnailInBackground(
-            Task task, Consumer<ThumbnailData> callback) {
+            Task task, @NonNull Consumer<ThumbnailData> callback) {
         Preconditions.assertUIThread();
 
         boolean lowResolution = !mHighResLoadingState.isEnabled();
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index 096ed2c..485d6c4 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -54,6 +54,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.List;
 
 public class FallbackRecentsView extends RecentsView<RecentsActivity, RecentsState>
         implements StateListener<RecentsState> {
@@ -179,7 +180,7 @@
     }
 
     @Override
-    protected void applyLoadPlan(ArrayList<GroupTask> taskGroups) {
+    protected void applyLoadPlan(List<GroupTask> taskGroups) {
         // When quick-switching on 3p-launcher, we add a "stub" tile corresponding to Launcher
         // as well. This tile is never shown as we have setCurrentTaskHidden, but allows use to
         // track the index of the next task appropriately, as if we are switching on any other app.
diff --git a/quickstep/src/com/android/quickstep/recents/data/RecentTasksDataSource.kt b/quickstep/src/com/android/quickstep/recents/data/RecentTasksDataSource.kt
new file mode 100644
index 0000000..6719099
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentTasksDataSource.kt
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.data
+
+import com.android.quickstep.util.GroupTask
+import java.util.function.Consumer
+
+interface RecentTasksDataSource {
+    fun getTasks(callback: Consumer<List<GroupTask>>?): Int
+}
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
new file mode 100644
index 0000000..ad8ae20
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -0,0 +1,114 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.data
+
+import com.android.quickstep.TaskIconCache
+import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource
+import com.android.quickstep.util.GroupTask
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import kotlin.coroutines.resume
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.suspendCancellableCoroutine
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class TasksRepository(
+    private val recentsModel: RecentTasksDataSource,
+    private val taskThumbnailDataSource: TaskThumbnailDataSource,
+    private val taskIconCache: TaskIconCache,
+) {
+    private val groupedTaskData = MutableStateFlow(emptyList<GroupTask>())
+    private val _taskData =
+        groupedTaskData.map { groupTaskList -> groupTaskList.flatMap { it.tasks } }
+    private val visibleTaskIds = MutableStateFlow(emptySet<Int>())
+
+    private val taskData: Flow<List<Task>> =
+        combine(_taskData, getThumbnailQueryResults()) { tasks, results ->
+            tasks.forEach { task ->
+                // Add retrieved thumbnails + remove unnecessary thumbnails
+                task.thumbnail = results[task.key.id]
+            }
+            tasks
+        }
+
+    fun getAllTaskData(forceRefresh: Boolean = false): Flow<List<Task>> {
+        if (forceRefresh) {
+            recentsModel.getTasks { groupedTaskData.value = it }
+        }
+        return taskData
+    }
+
+    fun getTaskDataById(taskId: Int): Flow<Task?> =
+        taskData.map { taskList -> taskList.firstOrNull { it.key.id == taskId } }
+
+    fun setVisibleTasks(visibleTaskIdList: List<Int>) {
+        this.visibleTaskIds.value = visibleTaskIdList.toSet()
+    }
+
+    /** Flow wrapper for [TaskThumbnailDataSource.updateThumbnailInBackground] api */
+    private fun getThumbnailDataRequest(task: Task): ThumbnailDataRequest =
+        flow {
+                emit(task.key.id to task.thumbnail)
+                val thumbnailDataResult: ThumbnailData? =
+                    suspendCancellableCoroutine { continuation ->
+                        val cancellableTask =
+                            taskThumbnailDataSource.updateThumbnailInBackground(task) {
+                                continuation.resume(it)
+                            }
+                        continuation.invokeOnCancellation { cancellableTask?.cancel() }
+                    }
+                emit(task.key.id to thumbnailDataResult)
+            }
+            .distinctUntilChanged()
+
+    /**
+     * This is a Flow that makes a query for thumbnail data to the [taskThumbnailDataSource] for
+     * each visible task. It then collects the responses and returns them in a Map as soon as they
+     * are available.
+     */
+    private fun getThumbnailQueryResults(): Flow<Map<Int, ThumbnailData?>> {
+        val visibleTasks =
+            combine(_taskData, visibleTaskIds) { tasks, visibleIds ->
+                tasks.filter { it.key.id in visibleIds }
+            }
+        val visibleThumbnailDataRequests: Flow<List<ThumbnailDataRequest>> =
+            visibleTasks.map {
+                it.map { visibleTask ->
+                    val taskCopy = Task(visibleTask).apply { thumbnail = visibleTask.thumbnail }
+                    getThumbnailDataRequest(taskCopy)
+                }
+            }
+        return visibleThumbnailDataRequests.flatMapLatest {
+            thumbnailRequestFlows: List<ThumbnailDataRequest> ->
+            if (thumbnailRequestFlows.isEmpty()) {
+                flowOf(emptyMap())
+            } else {
+                combine(thumbnailRequestFlows) { it.toMap() }
+            }
+        }
+    }
+}
+
+typealias ThumbnailDataRequest = Flow<Pair<Int, ThumbnailData?>>
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
new file mode 100644
index 0000000..55598f0
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/data/TaskThumbnailDataSource.kt
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.task.thumbnail.data
+
+import com.android.launcher3.util.CancellableTask
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import java.util.function.Consumer
+
+interface TaskThumbnailDataSource {
+    fun updateThumbnailInBackground(
+        task: Task,
+        callback: Consumer<ThumbnailData>
+    ): CancellableTask<ThumbnailData>?
+}
diff --git a/quickstep/src/com/android/quickstep/util/DesktopTask.java b/quickstep/src/com/android/quickstep/util/DesktopTask.java
index 07f2d68..8d99069 100644
--- a/quickstep/src/com/android/quickstep/util/DesktopTask.java
+++ b/quickstep/src/com/android/quickstep/util/DesktopTask.java
@@ -21,7 +21,7 @@
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
 
-import java.util.ArrayList;
+import java.util.List;
 
 /**
  * A {@link Task} container that can contain N number of tasks that are part of the desktop in
@@ -30,9 +30,9 @@
 public class DesktopTask extends GroupTask {
 
     @NonNull
-    public final ArrayList<Task> tasks;
+    public final List<Task> tasks;
 
-    public DesktopTask(@NonNull ArrayList<Task> tasks) {
+    public DesktopTask(@NonNull List<Task> tasks) {
         super(tasks.get(0), null, null, TaskView.Type.DESKTOP);
         this.tasks = tasks;
     }
@@ -53,6 +53,12 @@
     }
 
     @Override
+    @NonNull
+    public List<Task> getTasks() {
+        return tasks;
+    }
+
+    @Override
     public DesktopTask copy() {
         return new DesktopTask(tasks);
     }
diff --git a/quickstep/src/com/android/quickstep/util/GroupTask.java b/quickstep/src/com/android/quickstep/util/GroupTask.java
index 7dd6afc..945ffe3 100644
--- a/quickstep/src/com/android/quickstep/util/GroupTask.java
+++ b/quickstep/src/com/android/quickstep/util/GroupTask.java
@@ -23,6 +23,10 @@
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
 
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
 /**
  * A {@link Task} container that can contain one or two tasks, depending on if the two tasks
  * are represented as an app-pair in the recents task list.
@@ -62,6 +66,17 @@
     }
 
     /**
+     * Returns a List of all the Tasks in this GroupTask
+     */
+    public List<Task> getTasks() {
+        if (task2 == null) {
+            return Collections.singletonList(task1);
+        } else {
+            return Arrays.asList(task1, task2);
+        }
+    }
+
+    /**
      * Create a copy of this instance
      */
     public GroupTask copy() {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 4e5d646..8ee7dbc 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1715,7 +1715,7 @@
         return super.isPageScrollsInitialized() && mLoadPlanEverApplied;
     }
 
-    protected void applyLoadPlan(ArrayList<GroupTask> taskGroups) {
+    protected void applyLoadPlan(List<GroupTask> taskGroups) {
         if (mPendingAnimation != null) {
             mPendingAnimation.addEndListener(success -> applyLoadPlan(taskGroups));
             return;
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeRecentTasksDataSource.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeRecentTasksDataSource.kt
new file mode 100644
index 0000000..eaeb513
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeRecentTasksDataSource.kt
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.data
+
+import com.android.quickstep.util.GroupTask
+import java.util.function.Consumer
+
+class FakeRecentTasksDataSource : RecentTasksDataSource {
+    var taskList: List<GroupTask> = listOf()
+
+    override fun getTasks(callback: Consumer<List<GroupTask>>?): Int {
+        callback?.accept(taskList)
+        return 0
+    }
+
+    fun seedTasks(tasks: List<GroupTask>) {
+        taskList = tasks
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
new file mode 100644
index 0000000..b66b735
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTaskThumbnailDataSource.kt
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.data
+
+import android.graphics.Bitmap
+import com.android.launcher3.util.CancellableTask
+import com.android.quickstep.task.thumbnail.data.TaskThumbnailDataSource
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import java.util.function.Consumer
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+class FakeTaskThumbnailDataSource : TaskThumbnailDataSource {
+
+    val taskIdToBitmap: Map<Int, Bitmap> = (0..10).associateWith { mock() }
+    val taskIdToUpdatingTask: MutableMap<Int, () -> Unit> = mutableMapOf()
+    var shouldLoadSynchronously: Boolean = true
+
+    /** Retrieves and sets a thumbnail on [task] from [taskIdToBitmap]. */
+    override fun updateThumbnailInBackground(
+        task: Task,
+        callback: Consumer<ThumbnailData>
+    ): CancellableTask<ThumbnailData>? {
+        val thumbnailData = mock<ThumbnailData>()
+        whenever(thumbnailData.thumbnail).thenReturn(taskIdToBitmap[task.key.id])
+        val wrappedCallback = {
+            task.thumbnail = thumbnailData
+            callback.accept(thumbnailData)
+        }
+        if (shouldLoadSynchronously) {
+            wrappedCallback()
+        } else {
+            taskIdToUpdatingTask[task.key.id] = wrappedCallback
+        }
+        return null
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
new file mode 100644
index 0000000..c28a85a
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -0,0 +1,150 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.data
+
+import android.content.ComponentName
+import android.content.Intent
+import com.android.quickstep.TaskIconCache
+import com.android.quickstep.util.DesktopTask
+import com.android.quickstep.util.GroupTask
+import com.android.systemui.shared.recents.model.Task
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.drop
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.mockito.kotlin.mock
+
+@OptIn(ExperimentalCoroutinesApi::class)
+class TasksRepositoryTest {
+    private val tasks = (0..5).map(::createTaskWithId)
+    private val defaultTaskList =
+        listOf(
+            GroupTask(tasks[0]),
+            GroupTask(tasks[1], tasks[2], null),
+            DesktopTask(tasks.subList(3, 6))
+        )
+    private val recentsModel = FakeRecentTasksDataSource()
+    private val taskThumbnailDataSource = FakeTaskThumbnailDataSource()
+    private val taskIconCache = mock<TaskIconCache>()
+
+    private val systemUnderTest =
+        TasksRepository(recentsModel, taskThumbnailDataSource, taskIconCache)
+
+    @Test
+    fun getAllTaskDataReturnsFlattenedListOfTasks() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+
+        assertThat(systemUnderTest.getAllTaskData(forceRefresh = true).first()).isEqualTo(tasks)
+    }
+
+    @Test
+    fun getTaskDataByIdReturnsSpecificTask() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        assertThat(systemUnderTest.getTaskDataById(2).first()).isEqualTo(tasks[2])
+    }
+
+    @Test
+    fun setVisibleTasksPopulatesThumbnails() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getTaskDataById(1).drop(1).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap1)
+        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+    }
+
+    @Test
+    fun changingVisibleTasksContainsAlreadyPopulatedThumbnails() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getTaskDataById(2).drop(1).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+
+        // Prevent new loading of Bitmaps
+        taskThumbnailDataSource.shouldLoadSynchronously = false
+        systemUnderTest.setVisibleTasks(listOf(2, 3))
+
+        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+    }
+
+    @Test
+    fun retrievedThumbnailsAreDiscardedWhenTaskBecomesInvisible() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getTaskDataById(2).drop(1).first()!!.thumbnail!!.thumbnail)
+            .isEqualTo(bitmap2)
+
+        // Prevent new loading of Bitmaps
+        taskThumbnailDataSource.shouldLoadSynchronously = false
+        systemUnderTest.setVisibleTasks(listOf(0, 1))
+
+        assertThat(systemUnderTest.getTaskDataById(2).first()!!.thumbnail).isNull()
+    }
+
+    @Test
+    fun retrievedThumbnailsCauseEmissionOnTaskDataFlow() = runTest {
+        // Setup fakes
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        taskThumbnailDataSource.shouldLoadSynchronously = false
+
+        // Setup TasksRepository
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // Assert there is no bitmap in first emission
+        val taskFlow = systemUnderTest.getTaskDataById(2)
+        val taskFlowValuesList = mutableListOf<Task?>()
+        backgroundScope.launch(UnconfinedTestDispatcher(testScheduler)) {
+            taskFlow.toList(taskFlowValuesList)
+        }
+        assertThat(taskFlowValuesList[0]!!.thumbnail).isNull()
+
+        // Simulate bitmap loading after first emission
+        taskThumbnailDataSource.taskIdToUpdatingTask.getValue(2).invoke()
+
+        // Check for second emission
+        assertThat(taskFlowValuesList[1]!!.thumbnail!!.thumbnail).isEqualTo(bitmap2)
+    }
+
+    private fun createTaskWithId(taskId: Int) =
+        Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000))
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
index 0de5f19..aa08ca4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/SplitSelectStateControllerTest.kt
@@ -31,7 +31,6 @@
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.statehandlers.DepthController
 import com.android.launcher3.statemanager.StateManager
-import com.android.launcher3.statemanager.StatefulActivity
 import com.android.launcher3.util.ComponentKey
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.quickstep.RecentsModel
@@ -121,7 +120,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonMatchingComponent),
                         false /* findExactPairMatch */,
@@ -174,7 +173,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent),
                         false /* findExactPairMatch */,
@@ -215,7 +214,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonPrimaryUserComponent),
                         false /* findExactPairMatch */,
@@ -271,7 +270,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonPrimaryUserComponent),
                         false /* findExactPairMatch */,
@@ -324,7 +323,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent),
                         false /* findExactPairMatch */,
@@ -378,7 +377,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(nonMatchingComponent, matchingComponent),
                         false /* findExactPairMatch */,
@@ -431,7 +430,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent, matchingComponent),
                         false /* findExactPairMatch */,
@@ -497,7 +496,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent, matchingComponent),
                         false /* findExactPairMatch */,
@@ -549,7 +548,7 @@
 
         // Capture callback from recentsModel#getTasks()
         val consumer =
-            argumentCaptor<Consumer<ArrayList<GroupTask>>> {
+            argumentCaptor<Consumer<List<GroupTask>>> {
                     splitSelectStateController.findLastActiveTasksAndRunCallback(
                         listOf(matchingComponent2, matchingComponent),
                         true /* findExactPairMatch */,