diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index c4a3fb5..8644a12 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -1288,8 +1288,8 @@
             }
         }
 
-        return mLauncher.getFirstMatchForAppClose(launchCookieItemId,
-                packageName, UserHandle.of(runningTaskTarget.taskInfo.userId));
+        return mLauncher.getFirstMatchForAppClose(launchCookieItemId, packageName,
+                UserHandle.of(runningTaskTarget.taskInfo.userId), true /* supportsAllAppsState */);
     }
 
     private @NonNull RectF getDefaultWindowTargetRect() {
diff --git a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
index af6cb84..bc7a6ae 100644
--- a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
+++ b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
@@ -233,7 +233,8 @@
 
         return mActivity.getFirstMatchForAppClose(launchCookieItemId,
                 runningTaskView.getTask().key.getComponent().getPackageName(),
-                UserHandle.of(runningTaskView.getTask().key.userId));
+                UserHandle.of(runningTaskView.getTask().key.userId),
+                false /* supportsAllAppsState */);
     }
 
     @Override
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8fb2f53..9bb815c 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2743,8 +2743,11 @@
      * @param preferredItemId The id of the preferred item to match to if it exists.
      * @param packageName The package name of the app to match.
      * @param user The user of the app to match.
+     * @param supportsAllAppsState If true and we are in All Apps state, looks for view in All Apps.
+     *                             Else we only looks on the workspace.
      */
-    public View getFirstMatchForAppClose(int preferredItemId, String packageName, UserHandle user) {
+    public View getFirstMatchForAppClose(int preferredItemId, String packageName, UserHandle user,
+            boolean supportsAllAppsState) {
         final ItemInfoMatcher preferredItem = (info, cn) ->
                 info != null && info.id == preferredItemId;
         final ItemInfoMatcher packageAndUserAndApp = (info, cn) ->
@@ -2755,7 +2758,7 @@
                         && TextUtils.equals(info.getTargetComponent().getPackageName(),
                         packageName);
 
-        if (isInState(LauncherState.ALL_APPS)) {
+        if (supportsAllAppsState && isInState(LauncherState.ALL_APPS)) {
             return getFirstMatch(Collections.singletonList(mAppsView.getActiveRecyclerView()),
                     preferredItem, packageAndUserAndApp);
         } else {
diff --git a/src/com/android/launcher3/views/FloatingSurfaceView.java b/src/com/android/launcher3/views/FloatingSurfaceView.java
index 7f54d6d..19c28b4 100644
--- a/src/com/android/launcher3/views/FloatingSurfaceView.java
+++ b/src/com/android/launcher3/views/FloatingSurfaceView.java
@@ -159,7 +159,8 @@
             return;
         }
         View icon = mLauncher.getFirstMatchForAppClose(-1,
-                mContract.componentName.getPackageName(), mContract.user);
+                mContract.componentName.getPackageName(), mContract.user,
+                false /* supportsAllAppsState */);
 
         boolean iconChanged = mIcon != icon;
         if (iconChanged) {