Merge "4b/ Migrate TopTaskTracker to use visible running tasks from Shell" into main
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 97d7179..31aa489 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -667,7 +667,7 @@
         TopTaskTracker.CachedTaskInfo cachedTaskInfo = mGestureState.getRunningTask();
         if (mIsSwipeForSplit) {
             int[] splitTaskIds = TopTaskTracker.INSTANCE.get(mContext).getRunningSplitTaskIds();
-            runningTasks = cachedTaskInfo.getPlaceholderTasks(splitTaskIds);
+            runningTasks = cachedTaskInfo.getSplitPlaceholderTasks(splitTaskIds);
         } else {
             runningTasks = cachedTaskInfo.getPlaceholderTasks();
         }
diff --git a/quickstep/src/com/android/quickstep/GestureState.java b/quickstep/src/com/android/quickstep/GestureState.java
index cff352c..5190ec8 100644
--- a/quickstep/src/com/android/quickstep/GestureState.java
+++ b/quickstep/src/com/android/quickstep/GestureState.java
@@ -28,6 +28,7 @@
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.SET_END_TARGET_HOME;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.SET_END_TARGET_NEW_TASK;
 
+import android.app.TaskInfo;
 import android.content.Intent;
 import android.os.SystemClock;
 import android.view.MotionEvent;
@@ -45,6 +46,7 @@
 import com.android.quickstep.util.ActiveGestureProtoLogProxy;
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.systemui.shared.recents.model.ThumbnailData;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 
 import java.io.PrintWriter;
 import java.util.ArrayList;
@@ -330,13 +332,23 @@
         if (mRunningTask == null) {
             return new int[]{INVALID_TASK_ID, INVALID_TASK_ID};
         } else {
-            int cachedTasksSize = mRunningTask.mAllCachedTasks.size();
-            int count = Math.min(cachedTasksSize, getMultipleTasks ? 2 : 1);
-            int[] runningTaskIds = new int[count];
-            for (int i = 0; i < count; i++) {
-                runningTaskIds[i] = mRunningTask.mAllCachedTasks.get(i).taskId;
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                if (mRunningTask.getVisibleTasks().isEmpty()) {
+                    return new int[0];
+                }
+                GroupedTaskInfo topRunningTask = mRunningTask.getVisibleTasks().getFirst();
+                List<TaskInfo> groupedTasks = topRunningTask.getTaskInfoList();
+                return groupedTasks.stream().mapToInt(
+                        groupedTask -> groupedTask.taskId).toArray();
+            } else {
+                int cachedTasksSize = mRunningTask.mAllCachedTasks.size();
+                int count = Math.min(cachedTasksSize, getMultipleTasks ? 2 : 1);
+                int[] runningTaskIds = new int[count];
+                for (int i = 0; i < count; i++) {
+                    runningTaskIds[i] = mRunningTask.mAllCachedTasks.get(i).taskId;
+                }
+                return runningTaskIds;
             }
-            return runningTaskIds;
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 714838a..85e2b6e 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -114,12 +114,9 @@
             }
 
             @Override
-            public void onTaskMovedToFront(GroupedTaskInfo[] visibleTasks) {
+            public void onTaskMovedToFront(GroupedTaskInfo taskToFront) {
                 mMainThreadExecutor.execute(() -> {
-                    // TODO(b/346588978): We currently are only sending a single task, but this will
-                    //                    be updated once we send the full set of visible tasks
-                    final TaskInfo info = visibleTasks[0].getTaskInfo1();
-                    topTaskTracker.handleTaskMovedToFront(info);
+                    topTaskTracker.handleTaskMovedToFront(taskToFront.getTaskInfo1());
                 });
             }
 
@@ -127,6 +124,13 @@
             public void onTaskInfoChanged(RunningTaskInfo taskInfo) {
                 mMainThreadExecutor.execute(() -> topTaskTracker.onTaskChanged(taskInfo));
             }
+
+            @Override
+            public void onVisibleTasksChanged(GroupedTaskInfo[] visibleTasks) {
+                mMainThreadExecutor.execute(() -> {
+                    topTaskTracker.onVisibleTasksChanged(visibleTasks);
+                });
+            }
         });
         // We may receive onRunningTaskAppeared events later for tasks which have already been
         // included in the list returned by mSysUiProxy.getRunningTasks(), or may receive
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index de8be50..e296449 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -354,7 +354,11 @@
      * @return whether the given running task info matches the gesture-blocked task.
      */
     public boolean isGestureBlockedTask(CachedTaskInfo taskInfo) {
-        return taskInfo != null && taskInfo.getTaskId() == mGestureBlockingTaskId;
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return taskInfo != null && taskInfo.topGroupedTaskContainsTask(mGestureBlockingTaskId);
+        } else {
+            return taskInfo != null && taskInfo.getTaskId() == mGestureBlockingTaskId;
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index c9dfe6d..80d6137 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -18,17 +18,18 @@
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_HOME;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_RECENTS;
-import static android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM;
 import static android.content.Intent.ACTION_CHOOSER;
 import static android.content.Intent.FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS;
 import static android.view.Display.DEFAULT_DISPLAY;
 
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT;
+import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
 
-import android.annotation.UserIdInt;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
 import android.content.Context;
+import android.util.ArrayMap;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -46,8 +47,10 @@
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.TaskStackChangeListener;
 import com.android.systemui.shared.system.TaskStackChangeListeners;
+import com.android.wm.shell.shared.GroupedTaskInfo;
 import com.android.wm.shell.splitscreen.ISplitScreenListener;
 
+import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -61,37 +64,58 @@
  */
 public class TopTaskTracker extends ISplitScreenListener.Stub
         implements TaskStackChangeListener, SafeCloseable {
-
+    private static final String TAG = "TopTaskTracker";
     public static MainThreadInitializedObject<TopTaskTracker> INSTANCE =
             new MainThreadInitializedObject<>(TopTaskTracker::new);
 
     private static final int HISTORY_SIZE = 5;
 
+    private final Context mContext;
+
+    // Only used when Flags.enableShellTopTaskTracking() is disabled
     // Ordered list with first item being the most recent task.
     private final LinkedList<TaskInfo> mOrderedTaskList = new LinkedList<>();
-
-    private final Context mContext;
     private final SplitStageInfo mMainStagePosition = new SplitStageInfo();
     private final SplitStageInfo mSideStagePosition = new SplitStageInfo();
     private int mPinnedTaskId = INVALID_TASK_ID;
 
+    // Only used when Flags.enableShellTopTaskTracking() is enabled
+    // Mapping of display id to running tasks.  Running tasks are ordered from top most to
+    // bottom most.
+    private ArrayMap<Integer, ArrayList<GroupedTaskInfo>> mVisibleTasks = new ArrayMap<>();
+
     private TopTaskTracker(Context context) {
         mContext = context;
-        mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
-        mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
 
-        TaskStackChangeListeners.getInstance().registerTaskStackListener(this);
-        SystemUiProxy.INSTANCE.get(context).registerSplitScreenListener(this);
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // Just prepopulate a list for the default display tasks so we don't need to add null
+            // checks everywhere
+            mVisibleTasks.put(DEFAULT_DISPLAY, new ArrayList<>());
+        } else {
+            mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
+            mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
+
+            TaskStackChangeListeners.getInstance().registerTaskStackListener(this);
+            SystemUiProxy.INSTANCE.get(context).registerSplitScreenListener(this);
+        }
     }
 
     @Override
     public void close() {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         TaskStackChangeListeners.getInstance().unregisterTaskStackListener(this);
         SystemUiProxy.INSTANCE.get(mContext).unregisterSplitScreenListener(this);
     }
 
     @Override
     public void onTaskRemoved(int taskId) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskId);
     }
 
@@ -100,7 +124,11 @@
         handleTaskMovedToFront(taskInfo);
     }
 
-    public void handleTaskMovedToFront(TaskInfo taskInfo) {
+    void handleTaskMovedToFront(TaskInfo taskInfo) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskInfo.taskId);
         mOrderedTaskList.addFirst(taskInfo);
 
@@ -131,8 +159,39 @@
         }
     }
 
+    /**
+     * Called when the set of visible tasks have changed.
+     */
+    public void onVisibleTasksChanged(GroupedTaskInfo[] visibleTasks) {
+        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
+        // TODO(346588978): Per-display info, just have everything in order by display
+
+        // Clear existing tasks for each display
+        mVisibleTasks.forEach((displayId, visibleTasksOnDisplay) -> visibleTasksOnDisplay.clear());
+
+        // Update the visible tasks on each display
+        for (int i = 0; i < visibleTasks.length; i++) {
+            final int displayId = visibleTasks[i].getTaskInfo1().getDisplayId();
+            final ArrayList<GroupedTaskInfo> displayTasks;
+            if (mVisibleTasks.containsKey(displayId)) {
+                displayTasks = mVisibleTasks.get(displayId);
+            } else {
+                displayTasks = new ArrayList<>();
+                mVisibleTasks.put(displayId, displayTasks);
+            }
+            displayTasks.add(visibleTasks[i]);
+        }
+    }
+
     @Override
     public void onStagePositionChanged(@StageType int stage, @StagePosition int position) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         if (stage == SplitConfigurationOptions.STAGE_TYPE_MAIN) {
             mMainStagePosition.stagePosition = position;
         } else {
@@ -141,6 +200,10 @@
     }
 
     public void onTaskChanged(RunningTaskInfo taskInfo) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         for (int i = 0; i < mOrderedTaskList.size(); i++) {
             if (mOrderedTaskList.get(i).taskId == taskInfo.taskId) {
                 mOrderedTaskList.set(i, taskInfo);
@@ -151,6 +214,10 @@
 
     @Override
     public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         // If a task is not visible anymore or has been moved to undefined, stop tracking it.
         if (!visible || stage == SplitConfigurationOptions.STAGE_TYPE_UNDEFINED) {
             if (mMainStagePosition.taskId == taskId) {
@@ -170,11 +237,19 @@
 
     @Override
     public void onActivityPinned(String packageName, int userId, int taskId, int stackId) {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mPinnedTaskId = taskId;
     }
 
     @Override
     public void onActivityUnpinned() {
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
         mPinnedTaskId = INVALID_TASK_ID;
     }
 
@@ -183,21 +258,59 @@
      * Will return empty array if device is not in staged split
      */
     public int[] getRunningSplitTaskIds() {
-        if (mMainStagePosition.taskId == INVALID_TASK_ID
-                || mSideStagePosition.taskId == INVALID_TASK_ID) {
-            return new int[]{};
-        }
-        int[] out = new int[2];
-        if (mMainStagePosition.stagePosition == STAGE_POSITION_TOP_OR_LEFT) {
-            out[0] = mMainStagePosition.taskId;
-            out[1] = mSideStagePosition.taskId;
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // TODO(346588978): This assumes default display for now
+            final ArrayList<GroupedTaskInfo> visibleTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+            final GroupedTaskInfo splitTaskInfo = visibleTasks.stream()
+                    .filter(taskInfo -> taskInfo.getType() == TYPE_SPLIT)
+                    .findFirst().orElse(null);
+            if (splitTaskInfo != null && splitTaskInfo.getSplitBounds() != null) {
+                return new int[] {
+                        splitTaskInfo.getSplitBounds().leftTopTaskId,
+                        splitTaskInfo.getSplitBounds().rightBottomTaskId
+                };
+            }
+            return new int[0];
         } else {
-            out[1] = mMainStagePosition.taskId;
-            out[0] = mSideStagePosition.taskId;
+            if (mMainStagePosition.taskId == INVALID_TASK_ID
+                    || mSideStagePosition.taskId == INVALID_TASK_ID) {
+                return new int[]{};
+            }
+            int[] out = new int[2];
+            if (mMainStagePosition.stagePosition == STAGE_POSITION_TOP_OR_LEFT) {
+                out[0] = mMainStagePosition.taskId;
+                out[1] = mSideStagePosition.taskId;
+            } else {
+                out[1] = mMainStagePosition.taskId;
+                out[0] = mSideStagePosition.taskId;
+            }
+            return out;
         }
-        return out;
     }
 
+    /**
+     * Dumps the list of tasks in top task tracker.
+     */
+    public void dump(PrintWriter pw) {
+        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            return;
+        }
+
+        // TODO(346588978): This assumes default display for now
+        final ArrayList<GroupedTaskInfo> displayTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+        pw.println("TopTaskTracker:");
+        pw.println("  tasks: [");
+        for (GroupedTaskInfo taskInfo : displayTasks) {
+            final TaskInfo info = taskInfo.getTaskInfo1();
+            final boolean isExcluded = (info.baseIntent.getFlags()
+                    & FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS) != 0;
+            pw.println("    " + info.taskId + ": excluded=" + isExcluded
+                    + " visibleRequested=" + info.isVisibleRequested
+                    + " visible=" + info.isVisible
+                    + " " + info.baseIntent.getComponent());
+        }
+        pw.println("  ]");
+    }
 
     /**
      * Returns the CachedTaskInfo for the top most task
@@ -205,25 +318,35 @@
     @NonNull
     @UiThread
     public CachedTaskInfo getCachedTopTask(boolean filterOnlyVisibleRecents) {
-        if (filterOnlyVisibleRecents) {
-            // Since we only know about the top most task, any filtering may not be applied on the
-            // cache. The second to top task may change while the top task is still the same.
-            RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.true", () ->
-                    ActivityManagerWrapper.getInstance().getRunningTasks(true));
-            return new CachedTaskInfo(Arrays.asList(tasks));
-        }
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // TODO(346588978): Currently ignore filterOnlyVisibleRecents, but perhaps make this an
+            //  explicit filter For things to ignore (ie. PIP/Bubbles/Assistant/etc/so that this is
+            //  explicit)
+            // TODO(346588978): This assumes default display for now (as does all of Launcher)
+            final ArrayList<GroupedTaskInfo> displayTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+            return new CachedTaskInfo(new ArrayList<>(displayTasks));
+        } else {
+            if (filterOnlyVisibleRecents) {
+                // Since we only know about the top most task, any filtering may not be applied on
+                // the cache. The second to top task may change while the top task is still the
+                // same.
+                RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.true", () ->
+                        ActivityManagerWrapper.getInstance().getRunningTasks(true));
+                return new CachedTaskInfo(Arrays.asList(tasks));
+            }
 
-        if (mOrderedTaskList.isEmpty()) {
-            RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.false", () ->
-                    ActivityManagerWrapper.getInstance().getRunningTasks(
-                            false /* filterOnlyVisibleRecents */));
-            Collections.addAll(mOrderedTaskList, tasks);
-        }
+            if (mOrderedTaskList.isEmpty()) {
+                RunningTaskInfo[] tasks = TraceHelper.allowIpcs("getCachedTopTask.false", () ->
+                        ActivityManagerWrapper.getInstance().getRunningTasks(
+                                false /* filterOnlyVisibleRecents */));
+                Collections.addAll(mOrderedTaskList, tasks);
+            }
 
-        ArrayList<TaskInfo> tasks = new ArrayList<>(mOrderedTaskList);
-        // Strip the pinned task and recents task
-        tasks.removeIf(t -> t.taskId == mPinnedTaskId || isRecentsTask(t));
-        return new CachedTaskInfo(tasks);
+            ArrayList<TaskInfo> tasks = new ArrayList<>(mOrderedTaskList);
+            // Strip the pinned task and recents task
+            tasks.removeIf(t -> t.taskId == mPinnedTaskId || isRecentsTask(t));
+            return new CachedTaskInfo(tasks);
+        }
     }
 
     private static boolean isRecentsTask(TaskInfo task) {
@@ -237,24 +360,79 @@
      */
     public static class CachedTaskInfo {
 
+        // Only used when enableShellTopTaskTracking() is disabled
         @Nullable
         private final TaskInfo mTopTask;
+        @Nullable
         public final List<TaskInfo> mAllCachedTasks;
 
-        CachedTaskInfo(List<TaskInfo> allCachedTasks) {
+        // Only used when enableShellTopTaskTracking() is enabled
+        @Nullable
+        private final GroupedTaskInfo mTopGroupedTask;
+        @Nullable
+        private final ArrayList<GroupedTaskInfo> mVisibleTasks;
+
+
+        // Only used when enableShellTopTaskTracking() is enabled
+        CachedTaskInfo(@NonNull ArrayList<GroupedTaskInfo> visibleTasks) {
+            mAllCachedTasks = null;
+            mTopTask = null;
+            mVisibleTasks = visibleTasks;
+            mTopGroupedTask = !mVisibleTasks.isEmpty() ? mVisibleTasks.getFirst() : null;
+
+        }
+
+        // Only used when enableShellTopTaskTracking() is disabled
+        CachedTaskInfo(@NonNull List<TaskInfo> allCachedTasks) {
+            mVisibleTasks = null;
+            mTopGroupedTask = null;
             mAllCachedTasks = allCachedTasks;
             mTopTask = allCachedTasks.isEmpty() ? null : allCachedTasks.get(0);
         }
 
+        /**
+         * @return The list of visible tasks
+         */
+        public ArrayList<GroupedTaskInfo> getVisibleTasks() {
+            return mVisibleTasks;
+        }
+
+        /**
+         * @return The top task id
+         */
         public int getTaskId() {
-            return mTopTask == null ? INVALID_TASK_ID : mTopTask.taskId;
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // Callers should use topGroupedTaskContainsTask() instead
+                return INVALID_TASK_ID;
+            } else {
+                return mTopTask != null ? mTopTask.taskId : INVALID_TASK_ID;
+            }
+        }
+
+        /**
+         * @return Whether the top grouped task contains the given {@param taskId} if
+         *         Flags.enableShellTopTaskTracking() is true, otherwise it checks the top
+         *         task as reported from TaskStackListener.
+         */
+        public boolean topGroupedTaskContainsTask(int taskId) {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                return mTopGroupedTask != null && mTopGroupedTask.containsTask(taskId);
+            } else {
+                return mTopTask != null && mTopTask.taskId == taskId;
+            }
         }
 
         /**
          * Returns true if the root of the task chooser activity
          */
         public boolean isRootChooseActivity() {
-            return mTopTask != null && ACTION_CHOOSER.equals(mTopTask.baseIntent.getAction());
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                return mTopGroupedTask != null && ACTION_CHOOSER.equals(
+                        mTopGroupedTask.getTaskInfo1().baseIntent.getAction());
+            } else {
+                return mTopTask != null && ACTION_CHOOSER.equals(mTopTask.baseIntent.getAction());
+            }
         }
 
         /**
@@ -262,6 +440,10 @@
          * is another running task that is not excluded from recents, returns that underlying task.
          */
         public @Nullable CachedTaskInfo getVisibleNonExcludedTask() {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // Callers should not need this when the full set of visible tasks are provided
+                return null;
+            }
             if (mTopTask == null
                     || (mTopTask.baseIntent.getFlags() & FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS) == 0) {
                 // Not an excluded task.
@@ -278,24 +460,30 @@
         }
 
         /**
-         * Returns true if this represents the HOME task
+         * Returns true if this represents the HOME activity type task
          */
         public boolean isHomeTask() {
-            return mTopTask != null && mTopTask.configuration.windowConfiguration
-                    .getActivityType() == ACTIVITY_TYPE_HOME;
-        }
-
-        public boolean isRecentsTask() {
-            return TopTaskTracker.isRecentsTask(mTopTask);
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                return mTopGroupedTask != null
+                        && mTopGroupedTask.getTaskInfo1().getActivityType() == ACTIVITY_TYPE_HOME;
+            } else {
+                return mTopTask != null && mTopTask.configuration.windowConfiguration
+                        .getActivityType() == ACTIVITY_TYPE_HOME;
+            }
         }
 
         /**
-         * Returns {@code true} if this task windowing mode is set to {@link
-         * android.app.WindowConfiguration#WINDOWING_MODE_FREEFORM}
+         * Returns true if this represents the RECENTS activity type task
          */
-        public boolean isFreeformTask() {
-            return mTopTask != null && mTopTask.configuration.windowConfiguration.getWindowingMode()
-                    == WINDOWING_MODE_FREEFORM;
+        public boolean isRecentsTask() {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                return mTopGroupedTask != null
+                        && TopTaskTracker.isRecentsTask(mTopGroupedTask.getTaskInfo1());
+            } else {
+                return TopTaskTracker.isRecentsTask(mTopTask);
+            }
         }
 
         /**
@@ -303,43 +491,78 @@
          * is loaded by the model
          */
         public Task[] getPlaceholderTasks() {
-            return mTopTask == null ? new Task[0]
-                    : new Task[]{Task.from(new TaskKey(mTopTask), mTopTask, false)};
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to return more than a single task once the callers
+                //  are refactored
+                if (mVisibleTasks.isEmpty()) {
+                    return new Task[0];
+                }
+                final TaskInfo info = mVisibleTasks.getFirst().getTaskInfo1();
+                return new Task[]{Task.from(new TaskKey(info), info, false)};
+            } else {
+                return mTopTask == null ? new Task[0]
+                        : new Task[]{Task.from(new TaskKey(mTopTask), mTopTask, false)};
+            }
         }
 
         /**
          * Returns {@link Task} array corresponding to the provided task ids which can be used as a
          * placeholder until the true object is loaded by the model
          */
-        public Task[] getPlaceholderTasks(int[] taskIds) {
-            if (mTopTask == null) {
-                return new Task[0];
-            }
-            Task[] result = new Task[taskIds.length];
-            for (int i = 0; i < taskIds.length; i++) {
-                final int index = i;
-                int taskId = taskIds[i];
-                mAllCachedTasks.forEach(rti -> {
-                    if (rti.taskId == taskId) {
-                        result[index] = Task.from(new TaskKey(rti), rti, false);
-                    }
-                });
-            }
-            return result;
-        }
+        public Task[] getSplitPlaceholderTasks(int[] taskIds) {
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                if (mVisibleTasks.isEmpty()
+                        || mVisibleTasks.getFirst().getType() != TYPE_SPLIT) {
+                    return new Task[0];
+                }
 
-        @UserIdInt
-        @Nullable
-        public Integer getUserId() {
-            return mTopTask == null ? null : mTopTask.userId;
+                GroupedTaskInfo splitTask = mVisibleTasks.getFirst();
+                Task[] result = new Task[taskIds.length];
+                for (int i = 0; i < taskIds.length; i++) {
+                    TaskInfo info = splitTask.getTaskById(taskIds[i]);
+                    if (info == null) {
+                        Log.w(TAG, "Requested task (" + taskIds[i] + ") not found");
+                        return new Task[0];
+                    }
+                    result[i] = Task.from(new TaskKey(info), info, false);
+                }
+                return result;
+            } else {
+                if (mTopTask == null) {
+                    return new Task[0];
+                }
+                Task[] result = new Task[taskIds.length];
+                for (int i = 0; i < taskIds.length; i++) {
+                    final int index = i;
+                    int taskId = taskIds[i];
+                    mAllCachedTasks.forEach(rti -> {
+                        if (rti.taskId == taskId) {
+                            result[index] = Task.from(new TaskKey(rti), rti, false);
+                        }
+                    });
+                }
+                return result;
+            }
         }
 
         @Nullable
         public String getPackageName() {
-            if (mTopTask == null || mTopTask.baseActivity == null) {
-                return null;
+            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                // TODO(346588978): Update this to not make an assumption on a specific task info
+                if (mTopGroupedTask == null) {
+                    return null;
+                }
+                final TaskInfo info = mTopGroupedTask.getTaskInfo1();
+                if (info.baseActivity == null) {
+                    return null;
+                }
+                return info.baseActivity.getPackageName();
+            } else {
+                if (mTopTask == null || mTopTask.baseActivity == null) {
+                    return null;
+                }
+                return mTopTask.baseActivity.getPackageName();
             }
-            return mTopTask.baseActivity.getPackageName();
         }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index ad5720f..0242fb6 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -1362,15 +1362,17 @@
                 && runningTask != null
                 && runningTask.isRootChooseActivity();
 
-        // In the case where we are in an excluded, translucent overlay, ignore it and treat the
-        // running activity as the task behind the overlay.
-        TopTaskTracker.CachedTaskInfo otherVisibleTask = runningTask == null
-                ? null
-                : runningTask.getVisibleNonExcludedTask();
-        if (otherVisibleTask != null) {
-            ActiveGestureProtoLogProxy.logUpdateGestureStateRunningTask(
-                    otherVisibleTask.getPackageName(), runningTask.getPackageName());
-            gestureState.updateRunningTask(otherVisibleTask);
+        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            // In the case where we are in an excluded, translucent overlay, ignore it and treat the
+            // running activity as the task behind the overlay.
+            TopTaskTracker.CachedTaskInfo otherVisibleTask = runningTask == null
+                    ? null
+                    : runningTask.getVisibleNonExcludedTask();
+            if (otherVisibleTask != null) {
+                ActiveGestureProtoLogProxy.logUpdateGestureStateRunningTask(
+                        otherVisibleTask.getPackageName(), runningTask.getPackageName());
+                gestureState.updateRunningTask(otherVisibleTask);
+            }
         }
 
         boolean previousGestureAnimatedToLauncher =
@@ -1672,6 +1674,7 @@
         ContextualSearchStateManager.INSTANCE.get(this).dump("\t", pw);
         SystemUiProxy.INSTANCE.get(this).dump(pw);
         DeviceConfigWrapper.get().dump("   ", pw);
+        TopTaskTracker.INSTANCE.get(this).dump(pw);
     }
 
     private AbsSwipeUpHandler createLauncherSwipeHandler(
diff --git a/quickstep/src/com/android/quickstep/util/AppPairsController.java b/quickstep/src/com/android/quickstep/util/AppPairsController.java
index 7388d59..1312aa4 100644
--- a/quickstep/src/com/android/quickstep/util/AppPairsController.java
+++ b/quickstep/src/com/android/quickstep/util/AppPairsController.java
@@ -409,8 +409,8 @@
             );
         } else {
             // Tapped an app pair while in a single app
-            int runningTaskId = topTaskTracker
-                    .getCachedTopTask(false /* filterOnlyVisibleRecents */).getTaskId();
+            final TopTaskTracker.CachedTaskInfo runningTask = topTaskTracker
+                    .getCachedTopTask(false /* filterOnlyVisibleRecents */);
 
             mSplitSelectStateController.findLastActiveTasksAndRunCallback(
                     componentKeys,
@@ -418,10 +418,21 @@
                     foundTasks -> {
                         Task foundTask1 = foundTasks[0];
                         Task foundTask2 = foundTasks[1];
-                        boolean task1IsOnScreen =
-                                foundTask1 != null && foundTask1.getKey().getId() == runningTaskId;
-                        boolean task2IsOnScreen =
-                                foundTask2 != null && foundTask2.getKey().getId() == runningTaskId;
+                        boolean task1IsOnScreen;
+                        boolean task2IsOnScreen;
+                        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+                            task1IsOnScreen = foundTask1 != null
+                                    && runningTask.topGroupedTaskContainsTask(
+                                    foundTask1.getKey().getId());
+                            task2IsOnScreen = foundTask2 != null
+                                    && runningTask.topGroupedTaskContainsTask(
+                                    foundTask2.getKey().getId());
+                        } else {
+                            task1IsOnScreen = foundTask1 != null && foundTask1.getKey().getId()
+                                    == runningTask.getTaskId();
+                            task2IsOnScreen = foundTask2 != null && foundTask2.getKey().getId()
+                                    == runningTask.getTaskId();
+                        }
 
                         if (!task1IsOnScreen && !task2IsOnScreen) {
                             // Neither App A nor App B are on-screen, launch the app pair normally.
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 3a4e328..9d74bfb 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1966,7 +1966,7 @@
                     // We try to avoid this because it can cause a scroll jump, but it is needed
                     // for cases where the running task isn't included in this load plan (e.g. if
                     // the current running task is excludedFromRecents.)
-                    showCurrentTask(mActiveGestureRunningTasks);
+                    showCurrentTask(mActiveGestureRunningTasks, "applyLoadPlan");
                 } else {
                     setRunningTaskViewId(INVALID_TASK_ID);
                 }
@@ -2749,7 +2749,7 @@
             updateSizeAndPadding();
         }
 
-        showCurrentTask(mActiveGestureRunningTasks);
+        showCurrentTask(mActiveGestureRunningTasks, "onGestureAnimationStart");
         setEnableFreeScroll(false);
         setEnableDrawingLiveTile(false);
         setRunningTaskHidden(true);
@@ -2930,8 +2930,9 @@
      * All subsequent calls to reload will keep the task as the first item until {@link #reset()}
      * is called.  Also scrolls the view to this task.
      */
-    private void showCurrentTask(Task[] runningTasks) {
-        Log.d(TAG, "showCurrentTask - runningTasks: " + Arrays.toString(runningTasks));
+    private void showCurrentTask(Task[] runningTasks, String caller) {
+        Log.d(TAG, "showCurrentTask(" + caller + ") - runningTasks: "
+                + Arrays.toString(runningTasks));
         if (runningTasks.length == 0) {
             return;
         }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt
index 541a48d..ee70e0a 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/AppPairsControllerTest.kt
@@ -202,7 +202,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -226,7 +226,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -250,7 +250,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -274,7 +274,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -298,7 +298,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -322,7 +322,7 @@
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -341,12 +341,16 @@
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 1 already on screen
-        whenever(mockCachedTaskInfo.taskId).thenReturn(1)
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(1))).thenReturn(true)
+        } else {
+            whenever(mockCachedTaskInfo.taskId).thenReturn(1)
+        }
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -365,12 +369,16 @@
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 2 already on screen
-        whenever(mockCachedTaskInfo.taskId).thenReturn(2)
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(2))).thenReturn(true)
+        } else {
+            whenever(mockCachedTaskInfo.taskId).thenReturn(2)
+        }
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())
@@ -389,12 +397,16 @@
         whenever(mockTaskKey1.getId()).thenReturn(1)
         whenever(mockTaskKey2.getId()).thenReturn(2)
         // ... with app 3 already on screen
-        whenever(mockCachedTaskInfo.taskId).thenReturn(3)
+        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            whenever(mockCachedTaskInfo.topGroupedTaskContainsTask(eq(3))).thenReturn(true)
+        } else {
+            whenever(mockCachedTaskInfo.taskId).thenReturn(3)
+        }
 
         // Trigger app pair launch, capture and run callback from findLastActiveTasksAndRunCallback
         spyAppPairsController.handleAppPairLaunchInApp(
             mockAppPairIcon,
-            listOf(mockItemInfo1, mockItemInfo2)
+            listOf(mockItemInfo1, mockItemInfo2),
         )
         verify(splitSelectStateController)
             .findLastActiveTasksAndRunCallback(any(), any(), callbackCaptor.capture())