6a/ Group visible tasks into a single mixed grouped task

- When enable_shell_top_task_tracking is enabled, we will group visible
  tasks together into a single grouped task that represents the task
  in Overview and in TopTaskTracker.  This allows Launcher to prepare
  overview and gesture nav handling in cases where there are multiple
  stacked tasks (ie. fullscreen/split/desktop + fullscreen)

Bug: 346588978
Flag: com.android.wm.shell.enable_shell_top_task_tracking
Test: atest WMShellUnitTests
Change-Id: I07be12db791b82458d401006c204a11a4384b8aa

5 files changed, 373 insertions, 129 deletions

diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index 76705dcdd3d2..6936ddc5eeac 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -376,6 +376,14 @@ public class TaskInfo {
     }
 
     /**
+     * Returns the task id.
+     * @hide
+     */
+    public int getTaskId() {
+        return taskId;
+    }
+
+    /**
      * Whether this task is visible.
      */
     public boolean isVisible() {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
index 0182588398a4..2d4d458292ea 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/RecentTasksController.java
@@ -43,7 +43,6 @@ import android.graphics.Point;
 import android.os.Bundle;
 import android.os.RemoteException;
 import android.util.Slog;
-import android.util.SparseArray;
 import android.util.SparseIntArray;
 import android.window.DesktopModeFlags;
 import android.window.WindowContainerToken;
@@ -83,6 +82,7 @@ import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * Manages the recent task list from the system, caching it as necessary.
@@ -124,6 +124,10 @@ public class RecentTasksController implements TaskStackListenerCallback,
      * Cached list of the visible tasks, sorted from top most to bottom most.
      */
     private final List<RunningTaskInfo> mVisibleTasks = new ArrayList<>();
+    private final Map<Integer, TaskInfo> mVisibleTasksMap = new HashMap<>();
+
+    // Temporary vars used in `generateList()`
+    private final Map<Integer, TaskInfo> mTmpRemaining = new HashMap<>();
 
     /**
      * Creates {@link RecentTasksController}, returns {@code null} if the feature is not
@@ -348,8 +352,11 @@ public class RecentTasksController implements TaskStackListenerCallback,
     public void onVisibleTasksChanged(@NonNull List<? extends RunningTaskInfo> visibleTasks) {
         mVisibleTasks.clear();
         mVisibleTasks.addAll(visibleTasks);
+        mVisibleTasksMap.clear();
+        mVisibleTasksMap.putAll(mVisibleTasks.stream().collect(
+                Collectors.toMap(TaskInfo::getTaskId, task -> task)));
         // Notify with all the info and not just the running task info
-        notifyVisibleTasksChanged(visibleTasks);
+        notifyVisibleTasksChanged(mVisibleTasks);
     }
 
     @VisibleForTesting
@@ -458,7 +465,7 @@ public class RecentTasksController implements TaskStackListenerCallback,
         }
         try {
             // Compute the visible recent tasks in order, and move the task to the top
-            mListener.onVisibleTasksChanged(generateList(visibleTasks)
+            mListener.onVisibleTasksChanged(generateList(visibleTasks, "visibleTasksChanged")
                     .toArray(new GroupedTaskInfo[0]));
         } catch (RemoteException e) {
             Slog.w(TAG, "Failed call onVisibleTasksChanged", e);
@@ -494,40 +501,87 @@ public class RecentTasksController implements TaskStackListenerCallback,
     @VisibleForTesting
     ArrayList<GroupedTaskInfo> getRecentTasks(int maxNum, int flags, int userId) {
         // Note: the returned task list is ordered from the most-recent to least-recent order
-        return generateList(mActivityTaskManager.getRecentTasks(maxNum, flags, userId));
+        return generateList(mActivityTaskManager.getRecentTasks(maxNum, flags, userId),
+                "getRecentTasks");
     }
 
     /**
-     * Generates a list of GroupedTaskInfos for the given list of tasks.
+     * Returns whether the given task should be excluded from the generated list.
      */
-    private <T extends TaskInfo> ArrayList<GroupedTaskInfo> generateList(@NonNull List<T> tasks) {
-        // Make a mapping of task id -> task info
-        final SparseArray<TaskInfo> rawMapping = new SparseArray<>();
-        for (int i = 0; i < tasks.size(); i++) {
-            final TaskInfo taskInfo = tasks.get(i);
-            rawMapping.put(taskInfo.taskId, taskInfo);
+    private boolean excludeTaskFromGeneratedList(TaskInfo taskInfo) {
+        if (taskInfo.getWindowingMode() == WINDOWING_MODE_PINNED) {
+            // We don't current send pinned tasks as a part of recent or running tasks
+            return true;
+        }
+        if (isWallpaperTask(taskInfo)) {
+            // Don't add the fullscreen wallpaper task as an entry in grouped tasks
+            return true;
         }
+        return false;
+    }
 
-        ArrayList<TaskInfo> freeformTasks = new ArrayList<>();
-        Set<Integer> minimizedFreeformTasks = new HashSet<>();
+    /**
+     * Generates a list of GroupedTaskInfos for the given raw list of tasks (either recents or
+     * running tasks).
+     *
+     * The general flow is:
+     *  - Collect the desktop tasks
+     *  - Collect the visible tasks (in order), including the desktop tasks if visible
+     *  - Construct the final list with the visible tasks, followed by the subsequent tasks
+     *      - if enableShellTopTaskTracking() is enabled, the visible tasks will be grouped into
+     *        a single mixed task
+     *      - if the desktop tasks are not visible, they will be appended to the end of the list
+     *
+     * TODO(346588978): Generate list in per-display order
+     *
+     * @param tasks The list of tasks ordered from most recent to least recent
+     */
+    @VisibleForTesting
+    <T extends TaskInfo> ArrayList<GroupedTaskInfo> generateList(@NonNull List<T> tasks,
+            String reason) {
+        if (tasks.isEmpty()) {
+            return new ArrayList<>();
+        }
+
+        if (enableShellTopTaskTracking()) {
+            ProtoLog.v(WM_SHELL_TASK_OBSERVER, "RecentTasksController.generateList(%s)", reason);
+        }
 
-        int mostRecentFreeformTaskIndex = Integer.MAX_VALUE;
+        // Make a mapping of task id -> task info for the remaining tasks to be processed, this
+        // mapping is used to keep track of split tasks that may exist later in the task list that
+        // should be ignored because they've already been grouped
+        mTmpRemaining.clear();
+        mTmpRemaining.putAll(tasks.stream().collect(
+                Collectors.toMap(TaskInfo::getTaskId, task -> task)));
 
-        ArrayList<GroupedTaskInfo> groupedTasks = new ArrayList<>();
-        // Pull out the pairs as we iterate back in the list
+        // The final grouped tasks
+        ArrayList<GroupedTaskInfo> groupedTasks = new ArrayList<>(tasks.size());
+        ArrayList<GroupedTaskInfo> visibleGroupedTasks = new ArrayList<>();
+
+        // Phase 1: Extract the desktop and visible fullscreen/split tasks. We make new collections
+        // here as the GroupedTaskInfo can store them without copying
+        ArrayList<TaskInfo> desktopTasks = new ArrayList<>();
+        Set<Integer> minimizedDesktopTasks = new HashSet<>();
+        boolean desktopTasksVisible = false;
         for (int i = 0; i < tasks.size(); i++) {
             final TaskInfo taskInfo = tasks.get(i);
-            if (!rawMapping.contains(taskInfo.taskId)) {
-                // If it's not in the mapping, then it was already paired with another task
+            final int taskId = taskInfo.taskId;
+
+            if (!mTmpRemaining.containsKey(taskInfo.taskId)) {
+                // Skip if we've already processed it
                 continue;
             }
+
+            if (excludeTaskFromGeneratedList(taskInfo)) {
+                // Skip and update the list if we are excluding this task
+                mTmpRemaining.remove(taskId);
+                continue;
+            }
+
+            // Desktop tasks
             if (DesktopModeStatus.canEnterDesktopMode(mContext) &&
-                mDesktopUserRepositories.isPresent()
-                    && mDesktopUserRepositories.get().getCurrent().isActiveTask(taskInfo.taskId)) {
-                // Freeform tasks will be added as a separate entry
-                if (mostRecentFreeformTaskIndex == Integer.MAX_VALUE) {
-                    mostRecentFreeformTaskIndex = groupedTasks.size();
-                }
+                    mDesktopUserRepositories.isPresent()
+                    && mDesktopUserRepositories.get().getCurrent().isActiveTask(taskId)) {
                 // If task has their app bounds set to null which happens after reboot, set the
                 // app bounds to persisted lastFullscreenBounds. Also set the position in parent
                 // to the top left of the bounds.
@@ -538,49 +592,132 @@ public class RecentTasksController implements TaskStackListenerCallback,
                     taskInfo.positionInParent = new Point(taskInfo.lastNonFullscreenBounds.left,
                             taskInfo.lastNonFullscreenBounds.top);
                 }
-                freeformTasks.add(taskInfo);
-                if (mDesktopUserRepositories.get().getCurrent().isMinimizedTask(taskInfo.taskId)) {
-                    minimizedFreeformTasks.add(taskInfo.taskId);
+                desktopTasks.add(taskInfo);
+                if (mDesktopUserRepositories.get().getCurrent().isMinimizedTask(taskId)) {
+                    minimizedDesktopTasks.add(taskId);
                 }
+                desktopTasksVisible |= mVisibleTasksMap.containsKey(taskId);
+                mTmpRemaining.remove(taskId);
                 continue;
             }
 
-            final int pairedTaskId = mSplitTasks.get(taskInfo.taskId, INVALID_TASK_ID);
-            if (pairedTaskId != INVALID_TASK_ID && rawMapping.contains(pairedTaskId)) {
-                final TaskInfo pairedTaskInfo = rawMapping.get(pairedTaskId);
-                rawMapping.remove(pairedTaskId);
-                groupedTasks.add(GroupedTaskInfo.forSplitTasks(taskInfo, pairedTaskInfo,
-                        mTaskSplitBoundsMap.get(pairedTaskId)));
+            if (enableShellTopTaskTracking()) {
+                // Visible tasks
+                if (mVisibleTasksMap.containsKey(taskId)) {
+                    // Split tasks
+                    if (extractAndAddSplitGroupedTask(taskInfo, mTmpRemaining,
+                            visibleGroupedTasks)) {
+                        continue;
+                    }
+
+                    // Fullscreen tasks
+                    visibleGroupedTasks.add(GroupedTaskInfo.forFullscreenTasks(taskInfo));
+                    mTmpRemaining.remove(taskId);
+                }
             } else {
-                if (isWallpaperTask(taskInfo)) {
-                    // Don't add the wallpaper task as an entry in grouped tasks
+                // Split tasks
+                if (extractAndAddSplitGroupedTask(taskInfo, mTmpRemaining, groupedTasks)) {
                     continue;
                 }
-                // TODO(346588978): Consolidate multiple visible fullscreen tasks into the same
-                //  grouped task
+
+                // Fullscreen tasks
                 groupedTasks.add(GroupedTaskInfo.forFullscreenTasks(taskInfo));
             }
         }
 
-        // Add a special entry for freeform tasks
-        if (!freeformTasks.isEmpty()) {
-            groupedTasks.add(mostRecentFreeformTaskIndex,
-                    GroupedTaskInfo.forFreeformTasks(
-                            freeformTasks,
-                            minimizedFreeformTasks));
-        }
-
         if (enableShellTopTaskTracking()) {
-            // We don't current send pinned tasks as a part of recent or running tasks, so remove
-            // them from the list here
-            groupedTasks.removeIf(
-                    gti -> gti.getTaskInfo1().getWindowingMode() == WINDOWING_MODE_PINNED);
+            // Phase 2: If there were desktop tasks and they are visible, add them to the visible
+            //          list as well (the actual order doesn't matter for Overview)
+            if (!desktopTasks.isEmpty() && desktopTasksVisible) {
+                visibleGroupedTasks.add(
+                        GroupedTaskInfo.forFreeformTasks(desktopTasks, minimizedDesktopTasks));
+            }
+
+            if (!visibleGroupedTasks.isEmpty()) {
+                // Phase 3: Combine the visible tasks into a single mixed grouped task, only if
+                //          there are > 1 tasks to group, and add them to the final list
+                if (visibleGroupedTasks.size() > 1) {
+                    groupedTasks.add(GroupedTaskInfo.forMixed(visibleGroupedTasks));
+                } else {
+                    groupedTasks.addAll(visibleGroupedTasks);
+                }
+            }
+            dumpGroupedTasks(groupedTasks, "Phase 3");
+
+            // Phase 4: For the remaining non-visible split and fullscreen tasks, add grouped tasks
+            //          in order to the final list
+            for (int i = 0; i < tasks.size(); i++) {
+                final TaskInfo taskInfo = tasks.get(i);
+                if (!mTmpRemaining.containsKey(taskInfo.taskId)) {
+                    // Skip if we've already processed it
+                    continue;
+                }
+
+                // Split tasks
+                if (extractAndAddSplitGroupedTask(taskInfo, mTmpRemaining, groupedTasks)) {
+                    continue;
+                }
+
+                // Fullscreen tasks
+                groupedTasks.add(GroupedTaskInfo.forFullscreenTasks(taskInfo));
+            }
+            dumpGroupedTasks(groupedTasks, "Phase 4");
+
+            // Phase 5: If there were desktop tasks and they are not visible (ie. weren't added
+            //          above), add them to the end of the final list (the actual order doesn't
+            //          matter for Overview)
+            if (!desktopTasks.isEmpty() && !desktopTasksVisible) {
+                groupedTasks.add(
+                        GroupedTaskInfo.forFreeformTasks(desktopTasks, minimizedDesktopTasks));
+            }
+            dumpGroupedTasks(groupedTasks, "Phase 5");
+        } else {
+            // Add the desktop tasks at the end of the list
+            if (!desktopTasks.isEmpty()) {
+                groupedTasks.add(
+                        GroupedTaskInfo.forFreeformTasks(desktopTasks, minimizedDesktopTasks));
+            }
         }
 
         return groupedTasks;
     }
 
     /**
+     * Only to be called from `generateList()`. If the given {@param taskInfo} has a paired task,
+     * then a split grouped task with the pair is added to {@param tasksOut}.
+     *
+     * @return whether a split task was extracted and added to the given list
+     */
+    private boolean extractAndAddSplitGroupedTask(@NonNull TaskInfo taskInfo,
+            @NonNull Map<Integer, TaskInfo> remainingTasks,
+            @NonNull ArrayList<GroupedTaskInfo> tasksOut) {
+        final int pairedTaskId = mSplitTasks.get(taskInfo.taskId, INVALID_TASK_ID);
+        if (pairedTaskId == INVALID_TASK_ID || !remainingTasks.containsKey(pairedTaskId)) {
+            return false;
+        }
+
+        // Add both this task and its pair to the list, and mark the paired task to be
+        // skipped when it is encountered in the list
+        final TaskInfo pairedTaskInfo = remainingTasks.get(pairedTaskId);
+        remainingTasks.remove(taskInfo.taskId);
+        remainingTasks.remove(pairedTaskId);
+        tasksOut.add(GroupedTaskInfo.forSplitTasks(taskInfo, pairedTaskInfo,
+                mTaskSplitBoundsMap.get(pairedTaskId)));
+        return true;
+    }
+
+    /** Dumps the set of tasks to protolog */
+    private void dumpGroupedTasks(List<GroupedTaskInfo> groupedTasks, String reason) {
+        if (!WM_SHELL_TASK_OBSERVER.isEnabled()) {
+            return;
+        }
+        ProtoLog.v(WM_SHELL_TASK_OBSERVER, "    Tasks (%s):", reason);
+        for (GroupedTaskInfo task : groupedTasks) {
+            ProtoLog.v(WM_SHELL_TASK_OBSERVER, "        %s", task);
+        }
+    }
+
+    /**
      * Returns the top running leaf task ignoring {@param ignoreTaskToken} if it is specified.
      * NOTE: This path currently makes assumptions that ignoreTaskToken is for the top task.
      */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt
index 93f2e4cf0e45..11cd4031e8ef 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/recents/TaskStackTransitionObserver.kt
@@ -193,16 +193,18 @@ class TaskStackTransitionObserver(
     override fun onTransitionMerged(merged: IBinder, playing: IBinder) {}
 
     override fun onTransitionFinished(transition: IBinder, aborted: Boolean) {
-        if (enableShellTopTaskTracking()) {
-            if (pendingCloseTasks.isNotEmpty()) {
-                // Update the visible task list based on the pending close tasks
-                for (change in pendingCloseTasks) {
-                    visibleTasks.removeIf {
-                        it.taskId == change.taskId
-                    }
+        if (!enableShellTopTaskTracking()) {
+            return
+        }
+
+        if (pendingCloseTasks.isNotEmpty()) {
+            // Update the visible task list based on the pending close tasks
+            for (change in pendingCloseTasks) {
+                visibleTasks.removeIf {
+                    it.taskId == change.taskId
                 }
-                updateVisibleTasksList("transition-finished")
             }
+            updateVisibleTasksList("transition-finished")
         }
     }
 
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
index 28f4ea0c7ada..065fa219e8d0 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/recents/RecentTasksControllerTest.java
@@ -20,6 +20,7 @@ import static android.app.ActivityManager.RECENT_IGNORE_UNAVAILABLE;
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN;
 import static android.app.WindowConfiguration.WINDOWING_MODE_MULTI_WINDOW;
+import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession;
 import static com.android.window.flags.Flags.FLAG_ENABLE_DESKTOP_WINDOWING_PERSISTENCE;
@@ -28,7 +29,6 @@ import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_FULLSCREEN;
 import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
 import static com.android.wm.shell.shared.split.SplitScreenConstants.SNAP_TO_2_50_50;
 
-import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
@@ -48,9 +48,10 @@ import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import static java.lang.Integer.MAX_VALUE;
+import static java.util.stream.Collectors.joining;
 
-import android.app.ActivityManager;
 import android.app.ActivityManager.RecentTaskInfo;
+import android.app.ActivityManager.RunningTaskInfo;
 import android.app.ActivityTaskManager;
 import android.app.KeyguardManager;
 import android.content.ComponentName;
@@ -247,10 +248,10 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         ArrayList<GroupedTaskInfo> recentTasks =
                 mRecentTasksController.getRecentTasks(MAX_VALUE, RECENT_IGNORE_UNAVAILABLE, 0);
-        assertGroupedTasksListEquals(recentTasks,
-                t1.taskId, -1,
-                t2.taskId, -1,
-                t3.taskId, -1);
+        assertGroupedTasksListEquals(recentTasks, List.of(
+                List.of(t1.taskId),
+                List.of(t2.taskId),
+                List.of(t3.taskId)));
     }
 
     @Test
@@ -262,7 +263,9 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         ArrayList<GroupedTaskInfo> recentTasks =
                 mRecentTasksController.getRecentTasks(MAX_VALUE, RECENT_IGNORE_UNAVAILABLE, 0);
-        assertGroupedTasksListEquals(recentTasks, t1.taskId, -1, t3.taskId, -1);
+        assertGroupedTasksListEquals(recentTasks, List.of(
+                List.of(t1.taskId),
+                List.of(t3.taskId)));
     }
 
     @Test
@@ -286,11 +289,11 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         ArrayList<GroupedTaskInfo> recentTasks =
                 mRecentTasksController.getRecentTasks(MAX_VALUE, RECENT_IGNORE_UNAVAILABLE, 0);
-        assertGroupedTasksListEquals(recentTasks,
-                t1.taskId, -1,
-                t2.taskId, t4.taskId,
-                t3.taskId, t5.taskId,
-                t6.taskId, -1);
+        assertGroupedTasksListEquals(recentTasks, List.of(
+                List.of(t1.taskId),
+                List.of(t2.taskId, t4.taskId),
+                List.of(t3.taskId, t5.taskId),
+                List.of(t6.taskId)));
     }
 
     @Test
@@ -320,11 +323,11 @@ public class RecentTasksControllerTest extends ShellTestCase {
                         consumer);
         mMainExecutor.flushAll();
 
-        assertGroupedTasksListEquals(recentTasks[0],
-                t1.taskId, -1,
-                t2.taskId, t4.taskId,
-                t3.taskId, t5.taskId,
-                t6.taskId, -1);
+        assertGroupedTasksListEquals(recentTasks[0], List.of(
+                List.of(t1.taskId),
+                List.of(t2.taskId, t4.taskId),
+                List.of(t3.taskId, t5.taskId),
+                List.of(t6.taskId)));
     }
 
     @Test
@@ -343,9 +346,9 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         // 2 freeform tasks should be grouped into one, 3 total recents entries
         assertEquals(3, recentTasks.size());
-        GroupedTaskInfo freeformGroup = recentTasks.get(0);
-        GroupedTaskInfo singleGroup1 = recentTasks.get(1);
-        GroupedTaskInfo singleGroup2 = recentTasks.get(2);
+        GroupedTaskInfo singleGroup1 = recentTasks.get(0);
+        GroupedTaskInfo singleGroup2 = recentTasks.get(1);
+        GroupedTaskInfo freeformGroup = recentTasks.get(2);
 
         // Check that groups have expected types
         assertTrue(freeformGroup.isBaseType(TYPE_FREEFORM));
@@ -383,8 +386,8 @@ public class RecentTasksControllerTest extends ShellTestCase {
         // 2 split screen tasks grouped, 2 freeform tasks grouped, 3 total recents entries
         assertEquals(3, recentTasks.size());
         GroupedTaskInfo splitGroup = recentTasks.get(0);
-        GroupedTaskInfo freeformGroup = recentTasks.get(1);
-        GroupedTaskInfo singleGroup = recentTasks.get(2);
+        GroupedTaskInfo singleGroup = recentTasks.get(1);
+        GroupedTaskInfo freeformGroup = recentTasks.get(2);
 
         // Check that groups have expected types
         assertTrue(splitGroup.isBaseType(TYPE_SPLIT));
@@ -454,9 +457,9 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         // 3 freeform tasks should be grouped into one, 2 single tasks, 3 total recents entries
         assertEquals(3, recentTasks.size());
-        GroupedTaskInfo freeformGroup = recentTasks.get(0);
-        GroupedTaskInfo singleGroup1 = recentTasks.get(1);
-        GroupedTaskInfo singleGroup2 = recentTasks.get(2);
+        GroupedTaskInfo singleGroup1 = recentTasks.get(0);
+        GroupedTaskInfo singleGroup2 = recentTasks.get(1);
+        GroupedTaskInfo freeformGroup = recentTasks.get(2);
 
         // Check that groups have expected types
         assertTrue(freeformGroup.isBaseType(TYPE_FREEFORM));
@@ -523,7 +526,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         // Remove one of the tasks and ensure the pair is removed
         SurfaceControl mockLeash = mock(SurfaceControl.class);
-        ActivityManager.RunningTaskInfo rt2 = makeRunningTaskInfo(2);
+        RunningTaskInfo rt2 = makeRunningTaskInfo(2);
         mShellTaskOrganizer.onTaskAppeared(rt2, mockLeash);
         mShellTaskOrganizer.onTaskVanished(rt2);
 
@@ -537,13 +540,13 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
         // Remove one of the tasks and ensure the pair is removed
         SurfaceControl mockLeash = mock(SurfaceControl.class);
-        ActivityManager.RunningTaskInfo rt2Fullscreen = makeRunningTaskInfo(2);
+        RunningTaskInfo rt2Fullscreen = makeRunningTaskInfo(2);
         rt2Fullscreen.configuration.windowConfiguration.setWindowingMode(
                 WINDOWING_MODE_FULLSCREEN);
         mShellTaskOrganizer.onTaskAppeared(rt2Fullscreen, mockLeash);
 
         // Change the windowing mode and ensure the recent tasks change is notified
-        ActivityManager.RunningTaskInfo rt2MultiWIndow = makeRunningTaskInfo(2);
+        RunningTaskInfo rt2MultiWIndow = makeRunningTaskInfo(2);
         rt2MultiWIndow.configuration.windowConfiguration.setWindowingMode(
                 WINDOWING_MODE_MULTI_WINDOW);
         mShellTaskOrganizer.onTaskInfoChanged(rt2MultiWIndow);
@@ -557,7 +560,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void onTaskAdded_desktopModeRunningAppsEnabled_triggersOnRunningTaskAppeared()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskAdded(taskInfo);
 
@@ -570,7 +573,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void onTaskAdded_desktopModeRunningAppsDisabled_doesNotTriggerOnRunningTaskAppeared()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskAdded(taskInfo);
 
@@ -583,7 +586,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void taskWindowingModeChanged_desktopRunningAppsEnabled_triggersOnRunningTaskChanged()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskRunningInfoChanged(taskInfo);
 
@@ -597,7 +600,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
             taskWindowingModeChanged_desktopRunningAppsDisabled_doesNotTriggerOnRunningTaskChanged()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskRunningInfoChanged(taskInfo);
 
@@ -610,7 +613,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void onTaskRemoved_desktopModeRunningAppsEnabled_triggersOnRunningTaskVanished()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskRemoved(taskInfo);
 
@@ -623,7 +626,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void onTaskRemoved_desktopModeRunningAppsDisabled_doesNotTriggerOnRunningTaskVanished()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskRemoved(taskInfo);
 
@@ -632,10 +635,11 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
     @Test
     @EnableFlags(Flags.FLAG_ENABLE_TASK_STACK_OBSERVER_IN_SHELL)
+    @DisableFlags(com.android.wm.shell.Flags.FLAG_ENABLE_SHELL_TOP_TASK_TRACKING)
     public void onTaskMovedToFront_TaskStackObserverEnabled_triggersOnTaskMovedToFront()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskMovedToFrontThroughTransition(taskInfo);
 
@@ -648,7 +652,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     public void onTaskMovedToFront_TaskStackObserverEnabled_doesNotTriggersOnTaskMovedToFront()
             throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
 
         mRecentTasksControllerReal.onTaskMovedToFront(taskInfo);
 
@@ -700,7 +704,7 @@ public class RecentTasksControllerTest extends ShellTestCase {
     @EnableFlags(com.android.wm.shell.Flags.FLAG_ENABLE_SHELL_TOP_TASK_TRACKING)
     public void shellTopTaskTracker_onTaskRemoved_expectNoRecentsChanged() throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
         mRecentTasksControllerReal.onTaskRemoved(taskInfo);
         verify(mRecentTasksListener, never()).onRecentTasksChanged();
     }
@@ -710,22 +714,105 @@ public class RecentTasksControllerTest extends ShellTestCase {
     @EnableFlags(com.android.wm.shell.Flags.FLAG_ENABLE_SHELL_TOP_TASK_TRACKING)
     public void shellTopTaskTracker_onVisibleTasksChanged() throws Exception {
         mRecentTasksControllerReal.registerRecentTasksListener(mRecentTasksListener);
-        ActivityManager.RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
+        RunningTaskInfo taskInfo = makeRunningTaskInfo(/* taskId= */10);
         mRecentTasksControllerReal.onVisibleTasksChanged(List.of(taskInfo));
         verify(mRecentTasksListener, never()).onVisibleTasksChanged(any());
     }
 
+    @Test
+    public void generateList_emptyTaskList_expectNoGroupedTasks() throws Exception {
+        assertTrue(mRecentTasksControllerReal.generateList(List.of(), "test").isEmpty());
+    }
+
+    @Test
+    public void generateList_excludePipTask() throws Exception {
+        RunningTaskInfo task1 = makeRunningTaskInfo(1);
+        RunningTaskInfo pipTask = makeRunningTaskInfo(2);
+        pipTask.configuration.windowConfiguration.setWindowingMode(
+                WINDOWING_MODE_PINNED);
+
+        ArrayList<GroupedTaskInfo> groupedTasks = mRecentTasksControllerReal.generateList(
+                List.of(task1, pipTask),
+                "test");
+
+        assertGroupedTasksListEquals(groupedTasks, List.of(List.of(task1.taskId)));
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_STACK_OBSERVER_IN_SHELL)
+    @EnableFlags(com.android.wm.shell.Flags.FLAG_ENABLE_SHELL_TOP_TASK_TRACKING)
+    public void generateList_fullscreen_noVisibleTasks_expectNoGrouping() throws Exception {
+        RunningTaskInfo task1 = makeRunningTaskInfo(1);
+        RunningTaskInfo task2 = makeRunningTaskInfo(2);
+        RunningTaskInfo task3 = makeRunningTaskInfo(3);
+        // Reset visible tasks list
+        mRecentTasksControllerReal.onVisibleTasksChanged(List.of());
+
+        // Generate a list with a number of fullscreen tasks
+        ArrayList<GroupedTaskInfo> groupedTasks = mRecentTasksControllerReal.generateList(
+                List.of(task1, task2, task3),
+                "test");
+
+        assertGroupedTasksListEquals(groupedTasks, List.of(
+                List.of(task1.taskId),
+                List.of(task2.taskId),
+                List.of(task3.taskId)));
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_STACK_OBSERVER_IN_SHELL)
+    @EnableFlags(com.android.wm.shell.Flags.FLAG_ENABLE_SHELL_TOP_TASK_TRACKING)
+    public void generateList_fullscreen_singleVisibleTask_expectNoGrouping() throws Exception {
+        RunningTaskInfo task1 = makeRunningTaskInfo(1);
+        RunningTaskInfo task2 = makeRunningTaskInfo(2);
+        RunningTaskInfo task3 = makeRunningTaskInfo(3);
+
+        // Reset visible tasks list
+        mRecentTasksControllerReal.onVisibleTasksChanged(List.of(task1));
+
+        // Generate a list with a number of fullscreen tasks
+        ArrayList<GroupedTaskInfo> groupedTasks = mRecentTasksControllerReal.generateList(
+                List.of(task1, task2, task3),
+                "test");
+
+        assertGroupedTasksListEquals(groupedTasks, List.of(
+                List.of(task1.taskId),
+                List.of(task2.taskId),
+                List.of(task3.taskId)));
+    }
+
+    @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TASK_STACK_OBSERVER_IN_SHELL)
+    @EnableFlags(com.android.wm.shell.Flags.FLAG_ENABLE_SHELL_TOP_TASK_TRACKING)
+    public void generateList_fullscreen_multipleVisibleTasks_expectGrouping() throws Exception {
+        RunningTaskInfo task1 = makeRunningTaskInfo(1);
+        RunningTaskInfo task2 = makeRunningTaskInfo(2);
+        RunningTaskInfo task3 = makeRunningTaskInfo(3);
+
+        // Reset visible tasks list
+        mRecentTasksControllerReal.onVisibleTasksChanged(List.of(task1, task2));
+
+        // Generate a list with a number of fullscreen tasks
+        ArrayList<GroupedTaskInfo> groupedTasks = mRecentTasksControllerReal.generateList(
+                List.of(task1, task2, task3),
+                "test");
+
+        assertGroupedTasksListEquals(groupedTasks, List.of(
+                List.of(task1.taskId),
+                List.of(task2.taskId),
+                List.of(task3.taskId)));
+    }
+
     /**
      * Helper to create a task with a given task id.
      */
     private RecentTaskInfo makeTaskInfo(int taskId) {
         RecentTaskInfo info = new RecentTaskInfo();
         info.taskId = taskId;
-
+        info.realActivity = new ComponentName("testPackage", "testClass");
         Intent intent = new Intent();
         intent.setComponent(new ComponentName("com." + taskId, "Activity" + taskId));
         info.baseIntent = intent;
-
         info.lastNonFullscreenBounds = new Rect();
         return info;
     }
@@ -742,10 +829,14 @@ public class RecentTasksControllerTest extends ShellTestCase {
     /**
      * Helper to create a running task with a given task id.
      */
-    private ActivityManager.RunningTaskInfo makeRunningTaskInfo(int taskId) {
-        ActivityManager.RunningTaskInfo info = new ActivityManager.RunningTaskInfo();
+    private RunningTaskInfo makeRunningTaskInfo(int taskId) {
+        RunningTaskInfo info = new RunningTaskInfo();
         info.taskId = taskId;
         info.realActivity = new ComponentName("testPackage", "testClass");
+        Intent intent = new Intent();
+        intent.setComponent(new ComponentName("com." + taskId, "Activity" + taskId));
+        info.baseIntent = intent;
+        info.lastNonFullscreenBounds = new Rect();
         return info;
     }
 
@@ -759,37 +850,42 @@ public class RecentTasksControllerTest extends ShellTestCase {
 
     /**
      * Asserts that the recent tasks matches the given task ids.
+     * TODO(346588978): Separate out specific split verification during the iteration below
      *
-     * @param expectedTaskIds list of task ids that map to the flattened task ids of the tasks in
-     *                        the grouped task list
+     * @param expectedTaskIds a list of expected grouped task ids (itself a list of ints)
      */
-    private void assertGroupedTasksListEquals(List<GroupedTaskInfo> recentTasks,
-            int... expectedTaskIds) {
-        int[] flattenedTaskIds = new int[recentTasks.size() * 2];
-        for (int i = 0; i < recentTasks.size(); i++) {
-            GroupedTaskInfo pair = recentTasks.get(i);
-            int taskId1 = pair.getTaskInfo1().taskId;
-            flattenedTaskIds[2 * i] = taskId1;
-            flattenedTaskIds[2 * i + 1] = pair.getTaskInfo2() != null
-                    ? pair.getTaskInfo2().taskId
-                    : -1;
-
-            if (pair.getTaskInfo2() != null) {
-                assertNotNull(pair.getSplitBounds());
-                int leftTopTaskId = pair.getSplitBounds().leftTopTaskId;
-                int bottomRightTaskId = pair.getSplitBounds().rightBottomTaskId;
+    private void assertGroupedTasksListEquals(List<GroupedTaskInfo> groupedTasks,
+            List<List<Integer>> expectedTaskIds) {
+        List<List<Integer>> foundTaskIds = new ArrayList<>();
+        for (int i = 0; i < groupedTasks.size(); i++) {
+            GroupedTaskInfo groupedTask = groupedTasks.get(i);
+            List<Integer> groupedTaskIds = groupedTask.getTaskInfoList().stream()
+                    .map(taskInfo -> taskInfo.taskId)
+                    .toList();
+            foundTaskIds.add(groupedTaskIds);
+
+            if (groupedTask.isBaseType(TYPE_SPLIT)) {
+                assertNotNull(groupedTask.getSplitBounds());
+                int leftTopTaskId = groupedTask.getSplitBounds().leftTopTaskId;
+                int bottomRightTaskId = groupedTask.getSplitBounds().rightBottomTaskId;
                 // Unclear if pairs are ordered by split position, most likely not.
-                assertTrue(leftTopTaskId == taskId1
-                        || leftTopTaskId == pair.getTaskInfo2().taskId);
-                assertTrue(bottomRightTaskId == taskId1
-                        || bottomRightTaskId == pair.getTaskInfo2().taskId);
-            } else {
-                assertNull(pair.getSplitBounds());
+                assertTrue(leftTopTaskId == groupedTaskIds.getFirst()
+                        || leftTopTaskId == groupedTaskIds.getLast());
+                assertTrue(bottomRightTaskId == groupedTaskIds.getFirst()
+                        || bottomRightTaskId == groupedTaskIds.getLast());
             }
         }
-        assertArrayEquals("Expected: " + Arrays.toString(expectedTaskIds)
-                        + " Received: " + Arrays.toString(flattenedTaskIds),
-                flattenedTaskIds,
-                expectedTaskIds);
+        List<Integer> flattenedExpectedTaskIds = expectedTaskIds.stream()
+                .flatMap(List::stream)
+                .toList();
+        List<Integer> flattenedFoundTaskIds = foundTaskIds.stream()
+                .flatMap(List::stream)
+                .toList();
+        assertEquals("Expected: "
+                        + flattenedExpectedTaskIds.stream().map(String::valueOf).collect(joining())
+                        + " Received: "
+                        + flattenedFoundTaskIds.stream().map(String::valueOf).collect(joining()),
+                flattenedExpectedTaskIds,
+                flattenedFoundTaskIds);
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
index d94424c59376..e19d6e9e0179 100644
--- a/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
+++ b/packages/SystemUI/src/com/android/systemui/mediaprojection/appselector/data/RecentTaskListProvider.kt
@@ -57,6 +57,7 @@ constructor(
             // activity and a null second task, so the foreground task will be index 1, but when
             // opening the app selector in split screen mode, the foreground task will be the second
             // task in index 0.
+            // TODO(346588978): This needs to be updated for mixed groups
             val foregroundGroup =
                 if (groupedTasks.firstOrNull()?.splitBounds != null) groupedTasks.first()
                 else groupedTasks.elementAtOrNull(1)
@@ -69,7 +70,7 @@ constructor(
                         it.taskInfo1,
                         it.taskInfo1.taskId in foregroundTaskIds && it.taskInfo1.isVisible,
                         userManager.getUserInfo(it.taskInfo1.userId).toUserType(),
-                        it.splitBounds
+                        it.splitBounds,
                     )
 
                 val task2 =
@@ -78,7 +79,7 @@ constructor(
                             it.taskInfo2!!,
                             it.taskInfo2!!.taskId in foregroundTaskIds && it.taskInfo2!!.isVisible,
                             userManager.getUserInfo(it.taskInfo2!!.userId).toUserType(),
-                            it.splitBounds
+                            it.splitBounds,
                         )
                     } else null
 
@@ -92,7 +93,7 @@ constructor(
                 Integer.MAX_VALUE,
                 RECENT_IGNORE_UNAVAILABLE,
                 userTracker.userId,
-                backgroundExecutor
+                backgroundExecutor,
             ) { tasks ->
                 continuation.resume(tasks)
             }