Handle app switch in Overview

- Transfer screenshot from WM to launcher when recents animation gets cancelled due to stack order change
- Transform two live windows when app open animation happens from recents

Fixes: 139258979
Fixes: 139259253
Test: switch task in overview

Change-Id: I80bafb8d45b9250fda937223254e365596a7f538
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java
index b939898..89c34d3 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/AppToOverviewAnimationProvider.java
@@ -148,8 +148,10 @@
         valueAnimator.setDuration(RECENTS_LAUNCH_DURATION);
         valueAnimator.setInterpolator(TOUCH_RESPONSE_INTERPOLATOR);
         valueAnimator.addUpdateListener((v) -> {
-            params.setProgress((float) v.getAnimatedValue());
-            clipHelper.applyTransform(targetSet, params);
+            params.setProgress((float) v.getAnimatedValue())
+                    .setTargetSet(targetSet)
+                    .setLauncherOnTop(true);
+            clipHelper.applyTransform(params);
         });
 
         if (targetSet.isAnimatingHome()) {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
index d524b5d..ed6a7e8 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/BaseSwipeUpHandler.java
@@ -367,9 +367,12 @@
         float offsetX = mRecentsView == null ? 0 : mRecentsView.getScrollOffset();
         float offsetScale = getTaskCurveScaleForOffsetX(offsetX,
                 mClipAnimationHelper.getTargetRect().width());
-        mTransformParams.setProgress(shift).setOffsetX(offsetX).setOffsetScale(offsetScale);
-        mClipAnimationHelper.applyTransform(mRecentsAnimationWrapper.targetSet,
-                mTransformParams);
+        mTransformParams.setProgress(shift)
+                .setOffsetX(offsetX)
+                .setOffsetScale(offsetScale)
+                .setTargetSet(mRecentsAnimationWrapper.targetSet)
+                .setLauncherOnTop(true);
+        mClipAnimationHelper.applyTransform(mTransformParams);
     }
 
     private float getTaskCurveScaleForOffsetX(float offsetX, float taskWidth) {
@@ -386,8 +389,11 @@
     protected RectFSpringAnim createWindowAnimationToHome(float startProgress,
             HomeAnimationFactory homeAnimationFactory) {
         final RemoteAnimationTargetSet targetSet = mRecentsAnimationWrapper.targetSet;
-        final RectF startRect = new RectF(mClipAnimationHelper.applyTransform(targetSet,
-                mTransformParams.setProgress(startProgress), false /* launcherOnTop */));
+        final RectF startRect = new RectF(
+                mClipAnimationHelper.applyTransform(
+                        mTransformParams.setProgress(startProgress)
+                                .setTargetSet(targetSet)
+                                .setLauncherOnTop(false)));
         final RectF targetRect = homeAnimationFactory.getWindowTargetRect();
 
         final View floatingView = homeAnimationFactory.getFloatingView();
@@ -434,8 +440,7 @@
                     mTransformParams.setCornerRadius(endRadius * progress + startRadius
                             * (1f - progress));
                 }
-                mClipAnimationHelper.applyTransform(targetSet, mTransformParams,
-                        false /* launcherOnTop */);
+                mClipAnimationHelper.applyTransform(mTransformParams);
 
                 if (isFloatingIconView) {
                     ((FloatingIconView) floatingView).update(currentRect, 1f, progress,
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityControllerHelper.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityControllerHelper.java
index e7005a6..bb73521 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityControllerHelper.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/LauncherActivityControllerHelper.java
@@ -67,6 +67,7 @@
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.plugins.shared.LauncherOverlayManager;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 
 import java.util.function.BiPredicate;
@@ -491,4 +492,24 @@
             om.hideOverlay(150);
         }
     }
+
+    @Override
+    public void switchToScreenshot(ThumbnailData thumbnailData, Runnable runnable) {
+        Launcher launcher = getCreatedActivity();
+        RecentsView recentsView = launcher.getOverviewPanel();
+        if (recentsView == null) {
+            if (runnable != null) {
+                runnable.run();
+            }
+            return;
+        }
+        TaskView taskView = recentsView.getRunningTaskView();
+        if (taskView != null) {
+            taskView.setShowScreenshot(true);
+            taskView.getThumbnail().setThumbnail(taskView.getTask(), thumbnailData);
+            ViewUtils.postDraw(taskView, runnable);
+        } else if (runnable != null) {
+            runnable.run();
+        }
+    }
 }
\ No newline at end of file
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java
index 8783ee3..8301f36 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/SwipeSharedState.java
@@ -25,6 +25,7 @@
 import com.android.quickstep.util.RecentsAnimationListenerSet;
 import com.android.quickstep.util.SwipeAnimationTargetSet;
 import com.android.quickstep.util.SwipeAnimationTargetSet.SwipeAnimationListener;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.io.PrintWriter;
 
@@ -37,6 +38,7 @@
 
     private RecentsAnimationListenerSet mRecentsAnimationListener;
     private SwipeAnimationTargetSet mLastAnimationTarget;
+    private Runnable mRecentsAnimationCanceledCallback;
 
     private boolean mLastAnimationCancelled = false;
     private boolean mLastAnimationRunning = false;
@@ -67,11 +69,30 @@
     }
 
     @Override
-    public final void onRecentsAnimationCanceled() {
+    public final void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
+        if (thumbnailData != null) {
+            mOverviewComponentObserver.getActivityControlHelper().switchToScreenshot(thumbnailData,
+                    () -> {
+                        if (mRecentsAnimationCanceledCallback != null) {
+                            mRecentsAnimationCanceledCallback.run();
+                        }
+                        clearAnimationState();
+                    });
+        } else {
+            clearAnimationState();
+        }
+    }
+
+    public void setRecentsAnimationCanceledCallback(Runnable callback) {
+        mRecentsAnimationCanceledCallback = callback;
+    }
+
+    private void clearAnimationState() {
         clearAnimationTarget();
 
         mLastAnimationCancelled = true;
         mLastAnimationRunning = false;
+        mRecentsAnimationCanceledCallback = null;
     }
 
     private void clearListenerState(boolean finishAnimation) {
@@ -127,7 +148,7 @@
         if (mLastAnimationTarget != null) {
             listener.onRecentsAnimationStart(mLastAnimationTarget);
         } else if (mLastAnimationCancelled) {
-            listener.onRecentsAnimationCanceled();
+            listener.onRecentsAnimationCanceled(null);
         }
     }
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java
index 00fa0f2..b4a7996 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/TaskViewUtils.java
@@ -17,6 +17,7 @@
 
 import static com.android.launcher3.anim.Interpolators.LINEAR;
 import static com.android.launcher3.anim.Interpolators.TOUCH_RESPONSE_INTERPOLATOR;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_QUICKSTEP_LIVE_TILE;
 import static com.android.systemui.shared.system.RemoteAnimationTargetCompat.MODE_OPENING;
 
 import android.animation.Animator;
@@ -38,6 +39,11 @@
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat;
+import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
 
 /**
  * Utility class for helpful methods related to {@link TaskView} objects and their tasks.
@@ -115,12 +121,13 @@
             RemoteAnimationTargetCompat[] wallpaperTargets, final ClipAnimationHelper inOutHelper) {
         SyncRtSurfaceTransactionApplierCompat applier =
                 new SyncRtSurfaceTransactionApplierCompat(v);
-        ClipAnimationHelper.TransformParams params = new ClipAnimationHelper.TransformParams()
-                .setSyncTransactionApplier(applier);
-
         final RemoteAnimationTargetSet targetSet =
                 new RemoteAnimationTargetSet(appTargets, wallpaperTargets, MODE_OPENING);
         targetSet.addDependentTransactionApplier(applier);
+        ClipAnimationHelper.TransformParams params = new ClipAnimationHelper.TransformParams()
+                .setSyncTransactionApplier(applier)
+                .setTargetSet(targetSet)
+                .setLauncherOnTop(true);
 
         final RecentsView recentsView = v.getRecentsView();
         final ValueAnimator appAnimator = ValueAnimator.ofFloat(0, 1);
@@ -152,7 +159,33 @@
             public void onUpdate(float percent) {
                 // TODO: Take into account the current fullscreen progress for animating the insets
                 params.setProgress(1 - percent);
-                RectF taskBounds = inOutHelper.applyTransform(targetSet, params);
+                RectF taskBounds;
+                if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+                    List<SurfaceParams> surfaceParamsList = new ArrayList<>();
+                    // Append the surface transform params for the app that's being opened.
+                    Collections.addAll(surfaceParamsList, inOutHelper.getSurfaceParams(params));
+
+                    ClipAnimationHelper liveTileClipAnimationHelper =
+                            v.getRecentsView().getClipAnimationHelper();
+                    if (liveTileClipAnimationHelper != null) {
+                        // Append the surface transform params for the live tile app.
+                        ClipAnimationHelper.TransformParams liveTileParams =
+                                v.getRecentsView().getLiveTileParams(true /* mightNeedToRefill */);
+                        if (liveTileParams != null) {
+                            Collections.addAll(surfaceParamsList,
+                                    liveTileClipAnimationHelper.getSurfaceParams(liveTileParams));
+                        }
+                    }
+                    // Apply surface transform using the surface params list.
+                    ClipAnimationHelper.applySurfaceParams(params.syncTransactionApplier,
+                            surfaceParamsList.toArray(new SurfaceParams[surfaceParamsList.size()]));
+                    // Get the task bounds for the app that's being opened after surface transform
+                    // update.
+                    taskBounds = inOutHelper.updateCurrentRect(params);
+                } else {
+                    taskBounds = inOutHelper.applyTransform(params);
+                }
+
                 int taskIndex = recentsView.indexOfChild(v);
                 int centerTaskIndex = recentsView.getCurrentPage();
                 boolean parallaxCenterAndAdjacentTask = taskIndex != centerTaskIndex;
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
index c034579..7e99366 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -614,7 +614,7 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled() {
+    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
         mRecentsAnimationWrapper.setController(null);
         mActivityInitListener.unregister();
         setStateOnUiThread(STATE_GESTURE_CANCELLED | STATE_HANDLER_INVALIDATED);
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
index b24c788..d11afaf 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/DeviceLockedInputConsumer.java
@@ -46,6 +46,7 @@
 import com.android.quickstep.util.ClipAnimationHelper;
 import com.android.quickstep.util.RecentsAnimationListenerSet;
 import com.android.quickstep.util.SwipeAnimationTargetSet;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.InputMonitorCompat;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 
@@ -155,9 +156,7 @@
                     float dy = Math.max(mTouchDown.y - y, 0);
                     mProgress = dy / mDisplaySize.y;
                     mTransformParams.setProgress(mProgress);
-                    if (mTargetSet != null) {
-                        mClipAnimationHelper.applyTransform(mTargetSet, mTransformParams);
-                    }
+                    mClipAnimationHelper.applyTransform(mTransformParams);
                 }
                 break;
             }
@@ -228,13 +227,14 @@
         Utilities.scaleRectAboutCenter(displaySize, SCALE_DOWN);
         displaySize.offsetTo(displaySize.left, 0);
         mClipAnimationHelper.updateTargetRect(displaySize);
-        mClipAnimationHelper.applyTransform(mTargetSet, mTransformParams);
+        mTransformParams.setTargetSet(mTargetSet).setLauncherOnTop(true);
+        mClipAnimationHelper.applyTransform(mTransformParams);
 
         mStateCallback.setState(STATE_TARGET_RECEIVED);
     }
 
     @Override
-    public void onRecentsAnimationCanceled() {
+    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
         mTargetSet = null;
     }
 
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java
index 6275d21..629d240 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/inputconsumers/FallbackNoButtonInputConsumer.java
@@ -427,7 +427,7 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled() {
+    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
         mRecentsAnimationWrapper.setController(null);
         setStateOnUiThread(STATE_HANDLER_INVALIDATED);
     }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/ClipAnimationHelper.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/ClipAnimationHelper.java
index cae273a..587e9f2 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/ClipAnimationHelper.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/ClipAnimationHelper.java
@@ -157,12 +157,75 @@
         mUseRoundedCornersOnWindows = mSupportsRoundedCornersOnWindows && !dp.isMultiWindowMode;
     }
 
-    public RectF applyTransform(RemoteAnimationTargetSet targetSet, TransformParams params) {
-        return applyTransform(targetSet, params, true /* launcherOnTop */);
+    public RectF applyTransform(TransformParams params) {
+        SurfaceParams[] surfaceParams = getSurfaceParams(params);
+        if (surfaceParams == null) {
+            return null;
+        }
+        applySurfaceParams(params.syncTransactionApplier, surfaceParams);
+        return params.currentRect;
     }
 
-    public RectF applyTransform(RemoteAnimationTargetSet targetSet, TransformParams params,
-            boolean launcherOnTop) {
+    public SurfaceParams[] getSurfaceParams(TransformParams params) {
+        if (params.targetSet == null) {
+            return null;
+        }
+
+        float progress = params.progress;
+        updateCurrentRect(params);
+
+        SurfaceParams[] surfaceParams = new SurfaceParams[params.targetSet.unfilteredApps.length];
+        for (int i = 0; i < params.targetSet.unfilteredApps.length; i++) {
+            RemoteAnimationTargetCompat app = params.targetSet.unfilteredApps[i];
+            mTmpMatrix.setTranslate(app.position.x, app.position.y);
+            Rect crop = mTmpRect;
+            crop.set(app.sourceContainerBounds);
+            crop.offsetTo(0, 0);
+            float alpha;
+            int layer = RemoteAnimationProvider.getLayer(app, mBoostModeTargetLayers);
+            float cornerRadius = 0f;
+            float scale = Math.max(params.currentRect.width(), mTargetRect.width()) / crop.width();
+            if (app.mode == params.targetSet.targetMode) {
+                alpha = mTaskAlphaCallback.getAlpha(app, params.targetAlpha);
+                if (app.activityType != RemoteAnimationTargetCompat.ACTIVITY_TYPE_HOME) {
+                    mTmpMatrix.setRectToRect(mSourceRect, params.currentRect, ScaleToFit.FILL);
+                    mTmpMatrix.postTranslate(app.position.x, app.position.y);
+                    mClipRectF.roundOut(crop);
+                    if (mSupportsRoundedCornersOnWindows) {
+                        if (params.cornerRadius > -1) {
+                            cornerRadius = params.cornerRadius;
+                            scale = params.currentRect.width() / crop.width();
+                        } else {
+                            float windowCornerRadius = mUseRoundedCornersOnWindows
+                                    ? mWindowCornerRadius : 0;
+                            cornerRadius = Utilities.mapRange(progress, windowCornerRadius,
+                                    mTaskCornerRadius);
+                        }
+                        mCurrentCornerRadius = cornerRadius;
+                    }
+                } else if (params.targetSet.hasRecents) {
+                    // If home has a different target then recents, reverse anim the
+                    // home target.
+                    alpha = 1 - (progress * params.targetAlpha);
+                }
+            } else {
+                alpha = mBaseAlphaCallback.getAlpha(app, progress);
+                if (ENABLE_QUICKSTEP_LIVE_TILE.get() && params.launcherOnTop) {
+                    crop = null;
+                    layer = Integer.MAX_VALUE;
+                }
+            }
+
+            // Since radius is in Surface space, but we draw the rounded corners in screen space, we
+            // have to undo the scale.
+            surfaceParams[i] = new SurfaceParams(app.leash, alpha, mTmpMatrix, crop, layer,
+                    cornerRadius / scale);
+        }
+        applySurfaceParams(params.syncTransactionApplier, surfaceParams);
+        return surfaceParams;
+    }
+
+    public RectF updateCurrentRect(TransformParams params) {
         float progress = params.progress;
         if (params.currentRect == null) {
             RectF currentRect;
@@ -183,55 +246,6 @@
                     mSourceStackBounds.height() - (sourceWindowClipInsets.bottom * progress);
             params.setCurrentRectAndTargetAlpha(currentRect, 1);
         }
-
-        SurfaceParams[] surfaceParams = new SurfaceParams[targetSet.unfilteredApps.length];
-        for (int i = 0; i < targetSet.unfilteredApps.length; i++) {
-            RemoteAnimationTargetCompat app = targetSet.unfilteredApps[i];
-            mTmpMatrix.setTranslate(app.position.x, app.position.y);
-            Rect crop = mTmpRect;
-            crop.set(app.sourceContainerBounds);
-            crop.offsetTo(0, 0);
-            float alpha;
-            int layer = RemoteAnimationProvider.getLayer(app, mBoostModeTargetLayers);
-            float cornerRadius = 0f;
-            float scale = Math.max(params.currentRect.width(), mTargetRect.width()) / crop.width();
-            if (app.mode == targetSet.targetMode) {
-                alpha = mTaskAlphaCallback.getAlpha(app, params.targetAlpha);
-                if (app.activityType != RemoteAnimationTargetCompat.ACTIVITY_TYPE_HOME) {
-                    mTmpMatrix.setRectToRect(mSourceRect, params.currentRect, ScaleToFit.FILL);
-                    mTmpMatrix.postTranslate(app.position.x, app.position.y);
-                    mClipRectF.roundOut(crop);
-                    if (mSupportsRoundedCornersOnWindows) {
-                        if (params.cornerRadius > -1) {
-                            cornerRadius = params.cornerRadius;
-                            scale = params.currentRect.width() / crop.width();
-                        } else {
-                            float windowCornerRadius = mUseRoundedCornersOnWindows
-                                    ? mWindowCornerRadius : 0;
-                            cornerRadius = Utilities.mapRange(progress, windowCornerRadius,
-                                    mTaskCornerRadius);
-                        }
-                        mCurrentCornerRadius = cornerRadius;
-                    }
-                } else if (targetSet.hasRecents) {
-                    // If home has a different target then recents, reverse anim the
-                    // home target.
-                    alpha = 1 - (progress * params.targetAlpha);
-                }
-            } else {
-                alpha = mBaseAlphaCallback.getAlpha(app, progress);
-                if (ENABLE_QUICKSTEP_LIVE_TILE.get() && launcherOnTop) {
-                    crop = null;
-                    layer = Integer.MAX_VALUE;
-                }
-            }
-
-            // Since radius is in Surface space, but we draw the rounded corners in screen space, we
-            // have to undo the scale.
-            surfaceParams[i] = new SurfaceParams(app.leash, alpha, mTmpMatrix, crop, layer,
-                    cornerRadius / scale);
-        }
-        applySurfaceParams(params.syncTransactionApplier, surfaceParams);
         return params.currentRect;
     }
 
@@ -240,7 +254,7 @@
         return mCurrentRectWithInsets;
     }
 
-    private void applySurfaceParams(@Nullable SyncRtSurfaceTransactionApplierCompat
+    public static void applySurfaceParams(@Nullable SyncRtSurfaceTransactionApplierCompat
             syncTransactionApplier, SurfaceParams[] params) {
         if (syncTransactionApplier != null) {
             syncTransactionApplier.scheduleApply(params);
@@ -375,12 +389,14 @@
         float progress;
         public float offsetX;
         public float offsetScale;
-        @Nullable RectF currentRect;
+        public @Nullable RectF currentRect;
         float targetAlpha;
         boolean forLiveTile;
         float cornerRadius;
+        boolean launcherOnTop;
 
-        SyncRtSurfaceTransactionApplierCompat syncTransactionApplier;
+        public RemoteAnimationTargetSet targetSet;
+        public SyncRtSurfaceTransactionApplierCompat syncTransactionApplier;
 
         public TransformParams() {
             progress = 0;
@@ -390,6 +406,7 @@
             targetAlpha = 0;
             forLiveTile = false;
             cornerRadius = -1;
+            launcherOnTop = false;
         }
 
         public TransformParams setProgress(float progress) {
@@ -424,6 +441,16 @@
             return this;
         }
 
+        public TransformParams setLauncherOnTop(boolean launcherOnTop) {
+            this.launcherOnTop = launcherOnTop;
+            return this;
+        }
+
+        public TransformParams setTargetSet(RemoteAnimationTargetSet targetSet) {
+            this.targetSet = targetSet;
+            return this;
+        }
+
         public TransformParams setSyncTransactionApplier(
                 SyncRtSurfaceTransactionApplierCompat applier) {
             this.syncTransactionApplier = applier;
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java
index 71ad8ba..41017ca 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/RecentsAnimationListenerSet.java
@@ -25,6 +25,7 @@
 
 import com.android.launcher3.Utilities;
 import com.android.launcher3.util.Preconditions;
+import com.android.quickstep.TouchInteractionService;
 import com.android.quickstep.util.SwipeAnimationTargetSet.SwipeAnimationListener;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RecentsAnimationControllerCompat;
@@ -40,12 +41,6 @@
  */
 public class RecentsAnimationListenerSet implements RecentsAnimationListener {
 
-    // The actual app surface is replaced by a screenshot upon recents animation cancelation when
-    // the thumbnailData exists. Launcher takes the responsibility to clean up this screenshot
-    // after app transition is finished. This delay is introduced to cover the app transition
-    // period of time.
-    private final int TRANSITION_DELAY = 100;
-
     private final Set<SwipeAnimationListener> mListeners = new ArraySet<>();
     private final boolean mShouldMinimizeSplitScreen;
     private final Consumer<SwipeAnimationTargetSet> mOnFinishListener;
@@ -57,6 +52,8 @@
             Consumer<SwipeAnimationTargetSet> onFinishListener) {
         mShouldMinimizeSplitScreen = shouldMinimizeSplitScreen;
         mOnFinishListener = onFinishListener;
+        TouchInteractionService.getSwipeSharedState().setRecentsAnimationCanceledCallback(
+                () -> mController.cleanupScreenshot());
     }
 
     @UiThread
@@ -108,14 +105,9 @@
     public final void onAnimationCanceled(ThumbnailData thumbnailData) {
         Utilities.postAsyncCallback(MAIN_EXECUTOR.getHandler(), () -> {
             for (SwipeAnimationListener listener : getListeners()) {
-                listener.onRecentsAnimationCanceled();
+                listener.onRecentsAnimationCanceled(thumbnailData);
             }
         });
-        // TODO: handle the transition better instead of simply using a transition delay.
-        if (thumbnailData != null) {
-            MAIN_EXECUTOR.getHandler().postDelayed(() -> mController.cleanupScreenshot(),
-                    TRANSITION_DELAY);
-        }
     }
 
     private SwipeAnimationListener[] getListeners() {
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java
index 3da6b78..5ab335a 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/util/SwipeAnimationTargetSet.java
@@ -114,6 +114,10 @@
 
         void onRecentsAnimationStart(SwipeAnimationTargetSet targetSet);
 
-        void onRecentsAnimationCanceled();
+        /**
+         * Callback from the system when the recents animation is canceled. {@param thumbnailData}
+         * is passed back for rendering screenshot to replace live tile.
+         */
+        void onRecentsAnimationCanceled(ThumbnailData thumbnailData);
     }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
index 4c308fb..8651095 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -174,12 +174,11 @@
     @Override
     protected void onTaskLaunchAnimationUpdate(float progress, TaskView tv) {
         if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
-            if (mRecentsAnimationWrapper.targetSet != null && tv.isRunningTask()) {
+            if (tv.isRunningTask()) {
                 mTransformParams.setProgress(1 - progress)
                         .setSyncTransactionApplier(mSyncTransactionApplier)
                         .setForLiveTile(true);
-                mClipAnimationHelper.applyTransform(mRecentsAnimationWrapper.targetSet,
-                        mTransformParams);
+                mClipAnimationHelper.applyTransform(mTransformParams);
             } else {
                 redrawLiveTile(true);
             }
@@ -212,9 +211,18 @@
 
     @Override
     public void redrawLiveTile(boolean mightNeedToRefill) {
+        ClipAnimationHelper.TransformParams transformParams = getLiveTileParams(mightNeedToRefill);
+        if (transformParams != null) {
+            mClipAnimationHelper.applyTransform(transformParams);
+        }
+    }
+
+    @Override
+    public ClipAnimationHelper.TransformParams getLiveTileParams(
+            boolean mightNeedToRefill) {
         if (!mEnableDrawingLiveTile || mRecentsAnimationWrapper == null
                 || mClipAnimationHelper == null) {
-            return;
+            return null;
         }
         TaskView taskView = getRunningTaskView();
         if (taskView != null) {
@@ -236,12 +244,11 @@
             mTempRectF.set(mTempRect);
             mTransformParams.setProgress(1f)
                     .setCurrentRectAndTargetAlpha(mTempRectF, taskView.getAlpha())
-                    .setSyncTransactionApplier(mSyncTransactionApplier);
-            if (mRecentsAnimationWrapper.targetSet != null) {
-                mClipAnimationHelper.applyTransform(mRecentsAnimationWrapper.targetSet,
-                        mTransformParams);
-            }
+                    .setSyncTransactionApplier(mSyncTransactionApplier)
+                    .setTargetSet(mRecentsAnimationWrapper.targetSet)
+                    .setLauncherOnTop(true);
         }
+        return mTransformParams;
     }
 
     @Override
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
index 29e93aa..a659bdc 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/RecentsView.java
@@ -869,7 +869,9 @@
         setEnableFreeScroll(true);
         setEnableDrawingLiveTile(true);
         setOnScrollChangeListener(null);
-        setRunningTaskViewShowScreenshot(true);
+        if (!ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+            setRunningTaskViewShowScreenshot(true);
+        }
         setRunningTaskHidden(false);
         animateUpRunningTaskIconScale();
     }
@@ -1801,10 +1803,19 @@
         }
     }
 
+    public ClipAnimationHelper getClipAnimationHelper() {
+        return mClipAnimationHelper;
+    }
+
     public ClipAnimationHelper getTempClipAnimationHelper() {
         return mTempClipAnimationHelper;
     }
 
+    public ClipAnimationHelper.TransformParams getLiveTileParams(
+            boolean mightNeedToRefill) {
+        return null;
+    }
+
     private void updateEnabledOverlays() {
         int overlayEnabledPage = mOverlayEnabled ? getNextPage() : -1;
         int taskCount = getTaskViewCount();
diff --git a/quickstep/src/com/android/quickstep/ActivityControlHelper.java b/quickstep/src/com/android/quickstep/ActivityControlHelper.java
index 2d4b136..27b94f0 100644
--- a/quickstep/src/com/android/quickstep/ActivityControlHelper.java
+++ b/quickstep/src/com/android/quickstep/ActivityControlHelper.java
@@ -34,6 +34,7 @@
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.quickstep.util.ActivityInitListener;
 import com.android.quickstep.util.RemoteAnimationTargetSet;
+import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 
 import java.util.function.BiPredicate;
@@ -97,6 +98,8 @@
 
     default void closeOverlay() { }
 
+    default void switchToScreenshot(ThumbnailData thumbnailData, Runnable runnable) {}
+
     interface AnimationFactory {
 
         enum ShelfAnimState {