diff --git a/quickstep/protos_overrides/launcher_atom_extension.proto b/quickstep/protos_overrides/launcher_atom_extension.proto
index a1566f0..f5a277b 100644
--- a/quickstep/protos_overrides/launcher_atom_extension.proto
+++ b/quickstep/protos_overrides/launcher_atom_extension.proto
@@ -22,6 +22,7 @@
 // Wrapper message for containers used at the quickstep level.
 // Message name should match with launcher_atom_extension.proto message at
 // the AOSP level.
+// Next ID = 3
 message ExtendedContainers {
   reserved 2; // Deleted fields
 
@@ -31,10 +32,16 @@
 }
 
 // Represents on-device search result container.
+// Next ID = 4
 message DeviceSearchResultContainer{
   optional int32 query_length = 1;
   optional SearchAttributes search_attributes = 2;
+  // [0, m], m varies based on the display density and resolution
+  // To indicate the location of the tapped on-device search result.
+  // For application, it will be the column number in the apps row.
+  optional int32 grid_x = 3;
 
+  // Next ID = 4
   message SearchAttributes{
 
     // True if results are based on spell corrected query
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 3b9e2b2..45c8036 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -582,14 +582,18 @@
     }
 
     private static int getGridX(LauncherAtom.ItemInfo info, boolean parent) {
-        if (info.getContainerInfo().getContainerCase() == FOLDER) {
+        LauncherAtom.ContainerInfo containerInfo = info.getContainerInfo();
+        if (containerInfo.getContainerCase() == FOLDER) {
             if (parent) {
-                return info.getContainerInfo().getFolder().getWorkspace().getGridX();
+                return containerInfo.getFolder().getWorkspace().getGridX();
             } else {
-                return info.getContainerInfo().getFolder().getGridX();
+                return containerInfo.getFolder().getGridX();
             }
+        } else if (containerInfo.getContainerCase() == EXTENDED_CONTAINERS) {
+            return containerInfo.getExtendedContainers()
+                    .getDeviceSearchResultContainer().getGridX();
         } else {
-            return info.getContainerInfo().getWorkspace().getGridX();
+            return containerInfo.getWorkspace().getGridX();
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
index 19a48db..9bb3d56 100644
--- a/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
+++ b/quickstep/src/com/android/quickstep/util/SystemWindowManagerProxy.java
@@ -16,6 +16,7 @@
 package com.android.quickstep.util;
 
 import android.content.Context;
+import android.hardware.display.DisplayManager;
 import android.view.Display;
 
 import com.android.launcher3.util.window.WindowManagerProxy;
@@ -43,4 +44,10 @@
     public int getRotation(Context context) {
         return context.getResources().getConfiguration().windowConfiguration.getRotation();
     }
+
+    @Override
+    protected Display[] getDisplays(Context context) {
+        return context.getSystemService(DisplayManager.class).getDisplays(
+                DisplayManager.DISPLAY_CATEGORY_ALL_INCLUDING_DISABLED);
+    }
 }
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index 515f80a..641161b 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -209,7 +209,7 @@
         int oldMaxHeight = mMaxTranslation;
         updateExpectedHeight();
 
-        if (mMaxTranslation != oldMaxHeight) {
+        if (mMaxTranslation != oldMaxHeight || mCollapsed) {
             BaseAllAppsContainerView<?> parent = (BaseAllAppsContainerView<?>) getParent();
             if (parent != null) {
                 parent.setupHeader();
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 838564d..33beacd 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -265,8 +265,8 @@
     public static final BooleanFlag ENABLE_ONE_SEARCH_MOTION = new DeviceFlag(
             "ENABLE_ONE_SEARCH_MOTION", true, "Enables animations in OneSearch.");
 
-    public static final BooleanFlag ENABLE_SHOW_KEYBOARD_IN_ALL_APPS = new DeviceFlag(
-            "ENABLE_SHOW_KEYBOARD_IN_ALL_APPS", true,
+    public static final BooleanFlag ENABLE_SHOW_KEYBOARD_OPTION_IN_ALL_APPS = new DeviceFlag(
+            "ENABLE_SHOW_KEYBOARD_OPTION_IN_ALL_APPS", true,
             "Enable option to show keyboard when going to all-apps");
 
     public static final BooleanFlag USE_LOCAL_ICON_OVERRIDES = getDebugFlag(
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index 7d24fe8..3d7d235 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -735,7 +735,8 @@
             UNKNOWN(0),
             COLD(1),
             HOT(2),
-            TIMEOUT(3);
+            TIMEOUT(3),
+            FAIL(4);
 
             private final int mId;
 
diff --git a/src/com/android/launcher3/util/window/WindowManagerProxy.java b/src/com/android/launcher3/util/window/WindowManagerProxy.java
index 92f718e..9665bf9 100644
--- a/src/com/android/launcher3/util/window/WindowManagerProxy.java
+++ b/src/com/android/launcher3/util/window/WindowManagerProxy.java
@@ -90,7 +90,7 @@
      */
     public ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> estimateInternalDisplayBounds(
             Context context) {
-        Display[] displays = context.getSystemService(DisplayManager.class).getDisplays();
+        Display[] displays = getDisplays(context);
         ArrayMap<String, Pair<CachedDisplayInfo, WindowBounds[]>> result = new ArrayMap<>();
         for (Display display : displays) {
             if (isInternalDisplay(display)) {
@@ -355,4 +355,11 @@
         }
         return d.getRotation();
     }
+
+    /**
+     * Returns all currently valid logical displays.
+     */
+    protected Display[] getDisplays(Context context) {
+        return context.getSystemService(DisplayManager.class).getDisplays();
+    }
 }