Merge "Update view model when desk organizer updates tasks." into main

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 9875ec3a7ca8..46c9b07fb802 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
@@ -1054,7 +1054,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();
@@ -1072,7 +1073,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 0082d7971ad2..16fa5120d64b 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
@@ -125,6 +125,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;
@@ -210,6 +211,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<>();
@@ -308,7 +310,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
             DesktopModeCompatPolicy desktopModeCompatPolicy,
             DesktopTilingDecorViewModel desktopTilingDecorViewModel,
             MultiDisplayDragMoveIndicatorController multiDisplayDragMoveIndicatorController,
-            CompatUIHandler compatUI) {
+            CompatUIHandler compatUI,
+            DesksOrganizer desksOrganizer) {
         this(
                 context,
                 shellExecutor,
@@ -356,7 +359,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
                 desktopModeCompatPolicy,
                 desktopTilingDecorViewModel,
                 multiDisplayDragMoveIndicatorController,
-                compatUI);
+                compatUI,
+                desksOrganizer);
     }
 
     @VisibleForTesting
@@ -407,7 +411,8 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
             DesktopModeCompatPolicy desktopModeCompatPolicy,
             DesktopTilingDecorViewModel desktopTilingDecorViewModel,
             MultiDisplayDragMoveIndicatorController multiDisplayDragMoveIndicatorController,
-            CompatUIHandler compatUI) {
+            CompatUIHandler compatUI,
+            DesksOrganizer desksOrganizer) {
         mContext = context;
         mMainExecutor = shellExecutor;
         mMainHandler = mainHandler;
@@ -485,6 +490,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
         mDesktopTasksController.setSnapEventHandler(this);
         mMultiDisplayDragMoveIndicatorController = multiDisplayDragMoveIndicatorController;
         mLatencyTracker = LatencyTracker.getInstance(mContext);
+        mDesksOrganizer = desksOrganizer;
         shellInit.addInitCallback(this::onInit, this);
     }
 
@@ -523,6 +529,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 80dcd7d69f00..23994a2bd547 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)