diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 266ea52ae7..3b493bbd30 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -1298,8 +1298,8 @@ public class QuickstepTransitionManager implements OnDeviceProfileChangeListener
             }
         }
 
-        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 a72935b250..4fb03c4de4 100644
--- a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
+++ b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
@@ -246,7 +246,8 @@ public class LauncherSwipeHandlerV2 extends
 
         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 89dddb9359..0e41c797f5 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2713,8 +2713,11 @@ public class Launcher extends StatefulActivity<LauncherState> implements Launche
      * @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) ->
@@ -2725,7 +2728,7 @@ public class Launcher extends StatefulActivity<LauncherState> implements Launche
                         && 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 f32f904981..09c8c64087 100644
--- a/src/com/android/launcher3/views/FloatingSurfaceView.java
+++ b/src/com/android/launcher3/views/FloatingSurfaceView.java
@@ -159,7 +159,8 @@ public class FloatingSurfaceView extends AbstractFloatingView implements
             return;
         }
         View icon = mLauncher.getFirstMatchForAppClose(-1,
-                mContract.componentName.getPackageName(), mContract.user);
+                mContract.componentName.getPackageName(), mContract.user,
+                false /* supportsAllAppsState */);
 
         boolean iconChanged = mIcon != icon;
         if (iconChanged) {