Move hideKeyboard call to end of the all apps transition controller animation

Bug: 236685775
Bug: 235951830

Test: manual
Change-Id: Iae808a82eb4a10a554aef1cb5d87fc806efb0279
diff --git a/src/com/android/launcher3/allapps/AllAppsTransitionController.java b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
index e0f1b3d..a86faab 100644
--- a/src/com/android/launcher3/allapps/AllAppsTransitionController.java
+++ b/src/com/android/launcher3/allapps/AllAppsTransitionController.java
@@ -44,7 +44,6 @@
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.util.MultiAdditivePropertyFactory;
 import com.android.launcher3.util.MultiValueAlpha;
-import com.android.launcher3.util.UiThreadHelper;
 import com.android.launcher3.views.ScrimView;
 
 /**
@@ -229,7 +228,6 @@
     public void setStateWithAnimation(LauncherState toState,
             StateAnimationConfig config, PendingAnimation builder) {
         if (NORMAL.equals(toState) && mLauncher.isInState(ALL_APPS)) {
-            UiThreadHelper.hideKeyboardAsync(mLauncher, mLauncher.getAppsView().getWindowToken());
             builder.addEndListener(success -> {
                 // Reset pull back progress and alpha after switching states.
                 ALL_APPS_PULL_BACK_TRANSLATION.set(this, 0f);
@@ -311,6 +309,7 @@
         if (FeatureFlags.ENABLE_DEVICE_SEARCH.get()) return;
         if (Float.compare(mProgress, 1f) == 0) {
             mAppsView.reset(false /* animate */);
+            mLauncher.getAppsView().getSearchUiManager().getEditText().hideKeyboard();
         }
     }
 }