diff --git a/quickstep/tests/src/com/android/quickstep/StartLauncherViaGestureTests.java b/quickstep/tests/src/com/android/quickstep/StartLauncherViaGestureTests.java
index 0c5b9ad..664958a 100644
--- a/quickstep/tests/src/com/android/quickstep/StartLauncherViaGestureTests.java
+++ b/quickstep/tests/src/com/android/quickstep/StartLauncherViaGestureTests.java
@@ -48,7 +48,9 @@
         // b/143488140
         mLauncher.pressHome();
         // Start an activity where the gestures start.
-        startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
+        if (mLauncher.getRecentTasks().isEmpty()) {
+            startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
+        }
     }
 
     private void runTest(String... eventSequence) {
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index bf093fd..ed25cc7 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -69,14 +69,16 @@
     }
 
     private void startTestApps() throws Exception {
-        startAppFast(getAppPackageName());
-        startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
-        startTestActivity(2);
+        if (mLauncher.getRecentTasks().size() < 3) {
+            startAppFast(getAppPackageName());
+            startAppFast(resolveSystemApp(Intent.CATEGORY_APP_CALCULATOR));
+            startTestActivity(2);
 
-        executeOnLauncher(launcher -> assertTrue(
-                "Launcher activity is the top activity; expecting another activity to be the top "
-                        + "one",
-                isInBackground(launcher)));
+            executeOnLauncher(launcher -> assertTrue(
+                    "Launcher activity is the top activity; expecting another activity to be the "
+                            + "top one",
+                    isInBackground(launcher)));
+        }
     }
 
     @Test
diff --git a/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java b/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java
index 5af9113..d7af5f1 100644
--- a/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsInsetTransitionController.java
@@ -58,6 +58,17 @@
         mApps = appsView;
     }
 
+    public void hide() {
+        if (!BuildCompat.isAtLeastR()) return;
+
+        WindowInsets insets = mApps.getRootWindowInsets();
+        if (insets == null) return;
+
+        if (insets.isVisible(WindowInsets.Type.ime())) {
+            mApps.getWindowInsetsController().hide(WindowInsets.Type.ime());
+        }
+    }
+
     /**
      * Initializes member variables and requests for the {@link WindowInsetsAnimationController}
      * object.
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index 1224cb0..0268b96 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -36,6 +36,7 @@
 import android.util.FloatProperty;
 import android.view.View;
 import android.view.animation.Interpolator;
+import android.widget.EditText;
 
 import androidx.core.os.BuildCompat;
 
@@ -263,8 +264,16 @@
         }
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get() && BuildCompat.isAtLeastR()) {
             mInsetController.onAnimationEnd(mProgress);
+            if (Float.compare(mProgress, 0f) == 0) {
+                EditText editText = mAppsView.getSearchUiManager().getEditText();
+                if (editText != null) {
+                    editText.requestFocus();
+                }
+            }
             if (Float.compare(mProgress, 1f) == 0) {
-                mAppsView.getSearchUiManager().setTextSearchEnabled(true).requestFocus();
+                // Called when home gesture closes all apps container.
+                // TODO: should make the controller hide synchronously
+                mInsetController.hide();
             }
         }
     }
diff --git a/src/com/android/launcher3/allapps/SearchUiManager.java b/src/com/android/launcher3/allapps/SearchUiManager.java
index 7d5363f..aa056a0 100644
--- a/src/com/android/launcher3/allapps/SearchUiManager.java
+++ b/src/com/android/launcher3/allapps/SearchUiManager.java
@@ -66,12 +66,8 @@
     }
 
     /**
-     * Called to control how the search UI result should be handled.
-     *
-     * @param isEnabled when {@code true}, the search is all handled inside AOSP
-     *                  and is not overlayable.
-     * @return the searchbox edit text object
+     * @return the edit text object
      */
     @Nullable
-    EditText setTextSearchEnabled(boolean isEnabled);
+    EditText getEditText();
 }
diff --git a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
index 356c52c..2e89c89 100644
--- a/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
+++ b/src/com/android/launcher3/allapps/search/AppsSearchContainerLayout.java
@@ -216,7 +216,7 @@
     }
 
     @Override
-    public EditText setTextSearchEnabled(boolean isEnabled) {
+    public EditText getEditText() {
         return this;
     }
 }