Merge "Remove tasks from repo on ExitDesktop transitions." into main

2 files changed, 28 insertions, 5 deletions

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 0b1bb8f36fa8..d1534da9a078 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
@@ -29,6 +29,7 @@ import android.window.DesktopModeFlags
 import android.window.DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY
 import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.desktopmode.DesktopModeTransitionTypes.isExitDesktopModeTransition
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.TransitionUtil
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
@@ -79,7 +80,9 @@ class DesktopTasksTransitionObserver(
     private fun removeTaskIfNeeded(info: TransitionInfo) {
         // Since we are no longer removing all the tasks [onTaskVanished], we need to remove them by
         // checking the transitions.
-        if (!TransitionUtil.isOpeningType(info.type)) return
+        if (!(TransitionUtil.isOpeningType(info.type) || info.type.isExitDesktopModeTransition())) {
+            return
+        }
         // Remove a task from the repository if the app is launched outside of desktop.
         for (change in info.changes) {
             val taskInfo = change.taskInfo
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 fe87aa88a8db..737439ce3cfe 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
@@ -40,6 +40,7 @@ import com.android.window.flags.Flags
 import com.android.wm.shell.MockToken
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.desktopmode.DesktopModeTransitionTypes.TRANSIT_EXIT_DESKTOP_MODE_TASK_DRAG
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
@@ -126,14 +127,32 @@ class DesktopTasksTransitionObserverTest {
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION)
-    fun removeTasks_onTaskFullscreenLaunch_taskRemovedFromRepo() {
+    fun removeTasks_onTaskFullscreenLaunchWithOpenTransition_taskRemovedFromRepo() {
         val task = createTaskInfo(1, WINDOWING_MODE_FULLSCREEN)
         whenever(taskRepository.getVisibleTaskCount(any())).thenReturn(1)
         whenever(taskRepository.isActiveTask(task.taskId)).thenReturn(true)
 
         transitionObserver.onTransitionReady(
             transition = mock(),
-            info = createOpenTransition(task),
+            info = createOpenChangeTransition(task),
+            startTransaction = mock(),
+            finishTransaction = mock(),
+        )
+
+        verify(taskRepository, never()).minimizeTask(task.displayId, task.taskId)
+        verify(taskRepository).removeFreeformTask(task.displayId, task.taskId)
+    }
+
+    @Test
+    @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.isActiveTask(task.taskId)).thenReturn(true)
+
+        transitionObserver.onTransitionReady(
+            transition = mock(),
+            info = createOpenChangeTransition(task, TRANSIT_EXIT_DESKTOP_MODE_TASK_DRAG),
             startTransaction = mock(),
             finishTransaction = mock(),
         )
@@ -178,8 +197,9 @@ class DesktopTasksTransitionObserverTest {
         }
     }
 
-    private fun createOpenTransition(
-        task: RunningTaskInfo?
+    private fun createOpenChangeTransition(
+        task: RunningTaskInfo?,
+        type: Int = TRANSIT_OPEN
     ): TransitionInfo {
         return TransitionInfo(TRANSIT_OPEN, 0 /* flags */).apply {
             addChange(