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/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/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/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/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/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/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) {