Update view model when desk organizer updates tasks.

When an onTaskInfoChanged event reaches DesksOrganizer, it does not ultimately reach DesktopModeWindowDecorViewModel, meaning window decorations end up with outdated taskInfo. This CL fixes this by calling view model's onTaskInfoChanged during RootTaskDesksOrganizer#onTaskInfoChanged

Bug: 404292560
Test: Manual, logging
Flag: com.android.window.flags.enable_multiple_desktops_backend
Change-Id: Iac6f0afdfe3c4aff871886a650fc0b142c72dd2d

6 files changed, 71 insertions, 5 deletions

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 5fbbb0bf1e78..6b419029be44 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
@@ -1047,7 +1047,8 @@ public abstract class WMShellModule {
             DesktopModeCompatPolicy desktopModeCompatPolicy,
             DesktopTilingDecorViewModel desktopTilingDecorViewModel,
             MultiDisplayDragMoveIndicatorController multiDisplayDragMoveIndicatorController,
-            Optional<CompatUIHandler> compatUI
+            Optional<CompatUIHandler> compatUI,
+            DesksOrganizer desksOrganizer
     ) {
         if (!DesktopModeStatus.canEnterDesktopModeOrShowAppHandle(context)) {
             return Optional.empty();
@@ -1065,7 +1066,8 @@ public abstract class WMShellModule {
                 activityOrientationChangeHandler, focusTransitionObserver, desktopModeEventLogger,
                 desktopModeUiEventLogger, taskResourceLoader, recentsTransitionHandler,
                 desktopModeCompatPolicy, desktopTilingDecorViewModel,
-                multiDisplayDragMoveIndicatorController, compatUI.orElse(null)));
+                multiDisplayDragMoveIndicatorController, compatUI.orElse(null),
+                desksOrganizer));
     }
 
     @WMSingleton
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksOrganizer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksOrganizer.kt
index 5a988fcd1b77..1effcdb20505 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksOrganizer.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksOrganizer.kt
@@ -78,6 +78,9 @@ interface DesksOrganizer {
     /** Whether the desk is activate according to the given change at the end of a transition. */
     fun isDeskActiveAtEnd(change: TransitionInfo.Change, deskId: Int): Boolean
 
+    /** Allows for other classes to respond to task changes this organizer receives. */
+    fun setOnDesktopTaskInfoChangedListener(listener: (ActivityManager.RunningTaskInfo) -> Unit)
+
     /** A callback that is invoked when the desk container is created. */
     fun interface OnCreateCallback {
         /** Calls back when the [deskId] has been created. */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt
index 49ca58e7b32a..c30987ac7640 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt
@@ -54,6 +54,7 @@ class RootTaskDesksOrganizer(
         mutableListOf<CreateDeskMinimizationRootRequest>()
     @VisibleForTesting
     val deskMinimizationRootsByDeskId: MutableMap<Int, DeskMinimizationRoot> = mutableMapOf()
+    private var onTaskInfoChangedListener: ((RunningTaskInfo) -> Unit)? = null
 
     init {
         if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
@@ -213,6 +214,10 @@ class RootTaskDesksOrganizer(
             change.taskInfo?.isVisibleRequested == true &&
             change.mode == TRANSIT_TO_FRONT
 
+    override fun setOnDesktopTaskInfoChangedListener(listener: (RunningTaskInfo) -> Unit) {
+        onTaskInfoChangedListener = listener
+    }
+
     override fun onTaskAppeared(taskInfo: RunningTaskInfo, leash: SurfaceControl) {
         handleTaskAppeared(taskInfo, leash)
         updateLaunchAdjacentController()
@@ -220,6 +225,12 @@ class RootTaskDesksOrganizer(
 
     override fun onTaskInfoChanged(taskInfo: RunningTaskInfo) {
         handleTaskInfoChanged(taskInfo)
+        if (
+            taskInfo.taskId !in deskRootsByDeskId &&
+                deskMinimizationRootsByDeskId.values.none { it.rootId == taskInfo.taskId }
+        ) {
+            onTaskInfoChangedListener?.invoke(taskInfo)
+        }
         updateLaunchAdjacentController()
     }
 
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 69e1f36dec0b..f7617fb6c993 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
@@ -127,6 +127,7 @@ import com.android.wm.shell.desktopmode.common.ToggleTaskSizeInteraction;
 import com.android.wm.shell.desktopmode.common.ToggleTaskSizeUtilsKt;
 import com.android.wm.shell.desktopmode.education.AppHandleEducationController;
 import com.android.wm.shell.desktopmode.education.AppToWebEducationController;
+import com.android.wm.shell.desktopmode.multidesks.DesksOrganizer;
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter;
 import com.android.wm.shell.recents.RecentsTransitionHandler;
 import com.android.wm.shell.recents.RecentsTransitionStateListener;
@@ -212,6 +213,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
     private final AppHandleAndHeaderVisibilityHelper mAppHandleAndHeaderVisibilityHelper;
     private final AppHeaderViewHolder.Factory mAppHeaderViewHolderFactory;
     private final AppHandleViewHolder.Factory mAppHandleViewHolderFactory;
+    private final DesksOrganizer mDesksOrganizer;
     private boolean mTransitionDragActive;
 
     private SparseArray<EventReceiver> mEventReceiversByDisplay = new SparseArray<>();
@@ -310,7 +312,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
             DesktopModeCompatPolicy desktopModeCompatPolicy,
             DesktopTilingDecorViewModel desktopTilingDecorViewModel,
             MultiDisplayDragMoveIndicatorController multiDisplayDragMoveIndicatorController,
-            CompatUIHandler compatUI) {
+            CompatUIHandler compatUI,
+            DesksOrganizer desksOrganizer) {
         this(
                 context,
                 shellExecutor,
@@ -358,7 +361,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
                 desktopModeCompatPolicy,
                 desktopTilingDecorViewModel,
                 multiDisplayDragMoveIndicatorController,
-                compatUI);
+                compatUI,
+                desksOrganizer);
     }
 
     @VisibleForTesting
@@ -409,7 +413,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
             DesktopModeCompatPolicy desktopModeCompatPolicy,
             DesktopTilingDecorViewModel desktopTilingDecorViewModel,
             MultiDisplayDragMoveIndicatorController multiDisplayDragMoveIndicatorController,
-            CompatUIHandler compatUI) {
+            CompatUIHandler compatUI,
+            DesksOrganizer desksOrganizer) {
         mContext = context;
         mMainExecutor = shellExecutor;
         mMainHandler = mainHandler;
@@ -487,6 +492,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
         mDesktopTasksController.setSnapEventHandler(this);
         mMultiDisplayDragMoveIndicatorController = multiDisplayDragMoveIndicatorController;
         mLatencyTracker = LatencyTracker.getInstance(mContext);
+        mDesksOrganizer = desksOrganizer;
         shellInit.addInitCallback(this::onInit, this);
     }
 
@@ -525,6 +531,10 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
                     });
         }
         mFocusTransitionObserver.setLocalFocusTransitionListener(this, mMainExecutor);
+        mDesksOrganizer.setOnDesktopTaskInfoChangedListener((taskInfo) -> {
+            onTaskInfoChanged(taskInfo);
+            return Unit.INSTANCE;
+        });
     }
 
     @Override
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizerTest.kt
index 9af504797182..e57fc38e3607 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizerTest.kt
@@ -15,6 +15,7 @@
  */
 package com.android.wm.shell.desktopmode.multidesks
 
+import android.app.ActivityManager
 import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
 import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
 import android.testing.AndroidTestingRunner
@@ -48,7 +49,9 @@ import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mockito
+import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
+import org.mockito.kotlin.any
 import org.mockito.kotlin.argThat
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
@@ -67,6 +70,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     private val mockShellCommandHandler = mock<ShellCommandHandler>()
     private val mockShellTaskOrganizer = mock<ShellTaskOrganizer>()
     private val launchAdjacentController = LaunchAdjacentController(mock())
+    private val taskInfoChangedListener = mock<(ActivityManager.RunningTaskInfo) -> Unit>()
 
     private lateinit var organizer: RootTaskDesksOrganizer
 
@@ -79,6 +83,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
                 mockShellTaskOrganizer,
                 launchAdjacentController,
             )
+        organizer.setOnDesktopTaskInfoChangedListener(taskInfoChangedListener)
     }
 
     @Test fun testCreateDesk_createsDeskAndMinimizationRoots() = runTest { createDesk() }
@@ -652,6 +657,34 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
         assertThat(launchAdjacentController.launchAdjacentEnabled).isFalse()
     }
 
+    @Test
+    fun onTaskInfoChanged_taskNotRoot_invokesListener() = runTest {
+        createDesk()
+        val task = createFreeformTask().apply { taskId = TEST_CHILD_TASK_ID }
+
+        organizer.onTaskInfoChanged(task)
+
+        verify(taskInfoChangedListener).invoke(task)
+    }
+
+    @Test
+    fun onTaskInfoChanged_isDeskRoot_doesNotInvokeListener() = runTest {
+        val deskRoot = createDesk().deskRoot
+
+        organizer.onTaskInfoChanged(deskRoot.taskInfo)
+
+        verify(taskInfoChangedListener, never()).invoke(any())
+    }
+
+    @Test
+    fun onTaskInfoChanged_isMinimizationRoot_doesNotInvokeListener() = runTest {
+        val minimizationRoot = createDesk().minimizationRoot
+
+        organizer.onTaskInfoChanged(minimizationRoot.taskInfo)
+
+        verify(taskInfoChangedListener, never()).invoke(any())
+    }
+
     private data class DeskRoots(
         val deskRoot: DeskRoot,
         val minimizationRoot: DeskMinimizationRoot,
@@ -712,4 +745,8 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
                 hop.newParent == desk.deskRoot.token.asBinder() &&
                 hop.toTop
         }
+
+    companion object {
+        private const val TEST_CHILD_TASK_ID = 100
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt
index 2126d1d9b986..ccf3710d4b84 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTestsBase.kt
@@ -67,6 +67,7 @@ import com.android.wm.shell.desktopmode.DesktopUserRepositories
 import com.android.wm.shell.desktopmode.WindowDecorCaptionHandleRepository
 import com.android.wm.shell.desktopmode.education.AppHandleEducationController
 import com.android.wm.shell.desktopmode.education.AppToWebEducationController
+import com.android.wm.shell.desktopmode.multidesks.DesksOrganizer
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter
 import com.android.wm.shell.recents.RecentsTransitionHandler
 import com.android.wm.shell.recents.RecentsTransitionStateListener
@@ -146,6 +147,7 @@ open class DesktopModeWindowDecorViewModelTestsBase : ShellTestCase() {
     protected val mockMultiDisplayDragMoveIndicatorController =
         mock<MultiDisplayDragMoveIndicatorController>()
     protected val mockCompatUIHandler = mock<CompatUIHandler>()
+    protected val mockDesksOrganizer = mock<DesksOrganizer>()
     protected val mockInputManager = mock<InputManager>()
     private val mockTaskPositionerFactory =
         mock<DesktopModeWindowDecorViewModel.TaskPositionerFactory>()
@@ -246,6 +248,7 @@ open class DesktopModeWindowDecorViewModelTestsBase : ShellTestCase() {
             mockTilingWindowDecoration,
             mockMultiDisplayDragMoveIndicatorController,
             mockCompatUIHandler,
+            mockDesksOrganizer
         )
         desktopModeWindowDecorViewModel.setSplitScreenController(mockSplitScreenController)
         whenever(mockDisplayController.getDisplayLayout(any())).thenReturn(mockDisplayLayout)