diff --git a/go/quickstep/res/drawable/round_rect_dialog.xml b/go/quickstep/res/drawable/round_rect_dialog.xml
new file mode 100644
index 0000000..bbb7c5b
--- /dev/null
+++ b/go/quickstep/res/drawable/round_rect_dialog.xml
@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 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.
+-->
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+       android:shape="rectangle">
+    <corners android:radius="@dimen/modal_dialog_corner_radius" />
+</shape>
diff --git a/go/quickstep/res/layout/niu_actions_confirmation_dialog.xml b/go/quickstep/res/layout/niu_actions_confirmation_dialog.xml
new file mode 100644
index 0000000..db1531a
--- /dev/null
+++ b/go/quickstep/res/layout/niu_actions_confirmation_dialog.xml
@@ -0,0 +1,94 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2021 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.
+-->
+
+<LinearLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:id="@+id/niu_actions_confirmation_dialog_layout"
+    android:layout_width="match_parent"
+    android:layout_height="match_parent"
+    android:orientation="horizontal"
+    android:layout_gravity="center">
+
+    <Space
+        android:layout_width="0dp"
+        android:layout_height="1dp"
+        android:layout_weight="1"/>
+
+    <LinearLayout
+        android:layout_width="@dimen/modal_dialog_width"
+        android:layout_height="wrap_content"
+        android:background="@drawable/round_rect_dialog"
+        android:backgroundTint="?attr/modalDialogBackground"
+        android:orientation="vertical"
+        android:layout_gravity="center"
+        android:paddingTop="@dimen/modal_dialog_padding"
+        android:paddingLeft="@dimen/modal_dialog_padding"
+        android:paddingRight="@dimen/modal_dialog_padding"
+        android:paddingBottom="@dimen/modal_dialog_padding_bottom">
+
+        <TextView
+            style="@style/ModalDialogTitle"
+            android:id="@+id/niu_actions_confirmation_header"
+            android:text="@string/niu_actions_confirmation_title"/>
+
+        <Space
+            android:layout_width="0dp"
+            android:layout_height="@dimen/modal_dialog_vertical_spacer"/>
+
+        <ScrollView
+            android:layout_width="match_parent"
+            android:layout_height="@dimen/confirmation_dialog_text_height">
+
+            <TextView
+                style="@style/ModalDialogText"
+                android:id="@+id/niu_actions_confirmation_description"
+                android:text="@string/niu_actions_confirmation_text"/>
+        </ScrollView>
+
+        <Space
+            android:layout_width="0dp"
+            android:layout_height="@dimen/modal_dialog_vertical_spacer"/>
+
+        <LinearLayout
+            android:id="@+id/niu_actions_confirmation_buttons"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:orientation="horizontal">
+
+            <Space
+                android:layout_width="0dp"
+                android:layout_height="0dp"
+                android:layout_weight="1"/>
+
+            <Button
+                style="@style/ModalDialogButton"
+                android:id="@+id/niu_actions_confirmation_reject"
+                android:text="@string/niu_actions_confirmation_no"/>
+
+            <Button
+                style="@style/ModalDialogButton"
+                android:id="@+id/niu_actions_confirmation_accept"
+                android:text="@string/niu_actions_confirmation_yes"/>
+        </LinearLayout>
+
+    </LinearLayout>
+
+    <Space
+        android:layout_width="0dp"
+        android:layout_height="1dp"
+        android:layout_weight="1" />
+
+</LinearLayout>
\ No newline at end of file
diff --git a/go/quickstep/res/values-land/dimens.xml b/go/quickstep/res/values-land/dimens.xml
new file mode 100644
index 0000000..5097016
--- /dev/null
+++ b/go/quickstep/res/values-land/dimens.xml
@@ -0,0 +1,22 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 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.
+-->
+
+<resources>
+    <!-- Modal Dialogs -->
+    <dimen name="modal_dialog_width">360dp</dimen>
+    <dimen name="confirmation_dialog_text_height">168dp</dimen>
+</resources>
diff --git a/go/quickstep/res/values/attrs.xml b/go/quickstep/res/values/attrs.xml
index 03eab50..cdbdc2a 100644
--- a/go/quickstep/res/values/attrs.xml
+++ b/go/quickstep/res/values/attrs.xml
@@ -19,4 +19,6 @@
     <attr name="overviewButtonTextColor" format="color" />
     <attr name="overviewButtonIconColor" format="color" />
     <attr name="overviewButtonBackgroundColor" format="color" />
+    <!-- Modal dialog theming -->
+    <attr name="modalDialogBackground" format="color" />
 </resources>
\ No newline at end of file
diff --git a/go/quickstep/res/values/colors.xml b/go/quickstep/res/values/colors.xml
index ae72ef6..8034be2 100644
--- a/go/quickstep/res/values/colors.xml
+++ b/go/quickstep/res/values/colors.xml
@@ -20,4 +20,7 @@
     <color name="go_overview_text_color_dark">#F8F9FA</color>
     <color name="go_overview_button_color">#70FFFFFF</color>
     <color name="go_overview_button_color_dark">#474747</color>
+    <!-- Modal Dialogs -->
+    <color name="go_modal_dialog_background">#FFFFFF</color>
+    <color name="go_modal_dialog_background_dark">#424242</color>
 </resources>
diff --git a/go/quickstep/res/values/dimens.xml b/go/quickstep/res/values/dimens.xml
index 55cd138..0a7ac45 100644
--- a/go/quickstep/res/values/dimens.xml
+++ b/go/quickstep/res/values/dimens.xml
@@ -28,4 +28,12 @@
     <dimen name="overview_proactive_row_height">0dp</dimen>
     <dimen name="overview_proactive_row_bottom_margin">24dp</dimen>
     <dimen name="task_corner_radius_override">28dp</dimen>
+
+    <!-- Modal Dialogs -->
+    <dimen name="modal_dialog_width">288dp</dimen>
+    <dimen name="modal_dialog_padding">24dp</dimen>
+    <dimen name="modal_dialog_padding_bottom">8dp</dimen>
+    <dimen name="modal_dialog_vertical_spacer">12dp</dimen>
+    <dimen name="modal_dialog_corner_radius">8dp</dimen>
+    <dimen name="confirmation_dialog_text_height">216dp</dimen>
 </resources>
diff --git a/go/quickstep/res/values/strings.xml b/go/quickstep/res/values/strings.xml
index 61c8cd9..6e9e63e 100644
--- a/go/quickstep/res/values/strings.xml
+++ b/go/quickstep/res/values/strings.xml
@@ -11,4 +11,13 @@
     <string name="action_translate">Translate</string>
     <!-- Label for a button that triggers Search on a screenshot of the current app. [CHAR_LIMIT=40] -->
     <string name="action_search">Lens</string>
+    <!-- ******* NIU Actions First-Run Confirmation Dialog ******* -->
+    <!-- Dialog title -->
+    <string name="niu_actions_confirmation_title">Translate or listen to text on screen</string>
+    <!-- Dialog content -->
+    <string name="niu_actions_confirmation_text">Information such as text on your screen, web addresses, and screenshots may be shared with Google.\n\nTo change what information you share, go to <b>Settings > Apps > Default apps > Digital assistant app</b>.</string>
+    <!-- Label for a button that rejects the feature. [CHAR_LIMIT=40] -->
+    <string name="niu_actions_confirmation_no">CANCEL</string>
+    <!-- Label for a button that accepts the feature. [CHAR_LIMIT=40] -->
+    <string name="niu_actions_confirmation_yes">GOT IT</string>
 </resources>
diff --git a/go/quickstep/res/values/styles.xml b/go/quickstep/res/values/styles.xml
index ffe8f46..561531b 100644
--- a/go/quickstep/res/values/styles.xml
+++ b/go/quickstep/res/values/styles.xml
@@ -20,12 +20,14 @@
         <item name="overviewButtonTextColor">@color/go_overview_text_color</item>
         <item name="overviewButtonIconColor">@color/go_overview_text_color</item>
         <item name="overviewButtonBackgroundColor">@color/go_overview_button_color</item>
+        <item name="modalDialogBackground">@color/go_modal_dialog_background</item>
     </style>
 
     <style name="AppTheme.Dark" parent="@style/LauncherTheme.Dark">
         <item name="overviewButtonTextColor">@color/go_overview_text_color_dark</item>
         <item name="overviewButtonIconColor">@color/go_overview_text_color_dark</item>
         <item name="overviewButtonBackgroundColor">@color/go_overview_button_color_dark</item>
+        <item name="modalDialogBackground">@color/go_modal_dialog_background_dark</item>
     </style>
 
     <!-- Overview -->
@@ -56,4 +58,33 @@
         <item name="android:layout_height">wrap_content</item>
         <item name="android:orientation">vertical</item>
     </style>
+
+    <!-- Modal Dialogs -->
+    <style name="ModalDialogTitle">
+        <item name="android:fontFamily">sans-serif-medium</item>
+        <item name="android:textSize">20sp</item>
+        <item name="android:textColor">?android:attr/textColorPrimary</item>
+        <item name="android:lineHeight">24dp</item>
+        <item name="android:layout_width">match_parent</item>
+        <item name="android:layout_height">wrap_content</item>
+        <item name="android:layout_gravity">center_horizontal</item>
+    </style>
+
+    <style name="ModalDialogText">
+        <item name="android:fontFamily">sans-serif-medium</item>
+        <item name="android:textSize">16sp</item>
+        <item name="android:textColor">?android:attr/textColorSecondary</item>
+        <item name="android:lineHeight">24dp</item>
+        <item name="android:layout_width">match_parent</item>
+        <item name="android:layout_height">wrap_content</item>
+        <item name="android:layout_gravity">center_horizontal</item>
+    </style>
+
+    <style name="ModalDialogButton" parent="@style/Widget.AppCompat.Button.Borderless">
+        <item name="android:textSize">14sp</item>
+        <item name="android:textColor">?android:attr/colorAccent</item>
+        <item name="android:lineHeight">20dp</item>
+        <item name="android:layout_width">wrap_content</item>
+        <item name="android:layout_height">wrap_content</item>
+    </style>
 </resources>
\ No newline at end of file
diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
index b823c36..37f5352 100644
--- a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -20,20 +20,31 @@
 import static com.android.quickstep.views.OverviewActionsView.DISABLED_ROTATED;
 
 import android.annotation.SuppressLint;
+import android.app.AlertDialog;
 import android.app.assist.AssistContent;
 import android.content.ActivityNotFoundException;
 import android.content.ContentResolver;
 import android.content.Context;
 import android.content.Intent;
+import android.content.SharedPreferences;
+import android.graphics.Color;
 import android.graphics.Matrix;
+import android.graphics.drawable.ColorDrawable;
 import android.os.SystemClock;
 import android.os.UserManager;
 import android.provider.Settings;
 import android.text.TextUtils;
 import android.util.Log;
+import android.view.LayoutInflater;
+import android.view.View;
+import android.widget.Button;
 
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.BaseActivity;
+import com.android.launcher3.BaseDraggingActivity;
+import com.android.launcher3.R;
+import com.android.launcher3.Utilities;
 import com.android.quickstep.util.AssistContentRequester;
 import com.android.quickstep.views.OverviewActionsView;
 import com.android.quickstep.views.TaskThumbnailView;
@@ -53,6 +64,7 @@
     public static final String ACTIONS_ERROR_CODE = "niu_actions_app_error_code";
     public static final int ERROR_PERMISSIONS_STRUCTURE = 1;
     public static final int ERROR_PERMISSIONS_SCREENSHOT = 2;
+    private static final String NIU_ACTIONS_CONFIRMED = "launcher_go.niu_actions_confirmed";
     private static final String TAG = "TaskOverlayFactoryGo";
 
     private AssistContentRequester mContentRequester;
@@ -79,6 +91,9 @@
         private boolean mAssistStructurePermitted;
         private boolean mAssistScreenshotPermitted;
         private AssistContentRequester mFactoryContentRequester;
+        private SharedPreferences mSharedPreferences;
+        private String mPreviousAction;
+        private AlertDialog mConfirmationDialog;
 
         private TaskOverlayGo(TaskThumbnailView taskThumbnailView,
                 AssistContentRequester assistContentRequester) {
@@ -92,6 +107,12 @@
         @Override
         public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
                 boolean rotated) {
+            if (mConfirmationDialog != null && mConfirmationDialog.isShowing()) {
+                // Redraw the dialog in case the layout changed
+                mConfirmationDialog.dismiss();
+                showConfirmationDialog();
+            }
+
             getActionsView().updateDisabledFlags(DISABLED_NO_THUMBNAIL, thumbnail == null);
             checkSettings();
             if (thumbnail == null || TextUtils.isEmpty(mNIUPackageName)) {
@@ -105,6 +126,7 @@
             boolean isAllowedByPolicy = mThumbnailView.isRealSnapshot() && !isManagedProfileTask;
             getActionsView().setCallbacks(new OverlayUICallbacksGoImpl(isAllowedByPolicy, task));
             mTaskPackageName = task.key.getPackageName();
+            mSharedPreferences = Utilities.getPrefs(mApplicationContext);
 
             if (!mAssistStructurePermitted || !mAssistScreenshotPermitted) {
                 return;
@@ -131,6 +153,12 @@
          * Creates and sends an Intent corresponding to the button that was clicked
          */
         private void sendNIUIntent(String actionType) {
+            if (!mSharedPreferences.getBoolean(NIU_ACTIONS_CONFIRMED, false)) {
+                mPreviousAction = actionType;
+                showConfirmationDialog();
+                return;
+            }
+
             Intent intent = createNIUIntent(actionType);
             // Only add and send the image if the appropriate permissions are held
             if (mAssistStructurePermitted && mAssistScreenshotPermitted) {
@@ -218,6 +246,35 @@
         public void setImageActionsAPI(ImageActionsApi imageActionsApi) {
             mImageApi = imageActionsApi;
         }
+
+        private void showConfirmationDialog() {
+            BaseDraggingActivity activity = BaseActivity.fromContext(getActionsView().getContext());
+            LayoutInflater inflater = LayoutInflater.from(activity);
+            View view = inflater.inflate(R.layout.niu_actions_confirmation_dialog, /* root */ null);
+
+            Button acceptButton = view.findViewById(R.id.niu_actions_confirmation_accept);
+            acceptButton.setOnClickListener(this::onNiuActionsConfirmationAccept);
+
+            Button rejectButton = view.findViewById(R.id.niu_actions_confirmation_reject);
+            rejectButton.setOnClickListener(this::onNiuActionsConfirmationReject);
+
+            mConfirmationDialog = new AlertDialog.Builder(activity)
+                    .setView(view)
+                    .create();
+            mConfirmationDialog.getWindow()
+                    .setBackgroundDrawable(new ColorDrawable(Color.TRANSPARENT));
+            mConfirmationDialog.show();
+        }
+
+        private void onNiuActionsConfirmationAccept(View v) {
+            mConfirmationDialog.dismiss();
+            mSharedPreferences.edit().putBoolean(NIU_ACTIONS_CONFIRMED, true).apply();
+            sendNIUIntent(mPreviousAction);
+        }
+
+        private void onNiuActionsConfirmationReject(View v) {
+            mConfirmationDialog.cancel();
+        }
     }
 
     /**
diff --git a/protos/launcher_atom.proto b/protos/launcher_atom.proto
index 475b5be..6d49d75 100644
--- a/protos/launcher_atom.proto
+++ b/protos/launcher_atom.proto
@@ -141,6 +141,9 @@
   ALL_APPS_SEARCH_RESULT_NAVVYSITE = 25;
   ALL_APPS_SEARCH_RESULT_TIPS = 26;
   ALL_APPS_SEARCH_RESULT_PEOPLE_TILE = 27;
+
+  WIDGETS_BOTTOM_TRAY = 28;
+  WIDGETS_TRAY_PREDICTION = 29;
 }
 
 // Main app icons
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 22a8c9b..3d891e8 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+
 import android.app.prediction.AppTarget;
 import android.content.ComponentName;
 import android.text.TextUtils;
@@ -73,7 +75,9 @@
                     if (notAddedWidgets.size() > 0) {
                         // Even an apps have more than one widgets, we only include one widget.
                         fixedContainerItems.items.add(
-                                new PendingAddWidgetInfo(notAddedWidgets.get(0).widgetInfo));
+                                new PendingAddWidgetInfo(
+                                        notAddedWidgets.get(0).widgetInfo,
+                                        CONTAINER_WIDGETS_PREDICTION));
                     }
                 }
             }
@@ -90,7 +94,9 @@
                         new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
                 if (widgetItems.containsKey(targetWidget)) {
                     fixedContainerItems.items.add(
-                            new PendingAddWidgetInfo(widgetItems.get(targetWidget).widgetInfo));
+                            new PendingAddWidgetInfo(widgetItems.get(
+                                    targetWidget).widgetInfo,
+                                    CONTAINER_WIDGETS_PREDICTION));
                 }
             }
         }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 99ad8c0..e2e03ce 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -61,6 +61,7 @@
 import android.graphics.Matrix;
 import android.graphics.PointF;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.os.Build;
 import android.os.IBinder;
 import android.os.SystemClock;
@@ -1234,30 +1235,40 @@
         final RecentsOrientedState orientationState = mTaskViewSimulator.getOrientationState();
         final int windowRotation = orientationState.getDisplayRotation();
         final int homeRotation = orientationState.getRecentsActivityRotation();
+
+        final Matrix homeToWindowPositionMap = new Matrix();
+        final RectF startRect = updateProgressForStartRect(homeToWindowPositionMap, startProgress);
+        // Move the startRect to Launcher space as floatingIconView runs in Launcher
+        final Matrix windowToHomePositionMap = new Matrix();
+        homeToWindowPositionMap.invert(windowToHomePositionMap);
+        windowToHomePositionMap.mapRect(startRect);
+
         final Rect destinationBounds = SystemUiProxy.INSTANCE.get(mContext)
                 .startSwipePipToHome(taskInfo.topActivity,
                         TaskInfoCompat.getTopActivityInfo(taskInfo),
                         runningTaskTarget.taskInfo.pictureInPictureParams,
                         homeRotation,
                         mDp.hotseatBarSizePx);
-        final SwipePipToHomeAnimator swipePipToHomeAnimator = new SwipePipToHomeAnimator(
-                mContext,
-                runningTaskTarget.taskId,
-                taskInfo.topActivity,
-                runningTaskTarget.leash.getSurfaceControl(),
-                TaskInfoCompat.getPipSourceRectHint(
-                        runningTaskTarget.taskInfo.pictureInPictureParams),
-                TaskInfoCompat.getWindowConfigurationBounds(taskInfo),
-                updateProgressForStartRect(new Matrix(), startProgress),
-                destinationBounds,
-                mRecentsView.getPipCornerRadius(),
-                mRecentsView);
+        final SwipePipToHomeAnimator.Builder builder = new SwipePipToHomeAnimator.Builder()
+                .setContext(mContext)
+                .setTaskId(runningTaskTarget.taskId)
+                .setComponentName(taskInfo.topActivity)
+                .setLeash(runningTaskTarget.leash.getSurfaceControl())
+                .setSourceRectHint(TaskInfoCompat.getPipSourceRectHint(
+                        runningTaskTarget.taskInfo.pictureInPictureParams))
+                .setAppBounds(TaskInfoCompat.getWindowConfigurationBounds(taskInfo))
+                .setHomeToWindowPositionMap(homeToWindowPositionMap)
+                .setStartBounds(startRect)
+                .setDestinationBounds(destinationBounds)
+                .setCornerRadius(mRecentsView.getPipCornerRadius())
+                .setAttachedView(mRecentsView);
         // We would assume home and app window always in the same rotation While homeRotation
         // is not ROTATION_0 (which implies the rotation is turned on in launcher settings).
         if (homeRotation == ROTATION_0
                 && (windowRotation == ROTATION_90 || windowRotation == ROTATION_270)) {
-            swipePipToHomeAnimator.setFromRotation(mTaskViewSimulator, windowRotation);
+            builder.setFromRotation(mTaskViewSimulator, windowRotation);
         }
+        final SwipePipToHomeAnimator swipePipToHomeAnimator = builder.build();
         AnimatorPlaybackController activityAnimationToHome =
                 homeAnimFactory.createActivityAnimationToHome();
         swipePipToHomeAnimator.addAnimatorListener(new AnimatorListenerAdapter() {
@@ -1284,6 +1295,7 @@
                 mGestureState.setState(STATE_END_TARGET_ANIMATION_FINISHED);
             }
         });
+        setupWindowAnimation(swipePipToHomeAnimator);
         return swipePipToHomeAnimator;
     }
 
@@ -1314,6 +1326,11 @@
             HomeAnimationFactory homeAnimationFactory) {
         RectFSpringAnim anim =
                 super.createWindowAnimationToHome(startProgress, homeAnimationFactory);
+        setupWindowAnimation(anim);
+        return anim;
+    }
+
+    private void setupWindowAnimation(RectFSpringAnim anim) {
         anim.addOnUpdateListener((v, r, p) -> {
             updateSysUiFlags(Math.max(p, mCurrentShift.value));
         });
@@ -1331,7 +1348,6 @@
         if (mRecentsAnimationTargets != null) {
             mRecentsAnimationTargets.addReleaseCheck(anim);
         }
-        return anim;
     }
 
     public void onConsumerAboutToBeSwitched() {
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index bb5e493..ae2328e 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -95,7 +95,6 @@
     @UiThread
     public RecentsAnimationCallbacks startRecentsAnimation(GestureState gestureState,
             Intent intent, RecentsAnimationCallbacks.RecentsAnimationListener listener) {
-        Log.d("b/186444448", "startRecentsAnimation");
         // Notify if recents animation is still running
         if (mController != null) {
             String msg = "New recents animation started before old animation completed";
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index b1c9ed0..6575996 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -130,7 +130,8 @@
                 info.getAttribute().getNumber() /* origin */,
                 getCardinality(info) /* cardinality */,
                 info.getWidget().getSpanX(),
-                info.getWidget().getSpanY());
+                info.getWidget().getSpanY(),
+                getFeatures(info));
     }
 
     /**
@@ -365,15 +366,12 @@
                     atomInfo.getFolderIcon().getFromLabelState().getNumber() /* fromState */,
                     atomInfo.getFolderIcon().getToLabelState().getNumber() /* toState */,
                     atomInfo.getFolderIcon().getLabelInfo() /* edittext */,
-                    getCardinality(atomInfo) /* cardinality */);
+                    getCardinality(atomInfo) /* cardinality */,
+                    getFeatures(atomInfo) /* features */);
         }
     }
 
     private static int getCardinality(LauncherAtom.ItemInfo info) {
-        // TODO(b/187734511): Implement a unified solution for 1x1 widgets in folders/hotseat.
-        if (info.getItemCase().equals(LauncherAtom.ItemInfo.ItemCase.WIDGET)) {
-            return info.getWidget().getWidgetFeatures();
-        }
         switch (info.getContainerInfo().getContainerCase()) {
             case PREDICTED_HOTSEAT_CONTAINER:
                 return info.getContainerInfo().getPredictedHotseatContainer().getCardinality();
@@ -514,6 +512,13 @@
         }
     }
 
+    private static int getFeatures(LauncherAtom.ItemInfo info) {
+        if (info.getItemCase().equals(LauncherAtom.ItemInfo.ItemCase.WIDGET)) {
+            return info.getWidget().getWidgetFeatures();
+        }
+        return 0;
+    }
+
 
     /**
      * Interface to get stats log while it is dispatched to the system
diff --git a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
index 67a635b..7488649 100644
--- a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
+++ b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
@@ -56,7 +56,9 @@
     private final ComponentName mComponentName;
     private final SurfaceControl mLeash;
     private final Rect mAppBounds = new Rect();
+    private final Matrix mHomeToWindowPositionMap = new Matrix();
     private final Rect mStartBounds = new Rect();
+    private final RectF mCurrentBoundsF = new RectF();
     private final Rect mCurrentBounds = new Rect();
     private final Rect mDestinationBounds = new Rect();
     private final PipSurfaceTransactionHelper mSurfaceTransactionHelper;
@@ -66,10 +68,9 @@
     private final Rect mSourceHintRectInsets;
     private final Rect mSourceInsets = new Rect();
 
-    /** for rotation via {@link #setFromRotation(TaskViewSimulator, int)} */
-    private @RecentsOrientedState.SurfaceRotation int mFromRotation = Surface.ROTATION_0;
+    /** for rotation calculations */
+    private final @RecentsOrientedState.SurfaceRotation int mFromRotation;
     private final Rect mDestinationBoundsTransformed = new Rect();
-    private final Rect mDestinationBoundsAnimation = new Rect();
 
     /**
      * Flag to avoid the double-end problem since the leash would have been released
@@ -91,31 +92,39 @@
      * @param leash {@link SurfaceControl} this animator operates on
      * @param sourceRectHint See the definition in {@link android.app.PictureInPictureParams}
      * @param appBounds Bounds of the application, sourceRectHint is based on this bounds
+     * @param homeToWindowPositionMap {@link Matrix} to map a Rect from home to window space
      * @param startBounds Bounds of the application when this animator starts. This can be
      *                    different from the appBounds if user has swiped a certain distance and
      *                    Launcher has performed transform on the leash.
      * @param destinationBounds Bounds of the destination this animator ends to
+     * @param fromRotation From rotation if different from final rotation, ROTATION_0 otherwise
+     * @param destinationBoundsTransformed Destination bounds in window space
      * @param cornerRadius Corner radius in pixel value for PiP window
+     * @param view Attached view for logging purpose
      */
-    public SwipePipToHomeAnimator(@NonNull Context context,
+    private SwipePipToHomeAnimator(@NonNull Context context,
             int taskId,
             @NonNull ComponentName componentName,
             @NonNull SurfaceControl leash,
             @Nullable Rect sourceRectHint,
             @NonNull Rect appBounds,
+            @NonNull Matrix homeToWindowPositionMap,
             @NonNull RectF startBounds,
             @NonNull Rect destinationBounds,
+            @RecentsOrientedState.SurfaceRotation int fromRotation,
+            @NonNull Rect destinationBoundsTransformed,
             int cornerRadius,
             @NonNull View view) {
-        super(startBounds, new RectF(destinationBounds), context);
+        super(startBounds, new RectF(destinationBoundsTransformed), context);
         mTaskId = taskId;
         mComponentName = componentName;
         mLeash = leash;
         mAppBounds.set(appBounds);
+        mHomeToWindowPositionMap.set(homeToWindowPositionMap);
         startBounds.round(mStartBounds);
         mDestinationBounds.set(destinationBounds);
-        mDestinationBoundsTransformed.set(mDestinationBounds);
-        mDestinationBoundsAnimation.set(mDestinationBounds);
+        mFromRotation = fromRotation;
+        mDestinationBoundsTransformed.set(destinationBoundsTransformed);
         mSurfaceTransactionHelper = new PipSurfaceTransactionHelper(cornerRadius);
 
         if (sourceRectHint != null && (sourceRectHint.width() < destinationBounds.width()
@@ -191,37 +200,13 @@
         addOnUpdateListener(this::onAnimationUpdate);
     }
 
-    /** sets the from rotation if it's different from the target rotation. */
-    public void setFromRotation(TaskViewSimulator taskViewSimulator,
-            @RecentsOrientedState.SurfaceRotation int fromRotation) {
-        if (fromRotation != Surface.ROTATION_90 && fromRotation != Surface.ROTATION_270) {
-            Log.wtf(TAG, "Not a supported rotation, rotation=" + fromRotation);
-            return;
-        }
-        mFromRotation = fromRotation;
-        final Matrix matrix = new Matrix();
-        taskViewSimulator.applyWindowToHomeRotation(matrix);
-
-        // map the destination bounds into window space. mDestinationBounds is always calculated
-        // in the final home space and the animation runs in original window space.
-        final RectF transformed = new RectF(mDestinationBounds);
-        matrix.mapRect(transformed, new RectF(mDestinationBounds));
-        transformed.round(mDestinationBoundsTransformed);
-
-        // set the animation destination bounds for RectEvaluator calculation.
-        // bounds and insets are calculated as if the transition is from mAppBounds to
-        // mDestinationBoundsAnimation, separated from rotate / scale / position.
-        mDestinationBoundsAnimation.set(mAppBounds.left, mAppBounds.top,
-                mAppBounds.left + mDestinationBounds.width(),
-                mAppBounds.top + mDestinationBounds.height());
-    }
-
     private void onAnimationUpdate(@Nullable AppCloseConfig values, RectF currentRect,
             float progress) {
         if (mHasAnimationEnded) return;
         final SurfaceControl.Transaction tx =
                 PipSurfaceTransactionHelper.newSurfaceControlTransaction();
-        onAnimationUpdate(tx, currentRect, progress);
+        mHomeToWindowPositionMap.mapRect(mCurrentBoundsF, currentRect);
+        onAnimationUpdate(tx, mCurrentBoundsF, progress);
         tx.apply();
     }
 
@@ -309,6 +294,108 @@
         return new RotatedPosition(degree, positionX, positionY);
     }
 
+    /** Builder class for {@link SwipePipToHomeAnimator} */
+    public static class Builder {
+        private Context mContext;
+        private int mTaskId;
+        private ComponentName mComponentName;
+        private SurfaceControl mLeash;
+        private Rect mSourceRectHint;
+        private Rect mAppBounds;
+        private Matrix mHomeToWindowPositionMap;
+        private RectF mStartBounds;
+        private Rect mDestinationBounds;
+        private int mCornerRadius;
+        private View mAttachedView;
+        private @RecentsOrientedState.SurfaceRotation int mFromRotation = Surface.ROTATION_0;
+        private final Rect mDestinationBoundsTransformed = new Rect();
+
+        public Builder setContext(Context context) {
+            mContext = context;
+            return this;
+        }
+
+        public Builder setTaskId(int taskId) {
+            mTaskId = taskId;
+            return this;
+        }
+
+        public Builder setComponentName(ComponentName componentName) {
+            mComponentName = componentName;
+            return this;
+        }
+
+        public Builder setLeash(SurfaceControl leash) {
+            mLeash = leash;
+            return this;
+        }
+
+        public Builder setSourceRectHint(Rect sourceRectHint) {
+            mSourceRectHint = new Rect(sourceRectHint);
+            return this;
+        }
+
+        public Builder setAppBounds(Rect appBounds) {
+            mAppBounds = new Rect(appBounds);
+            return this;
+        }
+
+        public Builder setHomeToWindowPositionMap(Matrix homeToWindowPositionMap) {
+            mHomeToWindowPositionMap = new Matrix(homeToWindowPositionMap);
+            return this;
+        }
+
+        public Builder setStartBounds(RectF startBounds) {
+            mStartBounds = new RectF(startBounds);
+            return this;
+        }
+
+        public Builder setDestinationBounds(Rect destinationBounds) {
+            mDestinationBounds = new Rect(destinationBounds);
+            return this;
+        }
+
+        public Builder setCornerRadius(int cornerRadius) {
+            mCornerRadius = cornerRadius;
+            return this;
+        }
+
+        public Builder setAttachedView(View attachedView) {
+            mAttachedView = attachedView;
+            return this;
+        }
+
+        public Builder setFromRotation(TaskViewSimulator taskViewSimulator,
+                @RecentsOrientedState.SurfaceRotation int fromRotation) {
+            if (fromRotation != Surface.ROTATION_90 && fromRotation != Surface.ROTATION_270) {
+                Log.wtf(TAG, "Not a supported rotation, rotation=" + fromRotation);
+                return this;
+            }
+            final Matrix matrix = new Matrix();
+            taskViewSimulator.applyWindowToHomeRotation(matrix);
+
+            // map the destination bounds into window space. mDestinationBounds is always calculated
+            // in the final home space and the animation runs in original window space.
+            final RectF transformed = new RectF(mDestinationBounds);
+            matrix.mapRect(transformed, new RectF(mDestinationBounds));
+            transformed.round(mDestinationBoundsTransformed);
+
+            mFromRotation = fromRotation;
+            return this;
+        }
+
+        public SwipePipToHomeAnimator build() {
+            if (mDestinationBoundsTransformed.isEmpty()) {
+                mDestinationBoundsTransformed.set(mDestinationBounds);
+            }
+            return new SwipePipToHomeAnimator(mContext, mTaskId, mComponentName, mLeash,
+                    mSourceRectHint, mAppBounds,
+                    mHomeToWindowPositionMap, mStartBounds, mDestinationBounds,
+                    mFromRotation, mDestinationBoundsTransformed,
+                    mCornerRadius, mAttachedView);
+        }
+    }
+
     private static class RotatedPosition {
         private final float degree;
         private final float positionX;
diff --git a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
index 0e6ce87..7eee415 100644
--- a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
+++ b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
@@ -17,7 +17,6 @@
 
 import static com.android.launcher3.states.RotationHelper.deltaRotation;
 import static com.android.launcher3.touch.PagedOrientationHandler.MATRIX_POST_TRANSLATE;
-import static com.android.quickstep.util.NavigationModeFeatureFlag.LIVE_TILE;
 import static com.android.quickstep.util.RecentsOrientedState.postDisplayRotation;
 import static com.android.quickstep.util.RecentsOrientedState.preDisplayRotation;
 import static com.android.systemui.shared.system.WindowManagerWrapper.WINDOWING_MODE_FULLSCREEN;
@@ -62,7 +61,6 @@
     private final boolean mIsRecentsRtl;
 
     private final Rect mTaskRect = new Rect();
-    private boolean mDrawsBelowRecents;
     private final PointF mPivot = new PointF();
     private DeviceProfile mDp;
 
@@ -163,10 +161,6 @@
         recentsViewScroll.value = scroll;
     }
 
-    public void setDrawsBelowRecents(boolean drawsBelowRecents) {
-        mDrawsBelowRecents = drawsBelowRecents;
-    }
-
     /**
      * Adds animation for all the components corresponding to transition from an app to overview.
      */
@@ -301,12 +295,6 @@
         builder.withMatrix(mMatrix)
                 .withWindowCrop(mTmpCropRect)
                 .withCornerRadius(getCurrentCornerRadius());
-
-        if (LIVE_TILE.get() && params.getRecentsSurface() != null) {
-            // When relativeLayer = 0, it reverts the surfaces back to the original order.
-            builder.withRelativeLayerTo(params.getRecentsSurface(),
-                    mDrawsBelowRecents ? Integer.MIN_VALUE : 0);
-        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/views/FloatingWidgetView.java b/quickstep/src/com/android/quickstep/views/FloatingWidgetView.java
index 22ce942..88b11a0 100644
--- a/quickstep/src/com/android/quickstep/views/FloatingWidgetView.java
+++ b/quickstep/src/com/android/quickstep/views/FloatingWidgetView.java
@@ -152,6 +152,7 @@
             RectF widgetBackgroundPosition, Size windowSize, float windowCornerRadius,
             boolean appTargetIsTranslucent, int fallbackBackgroundColor) {
         mAppWidgetView = originalView;
+        // Deferrals must begin before GhostView is created. See b/190818220
         mAppWidgetView.beginDeferringUpdates();
         mBackgroundPosition = widgetBackgroundPosition;
         mAppTargetIsTranslucent = appTargetIsTranslucent;
@@ -240,6 +241,7 @@
         ((ViewGroup) dragLayer.getParent()).removeView(this);
         dragLayer.removeView(mListenerView);
         mBackgroundView.finish();
+        // Removing GhostView must occur before ending deferrals. See b/190818220
         mAppWidgetView.endDeferringUpdates();
         recycle();
         mLauncher.getViewCache().recycleView(R.layout.floating_widget_view, this);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index c7924e4..d59b459 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -82,7 +82,6 @@
 import android.text.TextPaint;
 import android.util.AttributeSet;
 import android.util.FloatProperty;
-import android.util.Log;
 import android.util.SparseBooleanArray;
 import android.view.Gravity;
 import android.view.HapticFeedbackConstants;
@@ -638,7 +637,6 @@
         mLiveTileTaskViewSimulator = new TaskViewSimulator(getContext(), getSizeStrategy());
         mLiveTileTaskViewSimulator.recentsViewScale.value = 1;
         mLiveTileTaskViewSimulator.setOrientationState(mOrientationState);
-        mLiveTileTaskViewSimulator.setDrawsBelowRecents(true);
 
         mTintingColor = getForegroundScrimDimColor(context);
     }
@@ -1206,6 +1204,8 @@
             mLiveTileTaskViewSimulator.taskSecondaryTranslation.value = 0;
             mLiveTileTaskViewSimulator.fullScreenProgress.value = 0;
             mLiveTileTaskViewSimulator.recentsViewScale.value = 1;
+
+            mLiveTileParams.setTargetAlpha(1);
         }
         if (mRunningTaskTileHidden) {
             setRunningTaskHidden(mRunningTaskTileHidden);
@@ -1591,7 +1591,6 @@
         }
         setEnableDrawingLiveTile(false);
         mLiveTileParams.setTargetSet(null);
-        mLiveTileTaskViewSimulator.setDrawsBelowRecents(true);
 
         // These are relatively expensive and don't need to be done this frame (RecentsView isn't
         // visible anyway), so defer by a frame to get off the critical path, e.g. app to home.
@@ -1669,7 +1668,6 @@
      * {@link #onGestureAnimationStart} and {@link #onGestureAnimationEnd()}.
      */
     public void onSwipeUpAnimationSuccess() {
-        Log.d("b/186444448", "onSwipeUpAnimationSuccess");
         animateUpTaskIconScale();
         setSwipeDownShouldLaunchApp(true);
     }
@@ -1683,7 +1681,6 @@
         pa.addListener(AnimatorListeners.forSuccessCallback(() -> {
             setLayoutRotation(newRotation, mOrientationState.getDisplayRotation());
             mActivity.getDragLayer().recreateControllers();
-            updateChildTaskOrientations();
             setRecentsChangedOrientation(false).start();
         }));
         pa.start();
@@ -1742,7 +1739,6 @@
      * Called when a gesture from an app has finished, and the animation to the target has ended.
      */
     public void onGestureAnimationEnd() {
-        Log.d("b/186444448", "onGestureEnd");
         mGestureActive = false;
         if (mOrientationState.setGestureActive(false)) {
             updateOrientationHandler();
@@ -1868,7 +1864,6 @@
 
     public void animateUpTaskIconScale() {
         mTaskIconScaledDown = false;
-        Log.d("b/186444448", "animateUpRunningTaskIconScale");
         int taskCount = getTaskViewCount();
         for (int i = 0; i < taskCount; i++) {
             TaskView taskView = getTaskViewAt(i);
@@ -2174,6 +2169,10 @@
         // Use setFloat instead of setViewAlpha as we want to keep the view visible even when it's
         // alpha is set to 0 so that it can be recycled in the view pool properly
         anim.setFloat(taskView, VIEW_ALPHA, 0, clampToProgress(ACCEL, 0, 0.5f));
+        if (LIVE_TILE.get() && taskView.isRunningTask()) {
+            anim.setFloat(mLiveTileParams, TransformParams.TARGET_ALPHA, 0,
+                    clampToProgress(ACCEL, 0, 0.5f));
+        }
         SplitSelectStateController splitController = mSplitPlaceholderView.getSplitController();
 
         ResourceProvider rp = DynamicResource.provider(mActivity);
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index c97225e..cb51d66 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -16,7 +16,6 @@
 
 package com.android.quickstep.views;
 
-import static com.android.quickstep.util.NavigationModeFeatureFlag.LIVE_TILE;
 import static com.android.quickstep.views.TaskThumbnailView.DIM_ALPHA;
 
 import android.animation.Animator;
@@ -59,6 +58,7 @@
 
     private static final int REVEAL_OPEN_DURATION = 150;
     private static final int REVEAL_CLOSE_DURATION = 100;
+    private final float mTaskInsetMargin;
 
     private BaseDraggingActivity mActivity;
     private TextView mTaskName;
@@ -75,6 +75,7 @@
 
         mActivity = BaseDraggingActivity.fromContext(context);
         setClipToOutline(true);
+        mTaskInsetMargin = getResources().getDimension(R.dimen.task_card_margin);
     }
 
     @Override
@@ -124,8 +125,13 @@
 
     private void setPosition(float x, float y, int overscrollShift) {
         PagedOrientationHandler pagedOrientationHandler = mTaskView.getPagedOrientationHandler();
+        // Inset due to margin
+        PointF additionalInset = pagedOrientationHandler
+                .getAdditionalInsetForTaskMenu(mTaskInsetMargin);
         int taskTopMargin = mActivity.getDeviceProfile().overviewTaskThumbnailTopMarginPx;
-        float adjustedY = y + taskTopMargin;
+
+        float adjustedY = y + taskTopMargin - additionalInset.y;
+        float adjustedX = x - additionalInset.x;
         // Changing pivot to make computations easier
         // NOTE: Changing the pivots means the rotated view gets rotated about the new pivots set,
         // which would render the X and Y position set here incorrect
@@ -137,7 +143,8 @@
             setPivotY(0);
         }
         setRotation(pagedOrientationHandler.getDegreesRotated());
-        setX(pagedOrientationHandler.getTaskMenuX(x, mTaskView.getThumbnail(), overscrollShift));
+        setX(pagedOrientationHandler.getTaskMenuX(adjustedX,
+                mTaskView.getThumbnail(), overscrollShift));
         setY(pagedOrientationHandler.getTaskMenuY(
                 adjustedY, mTaskView.getThumbnail(), overscrollShift));
     }
@@ -211,25 +218,14 @@
                 menuOptionView, mActivity.getDeviceProfile());
         menuOptionView.setEnabled(menuOption.isEnabled());
         menuOptionView.setAlpha(menuOption.isEnabled() ? 1 : 0.5f);
-        menuOptionView.setOnClickListener(view -> {
-            if (LIVE_TILE.get()) {
-                RecentsView recentsView = mTaskView.getRecentsView();
-                recentsView.switchToScreenshot(null,
-                        () -> recentsView.finishRecentsAnimation(true /* toRecents */,
-                                false /* shouldPip */,
-                                () -> menuOption.onClick(view)));
-            } else {
-                menuOption.onClick(view);
-            }
-        });
+        menuOptionView.setOnClickListener(menuOption::onClick);
         mOptionLayout.addView(menuOptionView);
     }
 
     private void orientAroundTaskView(TaskView taskView) {
         PagedOrientationHandler orientationHandler = taskView.getPagedOrientationHandler();
         measure(MeasureSpec.UNSPECIFIED, MeasureSpec.UNSPECIFIED);
-        float taskInsetMargin = getResources().getDimension(R.dimen.task_card_margin);
-        orientationHandler.setTaskMenuAroundTaskView(this, taskInsetMargin);
+        orientationHandler.setTaskMenuAroundTaskView(this, mTaskInsetMargin);
         mActivity.getDragLayer().getDescendantRectRelativeToSelf(taskView, sTempRect);
         Rect insets = mActivity.getDragLayer().getInsets();
         BaseDragLayer.LayoutParams params = (BaseDragLayer.LayoutParams) getLayoutParams();
@@ -243,9 +239,6 @@
         setScaleY(taskView.getScaleY());
         orientationHandler.setTaskOptionsMenuLayoutOrientation(
                 mActivity.getDeviceProfile(), mOptionLayout);
-        PointF additionalInset = orientationHandler.getAdditionalInsetForTaskMenu(taskInsetMargin);
-        insets.left += additionalInset.x;
-        insets.top += additionalInset.y;
         setPosition(sTempRect.left - insets.left, sTempRect.top - insets.top, 0);
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
index df195d7..ee6b94b 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailView.java
@@ -84,7 +84,6 @@
     private TaskOverlay mOverlay;
     private final Paint mPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
     private final Paint mBackgroundPaint = new Paint(Paint.ANTI_ALIAS_FLAG);
-    private final Paint mClearPaint = new Paint();
     private final Paint mDimmingPaintAfterClearing = new Paint();
     private final int mDimColor;
 
@@ -115,7 +114,6 @@
         super(context, attrs, defStyleAttr);
         mPaint.setFilterBitmap(true);
         mBackgroundPaint.setColor(Color.WHITE);
-        mClearPaint.setXfermode(new PorterDuffXfermode(PorterDuff.Mode.CLEAR));
         mActivity = BaseActivity.fromContext(context);
         // Initialize with placeholder value. It is overridden later by TaskView
         mFullscreenParams = TEMP_PARAMS.get(context);
@@ -308,7 +306,6 @@
             float cornerRadius) {
         if (LIVE_TILE.get()) {
             if (mTask != null && getTaskView().isRunningTask() && !getTaskView().showScreenshot()) {
-                canvas.drawRoundRect(x, y, width, height, cornerRadius, cornerRadius, mClearPaint);
                 canvas.drawRoundRect(x, y, width, height, cornerRadius, cornerRadius,
                         mDimmingPaintAfterClearing);
                 return;
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 57cc79c..2639f15 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -561,13 +561,7 @@
                     recentsView.getDepthController());
             anim.addListener(new AnimatorListenerAdapter() {
                 @Override
-                public void onAnimationStart(Animator animator) {
-                    recentsView.getLiveTileTaskViewSimulator().setDrawsBelowRecents(false);
-                }
-
-                @Override
                 public void onAnimationEnd(Animator animator) {
-                    recentsView.getLiveTileTaskViewSimulator().setDrawsBelowRecents(true);
                     mIsClickableAsLiveTile = true;
                 }
             });
@@ -731,7 +725,16 @@
     private void setIcon(Drawable icon) {
         if (icon != null) {
             mIconView.setDrawable(icon);
-            mIconView.setOnClickListener(v -> showTaskMenu());
+            mIconView.setOnClickListener(v -> {
+                if (LIVE_TILE.get() && isRunningTask()) {
+                    RecentsView recentsView = getRecentsView();
+                    recentsView.switchToScreenshot(
+                            () -> recentsView.finishRecentsAnimation(true /* toRecents */,
+                                    this::showTaskMenu));
+                } else {
+                    showTaskMenu();
+                }
+            });
             mIconView.setOnLongClickListener(v -> {
                 requestDisallowInterceptTouchEvent(true);
                 return showTaskMenu();
@@ -812,8 +815,6 @@
     }
 
     public void animateIconScaleAndDimIntoView() {
-        Log.d("b/186444448", "animateIconScaleAndDimIntoView: startProgress="
-                + mIconScaleAnimStartProgress);
         if (mIconAndDimAnimator != null) {
             mIconAndDimAnimator.cancel();
         }
@@ -823,7 +824,6 @@
         mIconAndDimAnimator.addListener(new AnimatorListenerAdapter() {
             @Override
             public void onAnimationEnd(Animator animation) {
-                Log.d("b/186444448", "animateIconScaleAndDimIntoView: end");
                 mIconAndDimAnimator = null;
             }
         });
diff --git a/res/drawable/work_card.xml b/res/drawable/work_card.xml
index 0e4b054..7048955 100644
--- a/res/drawable/work_card.xml
+++ b/res/drawable/work_card.xml
@@ -20,8 +20,5 @@
     android:shape="rectangle">
     <solid android:color="?androidprv:attr/colorSurface" />
     <corners android:radius="@dimen/work_edu_card_margin" />
-    <padding
-        android:left="@dimen/work_fab_radius"
-        android:right="@dimen/work_fab_radius" />
 </shape>
 
diff --git a/res/layout/work_apps_edu.xml b/res/layout/work_apps_edu.xml
index 919f1b2..84fdfdf 100644
--- a/res/layout/work_apps_edu.xml
+++ b/res/layout/work_apps_edu.xml
@@ -22,7 +22,8 @@
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:orientation="vertical"
-        android:padding="@dimen/work_edu_card_margin"
+        android:paddingHorizontal="@dimen/work_card_padding_horizontal"
+        android:paddingVertical="@dimen/work_card_padding_vertical"
         android:background="@drawable/work_card"
         android:layout_gravity="center_horizontal"
         android:gravity="center"
@@ -34,8 +35,7 @@
             android:id="@+id/work_apps_paused_title"
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
-            android:layout_marginTop="8dp"
-            android:layout_marginBottom="8dp"
+            android:layout_marginBottom="@dimen/work_card_padding_vertical"
             android:text="@string/work_profile_edu_work_apps"
             android:textAlignment="center"
             android:textSize="20sp" />
diff --git a/res/layout/work_apps_paused.xml b/res/layout/work_apps_paused.xml
index 02a50ca..841734c 100644
--- a/res/layout/work_apps_paused.xml
+++ b/res/layout/work_apps_paused.xml
@@ -28,18 +28,18 @@
         android:layout_marginTop="40dp"
         android:text="@string/work_apps_paused_title"
         android:textAlignment="center"
-        android:textSize="20sp" />
+        android:textSize="22sp" />
 
     <TextView
         android:layout_width="wrap_content"
         android:layout_height="wrap_content"
         android:id="@+id/work_apps_paused_content"
-        android:textColor="?attr/workProfileOverlayTextColor"
+        android:textColor="?android:attr/textColorSecondary"
         android:text="@string/work_apps_paused_body"
         android:textAlignment="center"
         android:layout_marginTop="16dp"
         android:layout_marginBottom="24dp"
-        android:textSize="16sp" />
+        android:textSize="14sp" />
 
     <Button
         android:layout_width="wrap_content"
@@ -49,5 +49,7 @@
         android:text="@string/work_apps_enable_btn_text"
         android:textAlignment="center"
         android:background="@drawable/rounded_action_button"
+        android:paddingStart="8dp"
+        android:paddingEnd="8dp"
         android:textSize="14sp" />
 </com.android.launcher3.allapps.WorkPausedCard>
\ No newline at end of file
diff --git a/res/layout/work_mode_fab.xml b/res/layout/work_mode_fab.xml
index 7183817..04faa15 100644
--- a/res/layout/work_mode_fab.xml
+++ b/res/layout/work_mode_fab.xml
@@ -25,7 +25,7 @@
     android:textColor="@color/all_apps_tab_text"
     android:textSize="14sp"
     android:background="@drawable/work_apps_toggle_background"
-    android:drawablePadding="16dp"
+    android:drawablePadding="8dp"
     android:drawableStart="@drawable/ic_corp_off"
     android:layout_marginBottom="@dimen/work_fab_margin"
     android:layout_marginEnd="@dimen/work_fab_margin"
diff --git a/res/raw/downgrade_schema.json b/res/raw/downgrade_schema.json
index 8f1780e..bc25cec 100644
--- a/res/raw/downgrade_schema.json
+++ b/res/raw/downgrade_schema.json
@@ -2,8 +2,14 @@
   // Note: Comments are not supported in JSON schema, but android parser is lenient.
 
   // Maximum DB version supported by this schema
-  "version" : 28,
+  "version" : 29,
 
+  "downgrade_to_28" : [
+    "ALTER TABLE favorites RENAME TO temp_favorites;",
+    "CREATE TABLE favorites(_id INTEGER PRIMARY KEY, title TEXT, intent TEXT, container INTEGER, screen INTEGER, cellX INTEGER, cellY INTEGER, spanX INTEGER, spanY INTEGER, itemType INTEGER, appWidgetId INTEGER NOT NULL DEFAULT - 1, iconPackage TEXT, iconResource TEXT, icon BLOB, appWidgetProvider TEXT, modified INTEGER NOT NULL DEFAULT 0, restored INTEGER NOT NULL DEFAULT 0, profileId INTEGER DEFAULT 0, rank INTEGER NOT NULL DEFAULT 0, options INTEGER NOT NULL DEFAULT 0);",
+    "INSERT INTO favorites SELECT _id, title, intent, container, screen, cellX, cellY, spanX, spanY, itemType, appWidgetId, iconPackage, iconResource, icon, appWidgetProvider, modified, restored, profileId, rank, options FROM temp_favorites;",
+    "DROP TABLE temp_favorites;"
+  ],
   "downgrade_to_27" : [
     "CREATE TABLE workspaceScreens (_id INTEGER PRIMARY KEY,screenRank INTEGER,modified INTEGER NOT NULL DEFAULT 0)",
     "insert into workspaceScreens (_id, screenRank) select screen as _id, screen as screenRank from favorites where container = -100 group by screen order by screen"
diff --git a/res/values-sw340dp/dimens.xml b/res/values-sw340dp/dimens.xml
index c9f2981..33b06f5 100644
--- a/res/values-sw340dp/dimens.xml
+++ b/res/values-sw340dp/dimens.xml
@@ -19,6 +19,5 @@
 
     <!-- Drag padding to add to the bottom of drop targets -->
     <dimen name="drop_target_drag_padding">20dp</dimen>
-    <dimen name="drop_target_text_size">16sp</dimen>
 
 </resources>
\ No newline at end of file
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index 9891ff5..a891e65 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -121,6 +121,8 @@
 <!-- Floating action button inside work tab to toggle work profile -->
     <dimen name="work_fab_height">48dp</dimen>
     <dimen name="work_fab_radius">24dp</dimen>
+    <dimen name="work_card_padding_horizontal">24dp</dimen>
+    <dimen name="work_card_padding_vertical">32dp</dimen>
     <dimen name="work_fab_margin">16dp</dimen>
     <dimen name="work_profile_footer_padding">20dp</dimen>
     <dimen name="work_profile_footer_text_size">16sp</dimen>
@@ -189,8 +191,9 @@
 <!-- Dragging -->
     <!-- Drag padding to add to the bottom of drop targets -->
     <dimen name="drop_target_drag_padding">14dp</dimen>
-    <dimen name="drop_target_text_size">20sp</dimen>
+    <dimen name="drop_target_text_size">16sp</dimen>
     <dimen name="drop_target_shadow_elevation">2dp</dimen>
+    <dimen name="drop_target_bar_margin_horizontal">4dp</dimen>
 
     <!-- the distance an icon must be dragged before button drop targets accept it -->
     <dimen name="drag_distanceThreshold">30dp</dimen>
diff --git a/res/values/styles.xml b/res/values/styles.xml
index 09c8b9b..a926cb5 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -258,7 +258,7 @@
     <!-- Drop targets -->
     <style name="DropTargetButtonBase" parent="@android:style/TextAppearance.DeviceDefault">
         <item name="android:drawablePadding">8dp</item>
-        <item name="android:padding">16dp</item>
+        <item name="android:padding">14dp</item>
         <item name="android:textColor">@color/drop_target_text</item>
         <item name="android:textSize">@dimen/drop_target_text_size</item>
         <item name="android:singleLine">true</item>
diff --git a/src/com/android/launcher3/ButtonDropTarget.java b/src/com/android/launcher3/ButtonDropTarget.java
index 0a674b5..0d33b6f 100644
--- a/src/com/android/launcher3/ButtonDropTarget.java
+++ b/src/com/android/launcher3/ButtonDropTarget.java
@@ -20,15 +20,12 @@
 
 import static com.android.launcher3.LauncherState.NORMAL;
 
-import android.animation.AnimatorSet;
 import android.content.Context;
-import android.content.res.ColorStateList;
 import android.content.res.Resources;
 import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
 import android.text.TextUtils;
 import android.util.AttributeSet;
-import android.util.Property;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.View.OnClickListener;
@@ -40,6 +37,7 @@
 import com.android.launcher3.dragndrop.DragController;
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.dragndrop.DragOptions;
+import com.android.launcher3.dragndrop.DragView;
 import com.android.launcher3.model.data.ItemInfo;
 
 /**
@@ -212,15 +210,20 @@
             return;
         }
         final DragLayer dragLayer = mLauncher.getDragLayer();
+        final DragView dragView = d.dragView;
         final Rect to = getIconRect(d);
-        final float scale = (float) to.width() / d.dragView.getMeasuredWidth();
-        d.dragView.detachContentView(/* reattachToPreviousParent= */ true);
+        final float scale = (float) to.width() / dragView.getMeasuredWidth();
+        dragView.disableColorExtraction();
+        dragView.detachContentView(/* reattachToPreviousParent= */ true);
         mDropTargetBar.deferOnDragEnd();
 
         Runnable onAnimationEndRunnable = () -> {
             completeDrop(d);
             mDropTargetBar.onDragEnd();
             mLauncher.getStateManager().goToState(NORMAL);
+            // Only re-enable updates once the workspace is back to normal, which will be after the
+            // current frame.
+            post(dragView::resumeColorExtraction);
         };
 
         dragLayer.animateView(d.dragView, to, scale, 0.1f, 0.1f,
diff --git a/src/com/android/launcher3/DropTargetBar.java b/src/com/android/launcher3/DropTargetBar.java
index 4a1b084..88f6c49 100644
--- a/src/com/android/launcher3/DropTargetBar.java
+++ b/src/com/android/launcher3/DropTargetBar.java
@@ -105,7 +105,8 @@
                         / (2 * (grid.inv.numColumns + 1)))
                         + grid.edgeMarginPx;
             } else {
-                gap = grid.desiredWorkspaceLeftRightMarginPx - grid.inv.defaultWidgetPadding.right;
+                gap = getContext().getResources()
+                        .getDimensionPixelSize(R.dimen.drop_target_bar_margin_horizontal);
             }
             lp.width = grid.availableWidthPx - 2 * gap;
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index bd20f1e..89b44a3 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1342,6 +1342,12 @@
         launcherInfo.minSpanX = itemInfo.minSpanX;
         launcherInfo.minSpanY = itemInfo.minSpanY;
         launcherInfo.user = appWidgetInfo.getProfile();
+        if (itemInfo instanceof PendingAddWidgetInfo) {
+            launcherInfo.sourceContainer = ((PendingAddWidgetInfo) itemInfo).sourceContainer;
+        } else if (itemInfo instanceof PendingRequestArgs) {
+            launcherInfo.sourceContainer =
+                    ((PendingRequestArgs) itemInfo).getWidgetSourceContainer();
+        }
 
         getModelWriter().addItemToDatabase(launcherInfo,
                 itemInfo.container, itemInfo.screenId, itemInfo.cellX, itemInfo.cellY);
@@ -2435,7 +2441,8 @@
 
                         // Also try to bind the widget. If the bind fails, the user will be shown
                         // a click to setup UI, which will ask for the bind permission.
-                        PendingAddWidgetInfo pendingInfo = new PendingAddWidgetInfo(appWidgetInfo);
+                        PendingAddWidgetInfo pendingInfo =
+                                new PendingAddWidgetInfo(appWidgetInfo, item.sourceContainer);
                         pendingInfo.spanX = item.spanX;
                         pendingInfo.spanY = item.spanY;
                         pendingInfo.minSpanX = item.minSpanX;
diff --git a/src/com/android/launcher3/LauncherProvider.java b/src/com/android/launcher3/LauncherProvider.java
index 39b16fd..440e9e3 100644
--- a/src/com/android/launcher3/LauncherProvider.java
+++ b/src/com/android/launcher3/LauncherProvider.java
@@ -61,6 +61,7 @@
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.DbDowngradeHelper;
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.LauncherDbUtils;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
@@ -98,7 +99,7 @@
      * Represents the schema of the database. Changes in scheme need not be backwards compatible.
      * When increasing the scheme version, ensure that downgrade_schema.json is updated
      */
-    public static final int SCHEMA_VERSION = 28;
+    public static final int SCHEMA_VERSION = 29;
 
     public static final String AUTHORITY = BuildConfig.APPLICATION_ID + ".settings";
     public static final String KEY_LAYOUT_PROVIDER_AUTHORITY = "KEY_LAYOUT_PROVIDER_AUTHORITY";
@@ -879,9 +880,18 @@
                     }
                     dropTable(db, "workspaceScreens");
                 }
-                case 28:
+                case 28: {
+                    boolean columnAdded = addIntegerColumn(
+                            db, Favorites.APPWIDGET_SOURCE, Favorites.CONTAINER_UNKNOWN);
+                    if (!columnAdded) {
+                        // Old version remains, which means we wipe old data
+                        break;
+                    }
+                }
+                case 29: {
                     // DB Upgraded successfully
                     return;
+                }
             }
 
             // DB was not upgraded
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 22c257a..d663480 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -197,6 +197,8 @@
         public static final int CONTAINER_HOTSEAT_PREDICTION = -103;
         public static final int CONTAINER_ALL_APPS = -104;
         public static final int CONTAINER_WIDGETS_TRAY = -105;
+        public static final int CONTAINER_BOTTOM_WIDGETS_TRAY = -112;
+        public static final int CONTAINER_PIN_WIDGETS = -113;
         // Represents search results view.
         public static final int CONTAINER_SEARCH_RESULTS = -106;
         public static final int CONTAINER_SHORTCUTS = -107;
@@ -207,6 +209,8 @@
         // Represents any of the extended containers implemented in non-AOSP variants.
         public static final int EXTENDED_CONTAINERS = -200;
 
+        public static final int CONTAINER_UNKNOWN = -1;
+
         public static final String containerToString(int container) {
             switch (container) {
                 case CONTAINER_DESKTOP: return "desktop";
@@ -306,6 +310,12 @@
          */
         public static final String OPTIONS = "options";
 
+        /**
+         * Stores the source container that the widget was added from.
+         * <p>Type: INTEGER</p>
+         */
+        public static final String APPWIDGET_SOURCE = "appWidgetSource";
+
         public static void addTableToDb(SQLiteDatabase db, long myProfileId, boolean optional) {
             addTableToDb(db, myProfileId, optional, TABLE_NAME);
         }
@@ -333,7 +343,8 @@
                     "restored INTEGER NOT NULL DEFAULT 0," +
                     "profileId INTEGER DEFAULT " + myProfileId + "," +
                     "rank INTEGER NOT NULL DEFAULT 0," +
-                    "options INTEGER NOT NULL DEFAULT 0" +
+                    "options INTEGER NOT NULL DEFAULT 0," +
+                    APPWIDGET_SOURCE + " INTEGER NOT NULL DEFAULT " + CONTAINER_UNKNOWN +
                     ");");
         }
     }
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index f412f92..5689394 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -1833,8 +1833,11 @@
             onDropExternal(touchXY, dropTargetLayout, d);
         } else {
             final View cell = mDragInfo.cell;
+            final DragView dragView = d.dragView;
             boolean droppedOnOriginalCellDuringTransition = false;
-            Runnable onCompleteRunnable = null;
+            Runnable onCompleteRunnable = dragView::resumeColorExtraction;
+
+            dragView.disableColorExtraction();
 
             if (dropTargetLayout != null && !d.cancelled) {
                 // Move internally
@@ -1945,7 +1948,9 @@
                         AppWidgetProviderInfo pInfo = hostView.getAppWidgetInfo();
                         if (pInfo != null && pInfo.resizeMode != AppWidgetProviderInfo.RESIZE_NONE
                                 && !options.isAccessibleDrag) {
+                            final Runnable previousRunnable = onCompleteRunnable;
                             onCompleteRunnable = () -> {
+                                previousRunnable.run();
                                 if (!isPageInTransition()) {
                                     AppWidgetResizeFrame.showForWidget(hostView, cellLayout);
                                 }
@@ -2014,7 +2019,7 @@
             parent.onDropChild(cell);
 
             mLauncher.getStateManager().goToState(NORMAL, SPRING_LOADED_EXIT_DELAY,
-                    onCompleteRunnable == null ? null : forSuccessCallback(onCompleteRunnable));
+                    forSuccessCallback(onCompleteRunnable));
             mStatsLogManager.logger().withItemInfo(d.dragInfo).withInstanceId(d.logInstanceId)
                     .log(LauncherEvent.LAUNCHER_ITEM_DROP_COMPLETED);
         }
diff --git a/src/com/android/launcher3/dragndrop/AddItemActivity.java b/src/com/android/launcher3/dragndrop/AddItemActivity.java
index 1503167..df97bfb 100644
--- a/src/com/android/launcher3/dragndrop/AddItemActivity.java
+++ b/src/com/android/launcher3/dragndrop/AddItemActivity.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.dragndrop;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PIN_WIDGETS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ADD_EXTERNAL_ITEM_BACK;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ADD_EXTERNAL_ITEM_CANCELLED;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ADD_EXTERNAL_ITEM_DRAGGED;
@@ -249,7 +250,8 @@
         mAppWidgetManager = new WidgetManagerHelper(this);
         mAppWidgetHost = new LauncherAppWidgetHost(this);
 
-        PendingAddWidgetInfo pendingInfo = new PendingAddWidgetInfo(widgetInfo);
+        PendingAddWidgetInfo pendingInfo =
+                new PendingAddWidgetInfo(widgetInfo, CONTAINER_PIN_WIDGETS);
         pendingInfo.spanX = Math.min(mIdp.numColumns, widgetInfo.spanX);
         pendingInfo.spanY = Math.min(mIdp.numRows, widgetInfo.spanY);
         mWidgetOptions = pendingInfo.getDefaultSizeOptions(this);
diff --git a/src/com/android/launcher3/dragndrop/DragView.java b/src/com/android/launcher3/dragndrop/DragView.java
index fb1a6be..1f93730 100644
--- a/src/com/android/launcher3/dragndrop/DragView.java
+++ b/src/com/android/launcher3/dragndrop/DragView.java
@@ -63,6 +63,7 @@
 import com.android.launcher3.util.RunnableList;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.BaseDragLayer;
+import com.android.launcher3.widget.LauncherAppWidgetHostView;
 
 /** A custom view for rendering an icon, folder, shortcut or widget during drag-n-drop. */
 public abstract class DragView<T extends Context & ActivityContext> extends FrameLayout {
@@ -447,12 +448,12 @@
             mContent.draw(picture.beginRecording(mWidth, mHeight));
             picture.endRecording();
             View view = new View(mActivity);
-            view.setClipToOutline(mContent.getClipToOutline());
-            view.setOutlineProvider(mContent.getOutlineProvider());
             view.setBackground(new PictureDrawable(picture));
             view.measure(makeMeasureSpec(mWidth, EXACTLY), makeMeasureSpec(mHeight, EXACTLY));
             view.layout(mContent.getLeft(), mContent.getTop(),
                     mContent.getRight(), mContent.getBottom());
+            setClipToOutline(mContent.getClipToOutline());
+            setOutlineProvider(mContent.getOutlineProvider());
             addViewInLayout(view, indexOfChild(mContent), mContent.getLayoutParams(), true);
 
             removeViewInLayout(mContent);
@@ -467,6 +468,24 @@
     }
 
     /**
+     * If the drag view uses color extraction, block it.
+     */
+    public void disableColorExtraction() {
+        if (mContent instanceof LauncherAppWidgetHostView) {
+            ((LauncherAppWidgetHostView) mContent).disableColorExtraction();
+        }
+    }
+
+    /**
+     * If the drag view uses color extraction, restores it.
+     */
+    public void resumeColorExtraction() {
+        if (mContent instanceof LauncherAppWidgetHostView) {
+            ((LauncherAppWidgetHostView) mContent).enableColorExtraction(/* updateColors= */ false);
+        }
+    }
+
+    /**
      * Removes this view from the {@link DragLayer}.
      *
      * <p>If the drag content is a {@link #mContent}, this call doesn't reattach the
diff --git a/src/com/android/launcher3/dragndrop/PinItemDragListener.java b/src/com/android/launcher3/dragndrop/PinItemDragListener.java
index 9f12e6e..2bdf8a0 100644
--- a/src/com/android/launcher3/dragndrop/PinItemDragListener.java
+++ b/src/com/android/launcher3/dragndrop/PinItemDragListener.java
@@ -17,6 +17,8 @@
 package com.android.launcher3.dragndrop;
 
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PIN_WIDGETS;
+
 import android.annotation.TargetApi;
 import android.appwidget.AppWidgetManager;
 import android.content.pm.LauncherApps.PinItemRequest;
@@ -84,7 +86,7 @@
                             mLauncher, mRequest.getAppWidgetProviderInfo(mLauncher));
             final PinWidgetFlowHandler flowHandler =
                     new PinWidgetFlowHandler(providerInfo, mRequest);
-            item = new PendingAddWidgetInfo(providerInfo) {
+            item = new PendingAddWidgetInfo(providerInfo, CONTAINER_PIN_WIDGETS) {
                 @Override
                 public WidgetAddFlowHandler getHandler() {
                     return flowHandler;
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index ff18297..d5b5452 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -381,6 +381,8 @@
                         LauncherSettings.Favorites.RANK);
                 final int optionsIndex = c.getColumnIndexOrThrow(
                         LauncherSettings.Favorites.OPTIONS);
+                final int sourceContainerIndex = c.getColumnIndexOrThrow(
+                        LauncherSettings.Favorites.APPWIDGET_SOURCE);
 
                 final LongSparseArray<Boolean> unlockedUsers = new LongSparseArray<>();
 
@@ -748,6 +750,7 @@
                                 appWidgetInfo.spanY = c.getInt(spanYIndex);
                                 appWidgetInfo.options = c.getInt(optionsIndex);
                                 appWidgetInfo.user = c.user;
+                                appWidgetInfo.sourceContainer = c.getInt(sourceContainerIndex);
 
                                 if (appWidgetInfo.spanX <= 0 || appWidgetInfo.spanY <= 0) {
                                     c.markDeleted("Widget has invalid size: "
diff --git a/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java b/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java
index 658c6e1..0283d5f 100644
--- a/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java
+++ b/src/com/android/launcher3/model/data/LauncherAppWidgetInfo.java
@@ -16,6 +16,11 @@
 
 package com.android.launcher3.model.data;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_ALL_APPS;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_BOTTOM_WIDGETS_TRAY;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PIN_WIDGETS;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_TRAY;
 import static com.android.launcher3.Utilities.ATLEAST_S;
 
 import android.appwidget.AppWidgetHostView;
@@ -139,6 +144,11 @@
 
     private boolean mHasNotifiedInitialWidgetSizeChanged;
 
+    /**
+     * The container from which this widget was added (e.g. widgets tray, pin widget, search)
+     */
+    public int sourceContainer = LauncherSettings.Favorites.CONTAINER_UNKNOWN;
+
     public LauncherAppWidgetInfo(int appWidgetId, ComponentName providerName) {
         this.appWidgetId = appWidgetId;
         this.providerName = providerName;
@@ -187,7 +197,8 @@
                 .put(LauncherSettings.Favorites.APPWIDGET_PROVIDER, providerName.flattenToString())
                 .put(LauncherSettings.Favorites.RESTORED, restoreStatus)
                 .put(LauncherSettings.Favorites.OPTIONS, options)
-                .put(LauncherSettings.Favorites.INTENT, bindOptions);
+                .put(LauncherSettings.Favorites.INTENT, bindOptions)
+                .put(LauncherSettings.Favorites.APPWIDGET_SOURCE, sourceContainer);
     }
 
     /**
@@ -255,11 +266,29 @@
         return widgetFeatures;
     }
 
+    public static LauncherAtom.Attribute getAttribute(int container) {
+        switch (container) {
+            case CONTAINER_WIDGETS_TRAY:
+                return LauncherAtom.Attribute.WIDGETS;
+            case CONTAINER_BOTTOM_WIDGETS_TRAY:
+                return LauncherAtom.Attribute.WIDGETS_BOTTOM_TRAY;
+            case CONTAINER_PIN_WIDGETS:
+                return LauncherAtom.Attribute.PINITEM;
+            case CONTAINER_WIDGETS_PREDICTION:
+                return LauncherAtom.Attribute.WIDGETS_TRAY_PREDICTION;
+            case CONTAINER_ALL_APPS:
+                return LauncherAtom.Attribute.ALL_APPS_SEARCH_RESULT_WIDGETS;
+            default:
+                return LauncherAtom.Attribute.UNKNOWN;
+        }
+    }
+
     @Override
     public LauncherAtom.ItemInfo buildProto(FolderInfo folderInfo) {
         LauncherAtom.ItemInfo info = super.buildProto(folderInfo);
         return info.toBuilder()
                 .setWidget(info.getWidget().toBuilder().setWidgetFeatures(widgetFeatures))
+                .setAttribute(getAttribute(sourceContainer))
                 .build();
     }
 }
diff --git a/src/com/android/launcher3/util/PendingRequestArgs.java b/src/com/android/launcher3/util/PendingRequestArgs.java
index 9b8c6a6..77c8c0c 100644
--- a/src/com/android/launcher3/util/PendingRequestArgs.java
+++ b/src/com/android/launcher3/util/PendingRequestArgs.java
@@ -20,7 +20,9 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 
+import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.widget.PendingAddWidgetInfo;
 import com.android.launcher3.widget.WidgetAddFlowHandler;
 
 /**
@@ -34,11 +36,13 @@
     private static final int TYPE_APP_WIDGET = 2;
 
     private final int mArg1;
+    private final int mArg2;
     private final int mObjectType;
     private final Parcelable mObject;
 
     public PendingRequestArgs(ItemInfo info) {
         mArg1 = 0;
+        mArg2 = 0;
         mObjectType = TYPE_NONE;
         mObject = null;
 
@@ -46,7 +50,12 @@
     }
 
     private PendingRequestArgs(int arg1, int objectType, Parcelable object) {
+        this(arg1, 0, objectType, object);
+    }
+
+    private PendingRequestArgs(int arg1, int arg2, int objectType, Parcelable object) {
         mArg1 = arg1;
+        mArg2 = arg2;
         mObjectType = objectType;
         mObject = object;
     }
@@ -56,6 +65,7 @@
         user = parcel.readParcelable(null);
 
         mArg1 = parcel.readInt();
+        mArg2 = parcel.readInt();
         mObjectType = parcel.readInt();
         mObject = parcel.readParcelable(getClass().getClassLoader());
     }
@@ -73,6 +83,7 @@
         dest.writeParcelable(user, flags);
 
         dest.writeInt(mArg1);
+        dest.writeInt(mArg2);
         dest.writeInt(mObjectType);
         dest.writeParcelable(mObject, flags);
     }
@@ -85,6 +96,10 @@
         return mObjectType == TYPE_APP_WIDGET ? mArg1 : 0;
     }
 
+    public int getWidgetSourceContainer() {
+        return mObjectType == TYPE_APP_WIDGET ? mArg2 : Favorites.CONTAINER_UNKNOWN;
+    }
+
     public Intent getPendingIntent() {
         return mObjectType == TYPE_INTENT ? (Intent) mObject : null;
     }
@@ -95,8 +110,13 @@
 
     public static PendingRequestArgs forWidgetInfo(
             int appWidgetId, WidgetAddFlowHandler widgetHandler, ItemInfo info) {
+        int sourceContainer = Favorites.CONTAINER_UNKNOWN;
+        if (info instanceof PendingAddWidgetInfo) {
+            sourceContainer = ((PendingAddWidgetInfo) info).sourceContainer;
+        }
         PendingRequestArgs args =
-                new PendingRequestArgs(appWidgetId, TYPE_APP_WIDGET, widgetHandler);
+                new PendingRequestArgs(
+                        appWidgetId, sourceContainer, TYPE_APP_WIDGET, widgetHandler);
         args.copyFrom(info);
         return args;
     }
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index fa50dfb..bfe1587 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -52,6 +52,7 @@
 import com.android.launcher3.widget.dragndrop.AppWidgetHostViewDragListener;
 
 import java.util.List;
+import java.util.Optional;
 
 /**
  * {@inheritDoc}
@@ -97,7 +98,9 @@
     private final Object mUpdateLock = new Object();
     private final ViewGroupFocusHelper mDragLayerRelativeCoordinateHelper;
     private long mDeferUpdatesUntilMillis = 0;
-    private RemoteViews mMostRecentRemoteViews;
+    private RemoteViews mDeferredRemoteViews;
+    private Optional<SparseIntArray> mDeferredColorChange = Optional.empty();
+    private boolean mEnableColorExtraction = true;
 
     public LauncherAppWidgetHostView(Context context) {
         super(context);
@@ -146,8 +149,11 @@
     @Override
     public void updateAppWidget(RemoteViews remoteViews) {
         synchronized (mUpdateLock) {
-            mMostRecentRemoteViews = remoteViews;
-            if (SystemClock.uptimeMillis() < mDeferUpdatesUntilMillis) return;
+            if (isDeferringUpdates()) {
+                mDeferredRemoteViews = remoteViews;
+                return;
+            }
+            mDeferredRemoteViews = null;
         }
 
         super.updateAppWidget(remoteViews);
@@ -184,10 +190,19 @@
     }
 
     /**
+     * Returns true if the application of {@link RemoteViews} through {@link #updateAppWidget} and
+     * colors through {@link #onColorsChanged} are currently being deferred.
+     * @see #beginDeferringUpdates()
+     */
+    private boolean isDeferringUpdates() {
+        return SystemClock.uptimeMillis() < mDeferUpdatesUntilMillis;
+    }
+
+    /**
      * Begin deferring the application of any {@link RemoteViews} updates made through
-     * {@link #updateAppWidget(RemoteViews)} until {@link #endDeferringUpdates()} has been called or
-     * the next {@link #updateAppWidget(RemoteViews)} call after {@link #UPDATE_LOCK_TIMEOUT_MILLIS}
-     * have elapsed.
+     * {@link #updateAppWidget} and color changes through {@link #onColorsChanged} until
+     * {@link #endDeferringUpdates()} has been called or the next {@link #updateAppWidget} or
+     * {@link #onColorsChanged} call after {@link #UPDATE_LOCK_TIMEOUT_MILLIS} have elapsed.
      */
     public void beginDeferringUpdates() {
         synchronized (mUpdateLock) {
@@ -197,18 +212,23 @@
 
     /**
      * Stop deferring the application of {@link RemoteViews} updates made through
-     * {@link #updateAppWidget(RemoteViews)} and apply the most recently received update.
+     * {@link #updateAppWidget} and color changes made through {@link #onColorsChanged} and apply
+     * any deferred updates.
      */
     public void endDeferringUpdates() {
         RemoteViews remoteViews;
+        Optional<SparseIntArray> deferredColors;
         synchronized (mUpdateLock) {
             mDeferUpdatesUntilMillis = 0;
-            remoteViews = mMostRecentRemoteViews;
-            mMostRecentRemoteViews = null;
+            remoteViews = mDeferredRemoteViews;
+            mDeferredRemoteViews = null;
+            deferredColors = mDeferredColorChange;
+            mDeferredColorChange = Optional.empty();
         }
         if (remoteViews != null) {
             updateAppWidget(remoteViews);
         }
+        deferredColors.ifPresent(colors -> onColorsChanged(null /* rectF */, colors));
     }
 
     public boolean onInterceptTouchEvent(MotionEvent ev) {
@@ -282,12 +302,7 @@
         super.onLayout(changed, left, top, right, bottom);
 
         mIsScrollable = checkScrollableRecursively(this);
-        if (!mIsInDragMode && getTag() instanceof LauncherAppWidgetInfo) {
-            LauncherAppWidgetInfo info = (LauncherAppWidgetInfo) getTag();
-            mDragLayerRelativeCoordinateHelper.viewToRect(this, mCurrentWidgetSize);
-            updateColorExtraction(mCurrentWidgetSize,
-                    mWorkspace.getPageIndexForScreenId(info.screenId));
-        }
+        updateColorExtraction();
     }
 
     /** Starts the drag mode. */
@@ -314,6 +329,7 @@
      * @param pageId The workspace page the widget is on.
      */
     private void updateColorExtraction(Rect rectInDragLayer, int pageId) {
+        if (!mEnableColorExtraction) return;
         mColorExtractor.getExtractedRectForViewRect(mLauncher, pageId, rectInDragLayer, mTempRectF);
 
         if (mTempRectF.isEmpty()) {
@@ -328,6 +344,38 @@
         }
     }
 
+    /**
+     * Update the color extraction, using the current position of the app widget.
+     */
+    private void updateColorExtraction() {
+        if (!mIsInDragMode && getTag() instanceof LauncherAppWidgetInfo) {
+            LauncherAppWidgetInfo info = (LauncherAppWidgetInfo) getTag();
+            mDragLayerRelativeCoordinateHelper.viewToRect(this, mCurrentWidgetSize);
+            updateColorExtraction(mCurrentWidgetSize,
+                    mWorkspace.getPageIndexForScreenId(info.screenId));
+        }
+    }
+
+    /**
+     * Enables the local color extraction.
+     *
+     * @param updateColors If true, this will update the color extraction using the current location
+     *                    of the App Widget.
+     */
+    public void enableColorExtraction(boolean updateColors) {
+        mEnableColorExtraction = true;
+        if (updateColors) {
+            updateColorExtraction();
+        }
+    }
+
+    /**
+     * Disables the local color extraction.
+     */
+    public void disableColorExtraction() {
+        mEnableColorExtraction = false;
+    }
+
     // Compare two location rectangles. Locations are always in the [0;1] range.
     private static boolean isSameLocation(@NonNull RectF rect1, @Nullable RectF rect2,
             float epsilon) {
@@ -344,6 +392,14 @@
 
     @Override
     public void onColorsChanged(RectF rectF, SparseIntArray colors) {
+        synchronized (mUpdateLock) {
+            if (isDeferringUpdates()) {
+                mDeferredColorChange = Optional.ofNullable(colors);
+                return;
+            }
+            mDeferredColorChange = Optional.empty();
+        }
+
         // setColorResources will reapply the view, which must happen in the UI thread.
         post(() -> setColorResources(colors));
     }
diff --git a/src/com/android/launcher3/widget/PendingAddWidgetInfo.java b/src/com/android/launcher3/widget/PendingAddWidgetInfo.java
index c04c8dc..cbec642 100644
--- a/src/com/android/launcher3/widget/PendingAddWidgetInfo.java
+++ b/src/com/android/launcher3/widget/PendingAddWidgetInfo.java
@@ -15,14 +15,15 @@
  */
 package com.android.launcher3.widget;
 
-import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_TRAY;
-
 import android.appwidget.AppWidgetHostView;
 import android.content.Context;
 import android.os.Bundle;
 
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.PendingAddItemInfo;
+import com.android.launcher3.logger.LauncherAtom;
+import com.android.launcher3.model.data.FolderInfo;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.widget.util.WidgetSizes;
 
 /**
@@ -36,8 +37,9 @@
     public LauncherAppWidgetProviderInfo info;
     public AppWidgetHostView boundWidget;
     public Bundle bindOptions = null;
+    public int sourceContainer;
 
-    public PendingAddWidgetInfo(LauncherAppWidgetProviderInfo i) {
+    public PendingAddWidgetInfo(LauncherAppWidgetProviderInfo i, int container) {
         if (i.isCustomWidget()) {
             itemType = LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET;
         } else {
@@ -53,7 +55,7 @@
         spanY = i.spanY;
         minSpanX = i.minSpanX;
         minSpanY = i.minSpanY;
-        this.container = CONTAINER_WIDGETS_TRAY;
+        this.sourceContainer = this.container = container;
     }
 
     public WidgetAddFlowHandler getHandler() {
@@ -63,4 +65,12 @@
     public Bundle getDefaultSizeOptions(Context context) {
         return WidgetSizes.getWidgetSizeOptions(context, componentName, spanX, spanY);
     }
+
+    @Override
+    public LauncherAtom.ItemInfo buildProto(FolderInfo folderInfo) {
+        LauncherAtom.ItemInfo info = super.buildProto(folderInfo);
+        return info.toBuilder()
+                .setAttribute(LauncherAppWidgetInfo.getAttribute(sourceContainer))
+                .build();
+    }
 }
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index a52adf6..8798332 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.widget;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_TRAY;
 import static com.android.launcher3.Utilities.ATLEAST_S;
 
 import android.content.Context;
@@ -99,6 +100,7 @@
 
     private RemoteViews mRemoteViewsPreview;
     private NavigableAppWidgetHostView mAppWidgetHostViewPreview;
+    private int mSourceContainer = CONTAINER_WIDGETS_TRAY;
 
     public WidgetCell(Context context) {
         this(context, null);
@@ -177,6 +179,10 @@
         mAppWidgetHostViewPreview = null;
     }
 
+    public void setSourceContainer(int sourceContainer) {
+        this.mSourceContainer = sourceContainer;
+    }
+
     public void applyFromCellItem(WidgetItem item, WidgetPreviewLoader loader) {
         applyPreviewOnAppWidgetHostView(item);
 
@@ -205,7 +211,7 @@
             mPreviewWidth += mShortcutPreviewPadding;
             mPreviewHeight += mShortcutPreviewPadding;
         } else {
-            setTag(new PendingAddWidgetInfo(item.widgetInfo));
+            setTag(new PendingAddWidgetInfo(item.widgetInfo, mSourceContainer));
         }
     }
 
diff --git a/src/com/android/launcher3/widget/WidgetsBottomSheet.java b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
index d7928c7..c045cf1 100644
--- a/src/com/android/launcher3/widget/WidgetsBottomSheet.java
+++ b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.widget;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_BOTTOM_WIDGETS_TRAY;
 import static com.android.launcher3.anim.Interpolators.FAST_OUT_SLOW_IN;
 
 import android.animation.PropertyValuesHolder;
@@ -229,6 +230,7 @@
         previewContainer.setOnClickListener(this);
         previewContainer.setOnLongClickListener(this);
         widget.setAnimatePreview(false);
+        widget.setSourceContainer(CONTAINER_BOTTOM_WIDGETS_TRAY);
 
         parent.addView(widget);
         return widget;
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
index fe42ddf..6fc6848 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+
 import android.content.Context;
 import android.util.AttributeSet;
 import android.util.Log;
@@ -132,6 +134,7 @@
         previewContainer.setOnClickListener(mWidgetCellOnClickListener);
         previewContainer.setOnLongClickListener(mWidgetCellOnLongClickListener);
         widget.setAnimatePreview(false);
+        widget.setSourceContainer(CONTAINER_WIDGETS_PREDICTION);
 
         parent.addView(widget);
         return widget;
diff --git a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
index 0ec0f02..2b2fef4 100644
--- a/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
+++ b/tests/src/com/android/launcher3/util/rule/ShellCommandRule.java
@@ -15,21 +15,27 @@
  */
 package com.android.launcher3.util.rule;
 
-import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
-
 import static androidx.test.InstrumentationRegistry.getInstrumentation;
 
+import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
+
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
-
-import org.junit.rules.TestRule;
-import org.junit.runner.Description;
-import org.junit.runners.model.Statement;
+import android.util.Log;
 
 import androidx.annotation.Nullable;
 import androidx.test.InstrumentationRegistry;
 import androidx.test.uiautomator.UiDevice;
 
+import com.android.systemui.shared.system.PackageManagerWrapper;
+
+import org.junit.Assert;
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+import java.util.ArrayList;
+
 /**
  * Test rule which executes a shell command at the start of the test.
  */
@@ -37,10 +43,19 @@
 
     private final String mCmd;
     private final String mRevertCommand;
+    private final boolean mCheckSuccess;
+    private final Runnable mAdditionalChecks;
 
-    public ShellCommandRule(String cmd, @Nullable String revertCommand) {
+    public ShellCommandRule(String cmd, @Nullable String revertCommand, boolean checkSuccess,
+            Runnable additionalChecks) {
         mCmd = cmd;
         mRevertCommand = revertCommand;
+        mCheckSuccess = checkSuccess;
+        mAdditionalChecks = additionalChecks;
+    }
+
+    public ShellCommandRule(String cmd, @Nullable String revertCommand) {
+        this(cmd, revertCommand, false, null);
     }
 
     @Override
@@ -48,12 +63,27 @@
         return new Statement() {
             @Override
             public void evaluate() throws Throwable {
-                UiDevice.getInstance(getInstrumentation()).executeShellCommand(mCmd);
+                final String result =
+                        UiDevice.getInstance(getInstrumentation()).executeShellCommand(mCmd);
+                if (mCheckSuccess) {
+                    Assert.assertTrue(
+                            "Failed command: " + mCmd + ", result: " + result,
+                            "Success".equals(result.replaceAll("\\s", "")));
+                }
+                if (mAdditionalChecks != null) mAdditionalChecks.run();
                 try {
                     base.evaluate();
                 } finally {
                     if (mRevertCommand != null) {
-                        UiDevice.getInstance(getInstrumentation()).executeShellCommand(mRevertCommand);
+                        final String revertResult = UiDevice.getInstance(
+                                getInstrumentation()).executeShellCommand(
+                                mRevertCommand);
+                        if (mCheckSuccess) {
+                            Assert.assertTrue(
+                                    "Failed command: " + mRevertCommand
+                                            + ", result: " + revertResult,
+                                    "Success".equals(result.replaceAll("\\s", "")));
+                        }
                     }
                 }
             }
@@ -72,7 +102,15 @@
      * Sets the target launcher as default launcher.
      */
     public static ShellCommandRule setDefaultLauncher() {
-        return new ShellCommandRule(getLauncherCommand(getLauncherInMyProcess()), null);
+        final ActivityInfo launcher = getLauncherInMyProcess();
+        Log.d("b/187080582", "Launcher: " + new ComponentName(launcher.packageName, launcher.name)
+                .flattenToString());
+        return new ShellCommandRule(getLauncherCommand(launcher), null, true, () ->
+                Assert.assertEquals("Setting default launcher failed",
+                        new ComponentName(launcher.packageName, launcher.name)
+                                .flattenToString(),
+                        PackageManagerWrapper.getInstance().getHomeActivities(new ArrayList<>())
+                                .flattenToString()));
     }
 
     public static String getLauncherCommand(ActivityInfo launcher) {
diff --git a/tests/src_common/com/android/launcher3/common/WidgetUtils.java b/tests/src_common/com/android/launcher3/common/WidgetUtils.java
index 5e17e0a..97500e3 100644
--- a/tests/src_common/com/android/launcher3/common/WidgetUtils.java
+++ b/tests/src_common/com/android/launcher3/common/WidgetUtils.java
@@ -56,7 +56,9 @@
         item.container = LauncherSettings.Favorites.CONTAINER_DESKTOP;
 
         if (bindWidget) {
-            PendingAddWidgetInfo pendingInfo = new PendingAddWidgetInfo(info);
+            PendingAddWidgetInfo pendingInfo =
+                    new PendingAddWidgetInfo(
+                            info, LauncherSettings.Favorites.CONTAINER_WIDGETS_TRAY);
             pendingInfo.spanX = item.spanX;
             pendingInfo.spanY = item.spanY;
             pendingInfo.minSpanX = item.minSpanX;
diff --git a/tests/tapl/com/android/launcher3/tapl/LogEventChecker.java b/tests/tapl/com/android/launcher3/tapl/LogEventChecker.java
index a809e2e..710e3cd 100644
--- a/tests/tapl/com/android/launcher3/tapl/LogEventChecker.java
+++ b/tests/tapl/com/android/launcher3/tapl/LogEventChecker.java
@@ -15,6 +15,10 @@
  */
 package com.android.launcher3.tapl;
 
+import static com.android.launcher3.testing.TestProtocol.SEQUENCE_MAIN;
+import static com.android.launcher3.testing.TestProtocol.SEQUENCE_PILFER;
+import static com.android.launcher3.testing.TestProtocol.SEQUENCE_TIS;
+
 import android.os.SystemClock;
 
 import com.android.launcher3.testing.TestProtocol;
@@ -87,6 +91,24 @@
         final ListMap<String> actualEvents = finishSync(waitForExpectedCountMs);
         if (actualEvents == null) return "null event sequences because launcher likely died";
 
+        final String lowLevelDiags = lowLevelMismatchDiagnostics(actualEvents);
+        // If we have a sequence mismatch for a successful gesture, we want to provide all low-level
+        // details.
+        if (successfulGesture) {
+            return lowLevelDiags;
+        }
+
+        final String sequenceMismatchInEnglish = highLevelMismatchDiagnostics(actualEvents);
+
+        if (sequenceMismatchInEnglish != null) {
+            LauncherInstrumentation.log(lowLevelDiags);
+            return "Hint: " + sequenceMismatchInEnglish;
+        } else {
+            return lowLevelDiags;
+        }
+    }
+
+    private String lowLevelMismatchDiagnostics(ListMap<String> actualEvents) {
         final StringBuilder sb = new StringBuilder();
         boolean hasMismatches = false;
         for (Map.Entry<String, List<Pattern>> expectedEvents : mExpectedEvents.entrySet()) {
@@ -118,6 +140,42 @@
         return hasMismatches ? "Mismatching events: " + sb.toString() : null;
     }
 
+    private String highLevelMismatchDiagnostics(ListMap<String> actualEvents) {
+        if (!mExpectedEvents.getNonNull(SEQUENCE_TIS).isEmpty()
+                && actualEvents.getNonNull(SEQUENCE_TIS).isEmpty()) {
+            return "TouchInteractionService didn't receive any of the touch events sent by the "
+                    + "test";
+        }
+        if (getMismatchPosition(mExpectedEvents.getNonNull(SEQUENCE_TIS),
+                actualEvents.getNonNull(SEQUENCE_TIS)) != -1) {
+            // If TIS has a mismatch that we can't convert to high-level diags, don't convert
+            // other sequences either.
+            return null;
+        }
+
+        if (mExpectedEvents.getNonNull(SEQUENCE_PILFER).size() == 1
+                && actualEvents.getNonNull(SEQUENCE_PILFER).isEmpty()) {
+            return "Launcher didn't detect the navigation gesture sent by the test";
+        }
+        if (mExpectedEvents.getNonNull(SEQUENCE_PILFER).isEmpty()
+                && actualEvents.getNonNull(SEQUENCE_PILFER).size() == 1) {
+            return "Launcher detected a navigation gesture, but the test didn't send one";
+        }
+        if (getMismatchPosition(mExpectedEvents.getNonNull(SEQUENCE_PILFER),
+                actualEvents.getNonNull(SEQUENCE_PILFER)) != -1) {
+            // If Pilfer has a mismatch that we can't convert to high-level diags, don't analyze
+            // other sequences.
+            return null;
+        }
+
+        if (!mExpectedEvents.getNonNull(SEQUENCE_MAIN).isEmpty()
+                && actualEvents.getNonNull(SEQUENCE_MAIN).isEmpty()) {
+            return "None of the touch or keyboard events sent by the test was received by "
+                    + "Launcher's main thread";
+        }
+        return null;
+    }
+
     // If the list of actual events matches the list of expected events, returns -1, otherwise
     // the position of the mismatch.
     private static int getMismatchPosition(List<Pattern> expected, List<String> actual) {