Merge changes Ia7519bcb,I0a60bd25,I5b8edf7b into main

* changes:
  Fix desktop immersive's state clean up
  Handle the case where transition may have empty immersive change
  Dump DesktopImmersiveController state
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index a472f79..44fce81 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -836,14 +836,21 @@
     @Provides
     static Optional<DesktopImmersiveController> provideDesktopImmersiveController(
             Context context,
+            ShellInit shellInit,
             Transitions transitions,
             @DynamicOverride DesktopRepository desktopRepository,
             DisplayController displayController,
-            ShellTaskOrganizer shellTaskOrganizer) {
+            ShellTaskOrganizer shellTaskOrganizer,
+            ShellCommandHandler shellCommandHandler) {
         if (DesktopModeStatus.canEnterDesktopMode(context)) {
             return Optional.of(
                     new DesktopImmersiveController(
-                            transitions, desktopRepository, displayController, shellTaskOrganizer));
+                            shellInit,
+                            transitions,
+                            desktopRepository,
+                            displayController,
+                            shellTaskOrganizer,
+                            shellCommandHandler));
         }
         return Optional.empty();
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopImmersiveController.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopImmersiveController.kt
index f69aa6d..1acde73 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopImmersiveController.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopImmersiveController.kt
@@ -34,10 +34,13 @@
 import com.android.wm.shell.ShellTaskOrganizer
 import com.android.wm.shell.common.DisplayController
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
+import com.android.wm.shell.sysui.ShellCommandHandler
+import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
 import com.android.wm.shell.transition.Transitions.TransitionHandler
 import com.android.wm.shell.transition.Transitions.TransitionObserver
 import com.android.wm.shell.windowdecor.OnTaskResizeAnimationListener
+import java.io.PrintWriter
 
 /**
  * A controller to move tasks in/out of desktop's full immersive state where the task
@@ -45,27 +48,34 @@
  * be transient below the status bar like in fullscreen immersive mode.
  */
 class DesktopImmersiveController(
+    shellInit: ShellInit,
     private val transitions: Transitions,
     private val desktopRepository: DesktopRepository,
     private val displayController: DisplayController,
     private val shellTaskOrganizer: ShellTaskOrganizer,
+    private val shellCommandHandler: ShellCommandHandler,
     private val transactionSupplier: () -> SurfaceControl.Transaction,
 ) : TransitionHandler, TransitionObserver {
 
     constructor(
+        shellInit: ShellInit,
         transitions: Transitions,
         desktopRepository: DesktopRepository,
         displayController: DisplayController,
         shellTaskOrganizer: ShellTaskOrganizer,
+        shellCommandHandler: ShellCommandHandler,
     ) : this(
+        shellInit,
         transitions,
         desktopRepository,
         displayController,
         shellTaskOrganizer,
+        shellCommandHandler,
         { SurfaceControl.Transaction() }
     )
 
-    private var state: TransitionState? = null
+    @VisibleForTesting
+    var state: TransitionState? = null
 
     @VisibleForTesting
     val pendingExternalExitTransitions = mutableListOf<ExternalPendingExit>()
@@ -79,10 +89,21 @@
     /** A listener to invoke on animation changes during entry/exit. */
     var onTaskResizeAnimationListener: OnTaskResizeAnimationListener? = null
 
+    init {
+        shellInit.addInitCallback({ onInit() }, this)
+    }
+
+    fun onInit() {
+        shellCommandHandler.addDumpCallback(this::dump, this)
+    }
+
     /** Starts a transition to enter full immersive state inside the desktop. */
     fun moveTaskToImmersive(taskInfo: RunningTaskInfo) {
         if (inProgress) {
-            logV("Cannot start entry because transition already in progress.")
+            logV(
+                "Cannot start entry because transition(s) already in progress: %s",
+                getRunningTransitions()
+            )
             return
         }
         val wct = WindowContainerTransaction().apply {
@@ -100,7 +121,10 @@
 
     fun moveTaskToNonImmersive(taskInfo: RunningTaskInfo) {
         if (inProgress) {
-            logV("Cannot start exit because transition already in progress.")
+            logV(
+                "Cannot start exit because transition(s) already in progress: %s",
+                getRunningTransitions()
+            )
             return
         }
 
@@ -225,14 +249,19 @@
         finishCallback: Transitions.TransitionFinishCallback
     ): Boolean {
         val state = requireState()
-        if (transition != state.transition) return false
+        check(state.transition == transition) {
+            "Transition $transition did not match expected state=$state"
+        }
         logD("startAnimation transition=%s", transition)
         animateResize(
             targetTaskId = state.taskId,
             info = info,
             startTransaction = startTransaction,
             finishTransaction = finishTransaction,
-            finishCallback = finishCallback
+            finishCallback = {
+                finishCallback.onTransitionFinished(/* wct= */ null)
+                clearState()
+            },
         )
         return true
     }
@@ -242,12 +271,18 @@
         info: TransitionInfo,
         startTransaction: SurfaceControl.Transaction,
         finishTransaction: SurfaceControl.Transaction,
-        finishCallback: Transitions.TransitionFinishCallback
+        finishCallback: Transitions.TransitionFinishCallback,
     ) {
         logD("animateResize for task#%d", targetTaskId)
-        val change = info.changes.first { c ->
+        val change = info.changes.firstOrNull { c ->
             val taskInfo = c.taskInfo
-            return@first taskInfo != null && taskInfo.taskId == targetTaskId
+            return@firstOrNull taskInfo != null && taskInfo.taskId == targetTaskId
+        }
+        if (change == null) {
+            logD("Did not find change for task#%d to animate", targetTaskId)
+            startTransaction.apply()
+            finishCallback.onTransitionFinished(/* wct= */ null)
+            return
         }
         animateResizeChange(change, startTransaction, finishTransaction, finishCallback)
     }
@@ -288,7 +323,6 @@
                         .apply()
                     onTaskResizeAnimationListener?.onAnimationEnd(taskId)
                     finishCallback.onTransitionFinished(null /* wct */)
-                    clearState()
                 }
             )
             addUpdateListener { animation ->
@@ -357,8 +391,17 @@
         // Check if this is a direct immersive enter/exit transition.
         if (transition == state?.transition) {
             val state = requireState()
-            val startBounds = info.changes.first { c -> c.taskInfo?.taskId == state.taskId }
-                .startAbsBounds
+            val immersiveChange = info.changes.firstOrNull { c ->
+                c.taskInfo?.taskId == state.taskId
+            }
+            if (immersiveChange == null) {
+                logV(
+                    "Direct move for task#%d in %s direction missing immersive change.",
+                    state.taskId, state.direction
+                )
+                return
+            }
+            val startBounds = immersiveChange.startAbsBounds
             logV("Direct move for task#%d in %s direction verified", state.taskId, state.direction)
             when (state.direction) {
                 Direction.ENTER -> {
@@ -446,11 +489,30 @@
     private fun requireState(): TransitionState =
         state ?: error("Expected non-null transition state")
 
+    private fun getRunningTransitions(): List<IBinder> {
+        val running = mutableListOf<IBinder>()
+        state?.let {
+            running.add(it.transition)
+        }
+        pendingExternalExitTransitions.forEach {
+            running.add(it.transition)
+        }
+        return running
+    }
+
     private fun TransitionInfo.hasTaskChange(taskId: Int): Boolean =
         changes.any { c -> c.taskInfo?.taskId == taskId }
 
+    private fun dump(pw: PrintWriter, prefix: String) {
+        val innerPrefix = "$prefix  "
+        pw.println("${prefix}DesktopImmersiveController")
+        pw.println(innerPrefix + "state=" + state)
+        pw.println(innerPrefix + "pendingExternalExitTransitions=" + pendingExternalExitTransitions)
+    }
+
     /** The state of the currently running transition. */
-    private data class TransitionState(
+    @VisibleForTesting
+    data class TransitionState(
         val transition: IBinder,
         val displayId: Int,
         val taskId: Int,
@@ -483,7 +545,8 @@
         fun asExit(): Exit? = if (this is Exit) this else null
     }
 
-    private enum class Direction {
+    @VisibleForTesting
+    enum class Direction {
         ENTER, EXIT
     }
 
@@ -495,9 +558,10 @@
         ProtoLog.d(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
     }
 
-    private companion object {
+    companion object {
         private const val TAG = "DesktopImmersive"
 
-        private const val FULL_IMMERSIVE_ANIM_DURATION_MS = 336L
+        @VisibleForTesting
+        const val FULL_IMMERSIVE_ANIM_DURATION_MS = 336L
     }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt
index e05a0b5..a4f4d05 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopImmersiveControllerTest.kt
@@ -15,6 +15,7 @@
  */
 package com.android.wm.shell.desktopmode
 
+import android.animation.AnimatorTestRule
 import android.app.ActivityManager.RunningTaskInfo
 import android.app.WindowConfiguration.WINDOW_CONFIG_BOUNDS
 import android.graphics.Rect
@@ -24,6 +25,7 @@
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
 import android.testing.AndroidTestingRunner
+import android.testing.TestableLooper
 import android.view.Display.DEFAULT_DISPLAY
 import android.view.Surface
 import android.view.SurfaceControl
@@ -43,6 +45,7 @@
 import com.android.wm.shell.desktopmode.DesktopTestHelpers.Companion.createFreeformTask
 import com.android.wm.shell.sysui.ShellInit
 import com.android.wm.shell.transition.Transitions
+import com.android.wm.shell.util.StubTransaction
 import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.Rule
@@ -64,17 +67,19 @@
  * Usage: atest WMShellUnitTests:DesktopImmersiveControllerTest
  */
 @SmallTest
+@TestableLooper.RunWithLooper
 @RunWith(AndroidTestingRunner::class)
 class DesktopImmersiveControllerTest : ShellTestCase() {
 
     @JvmField @Rule val setFlagsRule = SetFlagsRule()
+    @JvmField @Rule val animatorTestRule = AnimatorTestRule(this)
 
     @Mock private lateinit var mockTransitions: Transitions
     private lateinit var desktopRepository: DesktopRepository
     @Mock private lateinit var mockDisplayController: DisplayController
     @Mock private lateinit var mockShellTaskOrganizer: ShellTaskOrganizer
     @Mock private lateinit var mockDisplayLayout: DisplayLayout
-    private val transactionSupplier = { SurfaceControl.Transaction() }
+    private val transactionSupplier = { StubTransaction() }
 
     private lateinit var controller: DesktopImmersiveController
 
@@ -89,10 +94,12 @@
             (invocation.getArgument(0) as Rect).set(STABLE_BOUNDS)
         }
         controller = DesktopImmersiveController(
+            shellInit = mock(),
             transitions = mockTransitions,
             desktopRepository = desktopRepository,
             displayController = mockDisplayController,
             shellTaskOrganizer = mockShellTaskOrganizer,
+            shellCommandHandler = mock(),
             transactionSupplier = transactionSupplier,
         )
     }
@@ -672,6 +679,60 @@
         assertThat(controller.isImmersiveChange(transition, change)).isTrue()
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
+    fun externalAnimateResizeChange_doesNotCleanUpPendingTransitionState() {
+        val task = createFreeformTask()
+        val mockBinder = mock(IBinder::class.java)
+        whenever(mockTransitions.startTransition(eq(TRANSIT_CHANGE), any(), eq(controller)))
+            .thenReturn(mockBinder)
+        desktopRepository.setTaskInFullImmersiveState(
+            displayId = task.displayId,
+            taskId = task.taskId,
+            immersive = true
+        )
+
+        controller.moveTaskToNonImmersive(task)
+
+        controller.animateResizeChange(
+            change = TransitionInfo.Change(task.token, SurfaceControl()).apply {
+                taskInfo = task
+            },
+            startTransaction = StubTransaction(),
+            finishTransaction = StubTransaction(),
+            finishCallback = { }
+        )
+        animatorTestRule.advanceTimeBy(DesktopImmersiveController.FULL_IMMERSIVE_ANIM_DURATION_MS)
+
+        assertThat(controller.state).isNotNull()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
+    fun startAnimation_missingChange_clearsState() {
+        val task = createFreeformTask()
+        val mockBinder = mock(IBinder::class.java)
+        whenever(mockTransitions.startTransition(eq(TRANSIT_CHANGE), any(), eq(controller)))
+            .thenReturn(mockBinder)
+        desktopRepository.setTaskInFullImmersiveState(
+            displayId = task.displayId,
+            taskId = task.taskId,
+            immersive = false
+        )
+
+        controller.moveTaskToImmersive(task)
+
+        controller.startAnimation(
+            transition = mockBinder,
+            info = createTransitionInfo(changes = emptyList()),
+            startTransaction = StubTransaction(),
+            finishTransaction = StubTransaction(),
+            finishCallback = {}
+        )
+
+        assertThat(controller.state).isNull()
+    }
+
     private fun createTransitionInfo(
         @TransitionType type: Int = TRANSIT_CHANGE,
         @TransitionFlags flags: Int = 0,