Merge changes I05ef9baf,Idf38fef3 into main

* changes:
  [40/N] Desks: Fix failing WMShellUnitTests
  [39/N] Desks: Restore persisted desks on reboot

14 files changed, 930 insertions, 274 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 478fd6d83661..92c3020a14e6 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
@@ -759,6 +759,7 @@ public abstract class WMShellModule {
             ToggleResizeDesktopTaskTransitionHandler toggleResizeDesktopTaskTransitionHandler,
             DragToDesktopTransitionHandler dragToDesktopTransitionHandler,
             @DynamicOverride DesktopUserRepositories desktopUserRepositories,
+            DesktopRepositoryInitializer desktopRepositoryInitializer,
             Optional<DesktopImmersiveController> desktopImmersiveController,
             DesktopModeLoggerTransitionObserver desktopModeLoggerTransitionObserver,
             LaunchAdjacentController launchAdjacentController,
@@ -805,6 +806,7 @@ public abstract class WMShellModule {
                 dragToDesktopTransitionHandler,
                 desktopImmersiveController.get(),
                 desktopUserRepositories,
+                desktopRepositoryInitializer,
                 recentsTransitionHandler,
                 multiInstanceHelper,
                 mainExecutor,
@@ -1313,10 +1315,12 @@ public abstract class WMShellModule {
     static Optional<DesktopDisplayEventHandler> provideDesktopDisplayEventHandler(
             Context context,
             ShellInit shellInit,
+            @ShellMainThread CoroutineScope mainScope,
             DisplayController displayController,
             Optional<DesktopUserRepositories> desktopUserRepositories,
             Optional<DesktopTasksController> desktopTasksController,
-            Optional<DesktopDisplayModeController> desktopDisplayModeController
+            Optional<DesktopDisplayModeController> desktopDisplayModeController,
+            DesktopRepositoryInitializer desktopRepositoryInitializer
     ) {
         if (!DesktopModeStatus.canEnterDesktopMode(context)) {
             return Optional.empty();
@@ -1325,7 +1329,9 @@ public abstract class WMShellModule {
                 new DesktopDisplayEventHandler(
                         context,
                         shellInit,
+                        mainScope,
                         displayController,
+                        desktopRepositoryInitializer,
                         desktopUserRepositories.get(),
                         desktopTasksController.get(),
                         desktopDisplayModeController.get()));
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandler.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandler.kt
index afc48acad4f5..683b74392fa6 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandler.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandler.kt
@@ -23,15 +23,21 @@ import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.common.DisplayController
 import com.android.wm.shell.common.DisplayController.OnDisplaysChangedListener
 import com.android.wm.shell.desktopmode.multidesks.OnDeskRemovedListener
+import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.launch
 
 /** Handles display events in desktop mode */
 class DesktopDisplayEventHandler(
     private val context: Context,
     shellInit: ShellInit,
+    private val mainScope: CoroutineScope,
     private val displayController: DisplayController,
+    private val desktopRepositoryInitializer: DesktopRepositoryInitializer,
     private val desktopUserRepositories: DesktopUserRepositories,
     private val desktopTasksController: DesktopTasksController,
     private val desktopDisplayModeController: DesktopDisplayModeController,
@@ -61,15 +67,19 @@ class DesktopDisplayEventHandler(
             logV("Display #$displayId does not support desks")
             return
         }
-        logV("Creating new desk in new display#$displayId")
-        // TODO: b/362720497 - when SystemUI crashes with a freeform task open for any reason, the
-        //  task is recreated and received in [FreeformTaskListener] before this display callback
-        //  is invoked, which results in the repository trying to add the task to a desk before the
-        //  desk has been recreated here, which may result in a crash-loop if the repository is
-        //  checking that the desk exists before adding a task to it. See b/391984373.
-        desktopTasksController.createDesk(displayId)
-        // TODO: b/393978539 - consider activating the desk on creation when applicable, such as
-        //  for connected displays.
+
+        mainScope.launch {
+            desktopRepositoryInitializer.isInitialized.collect { initialized ->
+                if (!initialized) return@collect
+                if (desktopRepository.getNumberOfDesks(displayId) == 0) {
+                    logV("Creating new desk in new display#$displayId")
+                    // TODO: b/393978539 - consider activating the desk on creation when
+                    //  applicable, such as for connected displays.
+                    desktopTasksController.createDesk(displayId)
+                }
+                cancel()
+            }
+        }
     }
 
     override fun onDisplayRemoved(displayId: Int) {
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 f64bd757de3b..8636bc1f56c2 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
@@ -716,12 +716,15 @@ class DesktopRepository(
     }
 
     /**
-     * Returns the top transparent fullscreen task id for a given display's active desk, or null.
+     * Returns the top transparent fullscreen task id for a given display, or null.
      *
      * TODO: b/389960283 - add explicit [deskId] argument.
      */
     fun getTopTransparentFullscreenTaskId(displayId: Int): Int? =
-        desktopData.getActiveDesk(displayId)?.topTransparentFullscreenTaskId
+        desktopData
+            .desksSequence(displayId)
+            .mapNotNull { it.topTransparentFullscreenTaskId }
+            .firstOrNull()
 
     /**
      * Clears the top transparent fullscreen task id info for a given display's active desk.
@@ -818,7 +821,6 @@ class DesktopRepository(
     }
 
     /** Minimizes the task in its desk. */
-    @VisibleForTesting
     fun minimizeTaskInDesk(displayId: Int, deskId: Int, taskId: Int) {
         logD("MinimizeTaskInDesk: displayId=%d deskId=%d, task=%d", displayId, deskId, taskId)
         desktopData.getDesk(deskId)?.minimizedTasks?.add(taskId)
@@ -933,6 +935,12 @@ class DesktopRepository(
                 listener.onDeskRemoved(displayId = desk.displayId, deskId = desk.deskId)
             }
         }
+        if (
+            DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue &&
+                DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue
+        ) {
+            removeDeskFromPersistentRepository(desk)
+        }
         return activeTasks
     }
 
@@ -1031,6 +1039,24 @@ class DesktopRepository(
         }
     }
 
+    private fun removeDeskFromPersistentRepository(desk: Desk) {
+        mainCoroutineScope.launch {
+            try {
+                logD(
+                    "updatePersistentRepositoryForRemovedDesk user=%d desk=%d",
+                    userId,
+                    desk.deskId,
+                )
+                persistentRepository.removeDesktop(userId = userId, desktopId = desk.deskId)
+            } catch (throwable: Throwable) {
+                logE(
+                    "An exception occurred while updating the persistent repository \n%s",
+                    throwable.stackTrace,
+                )
+            }
+        }
+    }
+
     internal fun dump(pw: PrintWriter, prefix: String) {
         val innerPrefix = "$prefix  "
         pw.println("${prefix}DesktopRepository")
@@ -1049,6 +1075,7 @@ class DesktopRepository(
             }
             .forEach { (displayId, activeDeskId, desks) ->
                 pw.println("${prefix}Display #$displayId:")
+                pw.println("${innerPrefix}numOfDesks=${desks.size}")
                 pw.println("${innerPrefix}activeDesk=$activeDeskId")
                 pw.println("${innerPrefix}desks:")
                 val desksPrefix = "$innerPrefix  "
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 dbc599be57af..d0356d55035d 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
@@ -115,6 +115,9 @@ import com.android.wm.shell.desktopmode.multidesks.DeskTransition
 import com.android.wm.shell.desktopmode.multidesks.DesksOrganizer
 import com.android.wm.shell.desktopmode.multidesks.DesksTransitionObserver
 import com.android.wm.shell.desktopmode.multidesks.OnDeskRemovedListener
+import com.android.wm.shell.desktopmode.multidesks.createDesk
+import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer
+import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer.DeskRecreationFactory
 import com.android.wm.shell.draganddrop.DragAndDropController
 import com.android.wm.shell.freeform.FreeformTaskTransitionStarter
 import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
@@ -194,6 +197,7 @@ class DesktopTasksController(
     private val dragToDesktopTransitionHandler: DragToDesktopTransitionHandler,
     private val desktopImmersiveController: DesktopImmersiveController,
     private val userRepositories: DesktopUserRepositories,
+    desktopRepositoryInitializer: DesktopRepositoryInitializer,
     private val recentsTransitionHandler: RecentsTransitionHandler,
     private val multiInstanceHelper: MultiInstanceHelper,
     @ShellMainThread private val mainExecutor: ShellExecutor,
@@ -276,6 +280,19 @@ class DesktopTasksController(
         }
         userId = ActivityManager.getCurrentUser()
         taskRepository = userRepositories.getProfile(userId)
+
+        if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+            desktopRepositoryInitializer.deskRecreationFactory =
+                DeskRecreationFactory { deskUserId, destinationDisplayId, deskId ->
+                    if (deskUserId != userId) {
+                        // TODO: b/400984250 - add multi-user support for multi-desk restoration.
+                        logW("Tried to recreated desk of another user.")
+                        deskId
+                    } else {
+                        desksOrganizer.createDesk(destinationDisplayId)
+                    }
+                }
+        }
     }
 
     private fun onInit() {
@@ -1718,15 +1735,7 @@ class DesktopTasksController(
                     wct.reorder(runningTaskInfo.token, /* onTop= */ true)
                 } else if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) {
                     // Task is not running, start it
-                    wct.startTask(
-                        taskId,
-                        ActivityOptions.makeBasic()
-                            .apply {
-                                launchWindowingMode = WINDOWING_MODE_FREEFORM
-                                splashScreenStyle = SPLASH_SCREEN_STYLE_ICON
-                            }
-                            .toBundle(),
-                    )
+                    wct.startTask(taskId, createActivityOptionsForStartTask().toBundle())
                 }
             }
 
@@ -2826,9 +2835,6 @@ class DesktopTasksController(
         }
         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)
         )
@@ -2846,6 +2852,19 @@ class DesktopTasksController(
                 desksOrganizer.minimizeTask(wct, deskId, taskToMinimize)
             }
         }
+        if (DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {
+            expandedTasksOrderedFrontToBack
+                .filter { taskId -> taskId != taskIdToMinimize }
+                .reversed()
+                .forEach { taskId ->
+                    val runningTaskInfo = shellTaskOrganizer.getRunningTaskInfo(taskId)
+                    if (runningTaskInfo == null) {
+                        wct.startTask(taskId, createActivityOptionsForStartTask().toBundle())
+                    } else {
+                        desksOrganizer.reorderTaskToFront(wct, deskId, runningTaskInfo)
+                    }
+                }
+        }
         return { transition ->
             val activateDeskTransition =
                 if (newTaskIdInFront != null) {
@@ -3497,6 +3516,13 @@ class DesktopTasksController(
         }
     }
 
+    private fun createActivityOptionsForStartTask(): ActivityOptions {
+        return ActivityOptions.makeBasic().apply {
+            launchWindowingMode = WINDOWING_MODE_FREEFORM
+            splashScreenStyle = SPLASH_SCREEN_STYLE_ICON
+        }
+    }
+
     private fun dump(pw: PrintWriter, prefix: String) {
         val innerPrefix = "$prefix  "
         pw.println("${prefix}DesktopTasksController")
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 8c4bc2598dff..5a988fcd1b77 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
@@ -18,6 +18,8 @@ package com.android.wm.shell.desktopmode.multidesks
 import android.app.ActivityManager
 import android.window.TransitionInfo
 import android.window.WindowContainerTransaction
+import com.android.wm.shell.desktopmode.multidesks.DesksOrganizer.OnCreateCallback
+import kotlin.coroutines.suspendCoroutine
 
 /** An organizer of desk containers in which to host child desktop windows. */
 interface DesksOrganizer {
@@ -82,3 +84,9 @@ interface DesksOrganizer {
         fun onCreated(deskId: Int)
     }
 }
+
+/** Creates a new desk container in the given display. */
+suspend fun DesksOrganizer.createDesk(displayId: Int): Int = suspendCoroutine { cont ->
+    val onCreateCallback = OnCreateCallback { deskId -> cont.resumeWith(Result.success(deskId)) }
+    createDesk(displayId, onCreateCallback)
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
index 1566544f5303..f71eacab518d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopPersistentRepository.kt
@@ -132,10 +132,7 @@ class DesktopPersistentRepository(private val dataStore: DataStore<DesktopPersis
                     .toBuilder()
                     .putDesktopRepoByUser(
                         userId,
-                        currentRepository
-                            .toBuilder()
-                            .putDesktop(desktopId, desktop.build())
-                            .build(),
+                        currentRepository.toBuilder().putDesktop(desktopId, desktop.build()).build(),
                     )
                     .build()
             }
@@ -149,6 +146,33 @@ class DesktopPersistentRepository(private val dataStore: DataStore<DesktopPersis
         }
     }
 
+    /** Removes the desktop from the persistent repository. */
+    suspend fun removeDesktop(userId: Int, desktopId: Int) {
+        try {
+            dataStore.updateData { persistentRepositories: DesktopPersistentRepositories ->
+                val currentRepository =
+                    persistentRepositories.getDesktopRepoByUserOrDefault(
+                        userId,
+                        DesktopRepositoryState.getDefaultInstance(),
+                    )
+                persistentRepositories
+                    .toBuilder()
+                    .putDesktopRepoByUser(
+                        userId,
+                        currentRepository.toBuilder().removeDesktop(desktopId).build(),
+                    )
+                    .build()
+            }
+        } catch (throwable: Throwable) {
+            Log.e(
+                TAG,
+                "Error in removing desktop related data, data is " +
+                    "stored in a file named $DESKTOP_REPOSITORIES_DATASTORE_FILE",
+                throwable,
+            )
+        }
+    }
+
     suspend fun removeUsers(uids: List<Int>) {
         try {
             dataStore.updateData { persistentRepositories: DesktopPersistentRepositories ->
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializer.kt
index a26ebbf4c99a..8191181cac11 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializer.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializer.kt
@@ -17,8 +17,22 @@
 package com.android.wm.shell.desktopmode.persistence
 
 import com.android.wm.shell.desktopmode.DesktopUserRepositories
+import kotlinx.coroutines.flow.StateFlow
 
 /** Interface for initializing the [DesktopUserRepositories]. */
-fun interface DesktopRepositoryInitializer {
+interface DesktopRepositoryInitializer {
+    /** A factory used to recreate a desk from persistence. */
+    var deskRecreationFactory: DeskRecreationFactory
+
+    /** A flow that emits true when the repository has been initialized. */
+    val isInitialized: StateFlow<Boolean>
+
+    /** Initialize the user repositories from a persistent data store. */
     fun initialize(userRepositories: DesktopUserRepositories)
+
+    /** A factory for recreating desks. */
+    fun interface DeskRecreationFactory {
+        /** Recreates a restored desk and returns the new desk id. */
+        suspend fun recreateDesk(userId: Int, destinationDisplayId: Int, deskId: Int): Int
+    }
 }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerImpl.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerImpl.kt
index 0507e59c06e1..49cb7391fe97 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerImpl.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerImpl.kt
@@ -17,13 +17,19 @@
 package com.android.wm.shell.desktopmode.persistence
 
 import android.content.Context
+import android.view.Display
 import android.window.DesktopExperienceFlags
 import android.window.DesktopModeFlags
+import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.desktopmode.DesktopRepository
 import com.android.wm.shell.desktopmode.DesktopUserRepositories
+import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer.DeskRecreationFactory
+import com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_DESKTOP_MODE
 import com.android.wm.shell.shared.annotations.ShellMainThread
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.launch
 
 /**
@@ -37,62 +43,136 @@ class DesktopRepositoryInitializerImpl(
     private val persistentRepository: DesktopPersistentRepository,
     @ShellMainThread private val mainCoroutineScope: CoroutineScope,
 ) : DesktopRepositoryInitializer {
+
+    override var deskRecreationFactory: DeskRecreationFactory = DefaultDeskRecreationFactory()
+
+    private val _isInitialized = MutableStateFlow(false)
+    override val isInitialized: StateFlow<Boolean> = _isInitialized
+
     override fun initialize(userRepositories: DesktopUserRepositories) {
-        if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue()) return
+        if (!DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_PERSISTENCE.isTrue) {
+            _isInitialized.value = true
+            return
+        }
         //  TODO: b/365962554 - Handle the case that user moves to desktop before it's initialized
         mainCoroutineScope.launch {
-            val desktopUserPersistentRepositoryMap =
-                persistentRepository.getUserDesktopRepositoryMap() ?: return@launch
-            for (userId in desktopUserPersistentRepositoryMap.keys) {
-                val repository = userRepositories.getProfile(userId)
-                val desktopRepositoryState =
-                    persistentRepository.getDesktopRepositoryState(userId) ?: continue
-                val desktopByDesktopIdMap = desktopRepositoryState.desktopMap
-                for (desktopId in desktopByDesktopIdMap.keys) {
-                    val persistentDesktop =
-                        persistentRepository.readDesktop(userId, desktopId) ?: continue
-                    val maxTasks =
-                        DesktopModeStatus.getMaxTaskLimit(context).takeIf { it > 0 }
-                            ?: persistentDesktop.zOrderedTasksCount
-                    var visibleTasksCount = 0
-                    repository.addDesk(
-                        displayId = persistentDesktop.displayId,
-                        deskId =
-                            if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-                                persistentDesktop.desktopId
-                            } else {
-                                // When disabled, desk ids are always the display id.
-                                persistentDesktop.displayId
-                            },
+            try {
+                val desktopUserPersistentRepositoryMap =
+                    persistentRepository.getUserDesktopRepositoryMap() ?: return@launch
+                for (userId in desktopUserPersistentRepositoryMap.keys) {
+                    val repository = userRepositories.getProfile(userId)
+                    val desktopRepositoryState =
+                        persistentRepository.getDesktopRepositoryState(userId) ?: continue
+                    val desksToRestore = getDesksToRestore(desktopRepositoryState, userId)
+                    logV(
+                        "initialize() will restore desks=%s user=%d",
+                        desksToRestore.map { it.desktopId },
+                        userId,
                     )
-                    persistentDesktop.zOrderedTasksList
-                        // Reverse it so we initialize the repo from bottom to top.
-                        .reversed()
-                        .mapNotNull { taskId -> persistentDesktop.tasksByTaskIdMap[taskId] }
-                        // TODO: b/362720497 - add tasks to their respective desk when multi-desk
-                        //   persistence is implemented.
-                        .forEach { task ->
-                            if (
-                                task.desktopTaskState == DesktopTaskState.VISIBLE &&
-                                    visibleTasksCount < maxTasks
-                            ) {
-                                visibleTasksCount++
-                                repository.addTask(
-                                    persistentDesktop.displayId,
-                                    task.taskId,
-                                    isVisible = false,
-                                )
-                            } else {
-                                repository.addTask(
-                                    persistentDesktop.displayId,
-                                    task.taskId,
+                    desksToRestore.forEach { persistentDesktop ->
+                        val maxTasks = getTaskLimit(persistentDesktop)
+                        val displayId = persistentDesktop.displayId
+                        val deskId = persistentDesktop.desktopId
+                        // TODO: b/401107440 - Implement desk restoration to other displays.
+                        val newDisplayId = Display.DEFAULT_DISPLAY
+                        val newDeskId =
+                            deskRecreationFactory.recreateDesk(
+                                userId = userId,
+                                destinationDisplayId = newDisplayId,
+                                deskId = deskId,
+                            )
+                        logV(
+                            "Recreated desk=%d in display=%d using new deskId=%d and displayId=%d",
+                            deskId,
+                            displayId,
+                            newDeskId,
+                            newDisplayId,
+                        )
+                        if (newDeskId != deskId || newDisplayId != displayId) {
+                            logV("Removing obsolete desk from persistence under deskId=%d", deskId)
+                            persistentRepository.removeDesktop(userId, deskId)
+                        }
+
+                        // TODO: b/393961770 - [DesktopRepository] doesn't save desks to the
+                        //  persistent repository until a task is added to them. Update it so that
+                        //  empty desks can be restored too.
+                        repository.addDesk(displayId = displayId, deskId = newDeskId)
+                        var visibleTasksCount = 0
+                        persistentDesktop.zOrderedTasksList
+                            // Reverse it so we initialize the repo from bottom to top.
+                            .reversed()
+                            .mapNotNull { taskId -> persistentDesktop.tasksByTaskIdMap[taskId] }
+                            .forEach { task ->
+                                // Visible here means non-minimized a.k.a. expanded, it does not
+                                // mean
+                                // it is visible in WM (and |DesktopRepository|) terms.
+                                val isVisible =
+                                    task.desktopTaskState == DesktopTaskState.VISIBLE &&
+                                        visibleTasksCount < maxTasks
+
+                                repository.addTaskToDesk(
+                                    displayId = displayId,
+                                    deskId = newDeskId,
+                                    taskId = task.taskId,
                                     isVisible = false,
                                 )
-                                repository.minimizeTask(persistentDesktop.displayId, task.taskId)
+
+                                if (isVisible) {
+                                    visibleTasksCount++
+                                } else {
+                                    repository.minimizeTaskInDesk(
+                                        displayId = displayId,
+                                        deskId = newDeskId,
+                                        taskId = task.taskId,
+                                    )
+                                }
                             }
-                        }
+                    }
                 }
+            } finally {
+                _isInitialized.value = true
             }
         }
     }
+
+    private suspend fun getDesksToRestore(
+        state: DesktopRepositoryState,
+        userId: Int,
+    ): Set<Desktop> {
+        // TODO: b/365873835 - what about desks that won't be restored?
+        //  - invalid desk ids from multi-desk -> single-desk switching can be ignored / deleted.
+        val limitToSingleDeskPerDisplay =
+            !DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue
+        return state.desktopMap.keys
+            .mapNotNull { deskId ->
+                persistentRepository.readDesktop(userId, deskId)?.takeIf { desk ->
+                    // Do not restore invalid desks when multi-desks is disabled. This is
+                    // possible if the feature is disabled after having created multiple desks.
+                    val isValidSingleDesk = desk.desktopId == desk.displayId
+                    (!limitToSingleDeskPerDisplay || isValidSingleDesk)
+                }
+            }
+            .toSet()
+    }
+
+    private fun getTaskLimit(persistedDesk: Desktop): Int =
+        DesktopModeStatus.getMaxTaskLimit(context).takeIf { it > 0 }
+            ?: persistedDesk.zOrderedTasksCount
+
+    private fun logV(msg: String, vararg arguments: Any?) {
+        ProtoLog.v(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
+    /** A default implementation of [DeskRecreationFactory] that reuses the desk id. */
+    private class DefaultDeskRecreationFactory : DeskRecreationFactory {
+        override suspend fun recreateDesk(
+            userId: Int,
+            destinationDisplayId: Int,
+            deskId: Int,
+        ): Int = deskId
+    }
+
+    companion object {
+        private const val TAG = "DesktopRepositoryInitializerImpl"
+    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandlerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandlerTest.kt
index 8ad54f5a0bb4..275d7b73a112 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandlerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopDisplayEventHandlerTest.kt
@@ -28,14 +28,22 @@ import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.common.DisplayController
 import com.android.wm.shell.common.DisplayController.OnDisplaysChangedListener
 import com.android.wm.shell.common.ShellExecutor
+import com.android.wm.shell.desktopmode.persistence.DesktopRepositoryInitializer
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.sysui.ShellInit
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.runCurrent
+import kotlinx.coroutines.test.runTest
 import org.junit.After
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.mockito.Mockito.spy
+import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.whenever
@@ -46,15 +54,18 @@ import org.mockito.quality.Strictness
  *
  * Usage: atest WMShellUnitTests:DesktopDisplayEventHandlerTest
  */
+@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidTestingRunner::class)
 class DesktopDisplayEventHandlerTest : ShellTestCase() {
     @Mock lateinit var testExecutor: ShellExecutor
     @Mock lateinit var displayController: DisplayController
     @Mock private lateinit var mockDesktopUserRepositories: DesktopUserRepositories
+    @Mock private lateinit var mockDesktopRepositoryInitializer: DesktopRepositoryInitializer
     @Mock private lateinit var mockDesktopRepository: DesktopRepository
     @Mock private lateinit var mockDesktopTasksController: DesktopTasksController
     @Mock private lateinit var desktopDisplayModeController: DesktopDisplayModeController
+    private val testScope = TestScope()
 
     private lateinit var mockitoSession: StaticMockitoSession
     private lateinit var shellInit: ShellInit
@@ -77,7 +88,9 @@ class DesktopDisplayEventHandlerTest : ShellTestCase() {
             DesktopDisplayEventHandler(
                 context,
                 shellInit,
+                testScope.backgroundScope,
                 displayController,
+                mockDesktopRepositoryInitializer,
                 mockDesktopUserRepositories,
                 mockDesktopTasksController,
                 desktopDisplayModeController,
@@ -89,17 +102,66 @@ class DesktopDisplayEventHandlerTest : ShellTestCase() {
 
     @After
     fun tearDown() {
+        testScope.cancel()
         mockitoSession.finishMocking()
     }
 
     @Test
-    fun testDisplayAdded_supportsDesks_createsDesk() {
-        whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
+    fun testDisplayAdded_supportsDesks_desktopRepositoryInitialized_createsDesk() =
+        testScope.runTest {
+            val stateFlow = MutableStateFlow(false)
+            whenever(mockDesktopRepositoryInitializer.isInitialized).thenReturn(stateFlow)
+            whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
 
-        onDisplaysChangedListenerCaptor.lastValue.onDisplayAdded(DEFAULT_DISPLAY)
+            onDisplaysChangedListenerCaptor.lastValue.onDisplayAdded(DEFAULT_DISPLAY)
+            stateFlow.emit(true)
+            runCurrent()
 
-        verify(mockDesktopTasksController).createDesk(DEFAULT_DISPLAY)
-    }
+            verify(mockDesktopTasksController).createDesk(DEFAULT_DISPLAY)
+        }
+
+    @Test
+    fun testDisplayAdded_supportsDesks_desktopRepositoryNotInitialized_doesNotCreateDesk() =
+        testScope.runTest {
+            val stateFlow = MutableStateFlow(false)
+            whenever(mockDesktopRepositoryInitializer.isInitialized).thenReturn(stateFlow)
+            whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
+
+            onDisplaysChangedListenerCaptor.lastValue.onDisplayAdded(DEFAULT_DISPLAY)
+            runCurrent()
+
+            verify(mockDesktopTasksController, never()).createDesk(DEFAULT_DISPLAY)
+        }
+
+    @Test
+    fun testDisplayAdded_supportsDesks_desktopRepositoryInitializedTwice_createsDeskOnce() =
+        testScope.runTest {
+            val stateFlow = MutableStateFlow(false)
+            whenever(mockDesktopRepositoryInitializer.isInitialized).thenReturn(stateFlow)
+            whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
+
+            onDisplaysChangedListenerCaptor.lastValue.onDisplayAdded(DEFAULT_DISPLAY)
+            stateFlow.emit(true)
+            stateFlow.emit(true)
+            runCurrent()
+
+            verify(mockDesktopTasksController, times(1)).createDesk(DEFAULT_DISPLAY)
+        }
+
+    @Test
+    fun testDisplayAdded_supportsDesks_desktopRepositoryInitialized_deskExists_doesNotCreateDesk() =
+        testScope.runTest {
+            val stateFlow = MutableStateFlow(false)
+            whenever(mockDesktopRepositoryInitializer.isInitialized).thenReturn(stateFlow)
+            whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
+            whenever(mockDesktopRepository.getNumberOfDesks(DEFAULT_DISPLAY)).thenReturn(1)
+
+            onDisplaysChangedListenerCaptor.lastValue.onDisplayAdded(DEFAULT_DISPLAY)
+            stateFlow.emit(true)
+            runCurrent()
+
+            verify(mockDesktopTasksController, never()).createDesk(DEFAULT_DISPLAY)
+        }
 
     @Test
     fun testDisplayAdded_cannotEnterDesktopMode_doesNotCreateDesk() {
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 006c3cae121c..4c18ee1500b7 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
@@ -114,6 +114,8 @@ class DesktopImmersiveControllerTest : ShellTestCase() {
                 transactionSupplier = transactionSupplier,
             )
         desktopRepository = userRepositories.current
+        desktopRepository.addDesk(DEFAULT_DISPLAY, DEFAULT_DESK_ID)
+        desktopRepository.setActiveDesk(DEFAULT_DISPLAY, DEFAULT_DESK_ID)
     }
 
     @Test
@@ -835,5 +837,6 @@ class DesktopImmersiveControllerTest : ShellTestCase() {
     companion object {
         private val STABLE_BOUNDS = Rect(0, 100, 2000, 1900)
         private val DISPLAY_BOUNDS = Rect(0, 0, 2000, 2000)
+        private const val DEFAULT_DESK_ID = 0
     }
 }
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 de92d391645a..f84a1a38bdfc 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
@@ -1203,6 +1203,17 @@ class DesktopRepositoryTest(flags: FlagsParameterization) : ShellTestCase() {
     }
 
     @Test
+    @EnableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND, FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
+    fun removeDesk_removesFromPersistence() =
+        runTest(StandardTestDispatcher()) {
+            repo.addDesk(displayId = DEFAULT_DISPLAY, deskId = 2)
+
+            repo.removeDesk(deskId = 2)
+
+            verify(persistentRepository).removeDesktop(DEFAULT_USER_ID, 2)
+        }
+
+    @Test
     fun getTaskInFullImmersiveState_byDisplay() {
         repo.addDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
         repo.setActiveDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
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 75308442d76a..34c5ebd6d94d 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
@@ -22,6 +22,7 @@ import android.app.ActivityOptions
 import android.app.KeyguardManager
 import android.app.PendingIntent
 import android.app.PictureInPictureParams
+import android.app.WindowConfiguration
 import android.app.WindowConfiguration.ACTIVITY_TYPE_HOME
 import android.app.WindowConfiguration.ACTIVITY_TYPE_STANDARD
 import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
@@ -297,6 +298,10 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     private val wallpaperToken = MockToken().token()
     private val homeComponentName = ComponentName(HOME_LAUNCHER_PACKAGE_NAME, /* class */ "")
 
+    init {
+        mSetFlagsRule.setFlagsParameterization(flags)
+    }
+
     @Before
     fun setUp() {
         Dispatchers.setMain(StandardTestDispatcher())
@@ -433,6 +438,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
             dragToDesktopTransitionHandler,
             mMockDesktopImmersiveController,
             userRepositories,
+            repositoryInitializer,
             recentsTransitionHandler,
             multiInstanceHelper,
             shellExecutor,
@@ -627,11 +633,13 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isDesktopModeShowing_noTasks_returnsFalse() {
         assertThat(controller.isDesktopModeShowing(displayId = 0)).isFalse()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isDesktopModeShowing_noTasksVisible_returnsFalse() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
@@ -642,6 +650,15 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun isDesktopModeShowing_noActiveDesk_returnsFalse() {
+        taskRepository.setDeskInactive(deskId = 0)
+
+        assertThat(controller.isDesktopModeShowing(displayId = 0)).isFalse()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isDesktopModeShowing_tasksActiveAndVisible_returnsTrue() {
         val task1 = setUpFreeformTask()
         val task2 = setUpFreeformTask()
@@ -656,6 +673,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
         Flags.FLAG_INCLUDE_TOP_TRANSPARENT_FULLSCREEN_TASK_IN_DESKTOP_HEURISTIC,
     )
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isDesktopModeShowing_topTransparentFullscreenTask_returnsTrue() {
         val topTransparentTask = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
         taskRepository.setTopTransparentFullscreenTaskId(DEFAULT_DISPLAY, topTransparentTask.taskId)
@@ -665,6 +683,20 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+        Flags.FLAG_INCLUDE_TOP_TRANSPARENT_FULLSCREEN_TASK_IN_DESKTOP_HEURISTIC,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun isDesktopModeShowing_deskInactive_topTransparentFullscreenTask_returnsTrue() {
+        val topTransparentTask = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+        taskRepository.setTopTransparentFullscreenTaskId(DEFAULT_DISPLAY, topTransparentTask.taskId)
+        taskRepository.setDeskInactive(deskId = 0)
+
+        assertThat(controller.isDesktopModeShowing(displayId = DEFAULT_DISPLAY)).isTrue()
+    }
+
+    @Test
+    @EnableFlags(
         Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
         Flags.FLAG_ENABLE_PER_DISPLAY_DESKTOP_WALLPAPER_ACTIVITY,
     )
@@ -1053,11 +1085,29 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isAnyDeskActive_noTasks_returnsFalse() {
         assertThat(controller.isAnyDeskActive(DEFAULT_DISPLAY)).isFalse()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun isAnyDeskActive_noActiveDesk_returnsFalse() {
+        taskRepository.setDeskInactive(deskId = 0)
+
+        assertThat(controller.isAnyDeskActive(DEFAULT_DISPLAY)).isFalse()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun isAnyDeskActive_withActiveDesk_returnsTrue() {
+        taskRepository.setActiveDesk(displayId = DEFAULT_DISPLAY, deskId = 0)
+
+        assertThat(controller.isAnyDeskActive(DEFAULT_DISPLAY)).isTrue()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isAnyDeskActive_twoTasks_bothVisible_returnsTrue() {
         setUpHomeTask()
 
@@ -1068,6 +1118,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isInDesktop_twoTasks_oneVisible_returnsTrue() {
         setUpHomeTask()
 
@@ -1078,6 +1129,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun isAnyDeskActive_twoTasksVisibleOnDifferentDisplays_returnsTrue() {
         taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
         taskRepository.setActiveDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
@@ -1097,7 +1149,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         controller.addMoveToDeskTaskChanges(wct, task, deskId = 0)
 
         val finalBounds = findBoundsChange(wct, task)
-        assertThat(finalBounds).isEqualTo(Rect())
+        assertThat(finalBounds).isNull()
     }
 
     @Test
@@ -1108,7 +1160,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         controller.addMoveToDeskTaskChanges(wct, task, deskId = 0)
 
         val finalBounds = findBoundsChange(wct, task)
-        assertThat(finalBounds).isEqualTo(Rect())
+        assertThat(finalBounds).isNull()
     }
 
     @Test
@@ -1119,7 +1171,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         controller.addMoveToDeskTaskChanges(wct, task, deskId = 0)
 
         val finalBounds = findBoundsChange(wct, task)
-        assertThat(finalBounds).isEqualTo(Rect())
+        assertThat(finalBounds).isNull()
     }
 
     @Test
@@ -1130,7 +1182,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         controller.addMoveToDeskTaskChanges(wct, task, deskId = 0)
 
         val finalBounds = findBoundsChange(wct, task)
-        assertThat(finalBounds).isEqualTo(Rect())
+        assertThat(finalBounds).isNull()
     }
 
     @Test
@@ -1862,6 +1914,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         val wallpaperToken = MockToken().token()
         whenever(desktopWallpaperActivityTokenProvider.getToken(SECOND_DISPLAY))
             .thenReturn(wallpaperToken)
+        taskRepository.addDesk(SECOND_DISPLAY, deskId = 2)
         val task = setUpFreeformTask(displayId = SECOND_DISPLAY, deskId = 2, background = true)
 
         controller.moveTaskToDefaultDeskAndActivate(
@@ -2666,7 +2719,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveToNextDisplay_moveFromFirstToSecondDisplay() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveToNextDisplay_moveFromFirstToSecondDisplay_multiDesksDisabled() {
         // Set up two display ids
         taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
         whenever(rootTaskDisplayAreaOrganizer.displayIds)
@@ -2692,7 +2746,27 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveToNextDisplay_moveFromSecondToFirstDisplay() {
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveToNextDisplay_moveFromFirstToSecondDisplay_multiDesksEnabled() {
+        // Set up two display ids
+        val targetDeskId = 2
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = targetDeskId)
+        whenever(rootTaskDisplayAreaOrganizer.displayIds)
+            .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
+        // Create a mock for the target display area: second display
+        val secondDisplayArea = DisplayAreaInfo(MockToken().token(), SECOND_DISPLAY, 0)
+        whenever(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(SECOND_DISPLAY))
+            .thenReturn(secondDisplayArea)
+
+        val task = setUpFreeformTask(displayId = DEFAULT_DISPLAY)
+        controller.moveToNextDisplay(task.taskId)
+
+        verify(desksOrganizer).moveTaskToDesk(any(), eq(targetDeskId), eq(task))
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveToNextDisplay_moveFromSecondToFirstDisplay_multiDesksDisabled() {
         // Set up two display ids
         whenever(rootTaskDisplayAreaOrganizer.displayIds)
             .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
@@ -2718,6 +2792,25 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveToNextDisplay_moveFromSecondToFirstDisplay_multiDesksEnabled() {
+        // Set up two display ids
+        val targetDeskId = 0
+        whenever(rootTaskDisplayAreaOrganizer.displayIds)
+            .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
+        // Create a mock for the target display area: default display
+        val defaultDisplayArea = DisplayAreaInfo(MockToken().token(), DEFAULT_DISPLAY, 0)
+        whenever(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY))
+            .thenReturn(defaultDisplayArea)
+
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = SECOND_DISPLAY)
+        val task = setUpFreeformTask(displayId = SECOND_DISPLAY)
+        controller.moveToNextDisplay(task.taskId)
+
+        verify(desksOrganizer).moveTaskToDesk(any(), eq(targetDeskId), eq(task))
+    }
+
+    @Test
     @EnableFlags(
         FLAG_ENABLE_PER_DISPLAY_DESKTOP_WALLPAPER_ACTIVITY,
         Flags.FLAG_ENABLE_DESKTOP_WALLPAPER_ACTIVITY_FOR_SYSTEM_USER,
@@ -2781,6 +2874,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @EnableFlags(FLAG_ENABLE_MOVE_TO_NEXT_DISPLAY_SHORTCUT)
     fun moveToNextDisplay_sizeInDpPreserved() {
         // Set up two display ids
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = 2)
         whenever(rootTaskDisplayAreaOrganizer.displayIds)
             .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
         // Create a mock for the target display area: second display
@@ -2822,6 +2916,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @EnableFlags(FLAG_ENABLE_MOVE_TO_NEXT_DISPLAY_SHORTCUT)
     fun moveToNextDisplay_shiftWithinDestinationDisplayBounds() {
         // Set up two display ids
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = 2)
         whenever(rootTaskDisplayAreaOrganizer.displayIds)
             .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
         // Create a mock for the target display area: second display
@@ -2863,6 +2958,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         // Set up two display ids
         whenever(rootTaskDisplayAreaOrganizer.displayIds)
             .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = 2)
         // Create a mock for the target display area: second display
         val secondDisplayArea = DisplayAreaInfo(MockToken().token(), SECOND_DISPLAY, 0)
         whenever(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(SECOND_DISPLAY))
@@ -2971,7 +3067,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun moveToNextDisplay_toDesktopInOtherDisplay_bringsExistingTasksToFront() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveToNextDisplay_toDesktopInOtherDisplay_multiDesksDisabled_bringsExistingTasksToFront() {
         val transition = Binder()
         val sourceDeskId = 0
         val targetDeskId = 2
@@ -2997,6 +3094,34 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun moveToNextDisplay_toDesktopInOtherDisplay_multiDesksEnabled_bringsExistingTasksToFront() {
+        val transition = Binder()
+        val sourceDeskId = 0
+        val targetDeskId = 2
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = targetDeskId)
+        taskRepository.setDeskInactive(deskId = targetDeskId)
+        // Set up two display ids
+        whenever(rootTaskDisplayAreaOrganizer.displayIds)
+            .thenReturn(intArrayOf(DEFAULT_DISPLAY, SECOND_DISPLAY))
+        // Create a mock for the target display area: second display
+        val secondDisplayArea = DisplayAreaInfo(MockToken().token(), SECOND_DISPLAY, 0)
+        whenever(rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(SECOND_DISPLAY))
+            .thenReturn(secondDisplayArea)
+        whenever(transitions.startTransition(eq(TRANSIT_CHANGE), any(), anyOrNull()))
+            .thenReturn(transition)
+        val task1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = sourceDeskId)
+        val task2 = setUpFreeformTask(displayId = SECOND_DISPLAY, deskId = targetDeskId)
+
+        controller.moveToNextDisplay(task1.taskId)
+
+        // Existing desktop task in the target display is moved to front.
+        val wct = getLatestTransition()
+        assertNotNull(wct)
+        verify(desksOrganizer).reorderTaskToFront(wct, targetDeskId, task2)
+    }
+
+    @Test
     @EnableFlags(
         Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
         Flags.FLAG_ENABLE_PER_DISPLAY_DESKTOP_WALLPAPER_ACTIVITY,
@@ -3110,7 +3235,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @Test
     fun moveToNextDisplay_movingToDesktop_sendsTaskbarRoundingUpdate() {
         val transition = Binder()
-        val sourceDeskId = 1
+        val sourceDeskId = 0
         val targetDeskId = 2
         taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = targetDeskId)
         taskRepository.setDeskInactive(deskId = targetDeskId)
@@ -3627,7 +3752,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun handleRequest_fullscreenTask_freeformVisible_returnSwitchToFreeformWCT() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_fullscreenTask_freeformVisible_multiDesksDisabled_returnSwitchToFreeformWCT() {
         val homeTask = setUpHomeTask()
         val freeformTask = setUpFreeformTask()
         markTaskVisible(freeformTask)
@@ -3643,7 +3769,22 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun handleRequest_fullscreenTaskWithTaskOnHome_freeformVisible_returnSwitchToFreeformWCT() {
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_fullscreenTask_deskActive_multiDesksEnabled_movesToDesk() {
+        val deskId = 0
+        taskRepository.setActiveDesk(DEFAULT_DISPLAY, deskId = deskId)
+        setUpHomeTask()
+        val fullscreenTask = createFullscreenTask()
+
+        val wct = controller.handleRequest(Binder(), createTransition(fullscreenTask))
+
+        assertNotNull(wct, "should handle request")
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId, fullscreenTask)
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_fullscreenTaskWithTaskOnHome_freeformVisible_multiDesksDisabled_returnSwitchToFreeformWCT() {
         val homeTask = setUpHomeTask()
         val freeformTask = setUpFreeformTask()
         markTaskVisible(freeformTask)
@@ -3668,6 +3809,21 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_fullscreenTaskWithTaskOnHome_activeDesk_multiDesksEnabled_movesToDesk() {
+        val deskId = 0
+        setUpHomeTask()
+        setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = deskId)
+        val fullscreenTask = createFullscreenTask()
+        fullscreenTask.baseIntent.setFlags(Intent.FLAG_ACTIVITY_TASK_ON_HOME)
+
+        val wct = controller.handleRequest(Binder(), createTransition(fullscreenTask))
+
+        assertNotNull(wct, "should handle request")
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId, fullscreenTask)
+    }
+
+    @Test
     @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun handleRequest_fullscreenTaskToDesk_underTaskLimit_multiDesksDisabled_dontMinimize() {
         val freeformTask = setUpFreeformTask()
@@ -3899,6 +4055,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun handleRequest_fullscreenTask_noTasks_enforceDesktop_freeformDisplay_returnFreeformWCT() {
         whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
         whenever(DesktopModeStatus.enterDesktopByDefaultOnFreeformDisplay(context)).thenReturn(true)
@@ -3921,7 +4078,28 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    fun handleRequest_fullscreenTask_noTasks_enforceDesktop_fullscreenDisplay_returnNull() {
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun handleRequest_fullscreenTask_noInDesk_enforceDesktop_freeformDisplay_movesToDesk() {
+        val deskId = 0
+        taskRepository.setDeskInactive(deskId)
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
+        whenever(DesktopModeStatus.enterDesktopByDefaultOnFreeformDisplay(context)).thenReturn(true)
+        val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
+        tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FREEFORM
+
+        val fullscreenTask = createFullscreenTask()
+        val wct = controller.handleRequest(Binder(), createTransition(fullscreenTask))
+
+        assertNotNull(wct, "should handle request")
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId, fullscreenTask)
+    }
+
+    @Test
+    fun handleRequest_fullscreenTask_notInDesk_enforceDesktop_fullscreenDisplay_returnNull() {
+        taskRepository.setDeskInactive(deskId = 0)
         whenever(DesktopModeStatus.enterDesktopByDefaultOnFreeformDisplay(context)).thenReturn(true)
         val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
         tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FULLSCREEN
@@ -3933,6 +4111,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun handleRequest_fullscreenTask_freeformNotVisible_returnNull() {
         val freeformTask = setUpFreeformTask()
         markTaskHidden(freeformTask)
@@ -3941,12 +4120,23 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun handleRequest_fullscreenTask_noOtherTasks_returnNull() {
         val fullscreenTask = createFullscreenTask()
         assertThat(controller.handleRequest(Binder(), createTransition(fullscreenTask))).isNull()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_fullscreenTask_notInDesk_returnNull() {
+        taskRepository.setDeskInactive(deskId = 0)
+        val fullscreenTask = createFullscreenTask()
+
+        assertThat(controller.handleRequest(Binder(), createTransition(fullscreenTask))).isNull()
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun handleRequest_fullscreenTask_freeformTaskOnOtherDisplay_returnNull() {
         val fullscreenTaskDefaultDisplay = createFullscreenTask(displayId = DEFAULT_DISPLAY)
         createFreeformTask(displayId = SECOND_DISPLAY)
@@ -3957,6 +4147,20 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_fullscreenTask_deskInOtherDisplayActive_returnNull() {
+        taskRepository.setDeskInactive(deskId = 0)
+        val fullscreenTaskDefaultDisplay = createFullscreenTask(displayId = DEFAULT_DISPLAY)
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = 2)
+        taskRepository.setActiveDesk(displayId = SECOND_DISPLAY, deskId = 2)
+
+        val result =
+            controller.handleRequest(Binder(), createTransition(fullscreenTaskDefaultDisplay))
+
+        assertThat(result).isNull()
+    }
+
+    @Test
     @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun handleRequest_freeformTask_freeformVisible_aboveTaskLimit_multiDesksDisabled_minimize() {
         val deskId = 0
@@ -4008,7 +4212,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     fun handleRequest_freeformTask_relaunchActiveTask_taskBecomesUndefined() {
-        val freeformTask = setUpFreeformTask()
+        taskRepository.setDeskInactive(deskId = 0)
+        val freeformTask = setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = 0)
         markTaskHidden(freeformTask)
 
         val wct = controller.handleRequest(Binder(), createTransition(freeformTask))
@@ -4038,9 +4243,10 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         whenever(DesktopModeStatus.enterDesktopByDefaultOnFreeformDisplay(context)).thenReturn(true)
         val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
         tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FULLSCREEN
-
-        val freeformTask = setUpFreeformTask()
+        taskRepository.setDeskInactive(deskId = 0)
+        val freeformTask = setUpFreeformTask(DEFAULT_DISPLAY, deskId = 0)
         markTaskHidden(freeformTask)
+
         val wct = controller.handleRequest(Binder(), createTransition(freeformTask))
 
         assertNotNull(wct, "should handle request")
@@ -4049,7 +4255,10 @@ 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 handleRequest_freeformTask_desktopWallpaperDisabled_freeformNotVisible_reorderedToTop() {
         val freeformTask1 = setUpFreeformTask()
         val freeformTask2 = createFreeformTask()
@@ -4068,7 +4277,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-    fun handleRequest_freeformTask_desktopWallpaperEnabled_freeformNotVisible_reorderedToTop() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_freeformTask_desktopWallpaperEnabled_freeformNotVisible_multiDesksDisabled_reorderedToTop() {
         whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
         val freeformTask1 = setUpFreeformTask()
         val freeformTask2 = createFreeformTask()
@@ -4091,34 +4301,47 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-    fun handleRequest_freeformTask_desktopWallpaperDisabled_noOtherTasks_reorderedToTop() {
-        val task = createFreeformTask()
-        val result = controller.handleRequest(Binder(), createTransition(task))
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun handleRequest_freeformTask_desktopWallpaperEnabled_notInDesk_reorderedToTop() {
+        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
+        val deskId = 0
+        taskRepository.setDeskInactive(deskId)
+        val freeformTask1 = setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = deskId)
+        val freeformTask2 = createFreeformTask()
 
-        assertNotNull(result, "Should handle request")
-        assertThat(result.hierarchyOps?.size).isEqualTo(1)
-        result.assertReorderAt(0, task, toTop = true)
+        val wct =
+            controller.handleRequest(
+                Binder(),
+                createTransition(freeformTask2, type = TRANSIT_TO_FRONT),
+            )
+
+        assertNotNull(wct, "Should handle request")
+        verify(desksOrganizer).reorderTaskToFront(wct, deskId, freeformTask1)
+        wct.assertReorder(freeformTask2, toTop = true)
     }
 
     @Test
-    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
-    fun handleRequest_freeformTask_desktopWallpaperEnabled_noOtherTasks_reorderedToTop() {
-        whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
+    @DisableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun handleRequest_freeformTask_desktopWallpaperDisabled_noOtherTasks_reorderedToTop() {
         val task = createFreeformTask()
-
         val result = controller.handleRequest(Binder(), createTransition(task))
 
         assertNotNull(result, "Should handle request")
-        assertThat(result.hierarchyOps?.size).isEqualTo(2)
-        // Add desktop wallpaper activity
-        result.assertPendingIntentAt(0, desktopWallpaperIntent)
-        // Bring new task to front
-        result.assertReorderAt(1, task, toTop = true)
+        assertThat(result.hierarchyOps?.size).isEqualTo(1)
+        result.assertReorderAt(0, task, toTop = true)
     }
 
     @Test
-    @DisableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
+    @DisableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
     fun handleRequest_freeformTask_dskWallpaperDisabled_freeformOnOtherDisplayOnly_reorderedToTop() {
         val taskDefaultDisplay = createFreeformTask(displayId = DEFAULT_DISPLAY)
         // Second display task
@@ -4134,10 +4357,13 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_WALLPAPER_ACTIVITY)
     fun handleRequest_freeformTask_dskWallpaperEnabled_freeformOnOtherDisplayOnly_reorderedToTop() {
+        taskRepository.setDeskInactive(deskId = 0)
         whenever(desktopWallpaperActivityTokenProvider.getToken()).thenReturn(null)
         val taskDefaultDisplay = createFreeformTask(displayId = DEFAULT_DISPLAY)
         // Second display task
-        createFreeformTask(displayId = SECOND_DISPLAY)
+        taskRepository.addDesk(displayId = SECOND_DISPLAY, deskId = 2)
+        taskRepository.setActiveDesk(displayId = SECOND_DISPLAY, deskId = 2)
+        setUpFreeformTask(displayId = SECOND_DISPLAY, deskId = 2)
 
         val result = controller.handleRequest(Binder(), createTransition(taskDefaultDisplay))
 
@@ -4273,7 +4499,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
-    fun handleRequest_topActivityTransparentWithoutDisplay_returnSwitchToFreeformWCT() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_topActivityTransparentWithoutDisplay_multiDesksDisabled_returnSwitchToFreeformWCT() {
         val freeformTask = setUpFreeformTask()
         markTaskVisible(freeformTask)
 
@@ -4290,6 +4517,29 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun handleRequest_topActivityTransparentWithoutDisplay_multiDesksEnabled_returnSwitchToFreeformWCT() {
+        val deskId = 0
+        val freeformTask = setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = deskId)
+        markTaskVisible(freeformTask)
+
+        val task =
+            setUpFullscreenTask().apply {
+                isActivityStackTransparent = true
+                isTopActivityNoDisplay = true
+                numActivities = 1
+            }
+
+        val wct = controller.handleRequest(Binder(), createTransition(task))
+
+        assertNotNull(wct)
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId, task)
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
     @DisableFlags(Flags.FLAG_ENABLE_MODALS_FULLSCREEN_WITH_PERMISSION)
     fun handleRequest_topActivityTransparentWithDisplay_returnSwitchToFullscreenWCT() {
@@ -4348,7 +4598,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
         Flags.FLAG_INCLUDE_TOP_TRANSPARENT_FULLSCREEN_TASK_IN_DESKTOP_HEURISTIC,
     )
-    fun handleRequest_onlyTopTransparentFullscreenTask_returnSwitchToFreeformWCT() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_onlyTopTransparentFullscreenTask_multiDesksDisabled_returnSwitchToFreeformWCT() {
         val topTransparentTask = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
         taskRepository.setTopTransparentFullscreenTaskId(DEFAULT_DISPLAY, topTransparentTask.taskId)
 
@@ -4360,8 +4611,28 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+        Flags.FLAG_INCLUDE_TOP_TRANSPARENT_FULLSCREEN_TASK_IN_DESKTOP_HEURISTIC,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun handleRequest_onlyTopTransparentFullscreenTask_multiDesksEnabled_movesToDesktop() {
+        val deskId = 0
+        val topTransparentTask = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+        taskRepository.setTopTransparentFullscreenTaskId(DEFAULT_DISPLAY, topTransparentTask.taskId)
+        taskRepository.setDeskInactive(deskId = deskId)
+
+        val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+
+        val wct = controller.handleRequest(Binder(), createTransition(task))
+        assertNotNull(wct)
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId, task)
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
     fun handleRequest_desktopNotShowing_topTransparentFullscreenTask_returnNull() {
+        taskRepository.setDeskInactive(deskId = 0)
         val task = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
 
         assertThat(controller.handleRequest(Binder(), createTransition(task))).isNull()
@@ -4415,7 +4686,8 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
-    fun handleRequest_systemUIActivityWithoutDisplay_returnSwitchToFreeformWCT() {
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun handleRequest_systemUIActivityWithoutDisplay_multiDesksDisabled_returnSwitchToFreeformWCT() {
         val freeformTask = setUpFreeformTask()
         markTaskVisible(freeformTask)
 
@@ -4434,6 +4706,32 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
             .isEqualTo(WINDOWING_MODE_FREEFORM)
     }
 
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun handleRequest_systemUIActivityWithoutDisplay_multiDesksEnabled_movesTaskToDesk() {
+        val deskId = 0
+        val freeformTask = setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = deskId)
+        markTaskVisible(freeformTask)
+
+        // 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(displayId = DEFAULT_DISPLAY).apply {
+                baseActivity = baseComponent
+                isTopActivityNoDisplay = true
+            }
+
+        val wct = controller.handleRequest(Binder(), createTransition(task))
+
+        assertNotNull(wct)
+        verify(desksOrganizer).moveTaskToDesk(wct, deskId, task)
+    }
+
     @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
     fun handleRequest_defaultHomePackageWithDisplay_returnSwitchToFullscreenWCT() {
         val freeformTask = setUpFreeformTask()
@@ -4476,6 +4774,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
     @Test
     fun handleRequest_systemUIActivityWithDisplay_returnSwitchToFullscreenWCT_enforcedDesktop() {
+        taskRepository.setDeskInactive(deskId = 0)
         whenever(DesktopModeStatus.enterDesktopByDefaultOnFreeformDisplay(context)).thenReturn(true)
         val tda = rootTaskDisplayAreaOrganizer.getDisplayAreaInfo(DEFAULT_DISPLAY)!!
         tda.configuration.windowConfiguration.windowingMode = WINDOWING_MODE_FREEFORM
@@ -5109,6 +5408,55 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION,
         Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
     )
+    fun activateDesk_hasNonRunningTask_startsTask() {
+        val deskId = 0
+        val nonRunningTask =
+            setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = 0, background = true)
+
+        val transition = Binder()
+        val deskChange = mock(TransitionInfo.Change::class.java)
+        whenever(transitions.startTransition(eq(TRANSIT_TO_FRONT), any(), anyOrNull()))
+            .thenReturn(transition)
+        whenever(desksOrganizer.isDeskActiveAtEnd(deskChange, deskId)).thenReturn(true)
+        // Make desk inactive by activating another desk.
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 1)
+        taskRepository.setActiveDesk(DEFAULT_DISPLAY, deskId = 1)
+
+        controller.activateDesk(deskId, RemoteTransition(TestRemoteTransition()))
+
+        val wct = getLatestWct(TRANSIT_TO_FRONT, OneShotRemoteHandler::class.java)
+        assertNotNull(wct)
+        wct.assertLaunchTask(nonRunningTask.taskId, WINDOWING_MODE_FREEFORM)
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun activateDesk_hasRunningTask_reordersTask() {
+        val deskId = 0
+        val runningTask = setUpFreeformTask(displayId = DEFAULT_DISPLAY, deskId = 0)
+
+        val transition = Binder()
+        val deskChange = mock(TransitionInfo.Change::class.java)
+        whenever(transitions.startTransition(eq(TRANSIT_TO_FRONT), any(), anyOrNull()))
+            .thenReturn(transition)
+        whenever(desksOrganizer.isDeskActiveAtEnd(deskChange, deskId)).thenReturn(true)
+        // Make desk inactive by activating another desk.
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 1)
+        taskRepository.setActiveDesk(DEFAULT_DISPLAY, deskId = 1)
+
+        controller.activateDesk(deskId, RemoteTransition(TestRemoteTransition()))
+
+        verify(desksOrganizer).reorderTaskToFront(any(), eq(deskId), eq(runningTask))
+    }
+
+    @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)
@@ -6833,6 +7181,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
     fun shouldPlayDesktopAnimation_notShowingDesktop_doesNotPlay() {
+        taskRepository.setDeskInactive(deskId = 0)
         val triggerTask = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
         taskRepository.setTaskInFullImmersiveState(
             displayId = triggerTask.displayId,
@@ -6909,6 +7258,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
     fun shouldPlayDesktopAnimation_fullscreenStaysFullscreen_doesNotPlay() {
         val triggerTask = setUpFullscreenTask(displayId = DEFAULT_DISPLAY)
+        taskRepository.setDeskInactive(deskId = 0)
         assertThat(controller.isDesktopModeShowing(triggerTask.displayId)).isFalse()
 
         assertThat(
@@ -6944,6 +7294,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @EnableFlags(Flags.FLAG_ENABLE_FULLY_IMMERSIVE_IN_DESKTOP)
     fun shouldPlayDesktopAnimation_freeformExitsDesktop_doesNotPlay() {
         val triggerTask = setUpFreeformTask(displayId = DEFAULT_DISPLAY, active = false)
+        taskRepository.setDeskInactive(deskId = 0)
         assertThat(controller.isDesktopModeShowing(triggerTask.displayId)).isFalse()
 
         assertThat(
@@ -6973,6 +7324,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         Flags.FLAG_ENABLE_DESKTOP_WALLPAPER_ACTIVITY_FOR_SYSTEM_USER,
     )
     fun startLaunchTransition_desktopNotShowing_movesWallpaperToFront() {
+        taskRepository.setDeskInactive(deskId = 0)
         val launchingTask = createFreeformTask(displayId = DEFAULT_DISPLAY)
         val wct = WindowContainerTransaction()
         wct.reorder(launchingTask.token, /* onTop= */ true)
@@ -7442,7 +7794,15 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     }
 
     private fun findBoundsChange(wct: WindowContainerTransaction, task: RunningTaskInfo): Rect? =
-        wct.changes[task.token.asBinder()]?.configuration?.windowConfiguration?.bounds
+        wct.changes.entries
+            .find { (token, change) ->
+                token == task.token.asBinder() &&
+                    (change.windowSetMask and WindowConfiguration.WINDOW_CONFIG_BOUNDS) != 0
+            }
+            ?.value
+            ?.configuration
+            ?.windowConfiguration
+            ?.bounds
 
     private fun verifyWCTNotExecuted() {
         verify(transitions, never()).startTransition(anyInt(), any(), isNull())
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 6b2f90fc0a59..9af504797182 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_FREEFORM
 import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
 import android.testing.AndroidTestingRunner
 import android.view.Display
@@ -29,6 +30,7 @@ import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_R
 import android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_SET_LAUNCH_ROOT
 import androidx.test.filters.SmallTest
 import com.android.wm.shell.ShellTaskOrganizer
+import com.android.wm.shell.ShellTaskOrganizer.TaskListener
 import com.android.wm.shell.ShellTestCase
 import com.android.wm.shell.TestShellExecutor
 import com.android.wm.shell.common.LaunchAdjacentController
@@ -39,14 +41,17 @@ import com.android.wm.shell.sysui.ShellCommandHandler
 import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
 import kotlin.test.assertNotNull
+import kotlinx.coroutines.test.runTest
 import org.junit.Assert.assertEquals
 import org.junit.Assert.assertThrows
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.Mockito
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.argThat
 import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 /**
  * Tests for [RootTaskDesksOrganizer].
@@ -76,48 +81,17 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
             )
     }
 
-    @Test
-    fun testCreateDesk_callsBack() {
-        val callback = FakeOnCreateCallback()
-        organizer.createDesk(Display.DEFAULT_DISPLAY, callback)
-
-        val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
-
-        assertThat(callback.created).isTrue()
-        assertEquals(freeformRoot.taskId, callback.deskId)
-    }
-
-    @Test
-    fun testCreateDesk_createsMinimizationRoot() {
-        val callback = FakeOnCreateCallback()
-        organizer.createDesk(Display.DEFAULT_DISPLAY, callback)
-        val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
-
-        val minimizationRootTask = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(minimizationRootTask, SurfaceControl())
-
-        val minimizationRoot = organizer.deskMinimizationRootsByDeskId[freeformRoot.taskId]
-        assertNotNull(minimizationRoot)
-        assertThat(minimizationRoot.deskId).isEqualTo(freeformRoot.taskId)
-        assertThat(minimizationRoot.rootId).isEqualTo(minimizationRootTask.taskId)
-    }
+    @Test fun testCreateDesk_createsDeskAndMinimizationRoots() = runTest { createDesk() }
 
     @Test
-    fun testCreateMinimizationRoot_marksHidden() {
-        organizer.createDesk(Display.DEFAULT_DISPLAY, FakeOnCreateCallback())
-        val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
-
-        val minimizationRootTask = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(minimizationRootTask, SurfaceControl())
+    fun testCreateMinimizationRoot_marksHidden() = runTest {
+        val desk = createDesk()
 
         verify(mockShellTaskOrganizer)
             .applyTransaction(
                 argThat { wct ->
                     wct.changes.any { change ->
-                        change.key == minimizationRootTask.token.asBinder() &&
+                        change.key == desk.minimizationRoot.token.asBinder() &&
                             (change.value.changeMask and Change.CHANGE_HIDDEN != 0) &&
                             change.value.hidden
                     }
@@ -126,7 +100,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testOnTaskAppeared_withoutRequest_throws() {
+    fun testOnTaskAppeared_withoutRequest_throws() = runTest {
         val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
 
         assertThrows(Exception::class.java) {
@@ -135,41 +109,25 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testOnTaskAppeared_withRequestOnlyInAnotherDisplay_throws() {
-        organizer.createDesk(displayId = 2, FakeOnCreateCallback())
-        val freeformRoot = createFreeformTask(Display.DEFAULT_DISPLAY).apply { parentTaskId = -1 }
-
-        assertThrows(Exception::class.java) {
-            organizer.onTaskAppeared(freeformRoot, SurfaceControl())
-        }
-    }
-
-    @Test
-    fun testOnTaskAppeared_duplicateRoot_throws() {
-        organizer.createDesk(Display.DEFAULT_DISPLAY, FakeOnCreateCallback())
-        val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
+    fun testOnTaskAppeared_duplicateRoot_throws() = runTest {
+        val desk = createDesk()
 
         assertThrows(Exception::class.java) {
-            organizer.onTaskAppeared(freeformRoot, SurfaceControl())
+            organizer.onTaskAppeared(desk.deskRoot.taskInfo, SurfaceControl())
         }
     }
 
     @Test
-    fun testOnTaskAppeared_duplicateMinimizedRoot_throws() {
-        organizer.createDesk(Display.DEFAULT_DISPLAY, FakeOnCreateCallback())
-        val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
-        val minimizationRootTask = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
-        organizer.onTaskAppeared(minimizationRootTask, SurfaceControl())
+    fun testOnTaskAppeared_duplicateMinimizedRoot_throws() = runTest {
+        val desk = createDesk()
 
         assertThrows(Exception::class.java) {
-            organizer.onTaskAppeared(minimizationRootTask, SurfaceControl())
+            organizer.onTaskAppeared(desk.minimizationRoot.taskInfo, SurfaceControl())
         }
     }
 
     @Test
-    fun testOnTaskVanished_removesRoot() {
+    fun testOnTaskVanished_removesRoot() = runTest {
         val desk = createDesk()
 
         organizer.onTaskVanished(desk.deskRoot.taskInfo)
@@ -178,7 +136,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testOnTaskVanished_removesMinimizedRoot() {
+    fun testOnTaskVanished_removesMinimizedRoot() = runTest {
         val desk = createDesk()
 
         organizer.onTaskVanished(desk.deskRoot.taskInfo)
@@ -188,7 +146,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testDesktopWindowAppearsInDesk() {
+    fun testDesktopWindowAppearsInDesk() = runTest {
         val desk = createDesk()
         val child = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
 
@@ -198,7 +156,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testDesktopWindowAppearsInDeskMinimizationRoot() {
+    fun testDesktopWindowAppearsInDeskMinimizationRoot() = runTest {
         val desk = createDesk()
         val child = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
 
@@ -208,7 +166,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testDesktopWindowMovesToMinimizationRoot() {
+    fun testDesktopWindowMovesToMinimizationRoot() = runTest {
         val desk = createDesk()
         val child = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         organizer.onTaskAppeared(child, SurfaceControl())
@@ -221,7 +179,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testDesktopWindowDisappearsFromDesk() {
+    fun testDesktopWindowDisappearsFromDesk() = runTest {
         val desk = createDesk()
         val child = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
 
@@ -232,7 +190,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testDesktopWindowDisappearsFromDeskMinimizationRoot() {
+    fun testDesktopWindowDisappearsFromDeskMinimizationRoot() = runTest {
         val desk = createDesk()
         val child = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
 
@@ -243,7 +201,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testRemoveDesk_removesDeskRoot() {
+    fun testRemoveDesk_removesDeskRoot() = runTest {
         val desk = createDesk()
 
         val wct = WindowContainerTransaction()
@@ -259,7 +217,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testRemoveDesk_removesMinimizationRoot() {
+    fun testRemoveDesk_removesMinimizationRoot() = runTest {
         val desk = createDesk()
 
         val wct = WindowContainerTransaction()
@@ -275,7 +233,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testActivateDesk() {
+    fun testActivateDesk() = runTest {
         val desk = createDesk()
 
         val wct = WindowContainerTransaction()
@@ -299,7 +257,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testActivateDesk_didNotExist_throws() {
+    fun testActivateDesk_didNotExist_throws() = runTest {
         val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
 
         val wct = WindowContainerTransaction()
@@ -307,7 +265,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testMoveTaskToDesk() {
+    fun testMoveTaskToDesk() = runTest {
         val desk = createDesk()
 
         val desktopTask = createFreeformTask().apply { parentTaskId = -1 }
@@ -333,7 +291,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testMoveTaskToDesk_didNotExist_throws() {
+    fun testMoveTaskToDesk_didNotExist_throws() = runTest {
         val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
 
         val desktopTask = createFreeformTask().apply { parentTaskId = -1 }
@@ -344,7 +302,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testGetDeskAtEnd() {
+    fun testGetDeskAtEnd() = runTest {
         val desk = createDesk()
 
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
@@ -357,7 +315,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testGetDeskAtEnd_inMinimizationRoot() {
+    fun testGetDeskAtEnd_inMinimizationRoot() = runTest {
         val desk = createDesk()
 
         val task = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
@@ -370,27 +328,24 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun testIsDeskActiveAtEnd() {
-        organizer.createDesk(Display.DEFAULT_DISPLAY, FakeOnCreateCallback())
-        val freeformRoot = createFreeformTask().apply { parentTaskId = -1 }
-        freeformRoot.isVisibleRequested = true
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
+    fun testIsDeskActiveAtEnd() = runTest {
+        val desk = createDesk()
 
         val isActive =
             organizer.isDeskActiveAtEnd(
                 change =
-                    TransitionInfo.Change(freeformRoot.token, SurfaceControl()).apply {
-                        taskInfo = freeformRoot
+                    TransitionInfo.Change(desk.deskRoot.token, SurfaceControl()).apply {
+                        taskInfo = desk.deskRoot.taskInfo
                         mode = TRANSIT_TO_FRONT
                     },
-                deskId = freeformRoot.taskId,
+                deskId = desk.deskRoot.deskId,
             )
 
         assertThat(isActive).isTrue()
     }
 
     @Test
-    fun deactivateDesk_clearsLaunchRoot() {
+    fun deactivateDesk_clearsLaunchRoot() = runTest {
         val wct = WindowContainerTransaction()
         val desk = createDesk()
         organizer.activateDesk(wct, desk.deskRoot.deskId)
@@ -409,7 +364,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun isDeskChange_forDeskId() {
+    fun isDeskChange_forDeskId() = runTest {
         val desk = createDesk()
 
         assertThat(
@@ -424,7 +379,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun isDeskChange_forDeskId_inMinimizationRoot() {
+    fun isDeskChange_forDeskId_inMinimizationRoot() = runTest {
         val desk = createDesk()
 
         assertThat(
@@ -442,7 +397,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun isDeskChange_anyDesk() {
+    fun isDeskChange_anyDesk() = runTest {
         val desk = createDesk()
 
         assertThat(
@@ -456,7 +411,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun isDeskChange_anyDesk_inMinimizationRoot() {
+    fun isDeskChange_anyDesk_inMinimizationRoot() = runTest {
         val desk = createDesk()
 
         assertThat(
@@ -473,7 +428,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun minimizeTask() {
+    fun minimizeTask() = runTest {
         val desk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
@@ -486,7 +441,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun minimizeTask_alreadyMinimized_noOp() {
+    fun minimizeTask_alreadyMinimized_noOp() = runTest {
         val desk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
         val wct = WindowContainerTransaction()
@@ -498,7 +453,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun minimizeTask_inDifferentDesk_noOp() {
+    fun minimizeTask_inDifferentDesk_noOp() = runTest {
         val desk = createDesk()
         val otherDesk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = otherDesk.deskRoot.deskId }
@@ -511,7 +466,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun unminimizeTask() {
+    fun unminimizeTask() = runTest {
         val desk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
@@ -528,7 +483,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun unminimizeTask_alreadyUnminimized_noOp() {
+    fun unminimizeTask_alreadyUnminimized_noOp() = runTest {
         val desk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
@@ -542,7 +497,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun unminimizeTask_notInDesk_noOp() {
+    fun unminimizeTask_notInDesk_noOp() = runTest {
         val desk = createDesk()
         val task = createFreeformTask()
         val wct = WindowContainerTransaction()
@@ -553,7 +508,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun reorderTaskToFront() {
+    fun reorderTaskToFront() = runTest {
         val desk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
@@ -573,7 +528,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun reorderTaskToFront_notInDesk_noOp() {
+    fun reorderTaskToFront_notInDesk_noOp() = runTest {
         val desk = createDesk()
         val task = createFreeformTask()
         val wct = WindowContainerTransaction()
@@ -592,7 +547,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun reorderTaskToFront_minimized_unminimizesAndReorders() {
+    fun reorderTaskToFront_minimized_unminimizesAndReorders() = runTest {
         val desk = createDesk()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
@@ -615,7 +570,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskAppeared_visibleDesk_onlyDesk_disablesLaunchAdjacent() {
+    fun onTaskAppeared_visibleDesk_onlyDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
         createDesk(visible = true)
@@ -624,7 +579,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskAppeared_invisibleDesk_onlyDesk_enablesLaunchAdjacent() {
+    fun onTaskAppeared_invisibleDesk_onlyDesk_enablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = false
 
         createDesk(visible = false)
@@ -633,7 +588,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskAppeared_invisibleDesk_otherVisibleDesk_disablesLaunchAdjacent() {
+    fun onTaskAppeared_invisibleDesk_otherVisibleDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
         createDesk(visible = true)
@@ -643,7 +598,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskInfoChanged_deskBecomesVisible_onlyDesk_disablesLaunchAdjacent() {
+    fun onTaskInfoChanged_deskBecomesVisible_onlyDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
         val desk = createDesk(visible = false)
@@ -654,7 +609,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskInfoChanged_deskBecomesInvisible_onlyDesk_enablesLaunchAdjacent() {
+    fun onTaskInfoChanged_deskBecomesInvisible_onlyDesk_enablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = false
 
         val desk = createDesk(visible = true)
@@ -665,7 +620,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskInfoChanged_deskBecomesInvisible_otherVisibleDesk_disablesLaunchAdjacent() {
+    fun onTaskInfoChanged_deskBecomesInvisible_otherVisibleDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
         createDesk(visible = true)
@@ -677,7 +632,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskVanished_visibleDeskDisappears_onlyDesk_enablesLaunchAdjacent() {
+    fun onTaskVanished_visibleDeskDisappears_onlyDesk_enablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = false
 
         val desk = createDesk(visible = true)
@@ -687,7 +642,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
-    fun onTaskVanished_visibleDeskDisappears_otherDeskVisible_disablesLaunchAdjacent() {
+    fun onTaskVanished_visibleDeskDisappears_otherDeskVisible_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
         createDesk(visible = true)
@@ -702,20 +657,39 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
         val minimizationRoot: DeskMinimizationRoot,
     )
 
-    private fun createDesk(visible: Boolean = true): DeskRoots {
-        organizer.createDesk(Display.DEFAULT_DISPLAY, FakeOnCreateCallback())
-        val freeformRoot =
+    private suspend fun createDesk(visible: Boolean = true): DeskRoots {
+        val freeformRootTask =
             createFreeformTask().apply {
                 parentTaskId = -1
                 isVisible = visible
+                isVisibleRequested = visible
             }
-        organizer.onTaskAppeared(freeformRoot, SurfaceControl())
-        val minimizationRoot = createFreeformTask().apply { parentTaskId = -1 }
-        organizer.onTaskAppeared(minimizationRoot, SurfaceControl())
-        return DeskRoots(
-            organizer.deskRootsByDeskId[freeformRoot.taskId],
-            checkNotNull(organizer.deskMinimizationRootsByDeskId[freeformRoot.taskId]),
-        )
+        val minimizationRootTask = createFreeformTask().apply { parentTaskId = -1 }
+        Mockito.reset(mockShellTaskOrganizer)
+        whenever(
+                mockShellTaskOrganizer.createRootTask(
+                    Display.DEFAULT_DISPLAY,
+                    WINDOWING_MODE_FREEFORM,
+                    organizer,
+                    true,
+                )
+            )
+            .thenAnswer { invocation ->
+                val listener = (invocation.arguments[2] as TaskListener)
+                listener.onTaskAppeared(freeformRootTask, SurfaceControl())
+            }
+            .thenAnswer { invocation ->
+                val listener = (invocation.arguments[2] as TaskListener)
+                listener.onTaskAppeared(minimizationRootTask, SurfaceControl())
+            }
+        val deskId = organizer.createDesk(Display.DEFAULT_DISPLAY)
+        assertEquals(freeformRootTask.taskId, deskId)
+        val deskRoot = assertNotNull(organizer.deskRootsByDeskId.get(freeformRootTask.taskId))
+        val minimizationRoot =
+            assertNotNull(organizer.deskMinimizationRootsByDeskId[freeformRootTask.taskId])
+        assertThat(minimizationRoot.deskId).isEqualTo(freeformRootTask.taskId)
+        assertThat(minimizationRoot.rootId).isEqualTo(minimizationRootTask.taskId)
+        return DeskRoots(deskRoot, minimizationRoot)
     }
 
     private fun WindowContainerTransaction.hasMinimizationHops(
@@ -738,14 +712,4 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
                 hop.newParent == desk.deskRoot.token.asBinder() &&
                 hop.toTop
         }
-
-    private class FakeOnCreateCallback : DesksOrganizer.OnCreateCallback {
-        var deskId: Int? = null
-        val created: Boolean
-            get() = deskId != null
-
-        override fun onCreated(deskId: Int) {
-            this.deskId = deskId
-        }
-    }
 }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt
index dd9e6ca0deae..4440d4e801fe 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/persistence/DesktopRepositoryInitializerTest.kt
@@ -17,7 +17,6 @@
 package com.android.wm.shell.desktopmode.persistence
 
 import android.os.UserManager
-import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.testing.AndroidTestingRunner
 import android.view.Display.DEFAULT_DISPLAY
@@ -82,10 +81,27 @@ class DesktopRepositoryInitializerTest : ShellTestCase() {
     }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE, FLAG_ENABLE_DESKTOP_WINDOWING_HSUM)
-    /** TODO: b/362720497 - add multi-desk version when implemented. */
-    @DisableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
-    fun initWithPersistence_multipleUsers_addedCorrectly_multiDesksDisabled() =
+    fun init_updatesFlow() =
+        runTest(StandardTestDispatcher()) {
+            whenever(persistentRepository.getUserDesktopRepositoryMap())
+                .thenReturn(mapOf(USER_ID_1 to desktopRepositoryState1))
+            whenever(persistentRepository.getDesktopRepositoryState(USER_ID_1))
+                .thenReturn(desktopRepositoryState1)
+            whenever(persistentRepository.readDesktop(USER_ID_1, DESKTOP_ID_1)).thenReturn(desktop1)
+            whenever(persistentRepository.readDesktop(USER_ID_1, DESKTOP_ID_2)).thenReturn(desktop2)
+
+            repositoryInitializer.initialize(desktopUserRepositories)
+
+            assertThat(repositoryInitializer.isInitialized.value).isTrue()
+        }
+
+    @Test
+    @EnableFlags(
+        FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE,
+        FLAG_ENABLE_DESKTOP_WINDOWING_HSUM,
+        FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun initWithPersistence_multipleUsers_addedCorrectly() =
         runTest(StandardTestDispatcher()) {
             whenever(persistentRepository.getUserDesktopRepositoryMap())
                 .thenReturn(
@@ -104,50 +120,74 @@ class DesktopRepositoryInitializerTest : ShellTestCase() {
 
             repositoryInitializer.initialize(desktopUserRepositories)
 
-            // Desktop Repository currently returns all tasks across desktops for a specific user
-            // since the repository currently doesn't handle desktops. This test logic should be
-            // updated
-            // once the repository handles multiple desktops.
             assertThat(
-                    desktopUserRepositories.getProfile(USER_ID_1).getActiveTasks(DEFAULT_DISPLAY)
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getActiveTaskIdsInDesk(DESKTOP_ID_1)
                 )
-                .containsExactly(1, 3, 4, 5)
+                .containsExactly(1, 3)
                 .inOrder()
             assertThat(
                     desktopUserRepositories
                         .getProfile(USER_ID_1)
-                        .getExpandedTasksOrdered(DEFAULT_DISPLAY)
+                        .getActiveTaskIdsInDesk(DESKTOP_ID_2)
                 )
-                .containsExactly(5, 1)
+                .containsExactly(4, 5)
                 .inOrder()
             assertThat(
-                    desktopUserRepositories.getProfile(USER_ID_1).getMinimizedTasks(DEFAULT_DISPLAY)
+                    desktopUserRepositories
+                        .getProfile(USER_ID_2)
+                        .getActiveTaskIdsInDesk(DESKTOP_ID_3)
                 )
-                .containsExactly(3, 4)
+                .containsExactly(7, 8)
+                .inOrder()
+            assertThat(
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getExpandedTasksIdsInDeskOrdered(DESKTOP_ID_1)
+                )
+                .containsExactly(1)
                 .inOrder()
-
             assertThat(
-                    desktopUserRepositories.getProfile(USER_ID_2).getActiveTasks(DEFAULT_DISPLAY)
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getExpandedTasksIdsInDeskOrdered(DESKTOP_ID_2)
                 )
-                .containsExactly(7, 8)
+                .containsExactly(5)
                 .inOrder()
             assertThat(
                     desktopUserRepositories
                         .getProfile(USER_ID_2)
-                        .getExpandedTasksOrdered(DEFAULT_DISPLAY)
+                        .getExpandedTasksIdsInDeskOrdered(DESKTOP_ID_3)
+                )
+                .containsExactly(7)
+                .inOrder()
+            assertThat(
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getMinimizedTaskIdsInDesk(DESKTOP_ID_1)
+                )
+                .containsExactly(3)
+                .inOrder()
+            assertThat(
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getMinimizedTaskIdsInDesk(DESKTOP_ID_2)
                 )
-                .contains(7)
+                .containsExactly(4)
+                .inOrder()
             assertThat(
-                    desktopUserRepositories.getProfile(USER_ID_2).getMinimizedTasks(DEFAULT_DISPLAY)
+                    desktopUserRepositories
+                        .getProfile(USER_ID_2)
+                        .getMinimizedTaskIdsInDesk(DESKTOP_ID_3)
                 )
                 .containsExactly(8)
+                .inOrder()
         }
 
     @Test
-    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE)
-    /** TODO: b/362720497 - add multi-desk version when implemented. */
-    @DisableFlags(FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
-    fun initWithPersistence_singleUser_addedCorrectly_multiDesksDisabled() =
+    @EnableFlags(FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE, FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun initWithPersistence_singleUser_addedCorrectly() =
         runTest(StandardTestDispatcher()) {
             whenever(persistentRepository.getUserDesktopRepositoryMap())
                 .thenReturn(mapOf(USER_ID_1 to desktopRepositoryState1))
@@ -161,23 +201,44 @@ class DesktopRepositoryInitializerTest : ShellTestCase() {
             assertThat(
                     desktopUserRepositories
                         .getProfile(USER_ID_1)
-                        .getActiveTaskIdsInDesk(deskId = DEFAULT_DISPLAY)
+                        .getActiveTaskIdsInDesk(DESKTOP_ID_1)
+                )
+                .containsExactly(1, 3)
+                .inOrder()
+            assertThat(
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getActiveTaskIdsInDesk(DESKTOP_ID_2)
+                )
+                .containsExactly(4, 5)
+                .inOrder()
+            assertThat(
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getExpandedTasksIdsInDeskOrdered(DESKTOP_ID_1)
+                )
+                .containsExactly(1)
+                .inOrder()
+            assertThat(
+                    desktopUserRepositories
+                        .getProfile(USER_ID_1)
+                        .getExpandedTasksIdsInDeskOrdered(DESKTOP_ID_2)
                 )
-                .containsExactly(1, 3, 4, 5)
+                .containsExactly(5)
                 .inOrder()
             assertThat(
                     desktopUserRepositories
                         .getProfile(USER_ID_1)
-                        .getExpandedTasksIdsInDeskOrdered(deskId = DEFAULT_DISPLAY)
+                        .getMinimizedTaskIdsInDesk(DESKTOP_ID_1)
                 )
-                .containsExactly(5, 1)
+                .containsExactly(3)
                 .inOrder()
             assertThat(
                     desktopUserRepositories
                         .getProfile(USER_ID_1)
-                        .getMinimizedTaskIdsInDesk(deskId = DEFAULT_DISPLAY)
+                        .getMinimizedTaskIdsInDesk(DESKTOP_ID_2)
                 )
-                .containsExactly(3, 4)
+                .containsExactly(4)
                 .inOrder()
         }