Merge "Fix empty desktop mode after closing the last window" into main
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
index 7d01580..9bf244e 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepository.kt
@@ -46,6 +46,9 @@
         val activeTasks: ArraySet<Int> = ArraySet(),
         val visibleTasks: ArraySet<Int> = ArraySet(),
         val minimizedTasks: ArraySet<Int> = ArraySet(),
+        // Tasks that are closing, but are still visible
+        //  TODO(b/332682201): Remove when the repository state is updated via TransitionObserver
+        val closingTasks: ArraySet<Int> = ArraySet(),
         // Tasks currently in freeform mode, ordered from top to bottom (top is at index 0).
         val freeformTasksInZOrder: ArrayList<Int> = ArrayList(),
     )
@@ -169,6 +172,42 @@
         return result
     }
 
+    /**
+     * Mark a task with given [taskId] as closing on given [displayId]
+     *
+     * @return `true` if the task was not closing on given [displayId]
+     */
+    fun addClosingTask(displayId: Int, taskId: Int): Boolean {
+        val added = displayData.getOrCreate(displayId).closingTasks.add(taskId)
+        if (added) {
+            KtProtoLog.d(
+                WM_SHELL_DESKTOP_MODE,
+                "DesktopTaskRepo: added closing task=%d displayId=%d",
+                taskId,
+                displayId
+            )
+        }
+        return added
+    }
+
+    /**
+     * Remove task with given [taskId] from closing tasks.
+     *
+     * @return `true` if the task was closing
+     */
+    fun removeClosingTask(taskId: Int): Boolean {
+        var removed = false
+        displayData.forEach { _, data ->
+            if (data.closingTasks.remove(taskId)) {
+                removed = true
+            }
+        }
+        if (removed) {
+            KtProtoLog.d(WM_SHELL_DESKTOP_MODE, "DesktopTaskRepo: remove closing task=%d", taskId)
+        }
+        return removed
+    }
+
     /** Check if a task with the given [taskId] was marked as an active task */
     fun isActiveTask(taskId: Int): Boolean {
         return displayData.valueIterator().asSequence().any { data ->
@@ -176,6 +215,10 @@
         }
     }
 
+    /** Check if a task with the given [taskId] was marked as a closing task */
+    fun isClosingTask(taskId: Int): Boolean =
+        displayData.valueIterator().asSequence().any { data -> taskId in data.closingTasks }
+
     /** Whether a task is visible. */
     fun isVisibleTask(taskId: Int): Boolean {
         return displayData.valueIterator().asSequence().any { data ->
@@ -190,10 +233,16 @@
         }
     }
 
-    /** Check if a task with the given [taskId] is the only active task on its display */
-    fun isOnlyActiveTask(taskId: Int): Boolean {
+    /**
+     * Check if a task with the given [taskId] is the only active, non-closing, not-minimized task
+     * on its display
+     */
+    fun isOnlyActiveNonClosingTask(taskId: Int): Boolean {
         return displayData.valueIterator().asSequence().any { data ->
-            data.activeTasks.singleOrNull() == taskId
+            data.activeTasks
+                .subtract(data.closingTasks)
+                .subtract(data.minimizedTasks)
+                .singleOrNull() == taskId
         }
     }
 
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 c5111d6..d28cda3 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
@@ -442,12 +442,21 @@
      * active task.
      *
      * @param wct transaction to modify if the last active task is closed
+     * @param displayId display id of the window that's being closed
      * @param taskId task id of the window that's being closed
      */
-    fun onDesktopWindowClose(wct: WindowContainerTransaction, taskId: Int) {
-        if (desktopModeTaskRepository.isOnlyActiveTask(taskId)) {
+    fun onDesktopWindowClose(wct: WindowContainerTransaction, displayId: Int, taskId: Int) {
+        if (desktopModeTaskRepository.isOnlyActiveNonClosingTask(taskId)) {
             removeWallpaperActivity(wct)
         }
+        if (!desktopModeTaskRepository.addClosingTask(displayId, taskId)) {
+            // Could happen if the task hasn't been removed from closing list after it disappeared
+            KtProtoLog.w(
+                WM_SHELL_DESKTOP_MODE,
+                "DesktopTasksController: the task with taskId=%d is already closing!",
+                taskId
+            )
+        }
     }
 
     /** Move a task with given `taskId` to fullscreen */
@@ -871,7 +880,7 @@
                     false
                 }
                 // Handle back navigation for the last window if wallpaper available
-                shouldRemoveWallpaper(request) -> true
+                shouldHandleBackNavigation(request) -> true
                 // Only handle open or to front transitions
                 request.type != TRANSIT_OPEN && request.type != TRANSIT_TO_FRONT -> {
                     reason = "transition type not handled (${request.type})"
@@ -951,13 +960,9 @@
     private fun shouldLaunchAsModal(task: TaskInfo) =
         Flags.enableDesktopWindowingModalsPolicy() && isSingleTopActivityTranslucent(task)
 
-    private fun shouldRemoveWallpaper(request: TransitionRequestInfo): Boolean {
+    private fun shouldHandleBackNavigation(request: TransitionRequestInfo): Boolean {
         return Flags.enableDesktopWindowingWallpaperActivity() &&
-            request.type == TRANSIT_TO_BACK &&
-            request.triggerTask?.let { task ->
-                desktopModeTaskRepository.isOnlyActiveTask(task.taskId)
-            }
-                ?: false
+            request.type == TRANSIT_TO_BACK
     }
 
     private fun handleFreeformTaskLaunch(
@@ -1026,15 +1031,24 @@
 
     /** Handle back navigation by removing wallpaper activity if it's the last active task */
     private fun handleBackNavigation(task: RunningTaskInfo): WindowContainerTransaction? {
-        if (
-            desktopModeTaskRepository.isOnlyActiveTask(task.taskId) &&
+        val wct = if (
+            desktopModeTaskRepository.isOnlyActiveNonClosingTask(task.taskId) &&
                 desktopModeTaskRepository.wallpaperActivityToken != null
         ) {
             // Remove wallpaper activity when the last active task is removed
-            return WindowContainerTransaction().also { wct -> removeWallpaperActivity(wct) }
+            WindowContainerTransaction().also { wct -> removeWallpaperActivity(wct) }
         } else {
-            return null
+            null
         }
+        if (!desktopModeTaskRepository.addClosingTask(task.displayId, task.taskId)) {
+            // Could happen if the task hasn't been removed from closing list after it disappeared
+            KtProtoLog.w(
+                WM_SHELL_DESKTOP_MODE,
+                "DesktopTasksController: the task with taskId=%d is already closing!",
+                task.taskId
+            )
+        }
+        return wct
     }
 
     private fun addMoveToDesktopChanges(
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
index 7d2aa27..b88afb9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskListener.java
@@ -122,6 +122,10 @@
             mDesktopModeTaskRepository.ifPresent(repository -> {
                 repository.removeFreeformTask(taskInfo.displayId, taskInfo.taskId);
                 repository.unminimizeTask(taskInfo.displayId, taskInfo.taskId);
+                if (repository.removeClosingTask(taskInfo.taskId)) {
+                    ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
+                            "Removing closing freeform task: #%d", taskInfo.taskId);
+                }
                 if (repository.removeActiveTask(taskInfo.taskId)) {
                     ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                             "Removing active freeform task: #%d", taskInfo.taskId);
@@ -150,6 +154,10 @@
                         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
                                 "Adding active freeform task: #%d", taskInfo.taskId);
                     }
+                } else if (repository.isClosingTask(taskInfo.taskId)
+                        && repository.removeClosingTask(taskInfo.taskId)) {
+                    ProtoLog.v(ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE,
+                            "Removing closing freeform task: #%d", taskInfo.taskId);
                 }
                 repository.updateVisibleFreeformTasks(taskInfo.displayId, taskInfo.taskId,
                         taskInfo.isVisible);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
index 21b6db2..5e7e5e6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModel.java
@@ -435,7 +435,7 @@
                             SplitScreenController.EXIT_REASON_DESKTOP_MODE);
                 } else {
                     WindowContainerTransaction wct = new WindowContainerTransaction();
-                    mDesktopTasksController.onDesktopWindowClose(wct, mTaskId);
+                    mDesktopTasksController.onDesktopWindowClose(wct, mDisplayId, mTaskId);
                     mTaskOperations.closeTask(mTaskToken, wct);
                 }
             } else if (id == R.id.back_button) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
index 310ccc2..193d614 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopModeTaskRepositoryTest.kt
@@ -119,54 +119,86 @@
     }
 
     @Test
-    fun isOnlyActiveTask_noActiveTasks() {
+    fun isOnlyActiveNonClosingTask_noActiveNonClosingTasks() {
         // Not an active task
-        assertThat(repo.isOnlyActiveTask(1)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
+        assertThat(repo.isClosingTask(1)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveTask_singleActiveTask() {
+    fun isOnlyActiveNonClosingTask_singleActiveNonClosingTask() {
         repo.addActiveTask(DEFAULT_DISPLAY, 1)
         // The only active task
         assertThat(repo.isActiveTask(1)).isTrue()
-        assertThat(repo.isOnlyActiveTask(1)).isTrue()
+        assertThat(repo.isClosingTask(1)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(1)).isTrue()
         // Not an active task
         assertThat(repo.isActiveTask(99)).isFalse()
-        assertThat(repo.isOnlyActiveTask(99)).isFalse()
+        assertThat(repo.isClosingTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveTask_multipleActiveTasks() {
+    fun isOnlyActiveNonClosingTask_singleActiveClosingTask() {
+        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+        repo.addClosingTask(DEFAULT_DISPLAY, 1)
+        // The active task that's closing
+        assertThat(repo.isActiveTask(1)).isTrue()
+        assertThat(repo.isClosingTask(1)).isTrue()
+        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
+        // Not an active task
+        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
+    }
+
+    @Test
+    fun isOnlyActiveNonClosingTask_singleActiveMinimizedTask() {
+        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+        repo.minimizeTask(DEFAULT_DISPLAY, 1)
+        // The active task that's closing
+        assertThat(repo.isActiveTask(1)).isTrue()
+        assertThat(repo.isMinimizedTask(1)).isTrue()
+        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
+        // Not an active task
+        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
+    }
+
+    @Test
+    fun isOnlyActiveNonClosingTask_multipleActiveNonClosingTasks() {
         repo.addActiveTask(DEFAULT_DISPLAY, 1)
         repo.addActiveTask(DEFAULT_DISPLAY, 2)
         // Not the only task
         assertThat(repo.isActiveTask(1)).isTrue()
-        assertThat(repo.isOnlyActiveTask(1)).isFalse()
+        assertThat(repo.isClosingTask(1)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
         // Not the only task
         assertThat(repo.isActiveTask(2)).isTrue()
-        assertThat(repo.isOnlyActiveTask(2)).isFalse()
+        assertThat(repo.isClosingTask(2)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(2)).isFalse()
         // Not an active task
         assertThat(repo.isActiveTask(99)).isFalse()
-        assertThat(repo.isOnlyActiveTask(99)).isFalse()
+        assertThat(repo.isClosingTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveTask_multipleDisplays() {
+    fun isOnlyActiveNonClosingTask_multipleDisplays() {
         repo.addActiveTask(DEFAULT_DISPLAY, 1)
         repo.addActiveTask(DEFAULT_DISPLAY, 2)
         repo.addActiveTask(SECOND_DISPLAY, 3)
         // Not the only task on DEFAULT_DISPLAY
         assertThat(repo.isActiveTask(1)).isTrue()
-        assertThat(repo.isOnlyActiveTask(1)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
         // Not the only task on DEFAULT_DISPLAY
         assertThat(repo.isActiveTask(2)).isTrue()
-        assertThat(repo.isOnlyActiveTask(2)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(2)).isFalse()
         // The only active task on SECOND_DISPLAY
         assertThat(repo.isActiveTask(3)).isTrue()
-        assertThat(repo.isOnlyActiveTask(3)).isTrue()
+        assertThat(repo.isOnlyActiveNonClosingTask(3)).isTrue()
         // Not an active task
         assertThat(repo.isActiveTask(99)).isFalse()
-        assertThat(repo.isOnlyActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
     }
 
     @Test
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 e17f7f2..392161f 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
@@ -923,7 +923,7 @@
   @Test
   fun onDesktopWindowClose_noActiveTasks() {
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, 1 /* taskId */)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = 1)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -932,7 +932,7 @@
   fun onDesktopWindowClose_singleActiveTask_noWallpaperActivityToken() {
     val task = setUpFreeformTask()
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, task.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -944,12 +944,38 @@
     desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, task.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
     // Adds remove wallpaper operation
     wct.assertRemoveAt(index = 0, wallpaperToken)
   }
 
   @Test
+  fun onDesktopWindowClose_singleActiveTask_isClosing() {
+    val task = setUpFreeformTask()
+    val wallpaperToken = MockToken().token()
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.addClosingTask(DEFAULT_DISPLAY, task.taskId)
+
+    val wct = WindowContainerTransaction()
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
+    // Doesn't modify transaction
+    assertThat(wct.hierarchyOps).isEmpty()
+  }
+
+  @Test
+  fun onDesktopWindowClose_singleActiveTask_isMinimized() {
+    val task = setUpFreeformTask()
+    val wallpaperToken = MockToken().token()
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.minimizeTask(DEFAULT_DISPLAY, task.taskId)
+
+    val wct = WindowContainerTransaction()
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
+    // Doesn't modify transaction
+    assertThat(wct.hierarchyOps).isEmpty()
+  }
+
+  @Test
   fun onDesktopWindowClose_multipleActiveTasks() {
     val task1 = setUpFreeformTask()
     setUpFreeformTask()
@@ -957,12 +983,40 @@
     desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, task1.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task1.taskId)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
 
   @Test
+  fun onDesktopWindowClose_multipleActiveTasks_isOnlyNonClosingTask() {
+    val task1 = setUpFreeformTask()
+    val task2 = setUpFreeformTask()
+    val wallpaperToken = MockToken().token()
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.addClosingTask(DEFAULT_DISPLAY, task2.taskId)
+
+    val wct = WindowContainerTransaction()
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task1.taskId)
+    // Adds remove wallpaper operation
+    wct.assertRemoveAt(index = 0, wallpaperToken)
+  }
+
+  @Test
+  fun onDesktopWindowClose_multipleActiveTasks_hasMinimized() {
+    val task1 = setUpFreeformTask()
+    val task2 = setUpFreeformTask()
+    val wallpaperToken = MockToken().token()
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.minimizeTask(DEFAULT_DISPLAY, task2.taskId)
+
+    val wct = WindowContainerTransaction()
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task1.taskId)
+    // Adds remove wallpaper operation
+    wct.assertRemoveAt(index = 0, wallpaperToken)
+  }
+
+  @Test
   fun handleRequest_fullscreenTask_freeformVisible_returnSwitchToFreeformWCT() {
     assumeTrue(ENABLE_SHELL_TRANSITIONS)
 
@@ -1245,6 +1299,30 @@
   }
 
   @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_backTransition_multipleActiveTasks_singleNonClosing() {
+    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+
+    val task1 = setUpFreeformTask()
+    setUpFreeformTask()
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
+    // Doesn't handle request
+    assertThat(result).isNull()
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_backTransition_multipleActiveTasks_singleNonMinimized() {
+    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+
+    val task1 = setUpFreeformTask()
+    setUpFreeformTask()
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
+    // Doesn't handle request
+    assertThat(result).isNull()
+  }
+
+  @Test
   fun desktopTasksVisibilityChange_visible_setLaunchAdjacentDisabled() {
     val task = setUpFreeformTask()
     clearInvocations(launchAdjacentController)