Ensure Desktop Tasks show up in Taskbar multi instance menu

This change adds in desktop tasks to show up in the Manage Windows
Taskbar menu. Previously these were getting filtered out.

Fix: 380167542
Bug: 315989246
Test: Launch multiple tasks in desktop mode and see that they show up in
the multi instance viewer from the Taskbar menu
Flag: com.android.launcher3.enable_multi_instance_menu_taskbar

Change-Id: Ia07b3ad68f43c0f121b0608c5341d50970c8af4c
diff --git a/quickstep/src/com/android/launcher3/taskbar/ManageWindowsTaskbarShortcut.kt b/quickstep/src/com/android/launcher3/taskbar/ManageWindowsTaskbarShortcut.kt
index 8d1f4f5..a597ed6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/ManageWindowsTaskbarShortcut.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/ManageWindowsTaskbarShortcut.kt
@@ -32,7 +32,8 @@
 import com.android.launcher3.views.ActivityContext
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.SystemUiProxy
-import com.android.quickstep.util.GroupTask
+import com.android.quickstep.util.DesktopTask
+import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.android.wm.shell.shared.multiinstance.ManageWindowsViewContainer
 import java.util.Collections
@@ -60,20 +61,29 @@
     private val recentsModel = RecentsModel.INSTANCE[controllers.taskbarActivityContext]
 
     override fun onClick(v: View?) {
-        val filter =
-            Predicate<GroupTask> { task: GroupTask? ->
-                task != null && task.task1.key.packageName == itemInfo?.getTargetPackage()
-            }
-        recentsModel.getTasks(
-            { tasks: List<GroupTask> ->
-                // Since fetching thumbnails is asynchronous, use this set to gate until the tasks
-                // are ready to display
-                val pendingTaskIds =
-                    Collections.synchronizedSet(tasks.map { it.task1.key.id }.toMutableSet())
-                createAndShowTaskShortcutView(tasks, pendingTaskIds)
-            },
-            filter,
-        )
+        val targetPackage = itemInfo?.getTargetPackage()
+        val targetUserId = itemInfo?.user?.identifier
+        val isTargetPackageTask: (Task) -> Boolean = { task ->
+            task.key?.packageName == targetPackage && task.key.userId == targetUserId
+        }
+
+        recentsModel.getTasks { tasks ->
+            val desktopTask = tasks.filterIsInstance<DesktopTask>().firstOrNull()
+            val packageDesktopTasks =
+                (desktopTask?.tasks ?: emptyList()).filter(isTargetPackageTask)
+            val nonDesktopPackageTasks =
+                tasks.filter { isTargetPackageTask(it.task1) }.map { it.task1 }
+
+            // Add tasks from the fetched tasks, deduplicating by task ID
+            val packageTasks =
+                (packageDesktopTasks + nonDesktopPackageTasks).distinctBy { it.key.id }
+
+            // Since fetching thumbnails is asynchronous, use `awaitedTaskIds` to gate until the
+            // tasks are ready to display
+            val awaitedTaskIds = packageTasks.map { it.key.id }.toMutableSet()
+
+            createAndShowTaskShortcutView(packageTasks, awaitedTaskIds)
+        }
     }
 
     /**
@@ -83,25 +93,20 @@
      * thumbnails are processed, it creates a [TaskbarShortcutManageWindowsView] with the collected
      * thumbnails and positions it appropriately.
      */
-    private fun createAndShowTaskShortcutView(
-        tasks: List<GroupTask?>,
-        pendingTaskIds: MutableSet<Int>,
-    ) {
+    private fun createAndShowTaskShortcutView(tasks: List<Task>, pendingTaskIds: MutableSet<Int>) {
         val taskList = arrayListOf<Pair<Int, Bitmap?>>()
-        tasks.forEach { groupTask ->
-            groupTask?.task1?.let { task ->
-                recentsModel.thumbnailCache.getThumbnailInBackground(task) {
-                    thumbnailData: ThumbnailData ->
-                    pendingTaskIds.remove(task.key.id)
-                    // Add the current pair of task id and ThumbnailData to the list of all tasks
-                    if (thumbnailData.thumbnail != null) {
-                        taskList.add(task.key.id to thumbnailData.thumbnail)
-                    }
 
-                    // If the set is empty, all thumbnails have been fetched
-                    if (pendingTaskIds.isEmpty() && taskList.isNotEmpty()) {
-                        createAndPositionTaskbarShortcut(taskList)
-                    }
+        tasks.forEach { task ->
+            recentsModel.thumbnailCache.getThumbnailInBackground(task) {
+                thumbnailData: ThumbnailData ->
+                pendingTaskIds.remove(task.key.id)
+                // Add the current pair of task id and ThumbnailData to the list of all tasks
+                if (thumbnailData.thumbnail != null) {
+                    taskList.add(task.key.id to thumbnailData.thumbnail)
+                }
+                // If the set is empty, all thumbnails have been fetched
+                if (pendingTaskIds.isEmpty() && taskList.isNotEmpty()) {
+                    createAndPositionTaskbarShortcut(taskList)
                 }
             }
         }