Merge "[8/N] Desks: Implement move-to-desk for running tasks" into main

12 files changed, 636 insertions, 102 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt
index b0a38a85e0b2..03bc42f08d59 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopModeShellCommandHandler.kt
@@ -57,7 +57,7 @@ class DesktopModeShellCommandHandler(private val controller: DesktopTasksControl
                 return false
             }
         if (!Flags.enableMultipleDesktopsBackend()) {
-            return controller.moveTaskToDesktop(taskId, transitionSource = UNKNOWN)
+            return controller.moveTaskToDefaultDeskAndActivate(taskId, transitionSource = UNKNOWN)
         }
         if (args.size < 3) {
             pw.println("Error: desk id should be provided as arguments")
@@ -70,8 +70,9 @@ class DesktopModeShellCommandHandler(private val controller: DesktopTasksControl
                 pw.println("Error: desk id should be an integer")
                 return false
             }
+        controller.moveTaskToDesk(taskId = taskId, deskId = deskId, transitionSource = UNKNOWN)
         pw.println("Not implemented.")
-        return false
+        return true
     }
 
     private fun runMoveToNextDisplay(args: Array<String>, pw: PrintWriter): Boolean {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
index 9583f3da2ee3..4777e7f93bc9 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopRepository.kt
@@ -236,25 +236,32 @@ class DesktopRepository(
      * TODO: b/389960283 - add explicit [deskId] argument.
      */
     fun addTask(displayId: Int, taskId: Int, isVisible: Boolean) {
-        addOrMoveFreeformTaskToTop(displayId, taskId)
-        addActiveTask(displayId, taskId)
-        updateTask(displayId, taskId, isVisible)
+        val activeDesk =
+            checkNotNull(desktopData.getDefaultDesk(displayId)) {
+                "Expected desk in display: $displayId"
+            }
+        addTaskToDesk(displayId = displayId, deskId = activeDesk.deskId, taskId = taskId, isVisible)
     }
 
-    /**
-     * Adds task with [taskId] to the list of active tasks on [displayId]'s active desk.
-     *
-     * TODO: b/389960283 - add explicit [deskId] argument.
-     */
-    private fun addActiveTask(displayId: Int, taskId: Int) {
-        val activeDesk = desktopData.getDefaultDesk(displayId)
-        checkNotNull(activeDesk) { "Expected desk in display: $displayId" }
+    fun addTaskToDesk(displayId: Int, deskId: Int, taskId: Int, isVisible: Boolean) {
+        addOrMoveTaskToTopOfDesk(displayId = displayId, deskId = deskId, taskId = taskId)
+        addActiveTaskToDesk(displayId = displayId, deskId = deskId, taskId = taskId)
+        updateTaskInDesk(
+            displayId = displayId,
+            deskId = deskId,
+            taskId = taskId,
+            isVisible = isVisible,
+        )
+    }
+
+    private fun addActiveTaskToDesk(displayId: Int, deskId: Int, taskId: Int) {
+        val desk = checkNotNull(desktopData.getDesk(deskId)) { "Did not find desk: $deskId" }
 
-        // Removes task if it is active on another desk excluding [activeDesk].
-        removeActiveTask(taskId, excludedDeskId = activeDesk.deskId)
+        // Removes task if it is active on another desk excluding this desk.
+        removeActiveTask(taskId, excludedDeskId = deskId)
 
-        if (activeDesk.activeTasks.add(taskId)) {
-            logD("Adds active task=%d displayId=%d deskId=%d", taskId, displayId, activeDesk.deskId)
+        if (desk.activeTasks.add(taskId)) {
+            logD("Adds active task=%d displayId=%d deskId=%d", taskId, displayId, deskId)
             updateActiveTasksListeners(displayId)
         }
     }
@@ -405,10 +412,10 @@ class DesktopRepository(
                 emptySet()
             }
 
-    /** Removes task from visible tasks of all displays except [excludedDisplayId]. */
-    private fun removeVisibleTask(taskId: Int, excludedDisplayId: Int? = null) {
+    /** Removes task from visible tasks of all desks except [excludedDeskId]. */
+    private fun removeVisibleTask(taskId: Int, excludedDeskId: Int? = null) {
         desktopData.forAllDesks { displayId, desk ->
-            if (displayId != excludedDisplayId && desk.visibleTasks.remove(taskId)) {
+            if (desk.deskId != excludedDeskId && desk.visibleTasks.remove(taskId)) {
                 notifyVisibleTaskListeners(displayId, desk.visibleTasks.size)
             }
         }
@@ -423,30 +430,58 @@ class DesktopRepository(
      * TODO: b/389960283 - add explicit [deskId] argument.
      */
     fun updateTask(displayId: Int, taskId: Int, isVisible: Boolean) {
-        logD("updateTask taskId=%d, displayId=%d, isVisible=%b", taskId, displayId, isVisible)
+        val validDisplayId =
+            if (displayId == INVALID_DISPLAY) {
+                // When a task vanishes it doesn't have a displayId. Find the display of the task.
+                getDisplayIdForTask(taskId)
+            } else {
+                displayId
+            }
+        if (validDisplayId == null) {
+            logW("No display id found for task: taskId=%d", taskId)
+            return
+        }
+        val desk =
+            checkNotNull(desktopData.getDefaultDesk(validDisplayId)) {
+                "Expected a desk in display: $validDisplayId"
+            }
+        updateTaskInDesk(
+            displayId = validDisplayId,
+            deskId = desk.deskId,
+            taskId = taskId,
+            isVisible,
+        )
+    }
+
+    private fun updateTaskInDesk(displayId: Int, deskId: Int, taskId: Int, isVisible: Boolean) {
+        check(displayId != INVALID_DISPLAY) { "Display must be valid" }
+        logD(
+            "updateTaskInDesk taskId=%d, deskId=%d, displayId=%d, isVisible=%b",
+            taskId,
+            deskId,
+            displayId,
+            isVisible,
+        )
 
         if (isVisible) {
-            // If task is visible, remove it from any other display besides [displayId].
-            removeVisibleTask(taskId, excludedDisplayId = displayId)
-        } else if (displayId == INVALID_DISPLAY) {
-            // Task has vanished. Check which display to remove the task from.
-            removeVisibleTask(taskId)
-            return
+            // If task is visible, remove it from any other desk besides [deskId].
+            removeVisibleTask(taskId, excludedDeskId = deskId)
         }
-        val prevCount = getVisibleTaskCount(displayId)
+        val desk = checkNotNull(desktopData.getDesk(deskId)) { "Did not find desk: $deskId" }
+        val prevCount = getVisibleTaskCountInDesk(deskId)
         if (isVisible) {
-            desktopData.getDefaultDesk(displayId)?.visibleTasks?.add(taskId)
-                ?: error("Expected non-null desk in display $displayId")
+            desk.visibleTasks.add(taskId)
             unminimizeTask(displayId, taskId)
         } else {
-            desktopData.getDefaultDesk(displayId)?.visibleTasks?.remove(taskId)
+            desk.visibleTasks.remove(taskId)
         }
-        val newCount = getVisibleTaskCount(displayId)
+        val newCount = getVisibleTaskCount(deskId)
         if (prevCount != newCount) {
             logD(
-                "Update task visibility taskId=%d visible=%b displayId=%d",
+                "Update task visibility taskId=%d visible=%b deskId=%d displayId=%d",
                 taskId,
                 isVisible,
+                deskId,
                 displayId,
             )
             logD("VisibleTaskCount has changed from %d to %d", prevCount, newCount)
@@ -606,33 +641,32 @@ class DesktopRepository(
     /**
      * Gets number of visible freeform tasks on given [displayId]'s active desk.
      *
-     * TODO: b/389960283 - add explicit [deskId] argument.
+     * TODO: b/389960283 - migrate callers to [getVisibleTaskCountInDesk].
      */
     fun getVisibleTaskCount(displayId: Int): Int =
         (desktopData.getActiveDesk(displayId)?.visibleTasks?.size ?: 0).also {
             logD("getVisibleTaskCount=$it")
         }
 
+    /** Gets the number of visible tasks on the given desk. */
+    fun getVisibleTaskCountInDesk(deskId: Int): Int =
+        desktopData.getDesk(deskId)?.visibleTasks?.size ?: 0
+
     /**
      * Adds task (or moves if it already exists) to the top of the ordered list.
      *
      * Unminimizes the task if it is minimized.
-     *
-     * TODO: b/389960283 - add explicit [deskId] argument.
      */
-    private fun addOrMoveFreeformTaskToTop(displayId: Int, taskId: Int) {
-        val desk = getDefaultDesk(displayId) ?: error("Expected a desk in display: $displayId")
-        logD(
-            "Add or move task to top: display=%d taskId=%d deskId=%d",
-            taskId,
-            displayId,
-            desk.deskId,
-        )
+    private fun addOrMoveTaskToTopOfDesk(displayId: Int, deskId: Int, taskId: Int) {
+        val desk = desktopData.getDesk(deskId) ?: error("Could not find desk: $deskId")
+        logD("addOrMoveTaskToTopOfDesk: display=%d deskId=%d taskId=%d", displayId, deskId, taskId)
         desktopData.forAllDesks { _, desk1 -> desk1.freeformTasksInZOrder.remove(taskId) }
         desk.freeformTasksInZOrder.add(0, taskId)
+        // TODO: double check minimization logic.
         // Unminimize the task if it is minimized.
         unminimizeTask(displayId, taskId)
         if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {
+            // TODO: can probably just update the desk.
             updatePersistentRepository(displayId)
         }
     }
@@ -648,6 +682,7 @@ class DesktopRepository(
             // mark it as minimized.
             getDisplayIdForTask(taskId)?.let { minimizeTask(it, taskId) }
                 ?: logW("Minimize task: No display id found for task: taskId=%d", taskId)
+            return
         } else {
             logD("Minimize Task: display=%d, task=%d", displayId, taskId)
             desktopData.getActiveDesk(displayId)?.minimizedTasks?.add(taskId)
@@ -680,7 +715,7 @@ class DesktopRepository(
     private fun getDisplayIdForTask(taskId: Int): Int? {
         var displayForTask: Int? = null
         desktopData.forAllDesks { displayId, desk ->
-            if (taskId in desk.freeformTasksInZOrder) {
+            if (taskId in desk.activeTasks) {
                 displayForTask = displayId
             }
         }
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 645367780390..10ffd79717e3 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
@@ -357,15 +357,15 @@ class DesktopTasksController(
             0 -> return
             // Full screen case
             1 ->
-                moveRunningTaskToDesktop(
-                    allFocusedTasks.single(),
+                moveTaskToDefaultDeskAndActivate(
+                    allFocusedTasks.single().taskId,
                     transitionSource = transitionSource,
                 )
             // Split-screen case where there are two focused tasks, then we find the child
             // task to move to desktop.
             2 ->
-                moveRunningTaskToDesktop(
-                    getSplitFocusedTask(allFocusedTasks[0], allFocusedTasks[1]),
+                moveTaskToDefaultDeskAndActivate(
+                    getSplitFocusedTask(allFocusedTasks[0], allFocusedTasks[1]).taskId,
                     transitionSource = transitionSource,
                 )
             else ->
@@ -428,7 +428,7 @@ class DesktopTasksController(
 
     /** Moves task to desktop mode if task is running, else launches it in desktop mode. */
     @JvmOverloads
-    fun moveTaskToDesktop(
+    fun moveTaskToDefaultDeskAndActivate(
         taskId: Int,
         wct: WindowContainerTransaction = WindowContainerTransaction(),
         transitionSource: DesktopModeTransitionSource,
@@ -436,7 +436,49 @@ class DesktopTasksController(
         callback: IMoveToDesktopCallback? = null,
     ): Boolean {
         val runningTask = shellTaskOrganizer.getRunningTaskInfo(taskId)
-        if (runningTask == null) {
+        val backgroundTask = recentTasksController?.findTaskInBackground(taskId)
+        if (runningTask == null && backgroundTask == null) {
+            logW("moveTaskToDefaultDeskAndActivate taskId=%d not found", taskId)
+            return false
+        }
+        // TODO(342378842): Instead of using default display, support multiple displays
+        val displayId = runningTask?.displayId ?: DEFAULT_DISPLAY
+        val deskId =
+            checkNotNull(taskRepository.getDefaultDeskId(displayId)) {
+                "Expected a default desk to exist"
+            }
+        return moveTaskToDesk(
+            taskId = taskId,
+            deskId = deskId,
+            wct = wct,
+            transitionSource = transitionSource,
+            remoteTransition = remoteTransition,
+        )
+    }
+
+    /** Moves task to desktop mode if task is running, else launches it in desktop mode. */
+    fun moveTaskToDesk(
+        taskId: Int,
+        deskId: Int,
+        wct: WindowContainerTransaction = WindowContainerTransaction(),
+        transitionSource: DesktopModeTransitionSource,
+        remoteTransition: RemoteTransition? = null,
+        callback: IMoveToDesktopCallback? = null,
+    ): Boolean {
+        val runningTask = shellTaskOrganizer.getRunningTaskInfo(taskId)
+        if (runningTask != null) {
+            moveRunningTaskToDesk(
+                task = runningTask,
+                deskId = deskId,
+                wct = wct,
+                transitionSource = transitionSource,
+                remoteTransition = remoteTransition,
+                callback = callback,
+            )
+        }
+        val backgroundTask = recentTasksController?.findTaskInBackground(taskId)
+        if (backgroundTask != null) {
+            // TODO: b/391484662 - add support for |deskId|.
             return moveBackgroundTaskToDesktop(
                 taskId,
                 wct,
@@ -445,8 +487,8 @@ class DesktopTasksController(
                 callback,
             )
         }
-        moveRunningTaskToDesktop(runningTask, wct, transitionSource, remoteTransition, callback)
-        return true
+        logW("moveTaskToDesk taskId=%d not found", taskId)
+        return false
     }
 
     private fun moveBackgroundTaskToDesktop(
@@ -500,8 +542,9 @@ class DesktopTasksController(
     }
 
     /** Moves a running task to desktop. */
-    fun moveRunningTaskToDesktop(
+    private fun moveRunningTaskToDesk(
         task: RunningTaskInfo,
+        deskId: Int,
         wct: WindowContainerTransaction = WindowContainerTransaction(),
         transitionSource: DesktopModeTransitionSource,
         remoteTransition: RemoteTransition? = null,
@@ -511,20 +554,49 @@ class DesktopTasksController(
             logW("Cannot enter desktop for taskId %d, ineligible top activity found", task.taskId)
             return
         }
-        logV("moveRunningTaskToDesktop taskId=%d", task.taskId)
+        val displayId = taskRepository.getDisplayForDesk(deskId)
+        logV(
+            "moveRunningTaskToDesk taskId=%d deskId=%d displayId=%d",
+            task.taskId,
+            deskId,
+            displayId,
+        )
         exitSplitIfApplicable(wct, task)
         val exitResult =
             desktopImmersiveController.exitImmersiveIfApplicable(
                 wct = wct,
-                displayId = task.displayId,
+                displayId = displayId,
                 excludeTaskId = task.taskId,
                 reason = DesktopImmersiveController.ExitReason.TASK_LAUNCH,
             )
 
-        // Bring other apps to front first
         val taskIdToMinimize =
-            bringDesktopAppsToFrontBeforeShowingNewTask(task.displayId, wct, task.taskId)
-        addMoveToDesktopChanges(wct, task)
+            if (Flags.enableMultipleDesktopsBackend()) {
+                // Activate the desk first.
+                prepareForDeskActivation(displayId, wct)
+                desksOrganizer.activateDesk(wct, deskId)
+                if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {
+                    // TODO: 362720497 - do non-running tasks need to be restarted with
+                    // |wct#startTask|?
+                }
+                taskbarDesktopTaskListener?.onTaskbarCornerRoundingUpdate(
+                    doesAnyTaskRequireTaskbarRounding(displayId)
+                )
+                // TODO: 362720497 - activating a desk with the intention to move a new task to it
+                //  means we may need to minimize something in the activating desk. Do so here
+                // similar
+                //  to how it's done in #bringDesktopAppsToFrontBeforeShowingNewTask instead of
+                //  returning null.
+                null
+            } else {
+                // Bring other apps to front first.
+                bringDesktopAppsToFrontBeforeShowingNewTask(task.displayId, wct, task.taskId)
+            }
+        if (Flags.enableMultipleDesktopsBackend()) {
+            prepareMoveTaskToDesk(wct, task, deskId)
+        } else {
+            addMoveToDesktopChanges(wct, task)
+        }
 
         val transition: IBinder
         if (remoteTransition != null) {
@@ -543,6 +615,18 @@ class DesktopTasksController(
             addPendingMinimizeTransition(transition, it, MinimizeReason.TASK_LIMIT)
         }
         exitResult.asExit()?.runOnTransitionStart?.invoke(transition)
+        if (Flags.enableMultipleDesktopsBackend()) {
+            desksTransitionObserver.addPendingTransition(
+                DeskTransition.ActiveDeskWithTask(
+                    token = transition,
+                    displayId = displayId,
+                    deskId = deskId,
+                    enterTaskId = task.taskId,
+                )
+            )
+        } else {
+            taskRepository.setActiveDesk(displayId = displayId, deskId = deskId)
+        }
     }
 
     private fun invokeCallbackToOverview(transition: IBinder, callback: IMoveToDesktopCallback?) {
@@ -2152,6 +2236,7 @@ class DesktopTasksController(
      * different [displayId] if the task should be moved to a different display.
      */
     @VisibleForTesting
+    @Deprecated("Deprecated with multiple desks", ReplaceWith("prepareMoveTaskToDesk()"))
     fun addMoveToDesktopChanges(
         wct: WindowContainerTransaction,
         taskInfo: RunningTaskInfo,
@@ -2179,6 +2264,24 @@ class DesktopTasksController(
         }
     }
 
+    private fun prepareMoveTaskToDesk(
+        wct: WindowContainerTransaction,
+        taskInfo: RunningTaskInfo,
+        deskId: Int,
+    ) {
+        if (!Flags.enableMultipleDesktopsBackend()) return
+        val displayId = taskRepository.getDisplayForDesk(deskId)
+        val displayLayout = displayController.getDisplayLayout(displayId) ?: return
+        val initialBounds = getInitialBounds(displayLayout, taskInfo, displayId)
+        if (canChangeTaskPosition(taskInfo)) {
+            wct.setBounds(taskInfo.token, initialBounds)
+        }
+        desksOrganizer.moveTaskToDesk(wct, deskId = deskId, task = taskInfo)
+        if (useDesktopOverrideDensity()) {
+            wct.setDensityDpi(taskInfo.token, DESKTOP_DENSITY_OVERRIDE)
+        }
+    }
+
     /**
      * Apply changes to move a freeform task from one display to another, which includes handling
      * density changes between displays.
@@ -3173,7 +3276,7 @@ class DesktopTasksController(
             callback: IMoveToDesktopCallback?,
         ) {
             executeRemoteCallWithTaskPermission(controller, "moveTaskToDesktop") { c ->
-                c.moveTaskToDesktop(
+                c.moveTaskToDefaultDeskAndActivate(
                     taskId,
                     transitionSource = transitionSource,
                     remoteTransition = remoteTransition,
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DeskTransition.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DeskTransition.kt
index 4b658caf34ed..8c4fd9db050f 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DeskTransition.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DeskTransition.kt
@@ -34,4 +34,12 @@ sealed class DeskTransition {
     /** A transition to activate a desk in its display. */
     data class ActivateDesk(override val token: IBinder, val displayId: Int, val deskId: Int) :
         DeskTransition()
+
+    /** A transition to activate a desk by moving an outside task to it. */
+    data class ActiveDeskWithTask(
+        override val token: IBinder,
+        val displayId: Int,
+        val deskId: Int,
+        val enterTaskId: Int,
+    ) : DeskTransition()
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserver.kt
index c235f204ece4..6d88c3310a63 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserver.kt
@@ -68,6 +68,26 @@ class DesksTransitionObserver(
                     )
                 }
             }
+            is DeskTransition.ActiveDeskWithTask -> {
+                val withTask =
+                    info.changes.find { change ->
+                        change.taskInfo?.taskId == deskTransition.enterTaskId &&
+                            change.taskInfo?.isVisibleRequested == true &&
+                            desksOrganizer.getDeskAtEnd(change) == deskTransition.deskId
+                    }
+                withTask?.let {
+                    desktopRepository.setActiveDesk(
+                        displayId = deskTransition.displayId,
+                        deskId = deskTransition.deskId,
+                    )
+                    desktopRepository.addTaskToDesk(
+                        displayId = deskTransition.displayId,
+                        deskId = deskTransition.deskId,
+                        taskId = deskTransition.enterTaskId,
+                        isVisible = true,
+                    )
+                }
+            }
         }
     }
 }
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 65c0793aab51..5cda76e2f3e0 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
@@ -89,6 +89,7 @@ class RootTaskDesksOrganizer(
         task: RunningTaskInfo,
     ) {
         val root = roots[deskId] ?: error("Root not found for desk: $deskId")
+        wct.setWindowingMode(task.token, WINDOWING_MODE_UNDEFINED)
         wct.reparent(task.token, root.taskInfo.token, /* onTop= */ true)
     }
 
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 fb4ce13c441f..17d619c9bee8 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
@@ -755,7 +755,7 @@ public class DesktopModeWindowDecorViewModel implements WindowDecorViewModel,
         // App sometimes draws before the insets from WindowDecoration#relayout have
         // been added, so they must be added here
         decoration.addCaptionInset(wct);
-        mDesktopTasksController.moveTaskToDesktop(taskId, wct, source,
+        mDesktopTasksController.moveTaskToDefaultDeskAndActivate(taskId, wct, source,
                 /* remoteTransition= */ null, /* moveToDesktopCallback */ null);
         decoration.closeHandleMenu();
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
index 6a343c56d364..8510441c0557 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopRepositoryTest.kt
@@ -47,6 +47,7 @@ import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.runTest
 import kotlinx.coroutines.test.setMain
 import org.junit.After
+import org.junit.Assert.assertThrows
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -163,6 +164,69 @@ class DesktopRepositoryTest(flags: FlagsParameterization) : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addTask_deskDoesNotExist_throws() {
+        repo.removeDesk(deskId = 0)
+
+        assertThrows(Exception::class.java) {
+            repo.addTask(displayId = DEFAULT_DISPLAY, taskId = 5, isVisible = true)
+        }
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addTaskToDesk_deskDoesNotExist_throws() {
+        repo.removeDesk(deskId = 2)
+
+        assertThrows(Exception::class.java) {
+            repo.addTaskToDesk(
+                displayId = DEFAULT_DISPLAY,
+                deskId = 2,
+                taskId = 4,
+                isVisible = true,
+            )
+        }
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addTaskToDesk_addsToZOrderList() {
+        repo.addDesk(DEFAULT_DISPLAY, deskId = 2)
+        repo.addDesk(DEFAULT_DISPLAY, deskId = 3)
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 2, taskId = 5, isVisible = true)
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 2, taskId = 6, isVisible = true)
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 2, taskId = 7, isVisible = true)
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 3, taskId = 8, isVisible = true)
+
+        val orderedTasks = repo.getFreeformTasksIdsInDeskInZOrder(deskId = 2)
+        assertThat(orderedTasks[0]).isEqualTo(7)
+        assertThat(orderedTasks[1]).isEqualTo(6)
+        assertThat(orderedTasks[2]).isEqualTo(5)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addTaskToDesk_visible_addsToVisible() {
+        repo.addDesk(DEFAULT_DISPLAY, deskId = 2)
+
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 2, taskId = 5, isVisible = true)
+
+        assertThat(repo.isVisibleTask(5)).isTrue()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addTaskToDesk_removesFromAllOtherDesks() {
+        repo.addDesk(DEFAULT_DISPLAY, deskId = 2)
+        repo.addDesk(DEFAULT_DISPLAY, deskId = 3)
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 2, taskId = 7, isVisible = true)
+
+        repo.addTaskToDesk(displayId = DEFAULT_DISPLAY, deskId = 3, taskId = 7, isVisible = true)
+
+        assertThat(repo.getActiveTaskIdsInDesk(2)).doesNotContain(7)
+    }
+
+    @Test
     fun removeActiveTask_notifiesActiveTaskListener() {
         val listener = TestListener()
         repo.addActiveTaskListener(listener)
@@ -467,8 +531,8 @@ class DesktopRepositoryTest(flags: FlagsParameterization) : ShellTestCase() {
         val listener = TestVisibilityListener()
         val executor = TestShellExecutor()
         repo.addVisibleTasksListener(listener, executor)
-        repo.updateTask(DEFAULT_DISPLAY, taskId = 1, isVisible = true)
-        repo.updateTask(DEFAULT_DISPLAY, taskId = 2, isVisible = true)
+        repo.addTask(DEFAULT_DISPLAY, taskId = 1, isVisible = true)
+        repo.addTask(DEFAULT_DISPLAY, taskId = 2, isVisible = true)
         executor.flushAll()
 
         assertThat(listener.visibleTasksCountOnDefaultDisplay).isEqualTo(2)
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 f5e45be96fc7..aa7944cc837f 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
@@ -1405,11 +1405,12 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun moveToDesktop_tdaFullscreen_windowingModeSetToFreeform() {
         val task = setUpFullscreenTask()
         val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
         tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FULLSCREEN
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         val wct = getLatestEnterDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
             .isEqualTo(WINDOWING_MODE_FREEFORM)
@@ -1418,11 +1419,12 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun moveRunningTaskToDesktop_tdaFreeform_windowingModeSetToUndefined() {
         val task = setUpFullscreenTask()
         val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
         tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FREEFORM
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         val wct = getLatestEnterDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
             .isEqualTo(WINDOWING_MODE_UNDEFINED)
@@ -1431,11 +1433,78 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveTaskToDesktop_nonExistentTask_doesNothing() {
-        controller.moveTaskToDesktop(999, transitionSource = UNKNOWN)
-        verifyEnterDesktopWCTNotExecuted()
-        verify(desktopModeEnterExitTransitionListener, times(0))
-            .onEnterDesktopModeTransitionStarted(anyInt())
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_movesTaskToDefaultDesk() {
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, task)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_activatesDesk() {
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).activateDesk(wct, deskId = 0)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_triggersEnterDesktopListener() {
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
+
+        verify(desktopModeEnterExitTransitionListener)
+            .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveTaskToDesk_nonDefaultDesk_movesTaskToDesk() {
+        val transition = Binder()
+        whenever(enterDesktopTransitionHandler.moveToDesktop(any(), any())).thenReturn(transition)
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 3)
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+        task.isVisible = true
+
+        controller.moveTaskToDesk(taskId = task.taskId, deskId = 3, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 3, task)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveTaskToDesk_nonDefaultDesk_activatesDesk() {
+        val transition = Binder()
+        whenever(enterDesktopTransitionHandler.moveToDesktop(any(), any())).thenReturn(transition)
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 3)
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+        task.isVisible = true
+
+        controller.moveTaskToDesk(taskId = task.taskId, deskId = 3, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).activateDesk(wct, deskId = 3)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveTaskToDesk_nonDefaultDesk_triggersEnterDesktopListener() {
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 3)
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+
+        controller.moveTaskToDesk(taskId = task.taskId, deskId = 3, transitionSource = UNKNOWN)
+
+        verify(desktopModeEnterExitTransitionListener)
+            .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
     }
 
     @Test
@@ -1445,7 +1514,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
 
-        controller.moveTaskToDesktop(task.taskId, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
 
         with(getLatestEnterDesktopWct()) {
             assertLaunchTaskAt(0, task.taskId, WINDOWING_MODE_FREEFORM)
@@ -1459,7 +1528,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
 
-        controller.moveTaskToDesktop(task.taskId, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
 
         with(getLatestEnterDesktopWct()) {
             // Add desktop wallpaper activity
@@ -1471,7 +1540,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
-    fun moveRunningTaskToDesktop_topActivityTranslucentWithoutDisplay_taskIsMovedToDesktop() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_topActivityTranslucentWithoutDisplay_taskIsMovedToDesktop_multiDesksDisabled() {
         val task =
             setUpFullscreenTask().apply {
                 isActivityStackTransparent = true
@@ -1479,7 +1549,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
                 numActivities = 1
             }
 
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         val wct = getLatestEnterDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
             .isEqualTo(WINDOWING_MODE_FREEFORM)
@@ -1488,6 +1558,26 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun moveRunningTaskToDesktop_topActivityTranslucentWithoutDisplay_taskIsMovedToDesktop_multiDesksEnabled() {
+        val task =
+            setUpFullscreenTask().apply {
+                isActivityStackTransparent = true
+                isTopActivityNoDisplay = true
+                numActivities = 1
+            }
+
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, task = task)
+        verify(desktopModeEnterExitTransitionListener)
+            .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
     fun moveRunningTaskToDesktop_topActivityTranslucentWithDisplay_doesNothing() {
         val task =
@@ -1497,7 +1587,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
                 numActivities = 1
             }
 
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         verifyEnterDesktopWCTNotExecuted()
         verify(desktopModeEnterExitTransitionListener, times(0))
             .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
@@ -1516,13 +1606,14 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
                 isTopActivityNoDisplay = false
             }
 
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         verifyEnterDesktopWCTNotExecuted()
     }
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
-    fun moveRunningTaskToDesktop_systemUIActivityWithoutDisplay_doesNothing() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_systemUIActivityWithoutDisplay_doesNothing_multiDesksDisabled() {
         // Set task as systemUI package
         val systemUIPackageName =
             context.resources.getString(com.android.internal.R.string.config_systemUi)
@@ -1533,7 +1624,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
                 isTopActivityNoDisplay = true
             }
 
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
 
         val wct = getLatestEnterDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
@@ -1552,7 +1643,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         mContext.setMockPackageManager(packageManager)
         whenever(packageManager.getHomeActivities(any())).thenReturn(homeActivities)
 
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         verifyEnterDesktopWCTNotExecuted()
     }
 
@@ -1568,7 +1659,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         mContext.setMockPackageManager(packageManager)
         whenever(packageManager.getHomeActivities(any())).thenReturn(homeActivities)
 
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
 
         val wct = getLatestEnterDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
@@ -1576,6 +1667,28 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun moveRunningTaskToDesktop_systemUIActivityWithoutDisplay_doesNothing_multiDesksEnabled() {
+        // Set task as systemUI package
+        val systemUIPackageName =
+            context.resources.getString(com.android.internal.R.string.config_systemUi)
+        val baseComponent = ComponentName(systemUIPackageName, /* cls= */ "")
+        val task =
+            setUpFullscreenTask().apply {
+                baseActivity = baseComponent
+                isTopActivityNoDisplay = true
+            }
+
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, task = task)
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun moveBackgroundTaskToDesktop_remoteTransition_usesOneShotHandler() {
         val transitionHandlerArgCaptor = ArgumentCaptor.forClass(TransitionHandler::class.java)
@@ -1585,7 +1698,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         val task = createTaskInfo(1)
         whenever(shellTaskOrganizer.getRunningTaskInfo(anyInt())).thenReturn(null)
         whenever(recentTasksController.findTaskInBackground(anyInt())).thenReturn(task)
-        controller.moveTaskToDesktop(
+        controller.moveTaskToDefaultDeskAndActivate(
             taskId = task.taskId,
             transitionSource = UNKNOWN,
             remoteTransition = RemoteTransition(spy(TestRemoteTransition())),
@@ -1602,8 +1715,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         whenever(transitions.startTransition(anyInt(), any(), transitionHandlerArgCaptor.capture()))
             .thenReturn(Binder())
 
-        controller.moveRunningTaskToDesktop(
-            task = setUpFullscreenTask(),
+        controller.moveTaskToDefaultDeskAndActivate(
+            taskId = setUpFullscreenTask().taskId,
             transitionSource = UNKNOWN,
             remoteTransition = RemoteTransition(spy(TestRemoteTransition())),
         )
@@ -1614,14 +1727,20 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
     fun moveRunningTaskToDesktop_otherFreeformTasksBroughtToFront_desktopWallpaperDisabled() {
         val homeTask = setUpHomeTask()
         val freeformTask = setUpFreeformTask()
         val fullscreenTask = setUpFullscreenTask()
         markTaskHidden(freeformTask)
 
-        controller.moveRunningTaskToDesktop(fullscreenTask, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(
+            fullscreenTask.taskId,
+            transitionSource = UNKNOWN,
+        )
 
         with(getLatestEnterDesktopWct()) {
             // Operations should include home task, freeform task
@@ -1636,12 +1755,16 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun moveRunningTaskToDesktop_otherFreeformTasksBroughtToFront_desktopWallpaperEnabled() {
         val freeformTask = setUpFreeformTask()
         val fullscreenTask = setUpFullscreenTask()
         markTaskHidden(freeformTask)
 
-        controller.moveRunningTaskToDesktop(fullscreenTask, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(
+            fullscreenTask.taskId,
+            transitionSource = UNKNOWN,
+        )
 
         with(getLatestEnterDesktopWct()) {
             // Operations should include wallpaper intent, freeform task, fullscreen task
@@ -1657,6 +1780,43 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun moveRunningTaskToDesktop_desktopWallpaperEnabled_multiDesksEnabled() {
+        val freeformTask = setUpFreeformTask()
+        val fullscreenTask = setUpFullscreenTask()
+        markTaskHidden(freeformTask)
+
+        controller.moveTaskToDefaultDeskAndActivate(
+            fullscreenTask.taskId,
+            transitionSource = UNKNOWN,
+        )
+
+        val wct = getLatestEnterDesktopWct()
+        wct.assertPendingIntentAt(index = 0, desktopWallpaperIntent)
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, fullscreenTask)
+        verify(desksOrganizer).activateDesk(wct, deskId = 0)
+        verify(desktopModeEnterExitTransitionListener)
+            .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_activatesDesk_desktopWallpaperEnabled_multiDesksDisabled() {
+        val fullscreenTask = setUpFullscreenTask()
+
+        controller.moveTaskToDefaultDeskAndActivate(
+            fullscreenTask.taskId,
+            transitionSource = UNKNOWN,
+        )
+
+        assertThat(taskRepository.getActiveDeskId(DEFAULT_DISPLAY)).isEqualTo(DEFAULT_DISPLAY)
+    }
+
+    @Test
     fun moveRunningTaskToDesktop_onlyFreeformTasksFromCurrentDisplayBroughtToFront() {
         setUpHomeTask(displayId = DEFAULT_DISPLAY)
         val freeformTaskDefault = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
@@ -1668,7 +1828,10 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         val freeformTaskSecond = setUpFreeformTask(displayId = SECOND_DISPLAY)
         markTaskHidden(freeformTaskSecond)
 
-        controller.moveRunningTaskToDesktop(fullscreenTaskDefault, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(
+            fullscreenTaskDefault.taskId,
+            transitionSource = UNKNOWN,
+        )
 
         with(getLatestEnterDesktopWct()) {
             // Check that hierarchy operations do not include tasks from second display
@@ -1682,9 +1845,10 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveRunningTaskToDesktop_splitTaskExitsSplit() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_splitTaskExitsSplit_multiDesksDisabled() {
         val task = setUpSplitScreenTask()
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         val wct = getLatestEnterDesktopWct()
         assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
             .isEqualTo(WINDOWING_MODE_FREEFORM)
@@ -1699,12 +1863,27 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveRunningTaskToDesktop_splitTaskExitsSplit_multiDesksEnabled() {
+        val task = setUpSplitScreenTask()
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, task)
+        verify(desktopModeEnterExitTransitionListener)
+            .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
+        verify(splitScreenController)
+            .prepareExitSplitScreen(
+                any(),
+                anyInt(),
+                eq(SplitScreenController.EXIT_REASON_DESKTOP_MODE),
+            )
+    }
+
+    @Test
     fun moveRunningTaskToDesktop_fullscreenTaskDoesNotExitSplit() {
         val task = setUpFullscreenTask()
-        controller.moveRunningTaskToDesktop(task, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(task.taskId, transitionSource = UNKNOWN)
         val wct = getLatestEnterDesktopWct()
-        assertThat(wct.changes[task.token.asBinder()]?.windowingMode)
-            .isEqualTo(WINDOWING_MODE_FREEFORM)
         verify(desktopModeEnterExitTransitionListener)
             .onEnterDesktopModeTransitionStarted(FREEFORM_ANIMATION_DURATION)
         verify(splitScreenController, never())
@@ -1716,13 +1895,16 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
     fun moveRunningTaskToDesktop_desktopWallpaperDisabled_bringsTasksOver_dontShowBackTask() {
         val freeformTasks = (1..MAX_TASK_LIMIT).map { _ -> setUpFreeformTask() }
         val newTask = setUpFullscreenTask()
         val homeTask = setUpHomeTask()
 
-        controller.moveRunningTaskToDesktop(newTask, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(newTask.taskId, transitionSource = UNKNOWN)
 
         val wct = getLatestEnterDesktopWct()
         verify(desktopModeEnterExitTransitionListener)
@@ -1738,12 +1920,13 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun moveRunningTaskToDesktop_desktopWallpaperEnabled_bringsTasksOverLimit_dontShowBackTask() {
         val freeformTasks = (1..MAX_TASK_LIMIT).map { _ -> setUpFreeformTask() }
         val newTask = setUpFullscreenTask()
         val homeTask = setUpHomeTask()
 
-        controller.moveRunningTaskToDesktop(newTask, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(newTask.taskId, transitionSource = UNKNOWN)
 
         val wct = getLatestEnterDesktopWct()
         verify(desktopModeEnterExitTransitionListener)
@@ -3564,7 +3747,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveFocusedTaskToDesktop_fullscreenTaskIsMovedToDesktop() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveFocusedTaskToDesktop_fullscreenTaskIsMovedToDesktop_multiDesksDisabled() {
         val task1 = setUpFullscreenTask()
         val task2 = setUpFullscreenTask()
         val task3 = setUpFullscreenTask()
@@ -3581,7 +3765,25 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveFocusedTaskToDesktop_splitScreenTaskIsMovedToDesktop() {
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveFocusedTaskToDesktop_fullscreenTaskIsMovedToDesktop_multiDesksEnabled() {
+        val task1 = setUpFullscreenTask()
+        val task2 = setUpFullscreenTask()
+        val task3 = setUpFullscreenTask()
+
+        task1.isFocused = true
+        task2.isFocused = false
+        task3.isFocused = false
+
+        controller.moveFocusedTaskToDesktop(DEFAULT_DISPLAY, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, task1)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveFocusedTaskToDesktop_splitScreenTaskIsMovedToDesktop_multiDesksDisabled() {
         val task1 = setUpSplitScreenTask()
         val task2 = setUpFullscreenTask()
         val task3 = setUpFullscreenTask()
@@ -3608,6 +3810,33 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveFocusedTaskToDesktop_splitScreenTaskIsMovedToDesktop_multiDesksEnabled() {
+        val task1 = setUpSplitScreenTask()
+        val task2 = setUpFullscreenTask()
+        val task3 = setUpFullscreenTask()
+        val task4 = setUpSplitScreenTask()
+
+        task1.isFocused = true
+        task2.isFocused = false
+        task3.isFocused = false
+        task4.isFocused = true
+
+        task4.parentTaskId = task1.taskId
+
+        controller.moveFocusedTaskToDesktop(DEFAULT_DISPLAY, transitionSource = UNKNOWN)
+
+        val wct = getLatestEnterDesktopWct()
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId = 0, task4)
+        verify(splitScreenController)
+            .prepareExitSplitScreen(
+                any(),
+                anyInt(),
+                eq(SplitScreenController.EXIT_REASON_DESKTOP_MODE),
+            )
+    }
+
+    @Test
     fun moveFocusedTaskToFullscreen() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
@@ -3784,6 +4013,31 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun moveTaskToDesk_multipleDesks_addsPendingTransition() {
+        val transition = Binder()
+        whenever(enterDesktopTransitionHandler.moveToDesktop(any(), any())).thenReturn(transition)
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 3)
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+        task.isVisible = true
+
+        controller.moveTaskToDesk(taskId = task.taskId, deskId = 3, transitionSource = UNKNOWN)
+
+        verify(desksTransitionsObserver)
+            .addPendingTransition(
+                argThat {
+                    this is DeskTransition.ActiveDeskWithTask &&
+                        this.token == transition &&
+                        this.deskId == 3 &&
+                        this.enterTaskId == task.taskId
+                }
+            )
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_WINDOWING_DYNAMIC_INITIAL_BOUNDS)
     fun dragToDesktop_landscapeDevice_resizable_undefinedOrientation_defaultLandscapeBounds() {
         val spyController = spy(controller)
@@ -5154,7 +5408,11 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
             .thenReturn(ExitResult.Exit(exitingTask = 5, runOnTransitionStart = runOnStartTransit))
         whenever(enterDesktopTransitionHandler.moveToDesktop(wct, UNKNOWN)).thenReturn(transition)
 
-        controller.moveTaskToDesktop(taskId = task.taskId, wct = wct, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(
+            taskId = task.taskId,
+            wct = wct,
+            transitionSource = UNKNOWN,
+        )
 
         verify(mMockDesktopImmersiveController)
             .exitImmersiveIfApplicable(eq(wct), eq(task.displayId), eq(task.taskId), any())
@@ -5178,7 +5436,11 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
             .thenReturn(ExitResult.Exit(exitingTask = 5, runOnTransitionStart = runOnStartTransit))
         whenever(enterDesktopTransitionHandler.moveToDesktop(wct, UNKNOWN)).thenReturn(transition)
 
-        controller.moveTaskToDesktop(taskId = task.taskId, wct = wct, transitionSource = UNKNOWN)
+        controller.moveTaskToDefaultDeskAndActivate(
+            taskId = task.taskId,
+            wct = wct,
+            transitionSource = UNKNOWN,
+        )
 
         verify(mMockDesktopImmersiveController)
             .exitImmersiveIfApplicable(eq(wct), eq(task.displayId), eq(task.taskId), any())
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserverTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserverTest.kt
index 00c54c2ecc18..9f09e3f57927 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserverTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/multidesks/DesksTransitionObserverTest.kt
@@ -29,9 +29,12 @@ import com.android.window.flags.Flags
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.TestShellExecutor
 import com.android.wm.shell.desktopmode.DesktopRepository
+import com.android.wm.shell.desktopmode.DesktopTestHelpers.createFreeformTask
 import com.android.wm.shell.desktopmode.DesktopUserRepositories
 import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -51,6 +54,7 @@ class DesksTransitionObserverTest : ShellTestCase() {
     @JvmField @Rule val setFlagsRule = SetFlagsRule()
 
     private val mockDesksOrganizer = mock<DesksOrganizer>()
+    val testScope = TestScope()
 
     private lateinit var desktopUserRepositories: DesktopUserRepositories
     private lateinit var observer: DesksTransitionObserver
@@ -67,7 +71,7 @@ class DesksTransitionObserverTest : ShellTestCase() {
                 /* shellController= */ mock(),
                 /* persistentRepository= */ mock(),
                 /* repositoryInitializer= */ mock(),
-                /* mainCoroutineScope= */ mock(),
+                testScope,
                 /* userManager= */ mock(),
             )
         observer = DesksTransitionObserver(desktopUserRepositories, mockDesksOrganizer)
@@ -145,4 +149,32 @@ class DesksTransitionObserverTest : ShellTestCase() {
 
         assertThat(repository.getActiveDeskId(DEFAULT_DISPLAY)).isEqualTo(5)
     }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun onTransitionReady_activateDeskWithTask_updatesRepository() =
+        testScope.runTest {
+            val deskId = 5
+            val task = createFreeformTask(DEFAULT_DISPLAY).apply { isVisibleRequested = true }
+            val transition = Binder()
+            val change = Change(mock(), mock()).apply { taskInfo = task }
+            whenever(mockDesksOrganizer.getDeskAtEnd(change)).thenReturn(deskId)
+            val activateTransition =
+                DeskTransition.ActiveDeskWithTask(
+                    transition,
+                    displayId = DEFAULT_DISPLAY,
+                    deskId = deskId,
+                    enterTaskId = task.taskId,
+                )
+            repository.addDesk(DEFAULT_DISPLAY, deskId = deskId)
+
+            observer.addPendingTransition(activateTransition)
+            observer.onTransitionReady(
+                transition = transition,
+                info = TransitionInfo(TRANSIT_TO_FRONT, /* flags= */ 0).apply { addChange(change) },
+            )
+
+            assertThat(repository.getActiveDeskId(DEFAULT_DISPLAY)).isEqualTo(deskId)
+            assertThat(repository.getActiveTaskIdsInDesk(deskId)).contains(task.taskId)
+        }
 }
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 fbc78b21353e..4d4b15389eca 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.WindowConfiguration.WINDOWING_MODE_UNDEFINED
 import android.testing.AndroidTestingRunner
 import android.view.Display
 import android.view.SurfaceControl
@@ -217,6 +218,13 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
                 }
             )
             .isTrue()
+        assertThat(
+                wct.changes.any { change ->
+                    change.key == desktopTask.token.asBinder() &&
+                        change.value.windowingMode == WINDOWING_MODE_UNDEFINED
+                }
+            )
+            .isTrue()
     }
 
     @Test
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
index 4a91fb429f7b..f15418adf1e3 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/windowdecor/DesktopModeWindowDecorViewModelTests.kt
@@ -639,7 +639,7 @@ class DesktopModeWindowDecorViewModelTests : DesktopModeWindowDecorViewModelTest
 
         toDesktopListenerCaptor.value.accept(DesktopModeTransitionSource.APP_HANDLE_MENU_BUTTON)
 
-        verify(mockDesktopTasksController).moveTaskToDesktop(
+        verify(mockDesktopTasksController).moveTaskToDefaultDeskAndActivate(
             eq(decor.mTaskInfo.taskId),
             any(),
             eq(DesktopModeTransitionSource.APP_HANDLE_MENU_BUTTON),
@@ -877,7 +877,7 @@ class DesktopModeWindowDecorViewModelTests : DesktopModeWindowDecorViewModelTest
         )
 
         verify(mockDesktopTasksController, times(1))
-            .moveTaskToDesktop(any(), any(), any(), anyOrNull(), anyOrNull())
+            .moveTaskToDefaultDeskAndActivate(any(), any(), any(), anyOrNull(), anyOrNull())
     }
 
     @Test