Fix empty desktop mode after closing the last window (part 2)

- Changed `isOnlyActiveNonClosingTask` to `isOnlyVisibleNonClosingTask`, because back navigation doesn't remove task from active and doesn't call `onTaskVanished`
- Removed closing tasks removal from the `onTaskVanished` as it's handled in `onTaskInfoChanged`
- Changed `handleBackNavigation` to `handleTaskClosing`, because app can decide to close on back navigation and won't get handled

Bug: 343726497
Test: atest WMShellUnitTests:DesktopTasksControllerTest WMShellUnitTests:DesktopModeTaskRepositoryTest
Flag: com.android.window.flags.enable_desktop_windowing_wallpaper_activity
Change-Id: Ia4a8d8f1dbaf25536cb470a4501e899c61845f69

6 files changed, 239 insertions, 99 deletions

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 9bf244e1af1c..81891ce91e04 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
@@ -47,7 +47,7 @@ class DesktopModeTaskRepository {
         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
+        // 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(),
@@ -234,17 +234,16 @@ class DesktopModeTaskRepository {
     }
 
     /**
-     * Check if a task with the given [taskId] is the only active, non-closing, not-minimized task
+     * Check if a task with the given [taskId] is the only visible, non-closing, not-minimized task
      * on its display
      */
-    fun isOnlyActiveNonClosingTask(taskId: Int): Boolean {
-        return displayData.valueIterator().asSequence().any { data ->
-            data.activeTasks
+    fun isOnlyVisibleNonClosingTask(taskId: Int): Boolean =
+        displayData.valueIterator().asSequence().any { data ->
+            data.visibleTasks
                 .subtract(data.closingTasks)
                 .subtract(data.minimizedTasks)
                 .singleOrNull() == taskId
         }
-    }
 
     /** Get a set of the active tasks for given [displayId] */
     fun getActiveTasks(displayId: Int): ArraySet<Int> {
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 d28cda33a569..14ae3a7b9b27 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
@@ -40,7 +40,6 @@ import android.view.SurfaceControl
 import android.view.WindowManager.TRANSIT_CHANGE
 import android.view.WindowManager.TRANSIT_NONE
 import android.view.WindowManager.TRANSIT_OPEN
-import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.RemoteTransition
 import android.window.TransitionInfo
@@ -76,6 +75,7 @@ import com.android.wm.shell.recents.RecentsTransitionStateListener
 import com.android.wm.shell.shared.DesktopModeStatus
 import com.android.wm.shell.shared.DesktopModeStatus.DESKTOP_DENSITY_OVERRIDE
 import com.android.wm.shell.shared.DesktopModeStatus.useDesktopOverrideDensity
+import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.annotations.ExternalThread
 import com.android.wm.shell.shared.annotations.ShellMainThread
 import com.android.wm.shell.splitscreen.SplitScreenController
@@ -446,7 +446,7 @@ class DesktopTasksController(
      * @param taskId task id of the window that's being closed
      */
     fun onDesktopWindowClose(wct: WindowContainerTransaction, displayId: Int, taskId: Int) {
-        if (desktopModeTaskRepository.isOnlyActiveNonClosingTask(taskId)) {
+        if (desktopModeTaskRepository.isOnlyVisibleNonClosingTask(taskId)) {
             removeWallpaperActivity(wct)
         }
         if (!desktopModeTaskRepository.addClosingTask(displayId, taskId)) {
@@ -879,8 +879,8 @@ class DesktopTasksController(
                     reason = "recents animation is running"
                     false
                 }
-                // Handle back navigation for the last window if wallpaper available
-                shouldHandleBackNavigation(request) -> true
+                // Handle task closing for the last window if wallpaper is available
+                shouldHandleTaskClosing(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})"
@@ -918,7 +918,8 @@ class DesktopTasksController(
         val result =
             triggerTask?.let { task ->
                 when {
-                    request.type == TRANSIT_TO_BACK -> handleBackNavigation(task)
+                    // Check if the closing task needs to be handled
+                    TransitionUtil.isClosingType(request.type) -> handleTaskClosing(task)
                     // Check if the task has a top transparent activity
                     shouldLaunchAsModal(task) -> handleIncompatibleTaskLaunch(task)
                     // Check if the task has a top systemUI activity
@@ -960,9 +961,10 @@ class DesktopTasksController(
     private fun shouldLaunchAsModal(task: TaskInfo) =
         Flags.enableDesktopWindowingModalsPolicy() && isSingleTopActivityTranslucent(task)
 
-    private fun shouldHandleBackNavigation(request: TransitionRequestInfo): Boolean {
+    private fun shouldHandleTaskClosing(request: TransitionRequestInfo): Boolean {
         return Flags.enableDesktopWindowingWallpaperActivity() &&
-            request.type == TRANSIT_TO_BACK
+            TransitionUtil.isClosingType(request.type) &&
+            request.triggerTask != null
     }
 
     private fun handleFreeformTaskLaunch(
@@ -1029,10 +1031,10 @@ class DesktopTasksController(
         return WindowContainerTransaction().also { wct -> addMoveToFullscreenChanges(wct, task) }
     }
 
-    /** Handle back navigation by removing wallpaper activity if it's the last active task */
-    private fun handleBackNavigation(task: RunningTaskInfo): WindowContainerTransaction? {
+    /** Handle task closing by removing wallpaper activity if it's the last active task */
+    private fun handleTaskClosing(task: RunningTaskInfo): WindowContainerTransaction? {
         val wct = if (
-            desktopModeTaskRepository.isOnlyActiveNonClosingTask(task.taskId) &&
+            desktopModeTaskRepository.isOnlyVisibleNonClosingTask(task.taskId) &&
                 desktopModeTaskRepository.wallpaperActivityToken != null
         ) {
             // Remove wallpaper activity when the last active task is removed
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 b88afb9e7197..b48aee5ccd5e 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,10 +122,6 @@ public class FreeformTaskListener implements ShellTaskOrganizer.TaskListener,
             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);
diff --git a/libs/WindowManager/Shell/tests/unittest/Android.bp b/libs/WindowManager/Shell/tests/unittest/Android.bp
index 13f95ccea640..92be4f9f0374 100644
--- a/libs/WindowManager/Shell/tests/unittest/Android.bp
+++ b/libs/WindowManager/Shell/tests/unittest/Android.bp
@@ -46,6 +46,7 @@ android_test {
         "androidx.dynamicanimation_dynamicanimation",
         "dagger2",
         "frameworks-base-testutils",
+        "kotlin-test",
         "kotlinx-coroutines-android",
         "kotlinx-coroutines-core",
         "mockito-kotlin2",
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 193d61437b2b..6612aee0cd12 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,86 +119,91 @@ class DesktopModeTaskRepositoryTest : ShellTestCase() {
     }
 
     @Test
-    fun isOnlyActiveNonClosingTask_noActiveNonClosingTasks() {
-        // Not an active task
-        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
+    fun isOnlyVisibleNonClosingTask_noTasks() {
+        // No visible tasks
+        assertThat(repo.isOnlyVisibleNonClosingTask(1)).isFalse()
         assertThat(repo.isClosingTask(1)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveNonClosingTask_singleActiveNonClosingTask() {
-        repo.addActiveTask(DEFAULT_DISPLAY, 1)
-        // The only active task
-        assertThat(repo.isActiveTask(1)).isTrue()
+    fun isOnlyVisibleNonClosingTask_singleVisibleNonClosingTask() {
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
+
+        // The only visible task
+        assertThat(repo.isVisibleTask(1)).isTrue()
         assertThat(repo.isClosingTask(1)).isFalse()
-        assertThat(repo.isOnlyActiveNonClosingTask(1)).isTrue()
-        // Not an active task
-        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(1)).isTrue()
+        // Not a visible task
+        assertThat(repo.isVisibleTask(99)).isFalse()
         assertThat(repo.isClosingTask(99)).isFalse()
-        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveNonClosingTask_singleActiveClosingTask() {
-        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+    fun isOnlyVisibleNonClosingTask_singleVisibleClosingTask() {
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
         repo.addClosingTask(DEFAULT_DISPLAY, 1)
-        // The active task that's closing
-        assertThat(repo.isActiveTask(1)).isTrue()
+
+        // A visible task that's closing
+        assertThat(repo.isVisibleTask(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()
+        assertThat(repo.isOnlyVisibleNonClosingTask(1)).isFalse()
+        // Not a visible task
+        assertThat(repo.isVisibleTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveNonClosingTask_singleActiveMinimizedTask() {
-        repo.addActiveTask(DEFAULT_DISPLAY, 1)
+    fun isOnlyVisibleNonClosingTask_singleVisibleMinimizedTask() {
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
         repo.minimizeTask(DEFAULT_DISPLAY, 1)
-        // The active task that's closing
-        assertThat(repo.isActiveTask(1)).isTrue()
+
+        // The visible task that's closing
+        assertThat(repo.isVisibleTask(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()
+        assertThat(repo.isOnlyVisibleNonClosingTask(1)).isFalse()
+        // Not a visible task
+        assertThat(repo.isVisibleTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveNonClosingTask_multipleActiveNonClosingTasks() {
-        repo.addActiveTask(DEFAULT_DISPLAY, 1)
-        repo.addActiveTask(DEFAULT_DISPLAY, 2)
+    fun isOnlyVisibleNonClosingTask_multipleVisibleNonClosingTasks() {
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 2, visible = true)
+
         // Not the only task
-        assertThat(repo.isActiveTask(1)).isTrue()
+        assertThat(repo.isVisibleTask(1)).isTrue()
         assertThat(repo.isClosingTask(1)).isFalse()
-        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(1)).isFalse()
         // Not the only task
-        assertThat(repo.isActiveTask(2)).isTrue()
+        assertThat(repo.isVisibleTask(2)).isTrue()
         assertThat(repo.isClosingTask(2)).isFalse()
-        assertThat(repo.isOnlyActiveNonClosingTask(2)).isFalse()
-        // Not an active task
-        assertThat(repo.isActiveTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(2)).isFalse()
+        // Not a visible task
+        assertThat(repo.isVisibleTask(99)).isFalse()
         assertThat(repo.isClosingTask(99)).isFalse()
-        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(99)).isFalse()
     }
 
     @Test
-    fun isOnlyActiveNonClosingTask_multipleDisplays() {
-        repo.addActiveTask(DEFAULT_DISPLAY, 1)
-        repo.addActiveTask(DEFAULT_DISPLAY, 2)
-        repo.addActiveTask(SECOND_DISPLAY, 3)
+    fun isOnlyVisibleNonClosingTask_multipleDisplays() {
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 1, visible = true)
+        repo.updateVisibleFreeformTasks(DEFAULT_DISPLAY, taskId = 2, visible = true)
+        repo.updateVisibleFreeformTasks(SECOND_DISPLAY, taskId = 3, visible = true)
+
         // Not the only task on DEFAULT_DISPLAY
-        assertThat(repo.isActiveTask(1)).isTrue()
-        assertThat(repo.isOnlyActiveNonClosingTask(1)).isFalse()
+        assertThat(repo.isVisibleTask(1)).isTrue()
+        assertThat(repo.isOnlyVisibleNonClosingTask(1)).isFalse()
         // Not the only task on DEFAULT_DISPLAY
-        assertThat(repo.isActiveTask(2)).isTrue()
-        assertThat(repo.isOnlyActiveNonClosingTask(2)).isFalse()
-        // The only active task on SECOND_DISPLAY
-        assertThat(repo.isActiveTask(3)).isTrue()
-        assertThat(repo.isOnlyActiveNonClosingTask(3)).isTrue()
-        // Not an active task
-        assertThat(repo.isActiveTask(99)).isFalse()
-        assertThat(repo.isOnlyActiveNonClosingTask(99)).isFalse()
+        assertThat(repo.isVisibleTask(2)).isTrue()
+        assertThat(repo.isOnlyVisibleNonClosingTask(2)).isFalse()
+        // The only visible task on SECOND_DISPLAY
+        assertThat(repo.isVisibleTask(3)).isTrue()
+        assertThat(repo.isOnlyVisibleNonClosingTask(3)).isTrue()
+        // Not a visible task
+        assertThat(repo.isVisibleTask(99)).isFalse()
+        assertThat(repo.isOnlyVisibleNonClosingTask(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 392161fd8e66..a1a18a9b7c9d 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
@@ -45,6 +45,7 @@ import android.view.Display.DEFAULT_DISPLAY
 import android.view.SurfaceControl
 import android.view.WindowManager
 import android.view.WindowManager.TRANSIT_CHANGE
+import android.view.WindowManager.TRANSIT_CLOSE
 import android.view.WindowManager.TRANSIT_OPEN
 import android.view.WindowManager.TRANSIT_TO_BACK
 import android.view.WindowManager.TRANSIT_TO_FRONT
@@ -102,6 +103,8 @@ import com.google.common.truth.Truth.assertWithMessage
 import java.util.Optional
 import junit.framework.Assert.assertFalse
 import junit.framework.Assert.assertTrue
+import kotlin.test.assertNotNull
+import kotlin.test.assertNull
 import org.junit.After
 import org.junit.Assume.assumeTrue
 import org.junit.Before
@@ -1254,72 +1257,205 @@ class DesktopTasksControllerTest : ShellTestCase() {
   }
 
   @Test
+  @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_backTransition_singleActiveTaskNoTokenFlagDisabled_doesNotHandle() {
+    val task = setUpFreeformTask()
+
+    val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
+
+    assertNull(result, "Should not handle request")
+  }
+
+  @Test
   @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-  fun handleRequest_backTransition_singleActiveTask_noToken() {
+  fun handleRequest_backTransition_singleActiveTaskNoTokenFlagEnabled_doesNotHandle() {
     val task = setUpFreeformTask()
+
     val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
-    // Doesn't handle request
-    assertThat(result).isNull()
+
+    assertNull(result, "Should not handle request")
   }
 
   @Test
   @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-  fun handleRequest_backTransition_singleActiveTask_hasToken_desktopWallpaperDisabled() {
-    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
-
+  fun handleRequest_backTransition_singleActiveTaskWithTokenFlagDisabled_doesNotHandle() {
     val task = setUpFreeformTask()
+
+    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
     val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
-    // Doesn't handle request
-    assertThat(result).isNull()
+
+    assertNull(result, "Should not handle request")
   }
 
   @Test
   @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-  fun handleRequest_backTransition_singleActiveTask_hasToken_desktopWallpaperEnabled() {
+  fun handleRequest_backTransition_singleActiveTaskWithTokenFlagEnabled_handlesRequest() {
+    val task = setUpFreeformTask()
     val wallpaperToken = MockToken().token()
-    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
 
-    val task = setUpFreeformTask()
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
     val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_TO_BACK))
-    assertThat(result).isNotNull()
-    // Creates remove wallpaper transaction
-    result!!.assertRemoveAt(index = 0, wallpaperToken)
+
+    assertNotNull(result, "Should handle request")
+      // Should create remove wallpaper transaction
+      .assertRemoveAt(index = 0, wallpaperToken)
   }
 
   @Test
-  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-  fun handleRequest_backTransition_multipleActiveTasks() {
+  @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_backTransition_multipleActiveTasksFlagDisabled_doesNotHandle() {
+    val task1 = setUpFreeformTask()
+    setUpFreeformTask()
+
     desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
 
+    assertNull(result, "Should not handle request")
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_backTransition_multipleActiveTasksFlagEnabled_doesNotHandle() {
     val task1 = setUpFreeformTask()
     setUpFreeformTask()
+
+    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
     val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
-    // Doesn't handle request
-    assertThat(result).isNull()
+
+    assertNull(result, "Should not handle request")
   }
 
   @Test
   @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-  fun handleRequest_backTransition_multipleActiveTasks_singleNonClosing() {
-    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+  fun handleRequest_backTransition_multipleActiveTasksSingleNonClosing_handlesRequest() {
+    val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val wallpaperToken = MockToken().token()
 
-    val task1 = setUpFreeformTask()
-    setUpFreeformTask()
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.addClosingTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
     val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
-    // Doesn't handle request
-    assertThat(result).isNull()
+
+    assertNotNull(result, "Should handle request")
+      // Should create remove wallpaper transaction
+      .assertRemoveAt(index = 0, wallpaperToken)
   }
 
   @Test
   @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-  fun handleRequest_backTransition_multipleActiveTasks_singleNonMinimized() {
+  fun handleRequest_backTransition_multipleActiveTasksSingleNonMinimized_handlesRequest() {
+    val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val wallpaperToken = MockToken().token()
+
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
+
+    assertNotNull(result, "Should handle request")
+      // Should create remove wallpaper transaction
+      .assertRemoveAt(index = 0, wallpaperToken)
+  }
+
+  @Test
+  @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_singleActiveTaskNoTokenFlagDisabled_doesNotHandle() {
+    val task = setUpFreeformTask()
+
+    val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
+
+    assertNull(result, "Should not handle request")
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_singleActiveTaskNoTokenFlagEnabled_doesNotHandle() {
+    val task = setUpFreeformTask()
+
+    val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
+
+    assertNull(result, "Should not handle request")
+  }
+
+  @Test
+  @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_singleActiveTaskWithTokenFlagDisabled_doesNotHandle() {
+    val task = setUpFreeformTask()
+
     desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+    val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
+
+    assertNull(result, "Should not handle request")
+  }
 
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_singleActiveTaskWithTokenFlagEnabled_handlesRequest() {
+    val task = setUpFreeformTask()
+    val wallpaperToken = MockToken().token()
+
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    val result = controller.handleRequest(Binder(), createTransition(task, type = TRANSIT_CLOSE))
+
+    assertNotNull(result, "Should handle request")
+      // Should create remove wallpaper transaction
+      .assertRemoveAt(index = 0, wallpaperToken)
+  }
+
+  @Test
+  @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_multipleActiveTasksFlagDisabled_doesNotHandle() {
     val task1 = setUpFreeformTask()
     setUpFreeformTask()
-    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_TO_BACK))
-    // Doesn't handle request
-    assertThat(result).isNull()
+
+    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
+
+    assertNull(result, "Should not handle request")
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_multipleActiveTasksFlagEnabled_doesNotHandle() {
+    val task1 = setUpFreeformTask()
+    setUpFreeformTask()
+
+    desktopModeTaskRepository.wallpaperActivityToken = MockToken().token()
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
+
+    assertNull(result, "Should not handle request")
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_multipleActiveTasksSingleNonClosing_handlesRequest() {
+    val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val wallpaperToken = MockToken().token()
+
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.addClosingTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
+
+    assertNotNull(result, "Should handle request")
+      // Should create remove wallpaper transaction
+      .assertRemoveAt(index = 0, wallpaperToken)
+  }
+
+  @Test
+  @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+  fun handleRequest_closeTransition_multipleActiveTasksSingleNonMinimized_handlesRequest() {
+    val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val task2 = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+    val wallpaperToken = MockToken().token()
+
+    desktopModeTaskRepository.wallpaperActivityToken = wallpaperToken
+    desktopModeTaskRepository.minimizeTask(displayId = DEFAULT_DISPLAY, taskId = task2.taskId)
+    val result = controller.handleRequest(Binder(), createTransition(task1, type = TRANSIT_CLOSE))
+
+    assertNotNull(result, "Should handle request")
+      // Should create remove wallpaper transaction
+      .assertRemoveAt(index = 0, wallpaperToken)
   }
 
   @Test
@@ -1693,6 +1829,7 @@ class DesktopTasksControllerTest : ShellTestCase() {
     task.topActivityInfo = activityInfo
     whenever(shellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
     desktopModeTaskRepository.addActiveTask(displayId, task.taskId)
+    desktopModeTaskRepository.updateVisibleFreeformTasks(displayId, task.taskId, visible = true)
     desktopModeTaskRepository.addOrMoveFreeformTaskToTop(displayId, task.taskId)
     runningTasks.add(task)
     return task