diff --git a/quickstep/Android.bp b/quickstep/Android.bp
index f14cebd..1b9c661 100644
--- a/quickstep/Android.bp
+++ b/quickstep/Android.bp
@@ -52,6 +52,14 @@
         "tests/src/com/android/quickstep/TaplOverviewIconTest.java",
         "tests/src/com/android/quickstep/TaplTestsQuickstep.java",
         "tests/src/com/android/quickstep/TaplTestsSplitscreen.java",
-        "tests/src/com/android/launcher3/testcomponent/ExcludeFromRecentsTestActivity.java"
+        "tests/src/com/android/launcher3/testcomponent/ExcludeFromRecentsTestActivity.java",
+    ],
+}
+
+filegroup {
+    name: "launcher3-quickstep-screenshot-tests-src",
+    path: "tests/multivalentScreenshotTests",
+    srcs: [
+        "tests/multivalentScreenshotTests/src/**/*.kt",
     ],
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index d8145e5..e4f7262 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -145,6 +145,7 @@
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.util.DesktopTask;
 import com.android.quickstep.util.GroupTask;
+import com.android.quickstep.views.DesktopTaskView;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
@@ -1412,7 +1413,8 @@
                     if (foundTask != null) {
                         TaskView foundTaskView = recents.getTaskViewByTaskId(foundTask.key.id);
                         if (foundTaskView != null
-                                && foundTaskView.isVisibleToUser()) {
+                                && foundTaskView.isVisibleToUser()
+                                && !(foundTaskView instanceof DesktopTaskView)) {
                             TestLogging.recordEvent(
                                     TestProtocol.SEQUENCE_MAIN, "start: taskbarAppIcon");
                             foundTaskView.launchTasks();
@@ -1463,7 +1465,6 @@
                         return;
                     }
                 }
-
                 startActivity(intent);
             } else {
                 getSystemService(LauncherApps.class).startMainActivity(
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
index 4c5564e..9dee473 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
@@ -64,6 +64,10 @@
     companion object {
         private const val INDEX_LEFT = 0
         private const val INDEX_RIGHT = 1
+
+        private fun Region.addBoundsToRegion(bounds: Rect?) {
+            bounds?.let { op(it, Region.Op.UNION) }
+        }
     }
 
     /** The bottom insets taskbar provides to the IME when IME is visible. */
@@ -128,45 +132,25 @@
             }
         }
 
+        val bubbleControllers = controllers.bubbleControllers.getOrNull()
         val taskbarTouchableHeight = controllers.taskbarStashController.touchableHeight
         val bubblesTouchableHeight =
-            if (controllers.bubbleControllers.isPresent) {
-                controllers.bubbleControllers.get().bubbleStashController.getTouchableHeight()
-            } else {
-                0
-            }
+            bubbleControllers?.bubbleStashController?.getTouchableHeight() ?: 0
+        // add bounds for task bar and bubble bar stash controllers
         val touchableHeight = max(taskbarTouchableHeight, bubblesTouchableHeight)
-
-        if (
-            controllers.bubbleControllers.isPresent &&
-                controllers.bubbleControllers.get().bubbleStashController.isBubbleBarVisible()
-        ) {
-            val iconBounds =
-                controllers.bubbleControllers.get().bubbleBarViewController.bubbleBarBounds
-            defaultTouchableRegion.set(
-                iconBounds.left,
-                iconBounds.top,
-                iconBounds.right,
-                iconBounds.bottom
-            )
-        } else {
-            defaultTouchableRegion.set(
-                0,
-                windowLayoutParams.height - touchableHeight,
-                context.deviceProfile.widthPx,
-                windowLayoutParams.height
-            )
-
-            // if there's an animating bubble add it to the touch region so that it's clickable
-            val isAnimatingNewBubble =
-                controllers.bubbleControllers
-                    .getOrNull()
-                    ?.bubbleBarViewController
-                    ?.isAnimatingNewBubble ?: false
-            if (isAnimatingNewBubble) {
-                val iconBounds =
-                    controllers.bubbleControllers.get().bubbleBarViewController.bubbleBarBounds
-                defaultTouchableRegion.op(iconBounds, Region.Op.UNION)
+        defaultTouchableRegion.set(
+            0,
+            windowLayoutParams.height - touchableHeight,
+            context.deviceProfile.widthPx,
+            windowLayoutParams.height
+        )
+        if (bubbleControllers != null) {
+            val bubbleBarViewController = bubbleControllers.bubbleBarViewController
+            val isBubbleBarVisible = bubbleControllers.bubbleStashController.isBubbleBarVisible()
+            val isAnimatingNewBubble = bubbleBarViewController.isAnimatingNewBubble
+            // if bubble bar is visible or animating new bubble, add bar bounds to the touch region
+            if (isBubbleBarVisible || isAnimatingNewBubble) {
+                defaultTouchableRegion.addBoundsToRegion(bubbleBarViewController.bubbleBarBounds)
             }
         }
 
@@ -418,7 +402,7 @@
 
                 // Include the bounds of the bubble bar in the touchable region if they exist.
                 if (bubbleBarBounds != null) {
-                    region.op(bubbleBarBounds, Region.Op.UNION)
+                    region.addBoundsToRegion(bubbleBarBounds)
                 }
                 insetsInfo.touchableRegion.set(region)
                 debugTouchableRegion.lastSetTouchableReason = "Transient Taskbar is in Overview"
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index 6b5e51f..1a168a9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -27,7 +27,7 @@
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.util.DesktopTask
 import com.android.quickstep.util.GroupTask
-import com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps
+import com.android.wm.shell.shared.desktopmode.DesktopModeFlags
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import java.io.PrintWriter
 
@@ -45,7 +45,8 @@
 ) : LoggableTaskbarController {
 
     var canShowRunningApps =
-        DesktopModeStatus.canEnterDesktopMode(context) && enableDesktopWindowingTaskbarRunningApps()
+        DesktopModeStatus.canEnterDesktopMode(context) &&
+            DesktopModeFlags.TASKBAR_RUNNING_APPS.isEnabled(context)
         @VisibleForTesting
         set(isEnabledFromTest) {
             field = isEnabledFromTest
diff --git a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
index 92abebe..7d2d36d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarAllAppsButtonContainer.kt
@@ -63,10 +63,10 @@
         val padding = activityContext.taskbarSpecsEvaluator!!.taskbarIconPadding
 
         allAppsButton.setIconDrawable(drawable)
-        allAppsButton.setPadding(/* left= */ padding)
+        allAppsButton.setPadding(padding)
         allAppsButton.setForegroundTint(activityContext.getColor(R.color.all_apps_button_color))
 
-        // TODO(jagrutdesai) : add click listeners in future cl
+        // TODO(b/356465292) : add click listeners in future cl
         addView(allAppsButton)
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt
new file mode 100644
index 0000000..26e71f7
--- /dev/null
+++ b/quickstep/src/com/android/launcher3/taskbar/customization/TaskbarDividerContainer.kt
@@ -0,0 +1,65 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar.customization
+
+import android.annotation.SuppressLint
+import android.content.Context
+import android.util.AttributeSet
+import android.view.LayoutInflater
+import android.widget.LinearLayout
+import androidx.core.view.setPadding
+import com.android.launcher3.R
+import com.android.launcher3.Utilities.dpToPx
+import com.android.launcher3.taskbar.TaskbarActivityContext
+import com.android.launcher3.views.ActivityContext
+import com.android.launcher3.views.IconButtonView
+
+/** Taskbar divider view container for customizable taskbar. */
+class TaskbarDividerContainer
+@JvmOverloads
+constructor(
+    context: Context,
+    attrs: AttributeSet? = null,
+    defStyleAttr: Int = 0,
+) : LinearLayout(context, attrs), TaskbarContainer {
+
+    private val taskbarDivider: IconButtonView =
+        LayoutInflater.from(context).inflate(R.layout.taskbar_divider, this, false)
+            as IconButtonView
+    private val activityContext: TaskbarActivityContext = ActivityContext.lookupContext(context)
+
+    override val spaceNeeded: Int
+        get() {
+            return dpToPx(activityContext.taskbarSpecsEvaluator!!.taskbarIconSize.size.toFloat())
+        }
+
+    init {
+        setUpIcon()
+    }
+
+    @SuppressLint("UseCompatLoadingForDrawables")
+    fun setUpIcon() {
+        val drawable = resources.getDrawable(R.drawable.taskbar_divider_button)
+        val padding = activityContext.taskbarSpecsEvaluator!!.taskbarIconPadding
+
+        taskbarDivider.setIconDrawable(drawable)
+        taskbarDivider.setPadding(padding)
+
+        // TODO(b/356465292):: add click listeners in future cl
+        addView(taskbarDivider)
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index d4a9906..3c9bd0f 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -2100,7 +2100,6 @@
             // If there are no targets, then we don't need to capture anything
             mStateCallback.setStateOnUiThread(STATE_SCREENSHOT_CAPTURED);
         } else {
-            boolean finishTransitionPosted = false;
             // If we already have cached screenshot(s) from running tasks, skip update
             boolean shouldUpdate = false;
             int[] runningTaskIds = mGestureState.getRunningTaskIds(mIsSwipeForSplit);
@@ -2124,45 +2123,32 @@
                         }
 
                         MAIN_EXECUTOR.execute(() -> {
-                            if (!updateThumbnail(false /* refreshView */)) {
-                                setScreenshotCapturedState();
-                            }
+                            updateThumbnail();
+                            setScreenshotCapturedState();
                         });
                     });
                     return;
                 }
 
-                finishTransitionPosted = updateThumbnail(false /* refreshView */);
+                updateThumbnail();
             }
 
-            if (!finishTransitionPosted) {
-                setScreenshotCapturedState();
-            }
+            setScreenshotCapturedState();
         }
     }
 
     // Returns whether finish transition was posted.
-    private boolean updateThumbnail(boolean refreshView) {
+    private void updateThumbnail() {
         if (mGestureState.getEndTarget() == HOME
                 || mGestureState.getEndTarget() == NEW_TASK
                 || mGestureState.getEndTarget() == ALL_APPS
                 || mRecentsView == null) {
             // Capture the screenshot before finishing the transition to home or quickswitching to
             // ensure it's taken in the correct orientation, but no need to update the thumbnail.
-            return false;
+            return;
         }
 
-        boolean finishTransitionPosted = false;
-        TaskView updatedTaskView = mRecentsView.updateThumbnail(mTaskSnapshotCache, refreshView);
-        if (updatedTaskView != null && refreshView && !mCanceled) {
-            // Defer finishing the animation until the next launcher frame with the
-            // new thumbnail
-            finishTransitionPosted = ViewUtils.postFrameDrawn(updatedTaskView,
-                    () -> mStateCallback.setStateOnUiThread(STATE_SCREENSHOT_CAPTURED),
-                    this::isCanceled);
-        }
-
-        return finishTransitionPosted;
+        mRecentsView.updateThumbnail(mTaskSnapshotCache);
     }
 
     private void setScreenshotCapturedState() {
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
index da7a98f..32e2389 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
@@ -128,15 +128,11 @@
                     mController::finishAnimationToApp);
         } else {
             RemoteAnimationTarget[] nonAppTargets;
-            if (!TaskAnimationManager.ENABLE_SHELL_TRANSITIONS) {
-                nonAppTargets = mSystemUiProxy.onGoingToRecentsLegacy(appTargets);
-            } else {
-                final ArrayList<RemoteAnimationTarget> apps = new ArrayList<>();
-                final ArrayList<RemoteAnimationTarget> nonApps = new ArrayList<>();
-                classifyTargets(appTargets, apps, nonApps);
-                appTargets = apps.toArray(new RemoteAnimationTarget[apps.size()]);
-                nonAppTargets = nonApps.toArray(new RemoteAnimationTarget[nonApps.size()]);
-            }
+            final ArrayList<RemoteAnimationTarget> apps = new ArrayList<>();
+            final ArrayList<RemoteAnimationTarget> nonApps = new ArrayList<>();
+            classifyTargets(appTargets, apps, nonApps);
+            appTargets = apps.toArray(new RemoteAnimationTarget[apps.size()]);
+            nonAppTargets = nonApps.toArray(new RemoteAnimationTarget[nonApps.size()]);
             if (nonAppTargets == null) {
                 nonAppTargets = new RemoteAnimationTarget[0];
             }
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index baf669c..0c2f29b 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -23,7 +23,6 @@
 import static com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
 import static com.android.quickstep.util.ActiveGestureErrorDetector.GestureEvent.RECENT_TASKS_MISSING;
 import static com.android.quickstep.util.LogUtils.splitFailureMessage;
-import static com.android.window.flags.Flags.enableDesktopWindowingTaskbarRunningApps;
 
 import android.app.ActivityManager;
 import android.app.ActivityOptions;
@@ -95,6 +94,7 @@
 import com.android.wm.shell.recents.IRecentTasks;
 import com.android.wm.shell.recents.IRecentTasksListener;
 import com.android.wm.shell.shared.IShellTransitions;
+import com.android.wm.shell.shared.desktopmode.DesktopModeFlags;
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 import com.android.wm.shell.splitscreen.ISplitScreen;
 import com.android.wm.shell.splitscreen.ISplitScreenListener;
@@ -1000,77 +1000,6 @@
         }
     }
 
-    /**
-     * Start multiple tasks in split-screen simultaneously.
-     */
-    public void startTasksWithLegacyTransition(int taskId1, Bundle options1, int taskId2,
-            Bundle options2, @StagePosition int splitPosition,
-            @PersistentSnapPosition int snapPosition, RemoteAnimationAdapter adapter,
-            InstanceId instanceId) {
-        if (mSystemUiProxy != null) {
-            try {
-                mSplitScreen.startTasksWithLegacyTransition(taskId1, options1, taskId2, options2,
-                        splitPosition, snapPosition, adapter, instanceId);
-            } catch (RemoteException e) {
-                Log.w(TAG, splitFailureMessage(
-                        "startTasksWithLegacyTransition", "RemoteException"), e);
-            }
-        }
-    }
-
-    public void startIntentAndTaskWithLegacyTransition(PendingIntent pendingIntent, int userId1,
-            Bundle options1, int taskId, Bundle options2, @StagePosition int splitPosition,
-            @PersistentSnapPosition int snapPosition, RemoteAnimationAdapter adapter,
-            InstanceId instanceId) {
-        if (mSystemUiProxy != null) {
-            try {
-                mSplitScreen.startIntentAndTaskWithLegacyTransition(pendingIntent, userId1,
-                        options1, taskId, options2, splitPosition, snapPosition, adapter,
-                        instanceId);
-            } catch (RemoteException e) {
-                Log.w(TAG, splitFailureMessage(
-                        "startIntentAndTaskWithLegacyTransition", "RemoteException"), e);
-            }
-        }
-    }
-
-    public void startShortcutAndTaskWithLegacyTransition(ShortcutInfo shortcutInfo, Bundle options1,
-            int taskId, Bundle options2, @StagePosition int splitPosition,
-            @PersistentSnapPosition int snapPosition, RemoteAnimationAdapter adapter,
-            InstanceId instanceId) {
-        if (mSystemUiProxy != null) {
-            try {
-                mSplitScreen.startShortcutAndTaskWithLegacyTransition(shortcutInfo, options1,
-                        taskId, options2, splitPosition, snapPosition, adapter, instanceId);
-            } catch (RemoteException e) {
-                Log.w(TAG, splitFailureMessage(
-                        "startShortcutAndTaskWithLegacyTransition", "RemoteException"), e);
-            }
-        }
-    }
-
-    /**
-     * Starts a pair of intents or shortcuts in split-screen using legacy transition. Passing a
-     * non-null shortcut info means to start the app as a shortcut.
-     */
-    public void startIntentsWithLegacyTransition(PendingIntent pendingIntent1, int userId1,
-            @Nullable ShortcutInfo shortcutInfo1, @Nullable Bundle options1,
-            PendingIntent pendingIntent2, int userId2, @Nullable ShortcutInfo shortcutInfo2,
-            @Nullable Bundle options2, @StagePosition int sidePosition,
-            @PersistentSnapPosition int snapPosition, RemoteAnimationAdapter adapter,
-            InstanceId instanceId) {
-        if (mSystemUiProxy != null) {
-            try {
-                mSplitScreen.startIntentsWithLegacyTransition(pendingIntent1, userId1,
-                        shortcutInfo1, options1, pendingIntent2, userId2, shortcutInfo2, options2,
-                        sidePosition, snapPosition, adapter, instanceId);
-            } catch (RemoteException e) {
-                Log.w(TAG, splitFailureMessage(
-                        "startIntentsWithLegacyTransition", "RemoteException"), e);
-            }
-        }
-    }
-
     public void startShortcut(String packageName, String shortcutId, int position,
             Bundle options, UserHandle user, InstanceId instanceId) {
         if (mSplitScreen != null) {
@@ -1105,36 +1034,6 @@
         }
     }
 
-    /**
-     * Call this when going to recents so that shell can set-up and provide appropriate leashes
-     * for animation (eg. DividerBar).
-     *
-     * @return RemoteAnimationTargets of windows that need to animate but only exist in shell.
-     */
-    @Nullable
-    public RemoteAnimationTarget[] onGoingToRecentsLegacy(RemoteAnimationTarget[] apps) {
-        if (!TaskAnimationManager.ENABLE_SHELL_TRANSITIONS && mSplitScreen != null) {
-            try {
-                return mSplitScreen.onGoingToRecentsLegacy(apps);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call onGoingToRecentsLegacy");
-            }
-        }
-        return null;
-    }
-
-    @Nullable
-    public RemoteAnimationTarget[] onStartingSplitLegacy(RemoteAnimationTarget[] apps) {
-        if (mSplitScreen != null) {
-            try {
-                return mSplitScreen.onStartingSplitLegacy(apps);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call onStartingSplitLegacy");
-            }
-        }
-        return null;
-    }
-
     //
     // One handed
     //
@@ -1444,7 +1343,7 @@
 
     private boolean shouldEnableRunningTasksForDesktopMode() {
         return DesktopModeStatus.canEnterDesktopMode(mContext)
-                && enableDesktopWindowingTaskbarRunningApps();
+                && DesktopModeFlags.TASKBAR_RUNNING_APPS.isEnabled(mContext);
     }
 
     private boolean handleMessageAsync(Message msg) {
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 723aa03..f414399 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -264,12 +264,7 @@
                     }
                 }
 
-                RemoteAnimationTarget[] nonAppTargets = ENABLE_SHELL_TRANSITIONS
-                        ? null : getSystemUiProxy().onStartingSplitLegacy(
-                                appearedTaskTargets);
-                if (nonAppTargets == null) {
-                    nonAppTargets = new RemoteAnimationTarget[0];
-                }
+                RemoteAnimationTarget[] nonAppTargets = new RemoteAnimationTarget[0];
                 if ((containerInterface.isInLiveTileMode()
                             || mLastGestureState.getEndTarget() == RECENTS
                             || isNonRecentsStartedTasksAppeared(appearedTaskTargets))
diff --git a/quickstep/src/com/android/quickstep/TopTaskTracker.java b/quickstep/src/com/android/quickstep/TopTaskTracker.java
index 6ed05c8..3cf0542 100644
--- a/quickstep/src/com/android/quickstep/TopTaskTracker.java
+++ b/quickstep/src/com/android/quickstep/TopTaskTracker.java
@@ -28,7 +28,6 @@
 import android.annotation.UserIdInt;
 import android.app.ActivityManager.RunningTaskInfo;
 import android.content.Context;
-import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -48,7 +47,6 @@
 import com.android.systemui.shared.system.TaskStackChangeListeners;
 import com.android.wm.shell.splitscreen.ISplitScreenListener;
 
-import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
@@ -63,10 +61,6 @@
 public class TopTaskTracker extends ISplitScreenListener.Stub
         implements TaskStackChangeListener, SafeCloseable {
 
-    private static final String TAG = "TopTaskTracker";
-
-    private static final boolean DEBUG = true;
-
     public static MainThreadInitializedObject<TopTaskTracker> INSTANCE =
             new MainThreadInitializedObject<>(TopTaskTracker::new);
 
@@ -98,19 +92,10 @@
     @Override
     public void onTaskRemoved(int taskId) {
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskId);
-        if (DEBUG) {
-            Log.i(TAG, "onTaskRemoved: taskId=" + taskId);
-        }
     }
 
     @Override
     public void onTaskMovedToFront(RunningTaskInfo taskInfo) {
-        if (!mOrderedTaskList.isEmpty()
-                && mOrderedTaskList.getFirst().taskId != taskInfo.taskId
-                && DEBUG) {
-            Log.i(TAG, "onTaskMovedToFront: (moved taskInfo to front) taskId=" + taskInfo.taskId
-                    + ", baseIntent=" + taskInfo.baseIntent);
-        }
         mOrderedTaskList.removeIf(rto -> rto.taskId == taskInfo.taskId);
         mOrderedTaskList.addFirst(taskInfo);
 
@@ -121,11 +106,6 @@
             final RunningTaskInfo topTaskOnHomeDisplay = mOrderedTaskList.stream()
                     .filter(rto -> rto.displayId == DEFAULT_DISPLAY).findFirst().orElse(null);
             if (topTaskOnHomeDisplay != null) {
-                if (DEBUG) {
-                    Log.i(TAG, "onTaskMovedToFront: (removing top task on home display) taskId="
-                            + topTaskOnHomeDisplay.taskId
-                            + ", baseIntent=" + topTaskOnHomeDisplay.baseIntent);
-                }
                 mOrderedTaskList.removeIf(rto -> rto.taskId == topTaskOnHomeDisplay.taskId);
                 mOrderedTaskList.addFirst(topTaskOnHomeDisplay);
             }
@@ -139,10 +119,6 @@
                 if (info.taskId != taskInfo.taskId
                         && info.taskId != mMainStagePosition.taskId
                         && info.taskId != mSideStagePosition.taskId) {
-                    if (DEBUG) {
-                        Log.i(TAG, "onTaskMovedToFront: (removing task list overflow) taskId="
-                                + taskInfo.taskId + ", baseIntent=" + taskInfo.baseIntent);
-                    }
                     itr.remove();
                     return;
                 }
@@ -152,9 +128,6 @@
 
     @Override
     public void onStagePositionChanged(@StageType int stage, @StagePosition int position) {
-        if (DEBUG) {
-            Log.i(TAG, "onStagePositionChanged: stage=" + stage + ", position=" + position);
-        }
         if (stage == SplitConfigurationOptions.STAGE_TYPE_MAIN) {
             mMainStagePosition.stagePosition = position;
         } else {
@@ -164,10 +137,6 @@
 
     @Override
     public void onTaskStageChanged(int taskId, @StageType int stage, boolean visible) {
-        if (DEBUG) {
-            Log.i(TAG, "onTaskStageChanged: taskId=" + taskId
-                    + ", stage=" + stage + ", visible=" + visible);
-        }
         // If a task is not visible anymore or has been moved to undefined, stop tracking it.
         if (!visible || stage == SplitConfigurationOptions.STAGE_TYPE_UNDEFINED) {
             if (mMainStagePosition.taskId == taskId) {
@@ -187,18 +156,11 @@
 
     @Override
     public void onActivityPinned(String packageName, int userId, int taskId, int stackId) {
-        if (DEBUG) {
-            Log.i(TAG, "onActivityPinned: packageName=" + packageName
-                    + ", userId=" + userId + ", stackId=" + stackId);
-        }
         mPinnedTaskId = taskId;
     }
 
     @Override
     public void onActivityUnpinned() {
-        if (DEBUG) {
-            Log.i(TAG, "onActivityUnpinned");
-        }
         mPinnedTaskId = INVALID_TASK_ID;
     }
 
@@ -250,21 +212,6 @@
         return new CachedTaskInfo(tasks);
     }
 
-    public void dump(String prefix, PrintWriter writer) {
-        writer.println(prefix + "TopTaskTracker:");
-
-        writer.println(prefix + "\tmOrderedTaskList=[");
-        for (RunningTaskInfo taskInfo : mOrderedTaskList) {
-            writer.println(prefix + "\t\t(taskId=" + taskInfo.taskId
-                    + "; baseIntent=" + taskInfo.baseIntent
-                    + "; isRunning=" + taskInfo.isRunning + ")");
-        }
-        writer.println(prefix + "\t]");
-        writer.println(prefix + "\tmMainStagePosition=" + mMainStagePosition);
-        writer.println(prefix + "\tmSideStagePosition=" + mSideStagePosition);
-        writer.println(prefix + "\tmPinnedTaskId=" + mPinnedTaskId);
-    }
-
     /**
      * Class to provide information about a task which can be safely cached and do not change
      * during the lifecycle of the task.
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 2896979..5c940a3 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -1610,7 +1610,6 @@
         pw.println("\tmConsumer=" + mConsumer.getName());
         ActiveGestureLog.INSTANCE.dump("", pw);
         RecentsModel.INSTANCE.get(this).dump("", pw);
-        TopTaskTracker.INSTANCE.get(this).dump("", pw);
         if (mTaskAnimationManager != null) {
             mTaskAnimationManager.dump("", pw);
         }
diff --git a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
index d5aaed5..4f7a541 100644
--- a/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentTasksRepository.kt
@@ -46,5 +46,5 @@
      * Override [ThumbnailData] with a map of taskId to [ThumbnailData]. The override only applies
      * if the tasks are already visible, and will be invalidated when tasks become invisible.
      */
-    fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>)
+    fun addOrUpdateThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>)
 }
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 0714170..6acc940 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -82,12 +82,15 @@
 
     override fun setVisibleTasks(visibleTaskIdList: List<Int>) {
         this.visibleTaskIds.value = visibleTaskIdList.toSet()
-        setThumbnailOverride(thumbnailOverride.value)
+        addOrUpdateThumbnailOverride(emptyMap())
     }
 
-    override fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
+    override fun addOrUpdateThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
         this.thumbnailOverride.value =
-            thumbnailOverride.filterKeys(this.visibleTaskIds.value::contains).toMap()
+            this.thumbnailOverride.value
+                .toMutableMap()
+                .apply { putAll(thumbnailOverride) }
+                .filterKeys(this.visibleTaskIds.value::contains)
     }
 
     /** Flow wrapper for [TaskThumbnailDataSource.getThumbnailInBackground] api */
diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
index 7205fc8..54e34a0 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewModel.kt
@@ -58,8 +58,8 @@
         recentsViewData.thumbnailSplashProgress.value = taskThumbnailSplashAlpha
     }
 
-    fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
-        recentsTasksRepository.setThumbnailOverride(thumbnailOverride)
+    fun addOrUpdateThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
+        recentsTasksRepository.addOrUpdateThumbnailOverride(thumbnailOverride)
     }
 
     suspend fun waitForThumbnailsToUpdate(updatedThumbnails: Map<Int, ThumbnailData>) {
diff --git a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index d906bb3..54f2dd3 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -16,7 +16,6 @@
 
 package com.android.quickstep.util;
 
-import static com.android.launcher3.Utilities.postAsyncCallback;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_DESKTOP_MODE_SPLIT_LEFT_TOP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_DESKTOP_MODE_SPLIT_RIGHT_BOTTOM;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SPLIT_SELECTED_SECOND_APP;
@@ -60,8 +59,6 @@
 import android.os.UserHandle;
 import android.util.Log;
 import android.util.Pair;
-import android.view.RemoteAnimationAdapter;
-import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
 import android.window.IRemoteTransitionFinishedCallback;
 import android.window.RemoteTransition;
@@ -94,13 +91,11 @@
 import com.android.quickstep.RecentsModel;
 import com.android.quickstep.SplitSelectionListener;
 import com.android.quickstep.SystemUiProxy;
-import com.android.quickstep.TaskAnimationManager;
 import com.android.quickstep.views.FloatingTaskView;
 import com.android.quickstep.views.GroupedTaskView;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.RecentsViewContainer;
 import com.android.quickstep.views.SplitInstructionsView;
-import com.android.systemui.animation.RemoteAnimationRunnerCompat;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
@@ -276,17 +271,15 @@
                     // Loop through tasks in reverse, since they are ordered with recent tasks last
                     for (int j = taskGroups.size() - 1; j >= 0; j--) {
                         GroupTask groupTask = taskGroups.get(j);
-                        Task task1 = groupTask.task1;
-                        // Don't add duplicate Tasks
-                        if (isInstanceOfComponent(task1, key)
-                                && !Arrays.asList(lastActiveTasks).contains(task1)) {
-                            lastActiveTask = task1;
-                            break;
+                        // Account for desktop cases where there can be N tasks in the group
+                        for (Task task : groupTask.getTasks()) {
+                            if (isInstanceOfComponent(task, key)
+                                    && !Arrays.asList(lastActiveTasks).contains(task)) {
+                                lastActiveTask = task;
+                                break;
+                            }
                         }
-                        Task task2 = groupTask.task2;
-                        if (isInstanceOfComponent(task2, key)
-                                && !Arrays.asList(lastActiveTasks).contains(task2)) {
-                            lastActiveTask = task2;
+                        if (lastActiveTask != null) {
                             break;
                         }
                     }
@@ -460,77 +453,41 @@
         Bundle optionsBundle = options1.toBundle();
         Bundle extrasBundle = new Bundle(1);
         extrasBundle.putParcelable(KEY_EXTRA_WIDGET_INTENT, widgetIntent);
-        if (TaskAnimationManager.ENABLE_SHELL_TRANSITIONS) {
-            final RemoteTransition remoteTransition = getShellRemoteTransition(firstTaskId,
-                    secondTaskId, callback, "LaunchSplitPair");
-            switch (launchData.getSplitLaunchType()) {
-                case SPLIT_TASK_TASK ->
-                        mSystemUiProxy.startTasks(firstTaskId, optionsBundle, secondTaskId,
-                                null /* options2 */, initialStagePosition, snapPosition,
-                                remoteTransition, shellInstanceId);
+        final RemoteTransition remoteTransition = getRemoteTransition(firstTaskId,
+                secondTaskId, callback, "LaunchSplitPair");
+        switch (launchData.getSplitLaunchType()) {
+            case SPLIT_TASK_TASK ->
+                    mSystemUiProxy.startTasks(firstTaskId, optionsBundle, secondTaskId,
+                            null /* options2 */, initialStagePosition, snapPosition,
+                            remoteTransition, shellInstanceId);
 
-                case SPLIT_TASK_PENDINGINTENT ->
-                        mSystemUiProxy.startIntentAndTask(secondPI, secondUserId, optionsBundle,
-                                firstTaskId, extrasBundle, initialStagePosition, snapPosition,
-                                remoteTransition, shellInstanceId);
+            case SPLIT_TASK_PENDINGINTENT ->
+                    mSystemUiProxy.startIntentAndTask(secondPI, secondUserId, optionsBundle,
+                            firstTaskId, extrasBundle, initialStagePosition, snapPosition,
+                            remoteTransition, shellInstanceId);
 
-                case SPLIT_TASK_SHORTCUT ->
-                        mSystemUiProxy.startShortcutAndTask(secondShortcut, optionsBundle,
-                                firstTaskId, null /*options2*/, initialStagePosition, snapPosition,
-                                remoteTransition, shellInstanceId);
+            case SPLIT_TASK_SHORTCUT ->
+                    mSystemUiProxy.startShortcutAndTask(secondShortcut, optionsBundle,
+                            firstTaskId, null /*options2*/, initialStagePosition, snapPosition,
+                            remoteTransition, shellInstanceId);
 
-                case SPLIT_PENDINGINTENT_TASK ->
-                        mSystemUiProxy.startIntentAndTask(firstPI, firstUserId, optionsBundle,
-                                secondTaskId, null /*options2*/, initialStagePosition, snapPosition,
-                                remoteTransition, shellInstanceId);
+            case SPLIT_PENDINGINTENT_TASK ->
+                    mSystemUiProxy.startIntentAndTask(firstPI, firstUserId, optionsBundle,
+                            secondTaskId, null /*options2*/, initialStagePosition, snapPosition,
+                            remoteTransition, shellInstanceId);
 
-                case SPLIT_PENDINGINTENT_PENDINGINTENT ->
-                        mSystemUiProxy.startIntents(firstPI, firstUserId, firstShortcut,
-                                optionsBundle, secondPI, secondUserId, secondShortcut, extrasBundle,
-                                initialStagePosition, snapPosition, remoteTransition,
-                                shellInstanceId);
+            case SPLIT_PENDINGINTENT_PENDINGINTENT ->
+                    mSystemUiProxy.startIntents(firstPI, firstUserId, firstShortcut,
+                            optionsBundle, secondPI, secondUserId, secondShortcut, extrasBundle,
+                            initialStagePosition, snapPosition, remoteTransition,
+                            shellInstanceId);
 
-                case SPLIT_SHORTCUT_TASK ->
-                        mSystemUiProxy.startShortcutAndTask(firstShortcut, optionsBundle,
-                                secondTaskId, null /*options2*/, initialStagePosition, snapPosition,
-                                remoteTransition, shellInstanceId);
-            }
-        } else {
-            final RemoteAnimationAdapter adapter = getLegacyRemoteAdapter(firstTaskId, secondTaskId,
-                    callback);
-            switch (launchData.getSplitLaunchType()) {
-                case SPLIT_TASK_TASK ->
-                        mSystemUiProxy.startTasksWithLegacyTransition(firstTaskId, optionsBundle,
-                                secondTaskId, null /* options2 */, initialStagePosition,
-                                snapPosition, adapter, shellInstanceId);
-
-                case SPLIT_TASK_PENDINGINTENT ->
-                        mSystemUiProxy.startIntentAndTaskWithLegacyTransition(secondPI,
-                                secondUserId, optionsBundle, firstTaskId, null /*options2*/,
-                                initialStagePosition, snapPosition, adapter, shellInstanceId);
-
-                case SPLIT_TASK_SHORTCUT ->
-                        mSystemUiProxy.startShortcutAndTaskWithLegacyTransition(secondShortcut,
-                                optionsBundle, firstTaskId, null /*options2*/, initialStagePosition,
-                                snapPosition, adapter, shellInstanceId);
-
-                case SPLIT_PENDINGINTENT_TASK ->
-                        mSystemUiProxy.startIntentAndTaskWithLegacyTransition(firstPI, firstUserId,
-                                optionsBundle, secondTaskId, null /*options2*/,
-                                initialStagePosition, snapPosition, adapter, shellInstanceId);
-
-                case SPLIT_PENDINGINTENT_PENDINGINTENT ->
-                        mSystemUiProxy.startIntentsWithLegacyTransition(firstPI, firstUserId,
-                                firstShortcut, optionsBundle, secondPI, secondUserId,
-                                secondShortcut, null /*options2*/, initialStagePosition,
-                                snapPosition, adapter, shellInstanceId);
-
-                case SPLIT_SHORTCUT_TASK ->
-                        mSystemUiProxy.startShortcutAndTaskWithLegacyTransition(firstShortcut,
-                                optionsBundle, secondTaskId, null /*options2*/,
-                                initialStagePosition, snapPosition, adapter, shellInstanceId);
-            }
+            case SPLIT_SHORTCUT_TASK ->
+                    mSystemUiProxy.startShortcutAndTask(firstShortcut, optionsBundle,
+                            secondTaskId, null /*options2*/, initialStagePosition, snapPosition,
+                            remoteTransition, shellInstanceId);
         }
+
     }
 
     /**
@@ -576,20 +533,13 @@
         }
         Bundle optionsBundle = options1.toBundle();
 
-        if (TaskAnimationManager.ENABLE_SHELL_TRANSITIONS) {
-            final RemoteTransition transition = remoteTransition == null
-                    ? getShellRemoteTransition(
-                            firstTaskId, secondTaskId, callback, "LaunchExistingPair")
-                    : remoteTransition;
-            mSystemUiProxy.startTasks(firstTaskId, optionsBundle, secondTaskId, null /* options2 */,
-                    stagePosition, snapPosition, transition, null /*shellInstanceId*/);
-        } else {
-            final RemoteAnimationAdapter adapter = getLegacyRemoteAdapter(firstTaskId,
-                    secondTaskId, callback);
-            mSystemUiProxy.startTasksWithLegacyTransition(firstTaskId, optionsBundle, secondTaskId,
-                    null /* options2 */, stagePosition, snapPosition, adapter,
-                    null /*shellInstanceId*/);
-        }
+        final RemoteTransition transition = remoteTransition == null
+                ? getRemoteTransition(
+                firstTaskId, secondTaskId, callback, "LaunchExistingPair")
+                : remoteTransition;
+        mSystemUiProxy.startTasks(firstTaskId, optionsBundle, secondTaskId, null /* options2 */,
+                stagePosition, snapPosition, transition, null /*shellInstanceId*/);
+
     }
 
     /**
@@ -615,34 +565,16 @@
                 ActivityThread.currentActivityThread().getApplicationThread(),
                 "LaunchAppFullscreen");
         InstanceId instanceId = mSessionInstanceIds.first;
-        if (TaskAnimationManager.ENABLE_SHELL_TRANSITIONS) {
-            switch (launchData.getSplitLaunchType()) {
-                case SPLIT_SINGLE_TASK_FULLSCREEN -> mSystemUiProxy.startTasks(firstTaskId,
-                        optionsBundle, secondTaskId, null /* options2 */, initialStagePosition,
-                        SNAP_TO_50_50, remoteTransition, instanceId);
-                case SPLIT_SINGLE_INTENT_FULLSCREEN -> mSystemUiProxy.startIntentAndTask(firstPI,
-                        firstUserId, optionsBundle, secondTaskId, null /*options2*/,
-                        initialStagePosition, SNAP_TO_50_50, remoteTransition, instanceId);
-                case SPLIT_SINGLE_SHORTCUT_FULLSCREEN -> mSystemUiProxy.startShortcutAndTask(
-                        initialShortcut, optionsBundle, firstTaskId, null /* options2 */,
-                        initialStagePosition, SNAP_TO_50_50, remoteTransition, instanceId);
-            }
-        } else {
-            final RemoteAnimationAdapter adapter = getLegacyRemoteAdapter(firstTaskId,
-                    secondTaskId, callback);
-            switch (launchData.getSplitLaunchType()) {
-                case SPLIT_SINGLE_TASK_FULLSCREEN -> mSystemUiProxy.startTasksWithLegacyTransition(
-                        firstTaskId, optionsBundle, secondTaskId, null /* options2 */,
-                        initialStagePosition, SNAP_TO_50_50, adapter, instanceId);
-                case SPLIT_SINGLE_INTENT_FULLSCREEN ->
-                        mSystemUiProxy.startIntentAndTaskWithLegacyTransition(firstPI, firstUserId,
-                                optionsBundle, secondTaskId, null /*options2*/,
-                                initialStagePosition, SNAP_TO_50_50, adapter, instanceId);
-                case SPLIT_SINGLE_SHORTCUT_FULLSCREEN ->
-                        mSystemUiProxy.startShortcutAndTaskWithLegacyTransition(
-                                initialShortcut, optionsBundle, firstTaskId, null /* options2 */,
-                                initialStagePosition, SNAP_TO_50_50, adapter, instanceId);
-            }
+        switch (launchData.getSplitLaunchType()) {
+            case SPLIT_SINGLE_TASK_FULLSCREEN -> mSystemUiProxy.startTasks(firstTaskId,
+                    optionsBundle, secondTaskId, null /* options2 */, initialStagePosition,
+                    SNAP_TO_50_50, remoteTransition, instanceId);
+            case SPLIT_SINGLE_INTENT_FULLSCREEN -> mSystemUiProxy.startIntentAndTask(firstPI,
+                    firstUserId, optionsBundle, secondTaskId, null /*options2*/,
+                    initialStagePosition, SNAP_TO_50_50, remoteTransition, instanceId);
+            case SPLIT_SINGLE_SHORTCUT_FULLSCREEN -> mSystemUiProxy.startShortcutAndTask(
+                    initialShortcut, optionsBundle, firstTaskId, null /* options2 */,
+                    initialStagePosition, SNAP_TO_50_50, remoteTransition, instanceId);
         }
     }
 
@@ -660,7 +592,7 @@
         mSplitFromDesktopController = controller;
     }
 
-    private RemoteTransition getShellRemoteTransition(int firstTaskId, int secondTaskId,
+    private RemoteTransition getRemoteTransition(int firstTaskId, int secondTaskId,
             @Nullable Consumer<Boolean> callback, String transitionName) {
         final RemoteSplitLaunchTransitionRunner animationRunner =
                 new RemoteSplitLaunchTransitionRunner(firstTaskId, secondTaskId, callback);
@@ -668,14 +600,6 @@
                 ActivityThread.currentActivityThread().getApplicationThread(), transitionName);
     }
 
-    private RemoteAnimationAdapter getLegacyRemoteAdapter(int firstTaskId, int secondTaskId,
-            @Nullable Consumer<Boolean> callback) {
-        final RemoteSplitLaunchAnimationRunner animationRunner =
-                new RemoteSplitLaunchAnimationRunner(firstTaskId, secondTaskId, callback);
-        return new RemoteAnimationAdapter(animationRunner, 300, 150,
-                ActivityThread.currentActivityThread().getApplicationThread());
-    }
-
     /**
      * Will initialize {@link #mSessionInstanceIds} if null and log the first split event from
      * {@link #mSplitSelectDataHolder}
@@ -807,55 +731,6 @@
     }
 
     /**
-     * LEGACY
-     * Remote animation runner for animation to launch an app.
-     */
-    private class RemoteSplitLaunchAnimationRunner extends RemoteAnimationRunnerCompat {
-
-        private final int mInitialTaskId;
-        private final int mSecondTaskId;
-        private final Consumer<Boolean> mSuccessCallback;
-
-        RemoteSplitLaunchAnimationRunner(int initialTaskId, int secondTaskId,
-                @Nullable Consumer<Boolean> successCallback) {
-            mInitialTaskId = initialTaskId;
-            mSecondTaskId = secondTaskId;
-            mSuccessCallback = successCallback;
-        }
-
-        @Override
-        public void onAnimationStart(int transit, RemoteAnimationTarget[] apps,
-                RemoteAnimationTarget[] wallpapers, RemoteAnimationTarget[] nonApps,
-                Runnable finishedCallback) {
-            postAsyncCallback(mHandler,
-                    () -> mSplitAnimationController
-                            .playSplitLaunchAnimation(mLaunchingTaskView,
-                            mLaunchingIconView, mInitialTaskId, mSecondTaskId, apps, wallpapers,
-                            nonApps, mStateManager, mDepthController, null /* info */, null /* t */,
-                            () -> {
-                                finishedCallback.run();
-                                if (mSuccessCallback != null) {
-                                    mSuccessCallback.accept(true);
-                                }
-                                resetState();
-                            },
-                            QuickStepContract.getWindowCornerRadius(mContainer.asContext())));
-        }
-
-        @Override
-        public void onAnimationCancelled() {
-            postAsyncCallback(mHandler, () -> {
-                if (mSuccessCallback != null) {
-                    // Launching legacy tasks while recents animation is running will always cause
-                    // onAnimationCancelled to be called (should be fixed w/ shell transitions?)
-                    mSuccessCallback.accept(mRecentsAnimationRunning);
-                }
-                resetState();
-            });
-        }
-    }
-
-    /**
      * To be called whenever we exit split selection state. If
      * {@link FeatureFlags#enableSplitContextually()} is set, this should be the
      * central way split is getting reset, which should then go through the callbacks to reset
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 8b6bc39..3702f93 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1047,8 +1047,11 @@
     @Override
     @Nullable
     public Task onTaskThumbnailChanged(int taskId, ThumbnailData thumbnailData) {
+        if (enableRefactorTaskThumbnail()) {
+            // TODO(b/342560598): Listen in TaskRepository and reload
+            return null;
+        }
         if (mHandleTaskStackChanges) {
-            // TODO(b/342560598): Handle onTaskThumbnailChanged for new TTV.
             if (!enableRefactorTaskThumbnail()) {
                 TaskView taskView = getTaskViewByTaskId(taskId);
                 if (taskView != null) {
@@ -1067,6 +1070,7 @@
 
     @Override
     public void onTaskIconChanged(String pkg, UserHandle user) {
+        // TODO(b/342560598): Listen in TaskRepository and reload.
         for (int i = 0; i < getTaskViewCount(); i++) {
             TaskView tv = requireTaskViewAt(i);
             Task task = tv.getFirstTask();
@@ -1082,47 +1086,38 @@
 
     @Override
     public void onTaskIconChanged(int taskId) {
+        if (enableRefactorTaskThumbnail()) {
+            return;
+        }
         TaskView taskView = getTaskViewByTaskId(taskId);
         if (taskView != null) {
             taskView.refreshTaskThumbnailSplash();
         }
     }
 
-    /**
-     * Update the thumbnail(s) of the relevant TaskView.
-     *
-     * @param refreshNow Refresh immediately if it's true.
-     */
-    @Nullable
-    public TaskView updateThumbnail(
-            HashMap<Integer, ThumbnailData> thumbnailData, boolean refreshNow) {
+    /** Updates the thumbnail(s) of the relevant TaskView. */
+    public void updateThumbnail(Map<Integer, ThumbnailData> thumbnailData) {
         if (enableRefactorTaskThumbnail()) {
-            // TODO(b/342560598): Handle updateThumbnail for new TTV.
-            return null;
-        }
-        TaskView updatedTaskView = null;
-        for (Map.Entry<Integer, ThumbnailData> entry : thumbnailData.entrySet()) {
-            Integer id = entry.getKey();
-            ThumbnailData thumbnail = entry.getValue();
-            TaskView taskView = getTaskViewByTaskId(id);
-            if (taskView == null) {
-                continue;
+            mRecentsViewModel.addOrUpdateThumbnailOverride(thumbnailData);
+        } else {
+            for (Map.Entry<Integer, ThumbnailData> entry : thumbnailData.entrySet()) {
+                Integer id = entry.getKey();
+                ThumbnailData thumbnail = entry.getValue();
+                TaskView taskView = getTaskViewByTaskId(id);
+                if (taskView == null) {
+                    continue;
+                }
+                // taskView could be a GroupedTaskView, so select the relevant task by ID
+                TaskContainer taskContainer = taskView.getTaskContainerById(id);
+                if (taskContainer == null) {
+                    continue;
+                }
+                Task task = taskContainer.getTask();
+                TaskThumbnailViewDeprecated taskThumbnailViewDeprecated =
+                        taskContainer.getThumbnailViewDeprecated();
+                taskThumbnailViewDeprecated.setThumbnail(task, thumbnail, /*refreshNow=*/false);
             }
-            // taskView could be a GroupedTaskView, so select the relevant task by ID
-            TaskContainer taskAttributes = taskView.getTaskContainerById(id);
-            if (taskAttributes == null) {
-                continue;
-            }
-            Task task = taskAttributes.getTask();
-            TaskThumbnailViewDeprecated taskThumbnailViewDeprecated =
-                    taskAttributes.getThumbnailViewDeprecated();
-            taskThumbnailViewDeprecated.setThumbnail(task, thumbnail, refreshNow);
-            // thumbnailData can contain 1-2 ids, but they should correspond to the same
-            // TaskView, so overwriting is ok
-            updatedTaskView = taskView;
         }
-
-        return updatedTaskView;
     }
 
     @Override
@@ -2440,10 +2435,6 @@
                 List<Task> tasksToUpdate = containers.stream()
                         .map(TaskContainer::getTask)
                         .collect(Collectors.toCollection(ArrayList::new));
-                if (enableRefactorTaskThumbnail()) {
-                    visibleTaskIds.addAll(
-                            tasksToUpdate.stream().map((task) -> task.key.id).toList());
-                }
                 if (mTmpRunningTasks != null) {
                     for (Task t : mTmpRunningTasks) {
                         // Skip loading if this is the task that we are animating into
@@ -2451,6 +2442,10 @@
                         tasksToUpdate.removeIf(task -> task == t);
                     }
                 }
+                if (enableRefactorTaskThumbnail()) {
+                    visibleTaskIds.addAll(
+                            tasksToUpdate.stream().map((task) -> task.key.id).toList());
+                }
                 if (tasksToUpdate.isEmpty()) {
                     continue;
                 }
@@ -2507,6 +2502,11 @@
             mModel.preloadCacheIfNeeded();
         }
 
+        if (enableRefactorTaskThumbnail()) {
+            // TODO(b/342560598): Listen in TaskRepository and reload.
+            return;
+        }
+
         // Whenever the high res loading state changes, poke each of the visible tasks to see if
         // they want to updated their thumbnail state
         for (int i = 0; i < mHasVisibleTaskData.size(); i++) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt b/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
index 05c2462..a63ccec 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewHelper.kt
@@ -61,7 +61,7 @@
         // viewAttachedScope.
         recentsViewModel.setRunningTaskShowScreenshot(true)
         if (updatedThumbnails != null) {
-            recentsViewModel.setThumbnailOverride(updatedThumbnails)
+            recentsViewModel.addOrUpdateThumbnailOverride(updatedThumbnails)
         }
         viewAttachedScope.launch {
             recentsViewModel.waitForRunningTaskShowScreenshotToUpdate()
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index 292595c..e189d14 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -1398,7 +1398,6 @@
 
     protected open fun refreshTaskThumbnailSplash() {
         if (!enableRefactorTaskThumbnail()) {
-            // TODO(b/342560598) handle onTaskIconChanged
             taskContainers.forEach { it.thumbnailViewDeprecated.refreshSplashView() }
         }
     }
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt
new file mode 100644
index 0000000..3e0d6b5
--- /dev/null
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewScreenshotTest.kt
@@ -0,0 +1,120 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.taskbar.bubbles
+
+import android.content.Context
+import android.graphics.Bitmap
+import android.graphics.Canvas
+import android.graphics.Color
+import android.graphics.Paint
+import android.util.PathParser
+import android.view.LayoutInflater
+import androidx.test.core.app.ApplicationProvider
+import com.android.launcher3.R
+import com.android.wm.shell.common.bubbles.BubbleInfo
+import com.google.android.apps.nexuslauncher.imagecomparison.goldenpathmanager.ViewScreenshotGoldenPathManager
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import platform.test.runner.parameterized.ParameterizedAndroidJunit4
+import platform.test.runner.parameterized.Parameters
+import platform.test.screenshot.DeviceEmulationSpec
+import platform.test.screenshot.Displays
+import platform.test.screenshot.ViewScreenshotTestRule
+import platform.test.screenshot.getEmulatedDevicePathConfig
+
+/** Screenshot tests for [BubbleView]. */
+@RunWith(ParameterizedAndroidJunit4::class)
+class BubbleViewScreenshotTest(emulationSpec: DeviceEmulationSpec) {
+
+    private val context = ApplicationProvider.getApplicationContext<Context>()
+
+    companion object {
+        @Parameters(name = "{0}")
+        @JvmStatic
+        fun getTestSpecs() =
+            DeviceEmulationSpec.forDisplays(
+                Displays.Phone,
+                isDarkTheme = false,
+                isLandscape = false
+            )
+    }
+
+    @get:Rule
+    val screenshotRule =
+        ViewScreenshotTestRule(
+            emulationSpec,
+            ViewScreenshotGoldenPathManager(getEmulatedDevicePathConfig(emulationSpec))
+        )
+
+    @Test
+    fun bubbleView_hasUnseenContent() {
+        screenshotRule.screenshotTest("bubbleView_hasUnseenContent") { activity ->
+            activity.actionBar?.hide()
+            setupBubbleView()
+        }
+    }
+
+    @Test
+    fun bubbleView_seen() {
+        screenshotRule.screenshotTest("bubbleView_seen") { activity ->
+            activity.actionBar?.hide()
+            setupBubbleView().apply { markSeen() }
+        }
+    }
+
+    @Test
+    fun bubbleView_badgeHidden() {
+        screenshotRule.screenshotTest("bubbleView_badgeHidden") { activity ->
+            activity.actionBar?.hide()
+            setupBubbleView().apply { setBadgeScale(0f) }
+        }
+    }
+
+    private fun setupBubbleView(): BubbleView {
+        val inflater = LayoutInflater.from(context)
+
+        val iconSize = 100
+        // BubbleView uses launcher's badge to icon ratio and expects the badge image to already
+        // have the right size
+        val badgeToIconRatio = 0.444f
+        val badgeRadius = iconSize * badgeToIconRatio / 2
+        val icon = createCircleBitmap(radius = iconSize / 2, color = Color.LTGRAY)
+        val badge = createCircleBitmap(radius = badgeRadius.toInt(), color = Color.RED)
+
+        val bubbleInfo = BubbleInfo("key", 0, null, null, 0, context.packageName, null, null, false)
+        val bubbleView = inflater.inflate(R.layout.bubblebar_item_view, null) as BubbleView
+        val dotPath =
+            PathParser.createPathFromPathData(
+                context.resources.getString(com.android.internal.R.string.config_icon_mask)
+            )
+        val bubble =
+            BubbleBarBubble(bubbleInfo, bubbleView, badge, icon, Color.BLUE, dotPath, "test app")
+        bubbleView.setBubble(bubble)
+        bubbleView.showDotIfNeeded(1f)
+        return bubbleView
+    }
+
+    private fun createCircleBitmap(radius: Int, color: Int): Bitmap {
+        val bitmap = Bitmap.createBitmap(radius * 2, radius * 2, Bitmap.Config.ARGB_8888)
+        val canvas = Canvas(bitmap)
+        canvas.drawARGB(0, 0, 0, 0)
+        val paint = Paint()
+        paint.color = color
+        canvas.drawCircle(radius.toFloat(), radius.toFloat(), radius.toFloat(), paint)
+        return bitmap
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/logging/SettingsChangeLoggerTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/logging/SettingsChangeLoggerTest.kt
index ea2e484..d2479bc 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/logging/SettingsChangeLoggerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/logging/SettingsChangeLoggerTest.kt
@@ -20,8 +20,9 @@
 import androidx.test.core.app.ApplicationProvider
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.LauncherPrefs.Companion.ALLOW_ROTATION
 import com.android.launcher3.LauncherPrefs.Companion.THEMED_ICONS
-import com.android.launcher3.LauncherPrefs.Companion.backedUpItem
+import com.android.launcher3.SessionCommitReceiver.ADD_ICON_PREFERENCE_KEY
 import com.android.launcher3.logging.InstanceId
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ADD_NEW_APPS_TO_HOME_SCREEN_ENABLED
@@ -32,6 +33,10 @@
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_NAVIGATION_MODE_GESTURE_BUTTON
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_NOTIFICATION_DOT_ENABLED
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_THEMED_ICON_DISABLED
+import com.android.launcher3.states.RotationHelper.ALLOW_ROTATION_PREFERENCE_KEY
+import com.google.android.apps.nexuslauncher.PrefKey.KEY_ENABLE_MINUS_ONE
+import com.google.android.apps.nexuslauncher.PrefKey.OVERVIEW_SUGGESTED_ACTIONS
+import com.google.android.apps.nexuslauncher.PrefKey.SMARTSPACE_ON_HOME_SCREEN
 import com.google.common.truth.Truth.assertThat
 import org.junit.After
 import org.junit.Before
@@ -62,6 +67,7 @@
     @Captor private lateinit var mEventCaptor: ArgumentCaptor<StatsLogManager.EventEnum>
 
     private var mDefaultThemedIcons = false
+    private var mDefaultAllowRotation = false
 
     @Before
     fun setUp() {
@@ -70,8 +76,11 @@
         whenever(mStatsLogManager.logger()).doReturn(mMockLogger)
         whenever(mStatsLogManager.logger().withInstanceId(any())).doReturn(mMockLogger)
         mDefaultThemedIcons = LauncherPrefs.get(mContext).get(THEMED_ICONS)
+        mDefaultAllowRotation = LauncherPrefs.get(mContext).get(ALLOW_ROTATION)
         // To match the default value of THEMED_ICONS
         LauncherPrefs.get(mContext).put(THEMED_ICONS, false)
+        // To match the default value of ALLOW_ROTATION
+        LauncherPrefs.get(mContext).put(item = ALLOW_ROTATION, value = false)
 
         mSystemUnderTest = SettingsChangeLogger(mContext, mStatsLogManager)
     }
@@ -79,18 +88,19 @@
     @After
     fun tearDown() {
         LauncherPrefs.get(mContext).put(THEMED_ICONS, mDefaultThemedIcons)
-        mSystemUnderTest.close()
+        LauncherPrefs.get(mContext).put(ALLOW_ROTATION, mDefaultAllowRotation)
     }
 
     @Test
     fun loggingPrefs_correctDefaultValue() {
-        assertThat(mSystemUnderTest.loggingPrefs["pref_allowRotation"]!!.defaultValue).isFalse()
-        assertThat(mSystemUnderTest.loggingPrefs["pref_add_icon_to_home"]!!.defaultValue).isTrue()
-        assertThat(mSystemUnderTest.loggingPrefs["pref_overview_action_suggestions"]!!.defaultValue)
-            .isTrue()
-        assertThat(mSystemUnderTest.loggingPrefs["pref_smartspace_home_screen"]!!.defaultValue)
-            .isTrue()
-        assertThat(mSystemUnderTest.loggingPrefs["pref_enable_minus_one"]!!.defaultValue).isTrue()
+        val systemUnderTest = SettingsChangeLogger(mContext, mStatsLogManager)
+
+        assertThat(systemUnderTest.loggingPrefs[ALLOW_ROTATION_PREFERENCE_KEY]!!.defaultValue)
+            .isFalse()
+        assertThat(systemUnderTest.loggingPrefs[ADD_ICON_PREFERENCE_KEY]!!.defaultValue).isTrue()
+        assertThat(systemUnderTest.loggingPrefs[OVERVIEW_SUGGESTED_ACTIONS]!!.defaultValue).isTrue()
+        assertThat(systemUnderTest.loggingPrefs[SMARTSPACE_ON_HOME_SCREEN]!!.defaultValue).isTrue()
+        assertThat(systemUnderTest.loggingPrefs[KEY_ENABLE_MINUS_ONE]!!.defaultValue).isTrue()
     }
 
     @Test
@@ -101,24 +111,16 @@
         val capturedEvents = mEventCaptor.allValues
         assertThat(capturedEvents.isNotEmpty()).isTrue()
         verifyDefaultEvent(capturedEvents)
-        // pref_allowRotation false
         assertThat(capturedEvents.any { it.id == LAUNCHER_HOME_SCREEN_ROTATION_DISABLED.id })
             .isTrue()
     }
 
     @Test
-    fun logSnapshot_updateValue() {
-        LauncherPrefs.get(mContext)
-            .put(
-                item =
-                    backedUpItem(
-                        sharedPrefKey = "pref_allowRotation",
-                        defaultValue = false,
-                    ),
-                value = true
-            )
+    fun logSnapshot_updateAllowRotation() {
+        LauncherPrefs.get(mContext).put(item = ALLOW_ROTATION, value = true)
 
-        mSystemUnderTest.logSnapshot(mInstanceId)
+        // This a new object so the values of mLoggablePrefs will be different
+        SettingsChangeLogger(mContext, mStatsLogManager).logSnapshot(mInstanceId)
 
         verify(mMockLogger, atLeastOnce()).log(mEventCaptor.capture())
         val capturedEvents = mEventCaptor.allValues
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
index e488413..d94a351 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/FakeTasksRepository.kt
@@ -65,7 +65,7 @@
         setThumbnailOverrideInternal(thumbnailOverrideMap)
     }
 
-    override fun setThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
+    override fun addOrUpdateThumbnailOverride(thumbnailOverride: Map<Int, ThumbnailData>) {
         setThumbnailOverrideInternal(thumbnailOverride)
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
index aee5d1e..b34e156 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/TasksRepositoryTest.kt
@@ -193,57 +193,79 @@
     }
 
     @Test
-    fun setThumbnailOverrideOverrideThumbnails() = runTest {
+    fun addThumbnailOverrideOverrideThumbnails() = runTest {
         recentsModel.seedTasks(defaultTaskList)
         val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        val thumbnailOverride = createThumbnailData()
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1))
-        systemUnderTest.setThumbnailOverride(mapOf(2 to thumbnailOverride))
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
-            .isEqualTo(bitmap1)
-        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
-    }
-
-    @Test
-    fun setThumbnailOverrideClearedWhenTaskBecomeInvisible() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
-        val thumbnailOverride = createThumbnailData()
+        val thumbnailOverride2 = createThumbnailData()
         systemUnderTest.getAllTaskData(forceRefresh = true)
 
         systemUnderTest.setVisibleTasks(listOf(1, 2))
-        systemUnderTest.setThumbnailOverride(mapOf(2 to thumbnailOverride))
-        systemUnderTest.setVisibleTasks(listOf(1))
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-
-        // .drop(1) to ignore initial null content before from thumbnail was loaded.
-        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
-            .isEqualTo(bitmap1)
-        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
-    }
-
-    @Test
-    fun setThumbnailOverrideDoesNotOverrideInvisibleTasks() = runTest {
-        recentsModel.seedTasks(defaultTaskList)
-        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
-        val thumbnailOverride = createThumbnailData()
-        systemUnderTest.getAllTaskData(forceRefresh = true)
-
-        systemUnderTest.setVisibleTasks(listOf(1, 2))
-        systemUnderTest.setThumbnailOverride(mapOf(2 to thumbnailOverride))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
 
         // .drop(1) to ignore initial null content before from thumbnail was loaded.
         assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
             .isEqualTo(bitmap1)
         assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
-            .isEqualTo(thumbnailOverride.thumbnail)
+            .isEqualTo(thumbnailOverride2.thumbnail)
+    }
+
+    @Test
+    fun addThumbnailOverrideMultipleOverrides() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val thumbnailOverride1 = createThumbnailData()
+        val thumbnailOverride2 = createThumbnailData()
+        val thumbnailOverride3 = createThumbnailData()
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(1 to thumbnailOverride1))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride3))
+
+        assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail)
+            .isEqualTo(thumbnailOverride1.thumbnail)
+        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail)
+            .isEqualTo(thumbnailOverride3.thumbnail)
+    }
+
+    @Test
+    fun addThumbnailOverrideClearedWhenTaskBecomeInvisible() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        val thumbnailOverride1 = createThumbnailData()
+        val thumbnailOverride2 = createThumbnailData()
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(1 to thumbnailOverride1))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride2))
+        // Making task 2 invisible and visible again should clear the override
+        systemUnderTest.setVisibleTasks(listOf(1))
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getThumbnailById(1).first()!!.thumbnail)
+            .isEqualTo(thumbnailOverride1.thumbnail)
+        assertThat(systemUnderTest.getThumbnailById(2).drop(1).first()!!.thumbnail)
+            .isEqualTo(bitmap2)
+    }
+
+    @Test
+    fun addThumbnailOverrideDoesNotOverrideInvisibleTasks() = runTest {
+        recentsModel.seedTasks(defaultTaskList)
+        val bitmap1 = taskThumbnailDataSource.taskIdToBitmap[1]
+        val bitmap2 = taskThumbnailDataSource.taskIdToBitmap[2]
+        val thumbnailOverride = createThumbnailData()
+        systemUnderTest.getAllTaskData(forceRefresh = true)
+
+        systemUnderTest.setVisibleTasks(listOf(1))
+        systemUnderTest.addOrUpdateThumbnailOverride(mapOf(2 to thumbnailOverride))
+        systemUnderTest.setVisibleTasks(listOf(1, 2))
+
+        // .drop(1) to ignore initial null content before from thumbnail was loaded.
+        assertThat(systemUnderTest.getThumbnailById(1).drop(1).first()!!.thumbnail)
+            .isEqualTo(bitmap1)
+        assertThat(systemUnderTest.getThumbnailById(2).first()!!.thumbnail).isEqualTo(bitmap2)
     }
 
     private fun createTaskWithId(taskId: Int) =
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
index 00dbcc1..b3caf2d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/viewmodel/RecentsViewModelTest.kt
@@ -89,7 +89,7 @@
 
         val thumbnailUpdate = mapOf(2 to thumbnailDataOverride)
         systemUnderTest.setRunningTaskShowScreenshot(true)
-        systemUnderTest.setThumbnailOverride(thumbnailUpdate)
+        systemUnderTest.addOrUpdateThumbnailOverride(thumbnailUpdate)
 
         systemUnderTest.waitForRunningTaskShowScreenshotToUpdate()
         systemUnderTest.waitForThumbnailsToUpdate(thumbnailUpdate)
diff --git a/quickstep/tests/src/com/android/quickstep/TaskAnimationManagerTest.java b/quickstep/tests/src/com/android/quickstep/TaskAnimationManagerTest.java
index 2d79623..3a83ae3 100644
--- a/quickstep/tests/src/com/android/quickstep/TaskAnimationManagerTest.java
+++ b/quickstep/tests/src/com/android/quickstep/TaskAnimationManagerTest.java
@@ -16,7 +16,7 @@
 
 package com.android.quickstep;
 
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assume.assumeTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doReturn;
@@ -71,7 +71,7 @@
         final ArgumentCaptor<ActivityOptions> optionsCaptor =
                 ArgumentCaptor.forClass(ActivityOptions.class);
         verify(mSystemUiProxy).startRecentsActivity(any(), optionsCaptor.capture(), any());
-        assertTrue(optionsCaptor.getValue()
-                .isPendingIntentBackgroundActivityLaunchAllowedByPermission());
+        assertEquals(ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS,
+                optionsCaptor.getValue().getPendingIntentBackgroundActivityStartMode());
     }
 }
diff --git a/res/color-night-v31/material_color_surface_container_lowest.xml b/res/color-night-v31/material_color_surface_container_lowest.xml
new file mode 100644
index 0000000..4396f6d
--- /dev/null
+++ b/res/color-night-v31/material_color_surface_container_lowest.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:color="@android:color/system_neutral1_500" android:lStar="4" />
+</selector>
\ No newline at end of file
diff --git a/res/color-v31/material_color_surface_container_lowest.xml b/res/color-v31/material_color_surface_container_lowest.xml
new file mode 100644
index 0000000..f726aea
--- /dev/null
+++ b/res/color-v31/material_color_surface_container_lowest.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<selector xmlns:android="http://schemas.android.com/apk/res/android">
+    <item android:color="@android:color/system_neutral1_500" android:lStar="100" />
+</selector>
\ No newline at end of file
diff --git a/res/drawable/private_space_install_app_icon.xml b/res/drawable/private_space_install_app_icon.xml
index 12c4a82..cfec2b1 100644
--- a/res/drawable/private_space_install_app_icon.xml
+++ b/res/drawable/private_space_install_app_icon.xml
@@ -23,9 +23,9 @@
             android:pathData="M30 0H30A30 30 0 0 1 60 30V30A30 30 0 0 1 30 60H30A30 30 0 0 1 0 30V30A30 30 0 0 1 30 0Z" />
         <path
             android:pathData="M30 0H30A30 30 0 0 1 60 30V30A30 30 0 0 1 30 60H30A30 30 0 0 1 0 30V30A30 30 0 0 1 30 0Z"
-            android:fillColor="?attr/materialColorSurfaceContainerLowest" />
+            android:fillColor="@color/material_color_surface_container_lowest" />
         <path
             android:pathData="M29 31h-6v-2h6v-6h2v6h6v2h-6v6h-2v-6Z"
-            android:fillColor="?attr/materialColorOnSurface" />
+            android:fillColor="@color/material_color_on_surface" />
     </group>
 </vector>
diff --git a/res/values-night-v31/colors.xml b/res/values-night-v31/colors.xml
index 07c450e..0f630e5 100644
--- a/res/values-night-v31/colors.xml
+++ b/res/values-night-v31/colors.xml
@@ -57,4 +57,6 @@
         @android:color/system_accent1_200</color>
     <color name="work_fab_icon_color">
         @android:color/system_accent1_900</color>
+
+    <color name="material_color_on_surface">@android:color/system_neutral1_100</color>
 </resources>
\ No newline at end of file
diff --git a/res/values-night/colors.xml b/res/values-night/colors.xml
new file mode 100644
index 0000000..887a2a5
--- /dev/null
+++ b/res/values-night/colors.xml
@@ -0,0 +1,20 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<resources>
+    <color name="material_color_surface_container_lowest">#0D0E11</color>
+    <color name="material_color_on_surface">#E3E2E6</color>
+</resources>
\ No newline at end of file
diff --git a/res/values-v31/colors.xml b/res/values-v31/colors.xml
index 5c81d49..a5cdfc7 100644
--- a/res/values-v31/colors.xml
+++ b/res/values-v31/colors.xml
@@ -110,4 +110,6 @@
         @android:color/system_accent1_900</color>
 
     <color name="overview_foreground_scrim_color">@android:color/system_neutral1_1000</color>
+
+    <color name="material_color_on_surface">@android:color/system_neutral1_900</color>
 </resources>
diff --git a/res/values/colors.xml b/res/values/colors.xml
index c82358b..3f8bede 100644
--- a/res/values/colors.xml
+++ b/res/values/colors.xml
@@ -142,6 +142,9 @@
     <color name="widget_cell_title_color_dark">@color/system_on_surface_dark</color>
     <color name="widget_cell_subtitle_color_dark">@color/system_on_surface_variant_dark</color>
 
+    <color name="material_color_surface_container_lowest">#FFFFFF</color>
+    <color name="material_color_on_surface">#1B1B1F</color>
+
     <color name="system_primary_container_light">#D9E2FF</color>
     <color name="system_on_primary_container_light">#001945</color>
     <color name="system_primary_light">#475D92</color>
diff --git a/src/com/android/launcher3/model/ModelDbController.java b/src/com/android/launcher3/model/ModelDbController.java
index 7e1d40d..da1a221 100644
--- a/src/com/android/launcher3/model/ModelDbController.java
+++ b/src/com/android/launcher3/model/ModelDbController.java
@@ -83,6 +83,7 @@
 
 import org.xmlpull.v1.XmlPullParser;
 
+import java.io.File;
 import java.io.InputStream;
 import java.io.StringReader;
 
@@ -104,10 +105,30 @@
         mContext = context;
     }
 
+    private void printDBs(String prefix) {
+        try {
+            File directory = new File(
+                    mContext.getDatabasePath(InvariantDeviceProfile.INSTANCE.get(mContext).dbFile)
+                            .getParent()
+            );
+            if (directory.exists()) {
+                for (File file : directory.listFiles()) {
+                    Log.d("b/353505773", prefix + "Database file: " + file.getName());
+                }
+            } else {
+                Log.d("b/353505773", prefix + "No files found in the database directory");
+            }
+        } catch (Exception e) {
+            Log.e("b/353505773", prefix + e.getMessage());
+        }
+    }
+
     private synchronized void createDbIfNotExists() {
         if (mOpenHelper == null) {
             mOpenHelper = createDatabaseHelper(false /* forMigration */);
+            printDBs("before: ");
             RestoreDbTask.restoreIfNeeded(mContext, this);
+            printDBs("after: ");
         }
     }
 
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 35064cf..c949ce6 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -89,7 +89,7 @@
         if (!WIDGETS_ENABLED) {
             return Collections.emptyMap();
         }
-        return mWidgetsByPackageItem;
+        return new HashMap<>(mWidgetsByPackageItem);
     }
 
     /**
diff --git a/src/com/android/launcher3/recyclerview/AllAppsRecyclerViewPool.kt b/src/com/android/launcher3/recyclerview/AllAppsRecyclerViewPool.kt
index f895b30..0a7beab 100644
--- a/src/com/android/launcher3/recyclerview/AllAppsRecyclerViewPool.kt
+++ b/src/com/android/launcher3/recyclerview/AllAppsRecyclerViewPool.kt
@@ -17,6 +17,7 @@
 package com.android.launcher3.recyclerview
 
 import android.content.Context
+import android.util.Log
 import android.view.ViewGroup
 import androidx.annotation.VisibleForTesting
 import androidx.annotation.VisibleForTesting.Companion.PROTECTED
@@ -24,6 +25,7 @@
 import androidx.recyclerview.widget.RecyclerView.RecycledViewPool
 import androidx.recyclerview.widget.RecyclerView.ViewHolder
 import com.android.launcher3.BubbleTextView
+import com.android.launcher3.BuildConfig
 import com.android.launcher3.allapps.BaseAllAppsAdapter
 import com.android.launcher3.config.FeatureFlags
 import com.android.launcher3.util.CancellableTask
@@ -32,6 +34,7 @@
 import com.android.launcher3.util.Themes
 import com.android.launcher3.views.ActivityContext
 import com.android.launcher3.views.ActivityContext.ActivityContextDelegate
+import java.lang.IllegalStateException
 
 const val PREINFLATE_ICONS_ROW_COUNT = 4
 const val EXTRA_ICONS_COUNT = 2
@@ -47,6 +50,12 @@
     @VisibleForTesting(otherwise = PROTECTED)
     var mCancellableTask: CancellableTask<List<ViewHolder>>? = null
 
+    companion object {
+        private const val TAG = "AllAppsRecyclerViewPool"
+        private const val NULL_LAYOUT_MANAGER_ERROR_STRING =
+            "activeRv's layoutManager should not be null"
+    }
+
     /**
      * Preinflate app icons. If all apps RV cannot be scrolled down, we don't need to preinflate.
      */
@@ -54,6 +63,15 @@
         val appsView = context.appsView ?: return
         val activeRv: RecyclerView = appsView.activeRecyclerView ?: return
 
+        if (activeRv.layoutManager == null) {
+            if (BuildConfig.IS_STUDIO_BUILD) {
+                throw IllegalStateException(NULL_LAYOUT_MANAGER_ERROR_STRING)
+            } else {
+                Log.e(TAG, NULL_LAYOUT_MANAGER_ERROR_STRING)
+            }
+            return
+        }
+
         // Create a separate context dedicated for all apps preinflation thread. The goal is to
         // create a separate AssetManager obj internally to avoid lock contention with
         // AssetManager obj that is associated with the launcher context on the main thread.
diff --git a/src/com/android/launcher3/util/SplitConfigurationOptions.java b/src/com/android/launcher3/util/SplitConfigurationOptions.java
index 95624b1..837d7bc 100644
--- a/src/com/android/launcher3/util/SplitConfigurationOptions.java
+++ b/src/com/android/launcher3/util/SplitConfigurationOptions.java
@@ -186,12 +186,6 @@
         public int stagePosition = STAGE_POSITION_UNDEFINED;
         @StageType
         public int stageType = STAGE_TYPE_UNDEFINED;
-
-        @Override
-        public String toString() {
-            return "SplitStageInfo { taskId=" + taskId
-                    + ", stagePosition=" + stagePosition + ", stageType=" + stageType + " }";
-        }
     }
 
     public static StatsLogManager.EventEnum getLogEventForPosition(@StagePosition int position) {
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHost.java b/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
index 71d8503..91b899c 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHost.java
@@ -24,6 +24,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.util.Executors;
@@ -77,6 +78,11 @@
         mViewToRecycle = viewToRecycle;
     }
 
+    @VisibleForTesting
+    @Nullable ListenableHostView getViewToRecycle() {
+        return mViewToRecycle;
+    }
+
     @Override
     @NonNull
     public LauncherAppWidgetHostView onCreateView(Context context, int appWidgetId,
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
index 71f7d47..ff545fe 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
@@ -186,6 +186,35 @@
         assertThat(underTest.widgetsByComponentKey).isEmpty()
     }
 
+    @Test
+    fun getWidgetsByPackageItem_returnsACopyOfMap() {
+        loadWidgets()
+
+        val latch = CountDownLatch(1)
+        Executors.MODEL_EXECUTOR.execute {
+            var update = true
+
+            // each "widgetsByPackageItem" read returns a different copy of the map held internally.
+            // Modifying one shouldn't impact another.
+            for ((_, _) in underTest.widgetsByPackageItem.entries) {
+                underTest.widgetsByPackageItem.clear()
+                if (update) { // trigger update
+                    update = false
+                    // Similarly, model could update its code independently while a client is
+                    // iterating on the list.
+                    underTest.update(app, /* packageUser= */ null)
+                }
+            }
+
+            latch.countDown()
+        }
+        if (!latch.await(LOAD_WIDGETS_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+            fail("Timed out waiting for test")
+        }
+
+        // No exception
+    }
+
     private fun loadWidgets() {
         val latch = CountDownLatch(1)
         Executors.MODEL_EXECUTOR.execute {
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt
new file mode 100644
index 0000000..79b493a
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/LauncherAppWidgetHostTest.kt
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.Executors
+import java.util.function.IntConsumer
+import org.junit.Assert.assertNotSame
+import org.junit.Assert.assertNull
+import org.junit.Assert.assertSame
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class LauncherAppWidgetHostTest {
+
+    @Mock private lateinit var onAppWidgetRemovedCallback: IntConsumer
+
+    private val context = ActivityContextWrapper(getInstrumentation().targetContext)
+    private lateinit var underTest: LauncherAppWidgetHost
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        underTest = LauncherAppWidgetHost(context, onAppWidgetRemovedCallback, emptyList())
+    }
+
+    @Test
+    fun `Host set view to recycle`() {
+        val mockRecycleView = mock(ListenableHostView::class.java)
+
+        assertNull(underTest.viewToRecycle)
+        underTest.recycleViewForNextCreation(mockRecycleView)
+
+        assertSame(mockRecycleView, underTest.viewToRecycle)
+    }
+
+    @Test
+    fun `Host create view`() {
+        val mockRecycleView = mock(ListenableHostView::class.java)
+
+        var resultView = underTest.onCreateView(context, WIDGET_ID, null)
+
+        assertNotSame(mockRecycleView, resultView)
+
+        underTest.recycleViewForNextCreation(mockRecycleView)
+        resultView = underTest.onCreateView(context, WIDGET_ID, null)
+
+        assertSame(mockRecycleView, resultView)
+    }
+
+    @Test
+    fun `Runnable called when app widget removed`() {
+        underTest.onAppWidgetRemoved(WIDGET_ID)
+
+        Executors.MODEL_EXECUTOR.submit {}.get()
+        getInstrumentation().waitForIdleSync()
+
+        verify(onAppWidgetRemovedCallback).accept(WIDGET_ID)
+    }
+
+    companion object {
+        const val WIDGET_ID = 10001
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index c926ba9..1e2744c 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -22,7 +22,6 @@
 import static com.android.launcher3.testing.shared.TestProtocol.ICON_MISSING;
 import static com.android.launcher3.testing.shared.TestProtocol.WIDGET_CONFIG_NULL_EXTRA_INTENT;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
-import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
@@ -64,7 +63,6 @@
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.util.LooperExecutor;
-import com.android.launcher3.util.SimpleBroadcastReceiver;
 import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ExtendedLongPressTimeoutRule;
@@ -237,16 +235,12 @@
     }
 
     protected void clearPackageData(String pkg) throws IOException, InterruptedException {
-        final CountDownLatch count = new CountDownLatch(2);
-        final SimpleBroadcastReceiver broadcastReceiver =
-                new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, i -> count.countDown());
-        // We OK to make binder calls on main thread in test.
-        broadcastReceiver.registerPkgActions(mTargetContext, pkg,
-                Intent.ACTION_PACKAGE_RESTARTED, Intent.ACTION_PACKAGE_DATA_CLEARED);
-
-        mDevice.executeShellCommand("pm clear " + pkg);
-        assertTrue(pkg + " didn't restart", count.await(20, TimeUnit.SECONDS));
-        mTargetContext.unregisterReceiver(broadcastReceiver);
+        assertTrue("pm clear command failed",
+                mDevice.executeShellCommand("pm clear " + pkg)
+                .contains("Success"));
+        assertTrue("pm wait-for-handler command failed",
+                mDevice.executeShellCommand("pm wait-for-handler")
+                .contains("Success"));
     }
 
     protected TestRule getRulesInsideActivityMonitor() {