diff --git a/quickstep/src/com/android/launcher3/search/SearchResultIcon.java b/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
index 3dd4146..4c44479 100644
--- a/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
+++ b/quickstep/src/com/android/launcher3/search/SearchResultIcon.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.search;
 
+import static com.android.launcher3.LauncherSettings.Favorites.EXTENDED_CONTAINERS;
 import static com.android.launcher3.model.data.SearchActionItemInfo.FLAG_BADGE_WITH_PACKAGE;
 import static com.android.launcher3.model.data.SearchActionItemInfo.FLAG_PRIMARY_ICON_FROM_TITLE;
 import static com.android.launcher3.search.SearchTargetUtil.BUNDLE_EXTRA_PRIMARY_ICON_FROM_TITLE;
@@ -49,7 +50,6 @@
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.BitmapRenderer;
 import com.android.launcher3.icons.LauncherIcons;
-import com.android.launcher3.logger.LauncherAtom.ContainerInfo;
 import com.android.launcher3.logger.LauncherAtomExtensions.DeviceSearchResultContainer;
 import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
 import com.android.launcher3.model.data.AppInfo;
@@ -147,11 +147,10 @@
         SearchActionItemInfo itemInfo = new SearchActionItemInfo(searchAction.getIcon(),
                 searchTarget.getPackageName(), searchTarget.getUserHandle(),
                 searchAction.getTitle()) {
-            // Workaround to log ItemInfo with DeviceSearchResultContainer without
-            // updating ItemInfo.container field.
             @Override
-            public ContainerInfo getContainerInfo() {
-                return buildDeviceSearchResultContainer();
+            protected ExtendedContainers getExtendedContainer() {
+                return ExtendedContainers.newBuilder()
+                        .setDeviceSearchResultContainer(buildDeviceSearchResultContainer()).build();
             }
         };
         itemInfo.setIntent(searchAction.getIntent());
@@ -255,14 +254,13 @@
         AllAppsStore appsStore = mLauncher.getAppsView().getAppsStore();
         AppInfo appInfo = new AppInfo(
                 appsStore.getApp(new ComponentKey(componentName, userHandle))) {
-            // Workaround to log ItemInfo with DeviceSearchResultContainer without
-            // updating ItemInfo.container field.
             @Override
-            public ContainerInfo getContainerInfo() {
-                return buildDeviceSearchResultContainer();
+            protected ExtendedContainers getExtendedContainer() {
+                return ExtendedContainers.newBuilder()
+                        .setDeviceSearchResultContainer(buildDeviceSearchResultContainer()).build();
             }
         };
-
+        appInfo.container = EXTENDED_CONTAINERS;
         if (appInfo == null) {
             setVisibility(GONE);
             return;
@@ -273,13 +271,13 @@
 
     private void prepareUsingShortcutInfo(ShortcutInfo shortcutInfo) {
         WorkspaceItemInfo workspaceItemInfo = new WorkspaceItemInfo(shortcutInfo, getContext()) {
-            // Workaround to log ItemInfo with DeviceSearchResultContainer without
-            // updating ItemInfo.container field.
             @Override
-            public ContainerInfo getContainerInfo() {
-                return buildDeviceSearchResultContainer();
+            protected ExtendedContainers getExtendedContainer() {
+                return ExtendedContainers.newBuilder()
+                        .setDeviceSearchResultContainer(buildDeviceSearchResultContainer()).build();
             }
         };
+        workspaceItemInfo.container = EXTENDED_CONTAINERS;
         notifyItemInfoChanged(workspaceItemInfo);
         LauncherAppState launcherAppState = LauncherAppState.getInstance(getContext());
         MODEL_EXECUTOR.execute(() -> {
@@ -318,15 +316,10 @@
         }
     }
 
-    private ContainerInfo buildDeviceSearchResultContainer() {
-        return ContainerInfo.newBuilder().setExtendedContainers(
-                ExtendedContainers
-                        .newBuilder()
-                        .setDeviceSearchResultContainer(
-                                mSearchSessionTracker.getQueryLength()
-                                        .map(queryLength -> DeviceSearchResultContainer.newBuilder()
-                                                .setQueryLength(queryLength))
-                                        .orElse(DeviceSearchResultContainer.newBuilder())))
-                .build();
+    private DeviceSearchResultContainer buildDeviceSearchResultContainer() {
+        return mSearchSessionTracker.getQueryLength()
+                .map(queryLength -> DeviceSearchResultContainer.newBuilder()
+                        .setQueryLength(queryLength))
+                .orElse(DeviceSearchResultContainer.newBuilder()).build();
     }
 }
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index fe423ed..be27039 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -171,6 +171,9 @@
         public static final int CONTAINER_SETTINGS = -108;
         public static final int CONTAINER_TASKSWITCHER = -109;
 
+        // Represents any of the extended containers implemented in non-AOSP variants.
+        public static final int EXTENDED_CONTAINERS = -200;
+
         public static final String containerToString(int container) {
             switch (container) {
                 case CONTAINER_DESKTOP: return "desktop";
diff --git a/src/com/android/launcher3/model/data/ItemInfo.java b/src/com/android/launcher3/model/data/ItemInfo.java
index 3851ab0..00ac12f 100644
--- a/src/com/android/launcher3/model/data/ItemInfo.java
+++ b/src/com/android/launcher3/model/data/ItemInfo.java
@@ -26,6 +26,7 @@
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_SHORTCUTS;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_TASKSWITCHER;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_TRAY;
+import static com.android.launcher3.LauncherSettings.Favorites.EXTENDED_CONTAINERS;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
@@ -54,6 +55,7 @@
 import com.android.launcher3.logger.LauncherAtom.Shortcut;
 import com.android.launcher3.logger.LauncherAtom.ShortcutsContainer;
 import com.android.launcher3.logger.LauncherAtom.TaskSwitcherContainer;
+import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
 import com.android.launcher3.model.ModelWriter;
 import com.android.launcher3.util.ContentWriter;
 
@@ -403,12 +405,23 @@
                 return ContainerInfo.newBuilder()
                         .setTaskSwitcherContainer(TaskSwitcherContainer.getDefaultInstance())
                         .build();
-
+            case EXTENDED_CONTAINERS:
+                return ContainerInfo.newBuilder()
+                        .setExtendedContainers(getExtendedContainer())
+                        .build();
         }
         return ContainerInfo.getDefaultInstance();
     }
 
     /**
+     * Returns non-AOSP container wrapped by {@link ExtendedContainers} object. Should be overridden
+     * by build variants.
+     */
+    protected ExtendedContainers getExtendedContainer() {
+        return ExtendedContainers.getDefaultInstance();
+    }
+
+    /**
      * Returns shallow copy of the object.
      */
     public ItemInfo makeShallowCopy() {
diff --git a/src/com/android/launcher3/model/data/SearchActionItemInfo.java b/src/com/android/launcher3/model/data/SearchActionItemInfo.java
index 8469569..25355c9 100644
--- a/src/com/android/launcher3/model/data/SearchActionItemInfo.java
+++ b/src/com/android/launcher3/model/data/SearchActionItemInfo.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model.data;
 
+import static com.android.launcher3.LauncherSettings.Favorites.EXTENDED_CONTAINERS;
+
 import android.app.PendingIntent;
 import android.content.Intent;
 import android.graphics.drawable.Icon;
@@ -45,6 +47,7 @@
             CharSequence title) {
         this.user = user == null ? Process.myUserHandle() : user;
         this.title = title;
+        this.container = EXTENDED_CONTAINERS;
         mFallbackPackageName = packageName;
         mIcon = icon;
     }
@@ -55,6 +58,7 @@
         mFallbackPackageName = info.mFallbackPackageName;
         mFlags = info.mFlags;
         title = info.title;
+        this.container = EXTENDED_CONTAINERS;
     }
 
     /**