diff --git a/quickstep/src/com/android/quickstep/GestureState.java b/quickstep/src/com/android/quickstep/GestureState.java
index e0fa77a..e1d4536 100644
--- a/quickstep/src/com/android/quickstep/GestureState.java
+++ b/quickstep/src/com/android/quickstep/GestureState.java
@@ -28,7 +28,6 @@
 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;
@@ -47,7 +46,6 @@
 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;
@@ -334,13 +332,7 @@
             return new int[]{INVALID_TASK_ID, INVALID_TASK_ID};
         } else {
             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();
+                return mRunningTask.topGroupedTaskIds();
             } else {
                 int cachedTasksSize = mRunningTask.mAllCachedTasks.size();
                 int count = Math.min(cachedTasksSize, getMultipleTasks ? 2 : 1);
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 32ccd72..ee4ee38 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.quickstep.util.SplitScreenUtils.convertShellSplitBoundsToLauncher;
 import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_FREEFORM;
+import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
 
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.KeyguardManager;
@@ -39,6 +40,7 @@
 import com.android.quickstep.util.DesktopTask;
 import com.android.quickstep.util.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
+import com.android.wm.shell.Flags;
 import com.android.wm.shell.recents.IRecentTasksListener;
 import com.android.wm.shell.shared.GroupedTaskInfo;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
@@ -116,7 +118,8 @@
             @Override
             public void onTaskMovedToFront(GroupedTaskInfo taskToFront) {
                 mMainThreadExecutor.execute(() -> {
-                    topTaskTracker.handleTaskMovedToFront(taskToFront.getTaskInfo1());
+                    topTaskTracker.handleTaskMovedToFront(
+                            taskToFront.getBaseGroupedTask().getTaskInfo1());
                 });
             }
 
@@ -340,50 +343,74 @@
 
         int numVisibleTasks = 0;
         for (GroupedTaskInfo rawTask : rawTasks) {
-            if (rawTask.getType() == TYPE_FREEFORM) {
+            if (rawTask.isBaseType(TYPE_FREEFORM)) {
                 // TYPE_FREEFORM tasks is only created when desktop mode can be entered,
                 // leftover TYPE_FREEFORM tasks created when flag was on should be ignored.
                 if (DesktopModeStatus.canEnterDesktopMode(mContext)) {
-                    GroupTask desktopTask = createDesktopTask(rawTask);
+                    GroupTask desktopTask = createDesktopTask(rawTask.getBaseGroupedTask());
                     if (desktopTask != null) {
                         allTasks.add(desktopTask);
                     }
                 }
                 continue;
             }
-            TaskInfo taskInfo1 = rawTask.getTaskInfo1();
-            TaskInfo taskInfo2 = rawTask.getTaskInfo2();
-            Task.TaskKey task1Key = new Task.TaskKey(taskInfo1);
-            Task task1 = loadKeysOnly
-                    ? new Task(task1Key)
-                    : Task.from(task1Key, taskInfo1,
-                            tmpLockedUsers.get(task1Key.userId) /* isLocked */);
-            Task task2 = null;
-            if (taskInfo2 != null) {
-                // Is split task
-                Task.TaskKey task2Key = new Task.TaskKey(taskInfo2);
-                task2 = loadKeysOnly
-                        ? new Task(task2Key)
-                        : Task.from(task2Key, taskInfo2,
-                                tmpLockedUsers.get(task2Key.userId) /* isLocked */);
+
+            if (Flags.enableShellTopTaskTracking()) {
+                final TaskInfo taskInfo1 = rawTask.getBaseGroupedTask().getTaskInfo1();
+                final Task.TaskKey task1Key = new Task.TaskKey(taskInfo1);
+                final Task task1 = Task.from(task1Key, taskInfo1,
+                        tmpLockedUsers.get(task1Key.userId) /* isLocked */);
+                final Task task2;
+                final SplitConfigurationOptions.SplitBounds launcherSplitBounds;
+
+                if (rawTask.isBaseType(TYPE_SPLIT)) {
+                    final TaskInfo taskInfo2 = rawTask.getBaseGroupedTask().getTaskInfo2();
+                    final Task.TaskKey task2Key = new Task.TaskKey(taskInfo2);
+                    task2 = Task.from(task2Key, taskInfo2,
+                            tmpLockedUsers.get(task2Key.userId) /* isLocked */);
+                    launcherSplitBounds =
+                            convertShellSplitBoundsToLauncher(
+                                    rawTask.getBaseGroupedTask().getSplitBounds());
+                } else {
+                    task2 = null;
+                    launcherSplitBounds = null;
+                }
+                allTasks.add(new GroupTask(task1, task2, launcherSplitBounds));
             } else {
-                // Is fullscreen task
-                if (numVisibleTasks > 0) {
-                    boolean isExcluded = (taskInfo1.baseIntent.getFlags()
-                            & FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS) != 0;
-                    if (taskInfo1.isTopActivityTransparent && isExcluded) {
-                        // If there are already visible tasks, then ignore the excluded tasks and
-                        // don't add them to the returned list
-                        continue;
+                TaskInfo taskInfo1 = rawTask.getTaskInfo1();
+                TaskInfo taskInfo2 = rawTask.getTaskInfo2();
+                Task.TaskKey task1Key = new Task.TaskKey(taskInfo1);
+                Task task1 = loadKeysOnly
+                        ? new Task(task1Key)
+                        : Task.from(task1Key, taskInfo1,
+                                tmpLockedUsers.get(task1Key.userId) /* isLocked */);
+                Task task2 = null;
+                if (taskInfo2 != null) {
+                    // Is split task
+                    Task.TaskKey task2Key = new Task.TaskKey(taskInfo2);
+                    task2 = loadKeysOnly
+                            ? new Task(task2Key)
+                            : Task.from(task2Key, taskInfo2,
+                                    tmpLockedUsers.get(task2Key.userId) /* isLocked */);
+                } else {
+                    // Is fullscreen task
+                    if (numVisibleTasks > 0) {
+                        boolean isExcluded = (taskInfo1.baseIntent.getFlags()
+                                & FLAG_ACTIVITY_EXCLUDE_FROM_RECENTS) != 0;
+                        if (taskInfo1.isTopActivityTransparent && isExcluded) {
+                            // If there are already visible tasks, then ignore the excluded tasks
+                            // and don't add them to the returned list
+                            continue;
+                        }
                     }
                 }
+                if (taskInfo1.isVisible) {
+                    numVisibleTasks++;
+                }
+                final SplitConfigurationOptions.SplitBounds launcherSplitBounds =
+                        convertShellSplitBoundsToLauncher(rawTask.getSplitBounds());
+                allTasks.add(new GroupTask(task1, task2, launcherSplitBounds));
             }
-            if (taskInfo1.isVisible) {
-                numVisibleTasks++;
-            }
-            final SplitConfigurationOptions.SplitBounds launcherSplitBounds =
-                    convertShellSplitBoundsToLauncher(rawTask.getSplitBounds());
-            allTasks.add(new GroupTask(task1, task2, launcherSplitBounds));
         }
 
         return allTasks;
@@ -409,14 +436,6 @@
         return new DesktopTask(tasks);
     }
 
-    private ArrayList<GroupTask> copyOf(ArrayList<GroupTask> tasks) {
-        ArrayList<GroupTask> newTasks = new ArrayList<>();
-        for (int i = 0; i < tasks.size(); i++) {
-            newTasks.add(tasks.get(i).copy());
-        }
-        return newTasks;
-    }
-
     public void dump(String prefix, PrintWriter writer) {
         writer.println(prefix + "RecentTasksList:");
         writer.println(prefix + "  mChangeId=" + mChangeId);
@@ -439,14 +458,7 @@
         }
         writer.println(prefix + "  rawTasks=[");
         for (GroupedTaskInfo task : rawTasks) {
-            TaskInfo taskInfo1 = task.getTaskInfo1();
-            TaskInfo taskInfo2 = task.getTaskInfo2();
-            ComponentName cn1 = taskInfo1.topActivity;
-            ComponentName cn2 = taskInfo2 != null ? taskInfo2.topActivity : null;
-            writer.println(prefix + "    t1: (id=" + taskInfo1.taskId
-                    + "; package=" + (cn1 != null ? cn1.getPackageName() + ")" : "no package)")
-                    + " t2: (id=" + (taskInfo2 != null ? taskInfo2.taskId : "-1")
-                    + "; package=" + (cn2 != null ? cn2.getPackageName() + ")" : "no package)"));
+            writer.println(prefix + task);
         }
         writer.println(prefix + "  ]");
     }
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index bfd6107..b3d9da3 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -24,12 +24,12 @@
 
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_TYPE_A;
+import static com.android.wm.shell.Flags.enableShellTopTaskTracking;
 import static com.android.wm.shell.Flags.enableFlexibleSplit;
 import static com.android.wm.shell.shared.GroupedTaskInfo.TYPE_SPLIT;
 
 import android.app.ActivityManager.RunningTaskInfo;
 import android.app.TaskInfo;
-import android.content.Context;
 import android.util.ArrayMap;
 import android.util.Log;
 
@@ -37,7 +37,6 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
 
-import com.android.launcher3.dagger.ApplicationContext;
 import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.util.DaggerSingletonObject;
 import com.android.launcher3.util.DaggerSingletonTracker;
@@ -77,8 +76,6 @@
 
     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<>();
@@ -87,20 +84,13 @@
     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<>();
+    // Mapping of display id to visible tasks.  Visible tasks are ordered from top most to bottom
+    // most.
+    private ArrayMap<Integer, GroupedTaskInfo> mVisibleTasks = new ArrayMap<>();
 
     @Inject
-    public TopTaskTracker(@ApplicationContext Context context, DaggerSingletonTracker tracker,
-            SystemUiProxy systemUiProxy) {
-        mContext = context;
-
-        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 {
+    public TopTaskTracker(DaggerSingletonTracker tracker, SystemUiProxy systemUiProxy) {
+        if (!enableShellTopTaskTracking()) {
             mMainStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_MAIN;
             mSideStagePosition.stageType = SplitConfigurationOptions.STAGE_TYPE_SIDE;
 
@@ -109,7 +99,7 @@
         }
 
         tracker.addCloseable(() -> {
-            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            if (enableShellTopTaskTracking()) {
                 return;
             }
 
@@ -120,7 +110,7 @@
 
     @Override
     public void onTaskRemoved(int taskId) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -133,7 +123,7 @@
     }
 
     void handleTaskMovedToFront(TaskInfo taskInfo) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -187,32 +177,25 @@
      * Called when the set of visible tasks have changed.
      */
     public void onVisibleTasksChanged(GroupedTaskInfo[] visibleTasks) {
-        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (!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());
+        mVisibleTasks.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]);
+        Log.d(TAG, "onVisibleTasksChanged:");
+        for (GroupedTaskInfo groupedTask : visibleTasks) {
+            Log.d(TAG, "\t" + groupedTask);
+            final int displayId = groupedTask.getBaseGroupedTask().getTaskInfo1().getDisplayId();
+            mVisibleTasks.put(displayId, groupedTask);
         }
     }
 
     @Override
     public void onStagePositionChanged(@StageType int stage, @StagePosition int position) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -224,7 +207,7 @@
     }
 
     public void onTaskChanged(RunningTaskInfo taskInfo) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -238,7 +221,7 @@
 
     @Override
     public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -262,7 +245,7 @@
 
     @Override
     public void onActivityPinned(String packageName, int userId, int taskId, int stackId) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -271,7 +254,7 @@
 
     @Override
     public void onActivityUnpinned() {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (enableShellTopTaskTracking()) {
             return;
         }
 
@@ -279,16 +262,17 @@
     }
 
     /**
-     * @return index 0 will be task in left/top position, index 1 in right/bottom position.
-     * Will return empty array if device is not in staged split
+     * Return the running split task ids.  Index 0 will be task in left/top position, index 1 in
+     * right/bottom position, or and empty array if device is not in splitscreen.
      */
     public int[] getRunningSplitTaskIds() {
-        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 (enableShellTopTaskTracking()) {
+            // TODO(346588978): This assumes default display as splitscreen is only currently there
+            final GroupedTaskInfo visibleTasks = mVisibleTasks.get(DEFAULT_DISPLAY);
+            final GroupedTaskInfo splitTaskInfo =
+                    visibleTasks != null && visibleTasks.isBaseType(TYPE_SPLIT)
+                            ? visibleTasks.getBaseGroupedTask()
+                            : null;
             if (splitTaskInfo != null && splitTaskInfo.getSplitBounds() != null) {
                 return new int[] {
                         splitTaskInfo.getSplitBounds().leftTopTaskId,
@@ -317,24 +301,13 @@
      * Dumps the list of tasks in top task tracker.
      */
     public void dump(PrintWriter pw) {
-        if (!com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (!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("  ]");
+        mVisibleTasks.forEach((displayId, tasks) ->
+                pw.println("  visibleTasks(" + displayId + "): " + tasks));
     }
 
     /**
@@ -343,13 +316,12 @@
     @NonNull
     @UiThread
     public CachedTaskInfo getCachedTopTask(boolean filterOnlyVisibleRecents) {
-        if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+        if (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));
+            // TODO(346588978): This assumes default display as gesture nav is only supported there
+            return new CachedTaskInfo(mVisibleTasks.get(DEFAULT_DISPLAY));
         } else {
             if (filterOnlyVisibleRecents) {
                 // Since we only know about the top most task, any filtering may not be applied on
@@ -374,6 +346,11 @@
         }
     }
 
+    private static boolean isHomeTask(TaskInfo task) {
+        return task != null && task.configuration.windowConfiguration
+                .getActivityType() == ACTIVITY_TYPE_HOME;
+    }
+
     private static boolean isRecentsTask(TaskInfo task) {
         return task != null && task.configuration.windowConfiguration
                 .getActivityType() == ACTIVITY_TYPE_RECENTS;
@@ -384,7 +361,6 @@
      * during the lifecycle of the task.
      */
     public static class CachedTaskInfo {
-
         // Only used when enableShellTopTaskTracking() is disabled
         @Nullable
         private final TaskInfo mTopTask;
@@ -393,40 +369,48 @@
 
         // Only used when enableShellTopTaskTracking() is enabled
         @Nullable
-        private final GroupedTaskInfo mTopGroupedTask;
-        @Nullable
-        private final ArrayList<GroupedTaskInfo> mVisibleTasks;
+        private final GroupedTaskInfo mVisibleTasks;
 
 
         // Only used when enableShellTopTaskTracking() is enabled
-        CachedTaskInfo(@NonNull ArrayList<GroupedTaskInfo> visibleTasks) {
+        CachedTaskInfo(@Nullable 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
+         * Returns the "base" task that is used the as the representative running task of the set
+         * of tasks initially provided.
+         *
+         * Not for general use, as in other windowing modes (ie. split/desktop) the caller should
+         * not make assumptions about there being a single base task.
+         * TODO(346588978): Try to remove all usage of this if possible
          */
-        public ArrayList<GroupedTaskInfo> getVisibleTasks() {
-            return mVisibleTasks;
+        @Nullable
+        private TaskInfo getLegacyBaseTask() {
+            if (enableShellTopTaskTracking()) {
+                return mVisibleTasks != null
+                        ? mVisibleTasks.getBaseGroupedTask().getTaskInfo1()
+                        : null;
+            } else {
+                return mTopTask;
+            }
         }
 
         /**
-         * @return The top task id
+         * Returns the top task id.
          */
         public int getTaskId() {
-            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
+            if (enableShellTopTaskTracking()) {
                 // Callers should use topGroupedTaskContainsTask() instead
                 return INVALID_TASK_ID;
             } else {
@@ -435,29 +419,58 @@
         }
 
         /**
-         * @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.
+         * Returns the top grouped task ids if Flags.enableShellTopTaskTracking() is true, otherwise
+         * an empty array.
+         */
+        public int[] topGroupedTaskIds() {
+            if (enableShellTopTaskTracking()) {
+                if (mVisibleTasks == null) {
+                    return new int[0];
+                }
+                List<TaskInfo> groupedTasks = mVisibleTasks.getTaskInfoList();
+                return groupedTasks.stream().mapToInt(
+                        groupedTask -> groupedTask.taskId).toArray();
+            } else {
+                // Not used
+                return new int[0];
+            }
+        }
+
+        /**
+         * Returns 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);
+            if (enableShellTopTaskTracking()) {
+                return mVisibleTasks != null && mVisibleTasks.containsTask(taskId);
             } else {
                 return mTopTask != null && mTopTask.taskId == taskId;
             }
         }
 
         /**
-         * Returns true if the root of the task chooser activity
+         * Returns true if this represents the task chooser activity
          */
         public boolean isRootChooseActivity() {
-            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());
-            }
+            final TaskInfo baseTask = getLegacyBaseTask();
+            return baseTask != null && ACTION_CHOOSER.equals(baseTask.baseIntent.getAction());
+        }
+
+        /**
+         * Returns true if this represents the HOME activity type task
+         */
+        public boolean isHomeTask() {
+            final TaskInfo baseTask = getLegacyBaseTask();
+            return baseTask != null && TopTaskTracker.isHomeTask(baseTask);
+        }
+
+        /**
+         * Returns true if this represents the RECENTS activity type task
+         */
+        public boolean isRecentsTask() {
+            final TaskInfo baseTask = getLegacyBaseTask();
+            return baseTask != null && TopTaskTracker.isRecentsTask(baseTask);
         }
 
         /**
@@ -465,7 +478,7 @@
          * 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()) {
+            if (enableShellTopTaskTracking()) {
                 // Callers should not need this when the full set of visible tasks are provided
                 return null;
             }
@@ -485,49 +498,16 @@
         }
 
         /**
-         * Returns true if this represents the HOME activity type task
-         */
-        public boolean isHomeTask() {
-            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 true if this represents the RECENTS activity type task
-         */
-        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);
-            }
-        }
-
-        /**
          * Returns {@link Task} array which can be used as a placeholder until the true object
          * is loaded by the model
          */
         public Task[] getPlaceholderTasks() {
-            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)};
-            }
+            final TaskInfo baseTask = getLegacyBaseTask();
+            // TODO(346588978): Update this to return more than a single task once the callers
+            //  are refactored
+            return baseTask == null
+                    ? new Task[0]
+                    : new Task[]{Task.from(new TaskKey(baseTask), baseTask, false)};
         }
 
         /**
@@ -535,13 +515,12 @@
          * placeholder until the true object is loaded by the model
          */
         public Task[] getSplitPlaceholderTasks(int[] taskIds) {
-            if (com.android.wm.shell.Flags.enableShellTopTaskTracking()) {
-                if (mVisibleTasks.isEmpty()
-                        || mVisibleTasks.getFirst().getType() != TYPE_SPLIT) {
+            if (enableShellTopTaskTracking()) {
+                if (mVisibleTasks == null || !mVisibleTasks.isBaseType(TYPE_SPLIT)) {
                     return new Task[0];
                 }
 
-                GroupedTaskInfo splitTask = mVisibleTasks.getFirst();
+                GroupedTaskInfo splitTask = mVisibleTasks.getBaseGroupedTask();
                 Task[] result = new Task[taskIds.length];
                 for (int i = 0; i < taskIds.length; i++) {
                     TaskInfo info = splitTask.getTaskById(taskIds[i]);
@@ -572,22 +551,11 @@
 
         @Nullable
         public String getPackageName() {
-            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();
+            final TaskInfo baseTask = getLegacyBaseTask();
+            if (baseTask == null || baseTask.baseActivity == null) {
+                return null;
             }
+            return baseTask.baseActivity.getPackageName();
         }
     }
 }