Prevent task from exiting immersive twice on close

When a task is closing, the full immersive handler restores its bounds
and appends the transition to its list of tracked transitions. It also
add a change to prevent direct enter/exit while a pending transition is
in progress, to prevent the case where a onTaskInfoChanged changing the
immersive state of the app re-triggers an enter/exit while another
enter/exit is already in progress.
Finally, it adjusts the close animation to always animate from the start
bounds, since the taskInfo/end bounds can now be different even on close
transitions.

Flag: com.android.window.flags.enable_fully_immersive_in_desktop
Bug: 373987723
Test: enter desktop immersive, close the task using the header icon,
verify no crash and it animates correctly

Change-Id: I068b925ca45bd68d98434e1105bf0a3716247202
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandler.kt
index 320c003..9d4926b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandler.kt
@@ -71,7 +71,7 @@
 
     /** Whether there is an immersive transition that hasn't completed yet. */
     private val inProgress: Boolean
-        get() = state != null
+        get() = state != null || pendingExternalExitTransitions.isNotEmpty()
 
     private val rectEvaluator = RectEvaluator()
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
index 4350199..df9fc59 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandler.kt
@@ -58,9 +58,9 @@
         freeformTaskTransitionHandler.startMinimizedModeTransition(wct)
 
     /** Starts close transition and handles or delegates desktop task close animation. */
-    override fun startRemoveTransition(wct: WindowContainerTransaction?) {
+    override fun startRemoveTransition(wct: WindowContainerTransaction?): IBinder {
         requireNotNull(wct)
-        transitions.startTransition(WindowManager.TRANSIT_CLOSE, wct, /* handler= */ this)
+        return transitions.startTransition(WindowManager.TRANSIT_CLOSE, wct, /* handler= */ this)
     }
 
     /** Returns null, as it only handles transitions started from Shell. */
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 29e302a..1d17cd6 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
@@ -461,7 +461,12 @@
      * @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, displayId: Int, taskId: Int) {
+    fun onDesktopWindowClose(
+        wct: WindowContainerTransaction,
+        displayId: Int,
+        taskInfo: RunningTaskInfo,
+    ): ((IBinder) -> Unit)? {
+        val taskId = taskInfo.taskId
         if (taskRepository.isOnlyVisibleNonClosingTask(taskId)) {
             removeWallpaperActivity(wct)
         }
@@ -472,6 +477,7 @@
                 taskId
             )
         )
+        return immersiveTransitionHandler.exitImmersiveIfApplicable(wct, taskInfo)
     }
 
     fun minimizeTask(taskInfo: RunningTaskInfo) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
index 6aaf001..58337ec 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionHandler.java
@@ -99,9 +99,11 @@
 
 
     @Override
-    public void startRemoveTransition(WindowContainerTransaction wct) {
+    public IBinder startRemoveTransition(WindowContainerTransaction wct) {
         final int type = WindowManager.TRANSIT_CLOSE;
-        mPendingTransitionTokens.add(mTransitions.startTransition(type, wct, this));
+        final IBinder transition = mTransitions.startTransition(type, wct, this);
+        mPendingTransitionTokens.add(transition);
+        return transition;
     }
 
     @Override
@@ -229,8 +231,7 @@
         SurfaceControl.Transaction t = new SurfaceControl.Transaction();
         SurfaceControl sc = change.getLeash();
         finishT.hide(sc);
-        Rect startBounds = new Rect(change.getTaskInfo().configuration.windowConfiguration
-                .getBounds());
+        final Rect startBounds = new Rect(change.getStartAbsBounds());
         animator.addUpdateListener(animation -> {
             t.setPosition(sc, startBounds.left,
                     startBounds.top + (animation.getAnimatedFraction() * screenHeight));
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarter.java b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarter.java
index ea68a69..5984d48 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarter.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/freeform/FreeformTaskTransitionStarter.java
@@ -48,6 +48,7 @@
      *
      * @param wct the {@link WindowContainerTransaction} that closes the task
      *
+     * @return the started transition
      */
-    void startRemoveTransition(WindowContainerTransaction wct);
+    IBinder startRemoveTransition(WindowContainerTransaction wct);
 }
\ No newline at end of file
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 9e089b2..a06b4a2 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
@@ -56,6 +56,7 @@
 import android.graphics.Region;
 import android.hardware.input.InputManager;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.RemoteException;
 import android.os.UserHandle;
@@ -133,6 +134,7 @@
 
 import kotlin.Pair;
 import kotlin.Unit;
+import kotlin.jvm.functions.Function1;
 
 import kotlinx.coroutines.ExperimentalCoroutinesApi;
 
@@ -767,8 +769,13 @@
                             SplitScreenController.EXIT_REASON_DESKTOP_MODE);
                 } else {
                     WindowContainerTransaction wct = new WindowContainerTransaction();
-                    mDesktopTasksController.onDesktopWindowClose(wct, mDisplayId, mTaskId);
-                    mTaskOperations.closeTask(mTaskToken, wct);
+                    final Function1<IBinder, Unit> runOnTransitionStart =
+                            mDesktopTasksController.onDesktopWindowClose(
+                                    wct, mDisplayId, decoration.mTaskInfo);
+                    final IBinder transition = mTaskOperations.closeTask(mTaskToken, wct);
+                    if (transition != null && runOnTransitionStart != null) {
+                        runOnTransitionStart.invoke(transition);
+                    }
                 }
             } else if (id == R.id.back_button) {
                 mTaskOperations.injectBackKey(mDisplayId);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java
index 61b9393..bc85d2b 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/windowdecor/TaskOperations.java
@@ -76,13 +76,14 @@
         closeTask(taskToken, new WindowContainerTransaction());
     }
 
-    void closeTask(WindowContainerToken taskToken, WindowContainerTransaction wct) {
+    IBinder closeTask(WindowContainerToken taskToken, WindowContainerTransaction wct) {
         wct.removeTask(taskToken);
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
-            mTransitionStarter.startRemoveTransition(wct);
+            return mTransitionStarter.startRemoveTransition(wct);
         } else {
             mSyncQueue.queue(wct);
         }
+        return null;
     }
 
     IBinder minimizeTask(WindowContainerToken taskToken) {
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandlerTest.kt
index b137468..ef99b00 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopFullImmersiveTransitionHandlerTest.kt
@@ -52,6 +52,7 @@
 import org.mockito.kotlin.any
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.never
 import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
@@ -366,7 +367,7 @@
             immersive = false
         )
 
-        immersiveHandler.exitImmersiveIfApplicable(wct, task.taskId)
+        immersiveHandler.exitImmersiveIfApplicable(wct, task)
 
         assertThat(wct.hasBoundsChange(task.token)).isFalse()
     }
@@ -384,12 +385,12 @@
             immersive = true
         )
 
-        immersiveHandler.exitImmersiveIfApplicable(wct, task.taskId)?.invoke(transition)
+        immersiveHandler.exitImmersiveIfApplicable(wct, task)?.invoke(transition)
 
         assertThat(immersiveHandler.pendingExternalExitTransitions.any { exit ->
             exit.transition == transition && exit.displayId == DEFAULT_DISPLAY
                     && exit.taskId == task.taskId
-        }).isFalse()
+        }).isTrue()
     }
 
     @Test
@@ -405,7 +406,7 @@
             immersive = false
         )
 
-        immersiveHandler.exitImmersiveIfApplicable(wct, task.taskId)?.invoke(transition)
+        immersiveHandler.exitImmersiveIfApplicable(wct, task)?.invoke(transition)
 
         assertThat(immersiveHandler.pendingExternalExitTransitions.any { exit ->
             exit.transition == transition && exit.displayId == DEFAULT_DISPLAY
@@ -565,6 +566,24 @@
         ).isTrue()
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
+    fun exitImmersive_pendingExit_doesNotExitAgain() {
+        val task = createFreeformTask()
+        whenever(mockShellTaskOrganizer.getRunningTaskInfo(task.taskId)).thenReturn(task)
+        val wct = WindowContainerTransaction()
+        desktopRepository.setTaskInFullImmersiveState(
+            displayId = task.displayId,
+            taskId = task.taskId,
+            immersive = true
+        )
+        immersiveHandler.exitImmersiveIfApplicable(wct, task)?.invoke(Binder())
+
+        immersiveHandler.moveTaskToNonImmersive(task)
+
+        verify(mockTransitions, never()).startTransition(any(), any(), any())
+    }
+
     private fun createTransitionInfo(
         @TransitionType type: Int = TRANSIT_CHANGE,
         @TransitionFlags flags: Int = 0,
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
index 07de0716..81d59d5 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopMixedTransitionHandlerTest.kt
@@ -21,6 +21,7 @@
 import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
 import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
 import android.app.WindowConfiguration.WindowingMode
+import android.os.Binder
 import android.os.Handler
 import android.os.IBinder
 import android.testing.AndroidTestingRunner
@@ -107,6 +108,8 @@
     @Test
     fun startRemoveTransition_startsCloseTransition() {
         val wct = WindowContainerTransaction()
+        whenever(transitions.startTransition(WindowManager.TRANSIT_CLOSE, wct, mixedHandler))
+            .thenReturn(Binder())
 
         mixedHandler.startRemoveTransition(wct)
 
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 113990e..af51e32 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
@@ -1480,8 +1480,9 @@
 
   @Test
   fun onDesktopWindowClose_noActiveTasks() {
+    val task = setUpFreeformTask(active = false)
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = 1)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -1490,7 +1491,7 @@
   fun onDesktopWindowClose_singleActiveTask_noWallpaperActivityToken() {
     val task = setUpFreeformTask()
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -1502,7 +1503,7 @@
     taskRepository.wallpaperActivityToken = wallpaperToken
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
     // Adds remove wallpaper operation
     wct.assertRemoveAt(index = 0, wallpaperToken)
   }
@@ -1515,7 +1516,7 @@
     taskRepository.addClosingTask(DEFAULT_DISPLAY, task.taskId)
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -1528,7 +1529,7 @@
     taskRepository.minimizeTask(DEFAULT_DISPLAY, task.taskId)
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -1541,7 +1542,7 @@
     taskRepository.wallpaperActivityToken = wallpaperToken
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task1.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task1)
     // Doesn't modify transaction
     assertThat(wct.hierarchyOps).isEmpty()
   }
@@ -1555,7 +1556,7 @@
     taskRepository.addClosingTask(DEFAULT_DISPLAY, task2.taskId)
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task1.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task1)
     // Adds remove wallpaper operation
     wct.assertRemoveAt(index = 0, wallpaperToken)
   }
@@ -1569,7 +1570,7 @@
     taskRepository.minimizeTask(DEFAULT_DISPLAY, task2.taskId)
 
     val wct = WindowContainerTransaction()
-    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, taskId = task1.taskId)
+    controller.onDesktopWindowClose(wct, displayId = DEFAULT_DISPLAY, task1)
     // Adds remove wallpaper operation
     wct.assertRemoveAt(index = 0, wallpaperToken)
   }