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)
                 }
             }
         }