diff --git a/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java b/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
index ae8fd82..f2aa842 100644
--- a/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
+++ b/go/quickstep/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
@@ -89,5 +89,5 @@
         return RotationMode.NORMAL;
     }
 
-    public static void clearSwipeSharedState(boolean finishAnimation) {}
+    public static void clearSwipeSharedState(Launcher launcher, boolean finishAnimation) { }
 }
diff --git a/iconloaderlib/src/com/android/launcher3/icons/BaseIconFactory.java b/iconloaderlib/src/com/android/launcher3/icons/BaseIconFactory.java
index fc7d6b3..f491ed7 100644
--- a/iconloaderlib/src/com/android/launcher3/icons/BaseIconFactory.java
+++ b/iconloaderlib/src/com/android/launcher3/icons/BaseIconFactory.java
@@ -22,6 +22,7 @@
 import android.os.Build;
 import android.os.Process;
 import android.os.UserHandle;
+
 import androidx.annotation.NonNull;
 
 /**
@@ -35,6 +36,8 @@
     static final boolean ATLEAST_OREO = Build.VERSION.SDK_INT >= Build.VERSION_CODES.O;
     static final boolean ATLEAST_P = Build.VERSION.SDK_INT >= Build.VERSION_CODES.P;
 
+    private static final float ICON_BADGE_SCALE = 0.444f;
+
     private final Rect mOldBounds = new Rect();
     protected final Context mContext;
     private final Canvas mCanvas;
@@ -254,7 +257,7 @@
      * Adds the {@param badge} on top of {@param target} using the badge dimensions.
      */
     public void badgeWithDrawable(Canvas target, Drawable badge) {
-        int badgeSize = mContext.getResources().getDimensionPixelSize(R.dimen.profile_badge_size);
+        int badgeSize = getBadgeSizeForIconSize(mIconBitmapSize);
         badge.setBounds(mIconBitmapSize - badgeSize, mIconBitmapSize - badgeSize,
                 mIconBitmapSize, mIconBitmapSize);
         badge.draw(target);
@@ -335,6 +338,13 @@
     }
 
     /**
+     * Returns the correct badge size given an icon size
+     */
+    public static int getBadgeSizeForIconSize(int iconSize) {
+        return (int) (ICON_BADGE_SCALE * iconSize);
+    }
+
+    /**
      * An extension of {@link BitmapDrawable} which returns the bitmap pixel size as intrinsic size.
      * This allows the badging to be done based on the action bitmap size rather than
      * the scaled bitmap size.
diff --git a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
index 46e883a..2f8af44 100644
--- a/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
+++ b/quickstep/recents_ui_overrides/src/com/android/launcher3/uioverrides/RecentsUiFactory.java
@@ -18,6 +18,7 @@
 
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.LauncherState.OVERVIEW;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_QUICKSTEP_LIVE_TILE;
 import static com.android.quickstep.SysUINavigationMode.Mode.NO_BUTTON;
 
 import android.content.Context;
@@ -196,11 +197,15 @@
         return new RecentsViewStateController(launcher);
     }
 
-    /**
-     * Clears the swipe shared state for the current swipe gesture.
-     */
-    public static void clearSwipeSharedState(boolean finishAnimation) {
-        TouchInteractionService.getSwipeSharedState().clearAllState(finishAnimation);
+    /** Clears the swipe shared state for the current swipe gesture. */
+    public static void clearSwipeSharedState(Launcher launcher, boolean finishAnimation) {
+        if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+            launcher.<RecentsView>getOverviewPanel().switchToScreenshot(
+                    () -> TouchInteractionService.getSwipeSharedState().clearAllState(
+                            finishAnimation));
+        } else {
+            TouchInteractionService.getSwipeSharedState().clearAllState(finishAnimation);
+        }
     }
 
     /**
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/ViewUtils.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/ViewUtils.java
new file mode 100644
index 0000000..cbb6ad4
--- /dev/null
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/ViewUtils.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright (C) 2019 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;
+
+import android.graphics.Canvas;
+import android.view.View;
+
+import com.android.systemui.shared.system.WindowCallbacksCompat;
+
+import java.util.function.BooleanSupplier;
+
+/**
+ * Utility class for helpful methods related to {@link View} objects.
+ */
+public class ViewUtils {
+
+    /** See {@link #postDraw(View, Runnable, BooleanSupplier)}} */
+    public static boolean postDraw(View view, Runnable onFinishRunnable) {
+        return postDraw(view, onFinishRunnable, () -> false);
+    }
+
+    /**
+     * Inject some addition logic in order to make sure that the view is updated smoothly post
+     * draw, and allow addition task to be run after view update.
+     *
+     * @param onFinishRunnable runnable to be run right after the view finishes drawing.
+     */
+    public static boolean postDraw(View view, Runnable onFinishRunnable, BooleanSupplier canceled) {
+        // Defer finishing the animation until the next launcher frame with the
+        // new thumbnail
+        return new WindowCallbacksCompat(view) {
+            // The number of frames to defer until we actually finish the animation
+            private int mDeferFrameCount = 2;
+
+            @Override
+            public void onPostDraw(Canvas canvas) {
+                // If we were cancelled after this was attached, do not update
+                // the state.
+                if (canceled.getAsBoolean()) {
+                    detach();
+                    return;
+                }
+
+                if (mDeferFrameCount > 0) {
+                    mDeferFrameCount--;
+                    // Workaround, detach and reattach to invalidate the root node for
+                    // another draw
+                    detach();
+                    attach();
+                    view.invalidate();
+                    return;
+                }
+
+                if (onFinishRunnable != null) {
+                    onFinishRunnable.run();
+                }
+                detach();
+            }
+        }.attach();
+    }
+}
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 db8eb27..5c90d71 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/WindowTransformSwipeHandler.java
@@ -45,7 +45,6 @@
 import android.app.ActivityManager.RunningTaskInfo;
 import android.content.Context;
 import android.content.Intent;
-import android.graphics.Canvas;
 import android.graphics.PointF;
 import android.graphics.RectF;
 import android.os.Build;
@@ -90,12 +89,10 @@
 import com.android.systemui.shared.system.InputConsumerController;
 import com.android.systemui.shared.system.LatencyTrackerCompat;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
-import com.android.systemui.shared.system.WindowCallbacksCompat;
 
 @TargetApi(Build.VERSION_CODES.O)
 public class WindowTransformSwipeHandler<T extends BaseDraggingActivity>
-        extends BaseSwipeUpHandler<T, RecentsView>
-        implements OnApplyWindowInsetsListener {
+        extends BaseSwipeUpHandler<T, RecentsView> implements OnApplyWindowInsetsListener {
     private static final String TAG = WindowTransformSwipeHandler.class.getSimpleName();
 
     private static final String[] STATE_NAMES = DEBUG_STATES ? new String[16] : null;
@@ -617,7 +614,7 @@
     }
 
     @Override
-    public void onRecentsAnimationCanceled() {
+    public void onRecentsAnimationCanceled(ThumbnailData thumbnailData) {
         mRecentsAnimationWrapper.setController(null);
         mActivityInitListener.unregister();
         setStateOnUiThread(STATE_GESTURE_CANCELLED | STATE_HANDLER_INVALIDATED);
@@ -927,7 +924,15 @@
             windowAnim.addListener(new AnimationSuccessListener() {
                 @Override
                 public void onAnimationSuccess(Animator animator) {
-                    setStateOnUiThread(target.endState);
+                    if (target == NEW_TASK && mRecentsView != null
+                            && mRecentsView.getNextPage() == mRecentsView.getRunningTaskIndex()) {
+                        // We are about to launch the current running task, so use LAST_TASK state
+                        // instead of NEW_TASK. This could happen, for example, if our scroll is
+                        // aborted after we determined the target to be NEW_TASK.
+                        setStateOnUiThread(LAST_TASK.endState);
+                    } else {
+                        setStateOnUiThread(target.endState);
+                    }
                 }
             });
             windowAnim.start();
@@ -1007,6 +1012,10 @@
         }
     }
 
+    public boolean isCanceled() {
+        return mCanceled;
+    }
+
     @UiThread
     private void resumeLastTask() {
         mRecentsAnimationWrapper.finish(false /* toRecents */, null);
@@ -1099,14 +1108,21 @@
     }
 
     private void switchToScreenshot() {
+        SwipeAnimationTargetSet controller = mRecentsAnimationWrapper.getController();
         if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+            if (controller != null) {
+                // Update the screenshot of the task
+                if (mTaskSnapshot == null) {
+                    mTaskSnapshot = controller.screenshotTask(mRunningTaskId);
+                }
+                mRecentsView.updateThumbnail(mRunningTaskId, mTaskSnapshot, false /* refreshNow */);
+            }
             setStateOnUiThread(STATE_SCREENSHOT_CAPTURED);
         } else if (!mRecentsAnimationWrapper.hasTargets()) {
             // If there are no targets, then we don't need to capture anything
             setStateOnUiThread(STATE_SCREENSHOT_CAPTURED);
         } else {
             boolean finishTransitionPosted = false;
-            SwipeAnimationTargetSet controller = mRecentsAnimationWrapper.getController();
             if (controller != null) {
                 // Update the screenshot of the task
                 if (mTaskSnapshot == null) {
@@ -1123,34 +1139,8 @@
                 if (taskView != null && !mCanceled) {
                     // Defer finishing the animation until the next launcher frame with the
                     // new thumbnail
-                    finishTransitionPosted = new WindowCallbacksCompat(taskView) {
-
-                        // The number of frames to defer until we actually finish the animation
-                        private int mDeferFrameCount = 2;
-
-                        @Override
-                        public void onPostDraw(Canvas canvas) {
-                            // If we were cancelled after this was attached, do not update
-                            // the state.
-                            if (mCanceled) {
-                                detach();
-                                return;
-                            }
-
-                            if (mDeferFrameCount > 0) {
-                                mDeferFrameCount--;
-                                // Workaround, detach and reattach to invalidate the root node for
-                                // another draw
-                                detach();
-                                attach();
-                                taskView.invalidate();
-                                return;
-                            }
-
-                            setStateOnUiThread(STATE_SCREENSHOT_CAPTURED);
-                            detach();
-                        }
-                    }.attach();
+                    finishTransitionPosted = ViewUtils.postDraw(taskView,
+                            () -> setStateOnUiThread(STATE_SCREENSHOT_CAPTURED), this::isCanceled);
                 }
             }
             if (!finishTransitionPosted) {
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 0f9fc17..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
@@ -78,7 +78,12 @@
 
     @Override
     public void startHome() {
-        mActivity.getStateManager().goToState(NORMAL);
+        if (ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+            switchToScreenshot(() -> finishRecentsAnimation(true /* toRecents */,
+                    () -> mActivity.getStateManager().goToState(NORMAL)));
+        } else {
+            mActivity.getStateManager().goToState(NORMAL);
+        }
     }
 
     @Override
@@ -169,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);
             }
@@ -207,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) {
@@ -231,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 554f437..c69bcb6 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
@@ -17,7 +17,6 @@
 package com.android.quickstep.views;
 
 import static androidx.dynamicanimation.animation.DynamicAnimation.MIN_VISIBLE_CHANGE_PIXELS;
-
 import static com.android.launcher3.BaseActivity.STATE_HANDLER_INVISIBILITY_FLAGS;
 import static com.android.launcher3.InvariantDeviceProfile.CHANGE_FLAG_ICON_PARAMS;
 import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
@@ -107,6 +106,7 @@
 import com.android.quickstep.RecentsModel.TaskThumbnailChangeListener;
 import com.android.quickstep.TaskThumbnailCache;
 import com.android.quickstep.TaskUtils;
+import com.android.quickstep.ViewUtils;
 import com.android.quickstep.util.ClipAnimationHelper;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -380,14 +380,23 @@
         return null;
     }
 
-    public TaskView updateThumbnail(int taskId, ThumbnailData thumbnailData) {
+    /**
+     * Update the thumbnail of the task.
+     * @param refreshNow Refresh immediately if it's true.
+     */
+    public TaskView updateThumbnail(int taskId, ThumbnailData thumbnailData, boolean refreshNow) {
         TaskView taskView = getTaskView(taskId);
         if (taskView != null) {
-            taskView.getThumbnail().setThumbnail(taskView.getTask(), thumbnailData);
+            taskView.getThumbnail().setThumbnail(taskView.getTask(), thumbnailData, refreshNow);
         }
         return taskView;
     }
 
+    /** See {@link #updateThumbnail(int, ThumbnailData, boolean)} */
+    public TaskView updateThumbnail(int taskId, ThumbnailData thumbnailData) {
+        return updateThumbnail(taskId, thumbnailData, true /* refreshNow */);
+    }
+
     @Override
     protected void onWindowVisibilityChanged(int visibility) {
         super.onWindowVisibilityChanged(visibility);
@@ -573,11 +582,14 @@
             taskView.bind(task);
         }
 
-        TaskView runningTaskView = getRunningTaskView();
-        if (runningTaskView != null) {
-            setCurrentPage(indexOfChild(runningTaskView));
-        } else if (getTaskViewCount() > 0) {
-            setCurrentPage(indexOfChild(getTaskViewAt(0)));
+        if (mNextPage == INVALID_PAGE) {
+            // Set the current page to the running task, but not if settling on new task.
+            TaskView runningTaskView = getRunningTaskView();
+            if (runningTaskView != null) {
+                setCurrentPage(indexOfChild(runningTaskView));
+            } else if (getTaskViewCount() > 0) {
+                setCurrentPage(indexOfChild(getTaskViewAt(0)));
+            }
         }
 
         if (mIgnoreResetTaskId != -1 && getTaskView(mIgnoreResetTaskId) != ignoreResetTaskView) {
@@ -859,7 +871,9 @@
         setEnableFreeScroll(true);
         setEnableDrawingLiveTile(true);
         setOnScrollChangeListener(null);
-        setRunningTaskViewShowScreenshot(true);
+        if (!ENABLE_QUICKSTEP_LIVE_TILE.get()) {
+            setRunningTaskViewShowScreenshot(true);
+        }
         setRunningTaskHidden(false);
         animateUpRunningTaskIconScale();
     }
@@ -1791,10 +1805,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();
@@ -1819,4 +1842,19 @@
         final WindowInsets insets = getRootWindowInsets();
         return Math.max(insets.getSystemGestureInsets().right, insets.getSystemWindowInsetRight());
     }
+
+    /** If it's in the live tile mode, switch the running task into screenshot mode. */
+    public void switchToScreenshot(Runnable onFinishRunnable) {
+        TaskView taskView = getRunningTaskView();
+        if (taskView == null) {
+            if (onFinishRunnable != null) {
+                onFinishRunnable.run();
+            }
+            return;
+        }
+
+        taskView.setShowScreenshot(true);
+        taskView.getThumbnail().refresh();
+        ViewUtils.postDraw(taskView, onFinishRunnable);
+    }
 }
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
index 5799c01..adeb974 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -134,16 +134,35 @@
     }
 
     /**
-     * Updates this thumbnail.
+     * Updates the thumbnail.
+     * @param refreshNow whether the {@code thumbnailData} will be used to redraw immediately.
+     *                   In most cases, we use the {@link #setThumbnail(Task, ThumbnailData)}
+     *                   version with {@code refreshNow} is true. The only exception is
+     *                   in the live tile case that we grab a screenshot when user enters Overview
+     *                   upon swipe up so that a usable screenshot is accessible immediately when
+     *                   recents animation needs to be finished / cancelled.
      */
-    public void setThumbnail(Task task, ThumbnailData thumbnailData) {
+    public void setThumbnail(Task task, ThumbnailData thumbnailData, boolean refreshNow) {
         mTask = task;
-        if (thumbnailData != null && thumbnailData.thumbnail != null) {
-            Bitmap bm = thumbnailData.thumbnail;
+        mThumbnailData =
+                (thumbnailData != null && thumbnailData.thumbnail != null) ? thumbnailData : null;
+        if (refreshNow) {
+            refresh();
+        }
+    }
+
+    /** See {@link #setThumbnail(Task, ThumbnailData, boolean)} */
+    public void setThumbnail(Task task, ThumbnailData thumbnailData) {
+        setThumbnail(task, thumbnailData, true /* refreshNow */);
+    }
+
+    /** Updates the shader, paint, matrix to redraw. */
+    public void refresh() {
+        if (mThumbnailData != null && mThumbnailData.thumbnail != null) {
+            Bitmap bm = mThumbnailData.thumbnail;
             bm.prepareToDraw();
             mBitmapShader = new BitmapShader(bm, Shader.TileMode.CLAMP, Shader.TileMode.CLAMP);
             mPaint.setShader(mBitmapShader);
-            mThumbnailData = thumbnailData;
             updateThumbnailMatrix();
         } else {
             mBitmapShader = null;
diff --git a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
index 51802df..bfb9613 100644
--- a/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/recents_ui_overrides/src/com/android/quickstep/views/TaskView.java
@@ -53,7 +53,6 @@
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.logging.UserEventDispatcher;
-import com.android.launcher3.testing.TestProtocol;
 import com.android.launcher3.userevent.nano.LauncherLogProto;
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action.Direction;
 import com.android.launcher3.userevent.nano.LauncherLogProto.Action.Touch;
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 {
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index bf3cd8a..8e5ed1a 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -29,12 +29,16 @@
 import android.stats.launcher.nano.Launcher;
 import android.stats.launcher.nano.LauncherExtension;
 import android.stats.launcher.nano.LauncherTarget;
+import android.util.Log;
 import android.view.View;
 
 import com.android.launcher3.ItemInfo;
 import com.android.launcher3.logging.StatsLogManager;
 import com.android.launcher3.logging.StatsLogUtils;
 import com.android.launcher3.userevent.nano.LauncherLogProto.Target;
+import com.android.launcher3.userevent.nano.LauncherLogProto.ItemType;
+import com.android.launcher3.userevent.nano.LauncherLogProto.ContainerType;
+import com.android.launcher3.userevent.nano.LauncherLogProto.ControlType;
 import com.android.launcher3.util.ComponentKey;
 import com.android.systemui.shared.system.StatsLogCompat;
 import com.google.protobuf.nano.MessageNano;
@@ -50,6 +54,8 @@
 public class StatsLogCompatManager extends StatsLogManager {
 
     private static final int SUPPORTED_TARGET_DEPTH = 2;
+    private static final String TAG = "StatsLogCompatManager";
+    private static final boolean DEBUG = false;
 
     public StatsLogCompatManager(Context context) { }
 
@@ -59,6 +65,9 @@
         ext.srcTarget = new LauncherTarget[SUPPORTED_TARGET_DEPTH];
         int srcState = mStateProvider.getCurrentState();
         fillInLauncherExtension(v, ext);
+        if (ext.srcTarget[0] != null) {
+            ext.srcTarget[0].item = LauncherTarget.APP_ICON;
+        }
         StatsLogCompat.write(LAUNCH_APP, srcState, BACKGROUND /* dstState */,
                 MessageNano.toByteArray(ext), true);
     }
@@ -95,28 +104,132 @@
     }
 
     public static boolean fillInLauncherExtension(View v, LauncherExtension extension) {
+        if (DEBUG) {
+            Log.d(TAG, "fillInLauncherExtension");
+        }
+
         StatsLogUtils.LogContainerProvider provider = StatsLogUtils.getLaunchProviderRecursive(v);
         if (v == null || !(v.getTag() instanceof ItemInfo) || provider == null) {
+            if (DEBUG) {
+                Log.d(TAG, "View or provider is null, or view doesn't have an ItemInfo tag.");
+            }
+
             return false;
         }
         ItemInfo itemInfo = (ItemInfo) v.getTag();
         Target child = new Target();
         Target parent = new Target();
         provider.fillInLogContainerData(v, itemInfo, child, parent);
+        extension.srcTarget[0] = new LauncherTarget();
+        extension.srcTarget[1] = new LauncherTarget();
         copy(child, extension.srcTarget[0]);
         copy(parent, extension.srcTarget[1]);
         return true;
     }
 
     public static boolean fillInLauncherExtensionWithPageId(LauncherExtension ext, int pageId) {
+        if (DEBUG) {
+            Log.d(TAG, "fillInLauncherExtensionWithPageId, pageId = " + pageId);
+        }
+
         Target target = new Target();
         target.pageIndex = pageId;
+        ext.srcTarget[0] = new LauncherTarget();
         copy(target, ext.srcTarget[0]);
         return true;
     }
 
     private static void copy(Target src, LauncherTarget dst) {
-        // fill in
+        if (DEBUG) {
+            Log.d(TAG, "copy target information from clearcut Target to LauncherTarget.");
+        }
+
+        // Fill in type
+        switch (src.type) {
+            case Target.Type.ITEM:
+                dst.type = LauncherTarget.ITEM_TYPE;
+                break;
+            case Target.Type.CONTROL:
+                dst.type = LauncherTarget.CONTROL_TYPE;
+                break;
+            case Target.Type.CONTAINER:
+                dst.type = LauncherTarget.CONTAINER_TYPE;
+                break;
+            default:
+                dst.type = LauncherTarget.NONE;
+                break;
+        }
+
+        // Fill in item
+        switch (src.itemType) {
+            case ItemType.APP_ICON:
+                dst.item = LauncherTarget.APP_ICON;
+                break;
+            case ItemType.SHORTCUT:
+                dst.item = LauncherTarget.SHORTCUT;
+                break;
+            case ItemType.WIDGET:
+                dst.item = LauncherTarget.WIDGET;
+                break;
+            case ItemType.FOLDER_ICON:
+                dst.item = LauncherTarget.FOLDER_ICON;
+                break;
+            case ItemType.DEEPSHORTCUT:
+                dst.item = LauncherTarget.DEEPSHORTCUT;
+                break;
+            case ItemType.SEARCHBOX:
+                dst.item = LauncherTarget.SEARCHBOX;
+                break;
+            case ItemType.EDITTEXT:
+                dst.item = LauncherTarget.EDITTEXT;
+                break;
+            case ItemType.NOTIFICATION:
+                dst.item = LauncherTarget.NOTIFICATION;
+                break;
+            case ItemType.TASK:
+                dst.item = LauncherTarget.TASK;
+                break;
+            default:
+                dst.item = LauncherTarget.DEFAULT_ITEM;
+                break;
+        }
+
+        // Fill in container
+        switch (src.containerType) {
+            case ContainerType.HOTSEAT:
+                dst.container = LauncherTarget.HOTSEAT;
+                break;
+            case ContainerType.FOLDER:
+                dst.container = LauncherTarget.FOLDER;
+                break;
+            case ContainerType.PREDICTION:
+                dst.container = LauncherTarget.PREDICTION;
+                break;
+            case ContainerType.SEARCHRESULT:
+                dst.container = LauncherTarget.SEARCHRESULT;
+                break;
+            default:
+                dst.container = LauncherTarget.DEFAULT_CONTAINER;
+                break;
+        }
+
+        // Fill in control
+        switch (src.controlType) {
+            case ControlType.UNINSTALL_TARGET:
+                dst.control = LauncherTarget.UNINSTALL;
+                break;
+            case ControlType.REMOVE_TARGET:
+                dst.control = LauncherTarget.REMOVE;
+                break;
+            default:
+                dst.control = LauncherTarget.DEFAULT_CONTROL;
+                break;
+        }
+
+        // Fill in other fields
+        dst.pageId = src.pageIndex;
+        dst.gridX = src.gridX;
+        dst.gridY = src.gridY;
     }
 
     @Override
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index e295527..c2197ab 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -181,7 +181,7 @@
                                 SysUINavigationMode.INSTANCE.get(targetContext);
                         targetContext.getMainExecutor().execute(() ->
                                 sysUINavigationMode.addModeChangeListener(listener));
-                        latch.await(10, TimeUnit.SECONDS);
+                        latch.await(60, TimeUnit.SECONDS);
                         targetContext.getMainExecutor().execute(() ->
                                 sysUINavigationMode.removeModeChangeListener(listener));
                         assertTrue("Navigation mode didn't change to " + expectedMode,
diff --git a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
index a46617e..b7340cf 100644
--- a/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
+++ b/robolectric_tests/src/com/android/launcher3/model/DbDowngradeHelperTest.java
@@ -130,7 +130,7 @@
         }
         helper.close();
 
-        helper = new DatabaseHelper(mContext, null, DB_FILE) {
+        helper = new DatabaseHelper(mContext, DB_FILE) {
             @Override
             public void onOpen(SQLiteDatabase db) { }
         };
@@ -161,7 +161,7 @@
 
         DbDowngradeHelper.updateSchemaFile(mSchemaFile, LauncherProvider.SCHEMA_VERSION, mContext);
 
-        DatabaseHelper dbHelper = new DatabaseHelper(mContext, null, DB_FILE) {
+        DatabaseHelper dbHelper = new DatabaseHelper(mContext, DB_FILE) {
             @Override
             public void onOpen(SQLiteDatabase db) { }
         };
diff --git a/robolectric_tests/src/com/android/launcher3/util/TestLauncherProvider.java b/robolectric_tests/src/com/android/launcher3/util/TestLauncherProvider.java
index 9f833b1..a9c1a7c 100644
--- a/robolectric_tests/src/com/android/launcher3/util/TestLauncherProvider.java
+++ b/robolectric_tests/src/com/android/launcher3/util/TestLauncherProvider.java
@@ -29,7 +29,7 @@
 
     private static class MyDatabaseHelper extends DatabaseHelper {
         public MyDatabaseHelper(Context context) {
-            super(context, null, null);
+            super(context, null);
             initIds();
         }
 
diff --git a/src/com/android/launcher3/AbstractFloatingView.java b/src/com/android/launcher3/AbstractFloatingView.java
index 8fddf3c..af2cdc3 100644
--- a/src/com/android/launcher3/AbstractFloatingView.java
+++ b/src/com/android/launcher3/AbstractFloatingView.java
@@ -16,7 +16,6 @@
 
 package com.android.launcher3;
 
-import static android.view.accessibility.AccessibilityEvent.TYPE_VIEW_FOCUSED;
 import static android.view.accessibility.AccessibilityEvent.TYPE_WINDOW_CONTENT_CHANGED;
 import static android.view.accessibility.AccessibilityEvent.TYPE_WINDOW_STATE_CHANGED;
 
@@ -30,6 +29,7 @@
 import android.util.Pair;
 import android.view.MotionEvent;
 import android.view.View;
+import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.LinearLayout;
 
 import androidx.annotation.IntDef;
@@ -171,7 +171,7 @@
                 targetInfo.first, TYPE_WINDOW_STATE_CHANGED, targetInfo.second);
 
         if (mIsOpen) {
-            sendAccessibilityEvent(TYPE_VIEW_FOCUSED);
+            performAccessibilityAction(AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS, null);
         }
         ActivityContext.lookupContext(getContext()).getDragLayer()
                 .sendAccessibilityEvent(TYPE_WINDOW_CONTENT_CHANGED);
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index c034d2d..736142f 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -111,6 +111,7 @@
 
     // All apps
     public int allAppsCellHeightPx;
+    public int allAppsCellWidthPx;
     public int allAppsIconSizePx;
     public int allAppsIconDrawablePaddingPx;
     public float allAppsIconTextSizePx;
@@ -349,6 +350,7 @@
             allAppsIconDrawablePaddingPx = iconDrawablePaddingPx;
             allAppsCellHeightPx = getCellSize().y;
         }
+        allAppsCellWidthPx = allAppsIconSizePx + allAppsIconDrawablePaddingPx;
 
         if (isVerticalBarLayout()) {
             // Always hide the Workspace text with vertical bar layout.
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 9326e95..aff4c90 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -173,7 +173,7 @@
  * Default launcher application.
  */
 public class Launcher extends BaseDraggingActivity implements LauncherExterns,
-        Callbacks, LauncherProviderChangeListener, UserEventDelegate,
+        Callbacks, UserEventDelegate,
         InvariantDeviceProfile.OnIDPChangeListener, PluginListener<OverlayPlugin> {
     public static final String TAG = "Launcher";
 
@@ -614,13 +614,6 @@
         return mLauncherView.findViewById(id);
     }
 
-    @Override
-    public void onAppWidgetHostReset() {
-        if (mAppWidgetHost != null) {
-            mAppWidgetHost.startListening();
-        }
-    }
-
     private LauncherCallbacks mLauncherCallbacks;
 
     /**
@@ -1896,7 +1889,7 @@
             // recents animation into launcher. Defer launching the activity until Launcher is
             // next resumed.
             addOnResumeCallback(() -> startActivitySafely(v, intent, item, sourceContainer));
-            UiFactory.clearSwipeSharedState(true /* finishAnimation */);
+            UiFactory.clearSwipeSharedState(this, true /* finishAnimation */);
             return true;
         }
 
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index efb3d36..018b48f 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -21,7 +21,6 @@
 import static com.android.launcher3.util.SecureSettingsObserver.newNotificationSettingsObserver;
 
 import android.content.ComponentName;
-import android.content.ContentProviderClient;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
@@ -71,10 +70,6 @@
     }
 
     private LauncherAppState(Context context) {
-        if (getLocalProvider(context) == null) {
-            throw new RuntimeException(
-                    "Initializing LauncherAppState in the absence of LauncherProvider");
-        }
         Log.v(Launcher.TAG, "LauncherAppState initiated");
         Preconditions.assertUIThread();
         mContext = context;
@@ -99,7 +94,10 @@
         if (FeatureFlags.IS_DOGFOOD_BUILD) {
             filter.addAction(ACTION_FORCE_ROLOAD);
         }
+        // TODO: remove listener on terminate
         FeatureFlags.APP_SEARCH_IMPROVEMENTS.addChangeListener(context, mModel::forceReload);
+        CustomWidgetManager.INSTANCE.get(mContext)
+                .setWidgetRefreshCallback(mModel::refreshAndBindWidgetsAndShortcuts);
 
         mContext.registerReceiver(mModel, filter);
         UserManagerCompat.getInstance(mContext).enableAndResetCache();
@@ -149,16 +147,15 @@
         final LauncherAppsCompat launcherApps = LauncherAppsCompat.getInstance(mContext);
         launcherApps.removeOnAppsChangedCallback(mModel);
         mInstallSessionTracker.unregister();
+        CustomWidgetManager.INSTANCE.get(mContext).setWidgetRefreshCallback(null);
+
         if (mNotificationDotsObserver != null) {
             mNotificationDotsObserver.unregister();
         }
     }
 
     LauncherModel setLauncher(Launcher launcher) {
-        getLocalProvider(mContext).setLauncherProviderChangeListener(launcher);
         mModel.initialize(launcher);
-        CustomWidgetManager.INSTANCE.get(launcher)
-                .setWidgetRefreshCallback(mModel::refreshAndBindWidgetsAndShortcuts);
         return mModel;
     }
 
@@ -184,11 +181,4 @@
     public static InvariantDeviceProfile getIDP(Context context) {
         return InvariantDeviceProfile.INSTANCE.get(context);
     }
-
-    private static LauncherProvider getLocalProvider(Context context) {
-        try (ContentProviderClient cl = context.getContentResolver()
-                .acquireContentProviderClient(LauncherProvider.AUTHORITY)) {
-            return (LauncherProvider) cl.getLocalContentProvider();
-        }
-    }
 }
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index d78c1b3..42927ea 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -45,8 +45,6 @@
 import android.os.Binder;
 import android.os.Build;
 import android.os.Bundle;
-import android.os.Handler;
-import android.os.Message;
 import android.os.Process;
 import android.os.UserHandle;
 import android.provider.BaseColumns;
@@ -69,7 +67,6 @@
 import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.NoLocaleSQLiteHelper;
 import com.android.launcher3.util.PackageManagerHelper;
-import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.util.Thunk;
 
 import org.xmlpull.v1.XmlPullParser;
@@ -100,9 +97,6 @@
 
     static final String EMPTY_DATABASE_CREATED = "EMPTY_DATABASE_CREATED";
 
-    private final ChangeListenerWrapper mListenerWrapper = new ChangeListenerWrapper();
-    private Handler mListenerHandler;
-
     protected DatabaseHelper mOpenHelper;
 
     /**
@@ -122,7 +116,6 @@
         if (FeatureFlags.IS_DOGFOOD_BUILD) {
             Log.d(TAG, "Launcher process started");
         }
-        mListenerHandler = new Handler(mListenerWrapper);
 
         // The content provider exists for the entire duration of the launcher main process and
         // is the first component to get created.
@@ -130,14 +123,6 @@
         return true;
     }
 
-    /**
-     * Sets a provider listener.
-     */
-    public void setLauncherProviderChangeListener(LauncherProviderChangeListener listener) {
-        Preconditions.assertUIThread();
-        mListenerWrapper.mListener = listener;
-    }
-
     @Override
     public String getType(Uri uri) {
         SqlArguments args = new SqlArguments(uri, null, null);
@@ -153,7 +138,7 @@
      */
     protected synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
-            mOpenHelper = new DatabaseHelper(getContext(), mListenerHandler);
+            mOpenHelper = new DatabaseHelper(getContext());
 
             if (RestoreDbTask.isPending(getContext())) {
                 if (!RestoreDbTask.performRestore(getContext(), mOpenHelper,
@@ -552,15 +537,13 @@
      * The class is subclassed in tests to create an in-memory db.
      */
     public static class DatabaseHelper extends NoLocaleSQLiteHelper implements LayoutParserCallback {
-        private final BackupManager mBackupManager;
-        private final Handler mWidgetHostResetHandler;
         private final Context mContext;
         private int mMaxItemId = -1;
         private int mMaxScreenId = -1;
         private boolean mBackupTableExists;
 
-        DatabaseHelper(Context context, Handler widgetHostResetHandler) {
-            this(context, widgetHostResetHandler, LauncherFiles.LAUNCHER_DB);
+        DatabaseHelper(Context context) {
+            this(context, LauncherFiles.LAUNCHER_DB);
             // Table creation sometimes fails silently, which leads to a crash loop.
             // This way, we will try to create a table every time after crash, so the device
             // would eventually be able to recover.
@@ -577,12 +560,9 @@
         /**
          * Constructor used in tests and for restore.
          */
-        public DatabaseHelper(
-                Context context, Handler widgetHostResetHandler, String tableName) {
+        public DatabaseHelper(Context context, String tableName) {
             super(context, tableName, SCHEMA_VERSION);
             mContext = context;
-            mWidgetHostResetHandler = widgetHostResetHandler;
-            mBackupManager = new BackupManager(mContext);
         }
 
         protected void initIds() {
@@ -621,13 +601,6 @@
          * Overriden in tests.
          */
         protected void onEmptyDbCreated() {
-            // Database was just created, so wipe any previous widgets
-            if (mWidgetHostResetHandler != null) {
-                newLauncherWidgetHost().deleteHost();
-                mWidgetHostResetHandler.sendEmptyMessage(
-                        ChangeListenerWrapper.MSG_APP_WIDGET_HOST_RESET);
-            }
-
             // Set the flag for empty DB
             Utilities.getPrefs(mContext).edit().putBoolean(EMPTY_DATABASE_CREATED, true).commit();
         }
@@ -1027,23 +1000,4 @@
             }
         }
     }
-
-    private static class ChangeListenerWrapper implements Handler.Callback {
-
-        private static final int MSG_APP_WIDGET_HOST_RESET = 2;
-
-        private LauncherProviderChangeListener mListener;
-
-        @Override
-        public boolean handleMessage(Message msg) {
-            if (mListener != null) {
-                switch (msg.what) {
-                    case MSG_APP_WIDGET_HOST_RESET:
-                        mListener.onAppWidgetHostReset();
-                        break;
-                }
-            }
-            return true;
-        }
-    }
 }
diff --git a/src/com/android/launcher3/LauncherProviderChangeListener.java b/src/com/android/launcher3/LauncherProviderChangeListener.java
deleted file mode 100644
index 6afe885..0000000
--- a/src/com/android/launcher3/LauncherProviderChangeListener.java
+++ /dev/null
@@ -1,11 +0,0 @@
-package com.android.launcher3;
-
-/**
- * This class is a listener for {@link LauncherProvider} changes. It gets notified in the
- * sendNotify method. This listener is needed because by default the Launcher suppresses
- * standard data change callbacks.
- */
-public interface LauncherProviderChangeListener {
-
-    void onAppWidgetHostReset();
-}
diff --git a/src/com/android/launcher3/PagedView.java b/src/com/android/launcher3/PagedView.java
index 5c790f3..ff2b400 100644
--- a/src/com/android/launcher3/PagedView.java
+++ b/src/com/android/launcher3/PagedView.java
@@ -1088,7 +1088,9 @@
         if (mFreeScroll) {
             setCurrentPage(getNextPage());
         } else if (wasFreeScroll) {
-            snapToPage(getNextPage());
+            if (getScrollForPage(getNextPage()) != getScrollX()) {
+                snapToPage(getNextPage());
+            }
         }
     }
 
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 6ddebe7..c5f26aa 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -584,7 +584,7 @@
             LauncherIcons li = LauncherIcons.obtain(appState.getContext());
             Bitmap badge = li.getShortcutInfoBadge(si, appState.getIconCache()).iconBitmap;
             li.recycle();
-            float badgeSize = launcher.getResources().getDimension(R.dimen.profile_badge_size);
+            float badgeSize = iconSize * LauncherIcons.getBadgeSizeForIconSize(iconSize);
             float insetFraction = (iconSize - badgeSize) / iconSize;
             return new InsetDrawable(new FastBitmapDrawable(badge),
                     insetFraction, insetFraction, 0, 0);
diff --git a/src/com/android/launcher3/WidgetPreviewLoader.java b/src/com/android/launcher3/WidgetPreviewLoader.java
index abf6cbd..b25a79b 100644
--- a/src/com/android/launcher3/WidgetPreviewLoader.java
+++ b/src/com/android/launcher3/WidgetPreviewLoader.java
@@ -467,7 +467,7 @@
 
     private Bitmap generateShortcutPreview(BaseActivity launcher, ShortcutConfigActivityInfo info,
             int maxWidth, int maxHeight, Bitmap preview) {
-        int iconSize = launcher.getDeviceProfile().iconSizePx;
+        int iconSize = launcher.getDeviceProfile().allAppsIconSizePx;
         int padding = launcher.getResources()
                 .getDimensionPixelSize(R.dimen.widget_preview_shortcut_padding);
 
diff --git a/src/com/android/launcher3/graphics/DrawableFactory.java b/src/com/android/launcher3/graphics/DrawableFactory.java
index 288749f..837301f 100644
--- a/src/com/android/launcher3/graphics/DrawableFactory.java
+++ b/src/com/android/launcher3/graphics/DrawableFactory.java
@@ -80,26 +80,25 @@
      * Returns a drawable that can be used as a badge for the user or null.
      */
     @UiThread
-    public Drawable getBadgeForUser(UserHandle user, Context context) {
+    public Drawable getBadgeForUser(UserHandle user, Context context, int badgeSize) {
         if (mMyUser.equals(user)) {
             return null;
         }
 
-        Bitmap badgeBitmap = getUserBadge(user, context);
+        Bitmap badgeBitmap = getUserBadge(user, context, badgeSize);
         FastBitmapDrawable d = new FastBitmapDrawable(badgeBitmap);
         d.setFilterBitmap(true);
         d.setBounds(0, 0, badgeBitmap.getWidth(), badgeBitmap.getHeight());
         return d;
     }
 
-    protected synchronized Bitmap getUserBadge(UserHandle user, Context context) {
+    protected synchronized Bitmap getUserBadge(UserHandle user, Context context, int badgeSize) {
         Bitmap badgeBitmap = mUserBadges.get(user);
         if (badgeBitmap != null) {
             return badgeBitmap;
         }
 
         final Resources res = context.getApplicationContext().getResources();
-        int badgeSize = res.getDimensionPixelSize(R.dimen.profile_badge_size);
         badgeBitmap = Bitmap.createBitmap(badgeSize, badgeSize, Bitmap.Config.ARGB_8888);
 
         Drawable drawable = context.getPackageManager().getUserBadgedDrawableForDensity(
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index 05596bd..f9b2d5f 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -35,7 +35,9 @@
 import com.android.launcher3.util.ResourceBasedOverride;
 
 import java.util.LinkedList;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
 
 public class TestInformationHandler implements ResourceBasedOverride {
 
@@ -136,6 +138,7 @@
             }
 
             case TestProtocol.REQUEST_TOTAL_PSS_KB: {
+                runGcAndFinalizersSync();
                 Debug.MemoryInfo mem = new Debug.MemoryInfo();
                 Debug.getMemoryInfo(mem);
                 response.putInt(TestProtocol.TEST_INFO_RESPONSE_FIELD, mem.getTotalPss());
@@ -172,4 +175,29 @@
         return Launcher.ACTIVITY_TRACKER.getCreatedActivity() == null
                 || LauncherAppState.getInstance(mContext).getModel().isModelLoaded();
     }
+
+    private static void runGcAndFinalizersSync() {
+        Runtime.getRuntime().gc();
+        Runtime.getRuntime().runFinalization();
+
+        final CountDownLatch fence = new CountDownLatch(1);
+        new Object() {
+            @Override
+            protected void finalize() throws Throwable {
+                try {
+                    fence.countDown();
+                } finally {
+                    super.finalize();
+                }
+            }
+        };
+        try {
+            do {
+                Runtime.getRuntime().gc();
+                Runtime.getRuntime().runFinalization();
+            } while (!fence.await(100, TimeUnit.MILLISECONDS));
+        } catch (InterruptedException ex) {
+            throw new RuntimeException(ex);
+        }
+    }
 }
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index dc4af8c..6944879 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -37,6 +37,7 @@
 import com.android.launcher3.StylusEventHelper;
 import com.android.launcher3.WidgetPreviewLoader;
 import com.android.launcher3.graphics.DrawableFactory;
+import com.android.launcher3.icons.BaseIconFactory;
 import com.android.launcher3.model.WidgetItem;
 
 /**
@@ -80,6 +81,7 @@
     private Bitmap mDeferredBitmap;
 
     protected final BaseActivity mActivity;
+    protected DeviceProfile mDeviceProfile;
 
     public WidgetCell(Context context) {
         this(context, null);
@@ -93,6 +95,7 @@
         super(context, attrs, defStyle);
 
         mActivity = BaseActivity.fromContext(context);
+        mDeviceProfile = mActivity.getDeviceProfile();
         mStylusEventHelper = new StylusEventHelper(new SimpleOnStylusPressListener(this), this);
 
         setContainerWidth();
@@ -102,8 +105,7 @@
     }
 
     private void setContainerWidth() {
-        DeviceProfile profile = mActivity.getDeviceProfile();
-        mCellSize = (int) (profile.cellWidthPx * WIDTH_SCALE);
+        mCellSize = (int) (mDeviceProfile.allAppsCellWidthPx * WIDTH_SCALE);
         mPresetPreviewSize = (int) (mCellSize * PREVIEW_SCALE);
     }
 
@@ -180,8 +182,10 @@
             return;
         }
         if (bitmap != null) {
-            mWidgetImage.setBitmap(bitmap, DrawableFactory.INSTANCE.get(getContext())
-                    .getBadgeForUser(mItem.user, getContext()));
+            mWidgetImage.setBitmap(bitmap,
+                    DrawableFactory.INSTANCE.get(getContext()).getBadgeForUser(mItem.user,
+                            getContext(), BaseIconFactory.getBadgeSizeForIconSize(
+                                    mDeviceProfile.allAppsIconSizePx)));
             if (mAnimatePreview) {
                 mWidgetImage.setAlpha(0f);
                 ViewPropertyAnimator anim = mWidgetImage.animate();
diff --git a/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java b/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
index 467ae02..6d9ed88 100644
--- a/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
+++ b/src_ui_overrides/com/android/launcher3/uioverrides/UiFactory.java
@@ -96,7 +96,8 @@
 
     public static void resetPendingActivityResults(Launcher launcher, int requestCode) { }
 
-    public static void clearSwipeSharedState(boolean finishAnimation) {}
+    /** No-op. */
+    public static void clearSwipeSharedState(Launcher launcher, boolean finishAnimation) { }
 
     public static Person[] getPersons(ShortcutInfo si) {
         return Utilities.EMPTY_PERSON_ARRAY;
diff --git a/tests/AndroidManifest-common.xml b/tests/AndroidManifest-common.xml
index 0b74dc4..56eca6d 100644
--- a/tests/AndroidManifest-common.xml
+++ b/tests/AndroidManifest-common.xml
@@ -62,6 +62,12 @@
                 <action android:name="android.appwidget.action.APPWIDGET_CONFIGURE"/>
             </intent-filter>
         </activity>
+        <activity android:name="com.android.launcher3.testcomponent.CustomShortcutConfigActivity">
+            <intent-filter>
+                <action android:name="android.intent.action.CREATE_SHORTCUT" />
+                <category android:name="android.intent.category.DEFAULT" />
+            </intent-filter>
+        </activity>
         <activity
             android:name="com.android.launcher3.testcomponent.RequestPinItemActivity"
             android:icon="@drawable/test_drawable_pin_item"
diff --git a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 6fa8d62..27990f4 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -1,8 +1,11 @@
 package com.android.launcher3.provider;
 
+import static org.junit.Assert.assertEquals;
+
 import android.content.ContentValues;
 import android.database.Cursor;
 import android.database.sqlite.SQLiteDatabase;
+
 import androidx.test.InstrumentationRegistry;
 import androidx.test.filters.MediumTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -13,8 +16,6 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
-import static org.junit.Assert.assertEquals;
-
 /**
  * Tests for {@link RestoreDbTask}
  */
@@ -82,7 +83,7 @@
         private final long mProfileId;
 
         MyDatabaseHelper(long profileId) {
-            super(InstrumentationRegistry.getContext(), null, null);
+            super(InstrumentationRegistry.getContext(), null);
             mProfileId = profileId;
         }
 
diff --git a/tests/src/com/android/launcher3/testcomponent/CustomShortcutConfigActivity.java b/tests/src/com/android/launcher3/testcomponent/CustomShortcutConfigActivity.java
new file mode 100644
index 0000000..b673faa
--- /dev/null
+++ b/tests/src/com/android/launcher3/testcomponent/CustomShortcutConfigActivity.java
@@ -0,0 +1,66 @@
+/*
+ * Copyright (C) 2019 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.launcher3.testcomponent;
+
+import android.content.Context;
+import android.content.Intent;
+import android.content.pm.ShortcutInfo;
+import android.content.pm.ShortcutManager;
+import android.graphics.drawable.Icon;
+import android.os.Bundle;
+
+import com.android.launcher3.R;
+
+import java.util.UUID;
+
+/**
+ * A custom shortcut is a 1x1 widget that launches a specific intent when user tap on it.
+ * Custom shortcuts are replaced by deep shortcuts after api 25.
+ */
+public class CustomShortcutConfigActivity extends BaseTestingActivity {
+
+    @Override
+    protected void onCreate(Bundle savedInstanceState) {
+        super.onCreate(savedInstanceState);
+
+        Intent launchIntent = new Intent(this, BaseTestingActivity.class)
+                .setAction("com.android.launcher3.intent.action.test_shortcut");
+        Intent shortcutIntent = createShortcutResultIntent(
+                this, UUID.randomUUID().toString(), "Shortcut",
+                R.drawable.ic_widget, launchIntent);
+        setResult(RESULT_OK, shortcutIntent);
+        finish();
+    }
+
+    private static Intent createShortcutResultIntent(
+            Context context, String uniqueId, String name, int iconId, Intent launchIntent) {
+        ShortcutInfo shortcutInfo =
+                createShortcutInfo(context, uniqueId, name, iconId, launchIntent);
+        ShortcutManager sm = context.getSystemService(ShortcutManager.class);
+        return sm.createShortcutResultIntent(shortcutInfo);
+    }
+
+    private static ShortcutInfo createShortcutInfo(
+            Context context, String uniqueId, String name, int iconId, Intent launchIntent) {
+        return new ShortcutInfo.Builder(context, uniqueId)
+                .setShortLabel(name)
+                .setLongLabel(name)
+                .setIcon(Icon.createWithResource(context, iconId))
+                .setIntent(launchIntent)
+                .build();
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index 2cf6c2b..709822b 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -342,6 +342,22 @@
         }
     }
 
+    /**
+     * Test dragging a custom shortcut to the workspace and launch it.
+     *
+     * A custom shortcut is a 1x1 widget that launches a specific intent when user tap on it.
+     * Custom shortcuts are replaced by deep shortcuts after api 25.
+     */
+    @Test
+    @PortraitLandscape
+    public void testDragCustomShortcut() {
+        mLauncher.getWorkspace().openAllWidgets()
+                .getWidget("com.android.launcher3.testcomponent.CustomShortcutConfigActivity")
+                .dragToWorkspace();
+        mLauncher.getWorkspace().getWorkspaceAppIcon("Shortcut")
+                .launch(getAppPackageName());
+    }
+
     public static String getAppPackageName() {
         return getInstrumentation().getContext().getPackageName();
     }
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index 2f7e50a..afdabea 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -534,6 +534,9 @@
         // accessibility events prior to pressing Home.
         final String action;
         if (getNavigationModel() == NavigationModel.ZERO_BUTTON) {
+            final String anomaly = getAnomalyMessage();
+            if (anomaly != null) fail("Can't swipe up to Home: " + anomaly);
+
             final Point displaySize = getRealDisplaySize();
 
             if (hasLauncherObject("deep_shortcuts_container")) {
@@ -541,13 +544,17 @@
                         displaySize.x / 2, displaySize.y - 1,
                         displaySize.x / 2, 0,
                         ZERO_BUTTON_STEPS_FROM_BACKGROUND_TO_HOME);
-                assertTrue("Context menu is still visible after swiping up to home",
-                        !hasLauncherObject("deep_shortcuts_container"));
+                try (LauncherInstrumentation.Closable c = addContextLayer(
+                        "Swiped up from context menu to home")) {
+                    waitUntilGone("deep_shortcuts_container");
+                }
             }
             if (hasLauncherObject(WORKSPACE_RES_ID)) {
                 log(action = "already at home");
             } else {
-                log(action = "swiping up to home");
+                log("Hierarchy before swiping up to home");
+                dumpViewHierarchy();
+                log(action = "swiping up to home from " + getVisibleStateMessage());
                 final int finalState = mDevice.hasObject(By.pkg(getLauncherPackageName()))
                         ? NORMAL_STATE_ORDINAL : BACKGROUND_APP_STATE_ORDINAL;