[25/N] Desks: Replace visibleTaskCount checks with active desk checks

Many components use visibleTaskCount as a proxy for whether the system
is currently in desktop mode. This change abstracts that behind a
DesktopRepository#isAnyDeskActive function, which relies on activation
state rather than task count when multiple-desks are enabled.

Flag: com.android.window.flags.enable_multiple_desktops_backend
Bug: 362720497
Test: move a task to a desk, verify the taskbar becomes pinned
Change-Id: I8d0798d97444ffcffc8f7d467c1f26a66e989244
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java
index c6afc31..b507ca2 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/pip/PipDesktopState.java
@@ -88,7 +88,7 @@
             return false;
         }
         final int displayId = mPipDisplayLayoutState.getDisplayId();
-        return getDesktopRepository().getVisibleTaskCount(displayId) > 0
+        return getDesktopRepository().isAnyDeskActive(displayId)
                 || getDesktopWallpaperActivityTokenProvider().isWallpaperActivityVisible(displayId)
                 || isDisplayInFreeform();
     }
@@ -100,7 +100,7 @@
             return false;
         }
         final DesktopRepository desktopRepository = getDesktopRepository();
-        return desktopRepository.getVisibleTaskCount(pipTask.getDisplayId()) > 0
+        return desktopRepository.isAnyDeskActive(pipTask.getDisplayId())
                 || desktopRepository.isMinimizedPipPresentInDisplay(pipTask.getDisplayId());
     }
 
@@ -114,7 +114,7 @@
             return false;
         }
         final int displayId = mPipDisplayLayoutState.getDisplayId();
-        return getDesktopRepository().getVisibleTaskCount(displayId) == 0
+        return !getDesktopRepository().isAnyDeskActive(displayId)
                 && getDesktopWallpaperActivityTokenProvider().isWallpaperActivityVisible(displayId);
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java
index 201870f..9b11e4a 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/compatui/CompatUIController.java
@@ -872,7 +872,7 @@
             return false;
         }
         boolean isDesktopModeShowing = mDesktopUserRepositories.get().getCurrent()
-                .getVisibleTaskCount(taskInfo.displayId) > 0;
+                .isAnyDeskActive(taskInfo.displayId);
         return DesktopModeFlags.ENABLE_DESKTOP_SKIP_COMPAT_UI_EDUCATION_IN_DESKTOP_MODE_BUGFIX
                 .isTrue() && isDesktopModeShowing;
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt
index 6104e79..b8f4bb8 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopActivityOrientationChangeHandler.kt
@@ -84,7 +84,7 @@
         if (!Flags.respectOrientationChangeForUnresizeable()) return
         val task = shellTaskOrganizer.getRunningTaskInfo(taskId) ?: return
         val taskRepository = desktopUserRepositories.current
-        val isDesktopModeShowing = taskRepository.getVisibleTaskCount(task.displayId) > 0
+        val isDesktopModeShowing = taskRepository.isAnyDeskActive(task.displayId)
         if (!isDesktopModeShowing || !task.isFreeform || task.isResizeable) return
 
         val taskBounds = task.configuration.windowConfiguration.bounds
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index 3c694ae..2ce039e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -418,11 +418,15 @@
             .singleOrNull() == taskId
     }
 
-    /**
-     * Returns the active tasks in the given display's active desk.
-     *
-     * TODO: b/389960283 - migrate callers to [getActiveTaskIdsInDesk].
-     */
+    /** Whether the task is the only visible desktop task in the display. */
+    fun isOnlyVisibleTask(taskId: Int, displayId: Int): Boolean {
+        val desk = desktopData.getActiveDesk(displayId) ?: return false
+        return desk.visibleTasks.size == 1 && desk.visibleTasks.single() == taskId
+    }
+
+    /** Whether the display has only one visible desktop task. */
+    fun hasOnlyOneVisibleTask(displayId: Int): Boolean = getVisibleTaskCount(displayId) == 1
+
     @VisibleForTesting
     fun getActiveTasks(displayId: Int): ArraySet<Int> =
         ArraySet(desktopData.getActiveDesk(displayId)?.activeTasks)
@@ -551,7 +555,7 @@
         } else {
             desk.visibleTasks.remove(taskId)
         }
-        val newCount = getVisibleTaskCount(deskId)
+        val newCount = getVisibleTaskCountInDesk(deskId)
         if (prevCount != newCount) {
             logD(
                 "Update task visibility taskId=%d visible=%b deskId=%d displayId=%d",
@@ -735,18 +739,29 @@
         }
     }
 
-    /**
-     * Gets number of visible freeform tasks on given [displayId]'s active desk.
-     *
-     * TODO: b/389960283 - migrate callers to [getVisibleTaskCountInDesk].
-     */
+    /** Whether the display is currently showing any desk. */
+    fun isAnyDeskActive(displayId: Int): Boolean {
+        if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+            val desk = desktopData.getDefaultDesk(displayId)
+            if (desk == null) {
+                logE("Could not find default desk for display: $displayId")
+                return false
+            }
+            return desk.visibleTasks.isNotEmpty()
+        }
+        return desktopData.getActiveDesk(displayId) != null
+    }
+
+    /** Gets number of visible freeform tasks on given [displayId]'s active desk. */
+    @Deprecated("Use isAnyDeskActive() instead.", ReplaceWith("isAnyDeskActive()"))
+    @VisibleForTesting
     fun getVisibleTaskCount(displayId: Int): Int =
         (desktopData.getActiveDesk(displayId)?.visibleTasks?.size ?: 0).also {
             logD("getVisibleTaskCount=$it")
         }
 
     /** Gets the number of visible tasks on the given desk. */
-    fun getVisibleTaskCountInDesk(deskId: Int): Int =
+    private fun getVisibleTaskCountInDesk(deskId: Int): Int =
         desktopData.getDesk(deskId)?.visibleTasks?.size ?: 0
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
index bfad0b3..a6f8ebd 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksController.kt
@@ -335,8 +335,8 @@
         activateDefaultDeskInDisplay(displayId, remoteTransition)
     }
 
-    /** Gets number of visible freeform tasks in [displayId]. */
-    fun visibleTaskCount(displayId: Int): Int = taskRepository.getVisibleTaskCount(displayId)
+    /** Returns whether the given display has an active desk. */
+    fun isAnyDeskActive(displayId: Int): Boolean = taskRepository.isAnyDeskActive(displayId)
 
     /**
      * Returns true if any of the following is true:
@@ -344,9 +344,16 @@
      * - A transparent fullscreen task exists on top in Desktop Mode
      * - PiP on Desktop Windowing is enabled, there is an active PiP window and the desktop
      *   wallpaper is visible.
+     *
+     * TODO: b/362720497 - consolidate with [isAnyDeskActive].
+     *     - top-transparent-fullscreen case: should not be needed if we allow it to launch inside
+     *       the desk in fullscreen instead of force-exiting desktop and having to trick this method
+     *       into thinking it is in desktop mode when a task in this state exists.
+     *     - PIP case: a PIP presence should influence desk activation, so
+     *       [DesktopRepository#isAnyDeskActive] should be sufficient.
      */
     fun isDesktopModeShowing(displayId: Int): Boolean {
-        val hasVisibleTasks = visibleTaskCount(displayId) > 0
+        val hasVisibleTasks = taskRepository.isAnyDeskActive(displayId)
         val hasTopTransparentFullscreenTask =
             taskRepository.getTopTransparentFullscreenTaskId(displayId) != null
         val hasMinimizedPip =
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
index 9a97ae8..df2cf67 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserver.kt
@@ -26,6 +26,7 @@
 import android.view.WindowManager.TRANSIT_PIP
 import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TRANSIT_TO_FRONT
+import android.window.DesktopExperienceFlags
 import android.window.DesktopModeFlags
 import android.window.DesktopModeFlags.ENABLE_DESKTOP_WALLPAPER_ACTIVITY_FOR_SYSTEM_USER
 import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
@@ -162,18 +163,24 @@
                     continue
                 }
                 val desktopRepository = desktopUserRepositories.getProfile(taskInfo.userId)
-                val visibleTaskCount = desktopRepository.getVisibleTaskCount(taskInfo.displayId)
+                val isInDesktop = desktopRepository.isAnyDeskActive(taskInfo.displayId)
                 if (
-                    visibleTaskCount > 0 &&
+                    isInDesktop &&
                         change.mode == TRANSIT_TO_BACK &&
                         taskInfo.windowingMode == WINDOWING_MODE_FREEFORM
                 ) {
+                    val isLastTask =
+                        if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+                            desktopRepository.hasOnlyOneVisibleTask(taskInfo.displayId)
+                        } else {
+                            desktopRepository.isOnlyVisibleTask(taskInfo.taskId, taskInfo.displayId)
+                        }
                     desktopRepository.minimizeTask(taskInfo.displayId, taskInfo.taskId)
                     desktopMixedTransitionHandler.addPendingMixedTransition(
                         DesktopMixedTransitionHandler.PendingMixedTransition.Minimize(
                             transition,
                             taskInfo.taskId,
-                            visibleTaskCount == 1,
+                            isLastTask,
                         )
                     )
                 }
@@ -227,9 +234,9 @@
         taskInfo: ActivityManager.RunningTaskInfo
     ): Int? {
         val desktopRepository = desktopUserRepositories.getProfile(taskInfo.userId)
-        val visibleTaskCount = desktopRepository.getVisibleTaskCount(taskInfo.displayId)
+        val isInDesktop = desktopRepository.isAnyDeskActive(taskInfo.displayId)
         if (
-            visibleTaskCount > 0 &&
+            isInDesktop &&
                 taskInfo.windowingMode == WINDOWING_MODE_FREEFORM &&
                 backAnimationController.latestTriggerBackTask == taskInfo.taskId &&
                 !desktopRepository.isClosingTask(taskInfo.taskId)
@@ -253,7 +260,7 @@
 
             val desktopRepository = desktopUserRepositories.getProfile(taskInfo.userId)
             if (
-                desktopRepository.getVisibleTaskCount(taskInfo.displayId) == 0 &&
+                !desktopRepository.isAnyDeskActive(taskInfo.displayId) &&
                     change.mode == TRANSIT_CLOSE &&
                     taskInfo.windowingMode == WINDOWING_MODE_FREEFORM &&
                     desktopWallpaperActivityTokenProvider.getToken(taskInfo.displayId) != null
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandler.kt
index 224ff37..3576b25 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandler.kt
@@ -162,7 +162,7 @@
         }
 
     private fun isDesktopModeShowing(displayId: Int): Boolean =
-        desktopUserRepositories.current.getVisibleTaskCount(displayId) > 0
+        desktopUserRepositories.current.isAnyDeskActive(displayId)
 
     override fun handleRequest(
         transition: IBinder,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
index 61a193c..04f0336 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
@@ -782,7 +782,7 @@
     private boolean isPipExitingToDesktopMode() {
         DesktopRepository currentRepo = getCurrentRepo();
         return Flags.enableDesktopWindowingPip() && currentRepo != null
-                && (currentRepo.getVisibleTaskCount(mTaskInfo.displayId) > 0
+                && (currentRepo.isAnyDeskActive(mTaskInfo.displayId)
                     || isDisplayInFreeform());
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
index 347dcff..a0fb625 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
@@ -377,7 +377,7 @@
                 return this::setRecentsTransitionDuringKeyguard;
             } else if (mDesktopTasksController != null
                     // Check on the default display. Recents/gesture nav is only available there
-                    && mDesktopTasksController.visibleTaskCount(DEFAULT_DISPLAY) > 0) {
+                    && mDesktopTasksController.isAnyDeskActive(DEFAULT_DISPLAY)) {
                 return this::setRecentsTransitionDuringDesktop;
             }
         }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java
index 75ad621..e85d30f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/common/pip/PipDesktopStateTest.java
@@ -124,7 +124,7 @@
 
     @Test
     public void isPipEnteringInDesktopMode_visibleCountZero_minimizedPipPresent_returnsTrue() {
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(false);
         when(mMockDesktopRepository.isMinimizedPipPresentInDisplay(DISPLAY_ID)).thenReturn(true);
 
         assertTrue(mPipDesktopState.isPipEnteringInDesktopMode(mMockTaskInfo));
@@ -132,7 +132,7 @@
 
     @Test
     public void isPipEnteringInDesktopMode_visibleCountNonzero_minimizedPipAbsent_returnsTrue() {
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(true);
         when(mMockDesktopRepository.isMinimizedPipPresentInDisplay(DISPLAY_ID)).thenReturn(false);
 
         assertTrue(mPipDesktopState.isPipEnteringInDesktopMode(mMockTaskInfo));
@@ -140,7 +140,7 @@
 
     @Test
     public void isPipEnteringInDesktopMode_visibleCountZero_minimizedPipAbsent_returnsFalse() {
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(false);
         when(mMockDesktopRepository.isMinimizedPipPresentInDisplay(DISPLAY_ID)).thenReturn(false);
 
         assertFalse(mPipDesktopState.isPipEnteringInDesktopMode(mMockTaskInfo));
@@ -148,7 +148,7 @@
 
     @Test
     public void shouldExitPipExitDesktopMode_visibleCountZero_wallpaperInvisible_returnsFalse() {
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(false);
         when(mMockDesktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(
                 DISPLAY_ID)).thenReturn(false);
 
@@ -157,7 +157,7 @@
 
     @Test
     public void shouldExitPipExitDesktopMode_visibleCountNonzero_wallpaperVisible_returnsFalse() {
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(true);
         when(mMockDesktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(
                 DISPLAY_ID)).thenReturn(true);
 
@@ -166,7 +166,7 @@
 
     @Test
     public void shouldExitPipExitDesktopMode_visibleCountZero_wallpaperVisible_returnsTrue() {
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(false);
         when(mMockDesktopWallpaperActivityTokenProvider.isWallpaperActivityVisible(
                 DISPLAY_ID)).thenReturn(true);
 
@@ -176,7 +176,7 @@
     @Test
     public void getOutPipWindowingMode_exitToDesktop_displayFreeform_returnsUndefined() {
         // Set visible task count to 1 so isPipExitingToDesktopMode returns true
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(true);
         setDisplayWindowingMode(WINDOWING_MODE_FREEFORM);
 
         assertEquals(WINDOWING_MODE_UNDEFINED, mPipDesktopState.getOutPipWindowingMode());
@@ -185,7 +185,7 @@
     @Test
     public void getOutPipWindowingMode_exitToDesktop_displayFullscreen_returnsFreeform() {
         // Set visible task count to 1 so isPipExitingToDesktopMode returns true
-        when(mMockDesktopRepository.getVisibleTaskCount(DISPLAY_ID)).thenReturn(1);
+        when(mMockDesktopRepository.isAnyDeskActive(DISPLAY_ID)).thenReturn(true);
         setDisplayWindowingMode(WINDOWING_MODE_FULLSCREEN);
 
         assertEquals(WINDOWING_MODE_FREEFORM, mPipDesktopState.getOutPipWindowingMode());
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
index 2264ade..598a101 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/compatui/CompatUIControllerTest.java
@@ -67,8 +67,6 @@
 
 import dagger.Lazy;
 
-import java.util.Optional;
-
 import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
@@ -78,6 +76,8 @@
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
 
+import java.util.Optional;
+
 /**
  * Tests for {@link CompatUIController}.
  *
@@ -707,13 +707,13 @@
     @EnableFlags(Flags.FLAG_SKIP_COMPAT_UI_EDUCATION_IN_DESKTOP_MODE)
     public void testUpdateActiveTaskInfo_removeAllComponentWhenInDesktopModeFlagEnabled() {
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true);
-        when(mDesktopUserRepositories.getCurrent().getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mDesktopUserRepositories.getCurrent().isAnyDeskActive(DISPLAY_ID)).thenReturn(false);
 
         mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mController, never()).removeLayouts(taskInfo.taskId);
 
-        when(mDesktopUserRepositories.getCurrent().getVisibleTaskCount(DISPLAY_ID)).thenReturn(2);
+        when(mDesktopUserRepositories.getCurrent().isAnyDeskActive(DISPLAY_ID)).thenReturn(true);
 
         mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
@@ -724,14 +724,14 @@
     @RequiresFlagsDisabled(Flags.FLAG_APP_COMPAT_UI_FRAMEWORK)
     @DisableFlags(Flags.FLAG_SKIP_COMPAT_UI_EDUCATION_IN_DESKTOP_MODE)
     public void testUpdateActiveTaskInfo_removeAllComponentWhenInDesktopModeFlagDisabled() {
-        when(mDesktopUserRepositories.getCurrent().getVisibleTaskCount(DISPLAY_ID)).thenReturn(0);
+        when(mDesktopUserRepositories.getCurrent().isAnyDeskActive(DISPLAY_ID)).thenReturn(false);
         TaskInfo taskInfo = createTaskInfo(DISPLAY_ID, TASK_ID, /* hasSizeCompat= */ true);
 
         mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
         verify(mController, never()).removeLayouts(taskInfo.taskId);
 
-        when(mDesktopUserRepositories.getCurrent().getVisibleTaskCount(DISPLAY_ID)).thenReturn(2);
+        when(mDesktopUserRepositories.getCurrent().isAnyDeskActive(DISPLAY_ID)).thenReturn(true);
 
         mController.onCompatInfoChanged(new CompatUIInfo(taskInfo, mMockTaskListener));
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index 253740c..1e2fc57 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
@@ -1062,34 +1062,40 @@
     }
 
     @Test
-    fun visibleTaskCount_noTasks_returnsZero() {
-        assertThat(controller.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(0)
+    fun isAnyDeskActive_noTasks_returnsFalse() {
+        assertThat(controller.isAnyDeskActive(DEFAULT_DISPLAY)).isFalse()
     }
 
     @Test
-    fun visibleTaskCount_twoTasks_bothVisible_returnsTwo() {
+    fun isAnyDeskActive_twoTasks_bothVisible_returnsTrue() {
         setUpHomeTask()
+
         setUpFreeformTask().also(::markTaskVisible)
         setUpFreeformTask().also(::markTaskVisible)
-        assertThat(controller.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(2)
+
+        assertThat(controller.isAnyDeskActive(DEFAULT_DISPLAY)).isTrue()
     }
 
     @Test
-    fun visibleTaskCount_twoTasks_oneVisible_returnsOne() {
+    fun isInDesktop_twoTasks_oneVisible_returnsTrue() {
         setUpHomeTask()
+
         setUpFreeformTask().also(::markTaskVisible)
         setUpFreeformTask().also(::markTaskHidden)
-        assertThat(controller.visibleTaskCount(DEFAULT_DISPLAY)).isEqualTo(1)
+
+        assertThat(controller.isAnyDeskActive(DEFAULT_DISPLAY)).isTrue()
     }
 
     @Test
-    fun visibleTaskCount_twoTasksVisibleOnDifferentDisplays_returnsOne() {
+    fun isAnyDeskActive_twoTasksVisibleOnDifferentDisplays_returnsTrue() {
         taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
         taskRepository.setActiveDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
         setUpHomeTask()
+
         setUpFreeformTask(DEFAULT_DISPLAY).also(::markTaskVisible)
         setUpFreeformTask(SECOND_DISPLAY).also(::markTaskVisible)
-        assertThat(controller.visibleTaskCount(SECOND_DISPLAY)).isEqualTo(1)
+
+        assertThat(controller.isAnyDeskActive(SECOND_DISPLAY)).isTrue()
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
index dd577f4..5b0f94f 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksTransitionObserverTest.kt
@@ -128,7 +128,7 @@
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
     fun backNavigation_taskMinimized() {
         val task = createTaskInfo(1)
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
 
         transitionObserver.onTransitionReady(
             transition = mock(),
@@ -146,7 +146,9 @@
     fun backNavigation_withCloseTransitionNotLastTask_taskMinimized() {
         val task = createTaskInfo(1)
         val transition = mock<IBinder>()
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(2)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
+        whenever(taskRepository.isOnlyVisibleTask(task.taskId, task.displayId)).thenReturn(false)
+        whenever(taskRepository.hasOnlyOneVisibleTask(task.displayId)).thenReturn(false)
         whenever(taskRepository.isClosingTask(task.taskId)).thenReturn(false)
         whenever(backAnimationController.latestTriggerBackTask).thenReturn(task.taskId)
 
@@ -173,7 +175,7 @@
     fun backNavigation_withCloseTransitionLastTask_wallpaperActivityClosed_taskMinimized() {
         val task = createTaskInfo(1)
         val transition = mock<IBinder>()
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
         whenever(taskRepository.isClosingTask(task.taskId)).thenReturn(false)
         whenever(backAnimationController.latestTriggerBackTask).thenReturn(task.taskId)
 
@@ -202,7 +204,7 @@
     fun backNavigation_withCloseTransitionLastTask_wallpaperActivityReordered_taskMinimized() {
         val task = createTaskInfo(1)
         val transition = mock<IBinder>()
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
         whenever(taskRepository.isClosingTask(task.taskId)).thenReturn(false)
         whenever(backAnimationController.latestTriggerBackTask).thenReturn(task.taskId)
 
@@ -227,7 +229,7 @@
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
     fun backNavigation_nullTaskInfo_taskNotMinimized() {
         val task = createTaskInfo(1)
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
 
         transitionObserver.onTransitionReady(
             transition = mock(),
@@ -243,7 +245,7 @@
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
     fun removeTasks_onTaskFullscreenLaunchWithOpenTransition_taskRemovedFromRepo() {
         val task = createTaskInfo(1, WINDOWING_MODE_FULLSCREEN)
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
         whenever(taskRepository.isActiveTask(task.taskId)).thenReturn(true)
 
         transitionObserver.onTransitionReady(
@@ -261,7 +263,7 @@
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
     fun removeTasks_onTaskFullscreenLaunchExitDesktopTransition_taskRemovedFromRepo() {
         val task = createTaskInfo(1, WINDOWING_MODE_FULLSCREEN)
-        whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
+        whenever(taskRepository.isAnyDeskActive(any())).thenReturn(true)
         whenever(taskRepository.isActiveTask(task.taskId)).thenReturn(true)
 
         transitionObserver.onTransitionReady(
@@ -280,7 +282,7 @@
     fun closeLastTask_wallpaperTokenExists_wallpaperIsRemoved() {
         val mockTransition = Mockito.mock(IBinder::class.java)
         val task = createTaskInfo(1, WINDOWING_MODE_FREEFORM)
-        whenever(taskRepository.getVisibleTaskCount(task.displayId)).thenReturn(0)
+        whenever(taskRepository.isAnyDeskActive(task.displayId)).thenReturn(false)
 
         transitionObserver.onTransitionReady(
             transition = mockTransition,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandlerTest.kt
index 143d232..7560945 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/compatui/SystemModalsTransitionHandlerTest.kt
@@ -84,7 +84,7 @@
         spyContext = spy(mContext)
         // Simulate having one Desktop task so that we see Desktop Mode as active
         whenever(desktopUserRepositories.current).thenReturn(desktopRepository)
-        whenever(desktopRepository.getVisibleTaskCount(anyInt())).thenReturn(1)
+        whenever(desktopRepository.isAnyDeskActive(anyInt())).thenReturn(true)
         whenever(spyContext.packageManager).thenReturn(packageManager)
         whenever(componentName.packageName).thenReturn(HOME_LAUNCHER_PACKAGE_NAME)
         whenever(packageManager.getHomeActivities(ArrayList())).thenReturn(componentName)
@@ -111,7 +111,7 @@
 
     @Test
     fun startAnimation_desktopNotActive_doesNotAnimate() {
-        whenever(desktopUserRepositories.current.getVisibleTaskCount(anyInt())).thenReturn(1)
+        whenever(desktopUserRepositories.current.isAnyDeskActive(anyInt())).thenReturn(true)
         val info =
             TransitionInfoBuilder(TRANSIT_OPEN)
                 .addChange(TRANSIT_OPEN, createSystemModalTaskWithBaseActivity())