Merge "Desks: Allow obtaining non-existent desk instance in single-desk setups" into main

3 files changed, 18 insertions, 1 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 760d2124b845..e0bd00b40923 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
@@ -72,6 +72,11 @@ class DesktopDisplayEventHandler(
             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)
     }
 
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 6636770895fa..4ff1a5f1be31 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
@@ -926,7 +926,10 @@ class DesktopRepository(
         }
 
         override fun getDesk(deskId: Int): Desk =
-            checkNotNull(deskByDisplayId[deskId]) { "Expected desk $deskId to exist" }
+            // TODO: b/362720497 - consider enforcing that the desk has been created before trying
+            //  to use it. As of now, there are cases where a task may be created faster than a
+            //  desk is, so just create it here if needed. See b/391984373.
+            deskByDisplayId.getOrCreate(deskId)
 
         override fun getActiveDesk(displayId: Int): Desk {
             // TODO: 389787966 - consider migrating to an "active" state instead of checking the
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 5736793cd6a9..f5c93ee8ffe4 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
@@ -17,6 +17,7 @@
 package com.android.wm.shell.desktopmode
 
 import android.graphics.Rect
+import android.platform.test.annotations.DisableFlags
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.FlagsParameterization
 import android.platform.test.flag.junit.SetFlagsRule
@@ -1159,6 +1160,14 @@ class DesktopRepositoryTest(flags: FlagsParameterization) : ShellTestCase() {
         assertThat(repo.shouldDesktopBeActiveForPip(DEFAULT_DESKTOP_ID)).isFalse()
     }
 
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_MULTIPLE_DESKTOPS_BACKEND)
+    fun addTask_deskDoesNotExists_createsDesk() {
+        repo.addTask(displayId = 999, taskId = 6, isVisible = true)
+
+        assertThat(repo.getActiveTaskIdsInDesk(999)).contains(6)
+    }
+
     class TestListener : DesktopRepository.ActiveTasksListener {
         var activeChangesOnDefaultDisplay = 0
         var activeChangesOnSecondaryDisplay = 0