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/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 0da16b7..099f256 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -1339,6 +1339,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);
@@ -2416,7 +2422,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 ad2d7c2..28fd389 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -379,6 +379,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<>();
 
@@ -746,6 +748,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_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;