diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
index 54e34a0..b1f46a3 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
@@ -66,7 +66,7 @@
         combine(
                 updatedThumbnails.map {
                     recentsTasksRepository.getThumbnailById(it.key).filter { thumbnailData ->
-                        thumbnailData == it.value
+                        thumbnailData?.snapshotId == it.value.snapshotId
                     }
                 }
             ) {}
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 3702f93..fa9a282 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1048,7 +1048,7 @@
     @Nullable
     public Task onTaskThumbnailChanged(int taskId, ThumbnailData thumbnailData) {
         if (enableRefactorTaskThumbnail()) {
-            // TODO(b/342560598): Listen in TaskRepository and reload
+            mHelper.onTaskThumbnailChanged(taskId, thumbnailData);
             return null;
         }
         if (mHandleTaskStackChanges) {
@@ -1097,9 +1097,7 @@
 
     /** Updates the thumbnail(s) of the relevant TaskView. */
     public void updateThumbnail(Map<Integer, ThumbnailData> thumbnailData) {
-        if (enableRefactorTaskThumbnail()) {
-            mRecentsViewModel.addOrUpdateThumbnailOverride(thumbnailData);
-        } else {
+        if (!enableRefactorTaskThumbnail()) {
             for (Map.Entry<Integer, ThumbnailData> entry : thumbnailData.entrySet()) {
                 Integer id = entry.getKey();
                 ThumbnailData thumbnail = entry.getValue();
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt b/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
index a63ccec..e8c9718 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
@@ -60,9 +60,6 @@
         // Update recentsViewModel and apply the thumbnailOverride ASAP, before waiting inside
         // viewAttachedScope.
         recentsViewModel.setRunningTaskShowScreenshot(true)
-        if (updatedThumbnails != null) {
-            recentsViewModel.addOrUpdateThumbnailOverride(updatedThumbnails)
-        }
         viewAttachedScope.launch {
             recentsViewModel.waitForRunningTaskShowScreenshotToUpdate()
             if (updatedThumbnails != null) {
@@ -71,4 +68,8 @@
             ViewUtils.postFrameDrawn(taskView, onFinishRunnable)
         }
     }
+
+    fun onTaskThumbnailChanged(taskId: Int, thumbnailData: ThumbnailData) {
+        recentsViewModel.addOrUpdateThumbnailOverride(mapOf(taskId to thumbnailData))
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
index b3caf2d..dc16475 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
@@ -72,9 +72,8 @@
 
     @Test
     fun thumbnailOverrideWaitAndReset() = runTest {
-        val thumbnailData1 = createThumbnailData()
-        val thumbnailData2 = createThumbnailData()
-        val thumbnailDataOverride = createThumbnailData()
+        val thumbnailData1 = createThumbnailData().apply { snapshotId = 1 }
+        val thumbnailData2 = createThumbnailData().apply { snapshotId = 2 }
         tasksRepository.seedTasks(tasks)
         tasksRepository.seedThumbnailData(mapOf(1 to thumbnailData1, 2 to thumbnailData2))
 
@@ -87,15 +86,16 @@
         assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
         assertThat(thumbnailDataFlow2.first()).isEqualTo(thumbnailData2)
 
-        val thumbnailUpdate = mapOf(2 to thumbnailDataOverride)
         systemUnderTest.setRunningTaskShowScreenshot(true)
-        systemUnderTest.addOrUpdateThumbnailOverride(thumbnailUpdate)
+        val thumbnailOverride = mapOf(2 to createThumbnailData().apply { snapshotId = 3 })
+        systemUnderTest.addOrUpdateThumbnailOverride(thumbnailOverride)
 
         systemUnderTest.waitForRunningTaskShowScreenshotToUpdate()
-        systemUnderTest.waitForThumbnailsToUpdate(thumbnailUpdate)
+        val expectedUpdate = mapOf(2 to createThumbnailData().apply { snapshotId = 3 })
+        systemUnderTest.waitForThumbnailsToUpdate(expectedUpdate)
 
         assertThat(thumbnailDataFlow1.first()).isEqualTo(thumbnailData1)
-        assertThat(thumbnailDataFlow2.first()).isEqualTo(thumbnailDataOverride)
+        assertThat(thumbnailDataFlow2.first()?.snapshotId).isEqualTo(3)
 
         systemUnderTest.onReset()