diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index 2e78489..32e4097 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -44,7 +44,7 @@
 import com.android.quickstep.util.TISBindHelper;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer;
+import com.android.quickstep.views.TaskView.TaskContainer;
 import com.android.systemui.shared.recents.model.Task;
 
 import java.io.PrintWriter;
@@ -259,14 +259,14 @@
                         if (foundTaskView != null) {
                             // There is already a running app of this type, use that as second app.
                             // Get index of task (0 or 1), in case it's a GroupedTaskView
-                            TaskIdAttributeContainer taskAttributes =
-                                    foundTaskView.getTaskAttributesById(foundTask.key.id);
+                            TaskContainer taskContainer =
+                                    foundTaskView.getTaskContainerById(foundTask.key.id);
                             recents.confirmSplitSelect(
                                     foundTaskView,
                                     foundTask,
-                                    taskAttributes.getIconView().getDrawable(),
-                                    taskAttributes.getThumbnailView(),
-                                    taskAttributes.getThumbnailView().getThumbnail(),
+                                    taskContainer.getIconView().getDrawable(),
+                                    taskContainer.getThumbnailView(),
+                                    taskContainer.getThumbnailView().getThumbnail(),
                                     null /* intent */,
                                     null /* user */,
                                     info);
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
index 05a55d0..16185f5 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
@@ -151,7 +151,7 @@
             int sysuiFlags = 0;
             TaskView tv = mOverviewPanel.getTaskViewAt(0);
             if (tv != null) {
-                sysuiFlags = tv.getThumbnail().getSysUiStatusNavFlags();
+                sysuiFlags = tv.getFirstThumbnailView().getSysUiStatusNavFlags();
             }
             mLauncher.getSystemUiController().updateUiState(UI_STATE_FULLSCREEN_TASK, sysuiFlags);
         } else {
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewTouchController.java b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewTouchController.java
index 300d697..26b528c 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewTouchController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/TaskViewTouchController.java
@@ -252,7 +252,7 @@
                     mTaskBeingDragged, maxDuration, currentInterpolator);
 
             // Since the thumbnail is what is filling the screen, based the end displacement on it.
-            View thumbnailView = mTaskBeingDragged.getThumbnail();
+            View thumbnailView = mTaskBeingDragged.getFirstThumbnailView();
             mTempCords[1] = orientationHandler.getSecondaryDimension(thumbnailView);
             dl.getDescendantCoordRelativeToSelf(thumbnailView, mTempCords);
             mEndDisplacement = secondaryLayerDimension - mTempCords[1];
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index d6e4dbd..bf985af 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -139,7 +139,7 @@
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.quickstep.views.TaskView;
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer;
+import com.android.quickstep.views.TaskView.TaskContainer;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
@@ -923,7 +923,7 @@
             TaskView runningTask = mRecentsView.getRunningTaskView();
             TaskView centermostTask = mRecentsView.getTaskViewNearestToCenterOfScreen();
             int centermostTaskFlags = centermostTask == null ? 0
-                    : centermostTask.getThumbnail().getSysUiStatusNavFlags();
+                    : centermostTask.getFirstThumbnailView().getSysUiStatusNavFlags();
             boolean swipeUpThresholdPassed = windowProgress > 1 - UPDATE_SYSUI_FLAGS_THRESHOLD;
             boolean quickswitchThresholdPassed = centermostTask != runningTask;
 
@@ -2286,15 +2286,15 @@
                 int[] taskIds = nextTask.getTaskIds();
                 ActiveGestureLog.CompoundString nextTaskLog = new ActiveGestureLog.CompoundString(
                         "Launching task: ");
-                for (TaskIdAttributeContainer c : nextTask.getTaskIdAttributeContainers()) {
-                    if (c == null) {
+                for (TaskContainer container : nextTask.getTaskContainers()) {
+                    if (container == null) {
                         continue;
                     }
                     nextTaskLog
                             .append("[id: ")
-                            .append(c.getTask().key.id)
+                            .append(container.getTask().key.id)
                             .append(", pkg: ")
-                            .append(c.getTask().key.getPackageName())
+                            .append(container.getTask().key.getPackageName())
                             .append("] | ");
                 }
                 mGestureState.updateLastStartedTaskIds(taskIds);
@@ -2406,7 +2406,7 @@
         RemoteAnimationTarget taskTarget = taskTargetOptional.get();
         TaskView taskView = mRecentsView == null
                 ? null : mRecentsView.getTaskViewByTaskId(taskTarget.taskId);
-        if (taskView == null || !taskView.getThumbnail().shouldShowSplashView()) {
+        if (taskView == null || !taskView.getFirstThumbnailView().shouldShowSplashView()) {
             ActiveGestureLog.INSTANCE.addLog("Invalid task view splash state");
             finishRecentsAnimationOnTasksAppeared(null /* onFinishComplete */);
             return;
diff --git a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
index aab6aa1..e33ef7f 100644
--- a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
+++ b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
@@ -23,12 +23,12 @@
 import com.android.launcher3.popup.SystemShortcut
 import com.android.quickstep.views.RecentsView
 import com.android.quickstep.views.RecentsViewContainer
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
+import com.android.quickstep.views.TaskView.TaskContainer
 
 /** A menu item, "Desktop", that allows the user to bring the current app into Desktop Windowing. */
 class DesktopSystemShortcut(
     container: RecentsViewContainer,
-    private val mTaskContainer: TaskIdAttributeContainer,
+    private val mTaskContainer: TaskContainer,
     abstractFloatingViewHelper: AbstractFloatingViewHelper
 ) :
     SystemShortcut<RecentsViewContainer>(
@@ -59,7 +59,7 @@
             return object : TaskShortcutFactory {
                 override fun getShortcuts(
                     container: RecentsViewContainer,
-                    taskContainer: TaskIdAttributeContainer
+                    taskContainer: TaskContainer
                 ): List<DesktopSystemShortcut>? {
                     return if (!DesktopModeStatus.canEnterDesktopMode(container.asContext())) null
                     else if (!taskContainer.task.isDockable) null
diff --git a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
index 4b5a15d..6a9f509 100644
--- a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
+++ b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
@@ -305,8 +305,8 @@
             return null;
         }
         if (sourceTaskView == null
-                || sourceTaskView.getTask() == null
-                || sourceTaskView.getTask().key.getComponent() == null) {
+                || sourceTaskView.getFirstTask() == null
+                || sourceTaskView.getFirstTask().key.getComponent() == null) {
             // Disable if it's an invalid task
             return null;
         }
@@ -323,8 +323,8 @@
         }
 
         return mContainer.getFirstMatchForAppClose(launchCookieItemId,
-                sourceTaskView.getTask().key.getComponent().getPackageName(),
-                UserHandle.of(sourceTaskView.getTask().key.userId),
+                sourceTaskView.getFirstTask().key.getComponent().getPackageName(),
+                UserHandle.of(sourceTaskView.getFirstTask().key.userId),
                 false /* supportsAllAppsState */);
     }
 
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index d32c7a6..1a46fb6 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -46,7 +46,7 @@
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.quickstep.views.TaskThumbnailViewDeprecated;
 import com.android.quickstep.views.TaskView;
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer;
+import com.android.quickstep.views.TaskView.TaskContainer;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 
@@ -59,7 +59,7 @@
 public class TaskOverlayFactory implements ResourceBasedOverride {
 
     public static List<SystemShortcut> getEnabledShortcuts(TaskView taskView,
-            TaskIdAttributeContainer taskContainer) {
+            TaskContainer taskContainer) {
         final ArrayList<SystemShortcut> shortcuts = new ArrayList<>();
         final RecentsViewContainer container =
                 RecentsViewContainer.containerFromContext(taskView.getContext());
@@ -292,7 +292,7 @@
 
             @Override
             public void onClick(View view) {
-                saveScreenshot(mThumbnailView.getTaskView().getTask());
+                saveScreenshot(mThumbnailView.getTaskView().getFirstTask());
                 dismissTaskMenuView();
             }
         }
diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index 8df4bdd..cd9df26 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -58,7 +58,7 @@
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.quickstep.views.TaskThumbnailViewDeprecated;
 import com.android.quickstep.views.TaskView;
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer;
+import com.android.quickstep.views.TaskView.TaskContainer;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.view.AppTransitionAnimationSpecCompat;
 import com.android.systemui.shared.recents.view.AppTransitionAnimationSpecsFuture;
@@ -78,7 +78,7 @@
 public interface TaskShortcutFactory {
     @Nullable
     default List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-            TaskIdAttributeContainer taskContainer) {
+            TaskContainer taskContainer) {
         return null;
     }
 
@@ -108,7 +108,7 @@
     TaskShortcutFactory APP_INFO = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             TaskView taskView = taskContainer.getTaskView();
             AppInfo.SplitAccessibilityInfo accessibilityInfo =
                     new AppInfo.SplitAccessibilityInfo(taskView.containsMultipleTasks(),
@@ -176,7 +176,7 @@
         private final LauncherEvent mLauncherEvent;
 
         public FreeformSystemShortcut(int iconRes, int textRes, RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer, LauncherEvent launcherEvent) {
+                TaskContainer taskContainer, LauncherEvent launcherEvent) {
             super(iconRes, textRes, container, taskContainer.getItemInfo(),
                     taskContainer.getTaskView());
             mLauncherEvent = launcherEvent;
@@ -199,7 +199,7 @@
         }
 
         private void startActivity() {
-            final Task.TaskKey taskKey = mTaskView.getTask().key;
+            final Task.TaskKey taskKey = mTaskView.getFirstTask().key;
             final int taskId = taskKey.id;
             final ActivityOptions options = makeLaunchOptions(mTarget);
             if (options != null) {
@@ -292,7 +292,7 @@
     TaskShortcutFactory SPLIT_SELECT = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             DeviceProfile deviceProfile = container.getDeviceProfile();
             final Task task = taskContainer.getTask();
             final int intentFlags = task.key.baseIntent.getFlags();
@@ -327,7 +327,7 @@
         @Nullable
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             DeviceProfile deviceProfile = container.getDeviceProfile();
             final TaskView taskView = taskContainer.getTaskView();
             final RecentsView recentsView = taskView.getRecentsView();
@@ -336,7 +336,7 @@
                     recentsView.isTaskInExpectedScrollPosition(recentsView.indexOfChild(taskView));
             boolean shouldShowActionsButtonInstead =
                     isLargeTileFocusedTask && isInExpectedScrollPosition;
-            boolean hasUnpinnableApp = Arrays.stream(taskView.getTaskIdAttributeContainers())
+            boolean hasUnpinnableApp = Arrays.stream(taskView.getTaskContainers())
                     .anyMatch(att -> att != null && att.getItemInfo() != null
                             && ((att.getItemInfo().runtimeStatusFlags
                                 & ItemInfoWithIcon.FLAG_NOT_PINNABLE) != 0));
@@ -375,7 +375,7 @@
     TaskShortcutFactory FREE_FORM = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             final Task task = taskContainer.getTask();
             if (!task.isDockable) {
                 return null;
@@ -401,7 +401,7 @@
     TaskShortcutFactory PIN = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             if (!SystemUiProxy.INSTANCE.get(container.asContext()).isActive()) {
                 return null;
             }
@@ -423,7 +423,7 @@
         private final TaskView mTaskView;
 
         public PinSystemShortcut(RecentsViewContainer target,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             super(R.drawable.ic_pin, R.string.recent_task_option_pin, target,
                     taskContainer.getItemInfo(), taskContainer.getTaskView());
             mTaskView = taskContainer.getTaskView();
@@ -433,7 +433,7 @@
         public void onClick(View view) {
             if (mTaskView.launchTaskAnimated() != null) {
                 SystemUiProxy.INSTANCE.get(mTarget.asContext()).startScreenPinning(
-                        mTaskView.getTask().key.id);
+                        mTaskView.getFirstTask().key.id);
             }
             dismissTaskMenuView();
             mTarget.getStatsLogManager().logger().withItemInfo(mTaskView.getItemInfo())
@@ -444,7 +444,7 @@
     TaskShortcutFactory INSTALL = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             Task t = taskContainer.getTask();
             return InstantAppResolver.newInstance(container.asContext()).isInstantApp(
                     t.getTopComponent().getPackageName(), t.getKey().userId)
@@ -457,7 +457,7 @@
     TaskShortcutFactory WELLBEING = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             SystemShortcut<ActivityContext> wellbeingShortcut =
                     WellbeingModel.SHORTCUT_FACTORY.getShortcut(container,
                             taskContainer.getItemInfo(), taskContainer.getTaskView());
@@ -468,7 +468,7 @@
     TaskShortcutFactory SCREENSHOT = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             boolean isTablet = container.getDeviceProfile().isTablet;
             boolean isGridOnlyOverview = isTablet && Flags.enableGridOnlyOverview();
             // Extra conditions if it's not grid-only overview
@@ -498,7 +498,7 @@
     TaskShortcutFactory MODAL = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(RecentsViewContainer container,
-                TaskIdAttributeContainer taskContainer) {
+                TaskContainer taskContainer) {
             boolean isTablet = container.getDeviceProfile().isTablet;
             boolean isGridOnlyOverview = isTablet && Flags.enableGridOnlyOverview();
             // Extra conditions if it's not grid-only overview
diff --git a/quickstep/src/com/android/quickstep/TaskViewUtils.java b/quickstep/src/com/android/quickstep/TaskViewUtils.java
index d89d399..d2560e6 100644
--- a/quickstep/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskViewUtils.java
@@ -121,7 +121,7 @@
                 for (int i = 0; i < recentsView.getTaskViewCount(); i++) {
                     TaskView taskView = recentsView.getTaskViewAt(i);
                     if (recentsView.isTaskViewVisible(taskView)) {
-                        Task.TaskKey key = taskView.getTask().key;
+                        Task.TaskKey key = taskView.getFirstTask().key;
                         if (componentName.equals(key.getComponent()) && userId == key.userId) {
                             return taskView;
                         }
@@ -334,7 +334,7 @@
             // During animation we apply transformation on the thumbnailView (and not the rootView)
             // to follow the TaskViewSimulator. So the final matrix applied on the thumbnailView is:
             //    Mt K(0)` K(t) Mt`
-            TaskThumbnailViewDeprecated[] thumbnails = v.getThumbnails();
+            TaskThumbnailViewDeprecated[] thumbnails = v.getThumbnailViews();
 
             // In case simulator copies and thumbnail size do no match, ensure we get the lesser.
             // This ensures we do not create arrays with empty elements or attempt to references
diff --git a/quickstep/src/com/android/quickstep/util/AppPairsController.java b/quickstep/src/com/android/quickstep/util/AppPairsController.java
index a82031a..5c5ee7e 100644
--- a/quickstep/src/com/android/quickstep/util/AppPairsController.java
+++ b/quickstep/src/com/android/quickstep/util/AppPairsController.java
@@ -116,9 +116,9 @@
      */
     public void saveAppPair(GroupedTaskView gtv) {
         InteractionJankMonitorWrapper.begin(gtv, Cuj.CUJ_LAUNCHER_SAVE_APP_PAIR);
-        TaskView.TaskIdAttributeContainer[] attributes = gtv.getTaskIdAttributeContainers();
-        WorkspaceItemInfo recentsInfo1 = attributes[0].getItemInfo();
-        WorkspaceItemInfo recentsInfo2 = attributes[1].getItemInfo();
+        TaskView.TaskContainer[] containers = gtv.getTaskContainers();
+        WorkspaceItemInfo recentsInfo1 = containers[0].getItemInfo();
+        WorkspaceItemInfo recentsInfo2 = containers[1].getItemInfo();
         WorkspaceItemInfo app1 = lookupLaunchableItem(recentsInfo1.getComponentKey());
         WorkspaceItemInfo app2 = lookupLaunchableItem(recentsInfo2.getComponentKey());
 
diff --git a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
index ee2c2e1..40ea70f 100644
--- a/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
+++ b/quickstep/src/com/android/quickstep/util/SplitAnimationController.kt
@@ -69,7 +69,7 @@
 import com.android.quickstep.views.SplitInstructionsView
 import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
+import com.android.quickstep.views.TaskView.TaskContainer
 import com.android.quickstep.views.TaskViewIcon
 import com.android.wm.shell.shared.TransitionUtil
 import java.util.Optional
@@ -114,7 +114,7 @@
         } else if (splitSelectStateController.isDismissingFromSplitPair) {
             // Initiating split from overview, but on a split pair
             val taskView = taskViewSupplier.get()
-            for (container: TaskIdAttributeContainer in taskView.taskIdAttributeContainers) {
+            for (container: TaskContainer in taskView.taskContainers) {
                 if (container.task.getKey().getId() == splitSelectStateController.initialTaskId) {
                     val drawable = getDrawable(container.iconView, splitSelectSource)
                     return SplitAnimInitProps(
@@ -134,15 +134,17 @@
         } else {
             // Initiating split from overview on fullscreen task TaskView
             val taskView = taskViewSupplier.get()
-            val drawable = getDrawable(taskView.iconView, splitSelectSource)
-            return SplitAnimInitProps(
-                taskView.thumbnail,
-                taskView.thumbnail.thumbnail,
-                drawable!!,
-                fadeWithThumbnail = true,
-                isStagedTask = true,
-                taskView.iconView.asView()
-            )
+            taskView.taskContainers.first().let {
+                val drawable = getDrawable(it.iconView, splitSelectSource)
+                return SplitAnimInitProps(
+                    it.thumbnailView,
+                    it.thumbnailView.thumbnail,
+                    drawable!!,
+                    fadeWithThumbnail = true,
+                    isStagedTask = true,
+                    iconView = it.iconView.asView()
+                )
+            }
         }
     }
 
@@ -173,7 +175,7 @@
      *   (opposite of that representing [taskIdAttributeContainer])
      */
     fun addInitialSplitFromPair(
-        taskIdAttributeContainer: TaskIdAttributeContainer,
+        taskIdAttributeContainer: TaskContainer,
         builder: PendingAnimation,
         deviceProfile: DeviceProfile,
         taskViewWidth: Int,
@@ -561,8 +563,13 @@
                 // Launch split app pair animation
                 composeIconSplitLaunchAnimator(launchingIconView, info, t, finishCallback)
             } else {
-                composeFullscreenIconSplitLaunchAnimator(launchingIconView, info, t,
-                        finishCallback, appPairLaunchingAppIndex)
+                composeFullscreenIconSplitLaunchAnimator(
+                    launchingIconView,
+                    info,
+                    t,
+                    finishCallback,
+                    appPairLaunchingAppIndex
+                )
             }
         } else {
             // Fallback case: simple fade-in animation
@@ -629,18 +636,22 @@
 
     /**
      * @return -1 if [transitionInfo] contains both apps of the app pair to be animated, otherwise
-     *         the integer index corresponding to [launchingIconView]'s contents for the single app
-     *         to be animated
+     *   the integer index corresponding to [launchingIconView]'s contents for the single app to be
+     *   animated
      */
-    fun hasChangesForBothAppPairs(launchingIconView: AppPairIcon,
-                                          transitionInfo: TransitionInfo) : Int {
+    fun hasChangesForBothAppPairs(
+        launchingIconView: AppPairIcon,
+        transitionInfo: TransitionInfo
+    ): Int {
         val intent1 = launchingIconView.info.getFirstApp().intent.component?.packageName
         val intent2 = launchingIconView.info.getSecondApp().intent.component?.packageName
         var launchFullscreenAppIndex = -1
         for (change in transitionInfo.changes) {
             val taskInfo: RunningTaskInfo = change.taskInfo ?: continue
-            if (TransitionUtil.isOpeningType(change.mode) &&
-                    taskInfo.windowingMode == WINDOWING_MODE_FULLSCREEN) {
+            if (
+                TransitionUtil.isOpeningType(change.mode) &&
+                    taskInfo.windowingMode == WINDOWING_MODE_FULLSCREEN
+            ) {
                 val baseIntent = taskInfo.baseIntent.component?.packageName
                 if (baseIntent == intent1) {
                     if (launchFullscreenAppIndex > -1) {
@@ -695,8 +706,12 @@
         // If launching an app pair from Taskbar inside of an app context (no access to Launcher),
         // use the scale-up animation
         if (launchingIconView.context is TaskbarActivityContext) {
-            composeScaleUpLaunchAnimation(transitionInfo, t, finishCallback,
-                    WINDOWING_MODE_MULTI_WINDOW)
+            composeScaleUpLaunchAnimation(
+                transitionInfo,
+                t,
+                finishCallback,
+                WINDOWING_MODE_MULTI_WINDOW
+            )
             return
         }
 
@@ -767,28 +782,32 @@
         floatingView.bringToFront()
 
         launchAnimation.play(
-                getIconLaunchValueAnimator(t, dp, finishCallback, launcher, floatingView,
-                        rootCandidate))
+            getIconLaunchValueAnimator(t, dp, finishCallback, launcher, floatingView, rootCandidate)
+        )
         launchAnimation.start()
     }
 
     /**
-     * Similar to [composeIconSplitLaunchAnimator], but instructs [FloatingAppPairView] to animate
-     * a single fullscreen icon + background instead of for a pair
+     * Similar to [composeIconSplitLaunchAnimator], but instructs [FloatingAppPairView] to animate a
+     * single fullscreen icon + background instead of for a pair
      */
     @VisibleForTesting
     fun composeFullscreenIconSplitLaunchAnimator(
-            launchingIconView: AppPairIcon,
-            transitionInfo: TransitionInfo,
-            t: Transaction,
-            finishCallback: Runnable,
-            launchFullscreenIndex: Int
+        launchingIconView: AppPairIcon,
+        transitionInfo: TransitionInfo,
+        t: Transaction,
+        finishCallback: Runnable,
+        launchFullscreenIndex: Int
     ) {
         // If launching an app pair from Taskbar inside of an app context (no access to Launcher),
         // use the scale-up animation
         if (launchingIconView.context is TaskbarActivityContext) {
-            composeScaleUpLaunchAnimation(transitionInfo, t, finishCallback,
-                    WINDOWING_MODE_FULLSCREEN)
+            composeScaleUpLaunchAnimation(
+                transitionInfo,
+                t,
+                finishCallback,
+                WINDOWING_MODE_FULLSCREEN
+            )
             return
         }
 
@@ -799,16 +818,18 @@
         // Create an AnimatorSet that will run both shell and launcher transitions together
         val launchAnimation = AnimatorSet()
 
-        val appInfo = launchingIconView.info
-                .getContents()[launchFullscreenIndex] as WorkspaceItemInfo
+        val appInfo =
+            launchingIconView.info.getContents()[launchFullscreenIndex] as WorkspaceItemInfo
         val intentToLaunch = appInfo.intent.component?.packageName
         var rootCandidate: Change? = null
         for (change in transitionInfo.changes) {
             val taskInfo: RunningTaskInfo = change.taskInfo ?: continue
             val baseIntent = taskInfo.baseIntent.component?.packageName
-            if (TransitionUtil.isOpeningType(change.mode) &&
+            if (
+                TransitionUtil.isOpeningType(change.mode) &&
                     taskInfo.windowingMode == WINDOWING_MODE_FULLSCREEN &&
-                    baseIntent == intentToLaunch) {
+                    baseIntent == intentToLaunch
+            ) {
                 rootCandidate = change
             }
         }
@@ -832,26 +853,28 @@
         appIcon.setBounds(0, 0, dp.iconSizePx, dp.iconSizePx)
 
         val floatingView =
-                FloatingAppPairView.getFloatingAppPairView(
-                        launcher,
-                        drawableArea,
-                        appIcon,
-                        null /*appIcon2*/,
-                        0 /*dividerPos*/
-                )
+            FloatingAppPairView.getFloatingAppPairView(
+                launcher,
+                drawableArea,
+                appIcon,
+                null /*appIcon2*/,
+                0 /*dividerPos*/
+            )
         floatingView.bringToFront()
         launchAnimation.play(
-                getIconLaunchValueAnimator(t, dp, finishCallback, launcher, floatingView,
-                        rootCandidate))
+            getIconLaunchValueAnimator(t, dp, finishCallback, launcher, floatingView, rootCandidate)
+        )
         launchAnimation.start()
     }
 
-    private fun getIconLaunchValueAnimator(t: Transaction,
-                                           dp: com.android.launcher3.DeviceProfile,
-                                           finishCallback: Runnable,
-                                           launcher: QuickstepLauncher,
-                                           floatingView: FloatingAppPairView,
-                                           rootCandidate: Change) : ValueAnimator {
+    private fun getIconLaunchValueAnimator(
+        t: Transaction,
+        dp: com.android.launcher3.DeviceProfile,
+        finishCallback: Runnable,
+        launcher: QuickstepLauncher,
+        floatingView: FloatingAppPairView,
+        rootCandidate: Change
+    ): ValueAnimator {
         val progressUpdater = ValueAnimator.ofFloat(0f, 1f)
         val timings = AnimUtils.getDeviceAppPairLaunchTimings(dp.isTablet)
         progressUpdater.setDuration(timings.getDuration().toLong())
@@ -860,12 +883,12 @@
         // Shell animation: the apps are revealed toward end of the launch animation
         progressUpdater.addUpdateListener { valueAnimator: ValueAnimator ->
             val progress =
-                    Interpolators.clampToProgress(
-                            Interpolators.LINEAR,
-                            valueAnimator.animatedFraction,
-                            timings.appRevealStartOffset,
-                            timings.appRevealEndOffset
-                    )
+                Interpolators.clampToProgress(
+                    Interpolators.LINEAR,
+                    valueAnimator.animatedFraction,
+                    timings.appRevealStartOffset,
+                    timings.appRevealEndOffset
+                )
 
             // Set the alpha of the shell layer (2 apps + divider)
             t.setAlpha(rootCandidate.leash, progress)
@@ -873,65 +896,65 @@
         }
 
         progressUpdater.addUpdateListener(
-                object : MultiValueUpdateListener() {
-                    var mDx =
-                            FloatProp(
-                                    floatingView.startingPosition.left,
-                                    dp.widthPx / 2f - floatingView.startingPosition.width() / 2f,
-                                    Interpolators.clampToProgress(
-                                            timings.getStagedRectXInterpolator(),
-                                            timings.stagedRectSlideStartOffset,
-                                            timings.stagedRectSlideEndOffset
-                                    )
-                            )
-                    var mDy =
-                            FloatProp(
-                                    floatingView.startingPosition.top,
-                                    dp.heightPx / 2f - floatingView.startingPosition.height() / 2f,
-                                    Interpolators.clampToProgress(
-                                            Interpolators.EMPHASIZED,
-                                            timings.stagedRectSlideStartOffset,
-                                            timings.stagedRectSlideEndOffset
-                                    )
-                            )
-                    var mScaleX =
-                            FloatProp(
-                                    1f /* start */,
-                                    dp.widthPx / floatingView.startingPosition.width(),
-                                    Interpolators.clampToProgress(
-                                            Interpolators.EMPHASIZED,
-                                            timings.stagedRectSlideStartOffset,
-                                            timings.stagedRectSlideEndOffset
-                                    )
-                            )
-                    var mScaleY =
-                            FloatProp(
-                                    1f /* start */,
-                                    dp.heightPx / floatingView.startingPosition.height(),
-                                    Interpolators.clampToProgress(
-                                            Interpolators.EMPHASIZED,
-                                            timings.stagedRectSlideStartOffset,
-                                            timings.stagedRectSlideEndOffset
-                                    )
-                            )
+            object : MultiValueUpdateListener() {
+                var mDx =
+                    FloatProp(
+                        floatingView.startingPosition.left,
+                        dp.widthPx / 2f - floatingView.startingPosition.width() / 2f,
+                        Interpolators.clampToProgress(
+                            timings.getStagedRectXInterpolator(),
+                            timings.stagedRectSlideStartOffset,
+                            timings.stagedRectSlideEndOffset
+                        )
+                    )
+                var mDy =
+                    FloatProp(
+                        floatingView.startingPosition.top,
+                        dp.heightPx / 2f - floatingView.startingPosition.height() / 2f,
+                        Interpolators.clampToProgress(
+                            Interpolators.EMPHASIZED,
+                            timings.stagedRectSlideStartOffset,
+                            timings.stagedRectSlideEndOffset
+                        )
+                    )
+                var mScaleX =
+                    FloatProp(
+                        1f /* start */,
+                        dp.widthPx / floatingView.startingPosition.width(),
+                        Interpolators.clampToProgress(
+                            Interpolators.EMPHASIZED,
+                            timings.stagedRectSlideStartOffset,
+                            timings.stagedRectSlideEndOffset
+                        )
+                    )
+                var mScaleY =
+                    FloatProp(
+                        1f /* start */,
+                        dp.heightPx / floatingView.startingPosition.height(),
+                        Interpolators.clampToProgress(
+                            Interpolators.EMPHASIZED,
+                            timings.stagedRectSlideStartOffset,
+                            timings.stagedRectSlideEndOffset
+                        )
+                    )
 
-                    override fun onUpdate(percent: Float, initOnly: Boolean) {
-                        floatingView.progress = percent
-                        floatingView.x = mDx.value
-                        floatingView.y = mDy.value
-                        floatingView.scaleX = mScaleX.value
-                        floatingView.scaleY = mScaleY.value
-                        floatingView.invalidate()
-                    }
+                override fun onUpdate(percent: Float, initOnly: Boolean) {
+                    floatingView.progress = percent
+                    floatingView.x = mDx.value
+                    floatingView.y = mDy.value
+                    floatingView.scaleX = mScaleX.value
+                    floatingView.scaleY = mScaleY.value
+                    floatingView.invalidate()
                 }
+            }
         )
         progressUpdater.addListener(
-                object : AnimatorListenerAdapter() {
-                    override fun onAnimationEnd(animation: Animator) {
-                        safeRemoveViewFromDragLayer(launcher, floatingView)
-                        finishCallback.run()
-                    }
+            object : AnimatorListenerAdapter() {
+                override fun onAnimationEnd(animation: Animator) {
+                    safeRemoveViewFromDragLayer(launcher, floatingView)
+                    finishCallback.run()
                 }
+            }
         )
 
         return progressUpdater
@@ -939,9 +962,8 @@
 
     /**
      * This is a scale-up-and-fade-in animation (34% to 100%) for launching an app in Overview when
-     * there is no visible associated tile to expand from.
-     * [windowingMode] helps determine whether we are looking for a split or a single fullscreen
-     * [Change]
+     * there is no visible associated tile to expand from. [windowingMode] helps determine whether
+     * we are looking for a split or a single fullscreen [Change]
      */
     @VisibleForTesting
     fun composeScaleUpLaunchAnimation(
@@ -1095,12 +1117,12 @@
         animator.setDuration(QuickstepTransitionManager.SPLIT_LAUNCH_DURATION.toLong())
         animator.addUpdateListener { valueAnimator: ValueAnimator ->
             val progress =
-                    Interpolators.clampToProgress(
-                            Interpolators.LINEAR,
-                            valueAnimator.animatedFraction,
-                            0.8f,
-                            1f
-                    )
+                Interpolators.clampToProgress(
+                    Interpolators.LINEAR,
+                    valueAnimator.animatedFraction,
+                    0.8f,
+                    1f
+                )
             for (leash in openingTargets) {
                 animTransaction.setAlpha(leash, progress)
             }
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.java b/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
index 81c889c..5968901 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.java
@@ -31,7 +31,6 @@
 import android.graphics.drawable.shapes.RoundRectShape;
 import android.util.AttributeSet;
 import android.util.Log;
-import android.util.SparseArray;
 import android.view.View;
 import android.widget.FrameLayout;
 
@@ -70,14 +69,8 @@
 
     private static final boolean DEBUG = false;
 
-    @NonNull
-    private List<Task> mTasks = new ArrayList<>();
-
     private final ArrayList<TaskThumbnailViewDeprecated> mSnapshotViews = new ArrayList<>();
 
-    /** Maps {@code taskIds} to corresponding {@link TaskThumbnailViewDeprecated}s */
-    private final SparseArray<TaskThumbnailViewDeprecated> mSnapshotViewMap = new SparseArray<>();
-
     private final ArrayList<CancellableTask<?>> mPendingThumbnailRequests = new ArrayList<>();
 
     private final TaskView.FullscreenDrawParams mSnapshotDrawParams;
@@ -146,8 +139,7 @@
             mContainer.getDragLayer().getDescendantRectRelativeToSelf(mBackgroundView, bounds);
         } else {
             bounds.set(mBackgroundView.getLeft(), mBackgroundView.getTop(),
-                    mBackgroundView.getRight(),
-                    mBackgroundView.getBottom());
+                    mBackgroundView.getRight(), mBackgroundView.getBottom());
         }
         return Unit.INSTANCE;
     }
@@ -171,20 +163,19 @@
         }
         cancelPendingLoadTasks();
 
-        mTasks = new ArrayList<>(tasks);
-        mSnapshotViewMap.clear();
+        mTaskContainers = new TaskContainer[tasks.size()];
 
         // Ensure there are equal number of snapshot views and tasks.
         // More tasks than views, add views. More views than tasks, remove views.
         // TODO(b/251586230): use a ViewPool for creating TaskThumbnailViews
-        if (mSnapshotViews.size() > mTasks.size()) {
-            int diff = mSnapshotViews.size() - mTasks.size();
+        if (mSnapshotViews.size() > tasks.size()) {
+            int diff = mSnapshotViews.size() - tasks.size();
             for (int i = 0; i < diff; i++) {
                 TaskThumbnailViewDeprecated snapshotView = mSnapshotViews.remove(0);
                 removeView(snapshotView);
             }
-        } else if (mSnapshotViews.size() < mTasks.size()) {
-            int diff = mTasks.size() - mSnapshotViews.size();
+        } else if (mSnapshotViews.size() < tasks.size()) {
+            int diff = tasks.size() - mSnapshotViews.size();
             for (int i = 0; i < diff; i++) {
                 TaskThumbnailViewDeprecated snapshotView =
                         new TaskThumbnailViewDeprecated(getContext());
@@ -195,57 +186,20 @@
             }
         }
 
-        for (int i = 0; i < mTasks.size(); i++) {
-            Task task = mTasks.get(i);
-            TaskThumbnailViewDeprecated snapshotView = mSnapshotViews.get(i);
-            snapshotView.bind(task);
-            mSnapshotViewMap.put(task.key.id, snapshotView);
+        for (int i = 0; i < tasks.size(); i++) {
+            Task task = tasks.get(i);
+            TaskThumbnailViewDeprecated thumbnailView = mSnapshotViews.get(i);
+            thumbnailView.bind(task);
+            mTaskContainers[i] = createAttributeContainer(task, thumbnailView);
         }
 
-        updateTaskIdContainer();
-        updateTaskIdAttributeContainer();
-
         setOrientationState(orientedState);
     }
 
-    private void updateTaskIdContainer() {
-        mTaskIdContainer = new int[mTasks.size()];
-        for (int i = 0; i < mTasks.size(); i++) {
-            mTaskIdContainer[i] = mTasks.get(i).key.id;
-        }
-    }
-
-    private void updateTaskIdAttributeContainer() {
-        mTaskIdAttributeContainer = new TaskIdAttributeContainer[mTasks.size()];
-        for (int i = 0; i < mTasks.size(); i++) {
-            Task task = mTasks.get(i);
-            TaskThumbnailViewDeprecated thumbnailView = mSnapshotViewMap.get(task.key.id);
-            mTaskIdAttributeContainer[i] = createAttributeContainer(task, thumbnailView);
-        }
-    }
-
-    private TaskIdAttributeContainer createAttributeContainer(Task task,
+    private TaskContainer createAttributeContainer(Task task,
             TaskThumbnailViewDeprecated thumbnailView) {
-        return new TaskIdAttributeContainer(task, thumbnailView, mIconView,
-                STAGE_POSITION_UNDEFINED);
-    }
-
-    @Nullable
-    @Override
-    public Task getTask() {
-        // TODO(b/249371338): returning first task. This won't work well with multiple tasks.
-        return mTasks.size() > 0 ? mTasks.get(0) : null;
-    }
-
-    @Override
-    public TaskThumbnailViewDeprecated getThumbnail() {
-        // TODO(b/249371338): returning single thumbnail. This won't work well with multiple tasks.
-        Task task = getTask();
-        if (task != null) {
-            return mSnapshotViewMap.get(task.key.id);
-        }
-        // Return the place holder snapshot views. Callers expect this to be non-null
-        return mTaskThumbnailViewDeprecated;
+        return new TaskContainer(task, thumbnailView, mIconView,
+                STAGE_POSITION_UNDEFINED, /*digitalWellBeingToast=*/ null);
     }
 
     @Override
@@ -256,15 +210,12 @@
             TaskThumbnailCache thumbnailCache = model.getThumbnailCache();
 
             if (needsUpdate(changes, FLAG_UPDATE_THUMBNAIL)) {
-                for (Task task : mTasks) {
+                for (TaskContainer container : mTaskContainers) {
                     CancellableTask<?> thumbLoadRequest =
-                            thumbnailCache.updateThumbnailInBackground(task, thumbnailData -> {
-                                TaskThumbnailViewDeprecated thumbnailView =
-                                        mSnapshotViewMap.get(task.key.id);
-                                if (thumbnailView != null) {
-                                    thumbnailView.setThumbnail(task, thumbnailData);
-                                }
-                            });
+                            thumbnailCache.updateThumbnailInBackground(container.getTask(),
+                                    thumbnailData -> container.getThumbnailView().setThumbnail(
+                                            container.getTask(),
+                                            thumbnailData));
                     if (thumbLoadRequest != null) {
                         mPendingThumbnailRequests.add(thumbLoadRequest);
                     }
@@ -272,13 +223,10 @@
             }
         } else {
             if (needsUpdate(changes, FLAG_UPDATE_THUMBNAIL)) {
-                for (Task task : mTasks) {
-                    TaskThumbnailViewDeprecated thumbnailView = mSnapshotViewMap.get(task.key.id);
-                    if (thumbnailView != null) {
-                        thumbnailView.setThumbnail(null, null);
-                    }
+                for (TaskContainer container : mTaskContainers) {
+                    container.getThumbnailView().setThumbnail(null, null);
                     // Reset the task thumbnail ref
-                    task.thumbnail = null;
+                    container.getTask().thumbnail = null;
                 }
             }
         }
@@ -329,35 +277,17 @@
     @Override
     void refreshThumbnails(@Nullable HashMap<Integer, ThumbnailData> thumbnailDatas) {
         // Sets new thumbnails based on the incoming data and refreshes the rest.
-        // Create a copy of the thumbnail map, so we can track thumbnails that need refreshing.
-        SparseArray<TaskThumbnailViewDeprecated> thumbnailsToRefresh = mSnapshotViewMap.clone();
         if (thumbnailDatas != null) {
-            for (Task task : mTasks) {
-                int key = task.key.id;
-                TaskThumbnailViewDeprecated thumbnailView = thumbnailsToRefresh.get(key);
-                ThumbnailData thumbnailData = thumbnailDatas.get(key);
-                if (thumbnailView != null && thumbnailData != null) {
-                    thumbnailView.setThumbnail(task, thumbnailData);
-                    // Remove this thumbnail from the list that should be refreshed.
-                    thumbnailsToRefresh.remove(key);
+            for (TaskContainer container : mTaskContainers) {
+                ThumbnailData thumbnailData = thumbnailDatas.get(container.getTask().key.id);
+                if (thumbnailData != null) {
+                    container.getThumbnailView().setThumbnail(container.getTask(), thumbnailData);
+                } else {
+                    // Refresh the rest that were not updated.
+                    container.getThumbnailView().refresh();
                 }
             }
         }
-
-        // Refresh the rest that were not updated.
-        for (int i = 0; i < thumbnailsToRefresh.size(); i++) {
-            thumbnailsToRefresh.valueAt(i).refresh();
-        }
-    }
-
-    @Override
-    public TaskThumbnailViewDeprecated[] getThumbnails() {
-        TaskThumbnailViewDeprecated[] thumbnails =
-                new TaskThumbnailViewDeprecated[mSnapshotViewMap.size()];
-        for (int i = 0; i < thumbnails.length; i++) {
-            thumbnails[i] = mSnapshotViewMap.valueAt(i);
-        }
-        return thumbnails;
     }
 
     @Override
@@ -365,11 +295,8 @@
         resetPersistentViewTransforms();
         // Clear any references to the thumbnail (it will be re-read either from the cache or the
         // system on next bind)
-        for (Task task : mTasks) {
-            TaskThumbnailViewDeprecated thumbnailView = mSnapshotViewMap.get(task.key.id);
-            if (thumbnailView != null) {
-                thumbnailView.setThumbnail(task, null);
-            }
+        for (TaskContainer container : mTaskContainers) {
+            container.getThumbnailView().setThumbnail(container.getTask(), null);
         }
         setOverlayEnabled(false);
         onTaskListVisibilityChanged(false);
@@ -387,8 +314,7 @@
         int thumbnailTopMarginPx = mContainer.getDeviceProfile().overviewTaskThumbnailTopMarginPx;
         containerHeight -= thumbnailTopMarginPx;
 
-        int thumbnails = mSnapshotViewMap.size();
-        if (thumbnails == 0) {
+        if (mTaskContainers.length == 0) {
             return;
         }
 
@@ -408,8 +334,8 @@
         }
 
         // Desktop tile is a shrunk down version of launcher and freeform task thumbnails.
-        for (int i = 0; i < mTasks.size(); i++) {
-            Task task = mTasks.get(i);
+        for (TaskContainer container : mTaskContainers) {
+            Task task = container.getTask();
             Rect taskSize = task.appBounds;
             if (taskSize == null) {
                 // Default to quarter of the desktop if we did not get app bounds.
@@ -419,7 +345,7 @@
             int thumbWidth = (int) (taskSize.width() * scaleWidth);
             int thumbHeight = (int) (taskSize.height() * scaleHeight);
 
-            TaskThumbnailViewDeprecated thumbnailView = mSnapshotViewMap.get(task.key.id);
+            TaskThumbnailViewDeprecated thumbnailView = container.getThumbnailView();
             if (thumbnailView != null) {
                 thumbnailView.measure(MeasureSpec.makeMeasureSpec(thumbWidth, MeasureSpec.EXACTLY),
                         MeasureSpec.makeMeasureSpec(thumbHeight, MeasureSpec.EXACTLY));
@@ -461,9 +387,8 @@
         } else {
             mBackgroundView.setVisibility(VISIBLE);
         }
-        for (int i = 0; i < mSnapshotViewMap.size(); i++) {
-            TaskThumbnailViewDeprecated thumbnailView = mSnapshotViewMap.valueAt(i);
-            thumbnailView.getTaskOverlay().setFullscreenProgress(progress);
+        for (TaskContainer container : mTaskContainers) {
+            container.getThumbnailView().getTaskOverlay().setFullscreenProgress(progress);
         }
         // Animate icons and DWB banners in/out, except in QuickSwitch state, when tiles are
         // oversized and banner would look disproportionately large.
@@ -477,33 +402,30 @@
     @Override
     protected void updateSnapshotRadius() {
         super.updateSnapshotRadius();
-        for (int i = 0; i < mSnapshotViewMap.size(); i++) {
-            if (i == 0) {
-                // All snapshots share the same params. Only update it with the first snapshot.
-                updateFullscreenParams(mSnapshotDrawParams);
-            }
-            mSnapshotViewMap.valueAt(i).setFullscreenParams(mSnapshotDrawParams);
+        updateFullscreenParams(mSnapshotDrawParams);
+        for (TaskContainer container : mTaskContainers) {
+            container.getThumbnailView().setFullscreenParams(mSnapshotDrawParams);
         }
     }
 
     @Override
     public void setColorTint(float amount, int tintColor) {
-        for (int i = 0; i < mSnapshotViewMap.size(); i++) {
-            mSnapshotViewMap.valueAt(i).setDimAlpha(amount);
+        for (TaskContainer container : mTaskContainers) {
+            container.getThumbnailView().setDimAlpha(amount);
         }
     }
 
     @Override
     protected void applyThumbnailSplashAlpha() {
-        for (int i = 0; i < mSnapshotViewMap.size(); i++) {
-            mSnapshotViewMap.valueAt(i).setSplashAlpha(mTaskThumbnailSplashAlpha);
+        for (TaskContainer container : mTaskContainers) {
+            container.getThumbnailView().setSplashAlpha(mTaskThumbnailSplashAlpha);
         }
     }
 
     @Override
     void setThumbnailVisibility(int visibility, int taskId) {
-        for (int i = 0; i < mSnapshotViewMap.size(); i++) {
-            mSnapshotViewMap.valueAt(i).setVisibility(visibility);
+        for (TaskContainer container : mTaskContainers) {
+            container.getThumbnailView().setVisibility(visibility);
         }
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
index 82ba30b..f92b9dd 100644
--- a/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
+++ b/quickstep/src/com/android/quickstep/views/DigitalWellBeingToast.java
@@ -320,12 +320,12 @@
                 (FrameLayout.LayoutParams) mBanner.getLayoutParams();
         DeviceProfile deviceProfile = mContainer.getDeviceProfile();
         layoutParams.bottomMargin = ((ViewGroup.MarginLayoutParams)
-                mTaskView.getThumbnail().getLayoutParams()).bottomMargin;
+                mTaskView.getFirstThumbnailView().getLayoutParams()).bottomMargin;
         RecentsPagedOrientationHandler orientationHandler = mTaskView.getPagedOrientationHandler();
         Pair<Float, Float> translations = orientationHandler
                 .getDwbLayoutTranslations(mTaskView.getMeasuredWidth(),
                         mTaskView.getMeasuredHeight(), mSplitBounds, deviceProfile,
-                        mTaskView.getThumbnails(), mTask.key.id, mBanner);
+                        mTaskView.getThumbnailViews(), mTask.key.id, mBanner);
         mSplitOffsetTranslationX = translations.first;
         mSplitOffsetTranslationY = translations.second;
         updateTranslationY();
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index b2a8503..c26fc0c3 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -64,8 +64,6 @@
 public class GroupedTaskView extends TaskView {
 
     private static final String TAG = GroupedTaskView.class.getSimpleName();
-    @Nullable
-    private Task mSecondaryTask;
     // TODO(b/336612373): Support new TTV for GroupedTaskView
     private TaskThumbnailViewDeprecated mSnapshotView2;
     private TaskViewIcon mIconView2;
@@ -92,6 +90,17 @@
         mDigitalWellBeingToast2 = new DigitalWellBeingToast(mContainer, this);
     }
 
+    /**
+     * Returns the second task bound to this TaskView.
+     *
+     * @deprecated Use {@link #mTaskContainers} instead.
+     */
+    @Deprecated
+    @Nullable
+    private Task getSecondTask() {
+        return mTaskContainers.length > 1 ? mTaskContainers[1].getTask() : null;
+    }
+
     @Override
     public Unit getThumbnailBounds(@NonNull Rect bounds, boolean relativeToDragLayer) {
         if (mSplitBoundsConfig == null) {
@@ -141,13 +150,11 @@
     public void bind(Task primary, Task secondary, RecentsOrientedState orientedState,
             @Nullable SplitBounds splitBoundsConfig) {
         super.bind(primary, orientedState);
-        mSecondaryTask = secondary;
-        mTaskIdContainer = new int[]{mTaskIdContainer[0], secondary.key.id};
-        mTaskIdAttributeContainer = new TaskIdAttributeContainer[]{
-                mTaskIdAttributeContainer[0],
-                new TaskIdAttributeContainer(secondary, mSnapshotView2,
-                        mIconView2, STAGE_POSITION_BOTTOM_OR_RIGHT)};
-        mTaskIdAttributeContainer[0].setStagePosition(
+        mTaskContainers = new TaskContainer[]{
+                mTaskContainers[0],
+                new TaskContainer(secondary, findViewById(R.id.bottomright_snapshot),
+                        mIconView2, STAGE_POSITION_BOTTOM_OR_RIGHT, mDigitalWellBeingToast2)};
+        mTaskContainers[0].setStagePosition(
                 SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT);
         mSnapshotView2.bind(secondary);
         mSplitBoundsConfig = splitBoundsConfig;
@@ -169,12 +176,12 @@
     public void setUpShowAllInstancesListener() {
         // sets up the listener for the left/top task
         super.setUpShowAllInstancesListener();
-        if (mTaskIdAttributeContainer.length < 2) {
+        if (mTaskContainers.length < 2) {
             return;
         }
 
         // right/bottom task's base package name
-        String taskPackageName = mTaskIdAttributeContainer[1].getTask().key.getPackageName();
+        String taskPackageName = mTaskContainers[1].getTask().key.getPackageName();
 
         // icon of the right/bottom task
         View showWindowsView = findViewById(R.id.show_windows_right);
@@ -190,20 +197,21 @@
             TaskIconCache iconCache = model.getIconCache();
 
             if (needsUpdate(changes, FLAG_UPDATE_THUMBNAIL)) {
-                mThumbnailLoadRequest2 = thumbnailCache.updateThumbnailInBackground(mSecondaryTask,
-                        thumbnailData -> mSnapshotView2.setThumbnail(
-                                mSecondaryTask, thumbnailData
+                mThumbnailLoadRequest2 = thumbnailCache.updateThumbnailInBackground(
+                        getSecondTask(),
+                        thumbnailData -> mSnapshotView2.setThumbnail(getSecondTask(),
+                                thumbnailData
                         ));
             }
 
             if (needsUpdate(changes, FLAG_UPDATE_ICON)) {
-                mIconLoadRequest2 = iconCache.updateIconInBackground(mSecondaryTask,
+                mIconLoadRequest2 = iconCache.updateIconInBackground(getSecondTask(),
                         (task) -> {
                             setIcon(mIconView2, task.icon);
                             if (enableOverviewIconMenu()) {
                                 setText(mIconView2, task.title);
                             }
-                            mDigitalWellBeingToast2.initialize(mSecondaryTask);
+                            mDigitalWellBeingToast2.initialize(getSecondTask());
                             mDigitalWellBeingToast2.setSplitConfiguration(mSplitBoundsConfig);
                             mDigitalWellBeingToast.setSplitConfiguration(mSplitBoundsConfig);
                         });
@@ -213,7 +221,7 @@
                 mSnapshotView2.setThumbnail(null, null);
                 // Reset the task thumbnail reference as well (it will be fetched from the cache or
                 // reloaded next time we need it)
-                mSecondaryTask.thumbnail = null;
+                getSecondTask().thumbnail = null;
             }
             if (needsUpdate(changes, FLAG_UPDATE_ICON)) {
                 setIcon(mIconView2, null);
@@ -271,7 +279,7 @@
     @Nullable
     @Override
     public RunnableList launchTaskAnimated() {
-        if (mTask == null || mSecondaryTask == null) {
+        if (mTaskContainers.length == 0) {
             return null;
         }
 
@@ -304,8 +312,8 @@
     private void launchTaskInternal(@NonNull Consumer<Boolean> callback, boolean isQuickswitch,
             boolean launchingExistingTaskView) {
         getRecentsView().getSplitSelectController().launchExistingSplitPair(
-                launchingExistingTaskView ? this : null, mTask.key.id,
-                mSecondaryTask.key.id, SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT,
+                launchingExistingTaskView ? this : null, getFirstTask().key.id,
+                getSecondTask().key.id, SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT,
                 callback, isQuickswitch, getSnapPosition());
         Log.d(TAG, "launchTaskInternal - launchExistingSplitPair: " + Arrays.toString(
                 getTaskIds()));
@@ -314,10 +322,10 @@
     @Override
     void refreshThumbnails(@Nullable HashMap<Integer, ThumbnailData> thumbnailDatas) {
         super.refreshThumbnails(thumbnailDatas);
-        if (mSecondaryTask != null && thumbnailDatas != null) {
-            final ThumbnailData thumbnailData = thumbnailDatas.get(mSecondaryTask.key.id);
+        if (getSecondTask() != null && thumbnailDatas != null) {
+            final ThumbnailData thumbnailData = thumbnailDatas.get(getSecondTask().key.id);
             if (thumbnailData != null) {
-                mSnapshotView2.setThumbnail(mSecondaryTask, thumbnailData);
+                mSnapshotView2.setThumbnail(getSecondTask(), thumbnailData);
                 return;
             }
         }
@@ -325,11 +333,6 @@
         mSnapshotView2.refresh();
     }
 
-    @Override
-    public TaskThumbnailViewDeprecated[] getThumbnails() {
-        return new TaskThumbnailViewDeprecated[]{mTaskThumbnailViewDeprecated, mSnapshotView2};
-    }
-
     /**
      * Returns taskId that split selection was initiated with,
      * {@link ActivityTaskManager#INVALID_TASK_ID} if no tasks in this TaskView are part of
@@ -350,7 +353,7 @@
             // below aren't reliable since both of those views may be gone/transformed
             int initSplitTaskId = getThisTaskCurrentlyInSplitSelection();
             if (initSplitTaskId != INVALID_TASK_ID) {
-                return initSplitTaskId == mTask.key.id ? 1 : 0;
+                return initSplitTaskId == getFirstTask().key.id ? 1 : 0;
             }
         }
 
@@ -371,7 +374,7 @@
     @Override
     public void onRecycle() {
         super.onRecycle();
-        mSnapshotView2.setThumbnail(mSecondaryTask, null);
+        mSnapshotView2.setThumbnail(getSecondTask(), null);
         mSplitBoundsConfig = null;
     }
 
@@ -404,8 +407,8 @@
         } else {
             // Currently being split with this taskView, let the non-split selected thumbnail
             // take up full thumbnail area
-            Optional<TaskIdAttributeContainer> nonSplitContainer = Arrays.stream(
-                    mTaskIdAttributeContainer).filter(
+            Optional<TaskContainer> nonSplitContainer = Arrays.stream(
+                    mTaskContainers).filter(
                             container -> container.getTask().key.id != initSplitTaskId).findAny();
             nonSplitContainer.ifPresent(
                     taskIdAttributeContainer -> taskIdAttributeContainer.getThumbnailView().measure(
@@ -492,10 +495,10 @@
     }
 
     private void updateSecondaryDwbPlacement() {
-        if (mSecondaryTask == null) {
+        if (getSecondTask() == null) {
             return;
         }
-        mDigitalWellBeingToast2.initialize(mSecondaryTask);
+        mDigitalWellBeingToast2.initialize(getSecondTask());
     }
 
     @Override
@@ -548,17 +551,11 @@
      */
     @Override
     void setThumbnailVisibility(int visibility, int taskId) {
-        if (visibility == VISIBLE) {
-            mTaskThumbnailViewDeprecated.setVisibility(visibility);
-            mDigitalWellBeingToast.setBannerVisibility(visibility);
-            mSnapshotView2.setVisibility(visibility);
-            mDigitalWellBeingToast2.setBannerVisibility(visibility);
-        } else if (mTaskIdContainer.length > 0 && mTaskIdContainer[0] == taskId) {
-            mTaskThumbnailViewDeprecated.setVisibility(visibility);
-            mDigitalWellBeingToast.setBannerVisibility(visibility);
-        } else {
-            mSnapshotView2.setVisibility(visibility);
-            mDigitalWellBeingToast2.setBannerVisibility(visibility);
+        for (TaskContainer container : mTaskContainers) {
+            if (visibility == VISIBLE || container.getTask().key.id == taskId) {
+                container.getThumbnailView().setVisibility(visibility);
+                container.getDigitalWellBeingToast().setBannerVisibility(visibility);
+            }
         }
     }
 }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 077cd1b..075f159 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -204,7 +204,7 @@
 import com.android.quickstep.util.TaskVisualsChangeListener;
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.util.VibrationConstants;
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer;
+import com.android.quickstep.views.TaskView.TaskContainer;
 import com.android.systemui.plugins.ResourceProvider;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -591,7 +591,7 @@
             if (taskView == null) {
                 return;
             }
-            Task.TaskKey taskKey = taskView.getTask().key;
+            Task.TaskKey taskKey = taskView.getFirstTask().key;
             UI_HELPER_EXECUTOR.execute(new CancellableTask<>(
                     () -> PackageManagerWrapper.getInstance()
                             .getActivityInfo(taskKey.getComponent(), taskKey.userId) == null,
@@ -991,8 +991,7 @@
         if (mHandleTaskStackChanges) {
             TaskView taskView = getTaskViewByTaskId(taskId);
             if (taskView != null) {
-                for (TaskIdAttributeContainer container :
-                        taskView.getTaskIdAttributeContainers()) {
+                for (TaskContainer container : taskView.getTaskContainers()) {
                     if (container == null || taskId != container.getTask().key.id) {
                         continue;
                     }
@@ -1007,11 +1006,12 @@
     public void onTaskIconChanged(String pkg, UserHandle user) {
         for (int i = 0; i < getTaskViewCount(); i++) {
             TaskView tv = requireTaskViewAt(i);
-            Task task = tv.getTask();
+            Task task = tv.getFirstTask();
             if (task != null && task.key != null && pkg.equals(task.key.getPackageName())
                     && task.key.userId == user.getIdentifier()) {
                 task.icon = null;
-                if (tv.getIconView().getDrawable() != null) {
+                TaskViewIcon firstIconView = tv.getFirstIconView();
+                if (firstIconView != null && firstIconView.getDrawable() != null) {
                     tv.onTaskListVisibilityChanged(true /* visible */);
                 }
             }
@@ -1042,7 +1042,7 @@
                 continue;
             }
             // taskView could be a GroupedTaskView, so select the relevant task by ID
-            TaskIdAttributeContainer taskAttributes = taskView.getTaskAttributesById(id);
+            TaskContainer taskAttributes = taskView.getTaskContainerById(id);
             if (taskAttributes == null) {
                 continue;
             }
@@ -1731,7 +1731,7 @@
 
         int[] currentTaskIds;
         TaskView currentTaskView = getTaskViewAt(mCurrentPage);
-        if (currentTaskView != null && currentTaskView.getTask() != null) {
+        if (currentTaskView != null && currentTaskView.getFirstTask() != null) {
             currentTaskIds = currentTaskView.getTaskIds();
         } else {
             currentTaskIds = new int[0];
@@ -2354,7 +2354,7 @@
         // Update the task data for the in/visible children
         for (int i = 0; i < getTaskViewCount(); i++) {
             TaskView taskView = requireTaskViewAt(i);
-            TaskIdAttributeContainer[] containers = taskView.getTaskIdAttributeContainers();
+            TaskContainer[] containers = taskView.getTaskContainers();
             if (containers.length == 0) {
                 continue;
             }
@@ -2368,7 +2368,7 @@
             if (visible) {
                 // Default update all non-null tasks, then remove running ones
                 List<Task> tasksToUpdate = Arrays.stream(containers).filter(Objects::nonNull)
-                        .map(TaskIdAttributeContainer::getTask)
+                        .map(TaskContainer::getTask)
                         .collect(Collectors.toCollection(ArrayList::new));
                 if (mTmpRunningTasks != null) {
                     for (Task t : mTmpRunningTasks) {
@@ -2393,7 +2393,7 @@
                     mHasVisibleTaskData.put(task.key.id, visible);
                 }
             } else {
-                for (TaskIdAttributeContainer container : containers) {
+                for (TaskContainer container : containers) {
                     if (container == null) {
                         continue;
                     }
@@ -3816,7 +3816,7 @@
 
                 if (success) {
                     if (shouldRemoveTask) {
-                        if (dismissedTaskView.getTask() != null) {
+                        if (dismissedTaskView.getFirstTask() != null) {
                             if (dismissedTaskView.isRunningTask()) {
                                 finishRecentsAnimation(true /* toRecents */, false /* shouldPip */,
                                         () -> removeTaskInternal(dismissedTaskViewId));
@@ -4730,7 +4730,7 @@
      */
     public void resetModalVisuals() {
         if (mSelectedTask != null) {
-            mSelectedTask.getThumbnail().getTaskOverlay().resetModalVisuals();
+            mSelectedTask.getFirstThumbnailView().getTaskOverlay().resetModalVisuals();
         }
     }
 
@@ -4749,7 +4749,7 @@
             StatsLogManager.EventEnum splitEvent) {
         mSplitHiddenTaskView = taskView;
         mSplitSelectStateController.setInitialTaskSelect(null /*intent*/,
-                stagePosition, taskView.getItemInfo(), splitEvent, taskView.mTask.key.id);
+                stagePosition, taskView.getItemInfo(), splitEvent, taskView.getFirstTask().key.id);
         mSplitSelectStateController.setAnimateCurrentTaskDismissal(
                 true /*animateCurrentTaskDismissal*/);
         mSplitHiddenTaskViewIndex = indexOfChild(taskView);
@@ -4800,11 +4800,11 @@
             // Animate pair thumbnail into full thumbnail
             boolean primaryTaskSelected = mSplitHiddenTaskView.getTaskIds()[0]
                     == mSplitSelectStateController.getInitialTaskId();
-            TaskIdAttributeContainer taskIdAttributeContainer = mSplitHiddenTaskView
-                    .getTaskIdAttributeContainers()[primaryTaskSelected ? 1 : 0];
-            TaskThumbnailViewDeprecated thumbnail = taskIdAttributeContainer.getThumbnailView();
+            TaskContainer taskContainer = mSplitHiddenTaskView
+                    .getTaskContainers()[primaryTaskSelected ? 1 : 0];
+            TaskThumbnailViewDeprecated thumbnail = taskContainer.getThumbnailView();
             mSplitSelectStateController.getSplitAnimationController()
-                    .addInitialSplitFromPair(taskIdAttributeContainer, builder,
+                    .addInitialSplitFromPair(taskContainer, builder,
                             mContainer.getDeviceProfile(),
                             mSplitHiddenTaskView.getWidth(), mSplitHiddenTaskView.getHeight(),
                             primaryTaskSelected);
@@ -5202,7 +5202,7 @@
         updateGridProperties();
         updateScrollSynchronously();
 
-        int targetSysUiFlags = tv.getThumbnail().getSysUiStatusNavFlags();
+        int targetSysUiFlags = tv.getFirstThumbnailView().getSysUiStatusNavFlags();
         final boolean[] passedOverviewThreshold = new boolean[] {false};
         ValueAnimator progressAnim = ValueAnimator.ofFloat(0, 1);
         progressAnim.addUpdateListener(animator -> {
@@ -5266,7 +5266,7 @@
                 } else {
                     tv.launchTask(this::onTaskLaunchAnimationEnd);
                 }
-                Task task = tv.getTask();
+                Task task = tv.getFirstTask();
                 if (task != null) {
                     mContainer.getStatsLogManager().logger().withItemInfo(tv.getItemInfo())
                             .log(LAUNCHER_TASK_LAUNCH_SWIPE_DOWN);
@@ -5899,7 +5899,7 @@
         }
 
         taskView.setShowScreenshot(true);
-        for (TaskIdAttributeContainer container : taskView.getTaskIdAttributeContainers()) {
+        for (TaskContainer container : taskView.getTaskContainers()) {
             if (container == null) {
                 continue;
             }
@@ -6269,7 +6269,7 @@
     /**
      * Moves the provided task into desktop mode, and invoke {@code successCallback} if succeeded.
      */
-    public void moveTaskToDesktop(TaskIdAttributeContainer taskContainer,
+    public void moveTaskToDesktop(TaskContainer taskContainer,
             Runnable successCallback) {
         if (!DesktopModeStatus.canEnterDesktopMode(mContext)) {
             return;
@@ -6278,7 +6278,7 @@
                 () -> moveTaskToDesktopInternal(taskContainer, successCallback)));
     }
 
-    private void moveTaskToDesktopInternal(TaskIdAttributeContainer taskContainer,
+    private void moveTaskToDesktopInternal(TaskContainer taskContainer,
             Runnable successCallback) {
         if (mDesktopRecentsTransitionController == null) {
             return;
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index 443f83c..d0bf2c2 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -56,7 +56,7 @@
 import com.android.quickstep.TaskUtils;
 import com.android.quickstep.orientation.RecentsPagedOrientationHandler;
 import com.android.quickstep.util.TaskCornerRadius;
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer;
+import com.android.quickstep.views.TaskView.TaskContainer;
 
 /**
  * Contains options for a recent task when long-pressing its icon.
@@ -76,7 +76,7 @@
     private ValueAnimator mRevealAnimator;
     @Nullable private Runnable mOnClosingStartCallback;
     private TaskView mTaskView;
-    private TaskIdAttributeContainer mTaskContainer;
+    private TaskContainer mTaskContainer;
     private LinearLayout mOptionLayout;
     private float mMenuTranslationYBeforeOpen;
     private float mMenuTranslationXBeforeOpen;
@@ -163,7 +163,10 @@
         }
     }
 
-    public static boolean showForTask(TaskIdAttributeContainer taskContainer,
+    /**
+     * Show a task menu for the given taskContainer.
+     */
+    public static boolean showForTask(TaskContainer taskContainer,
             @Nullable Runnable onClosingStartCallback) {
         RecentsViewContainer container = RecentsViewContainer.containerFromContext(
                 taskContainer.getTaskView().getContext());
@@ -173,11 +176,14 @@
         return taskMenuView.populateAndShowForTask(taskContainer);
     }
 
-    public static boolean showForTask(TaskIdAttributeContainer taskContainer) {
+    /**
+     * Show a task menu for the given taskContainer.
+     */
+    public static boolean showForTask(TaskContainer taskContainer) {
         return showForTask(taskContainer, null);
     }
 
-    private boolean populateAndShowForTask(TaskIdAttributeContainer taskContainer) {
+    private boolean populateAndShowForTask(TaskContainer taskContainer) {
         if (isAttachedToWindow()) {
             return false;
         }
@@ -198,7 +204,7 @@
         return true;
     }
 
-    private void addMenuOptions(TaskIdAttributeContainer taskContainer) {
+    private void addMenuOptions(TaskContainer taskContainer) {
         if (enableOverviewIconMenu()) {
             removeView(mTaskName);
         } else {
@@ -226,7 +232,7 @@
         mOptionLayout.addView(menuOptionView);
     }
 
-    private void orientAroundTaskView(TaskIdAttributeContainer taskContainer) {
+    private void orientAroundTaskView(TaskContainer taskContainer) {
         RecentsView recentsView = mContainer.getOverviewPanel();
         RecentsPagedOrientationHandler orientationHandler =
                 recentsView.getPagedOrientationHandler();
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
index a138db0..7adc32e 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
@@ -37,14 +37,14 @@
 import com.android.launcher3.popup.SystemShortcut
 import com.android.launcher3.util.Themes
 import com.android.quickstep.TaskOverlayFactory
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
+import com.android.quickstep.views.TaskView.TaskContainer
 
 class TaskMenuViewWithArrow<T> : ArrowPopup<T> where T : RecentsViewContainer, T : Context {
     companion object {
         const val TAG = "TaskMenuViewWithArrow"
 
         fun <T> showForTask(
-            taskContainer: TaskIdAttributeContainer,
+            taskContainer: TaskContainer,
             alignedOptionIndex: Int = 0
         ): Boolean where T : RecentsViewContainer, T : Context {
             val container: RecentsViewContainer =
@@ -87,7 +87,7 @@
 
     private lateinit var taskView: TaskView
     private lateinit var optionLayout: LinearLayout
-    private lateinit var taskContainer: TaskIdAttributeContainer
+    private lateinit var taskContainer: TaskContainer
 
     private var optionMeasuredHeight = 0
     private val arrowHorizontalPadding: Int
@@ -141,7 +141,7 @@
     }
 
     private fun populateAndShowForTask(
-        taskContainer: TaskIdAttributeContainer,
+        taskContainer: TaskContainer,
         alignedOptionIndex: Int
     ): Boolean {
         if (isAttachedToWindow) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 93dd44f..4046a89 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -326,9 +326,6 @@
                 }
             };
 
-    @Nullable
-    protected Task mTask;
-    @Nullable // can be null when enableRefactorTaskThumbnail() == true
     protected TaskThumbnailViewDeprecated mTaskThumbnailViewDeprecated;
     protected TaskThumbnailView mTaskThumbnailView;
     protected TaskViewIcon mIconView;
@@ -371,9 +368,7 @@
     private float mStableAlpha = 1;
 
     private int mTaskViewId = -1;
-    protected int[] mTaskIdContainer = new int[0];
-    protected TaskIdAttributeContainer[] mTaskIdAttributeContainer =
-            new TaskIdAttributeContainer[0];
+    protected TaskContainer[] mTaskContainers = new TaskContainer[0];
 
     private boolean mShowScreenshot;
     private boolean mBorderEnabled;
@@ -501,16 +496,19 @@
     public void notifyIsRunningTaskUpdated() {
         // TODO(b/335649589): TaskView's VM will already have access to TaskThumbnailView's VM
         //  so there will be no need to access TaskThumbnailView's VM through the TaskThumbnailView
-        if (mTask != null) {
+        if (mTaskContainers.length > 0) {
             bindTaskThumbnailView();
         }
     }
 
     /**
-     * Builds proto for logging
+     * Builds proto for logging.
+     *
+     * @deprecated Use {@link #getItemInfo(Task)} instead.
      */
+    @Deprecated
     public WorkspaceItemInfo getItemInfo() {
-        return getItemInfo(mTask);
+        return getItemInfo(getFirstTask());
     }
 
     /**
@@ -687,11 +685,9 @@
      */
     public void bind(Task task, RecentsOrientedState orientedState) {
         cancelPendingLoadTasks();
-        mTask = task;
-        mTaskIdContainer = new int[]{mTask.key.id};
-        mTaskIdAttributeContainer = new TaskIdAttributeContainer[]{
-                new TaskIdAttributeContainer(task, mTaskThumbnailViewDeprecated, mIconView,
-                        STAGE_POSITION_UNDEFINED)};
+        mTaskContainers = new TaskContainer[]{
+                new TaskContainer(task, mTaskThumbnailViewDeprecated, mIconView,
+                        STAGE_POSITION_UNDEFINED, mDigitalWellBeingToast)};
         if (enableRefactorTaskThumbnail()) {
             bindTaskThumbnailView();
         } else {
@@ -703,17 +699,17 @@
     private void bindTaskThumbnailView() {
         // TODO(b/335649589): TaskView's VM will already have access to TaskThumbnailView's VM
         //  so there will be no need to access TaskThumbnailView's VM through the TaskThumbnailView
-        mTaskThumbnailView.getViewModel().bind(new TaskThumbnail(mTask, isRunningTask()));
+        mTaskThumbnailView.getViewModel().bind(new TaskThumbnail(getFirstTask(), isRunningTask()));
     }
 
     /**
      * Sets up an on-click listener and the visibility for show_windows icon on top of the task.
      */
     public void setUpShowAllInstancesListener() {
-        if (mTaskIdAttributeContainer.length == 0) {
+        if (mTaskContainers.length == 0) {
             return;
         }
-        String taskPackageName = mTaskIdAttributeContainer[0].mTask.key.getPackageName();
+        String taskPackageName = mTaskContainers[0].getTask().key.getPackageName();
 
         // icon of the top/left task
         View showWindowsView = findViewById(R.id.show_windows);
@@ -755,50 +751,62 @@
         filterView.setOnClickListener(callback);
     }
 
-    public TaskIdAttributeContainer[] getTaskIdAttributeContainers() {
-        return mTaskIdAttributeContainer;
+    /**
+     * Returns an array of all TaskContainers in the TaskView.
+     */
+    public TaskContainer[] getTaskContainers() {
+        return mTaskContainers;
     }
 
+    /**
+     * Returns the first task bound to this TaskView.
+     *
+     * @deprecated Use {@link #mTaskContainers} instead.
+     */
+    @Deprecated
     @Nullable
-    public Task getTask() {
-        return mTask;
+    public Task getFirstTask() {
+        return mTaskContainers.length > 0 ? mTaskContainers[0].getTask() : null;
     }
 
     /**
      * Check if given {@code taskId} is tracked in this view
      */
     public boolean containsTaskId(int taskId) {
-        return Arrays.stream(mTaskIdContainer).anyMatch(myTaskId -> myTaskId == taskId);
+        return getTaskContainerById(taskId) != null;
     }
 
     /**
      * Returns a copy of integer array containing taskIds of all tasks in the TaskView.
      */
     public int[] getTaskIds() {
-        return Arrays.copyOf(mTaskIdContainer, mTaskIdContainer.length);
+        return Arrays.stream(mTaskContainers).mapToInt(
+                container -> container.getTask().key.id).toArray();
     }
 
     public boolean containsMultipleTasks() {
-        return mTaskIdContainer.length > 1;
+        return mTaskContainers.length > 1;
     }
 
     /**
-     * Returns the TaskIdAttributeContainer corresponding to a given taskId, or null if the TaskView
-     * does not contain a Task with that ID.
+     * Returns the TaskContainer corresponding to a given taskId, or null if the TaskView does
+     * not contain a Task with that ID.
      */
     @Nullable
-    public TaskIdAttributeContainer getTaskAttributesById(int taskId) {
-        for (TaskIdAttributeContainer attributes : mTaskIdAttributeContainer) {
-            if (attributes.getTask().key.id == taskId) {
-                return attributes;
-            }
-        }
-
-        return null;
+    public TaskContainer getTaskContainerById(int taskId) {
+        return Arrays.stream(mTaskContainers).filter(
+                container -> container.getTask().key.id == taskId).findFirst().orElse(null);
     }
 
-    public TaskThumbnailViewDeprecated getThumbnail() {
-        return mTaskThumbnailViewDeprecated;
+    /**
+     * Returns the first thumbnailView of the TaskView.
+     *
+     * @deprecated Use {@link #mTaskContainers} instead.
+     */
+    @Deprecated
+    public TaskThumbnailViewDeprecated getFirstThumbnailView() {
+        return mTaskContainers.length > 0 ? mTaskContainers[0].getThumbnailView()
+                : mTaskThumbnailViewDeprecated;
     }
 
     void refreshThumbnails(@Nullable HashMap<Integer, ThumbnailData> thumbnailDatas) {
@@ -806,10 +814,10 @@
             // TODO(b/334825222) add thumbnail logic
             return;
         }
-        if (mTask != null && thumbnailDatas != null) {
-            final ThumbnailData thumbnailData = thumbnailDatas.get(mTask.key.id);
+        if (getFirstTask() != null && thumbnailDatas != null) {
+            final ThumbnailData thumbnailData = thumbnailDatas.get(getFirstTask().key.id);
             if (thumbnailData != null) {
-                mTaskThumbnailViewDeprecated.setThumbnail(mTask, thumbnailData);
+                mTaskThumbnailViewDeprecated.setThumbnail(getFirstTask(), thumbnailData);
                 return;
             }
         }
@@ -817,19 +825,27 @@
         mTaskThumbnailViewDeprecated.refresh();
     }
 
-    /** TODO(b/197033698) Remove all usages of above method and migrate to this one */
-    public TaskThumbnailViewDeprecated[] getThumbnails() {
-        return new TaskThumbnailViewDeprecated[]{mTaskThumbnailViewDeprecated};
+    public TaskThumbnailViewDeprecated[] getThumbnailViews() {
+        return Arrays.stream(mTaskContainers).map(
+                TaskContainer::getThumbnailView).toArray(
+                TaskThumbnailViewDeprecated[]::new);
     }
 
-    public TaskViewIcon getIconView() {
-        return mIconView;
+    /**
+     * Returns the first iconView of the TaskView.
+     *
+     * @deprecated Use {@link #mTaskContainers} instead.
+     */
+    @Deprecated
+    @Nullable
+    public TaskViewIcon getFirstIconView() {
+        return mTaskContainers.length > 0 ? mTaskContainers[0].getIconView() : null;
     }
 
     @Override
     public boolean dispatchTouchEvent(MotionEvent ev) {
         RecentsView recentsView = getRecentsView();
-        if (recentsView == null || getTask() == null) {
+        if (recentsView == null || getFirstTask() == null) {
             return false;
         }
         SplitSelectStateController splitSelectStateController =
@@ -837,7 +853,7 @@
         // Disable taps for split selection animation unless we have multiple tasks
         boolean disableTapsForSplitSelect =
                 splitSelectStateController.isSplitSelectActive()
-                        && splitSelectStateController.getInitialTaskId() == getTask().key.id
+                        && splitSelectStateController.getInitialTaskId() == getFirstTask().key.id
                         && !containsMultipleTasks();
         if (disableTapsForSplitSelect) {
             return false;
@@ -850,19 +866,21 @@
     }
 
     private void onClick(View view) {
-        if (getTask() == null) {
+        if (getFirstTask() == null) {
             Log.d("b/310064698", "onClick - task is null");
             return;
         }
         if (confirmSecondSplitSelectApp()) {
-            Log.d("b/310064698", mTask + " - onClick - split select is active");
+            Log.d("b/310064698",
+                    Arrays.toString(getTaskIds()) + " - onClick - split select is active");
             return;
         }
         RunnableList callbackList = launchTasks();
-        Log.d("b/310064698", mTask + " - onClick - callbackList: " + callbackList);
+        Log.d("b/310064698",
+                Arrays.toString(getTaskIds()) + " - onClick - callbackList: " + callbackList);
         if (callbackList != null) {
-            callbackList.add(() -> Log.d("b/310064698", Arrays.toString(
-                    getTaskIds()) + " - onClick - launchCompleted"));
+            callbackList.add(() -> Log.d("b/310064698",
+                    Arrays.toString(getTaskIds()) + " - onClick - launchCompleted"));
         }
         mContainer.getStatsLogManager().logger().withItemInfo(getItemInfo())
                 .log(LAUNCHER_TASK_LAUNCH_TAP);
@@ -874,10 +892,10 @@
      */
     protected boolean confirmSecondSplitSelectApp() {
         int index = getLastSelectedChildTaskIndex();
-        if (index >= mTaskIdAttributeContainer.length) {
+        if (index >= mTaskContainers.length) {
             return false;
         }
-        TaskIdAttributeContainer container = mTaskIdAttributeContainer[index];
+        TaskContainer container = mTaskContainers[index];
         if (container != null) {
             return getRecentsView().confirmSplitSelect(this, container.getTask(),
                     container.getIconView().getDrawable(), container.getThumbnailView(),
@@ -903,14 +921,15 @@
      */
     @Nullable
     public RunnableList launchTaskAnimated() {
-        if (mTask != null) {
+        if (getFirstTask() != null) {
             TestLogging.recordEvent(
-                    TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
+                    TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", Arrays.toString(
+                            getTaskIds()));
             ActivityOptionsWrapper opts = mContainer.getActivityLaunchOptions(this, null);
             opts.options.setLaunchDisplayId(
                     getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
             if (ActivityManagerWrapper.getInstance()
-                    .startActivityFromRecents(mTask.key, opts.options)) {
+                    .startActivityFromRecents(getFirstTask().key, opts.options)) {
                 Log.d(TAG, "launchTaskAnimated - startActivityFromRecents: " + Arrays.toString(
                         getTaskIds()));
                 ActiveGestureLog.INSTANCE.trackEvent(EXPECTING_TASK_APPEARED);
@@ -938,7 +957,7 @@
                 return null;
             }
         } else {
-            Log.d(TAG, "launchTaskAnimated - mTask is null" + Arrays.toString(getTaskIds()));
+            Log.d(TAG, "launchTaskAnimated - getTask() is null" + Arrays.toString(getTaskIds()));
             return null;
         }
     }
@@ -954,9 +973,10 @@
      * Starts the task associated with this view without any animation
      */
     public void launchTask(@NonNull Consumer<Boolean> callback, boolean isQuickswitch) {
-        if (mTask != null) {
+        if (getFirstTask() != null) {
             TestLogging.recordEvent(
-                    TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
+                    TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", Arrays.toString(
+                            getTaskIds()));
 
             TaskRemovedDuringLaunchListener failureListener = new TaskRemovedDuringLaunchListener(
                     getContext().getApplicationContext());
@@ -964,7 +984,7 @@
                 // We only listen for failures to launch in quickswitch because the during this
                 // gesture launcher is in the background state, vs other launches which are in
                 // the actual overview state
-                failureListener.register(mContainer, mTask.key.id, () -> {
+                failureListener.register(mContainer, getFirstTask().key.id, () -> {
                     notifyTaskLaunchFailed(TAG);
                     RecentsView rv = getRecentsView();
                     if (rv != null) {
@@ -997,7 +1017,7 @@
             if (!enableRefactorTaskThumbnail()) {
                 opts.setDisableStartingWindow(mTaskThumbnailViewDeprecated.shouldShowSplashView());
             }
-            Task.TaskKey key = mTask.key;
+            Task.TaskKey key = getFirstTask().key;
             UI_HELPER_EXECUTOR.execute(() -> {
                 if (!ActivityManagerWrapper.getInstance().startActivityFromRecents(key, opts)) {
                     // If the call to start activity failed, then post the result immediately,
@@ -1013,7 +1033,7 @@
             });
         } else {
             callback.accept(false);
-            Log.d(TAG, "launchTask - mTask is null" + Arrays.toString(getTaskIds()));
+            Log.d(TAG, "launchTask - getTask() is null" + Arrays.toString(getTaskIds()));
         }
     }
 
@@ -1069,7 +1089,8 @@
             anim.addListener(new AnimatorListenerAdapter() {
                 @Override
                 public void onAnimationEnd(Animator animator) {
-                    if (mTask != null && mTask.key.displayId != getRootViewDisplayId()) {
+                    if (getFirstTask() != null
+                            && getFirstTask().key.displayId != getRootViewDisplayId()) {
                         launchTaskAnimated();
                     }
                     mIsClickableAsLiveTile = true;
@@ -1110,7 +1131,7 @@
      * @param visible If this task view will be visible to the user in overview or hidden
      */
     public void onTaskListVisibilityChanged(boolean visible, @TaskDataChanges int changes) {
-        if (mTask == null) {
+        if (getFirstTask() == null) {
             return;
         }
         cancelPendingLoadTasks();
@@ -1123,15 +1144,16 @@
 
             if (needsUpdate(changes, FLAG_UPDATE_THUMBNAIL)) {
                 mThumbnailLoadRequest = thumbnailCache.updateThumbnailInBackground(
-                        mTask, thumbnail -> {
+                        getFirstTask(), thumbnail -> {
                             if (!enableRefactorTaskThumbnail()) {
                                 // TODO(b/334825222) add thumbnail state
-                                mTaskThumbnailViewDeprecated.setThumbnail(mTask, thumbnail);
+                                mTaskThumbnailViewDeprecated.setThumbnail(getFirstTask(),
+                                        thumbnail);
                             }
                         });
             }
             if (needsUpdate(changes, FLAG_UPDATE_ICON)) {
-                mIconLoadRequest = iconCache.updateIconInBackground(mTask,
+                mIconLoadRequest = iconCache.updateIconInBackground(getFirstTask(),
                         (task) -> {
                             setIcon(mIconView, task.icon);
                             if (enableOverviewIconMenu()) {
@@ -1151,7 +1173,7 @@
                 }
                 // Reset the task thumbnail reference as well (it will be fetched from the cache or
                 // reloaded next time we need it)
-                mTask.thumbnail = null;
+                getFirstTask().thumbnail = null;
             }
             if (needsUpdate(changes, FLAG_UPDATE_ICON)) {
                 setIcon(mIconView, null);
@@ -1195,8 +1217,8 @@
     }
 
     protected boolean showTaskMenuWithContainer(TaskViewIcon iconView) {
-        Optional<TaskIdAttributeContainer> menuContainer = Arrays.stream(
-                mTaskIdAttributeContainer).filter(
+        Optional<TaskContainer> menuContainer = Arrays.stream(
+                mTaskContainers).filter(
                         container -> container.getIconView() == iconView).findAny();
         if (menuContainer.isEmpty()) {
             return false;
@@ -1270,7 +1292,7 @@
         if (!enableRefactorTaskThumbnail()) {
             mTaskThumbnailViewDeprecated.getTaskOverlay().updateOrientationState(orientationState);
         }
-        mDigitalWellBeingToast.initialize(mTask);
+        mDigitalWellBeingToast.initialize(getFirstTask());
     }
 
     /**
@@ -1373,7 +1395,7 @@
         if (enableRefactorTaskThumbnail()) {
             notifyIsRunningTaskUpdated();
         } else {
-            mTaskThumbnailViewDeprecated.setThumbnail(mTask, null);
+            mTaskThumbnailViewDeprecated.setThumbnail(getFirstTask(), null);
         }
         setOverlayEnabled(false);
         onTaskListVisibilityChanged(false);
@@ -1670,7 +1692,7 @@
                         getContext().getText(R.string.accessibility_close)));
 
         final Context context = getContext();
-        for (TaskIdAttributeContainer taskContainer : mTaskIdAttributeContainer) {
+        for (TaskContainer taskContainer : mTaskContainers) {
             for (SystemShortcut s : TraceHelper.allowIpcs(
                     "TV.a11yInfo", () -> getEnabledShortcuts(this, taskContainer))) {
                 info.addAction(s.createAccessibilityAction(context));
@@ -1705,7 +1727,7 @@
             return true;
         }
 
-        for (TaskIdAttributeContainer taskContainer : mTaskIdAttributeContainer) {
+        for (TaskContainer taskContainer : mTaskContainers) {
             for (SystemShortcut s : getEnabledShortcuts(this,
                     taskContainer)) {
                 if (s.hasHandlerForAction(action)) {
@@ -1719,8 +1741,8 @@
     }
 
     @Nullable
-    public RecentsView getRecentsView() {
-        return (RecentsView) getParent();
+    public RecentsView<?, ?> getRecentsView() {
+        return (RecentsView<?, ?>) getParent();
     }
 
     RecentsPagedOrientationHandler getPagedOrientationHandler() {
@@ -1729,8 +1751,9 @@
 
     private void notifyTaskLaunchFailed(String tag) {
         String msg = "Failed to launch task";
-        if (mTask != null) {
-            msg += " (task=" + mTask.key.baseIntent + " userId=" + mTask.key.userId + ")";
+        if (getFirstTask() != null) {
+            msg += " (task=" + getFirstTask().key.baseIntent + " userId="
+                    + getFirstTask().key.userId + ")";
         }
         Log.w(tag, msg);
         Toast.makeText(getContext(), R.string.activity_not_available, LENGTH_SHORT).show();
@@ -1971,7 +1994,10 @@
         }
     }
 
-    public class TaskIdAttributeContainer {
+    /**
+     * Holder for all Task dependent information.
+     */
+    public class TaskContainer {
         private final TaskThumbnailViewDeprecated mThumbnailView;
         private final Task mTask;
         private final TaskViewIcon mIconView;
@@ -1979,15 +2005,18 @@
         private @SplitConfigurationOptions.StagePosition int mStagePosition;
         @IdRes
         private final int mA11yNodeId;
+        private final DigitalWellBeingToast mDigitalWellBeingToast;
 
-        public TaskIdAttributeContainer(Task task, TaskThumbnailViewDeprecated thumbnailView,
-                TaskViewIcon iconView, int stagePosition) {
+        public TaskContainer(Task task, TaskThumbnailViewDeprecated thumbnailView,
+                TaskViewIcon iconView, int stagePosition,
+                DigitalWellBeingToast digitalWellBeingToast) {
             this.mTask = task;
             this.mThumbnailView = thumbnailView;
             this.mIconView = iconView;
             this.mStagePosition = stagePosition;
             this.mA11yNodeId = (stagePosition == STAGE_POSITION_BOTTOM_OR_RIGHT) ?
                     R.id.split_bottomRight_appInfo : R.id.split_topLeft_appInfo;
+            this.mDigitalWellBeingToast = digitalWellBeingToast;
         }
 
         public TaskThumbnailViewDeprecated getThumbnailView() {
@@ -2021,5 +2050,9 @@
         public int getA11yNodeId() {
             return mA11yNodeId;
         }
+
+        public DigitalWellBeingToast getDigitalWellBeingToast() {
+            return mDigitalWellBeingToast;
+        }
     }
 }
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index 0f9d96c..d59aafb 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -16,12 +16,12 @@
 
 package com.android.quickstep
 
-import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
-import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
-import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import android.content.ComponentName
 import android.content.Intent
 import android.platform.test.flag.junit.SetFlagsRule
+import com.android.dx.mockito.inline.extended.ExtendedMockito.doReturn
+import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
+import com.android.dx.mockito.inline.extended.StaticMockitoSession
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingViewHelper
 import com.android.launcher3.logging.StatsLogManager
@@ -39,12 +39,12 @@
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
-import org.mockito.quality.Strictness
 import org.mockito.kotlin.any
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
+import org.mockito.quality.Strictness
 
 /** Test for DesktopSystemShortcut */
 class DesktopSystemShortcutTest {
@@ -64,15 +64,18 @@
     private lateinit var mockitoSession: StaticMockitoSession
 
     @Before
-    fun setUp(){
-        mockitoSession = mockitoSession().strictness(Strictness.LENIENT)
-                .spyStatic(DesktopModeStatus::class.java).startMocking()
+    fun setUp() {
+        mockitoSession =
+            mockitoSession()
+                .strictness(Strictness.LENIENT)
+                .spyStatic(DesktopModeStatus::class.java)
+                .startMocking()
         doReturn(true).`when` { DesktopModeStatus.enforceDeviceRestrictions() }
         doReturn(true).`when` { DesktopModeStatus.isDesktopModeSupported(any()) }
     }
 
     @After
-    fun tearDown(){
+    fun tearDown() {
         mockitoSession.finishMocking()
     }
 
@@ -85,11 +88,12 @@
                 isDockable = true
             }
         val taskContainer =
-            taskView.TaskIdAttributeContainer(
+            taskView.TaskContainer(
                 task,
                 null,
                 null,
-                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED
+                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                null
             )
 
         val shortcuts = factory.getShortcuts(launcher, taskContainer)
@@ -106,11 +110,12 @@
                 isDockable = true
             }
         val taskContainer =
-            taskView.TaskIdAttributeContainer(
+            taskView.TaskContainer(
                 task,
                 null,
                 null,
-                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED
+                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                null
             )
 
         val shortcuts = factory.getShortcuts(launcher, taskContainer)
@@ -128,11 +133,12 @@
                 isDockable = true
             }
         val taskContainer =
-            taskView.TaskIdAttributeContainer(
+            taskView.TaskContainer(
                 task,
                 null,
                 null,
-                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED
+                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                null
             )
 
         val shortcuts = factory.getShortcuts(launcher, taskContainer)
@@ -148,11 +154,12 @@
                 isDockable = false
             }
         val taskContainer =
-            taskView.TaskIdAttributeContainer(
+            taskView.TaskContainer(
                 task,
                 null,
                 null,
-                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED
+                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                null
             )
 
         val shortcuts = factory.getShortcuts(launcher, taskContainer)
@@ -168,11 +175,12 @@
                 isDockable = true
             }
         val taskContainer =
-            taskView.TaskIdAttributeContainer(
+            taskView.TaskContainer(
                 task,
                 null,
                 null,
-                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED
+                SplitConfigurationOptions.STAGE_POSITION_UNDEFINED,
+                null
             )
 
         whenever(launcher.getOverviewPanel<LauncherRecentsView>()).thenReturn(recentsView)
diff --git a/quickstep/tests/src/com/android/quickstep/TaplDigitalWellBeingToastTest.java b/quickstep/tests/src/com/android/quickstep/TaplDigitalWellBeingToastTest.java
index 4aa7cb0..37ab131 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplDigitalWellBeingToastTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplDigitalWellBeingToastTest.java
@@ -86,8 +86,8 @@
         final TaskView task = getOnceNotNull("No latest task", launcher -> getLatestTask(launcher));
 
         return getFromLauncher(launcher -> {
-            assertTrue("Latest task is not Calculator",
-                    CALCULATOR_PACKAGE.equals(task.getTask().getTopComponent().getPackageName()));
+            assertTrue("Latest task is not Calculator", CALCULATOR_PACKAGE.equals(
+                    task.getFirstTask().getTopComponent().getPackageName()));
             return task.getDigitalWellBeingToast();
         });
     }
diff --git a/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt b/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
index de98703..eb20de2 100644
--- a/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/util/SplitAnimationControllerTest.kt
@@ -35,7 +35,7 @@
 import com.android.quickstep.views.IconView
 import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
-import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
+import com.android.quickstep.views.TaskView.TaskContainer
 import com.android.systemui.shared.recents.model.Task
 import org.junit.Assert.assertEquals
 import org.junit.Before
@@ -67,7 +67,7 @@
     private val mockGroupedTaskView: GroupedTaskView = mock()
     private val mockTask: Task = mock()
     private val mockTaskKey: Task.TaskKey = mock()
-    private val mockTaskIdAttributeContainer: TaskIdAttributeContainer = mock()
+    private val mockTaskContainer: TaskContainer = mock()
     // AppPairIcon
     private val mockAppPairIcon: AppPairIcon = mock()
     private val mockContextThemeWrapper: ContextThemeWrapper = mock()
@@ -83,14 +83,15 @@
     private val transitionInfo: TransitionInfo = mock()
     private val transaction: Transaction = mock()
 
-    lateinit var splitAnimationController: SplitAnimationController
+    private lateinit var splitAnimationController: SplitAnimationController
 
     @Before
     fun setup() {
-        whenever(mockTaskView.thumbnail).thenReturn(mockThumbnailView)
+        whenever(mockTaskContainer.thumbnailView).thenReturn(mockThumbnailView)
         whenever(mockThumbnailView.thumbnail).thenReturn(mockBitmap)
-        whenever(mockTaskView.iconView).thenReturn(mockIconView)
+        whenever(mockTaskContainer.iconView).thenReturn(mockIconView)
         whenever(mockIconView.drawable).thenReturn(mockTaskViewDrawable)
+        whenever(mockTaskView.taskContainers).thenReturn(Array(1) { mockTaskContainer })
 
         whenever(splitSelectSource.drawable).thenReturn(mockSplitSourceDrawable)
         whenever(splitSelectSource.view).thenReturn(mockSplitSourceView)
@@ -177,14 +178,13 @@
         // Remove icon view from GroupedTaskView
         whenever(mockIconView.drawable).thenReturn(null)
 
-        whenever(mockTaskIdAttributeContainer.task).thenReturn(mockTask)
-        whenever(mockTaskIdAttributeContainer.iconView).thenReturn(mockIconView)
-        whenever(mockTaskIdAttributeContainer.thumbnailView).thenReturn(mockThumbnailView)
+        whenever(mockTaskContainer.task).thenReturn(mockTask)
+        whenever(mockTaskContainer.iconView).thenReturn(mockIconView)
+        whenever(mockTaskContainer.thumbnailView).thenReturn(mockThumbnailView)
         whenever(mockTask.getKey()).thenReturn(mockTaskKey)
         whenever(mockTaskKey.getId()).thenReturn(taskId)
         whenever(mockSplitSelectStateController.initialTaskId).thenReturn(taskId)
-        whenever(mockGroupedTaskView.taskIdAttributeContainers)
-            .thenReturn(Array(1) { mockTaskIdAttributeContainer })
+        whenever(mockGroupedTaskView.taskContainers).thenReturn(Array(1) { mockTaskContainer })
         val splitAnimInitProps: SplitAnimationController.Companion.SplitAnimInitProps =
             splitAnimationController.getFirstAnimInitViews(
                 { mockGroupedTaskView },
@@ -277,9 +277,7 @@
         doNothing()
             .whenever(spySplitAnimationController)
             .composeIconSplitLaunchAnimator(any(), any(), any(), any())
-        doReturn(-1)
-                .whenever(spySplitAnimationController)
-                .hasChangesForBothAppPairs(any(), any())
+        doReturn(-1).whenever(spySplitAnimationController).hasChangesForBothAppPairs(any(), any())
 
         spySplitAnimationController.playSplitLaunchAnimation(
             null /* launchingTaskView */,
@@ -305,41 +303,10 @@
         val spySplitAnimationController = spy(splitAnimationController)
         whenever(mockAppPairIcon.context).thenReturn(mockContextThemeWrapper)
         doNothing()
-                .whenever(spySplitAnimationController)
-                .composeFullscreenIconSplitLaunchAnimator(any(), any(), any(), any(), any())
-        doReturn(0)
-                .whenever(spySplitAnimationController)
-                .hasChangesForBothAppPairs(any(), any())
-
-        spySplitAnimationController.playSplitLaunchAnimation(
-                null /* launchingTaskView */,
-                mockAppPairIcon,
-                taskId,
-                taskId2,
-                null /* apps */,
-                null /* wallpapers */,
-                null /* nonApps */,
-                stateManager,
-                depthController,
-                transitionInfo,
-                transaction,
-                {} /* finishCallback */
-        )
-
-        verify(spySplitAnimationController)
-                .composeFullscreenIconSplitLaunchAnimator(any(), any(), any(), any(), eq(0))
-    }
-
-    @Test
-    fun playsAppropriateSplitLaunchAnimation_playsIconLaunchFromTaskbarCMultiWindow() {
-        val spySplitAnimationController = spy(splitAnimationController)
-        whenever(mockAppPairIcon.context).thenReturn(mockTaskbarActivityContext)
-        doNothing()
             .whenever(spySplitAnimationController)
-            .composeScaleUpLaunchAnimation(any(), any(), any(), any())
-        doReturn(-1)
-                .whenever(spySplitAnimationController)
-                .hasChangesForBothAppPairs(any(), any())
+            .composeFullscreenIconSplitLaunchAnimator(any(), any(), any(), any(), any())
+        doReturn(0).whenever(spySplitAnimationController).hasChangesForBothAppPairs(any(), any())
+
         spySplitAnimationController.playSplitLaunchAnimation(
             null /* launchingTaskView */,
             mockAppPairIcon,
@@ -355,8 +322,35 @@
             {} /* finishCallback */
         )
 
-        verify(spySplitAnimationController).composeScaleUpLaunchAnimation(any(), any(), any(),
-                eq(WINDOWING_MODE_MULTI_WINDOW))
+        verify(spySplitAnimationController)
+            .composeFullscreenIconSplitLaunchAnimator(any(), any(), any(), any(), eq(0))
+    }
+
+    @Test
+    fun playsAppropriateSplitLaunchAnimation_playsIconLaunchFromTaskbarCMultiWindow() {
+        val spySplitAnimationController = spy(splitAnimationController)
+        whenever(mockAppPairIcon.context).thenReturn(mockTaskbarActivityContext)
+        doNothing()
+            .whenever(spySplitAnimationController)
+            .composeScaleUpLaunchAnimation(any(), any(), any(), any())
+        doReturn(-1).whenever(spySplitAnimationController).hasChangesForBothAppPairs(any(), any())
+        spySplitAnimationController.playSplitLaunchAnimation(
+            null /* launchingTaskView */,
+            mockAppPairIcon,
+            taskId,
+            taskId2,
+            null /* apps */,
+            null /* wallpapers */,
+            null /* nonApps */,
+            stateManager,
+            depthController,
+            transitionInfo,
+            transaction,
+            {} /* finishCallback */
+        )
+
+        verify(spySplitAnimationController)
+            .composeScaleUpLaunchAnimation(any(), any(), any(), eq(WINDOWING_MODE_MULTI_WINDOW))
     }
 
     @Test
@@ -364,28 +358,26 @@
         val spySplitAnimationController = spy(splitAnimationController)
         whenever(mockAppPairIcon.context).thenReturn(mockTaskbarActivityContext)
         doNothing()
-                .whenever(spySplitAnimationController)
-                .composeScaleUpLaunchAnimation(any(), any(), any(), any())
-        doReturn(0)
-                .whenever(spySplitAnimationController)
-                .hasChangesForBothAppPairs(any(), any())
+            .whenever(spySplitAnimationController)
+            .composeScaleUpLaunchAnimation(any(), any(), any(), any())
+        doReturn(0).whenever(spySplitAnimationController).hasChangesForBothAppPairs(any(), any())
         spySplitAnimationController.playSplitLaunchAnimation(
-                null /* launchingTaskView */,
-                mockAppPairIcon,
-                taskId,
-                taskId2,
-                null /* apps */,
-                null /* wallpapers */,
-                null /* nonApps */,
-                stateManager,
-                depthController,
-                transitionInfo,
-                transaction,
-                {} /* finishCallback */
+            null /* launchingTaskView */,
+            mockAppPairIcon,
+            taskId,
+            taskId2,
+            null /* apps */,
+            null /* wallpapers */,
+            null /* nonApps */,
+            stateManager,
+            depthController,
+            transitionInfo,
+            transaction,
+            {} /* finishCallback */
         )
 
-        verify(spySplitAnimationController).composeScaleUpLaunchAnimation(any(), any(), any(),
-                eq(WINDOWING_MODE_FULLSCREEN))
+        verify(spySplitAnimationController)
+            .composeScaleUpLaunchAnimation(any(), any(), any(), eq(WINDOWING_MODE_FULLSCREEN))
     }
 
     @Test