Merge "Desks: Add multi-user support" into main

7 files changed, 335 insertions, 141 deletions

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 25737c4950d6..127a8a9c41ad 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
@@ -49,9 +49,6 @@ class DesktopDisplayEventHandler(
     private val desktopDisplayModeController: DesktopDisplayModeController,
 ) : OnDisplaysChangedListener, OnDeskRemovedListener {
 
-    private val desktopRepository: DesktopRepository
-        get() = desktopUserRepositories.current
-
     init {
         shellInit.addInitCallback({ onInit() }, this)
     }
@@ -66,7 +63,7 @@ class DesktopDisplayEventHandler(
                 object : UserChangeListener {
                     override fun onUserChanged(newUserId: Int, userContext: Context) {
                         val displayIds = rootTaskDisplayAreaOrganizer.displayIds
-                        createDefaultDesksIfNeeded(displayIds.toSet())
+                        createDefaultDesksIfNeeded(displayIds.toSet(), newUserId)
                     }
                 }
             )
@@ -78,7 +75,7 @@ class DesktopDisplayEventHandler(
             desktopDisplayModeController.refreshDisplayWindowingMode()
         }
 
-        createDefaultDesksIfNeeded(displayIds = setOf(displayId))
+        createDefaultDesksIfNeeded(displayIds = setOf(displayId), userId = null)
     }
 
     override fun onDisplayRemoved(displayId: Int) {
@@ -99,23 +96,22 @@ class DesktopDisplayEventHandler(
     }
 
     override fun onDeskRemoved(lastDisplayId: Int, deskId: Int) {
-        val remainingDesks = desktopRepository.getNumberOfDesks(lastDisplayId)
-        if (remainingDesks == 0) {
-            logV("All desks removed from display#$lastDisplayId")
-            createDefaultDesksIfNeeded(setOf(lastDisplayId))
-        }
+        createDefaultDesksIfNeeded(setOf(lastDisplayId), userId = null)
     }
 
-    private fun createDefaultDesksIfNeeded(displayIds: Set<Int>) {
+    private fun createDefaultDesksIfNeeded(displayIds: Set<Int>, userId: Int?) {
         if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) return
         logV("createDefaultDesksIfNeeded displays=%s", displayIds)
         mainScope.launch {
             desktopRepositoryInitializer.isInitialized.collect { initialized ->
                 if (!initialized) return@collect
+                val repository =
+                    userId?.let { desktopUserRepositories.getProfile(userId) }
+                        ?: desktopUserRepositories.current
                 displayIds
                     .filter { displayId -> displayId != Display.INVALID_DISPLAY }
                     .filter { displayId -> supportsDesks(displayId) }
-                    .filter { displayId -> desktopRepository.getNumberOfDesks(displayId) == 0 }
+                    .filter { displayId -> repository.getNumberOfDesks(displayId) == 0 }
                     .also { displaysNeedingDesk ->
                         logV(
                             "createDefaultDesksIfNeeded creating default desks in displays=%s",
@@ -125,7 +121,7 @@ class DesktopDisplayEventHandler(
                     .forEach { displayId ->
                         // TODO: b/393978539 - consider activating the desk on creation when
                         //  applicable, such as for connected displays.
-                        desktopTasksController.createDesk(displayId)
+                        desktopTasksController.createDesk(displayId, repository.userId)
                     }
                 cancel()
             }
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 6214f329e0fd..2d6253b642be 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
@@ -42,6 +42,7 @@ import android.os.Handler
 import android.os.IBinder
 import android.os.SystemProperties
 import android.os.UserHandle
+import android.os.UserManager
 import android.util.Slog
 import android.view.Display
 import android.view.Display.DEFAULT_DISPLAY
@@ -115,7 +116,6 @@ 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
@@ -163,6 +163,7 @@ import java.util.Optional
 import java.util.concurrent.Executor
 import java.util.concurrent.TimeUnit
 import java.util.function.Consumer
+import kotlin.coroutines.suspendCoroutine
 import kotlin.jvm.optionals.getOrNull
 
 /**
@@ -283,14 +284,8 @@ class DesktopTasksController(
 
         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 re-create desk of another user.")
-                        null
-                    } else {
-                        desksOrganizer.createDesk(destinationDisplayId)
-                    }
+                DeskRecreationFactory { deskUserId, destinationDisplayId, _ ->
+                    createDeskSuspending(displayId = destinationDisplayId, userId = deskUserId)
                 }
         }
     }
@@ -493,22 +488,55 @@ class DesktopTasksController(
         runOnTransitStart?.invoke(transition)
     }
 
-    /** Creates a new desk in the given display. */
-    fun createDesk(displayId: Int) {
+    /** Adds a new desk to the given display for the given user. */
+    fun createDesk(displayId: Int, userId: Int = this.userId) {
+        logV("addDesk displayId=%d, userId=%d", displayId, userId)
+        val repository = userRepositories.getProfile(userId)
+        createDesk(displayId, userId) { deskId ->
+            if (deskId == null) {
+                logW("Failed to add desk in displayId=%d for userId=%d", displayId, userId)
+            } else {
+                repository.addDesk(displayId = displayId, deskId = deskId)
+            }
+        }
+    }
+
+    private fun createDesk(displayId: Int, userId: Int = this.userId, onResult: (Int?) -> Unit) {
         if (displayId == Display.INVALID_DISPLAY) {
             logW("createDesk attempt with invalid displayId", displayId)
+            onResult(null)
             return
         }
-        if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            desksOrganizer.createDesk(displayId) { deskId ->
-                taskRepository.addDesk(displayId = displayId, deskId = deskId)
-            }
-        } else {
+        if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
             // In single-desk, the desk reuses the display id.
-            taskRepository.addDesk(displayId = displayId, deskId = displayId)
+            logD("createDesk reusing displayId=%d for single-desk", displayId)
+            onResult(displayId)
+            return
+        }
+        if (
+            DesktopModeFlags.ENABLE_DESKTOP_WINDOWING_HSUM.isTrue &&
+                UserManager.isHeadlessSystemUserMode() &&
+                UserHandle.USER_SYSTEM == userId
+        ) {
+            logW("createDesk ignoring attempt for system user")
+            return
+        }
+        desksOrganizer.createDesk(displayId, userId) { deskId ->
+            logD(
+                "createDesk obtained deskId=%d for displayId=%d and userId=%d",
+                deskId,
+                displayId,
+                userId,
+            )
+            onResult(deskId)
         }
     }
 
+    private suspend fun createDeskSuspending(displayId: Int, userId: Int = this.userId): Int? =
+        suspendCoroutine { cont ->
+            createDesk(displayId, userId) { deskId -> cont.resumeWith(Result.success(deskId)) }
+        }
+
     /** Moves task to desktop mode if task is running, else launches it in desktop mode. */
     @JvmOverloads
     fun moveTaskToDefaultDeskAndActivate(
@@ -3024,18 +3052,17 @@ class DesktopTasksController(
             }
 
         val wct = WindowContainerTransaction()
-        if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
-            tasksToRemove.forEach {
-                val task = shellTaskOrganizer.getRunningTaskInfo(it)
-                if (task != null) {
-                    wct.removeTask(task.token)
-                } else {
-                    recentTasksController?.removeBackgroundTask(it)
-                }
+        tasksToRemove.forEach {
+            // TODO: b/404595635 - consider moving this block into [DesksOrganizer].
+            val task = shellTaskOrganizer.getRunningTaskInfo(it)
+            if (task != null) {
+                wct.removeTask(task.token)
+            } else {
+                recentTasksController?.removeBackgroundTask(it)
             }
-        } else {
-            // TODO: 362720497 - double check background tasks are also removed.
-            desksOrganizer.removeDesk(wct, deskId)
+        }
+        if (DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue) {
+            desksOrganizer.removeDesk(wct, deskId, userId)
         }
         if (!DesktopExperienceFlags.ENABLE_MULTIPLE_DESKTOPS_BACKEND.isTrue && wct.isEmpty) return
         val transition = transitions.startTransition(TRANSIT_CLOSE, wct, /* handler= */ null)
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 1effcdb20505..605465b15468 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,13 +18,11 @@ 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 {
-    /** Creates a new desk container in the given display. */
-    fun createDesk(displayId: Int, callback: OnCreateCallback)
+    /** Creates a new desk container to use in the given display for the given user. */
+    fun createDesk(displayId: Int, userId: Int, callback: OnCreateCallback)
 
     /** Activates the given desk, making it visible in its display. */
     fun activateDesk(wct: WindowContainerTransaction, deskId: Int)
@@ -32,8 +30,8 @@ interface DesksOrganizer {
     /** Deactivates the given desk, removing it as the default launch container for new tasks. */
     fun deactivateDesk(wct: WindowContainerTransaction, deskId: Int)
 
-    /** Removes the given desk and its desktop windows. */
-    fun removeDesk(wct: WindowContainerTransaction, deskId: Int)
+    /** Removes the given desk of the given user. */
+    fun removeDesk(wct: WindowContainerTransaction, deskId: Int, userId: Int)
 
     /** Moves the given task to the given desk. */
     fun moveTaskToDesk(
@@ -87,9 +85,3 @@ 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/multidesks/RootTaskDesksOrganizer.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt
index c30987ac7640..e4edeb95be6d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/multidesks/RootTaskDesksOrganizer.kt
@@ -30,6 +30,7 @@ import android.window.TransitionInfo
 import android.window.WindowContainerToken
 import android.window.WindowContainerTransaction
 import androidx.core.util.forEach
+import androidx.core.util.valueIterator
 import com.android.internal.annotations.VisibleForTesting
 import com.android.internal.protolog.ProtoLog
 import com.android.wm.shell.ShellTaskOrganizer
@@ -40,7 +41,13 @@ import com.android.wm.shell.sysui.ShellCommandHandler
 import com.android.wm.shell.sysui.ShellInit
 import java.io.PrintWriter
 
-/** A [DesksOrganizer] that uses root tasks as the container of each desk. */
+/**
+ * A [DesksOrganizer] that uses root tasks as the container of each desk.
+ *
+ * Note that root tasks are reusable between multiple users at the same time, and may also be
+ * pre-created to have one ready for the first entry to the default desk, so root-task existence
+ * does not imply a formal desk exists to the user.
+ */
 class RootTaskDesksOrganizer(
     shellInit: ShellInit,
     shellCommandHandler: ShellCommandHandler,
@@ -65,9 +72,26 @@ class RootTaskDesksOrganizer(
         }
     }
 
-    override fun createDesk(displayId: Int, callback: OnCreateCallback) {
-        logV("createDesk in display: %d", displayId)
-        createDeskRootRequests += CreateDeskRequest(displayId, callback)
+    override fun createDesk(displayId: Int, userId: Int, callback: OnCreateCallback) {
+        logV("createDesk in displayId=%d userId=%s", displayId, userId)
+        // Find an existing desk that is not yet used by this user.
+        val unassignedDesk =
+            deskRootsByDeskId
+                .valueIterator()
+                .asSequence()
+                .filterNot { desk -> userId in desk.users }
+                .firstOrNull()
+        if (unassignedDesk != null) {
+            unassignedDesk.users.add(userId)
+            callback.onCreated(unassignedDesk.deskId)
+            return
+        }
+        createDeskRoot(displayId, userId, callback)
+    }
+
+    private fun createDeskRoot(displayId: Int, userId: Int, callback: OnCreateCallback) {
+        logV("createDeskRoot in display: %d for user: %d", displayId, userId)
+        createDeskRootRequests += CreateDeskRequest(displayId, userId, callback)
         shellTaskOrganizer.createRootTask(
             displayId,
             WINDOWING_MODE_FREEFORM,
@@ -76,31 +100,52 @@ class RootTaskDesksOrganizer(
         )
     }
 
-    override fun removeDesk(wct: WindowContainerTransaction, deskId: Int) {
-        logV("removeDesk %d", deskId)
-        deskRootsByDeskId[deskId]?.let { root -> wct.removeRootTask(root.token) }
-        deskMinimizationRootsByDeskId[deskId]?.let { root -> wct.removeRootTask(root.token) }
+    override fun removeDesk(wct: WindowContainerTransaction, deskId: Int, userId: Int) {
+        logV("removeDesk %d for userId=%d", deskId, userId)
+        val deskRoot = deskRootsByDeskId[deskId]
+        if (deskRoot == null) {
+            logW("removeDesk attempted to remove non-existent desk=%d", deskId)
+            return
+        }
+        updateLaunchRoot(wct, deskId, enabled = false)
+        deskRoot.users.remove(userId)
+        if (deskRoot.users.isEmpty()) {
+            // No longer in use by any users, remove it completely.
+            logD("removeDesk %d is no longer used by any users, removing it completely", deskId)
+            wct.removeRootTask(deskRoot.token)
+            deskMinimizationRootsByDeskId[deskId]?.let { root -> wct.removeRootTask(root.token) }
+        }
     }
 
     override fun activateDesk(wct: WindowContainerTransaction, deskId: Int) {
         logV("activateDesk %d", deskId)
         val root = checkNotNull(deskRootsByDeskId[deskId]) { "Root not found for desk: $deskId" }
         wct.reorder(root.token, /* onTop= */ true)
-        wct.setLaunchRoot(
-            /* container= */ root.taskInfo.token,
-            /* windowingModes= */ intArrayOf(WINDOWING_MODE_FREEFORM, WINDOWING_MODE_UNDEFINED),
-            /* activityTypes= */ intArrayOf(ACTIVITY_TYPE_UNDEFINED, ACTIVITY_TYPE_STANDARD),
-        )
+        updateLaunchRoot(wct, deskId, enabled = true)
     }
 
     override fun deactivateDesk(wct: WindowContainerTransaction, deskId: Int) {
         logV("deactivateDesk %d", deskId)
+        updateLaunchRoot(wct, deskId, enabled = false)
+    }
+
+    private fun updateLaunchRoot(wct: WindowContainerTransaction, deskId: Int, enabled: Boolean) {
         val root = checkNotNull(deskRootsByDeskId[deskId]) { "Root not found for desk: $deskId" }
-        wct.setLaunchRoot(
-            /* container= */ root.taskInfo.token,
-            /* windowingModes= */ null,
-            /* activityTypes= */ null,
-        )
+        root.isLaunchRootRequested = enabled
+        logD("updateLaunchRoot deskId=%d enabled=%b", deskId, enabled)
+        if (enabled) {
+            wct.setLaunchRoot(
+                /* container= */ root.taskInfo.token,
+                /* windowingModes= */ intArrayOf(WINDOWING_MODE_FREEFORM, WINDOWING_MODE_UNDEFINED),
+                /* activityTypes= */ intArrayOf(ACTIVITY_TYPE_UNDEFINED, ACTIVITY_TYPE_STANDARD),
+            )
+        } else {
+            wct.setLaunchRoot(
+                /* container= */ root.taskInfo.token,
+                /* windowingModes= */ null,
+                /* activityTypes= */ null,
+            )
+        }
     }
 
     override fun moveTaskToDesk(
@@ -275,7 +320,13 @@ class RootTaskDesksOrganizer(
             // Appearing root matches desk request.
             val deskId = taskInfo.taskId
             logV("Desk #$deskId appeared")
-            deskRootsByDeskId[deskId] = DeskRoot(deskId, taskInfo, leash)
+            deskRootsByDeskId[deskId] =
+                DeskRoot(
+                    deskId = deskId,
+                    taskInfo = taskInfo,
+                    leash = leash,
+                    users = mutableSetOf(deskRequest.userId),
+                )
             createDeskRootRequests.remove(deskRequest)
             deskRequest.onCreateCallback.onCreated(deskId)
             createDeskMinimizationRoot(displayId = appearingInDisplayId, deskId = deskId)
@@ -430,6 +481,8 @@ class RootTaskDesksOrganizer(
         val taskInfo: RunningTaskInfo,
         val leash: SurfaceControl,
         val children: MutableSet<Int> = mutableSetOf(),
+        val users: MutableSet<Int> = mutableSetOf(),
+        var isLaunchRootRequested: Boolean = false,
     ) {
         val token: WindowContainerToken = taskInfo.token
     }
@@ -449,15 +502,24 @@ class RootTaskDesksOrganizer(
 
     private data class CreateDeskRequest(
         val displayId: Int,
+        val userId: Int,
         val onCreateCallback: OnCreateCallback,
     )
 
     private data class CreateDeskMinimizationRootRequest(val displayId: Int, val deskId: Int)
 
+    private fun logD(msg: String, vararg arguments: Any?) {
+        ProtoLog.d(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
     private fun logV(msg: String, vararg arguments: Any?) {
         ProtoLog.v(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
     }
 
+    private fun logW(msg: String, vararg arguments: Any?) {
+        ProtoLog.w(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
+    }
+
     private fun logE(msg: String, vararg arguments: Any?) {
         ProtoLog.e(WM_SHELL_DESKTOP_MODE, "%s: $msg", TAG, *arguments)
     }
@@ -473,7 +535,9 @@ class RootTaskDesksOrganizer(
             val minimizationRoot = deskMinimizationRootsByDeskId[deskId]
             pw.println("$innerPrefix  #$deskId visible=${root.taskInfo.isVisible}")
             pw.println("$innerPrefix    displayId=${root.taskInfo.displayId}")
+            pw.println("$innerPrefix    isLaunchRootRequested=${root.isLaunchRootRequested}")
             pw.println("$innerPrefix    children=${root.children}")
+            pw.println("$innerPrefix    users=${root.users}")
             pw.println("$innerPrefix    minimization root:")
             pw.println("$innerPrefix      rootId=${minimizationRoot?.rootId}")
             if (minimizationRoot != null) {
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 85a431be8e8b..565145b8616e 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
@@ -51,6 +51,7 @@ import org.mockito.Mockito.times
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.clearInvocations
+import org.mockito.kotlin.mock
 import org.mockito.kotlin.whenever
 import org.mockito.quality.Strictness
 
@@ -213,12 +214,15 @@ class DesktopDisplayEventHandlerTest : ShellTestCase() {
     @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun testUserChanged_createsDeskWhenNeeded() =
         testScope.runTest {
+            val userId = 11
             whenever(DesktopModeStatus.canEnterDesktopMode(context)).thenReturn(true)
             val userChangeListenerCaptor = argumentCaptor<UserChangeListener>()
             verify(mockShellController).addUserChangeListener(userChangeListenerCaptor.capture())
-            whenever(mockDesktopRepository.getNumberOfDesks(displayId = 2)).thenReturn(0)
-            whenever(mockDesktopRepository.getNumberOfDesks(displayId = 3)).thenReturn(0)
-            whenever(mockDesktopRepository.getNumberOfDesks(displayId = 4)).thenReturn(1)
+            val mockRepository = mock<DesktopRepository>()
+            whenever(mockDesktopUserRepositories.getProfile(userId)).thenReturn(mockRepository)
+            whenever(mockRepository.getNumberOfDesks(displayId = 2)).thenReturn(0)
+            whenever(mockRepository.getNumberOfDesks(displayId = 3)).thenReturn(0)
+            whenever(mockRepository.getNumberOfDesks(displayId = 4)).thenReturn(1)
             whenever(mockRootTaskDisplayAreaOrganizer.displayIds).thenReturn(intArrayOf(2, 3, 4))
             desktopRepositoryInitializer.initialize(mockDesktopUserRepositories)
             handler.onDisplayAdded(displayId = 2)
@@ -227,7 +231,7 @@ class DesktopDisplayEventHandlerTest : ShellTestCase() {
             runCurrent()
 
             clearInvocations(mockDesktopTasksController)
-            userChangeListenerCaptor.lastValue.onUserChanged(1, context)
+            userChangeListenerCaptor.lastValue.onUserChanged(userId, context)
             runCurrent()
 
             verify(mockDesktopTasksController).createDesk(displayId = 2)
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/DesktopTasksControllerTest.kt
index a222d392848b..d81786b5e6a5 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
@@ -50,6 +50,7 @@ import android.os.Binder
 import android.os.Bundle
 import android.os.Handler
 import android.os.IBinder
+import android.os.UserHandle
 import android.os.UserManager
 import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
@@ -5537,7 +5538,7 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
 
         controller.removeDesk(deskId = 2)
 
-        verify(desksOrganizer).removeDesk(any(), eq(2))
+        verify(desksOrganizer).removeDesk(any(), eq(2), any())
         verify(desksTransitionsObserver)
             .addPendingTransition(
                 argThat {
@@ -5553,6 +5554,49 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
         Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION,
         Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
     )
+    fun removeDesk_multipleDesks_removesRunningTasks() {
+        val transition = Binder()
+        whenever(transitions.startTransition(eq(TRANSIT_CLOSE), any(), anyOrNull()))
+            .thenReturn(transition)
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 2)
+        val task1 = setUpFreeformTask(deskId = 2)
+        val task2 = setUpFreeformTask(deskId = 2)
+        val task3 = setUpFreeformTask(deskId = 2)
+
+        controller.removeDesk(deskId = 2)
+
+        val wct = getLatestWct(TRANSIT_CLOSE)
+        wct.assertRemove(task1.token)
+        wct.assertRemove(task2.token)
+        wct.assertRemove(task3.token)
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
+    fun removeDesk_multipleDesks_removesRecentTasks() {
+        val transition = Binder()
+        whenever(transitions.startTransition(eq(TRANSIT_CLOSE), any(), anyOrNull()))
+            .thenReturn(transition)
+        taskRepository.addDesk(DEFAULT_DISPLAY, deskId = 2)
+        val task1 = setUpFreeformTask(deskId = 2, background = true)
+        val task2 = setUpFreeformTask(deskId = 2, background = true)
+        val task3 = setUpFreeformTask(deskId = 2, background = true)
+
+        controller.removeDesk(deskId = 2)
+
+        verify(recentTasksController).removeBackgroundTask(task1.taskId)
+        verify(recentTasksController).removeBackgroundTask(task2.taskId)
+        verify(recentTasksController).removeBackgroundTask(task3.taskId)
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_BACK_NAVIGATION,
+        Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND,
+    )
     fun activateDesk_multipleDesks_addsPendingTransition() {
         val deskId = 0
         val transition = Binder()
@@ -7509,11 +7553,12 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
     fun testCreateDesk() {
         val currentDeskCount = taskRepository.getNumberOfDesks(DEFAULT_DISPLAY)
-        whenever(desksOrganizer.createDesk(eq(DEFAULT_DISPLAY), any())).thenAnswer { invocation ->
-            (invocation.arguments[1] as DesksOrganizer.OnCreateCallback).onCreated(deskId = 5)
+        whenever(desksOrganizer.createDesk(eq(DEFAULT_DISPLAY), any(), any())).thenAnswer {
+            invocation ->
+            (invocation.arguments[2] as DesksOrganizer.OnCreateCallback).onCreated(deskId = 5)
         }
 
-        controller.createDesk(DEFAULT_DISPLAY)
+        controller.createDesk(DEFAULT_DISPLAY, taskRepository.userId)
 
         assertThat(taskRepository.getNumberOfDesks(DEFAULT_DISPLAY)).isEqualTo(currentDeskCount + 1)
     }
@@ -7523,7 +7568,17 @@ class DesktopTasksControllerTest(flags: FlagsParameterization) : ShellTestCase()
     fun testCreateDesk_invalidDisplay_dropsRequest() {
         controller.createDesk(INVALID_DISPLAY)
 
-        verify(desksOrganizer, never()).createDesk(any(), any())
+        verify(desksOrganizer, never()).createDesk(any(), any(), any())
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun testCreateDesk_systemUser_dropsRequest() {
+        assumeTrue(UserManager.isHeadlessSystemUserMode())
+
+        controller.createDesk(DEFAULT_DISPLAY, UserHandle.USER_SYSTEM)
+
+        verify(desksOrganizer, never()).createDesk(any(), any(), any())
     }
 
     @Test
@@ -8256,6 +8311,15 @@ private fun WindowContainerTransaction.assertReorderSequenceInRange(
         .inOrder()
 }
 
+private fun WindowContainerTransaction.assertRemove(token: WindowContainerToken) {
+    assertThat(
+            hierarchyOps.any { hop ->
+                hop.container == token.asBinder() && hop.type == HIERARCHY_OP_TYPE_REMOVE_TASK
+            }
+        )
+        .isTrue()
+}
+
 private fun WindowContainerTransaction.assertRemoveAt(index: Int, token: WindowContainerToken) {
     assertIndexInBounds(index)
     val op = hierarchyOps[index]
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 e57fc38e3607..34f832b4fba4 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
@@ -19,7 +19,7 @@ import android.app.ActivityManager
 import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
 import android.app.WindowConfiguration.WINDOWING_MODE_UNDEFINED
 import android.testing.AndroidTestingRunner
-import android.view.Display
+import android.view.Display.DEFAULT_DISPLAY
 import android.view.SurfaceControl
 import android.view.WindowManager.TRANSIT_TO_FRONT
 import android.window.TransitionInfo
@@ -41,9 +41,9 @@ import com.android.wm.shell.desktopmode.multidesks.RootTaskDesksOrganizer.DeskRo
 import com.android.wm.shell.sysui.ShellCommandHandler
 import com.android.wm.shell.sysui.ShellInit
 import com.google.common.truth.Truth.assertThat
+import kotlin.coroutines.suspendCoroutine
 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
@@ -86,11 +86,21 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
         organizer.setOnDesktopTaskInfoChangedListener(taskInfoChangedListener)
     }
 
-    @Test fun testCreateDesk_createsDeskAndMinimizationRoots() = runTest { createDesk() }
+    @Test fun testCreateDesk_createsDeskAndMinimizationRoots() = runTest { createDeskSuspending() }
+
+    @Test
+    fun testCreateDesk_rootExistsForOtherUser_reusesRoot() = runTest {
+        val desk = createDeskSuspending(userId = PRIMARY_USER_ID)
+
+        val deskId =
+            organizer.createDeskSuspending(displayId = DEFAULT_DISPLAY, userId = SECONDARY_USER_ID)
+
+        assertThat(deskId).isEqualTo(desk.deskRoot.deskId)
+    }
 
     @Test
     fun testCreateMinimizationRoot_marksHidden() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         verify(mockShellTaskOrganizer)
             .applyTransaction(
@@ -115,7 +125,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testOnTaskAppeared_duplicateRoot_throws() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         assertThrows(Exception::class.java) {
             organizer.onTaskAppeared(desk.deskRoot.taskInfo, SurfaceControl())
@@ -124,7 +134,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testOnTaskAppeared_duplicateMinimizedRoot_throws() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         assertThrows(Exception::class.java) {
             organizer.onTaskAppeared(desk.minimizationRoot.taskInfo, SurfaceControl())
@@ -133,7 +143,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testOnTaskVanished_removesRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         organizer.onTaskVanished(desk.deskRoot.taskInfo)
 
@@ -142,7 +152,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testOnTaskVanished_removesMinimizedRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         organizer.onTaskVanished(desk.deskRoot.taskInfo)
         organizer.onTaskVanished(desk.minimizationRoot.taskInfo)
@@ -152,7 +162,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testDesktopWindowAppearsInDesk() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val child = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
 
         organizer.onTaskAppeared(child, SurfaceControl())
@@ -162,7 +172,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testDesktopWindowAppearsInDeskMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val child = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
 
         organizer.onTaskAppeared(child, SurfaceControl())
@@ -172,7 +182,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testDesktopWindowMovesToMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val child = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         organizer.onTaskAppeared(child, SurfaceControl())
 
@@ -185,7 +195,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testDesktopWindowDisappearsFromDesk() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val child = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
 
         organizer.onTaskAppeared(child, SurfaceControl())
@@ -196,7 +206,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testDesktopWindowDisappearsFromDeskMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val child = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
 
         organizer.onTaskAppeared(child, SurfaceControl())
@@ -206,11 +216,23 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
+    fun testRemoveDesk_disablesAsLaunchRoot() = runTest {
+        val desk = createDeskSuspending(userId = PRIMARY_USER_ID)
+        val wct = WindowContainerTransaction()
+        organizer.activateDesk(wct, desk.deskRoot.deskId)
+        assertThat(desk.deskRoot.isLaunchRootRequested).isTrue()
+
+        organizer.removeDesk(wct, desk.deskRoot.deskId, userId = PRIMARY_USER_ID)
+
+        assertThat(desk.deskRoot.isLaunchRootRequested).isFalse()
+    }
+
+    @Test
     fun testRemoveDesk_removesDeskRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending(userId = PRIMARY_USER_ID)
 
         val wct = WindowContainerTransaction()
-        organizer.removeDesk(wct, desk.deskRoot.deskId)
+        organizer.removeDesk(wct, desk.deskRoot.deskId, userId = PRIMARY_USER_ID)
 
         assertThat(
                 wct.hierarchyOps.any { hop ->
@@ -223,10 +245,10 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testRemoveDesk_removesMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending(userId = PRIMARY_USER_ID)
 
         val wct = WindowContainerTransaction()
-        organizer.removeDesk(wct, desk.deskRoot.deskId)
+        organizer.removeDesk(wct, desk.deskRoot.deskId, userId = PRIMARY_USER_ID)
 
         assertThat(
                 wct.hierarchyOps.any { hop ->
@@ -238,8 +260,27 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     }
 
     @Test
+    fun testRemoveDesk_rootUsedByOtherUser_keepsDeskRoot() = runTest {
+        val primaryUserDesk = createDeskSuspending(userId = PRIMARY_USER_ID)
+        val secondaryUserDesk = createDeskSuspending(userId = SECONDARY_USER_ID)
+        assertThat(primaryUserDesk).isEqualTo(secondaryUserDesk)
+
+        val wct = WindowContainerTransaction()
+        organizer.removeDesk(wct, primaryUserDesk.deskRoot.deskId, userId = PRIMARY_USER_ID)
+
+        assertThat(
+                wct.hierarchyOps.any { hop ->
+                    hop.type == HierarchyOp.HIERARCHY_OP_TYPE_REMOVE_ROOT_TASK &&
+                        hop.container == primaryUserDesk.deskRoot.token.asBinder()
+                }
+            )
+            .isFalse()
+        assertThat(primaryUserDesk.deskRoot.users).containsExactly(SECONDARY_USER_ID)
+    }
+
+    @Test
     fun testActivateDesk() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         val wct = WindowContainerTransaction()
         organizer.activateDesk(wct, desk.deskRoot.deskId)
@@ -271,7 +312,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testMoveTaskToDesk() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         val desktopTask = createFreeformTask().apply { parentTaskId = -1 }
         val wct = WindowContainerTransaction()
@@ -308,7 +349,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testGetDeskAtEnd() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val endDesk =
@@ -321,7 +362,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testGetDeskAtEnd_inMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         val task = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
         val endDesk =
@@ -334,7 +375,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun testIsDeskActiveAtEnd() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         val isActive =
             organizer.isDeskActiveAtEnd(
@@ -352,7 +393,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     @Test
     fun deactivateDesk_clearsLaunchRoot() = runTest {
         val wct = WindowContainerTransaction()
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         organizer.activateDesk(wct, desk.deskRoot.deskId)
 
         organizer.deactivateDesk(wct, desk.deskRoot.deskId)
@@ -370,7 +411,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun isDeskChange_forDeskId() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         assertThat(
                 organizer.isDeskChange(
@@ -385,7 +426,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun isDeskChange_forDeskId_inMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         assertThat(
                 organizer.isDeskChange(
@@ -403,7 +444,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun isDeskChange_anyDesk() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         assertThat(
                 organizer.isDeskChange(
@@ -417,7 +458,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun isDeskChange_anyDesk_inMinimizationRoot() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
 
         assertThat(
                 organizer.isDeskChange(
@@ -434,7 +475,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun minimizeTask() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
         organizer.moveTaskToDesk(wct, desk.deskRoot.deskId, task)
@@ -447,7 +488,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun minimizeTask_alreadyMinimized_noOp() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = desk.minimizationRoot.rootId }
         val wct = WindowContainerTransaction()
         organizer.onTaskAppeared(task, SurfaceControl())
@@ -459,8 +500,8 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun minimizeTask_inDifferentDesk_noOp() = runTest {
-        val desk = createDesk()
-        val otherDesk = createDesk()
+        val desk = createDeskSuspending()
+        val otherDesk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = otherDesk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
         organizer.onTaskAppeared(task, SurfaceControl())
@@ -472,7 +513,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun unminimizeTask() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
         organizer.moveTaskToDesk(wct, desk.deskRoot.deskId, task)
@@ -489,7 +530,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun unminimizeTask_alreadyUnminimized_noOp() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
         organizer.moveTaskToDesk(wct, desk.deskRoot.deskId, task)
@@ -503,7 +544,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun unminimizeTask_notInDesk_noOp() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask()
         val wct = WindowContainerTransaction()
 
@@ -514,7 +555,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun reorderTaskToFront() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
         organizer.onTaskAppeared(task, SurfaceControl())
@@ -534,7 +575,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun reorderTaskToFront_notInDesk_noOp() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask()
         val wct = WindowContainerTransaction()
 
@@ -553,7 +594,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun reorderTaskToFront_minimized_unminimizesAndReorders() = runTest {
-        val desk = createDesk()
+        val desk = createDeskSuspending()
         val task = createFreeformTask().apply { parentTaskId = desk.deskRoot.deskId }
         val wct = WindowContainerTransaction()
         organizer.onTaskAppeared(task, SurfaceControl())
@@ -578,7 +619,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskAppeared_visibleDesk_onlyDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
-        createDesk(visible = true)
+        createDeskSuspending(visible = true)
 
         assertThat(launchAdjacentController.launchAdjacentEnabled).isFalse()
     }
@@ -587,7 +628,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskAppeared_invisibleDesk_onlyDesk_enablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = false
 
-        createDesk(visible = false)
+        createDeskSuspending(visible = false)
 
         assertThat(launchAdjacentController.launchAdjacentEnabled).isTrue()
     }
@@ -596,8 +637,8 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskAppeared_invisibleDesk_otherVisibleDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
-        createDesk(visible = true)
-        createDesk(visible = false)
+        createDeskSuspending(visible = true)
+        createDeskSuspending(visible = false)
 
         assertThat(launchAdjacentController.launchAdjacentEnabled).isFalse()
     }
@@ -606,7 +647,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskInfoChanged_deskBecomesVisible_onlyDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
-        val desk = createDesk(visible = false)
+        val desk = createDeskSuspending(visible = false)
         desk.deskRoot.taskInfo.isVisible = true
         organizer.onTaskInfoChanged(desk.deskRoot.taskInfo)
 
@@ -617,7 +658,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskInfoChanged_deskBecomesInvisible_onlyDesk_enablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = false
 
-        val desk = createDesk(visible = true)
+        val desk = createDeskSuspending(visible = true)
         desk.deskRoot.taskInfo.isVisible = false
         organizer.onTaskInfoChanged(desk.deskRoot.taskInfo)
 
@@ -628,8 +669,8 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskInfoChanged_deskBecomesInvisible_otherVisibleDesk_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
-        createDesk(visible = true)
-        val desk = createDesk(visible = true)
+        createDeskSuspending(visible = true)
+        val desk = createDeskSuspending(visible = true)
         desk.deskRoot.taskInfo.isVisible = false
         organizer.onTaskInfoChanged(desk.deskRoot.taskInfo)
 
@@ -640,7 +681,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskVanished_visibleDeskDisappears_onlyDesk_enablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = false
 
-        val desk = createDesk(visible = true)
+        val desk = createDeskSuspending(visible = true)
         organizer.onTaskVanished(desk.deskRoot.taskInfo)
 
         assertThat(launchAdjacentController.launchAdjacentEnabled).isTrue()
@@ -650,8 +691,8 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
     fun onTaskVanished_visibleDeskDisappears_otherDeskVisible_disablesLaunchAdjacent() = runTest {
         launchAdjacentController.launchAdjacentEnabled = true
 
-        createDesk(visible = true)
-        val desk = createDesk(visible = true)
+        createDeskSuspending(visible = true)
+        val desk = createDeskSuspending(visible = true)
         organizer.onTaskVanished(desk.deskRoot.taskInfo)
 
         assertThat(launchAdjacentController.launchAdjacentEnabled).isFalse()
@@ -659,7 +700,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun onTaskInfoChanged_taskNotRoot_invokesListener() = runTest {
-        createDesk()
+        createDeskSuspending()
         val task = createFreeformTask().apply { taskId = TEST_CHILD_TASK_ID }
 
         organizer.onTaskInfoChanged(task)
@@ -669,7 +710,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun onTaskInfoChanged_isDeskRoot_doesNotInvokeListener() = runTest {
-        val deskRoot = createDesk().deskRoot
+        val deskRoot = createDeskSuspending().deskRoot
 
         organizer.onTaskInfoChanged(deskRoot.taskInfo)
 
@@ -678,7 +719,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
 
     @Test
     fun onTaskInfoChanged_isMinimizationRoot_doesNotInvokeListener() = runTest {
-        val minimizationRoot = createDesk().minimizationRoot
+        val minimizationRoot = createDeskSuspending().minimizationRoot
 
         organizer.onTaskInfoChanged(minimizationRoot.taskInfo)
 
@@ -690,7 +731,10 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
         val minimizationRoot: DeskMinimizationRoot,
     )
 
-    private suspend fun createDesk(visible: Boolean = true): DeskRoots {
+    private suspend fun createDeskSuspending(
+        visible: Boolean = true,
+        userId: Int = PRIMARY_USER_ID,
+    ): DeskRoots {
         val freeformRootTask =
             createFreeformTask().apply {
                 parentTaskId = -1
@@ -701,7 +745,7 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
         Mockito.reset(mockShellTaskOrganizer)
         whenever(
                 mockShellTaskOrganizer.createRootTask(
-                    Display.DEFAULT_DISPLAY,
+                    DEFAULT_DISPLAY,
                     WINDOWING_MODE_FREEFORM,
                     organizer,
                     true,
@@ -715,13 +759,9 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
                 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)
+        val deskId = organizer.createDeskSuspending(DEFAULT_DISPLAY, userId)
+        val deskRoot = assertNotNull(organizer.deskRootsByDeskId.get(deskId))
+        val minimizationRoot = assertNotNull(organizer.deskMinimizationRootsByDeskId[deskId])
         return DeskRoots(deskRoot, minimizationRoot)
     }
 
@@ -746,7 +786,14 @@ class RootTaskDesksOrganizerTest : ShellTestCase() {
                 hop.toTop
         }
 
+    private suspend fun DesksOrganizer.createDeskSuspending(displayId: Int, userId: Int): Int =
+        suspendCoroutine { cont ->
+            createDesk(displayId, userId) { deskId -> cont.resumeWith(Result.success(deskId)) }
+        }
+
     companion object {
+        private const val PRIMARY_USER_ID = 10
+        private const val SECONDARY_USER_ID = 11
         private const val TEST_CHILD_TASK_ID = 100
     }
 }