diff --git a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
index bc6348d..aa31261 100644
--- a/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/LauncherTaskbarUIController.java
@@ -63,6 +63,7 @@
 import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Stream;
@@ -459,7 +460,7 @@
         }
 
         @Override
-        public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+        public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
             endGestureStateOverride(true);
         }
 
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 30a0fdf..a0a940b 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -121,6 +121,7 @@
 import com.android.systemui.shared.system.TaskStackChangeListeners;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.function.Consumer;
 
 /**
@@ -399,9 +400,10 @@
         // Set up a entire animation lifecycle callback to notify the current recents view when
         // the animation is canceled
         mGestureState.runOnceAtState(STATE_RECENTS_ANIMATION_CANCELED, () -> {
-                ThumbnailData snapshot = mGestureState.consumeRecentsAnimationCanceledSnapshot();
-                if (snapshot != null) {
-                    mRecentsView.switchToScreenshot(snapshot, () -> {
+                HashMap<Integer, ThumbnailData> snapshots =
+                        mGestureState.consumeRecentsAnimationCanceledSnapshot();
+                if (snapshots != null) {
+                    mRecentsView.switchToScreenshot(snapshots, () -> {
                         if (mRecentsAnimationController != null) {
                             mRecentsAnimationController.cleanupScreenshot();
                         }
@@ -809,7 +811,7 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+    public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
         ActiveGestureLog.INSTANCE.addLog("cancelRecentsAnimation");
         mActivityInitListener.unregister();
         mStateCallback.setStateOnUiThread(STATE_GESTURE_CANCELLED | STATE_HANDLER_INVALIDATED);
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 944fb48..834bc55 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -66,6 +66,7 @@
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 
+import java.util.HashMap;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
@@ -192,7 +193,8 @@
 
     public void closeOverlay() { }
 
-    public void switchRunningTaskViewToScreenshot(ThumbnailData thumbnailData, Runnable runnable) {
+    public void switchRunningTaskViewToScreenshot(HashMap<Integer, ThumbnailData> thumbnailDatas,
+            Runnable runnable) {
         ACTIVITY_TYPE activity = getCreatedActivity();
         if (activity == null) {
             return;
@@ -204,7 +206,7 @@
             }
             return;
         }
-        recentsView.switchToScreenshot(thumbnailData, runnable);
+        recentsView.switchToScreenshot(thumbnailDatas, runnable);
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/GestureState.java b/quickstep/src/com/android/quickstep/GestureState.java
index 99f1dc7..9bc98d6 100644
--- a/quickstep/src/com/android/quickstep/GestureState.java
+++ b/quickstep/src/com/android/quickstep/GestureState.java
@@ -35,6 +35,7 @@
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Set;
 
@@ -139,7 +140,7 @@
     private Set<Integer> mPreviouslyAppearedTaskIds = new HashSet<>();
     private int mLastStartedTaskId = -1;
     private RecentsAnimationController mRecentsAnimationController;
-    private ThumbnailData mRecentsAnimationCanceledSnapshot;
+    private HashMap<Integer, ThumbnailData> mRecentsAnimationCanceledSnapshots;
 
     /** The time when the swipe up gesture is triggered. */
     private long mSwipeUpStartTimeMs;
@@ -354,16 +355,16 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
-        mRecentsAnimationCanceledSnapshot = thumbnailData;
+    public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
+        mRecentsAnimationCanceledSnapshots = thumbnailDatas;
         mStateCallback.setState(STATE_RECENTS_ANIMATION_CANCELED);
         mStateCallback.setState(STATE_RECENTS_ANIMATION_ENDED);
-        if (mRecentsAnimationCanceledSnapshot != null) {
+        if (mRecentsAnimationCanceledSnapshots != null) {
             // Clean up the screenshot to finalize the recents animation cancel
             if (mRecentsAnimationController != null) {
                 mRecentsAnimationController.cleanupScreenshot();
             }
-            mRecentsAnimationCanceledSnapshot = null;
+            mRecentsAnimationCanceledSnapshots = null;
         }
     }
 
@@ -378,9 +379,10 @@
      * while STATE_RECENTS_ANIMATION_CANCELED state is being set, and the caller is responsible for
      * calling {@link RecentsAnimationController#cleanupScreenshot()}.
      */
-    ThumbnailData consumeRecentsAnimationCanceledSnapshot() {
-        ThumbnailData data = mRecentsAnimationCanceledSnapshot;
-        mRecentsAnimationCanceledSnapshot = null;
+    HashMap<Integer, ThumbnailData> consumeRecentsAnimationCanceledSnapshot() {
+        HashMap<Integer, ThumbnailData> data =
+                new HashMap<Integer, ThumbnailData>(mRecentsAnimationCanceledSnapshots);
+        mRecentsAnimationCanceledSnapshots = null;
         return data;
     }
 
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
index b232464..26d935d 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.java
@@ -39,6 +39,7 @@
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 
 /**
  * Helper class to handle various atomic commands for switching between Overview.
@@ -211,7 +212,7 @@
             }
 
             @Override
-            public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+            public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
                 interactionHandler.onGestureCancelled();
                 cmd.removeListener(this);
 
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
index e948221..750985a 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
@@ -32,6 +32,7 @@
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.Set;
 
 /**
@@ -125,10 +126,10 @@
 
     @BinderThread
     @Override
-    public final void onAnimationCanceled(ThumbnailData thumbnailData) {
+    public final void onAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
         Utilities.postAsyncCallback(MAIN_EXECUTOR.getHandler(), () -> {
             for (RecentsAnimationListener listener : getListeners()) {
-                listener.onRecentsAnimationCanceled(thumbnailData);
+                listener.onRecentsAnimationCanceled(thumbnailDatas);
             }
         });
     }
@@ -166,7 +167,7 @@
          * Callback from the system when the recents animation is canceled. {@param thumbnailData}
          * is passed back for rendering screenshot to replace live tile.
          */
-        default void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {}
+        default void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {}
 
         /**
          * Callback made whenever the recents animation is finished.
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 258a0c2..ed7f0b3 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -42,6 +42,8 @@
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
 
+import java.util.HashMap;
+
 public class TaskAnimationManager implements RecentsAnimationCallbacks.RecentsAnimationListener {
     public static final boolean ENABLE_SHELL_TRANSITIONS =
             SystemProperties.getBoolean("persist.debug.shell_transit", false);
@@ -135,7 +137,7 @@
             }
 
             @Override
-            public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+            public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
                 cleanUpRecentsAnimation();
             }
 
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
index fcc0217..3d737ca 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
@@ -60,6 +60,8 @@
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams.Builder;
 
+import java.util.HashMap;
+
 /**
  * A placeholder input consumer used when the device is still locked, e.g. from secure camera.
  */
@@ -239,7 +241,7 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+    public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
         mRecentsAnimationController = null;
         mTransformParams.setTargetSet(null);
     }
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/ProgressDelegateInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/ProgressDelegateInputConsumer.java
index c69b510..033fd85 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/ProgressDelegateInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/ProgressDelegateInputConsumer.java
@@ -44,6 +44,8 @@
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.InputMonitorCompat;
 
+import java.util.HashMap;
+
 /**
  * Input consumer which delegates the swipe-progress handling
  */
@@ -174,7 +176,7 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+    public void onRecentsAnimationCanceled(HashMap<Integer, ThumbnailData> thumbnailDatas) {
         mRecentsAnimationController = null;
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index 72d3731..357fb67 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -24,6 +24,7 @@
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 
+import java.util.HashMap;
 import java.util.function.Consumer;
 
 /**
@@ -162,6 +163,20 @@
     }
 
     @Override
+    void refreshThumbnails(HashMap<Integer, ThumbnailData> thumbnailDatas) {
+        super.refreshThumbnails(thumbnailDatas);
+        if (mSecondaryTask != null && thumbnailDatas != null) {
+            final ThumbnailData thumbnailData = thumbnailDatas.get(mSecondaryTask.key.id);
+            if (thumbnailData != null) {
+                mSnapshotView2.setThumbnail(mSecondaryTask, thumbnailData);
+                return;
+            }
+        }
+
+        mSnapshotView2.refresh();
+    }
+
+    @Override
     public TaskThumbnailView[] getThumbnails() {
         return new TaskThumbnailView[]{mSnapshotView, mSnapshotView2};
     }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 33ba7dd..c721e50 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -180,6 +180,7 @@
 import com.android.wm.shell.pip.IPipAnimationListener;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.function.Consumer;
 
@@ -4748,15 +4749,12 @@
      *  from gesture state, which is a larger change of it having to keep track of multiple tasks.
      *  OR. Maybe it doesn't need to pass in a thumbnail and we can use the exact same flow as above
      */
-    public void switchToScreenshot(ThumbnailData thumbnailData, Runnable onFinishRunnable) {
-        TaskView taskView = getRunningTaskView();
+    public void switchToScreenshot(@Nullable HashMap<Integer, ThumbnailData> thumbnailDatas,
+            Runnable onFinishRunnable) {
+        final TaskView taskView = getRunningTaskView();
         if (taskView != null) {
             taskView.setShowScreenshot(true);
-            if (thumbnailData != null) {
-                taskView.getThumbnail().setThumbnail(taskView.getTask(), thumbnailData);
-            } else {
-                taskView.getThumbnail().refresh();
-            }
+            taskView.refreshThumbnails(thumbnailDatas);
             ViewUtils.postFrameDrawn(taskView, onFinishRunnable);
         } else {
             onFinishRunnable.run();
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index f88b243..97c4748 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -61,6 +61,7 @@
 
 import androidx.annotation.IntDef;
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.DeviceProfile;
@@ -98,6 +99,7 @@
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.views.TaskThumbnailView.PreviewPositionHelper;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.ActivityOptionsCompat;
 import com.android.systemui.shared.system.QuickStepContract;
@@ -106,6 +108,7 @@
 import java.lang.annotation.Retention;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
@@ -559,6 +562,18 @@
         return mSnapshotView;
     }
 
+    void refreshThumbnails(@Nullable HashMap<Integer, ThumbnailData> thumbnailDatas) {
+        if (mTask != null && thumbnailDatas != null) {
+            final ThumbnailData thumbnailData = thumbnailDatas.get(mTask.key.id);
+            if (thumbnailData != null) {
+                mSnapshotView.setThumbnail(mTask, thumbnailData);
+                return;
+            }
+        }
+
+        mSnapshotView.refresh();
+    }
+
     /** TODO(b/197033698) Remove all usages of above method and migrate to this one */
     public TaskThumbnailView[] getThumbnails() {
         return new TaskThumbnailView[]{mSnapshotView};