Merge "Change thumbnail/icon to get the image rather than updating it" into main
diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index 31a9009..88804b7 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -23,13 +23,6 @@
 }
 
 flag {
-    name: "enable_grid_only_overview"
-    namespace: "launcher"
-    description: "Enable a grid-only overview without a focused task."
-    bug: "257950105"
-}
-
-flag {
     name: "enable_cursor_hover_states"
     namespace: "launcher"
     description: "Enables cursor hover states for certain elements."
@@ -44,13 +37,6 @@
 }
 
 flag {
-    name: "enable_overview_icon_menu"
-    namespace: "launcher"
-    description: "Enable updated overview icon and menu within task."
-    bug: "257950105"
-}
-
-flag {
     name: "enable_focus_outline"
     namespace: "launcher"
     description: "Enables focus states outline for launcher."
@@ -238,13 +224,6 @@
 }
 
 flag {
-    name: "enable_refactor_task_thumbnail"
-    namespace: "launcher"
-    description: "Enables rewritten version of TaskThumbnailViews in Overview"
-    bug: "331753115"
-}
-
-flag {
   name: "enable_handle_delayed_gesture_callbacks"
   namespace: "launcher"
   description: "Enables additional handling for delayed mid-gesture callbacks"
diff --git a/aconfig/launcher_overview.aconfig b/aconfig/launcher_overview.aconfig
new file mode 100644
index 0000000..f9327fe
--- /dev/null
+++ b/aconfig/launcher_overview.aconfig
@@ -0,0 +1,23 @@
+package: "com.android.launcher3"
+container: "system_ext"
+
+flag {
+    name: "enable_grid_only_overview"
+    namespace: "launcher_overview"
+    description: "Enable a grid-only overview without a focused task."
+    bug: "257950105"
+}
+
+flag {
+    name: "enable_overview_icon_menu"
+    namespace: "launcher_overview"
+    description: "Enable updated overview icon and menu within task."
+    bug: "257950105"
+}
+
+flag {
+    name: "enable_refactor_task_thumbnail"
+    namespace: "launcher_overview"
+    description: "Enables rewritten version of TaskThumbnailViews in Overview"
+    bug: "331753115"
+}
diff --git a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
index 68558fa..0fb9718 100644
--- a/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
+++ b/go/quickstep/src/com/android/quickstep/TaskOverlayFactoryGo.java
@@ -31,6 +31,7 @@
 import android.content.Context;
 import android.content.Intent;
 import android.content.SharedPreferences;
+import android.graphics.Bitmap;
 import android.graphics.Color;
 import android.graphics.Matrix;
 import android.graphics.drawable.ColorDrawable;
@@ -47,6 +48,7 @@
 import android.widget.TextView;
 
 import androidx.annotation.IntDef;
+import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.BaseActivity;
@@ -58,7 +60,6 @@
 import com.android.quickstep.views.GoOverviewActionsView;
 import com.android.quickstep.views.TaskContainer;
 import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.lang.annotation.Retention;
 
@@ -131,7 +132,7 @@
          * Called when the current task is interactive for the user
          */
         @Override
-        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
+        public void initOverlay(Task task, @Nullable Bitmap thumbnail, Matrix matrix,
                 boolean rotated) {
             if (mDialog != null && mDialog.isShowing()) {
                 // Redraw the dialog in case the layout changed
diff --git a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
index 7cdca74..7235b63 100644
--- a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
+++ b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
@@ -51,9 +51,11 @@
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -86,6 +88,12 @@
     private static final String EXTRA_UI_SURFACE = "ui_surface";
     private static final Pattern UI_SURFACE_PATTERN =
             Pattern.compile("^(widgets|widgets_hub)$");
+
+    /**
+     * User ids that should be filtered out of the widget lists created by this activity.
+     */
+    private static final String EXTRA_USER_ID_FILTER = "filtered_user_ids";
+
     private SimpleDragLayer<WidgetPickerActivity> mDragLayer;
     private WidgetsModel mModel;
     private LauncherAppState mApp;
@@ -101,6 +109,10 @@
     @NonNull
     private List<AppWidgetProviderInfo> mAddedWidgets = new ArrayList<>();
 
+    /** A set of user ids that should be filtered out from the selected widgets. */
+    @NonNull
+    Set<Integer> mFilteredUserIds = new HashSet<>();
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
@@ -145,6 +157,12 @@
         if (addedWidgets != null) {
             mAddedWidgets = addedWidgets;
         }
+        ArrayList<Integer> filteredUsers = getIntent().getIntegerArrayListExtra(
+                EXTRA_USER_ID_FILTER);
+        mFilteredUserIds.clear();
+        if (filteredUsers != null) {
+            mFilteredUserIds.addAll(filteredUsers);
+        }
     }
 
     @NonNull
@@ -289,6 +307,13 @@
             return rejectWidget(widget, "shortcut");
         }
 
+        if (mFilteredUserIds.contains(widget.user.getIdentifier())) {
+            return rejectWidget(
+                    widget,
+                    "widget user: %d is being filtered",
+                    widget.user.getIdentifier());
+        }
+
         if (mWidgetCategoryFilter > 0 && (info.widgetCategory & mWidgetCategoryFilter) == 0) {
             return rejectWidget(
                     widget,
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index b647a3e..0fa3fbc 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -247,7 +247,9 @@
                 ? context.getColor(R.color.taskbar_nav_icon_light_color)
                 : context.getColor(R.color.taskbar_nav_icon_dark_color);
 
-        mTaskbarTransitions = new TaskbarTransitions(mContext, mNavButtonsView);
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions = new TaskbarTransitions(mContext, mNavButtonsView);
+        }
     }
 
     /**
@@ -359,7 +361,9 @@
                 R.bool.floating_rotation_button_position_left);
         mControllers.rotationButtonController.setRotationButton(mFloatingRotationButton,
                 mRotationButtonListener);
-        mTaskbarTransitions.init();
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.init();
+        }
 
         applyState();
         mPropertyHolders.forEach(StatePropertyHolder::endAnimation);
@@ -621,7 +625,9 @@
     }
 
     public void setWallpaperVisible(boolean isVisible) {
-        mTaskbarTransitions.setWallpaperVisibility(isVisible);
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.setWallpaperVisibility(isVisible);
+        }
     }
 
     public void onTransitionModeUpdated(int barMode, boolean checkBarModes) {
@@ -632,25 +638,32 @@
     }
 
     public void checkNavBarModes() {
-        boolean isBarHidden = (mSysuiStateFlags & SYSUI_STATE_NAV_BAR_HIDDEN) != 0;
-        mTaskbarTransitions.transitionTo(mTransitionMode, !isBarHidden);
-    }
-
-    public void finishBarAnimations() {
-        mTaskbarTransitions.finishAnimations();
-    }
-
-    public void touchAutoDim(boolean reset) {
-        mTaskbarTransitions.setAutoDim(false);
-        mHandler.removeCallbacks(mAutoDim);
-        if (reset) {
-            mHandler.postDelayed(mAutoDim, AUTODIM_TIMEOUT_MS);
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            boolean isBarHidden = (mSysuiStateFlags & SYSUI_STATE_NAV_BAR_HIDDEN) != 0;
+            mTaskbarTransitions.transitionTo(mTransitionMode, !isBarHidden);
         }
     }
 
-    public void transitionTo(@BarTransitions.TransitionMode int barMode,
-            boolean animate) {
-        mTaskbarTransitions.transitionTo(barMode, animate);
+    public void finishBarAnimations() {
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.finishAnimations();
+        }
+    }
+
+    public void touchAutoDim(boolean reset) {
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.setAutoDim(false);
+            mHandler.removeCallbacks(mAutoDim);
+            if (reset) {
+                mHandler.postDelayed(mAutoDim, AUTODIM_TIMEOUT_MS);
+            }
+        }
+    }
+
+    public void transitionTo(@BarTransitions.TransitionMode int barMode, boolean animate) {
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.transitionTo(barMode, animate);
+        }
     }
 
     /** Use to set the translationY for the all nav+contextual buttons */
@@ -752,7 +765,9 @@
 
     private void onDarkIntensityChanged() {
         updateNavButtonColor();
-        mTaskbarTransitions.onDarkIntensityChanged(mTaskbarNavButtonDarkIntensity.value);
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.onDarkIntensityChanged(mTaskbarNavButtonDarkIntensity.value);
+        }
     }
 
     protected ImageView addButton(@DrawableRes int drawableId, @TaskbarButton int buttonType,
@@ -1100,7 +1115,9 @@
                 + mOnBackgroundNavButtonColorOverrideMultiplier.value);
 
         mNavButtonsView.dumpLogs(prefix + "\t", pw);
-        mTaskbarTransitions.dumpLogs(prefix + "\t", pw);
+        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+            mTaskbarTransitions.dumpLogs(prefix + "\t", pw);
+        }
     }
 
     private static String getStateString(int flags) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 21a8268..6c3b4ad 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -43,6 +43,8 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
 import static com.android.wm.shell.Flags.enableTinyTaskbar;
 
+import static java.lang.invoke.MethodHandles.Lookup.PROTECTED;
+
 import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
 import android.app.ActivityOptions;
@@ -1515,7 +1517,8 @@
         return mIsNavBarKidsMode && isThreeButtonNav();
     }
 
-    protected boolean isNavBarForceVisible() {
+    @VisibleForTesting(otherwise = PROTECTED)
+    public boolean isNavBarForceVisible() {
         return mIsNavBarForceVisible;
     }
 
@@ -1649,6 +1652,10 @@
         return mControllers.uiController.canToggleHomeAllApps();
     }
 
+    boolean isIconAlignedWithHotseat() {
+        return mControllers.uiController.isIconAlignedWithHotseat();
+    }
+
     @VisibleForTesting
     public TaskbarControllers getControllers() {
         return mControllers;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarHoverToolTipController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarHoverToolTipController.java
index 0443197..dd14109 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarHoverToolTipController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarHoverToolTipController.java
@@ -40,6 +40,7 @@
 import com.android.launcher3.BubbleTextView;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.compat.AccessibilityManagerCompat;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.views.ArrowTipView;
@@ -73,6 +74,8 @@
         } else if (mHoverView instanceof FolderIcon
                 && ((FolderIcon) mHoverView).mInfo.title != null) {
             mToolTipText = ((FolderIcon) mHoverView).mInfo.title.toString();
+        } else if (mHoverView instanceof AppPairIcon) {
+            mToolTipText = ((AppPairIcon) mHoverView).getTitleTextView().getText().toString();
         } else {
             mToolTipText = null;
         }
@@ -156,6 +159,10 @@
         if (mHoverView == null || mToolTipText == null) {
             return;
         }
+        // Do not show tooltip if taskbar icons are transitioning to hotseat.
+        if (mActivity.isIconAlignedWithHotseat()) {
+            return;
+        }
         if (mHoverView instanceof FolderIcon && !((FolderIcon) mHoverView).getIconVisible()) {
             return;
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index b294208..b90e5fd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -115,6 +115,7 @@
     private WindowManager mWindowManager;
     private FrameLayout mTaskbarRootLayout;
     private boolean mAddedWindow;
+    private boolean mIsSuspended;
     private final TaskbarNavButtonController mNavButtonController;
     private final ComponentCallbacks mComponentCallbacks;
 
@@ -443,6 +444,8 @@
      */
     @VisibleForTesting
     public synchronized void recreateTaskbar() {
+        if (mIsSuspended) return;
+
         Trace.beginSection("recreateTaskbar");
         try {
             DeviceProfile dp = mUserUnlocked ?
@@ -648,8 +651,22 @@
         }
     }
 
+    /**
+     * Removes Taskbar from the window manager and prevents recreation if {@code true}.
+     * <p>
+     * Suspending is for testing purposes only; avoid calling this method in production.
+     */
     @VisibleForTesting
-    void addTaskbarRootViewToWindow() {
+    public void setSuspended(boolean isSuspended) {
+        mIsSuspended = isSuspended;
+        if (mIsSuspended) {
+            removeTaskbarRootViewFromWindow();
+        } else {
+            addTaskbarRootViewToWindow();
+        }
+    }
+
+    private void addTaskbarRootViewToWindow() {
         if (enableTaskbarNoRecreate() && !mAddedWindow && mTaskbarActivityContext != null) {
             mWindowManager.addView(mTaskbarRootLayout,
                     mTaskbarActivityContext.getWindowLayoutParams());
@@ -657,8 +674,7 @@
         }
     }
 
-    @VisibleForTesting
-    void removeTaskbarRootViewFromWindow() {
+    private void removeTaskbarRootViewFromWindow() {
         if (enableTaskbarNoRecreate() && mAddedWindow) {
             mWindowManager.removeViewImmediate(mTaskbarRootLayout);
             mAddedWindow = false;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index dbebdb2..36828a8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -147,20 +147,39 @@
             taskListChangeId =
                 recentsModel.getTasks { tasks ->
                     allRecentTasks = tasks
+                    val oldRunningPackages = runningAppPackages
+                    val oldMinimizedPackages = minimizedAppPackages
                     desktopTask = allRecentTasks.filterIsInstance<DesktopTask>().firstOrNull()
-                    onRecentsOrHotseatChanged()
-                    controllers.taskbarViewController.commitRunningAppsToUI()
+                    val runningPackagesChanged = oldRunningPackages != runningAppPackages
+                    val minimizedPackagessChanged = oldMinimizedPackages != minimizedAppPackages
+                    if (
+                        onRecentsOrHotseatChanged() ||
+                            runningPackagesChanged ||
+                            minimizedPackagessChanged
+                    ) {
+                        controllers.taskbarViewController.commitRunningAppsToUI()
+                    }
                 }
         }
     }
 
-    private fun onRecentsOrHotseatChanged() {
+    /**
+     * Updates [shownTasks] when Recents or Hotseat changes.
+     *
+     * @return Whether [shownTasks] changed.
+     */
+    private fun onRecentsOrHotseatChanged(): Boolean {
+        val oldShownTasks = shownTasks
         shownTasks =
             if (isInDesktopMode) {
                 computeShownRunningTasks()
             } else {
                 computeShownRecentTasks()
             }
+        val shownTasksChanged = oldShownTasks != shownTasks
+        if (!shownTasksChanged) {
+            return shownTasksChanged
+        }
 
         for (groupTask in shownTasks) {
             for (task in groupTask.tasks) {
@@ -179,6 +198,7 @@
                 }
             }
         }
+        return shownTasksChanged
     }
 
     private fun computeShownRunningTasks(): List<GroupTask> {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index fa2d907..64fb04b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -31,6 +31,7 @@
 import static com.android.launcher3.util.FlagDebugUtils.formatFlagChange;
 import static com.android.quickstep.util.SystemActionConstants.SYSTEM_ACTION_ID_TASKBAR;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DIALOG_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE;
@@ -1018,7 +1019,7 @@
         long startDelay = 0;
 
         updateStateForFlag(FLAG_STASHED_IN_APP_SYSUI, hasAnyFlag(systemUiStateFlags,
-                SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE));
+                SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE | SYSUI_STATE_DIALOG_SHOWING));
 
         boolean stashForBubbles = hasAnyFlag(FLAG_IN_OVERVIEW)
                 && hasAnyFlag(systemUiStateFlags, SYSUI_STATE_BUBBLES_EXPANDED)
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index f6b1328..4100e51 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -445,6 +445,13 @@
         }
     }
 
+    /**
+     * Removes the given bubble from the backing list of bubbles after it was dismissed by the user.
+     */
+    public void onBubbleDismissed(BubbleView bubble) {
+        mBubbles.remove(bubble.getBubble().getKey());
+    }
+
     /** Tells WMShell to show the currently selected bubble. */
     public void showSelectedBubble() {
         if (getSelectedBubbleKey() != null) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index 753237a..402b091 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -182,6 +182,8 @@
 
     @Nullable
     private BubbleView mDraggedBubbleView;
+    @Nullable
+    private BubbleView mDismissedByDragBubbleView;
     private float mAlphaDuringDrag = 1f;
 
     private Controller mController;
@@ -767,6 +769,10 @@
     public void removeBubble(View bubble) {
         if (isExpanded()) {
             // TODO b/347062801 - animate the bubble bar if the last bubble is removed
+            final boolean dismissedByDrag = mDraggedBubbleView == bubble;
+            if (dismissedByDrag) {
+                mDismissedByDragBubbleView = mDraggedBubbleView;
+            }
             int bubbleCount = getChildCount();
             mBubbleAnimator = new BubbleAnimator(mIconSize, mExpandedBarIconsSpacing,
                     bubbleCount, mBubbleBarLocation.isOnLeft(isLayoutRtl()));
@@ -786,8 +792,11 @@
 
                 @Override
                 public void onAnimationUpdate(float animatedFraction) {
-                    bubble.setScaleX(1 - animatedFraction);
-                    bubble.setScaleY(1 - animatedFraction);
+                    // don't update the scale if this bubble was dismissed by drag
+                    if (!dismissedByDrag) {
+                        bubble.setScaleX(1 - animatedFraction);
+                        bubble.setScaleY(1 - animatedFraction);
+                    }
                     updateBubblesLayoutProperties(mBubbleBarLocation);
                     invalidate();
                 }
@@ -818,6 +827,7 @@
         updateWidth();
         updateBubbleAccessibilityStates();
         updateContentDescription();
+        mDismissedByDragBubbleView = null;
     }
 
     private void updateWidth() {
@@ -864,7 +874,7 @@
         float elevationState = (1 - widthState);
         for (int i = 0; i < bubbleCount; i++) {
             BubbleView bv = (BubbleView) getChildAt(i);
-            if (bv == mDraggedBubbleView) {
+            if (bv == mDraggedBubbleView || bv == mDismissedByDragBubbleView) {
                 // Skip the dragged bubble. Its translation is managed by the drag controller.
                 continue;
             }
@@ -1048,6 +1058,8 @@
         mDraggedBubbleView = view;
         if (view != null) {
             view.setZ(mDragElevation);
+            // we started dragging a bubble. reset the bubble that was previously dismissed by drag
+            mDismissedByDragBubbleView = null;
         }
         setIsDragging(view != null);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index dbc78db..40e5b64 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -526,6 +526,12 @@
         mSystemUiProxy.stopBubbleDrag(location, mBarView.getRestingTopPositionOnScreen());
     }
 
+    /** Notifies {@link BubbleBarView} that the dragged bubble was dismissed. */
+    public void onBubbleDragDismissed(BubbleView bubble) {
+        mBubbleBarController.onBubbleDismissed(bubble);
+        mBarView.removeBubble(bubble);
+    }
+
     /**
      * Notifies {@link BubbleBarView} that drag and all animations are finished.
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java
index efc747c..8316b5b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleDragController.java
@@ -153,6 +153,7 @@
             @Override
             protected void onDragDismiss() {
                 mBubblePinController.onDragEnd();
+                mBubbleBarViewController.onBubbleDragDismissed(bubbleView);
                 mBubbleBarViewController.onBubbleDragEnd();
             }
 
diff --git a/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt b/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt
index dc6365b..181cba0 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt
+++ b/quickstep/src/com/android/launcher3/uioverrides/flags/DevOptionsUiHelper.kt
@@ -79,6 +79,7 @@
 import com.android.launcher3.util.OnboardingPrefs.HOTSEAT_DISCOVERY_TIP_COUNT
 import com.android.launcher3.util.OnboardingPrefs.HOTSEAT_LONGPRESS_TIP_SEEN
 import com.android.launcher3.util.OnboardingPrefs.TASKBAR_EDU_TOOLTIP_STEP
+import com.android.launcher3.util.OnboardingPrefs.TASKBAR_SEARCH_EDU_SEEN
 import com.android.launcher3.util.PluginManagerWrapper
 import com.android.launcher3.util.StartActivityParams
 import com.android.launcher3.util.UserIconInfo
@@ -394,6 +395,7 @@
                 HOTSEAT_LONGPRESS_TIP_SEEN.sharedPrefKey
             )
             addOnboardPref("Taskbar Education", TASKBAR_EDU_TOOLTIP_STEP.sharedPrefKey)
+            addOnboardPref("Taskbar Search Education", TASKBAR_SEARCH_EDU_SEEN.sharedPrefKey)
             addOnboardPref("All Apps Visited Count", ALL_APPS_VISITED_COUNT.sharedPrefKey)
         }
     }
diff --git a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
index 93e4fbd..31e4e33 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/touchcontrollers/QuickSwitchTouchController.java
@@ -151,7 +151,7 @@
             int sysuiFlags = 0;
             TaskView tv = mOverviewPanel.getTaskViewAt(0);
             if (tv != null) {
-                sysuiFlags = tv.getFirstThumbnailViewDeprecated().getSysUiStatusNavFlags();
+                sysuiFlags = tv.getTaskContainers().getFirst().getSysUiStatusNavFlags();
             }
             mLauncher.getSystemUiController().updateUiState(UI_STATE_FULLSCREEN_TASK, sysuiFlags);
         } else {
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index bdbe826..867533c 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -926,7 +926,7 @@
             TaskView runningTask = mRecentsView.getRunningTaskView();
             TaskView centermostTask = mRecentsView.getTaskViewNearestToCenterOfScreen();
             int centermostTaskFlags = centermostTask == null ? 0
-                    : centermostTask.getFirstThumbnailViewDeprecated().getSysUiStatusNavFlags();
+                    : centermostTask.getTaskContainers().getFirst().getSysUiStatusNavFlags();
             boolean swipeUpThresholdPassed = windowProgress > 1 - UPDATE_SYSUI_FLAGS_THRESHOLD;
             boolean quickswitchThresholdPassed = centermostTask != runningTask;
 
@@ -2445,21 +2445,20 @@
             finishRecentsAnimationOnTasksAppeared(null /* onFinishComplete */);
             return;
         }
-        Optional<RemoteAnimationTarget> taskTargetOptional =
-                Arrays.stream(appearedTaskTargets)
-                        .filter(mGestureState.mLastStartedTaskIdPredicate)
-                        .findFirst();
-        if (!taskTargetOptional.isPresent()) {
+        RemoteAnimationTarget[] taskTargets = Arrays.stream(appearedTaskTargets)
+                .filter(mGestureState.mLastStartedTaskIdPredicate)
+                .toArray(RemoteAnimationTarget[]::new);
+        if (taskTargets.length == 0) {
             ActiveGestureLog.INSTANCE.addLog("No appeared task matching started task id");
             finishRecentsAnimationOnTasksAppeared(null /* onFinishComplete */);
             return;
         }
-        RemoteAnimationTarget taskTarget = taskTargetOptional.get();
+        RemoteAnimationTarget taskTarget = taskTargets[0];
         TaskView taskView = mRecentsView == null
                 ? null : mRecentsView.getTaskViewByTaskId(taskTarget.taskId);
-        if (taskView == null
-                || !taskView.getFirstThumbnailViewDeprecated().shouldShowSplashView()) {
-            ActiveGestureLog.INSTANCE.addLog("Invalid task view splash state");
+        if (taskView == null || taskView.getTaskContainers().stream().noneMatch(
+                TaskContainer::getShouldShowSplashView)) {
+            ActiveGestureLog.INSTANCE.addLog("Splash not needed");
             finishRecentsAnimationOnTasksAppeared(null /* onFinishComplete */);
             return;
         }
@@ -2468,13 +2467,13 @@
             finishRecentsAnimationOnTasksAppeared(null /* onFinishComplete */);
             return;
         }
-        animateSplashScreenExit(mContainer, appearedTaskTargets, taskTarget.leash);
+        animateSplashScreenExit(mContainer, appearedTaskTargets, taskTargets);
     }
 
     private void animateSplashScreenExit(
             @NonNull T activity,
             @NonNull RemoteAnimationTarget[] appearedTaskTargets,
-            @NonNull SurfaceControl leash) {
+            @NonNull RemoteAnimationTarget[] animatingTargets) {
         ViewGroup splashView = activity.getDragLayer();
         final QuickstepLauncher quickstepLauncher = activity instanceof QuickstepLauncher
                 ? (QuickstepLauncher) activity : null;
@@ -2492,26 +2491,28 @@
         }
         surfaceApplier.scheduleApply(transaction);
 
-        SplashScreenExitAnimationUtils.startAnimations(splashView, leash,
-                mSplashMainWindowShiftLength, new TransactionPool(), new Rect(),
-                SPLASH_ANIMATION_DURATION, SPLASH_FADE_OUT_DURATION,
-                /* iconStartAlpha= */ 0, /* brandingStartAlpha= */ 0,
-                SPLASH_APP_REVEAL_DELAY, SPLASH_APP_REVEAL_DURATION,
-                new AnimatorListenerAdapter() {
-                    @Override
-                    public void onAnimationEnd(Animator animation) {
-                        // Hiding launcher which shows the app surface behind, then
-                        // finishing recents to the app. After transition finish, showing
-                        // the views on launcher again, so it can be visible when next
-                        // animation starts.
-                        splashView.setAlpha(0);
-                        if (quickstepLauncher != null) {
-                            quickstepLauncher.getDepthController()
-                                    .pauseBlursOnWindows(false);
+        for (RemoteAnimationTarget target : animatingTargets) {
+            SplashScreenExitAnimationUtils.startAnimations(splashView, target.leash,
+                    mSplashMainWindowShiftLength, new TransactionPool(), target.screenSpaceBounds,
+                    SPLASH_ANIMATION_DURATION, SPLASH_FADE_OUT_DURATION,
+                    /* iconStartAlpha= */ 0, /* brandingStartAlpha= */ 0,
+                    SPLASH_APP_REVEAL_DELAY, SPLASH_APP_REVEAL_DURATION,
+                    new AnimatorListenerAdapter() {
+                        @Override
+                        public void onAnimationEnd(Animator animation) {
+                            // Hiding launcher which shows the app surface behind, then
+                            // finishing recents to the app. After transition finish, showing
+                            // the views on launcher again, so it can be visible when next
+                            // animation starts.
+                            splashView.setAlpha(0);
+                            if (quickstepLauncher != null) {
+                                quickstepLauncher.getDepthController()
+                                        .pauseBlursOnWindows(false);
+                            }
+                            finishRecentsAnimationOnTasksAppeared(() -> splashView.setAlpha(1));
                         }
-                        finishRecentsAnimationOnTasksAppeared(() -> splashView.setAlpha(1));
-                    }
-                });
+                    });
+        }
     }
 
     private void finishRecentsAnimationOnTasksAppeared(Runnable onFinishComplete) {
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index 13e9844..18461a6 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -132,6 +132,13 @@
      * Init drag layer and overview panel views.
      */
     protected void setupViews() {
+        SystemUiProxy systemUiProxy = SystemUiProxy.INSTANCE.get(this);
+        // SplitSelectStateController needs to be created before setContentView()
+        mSplitSelectStateController =
+                new SplitSelectStateController(this, mHandler, getStateManager(),
+                        null /* depthController */, getStatsLogManager(),
+                        systemUiProxy, RecentsModel.INSTANCE.get(this),
+                        null /*activityBackCallback*/);
         inflateRootView(R.layout.fallback_recents_activity);
         setContentView(getRootView());
         mDragLayer = findViewById(R.id.drag_layer);
@@ -139,12 +146,6 @@
         mFallbackRecentsView = findViewById(R.id.overview_panel);
         mActionsView = findViewById(R.id.overview_actions_view);
         getRootView().getSysUiScrim().getSysUIProgress().updateValue(0);
-        SystemUiProxy systemUiProxy = SystemUiProxy.INSTANCE.get(this);
-        mSplitSelectStateController =
-                new SplitSelectStateController(this, mHandler, getStateManager(),
-                        null /* depthController */, getStatsLogManager(),
-                        systemUiProxy, RecentsModel.INSTANCE.get(this),
-                        null /*activityBackCallback*/);
         mDragLayer.recreateControllers();
         if (enableDesktopWindowingMode()) {
             mDesktopRecentsTransitionController = new DesktopRecentsTransitionController(
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index b7f3f65..80902e3 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -16,18 +16,22 @@
 
 package com.android.quickstep;
 
+import static com.android.launcher3.Flags.enableRefactorTaskThumbnail;
 import static com.android.quickstep.views.OverviewActionsView.DISABLED_NO_THUMBNAIL;
 import static com.android.quickstep.views.OverviewActionsView.DISABLED_ROTATED;
 
 import android.annotation.SuppressLint;
 import android.content.Context;
+import android.graphics.Bitmap;
 import android.graphics.Insets;
 import android.graphics.Matrix;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.os.Build;
 import android.view.View;
 
 import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
 import androidx.annotation.RequiresApi;
 
 import com.android.launcher3.BaseActivity;
@@ -38,6 +42,7 @@
 import com.android.launcher3.util.ResourceBasedOverride;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.Snackbar;
+import com.android.quickstep.task.util.TaskOverlayHelper;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.quickstep.views.DesktopTaskView;
 import com.android.quickstep.views.GroupedTaskView;
@@ -47,7 +52,6 @@
 import com.android.quickstep.views.TaskContainer;
 import com.android.quickstep.views.TaskView;
 import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.recents.model.ThumbnailData;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -128,12 +132,43 @@
 
         private T mActionsView;
         protected ImageActionsApi mImageApi;
+        protected TaskOverlayHelper mHelper;
 
         protected TaskOverlay(TaskContainer taskContainer) {
             mApplicationContext = taskContainer.getTaskView().getContext().getApplicationContext();
             mTaskContainer = taskContainer;
-            mImageApi = new ImageActionsApi(
-                    mApplicationContext, mTaskContainer::getThumbnail);
+            if (enableRefactorTaskThumbnail()) {
+                mHelper = new TaskOverlayHelper(mTaskContainer.getTask(), this);
+            }
+            mImageApi = new ImageActionsApi(mApplicationContext, this::getThumbnail);
+        }
+
+        /**
+         * Initialize the overlay when a Task is bound to the TaskView.
+         */
+        public void init() {
+            if (enableRefactorTaskThumbnail()) {
+                mHelper.init();
+            }
+        }
+
+        /**
+         * Destroy the overlay when the TaskView is recycled.
+         */
+        public void destroy() {
+            if (enableRefactorTaskThumbnail()) {
+                mHelper.destroy();
+            }
+        }
+
+        protected @Nullable Bitmap getThumbnail() {
+            return enableRefactorTaskThumbnail() ? mHelper.getEnabledState().getThumbnail()
+                    : mTaskContainer.getThumbnailViewDeprecated().getThumbnail();
+        }
+
+        protected boolean isRealSnapshot() {
+            return enableRefactorTaskThumbnail() ? mHelper.getEnabledState().isRealSnapshot()
+                    : mTaskContainer.getThumbnailViewDeprecated().isRealSnapshot();
         }
 
         protected T getActionsView() {
@@ -152,14 +187,13 @@
         /**
          * Called when the current task is interactive for the user
          */
-        public void initOverlay(Task task, ThumbnailData thumbnail, Matrix matrix,
+        public void initOverlay(Task task, @Nullable Bitmap thumbnail, Matrix matrix,
                 boolean rotated) {
             getActionsView().updateDisabledFlags(DISABLED_NO_THUMBNAIL, thumbnail == null);
 
             if (thumbnail != null) {
                 getActionsView().updateDisabledFlags(DISABLED_ROTATED, rotated);
-                boolean isAllowedByPolicy = mTaskContainer.isRealSnapshot();
-                getActionsView().setCallbacks(new OverlayUICallbacksImpl(isAllowedByPolicy, task));
+                getActionsView().setCallbacks(new OverlayUICallbacksImpl(isRealSnapshot(), task));
             }
         }
 
@@ -183,8 +217,8 @@
          */
         @SuppressLint("NewApi")
         protected void saveScreenshot(Task task) {
-            if (mTaskContainer.isRealSnapshot()) {
-                mImageApi.saveScreenshot(mTaskContainer.getThumbnail(),
+            if (isRealSnapshot()) {
+                mImageApi.saveScreenshot(getThumbnail(),
                         getTaskSnapshotBounds(), getTaskSnapshotInsets(), task.key);
             } else {
                 showBlockedByPolicyMessage();
@@ -259,7 +293,36 @@
          */
         @RequiresApi(api = Build.VERSION_CODES.Q)
         public Insets getTaskSnapshotInsets() {
-            return mTaskContainer.getScaledInsets();
+            Bitmap thumbnail = getThumbnail();
+            if (thumbnail == null) {
+                return Insets.NONE;
+            }
+
+            RectF bitmapRect = new RectF(
+                    0,
+                    0,
+                    thumbnail.getWidth(),
+                    thumbnail.getHeight());
+            View snapshotView = mTaskContainer.getSnapshotView();
+            RectF viewRect = new RectF(0, 0, snapshotView.getMeasuredWidth(),
+                    snapshotView.getMeasuredHeight());
+
+            // The position helper matrix tells us how to transform the bitmap to fit the view, the
+            // inverse tells us where the view would be in the bitmaps coordinates. The insets are
+            // the difference between the bitmap bounds and the projected view bounds.
+            Matrix boundsToBitmapSpace = new Matrix();
+            Matrix thumbnailMatrix = enableRefactorTaskThumbnail()
+                    ? mHelper.getEnabledState().getThumbnailMatrix()
+                    : mTaskContainer.getThumbnailViewDeprecated().getThumbnailMatrix();
+            thumbnailMatrix.invert(boundsToBitmapSpace);
+            RectF boundsInBitmapSpace = new RectF();
+            boundsToBitmapSpace.mapRect(boundsInBitmapSpace, viewRect);
+
+            RecentsViewContainer container = RecentsViewContainer.containerFromContext(
+                    getTaskView().getContext());
+            int bottomInset = container.getDeviceProfile().isTablet
+                    ? Math.round(bitmapRect.bottom - boundsInBitmapSpace.bottom) : 0;
+            return Insets.of(0, 0, 0, bottomInset);
         }
 
         /**
diff --git a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
index 28212cf..fdb62df 100644
--- a/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
+++ b/quickstep/src/com/android/quickstep/recents/viewmodel/RecentsViewData.kt
@@ -24,4 +24,11 @@
 
     // This is typically a View concern but it is used to invalidate rendering in other Views
     val scale = MutableStateFlow(1f)
+
+    // Whether the current RecentsView state supports task overlays.
+    // TODO(b/331753115): Derive from RecentsView state flow once migrated to MVVM.
+    val overlayEnabled = MutableStateFlow(false)
+
+    // The settled set of visible taskIds that is updated after RecentsView scroll settles.
+    val settledFullyVisibleTaskIds = MutableStateFlow(emptySet<Int>())
 }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
new file mode 100644
index 0000000..feee11f
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
@@ -0,0 +1,31 @@
+/*
+ * 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.quickstep.task.thumbnail
+
+import android.graphics.Bitmap
+import android.graphics.Matrix
+
+/** Ui state for [com.android.quickstep.TaskOverlayFactory.TaskOverlay] */
+sealed class TaskOverlayUiState {
+    data object Disabled : TaskOverlayUiState()
+
+    data class Enabled(
+        val isRealSnapshot: Boolean,
+        val thumbnail: Bitmap?,
+        val thumbnailMatrix: Matrix
+    ) : TaskOverlayUiState()
+}
diff --git a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
new file mode 100644
index 0000000..de39584
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
@@ -0,0 +1,90 @@
+/*
+ * 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.quickstep.task.util
+
+import android.util.Log
+import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
+import com.android.quickstep.task.viewmodel.TaskOverlayViewModel
+import com.android.quickstep.views.RecentsView
+import com.android.quickstep.views.RecentsViewContainer
+import com.android.systemui.shared.recents.model.Task
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.MainScope
+import kotlinx.coroutines.launch
+
+/**
+ * Helper for [TaskOverlayFactory.TaskOverlay] to interact with [TaskOverlayViewModel], this helper
+ * should merge with [TaskOverlayFactory.TaskOverlay] when it's migrated to MVVM.
+ */
+class TaskOverlayHelper(val task: Task, val overlay: TaskOverlayFactory.TaskOverlay<*>) {
+    private lateinit var job: Job
+    private var uiState: TaskOverlayUiState = Disabled
+
+    // TODO(b/335649589): Ideally create and obtain this from DI. This ViewModel should be scoped
+    //  to [TaskView], and also shared between [TaskView] and [TaskThumbnailView]
+    //  This is using a lazy for now because the dependencies cannot be obtained without DI.
+    private val taskOverlayViewModel by lazy {
+        val recentsView =
+            RecentsViewContainer.containerFromContext<RecentsViewContainer>(
+                    overlay.taskView.context
+                )
+                .getOverviewPanel<RecentsView<*, *>>()
+        TaskOverlayViewModel(task, recentsView.mRecentsViewData, recentsView.mTasksRepository)
+    }
+
+    // TODO(b/331753115): TaskOverlay should listen for state changes and react.
+    val enabledState: Enabled
+        get() = uiState as Enabled
+
+    fun init() {
+        // TODO(b/335396935): This should be changed to TaskView's scope.
+        job =
+            MainScope().launch {
+                taskOverlayViewModel.overlayState.collect {
+                    uiState = it
+                    if (it is Enabled) {
+                        Log.d(
+                            TAG,
+                            "initOverlay - taskId: ${task.key.id}, thumbnail: ${it.thumbnail}"
+                        )
+                        overlay.initOverlay(
+                            task,
+                            it.thumbnail,
+                            it.thumbnailMatrix,
+                            /* rotated= */ false
+                        )
+                    } else {
+                        Log.d(TAG, "reset - taskId: ${task.key.id}")
+                        overlay.reset()
+                    }
+                }
+            }
+    }
+
+    fun destroy() {
+        job.cancel()
+        uiState = Disabled
+        overlay.reset()
+    }
+
+    companion object {
+        private const val TAG = "TaskOverlayHelper"
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
new file mode 100644
index 0000000..4682323
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
@@ -0,0 +1,59 @@
+/*
+ * 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.quickstep.task.viewmodel
+
+import android.graphics.Matrix
+import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
+import com.android.systemui.shared.recents.model.Task
+import kotlinx.coroutines.flow.combine
+import kotlinx.coroutines.flow.distinctUntilChanged
+import kotlinx.coroutines.flow.distinctUntilChangedBy
+import kotlinx.coroutines.flow.map
+
+/** View model for TaskOverlay */
+class TaskOverlayViewModel(
+    task: Task,
+    recentsViewData: RecentsViewData,
+    tasksRepository: RecentTasksRepository,
+) {
+    val overlayState =
+        combine(
+                recentsViewData.overlayEnabled,
+                recentsViewData.settledFullyVisibleTaskIds.map { it.contains(task.key.id) },
+                tasksRepository
+                    .getTaskDataById(task.key.id)
+                    .map { it?.thumbnail }
+                    .distinctUntilChangedBy { it?.snapshotId }
+            ) { isOverlayEnabled, isFullyVisible, thumbnailData ->
+                if (isOverlayEnabled && isFullyVisible) {
+                    Enabled(
+                        isRealSnapshot = (thumbnailData?.isRealSnapshot ?: false) && !task.isLocked,
+                        thumbnailData?.thumbnail,
+                        // TODO(b/343101424): Use PreviewPositionHelper, listen from a common source
+                        // with
+                        //  TaskThumbnailView.
+                        Matrix.IDENTITY_MATRIX
+                    )
+                } else {
+                    Disabled
+                }
+            }
+            .distinctUntilChanged()
+}
diff --git a/quickstep/src/com/android/quickstep/util/DesktopTask.java b/quickstep/src/com/android/quickstep/util/DesktopTask.java
index 8d99069..307b2fa 100644
--- a/quickstep/src/com/android/quickstep/util/DesktopTask.java
+++ b/quickstep/src/com/android/quickstep/util/DesktopTask.java
@@ -22,6 +22,7 @@
 import com.android.systemui.shared.recents.model.Task;
 
 import java.util.List;
+import java.util.Objects;
 
 /**
  * A {@link Task} container that can contain N number of tasks that are part of the desktop in
@@ -68,4 +69,16 @@
         return "type=" + taskViewType + " tasks=" + tasks;
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof DesktopTask that)) return false;
+        if (!super.equals(o)) return false;
+        return Objects.equals(tasks, that.tasks);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(super.hashCode(), tasks);
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/util/GroupTask.java b/quickstep/src/com/android/quickstep/util/GroupTask.java
index 945ffe3..e8b611c 100644
--- a/quickstep/src/com/android/quickstep/util/GroupTask.java
+++ b/quickstep/src/com/android/quickstep/util/GroupTask.java
@@ -26,6 +26,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * A {@link Task} container that can contain one or two tasks, depending on if the two tasks
@@ -91,4 +92,17 @@
         return "type=" + taskViewType + " task1=" + task1 + " task2=" + task2;
     }
 
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (!(o instanceof GroupTask that)) return false;
+        return taskViewType == that.taskViewType && Objects.equals(task1,
+                that.task1) && Objects.equals(task2, that.task2)
+                && Objects.equals(mSplitBounds, that.mSplitBounds);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(task1, task2, mSplitBounds, taskViewType);
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
index 88c3a08..48ed67b 100644
--- a/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
+++ b/quickstep/src/com/android/quickstep/util/SwipePipToHomeAnimator.java
@@ -453,8 +453,9 @@
             }
             // adjust the mSourceRectHint / mAppBounds by display cutout if applicable.
             if (mSourceRectHint != null && mDisplayCutoutInsets != null) {
-                if (mFromRotation == Surface.ROTATION_0 && mDisplayCutoutInsets.top >= 0) {
-                    // TODO: this is to special case the issues on Pixel Foldable device(s).
+                if (mFromRotation == Surface.ROTATION_0) {
+                    // TODO: this is to special case the issues on Foldable device
+                    // with display cutout.
                     mSourceRectHint.offset(mDisplayCutoutInsets.left, mDisplayCutoutInsets.top);
                 } else if (mFromRotation == Surface.ROTATION_90) {
                     mSourceRectHint.offset(mDisplayCutoutInsets.left, mDisplayCutoutInsets.top);
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 8553635..ee1b3e7 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -135,7 +135,6 @@
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.BaseActivity.MultiWindowModeChangedListener;
 import com.android.launcher3.DeviceProfile;
-import com.android.launcher3.Flags;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.PagedView;
@@ -229,10 +228,12 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
@@ -1015,14 +1016,17 @@
     @Nullable
     public Task onTaskThumbnailChanged(int taskId, ThumbnailData thumbnailData) {
         if (mHandleTaskStackChanges) {
-            TaskView taskView = getTaskViewByTaskId(taskId);
-            if (taskView != null) {
-                for (TaskContainer container : taskView.getTaskContainers()) {
-                    if (container == null || taskId != container.getTask().key.id) {
-                        continue;
+            // TODO(b/342560598): Handle onTaskThumbnailChanged for new TTV.
+            if (!enableRefactorTaskThumbnail()) {
+                TaskView taskView = getTaskViewByTaskId(taskId);
+                if (taskView != null) {
+                    for (TaskContainer container : taskView.getTaskContainers()) {
+                        if (container == null || taskId != container.getTask().key.id) {
+                            continue;
+                        }
+                        container.getThumbnailViewDeprecated().setThumbnail(container.getTask(),
+                                thumbnailData);
                     }
-                    container.getThumbnailViewDeprecated().setThumbnail(container.getTask(),
-                            thumbnailData);
                 }
             }
         }
@@ -1059,6 +1063,10 @@
     @Nullable
     public TaskView updateThumbnail(
             HashMap<Integer, ThumbnailData> thumbnailData, boolean refreshNow) {
+        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();
@@ -2390,12 +2398,11 @@
             if (containers.isEmpty()) {
                 continue;
             }
-            int index = indexOfChild(taskView);
             boolean visible;
             if (showAsGrid()) {
                 visible = isTaskViewWithinBounds(taskView, visibleStart, visibleEnd);
             } else {
-                visible = lower <= index && index <= upper;
+                visible = lower <= i && i <= upper;
             }
             if (visible) {
                 // Default update all non-null tasks, then remove running ones
@@ -2426,7 +2433,7 @@
                         }
                         taskView.onTaskListVisibilityChanged(true /* visible */, changes);
                     }
-                    mHasVisibleTaskData.put(task.key.id, visible);
+                    mHasVisibleTaskData.put(task.key.id, true);
                 }
             } else {
                 for (TaskContainer container : containers) {
@@ -2914,7 +2921,7 @@
         int prevRunningTaskViewId = mRunningTaskViewId;
         mRunningTaskViewId = runningTaskViewId;
 
-        if (Flags.enableRefactorTaskThumbnail()) {
+        if (enableRefactorTaskThumbnail()) {
             TaskView previousRunningTaskView = getTaskViewFromTaskViewId(prevRunningTaskViewId);
             if (previousRunningTaskView != null) {
                 previousRunningTaskView.notifyIsRunningTaskUpdated();
@@ -4860,14 +4867,16 @@
                     == mSplitSelectStateController.getInitialTaskId();
             TaskContainer taskContainer = mSplitHiddenTaskView
                     .getTaskContainers().get(primaryTaskSelected ? 1 : 0);
-            TaskThumbnailViewDeprecated thumbnail = taskContainer.getThumbnailViewDeprecated();
             mSplitSelectStateController.getSplitAnimationController()
                     .addInitialSplitFromPair(taskContainer, builder,
                             mContainer.getDeviceProfile(),
                             mSplitHiddenTaskView.getWidth(), mSplitHiddenTaskView.getHeight(),
                             primaryTaskSelected);
             builder.addOnFrameCallback(() ->{
-                thumbnail.refreshSplashView();
+                // TODO(b/334826842): Handle splash icon for new TTV.
+                if (!enableRefactorTaskThumbnail()) {
+                    taskContainer.getThumbnailViewDeprecated().refreshSplashView();
+                }
                 mSplitHiddenTaskView.updateSnapshotRadius();
             });
         } else if (isInitiatingSplitFromTaskView) {
@@ -5259,7 +5268,7 @@
         updateGridProperties();
         updateScrollSynchronously();
 
-        int targetSysUiFlags = tv.getFirstThumbnailViewDeprecated().getSysUiStatusNavFlags();
+        int targetSysUiFlags = tv.getTaskContainers().getFirst().getSysUiStatusNavFlags();
         final boolean[] passedOverviewThreshold = new boolean[] {false};
         ValueAnimator progressAnim = ValueAnimator.ofFloat(0, 1);
         progressAnim.addUpdateListener(animator -> {
@@ -5346,6 +5355,43 @@
         updateCurrentTaskActionsVisibility();
         loadVisibleTaskData(TaskView.FLAG_UPDATE_ALL);
         updateEnabledOverlays();
+
+        if (enableRefactorTaskThumbnail()) {
+            int screenStart = 0;
+            int screenEnd = 0;
+            int centerPageIndex = 0;
+            if (showAsGrid()) {
+                screenStart = getPagedOrientationHandler().getPrimaryScroll(this);
+                int pageOrientedSize = getPagedOrientationHandler().getMeasuredSize(this);
+                screenEnd = screenStart + pageOrientedSize;
+            } else {
+                centerPageIndex = getPageNearestToCenterOfScreen();
+            }
+
+            Set<Integer> fullyVisibleTaskIds = new HashSet<>();
+
+            // Update the task data for the in/visible children
+            for (int i = 0; i < getTaskViewCount(); i++) {
+                TaskView taskView = requireTaskViewAt(i);
+                List<TaskContainer> containers = taskView.getTaskContainers();
+                if (containers.isEmpty()) {
+                    continue;
+                }
+                boolean isFullyVisible;
+                if (showAsGrid()) {
+                    isFullyVisible = isTaskViewFullyWithinBounds(taskView, screenStart,
+                            screenEnd);
+                } else {
+                    isFullyVisible = i == centerPageIndex;
+                }
+                if (isFullyVisible) {
+                    List<Integer> taskIds = containers.stream().map(
+                            taskContainer -> taskContainer.getTask().key.id).toList();
+                    fullyVisibleTaskIds.addAll(taskIds);
+                }
+            }
+            mRecentsViewData.getSettledFullyVisibleTaskIds().setValue(fullyVisibleTaskIds);
+        }
     }
 
     @Override
@@ -5896,6 +5942,7 @@
         if (mOverlayEnabled != overlayEnabled) {
             mOverlayEnabled = overlayEnabled;
             updateEnabledOverlays();
+            mRecentsViewData.getOverlayEnabled().setValue(overlayEnabled);
         }
     }
 
@@ -5947,6 +5994,12 @@
     }
 
     private void switchToScreenshotInternal(Runnable onFinishRunnable) {
+        // TODO(b/342560598): Handle switchToScreenshot for new TTV.
+        if (enableRefactorTaskThumbnail()) {
+            onFinishRunnable.run();
+            return;
+        }
+
         TaskView taskView = getRunningTaskView();
         if (taskView == null) {
             onFinishRunnable.run();
@@ -6034,7 +6087,7 @@
      * tasks to be dimmed while other elements in the recents view are left alone.
      */
     public void showForegroundScrim(boolean show) {
-        // TODO(b/335606129) Add scrim response into new TTV - this is called from overlay
+        // TODO(b/349601769) Add scrim response into new TTV - this is called from overlay
         if (!show && mColorTint == 0) {
             if (mTintingAnimator != null) {
                 mTintingAnimator.cancel();
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index cfdee6c..2e01e7e 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -18,9 +18,9 @@
 
 import android.content.Intent
 import android.graphics.Bitmap
-import android.graphics.Insets
 import android.view.View
-import com.android.launcher3.Flags
+import com.android.launcher3.Flags.enableRefactorTaskThumbnail
+import com.android.launcher3.Flags.privateSpaceRestrictAccessibilityDrag
 import com.android.launcher3.LauncherSettings
 import com.android.launcher3.model.data.ItemInfoWithIcon
 import com.android.launcher3.model.data.WorkspaceItemInfo
@@ -64,13 +64,16 @@
     val thumbnail: Bitmap?
         get() = thumbnailViewDeprecated.thumbnail
 
-    // TODO(b/349120849): Extract ThumbnailData from TaskContainerData/TaskThumbnailViewModel
-    val isRealSnapshot: Boolean
-        get() = thumbnailViewDeprecated.isRealSnapshot()
+    // TODO(b/334826842): Support shouldShowSplashView for new TTV.
+    val shouldShowSplashView: Boolean
+        get() =
+            if (enableRefactorTaskThumbnail()) false
+            else thumbnailViewDeprecated.shouldShowSplashView()
 
-    // TODO(b/349120849): Extract ThumbnailData from TaskContainerData/TaskThumbnailViewModel
-    val scaledInsets: Insets
-        get() = thumbnailViewDeprecated.scaledInsets
+    // TODO(b/350743460) Support sysUiStatusNavFlags for new TTV.
+    val sysUiStatusNavFlags: Int
+        get() =
+            if (enableRefactorTaskThumbnail()) 0 else thumbnailViewDeprecated.sysUiStatusNavFlags
 
     /** Builds proto for logging */
     val itemInfo: WorkspaceItemInfo
@@ -83,7 +86,7 @@
                 intent = Intent().setComponent(componentKey.componentName)
                 title = task.title
                 taskView.recentsView?.let { screenId = it.indexOfChild(taskView) }
-                if (Flags.privateSpaceRestrictAccessibilityDrag()) {
+                if (privateSpaceRestrictAccessibilityDrag()) {
                     if (
                         UserCache.getInstance(taskView.context)
                             .getUserInfo(componentKey.user)
@@ -98,12 +101,14 @@
     fun destroy() {
         digitalWellBeingToast?.destroy()
         thumbnailView?.let { taskView.removeView(it) }
+        overlay.destroy()
     }
 
     fun bind() {
-        if (Flags.enableRefactorTaskThumbnail() && thumbnailView != null) {
+        if (enableRefactorTaskThumbnail() && thumbnailView != null) {
             thumbnailViewDeprecated.setTaskOverlay(overlay)
             bindThumbnailView()
+            overlay.init()
         } else {
             thumbnailViewDeprecated.bind(task, overlay)
         }
@@ -116,4 +121,10 @@
         //  this should be decided inside TaskThumbnailViewModel.
         thumbnailView?.viewModel?.bind(TaskThumbnail(task.key.id, taskView.isRunningTask))
     }
+
+    fun setOverlayEnabled(enabled: Boolean) {
+        if (!enableRefactorTaskThumbnail()) {
+            thumbnailViewDeprecated.setOverlayEnabled(enabled)
+        }
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
index 2afb6a6..5b7e6c7 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
@@ -28,16 +28,13 @@
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.ColorFilter;
-import android.graphics.Insets;
 import android.graphics.Matrix;
 import android.graphics.Paint;
 import android.graphics.PorterDuff;
 import android.graphics.PorterDuffXfermode;
 import android.graphics.Rect;
-import android.graphics.RectF;
 import android.graphics.Shader;
 import android.graphics.drawable.Drawable;
-import android.os.Build;
 import android.util.AttributeSet;
 import android.util.FloatProperty;
 import android.util.Property;
@@ -45,7 +42,6 @@
 import android.widget.ImageView;
 
 import androidx.annotation.Nullable;
-import androidx.annotation.RequiresApi;
 import androidx.core.graphics.ColorUtils;
 
 import com.android.launcher3.DeviceProfile;
@@ -266,40 +262,6 @@
         return mDimAlpha;
     }
 
-    /**
-     * Get the scaled insets that are being used to draw the task view. This is a subsection of
-     * the full snapshot.
-     *
-     * @return the insets in snapshot bitmap coordinates.
-     */
-    @RequiresApi(api = Build.VERSION_CODES.Q)
-    public Insets getScaledInsets() {
-        if (mThumbnailData == null) {
-            return Insets.NONE;
-        }
-
-        RectF bitmapRect = new RectF(
-                0,
-                0,
-                mThumbnailData.getThumbnail().getWidth(),
-                mThumbnailData.getThumbnail().getHeight());
-        RectF viewRect = new RectF(0, 0, getMeasuredWidth(), getMeasuredHeight());
-
-        // The position helper matrix tells us how to transform the bitmap to fit the view, the
-        // inverse tells us where the view would be in the bitmaps coordinates. The insets are the
-        // difference between the bitmap bounds and the projected view bounds.
-        Matrix boundsToBitmapSpace = new Matrix();
-        mPreviewPositionHelper.getMatrix().invert(boundsToBitmapSpace);
-        RectF boundsInBitmapSpace = new RectF();
-        boundsToBitmapSpace.mapRect(boundsInBitmapSpace, viewRect);
-
-        DeviceProfile dp = mContainer.getDeviceProfile();
-        int bottomInset = dp.isTablet
-                ? Math.round(bitmapRect.bottom - boundsInBitmapSpace.bottom) : 0;
-        return Insets.of(0, 0, 0, bottomInset);
-    }
-
-
     @SystemUiControllerFlags
     public int getSysUiStatusNavFlags() {
         if (mThumbnailData != null) {
@@ -487,7 +449,9 @@
      */
     private void refreshOverlay() {
         if (mOverlayEnabled) {
-            mOverlay.initOverlay(mTask, mThumbnailData, mPreviewPositionHelper.getMatrix(),
+            mOverlay.initOverlay(mTask,
+                    mThumbnailData != null ? mThumbnailData.getThumbnail() : null,
+                    mPreviewPositionHelper.getMatrix(),
                     mPreviewPositionHelper.isOrientationChanged());
         } else {
             mOverlay.reset();
@@ -560,6 +524,10 @@
         return mThumbnailData.isRealSnapshot && !mTask.isLocked;
     }
 
+    public Matrix getThumbnailMatrix() {
+        return mPreviewPositionHelper.getMatrix();
+    }
+
     @Override
     public void onRecycle() {
         // Do nothing
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index c7731f1..9977d30 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -156,11 +156,6 @@
         get() = taskContainers[0].task
 
     @get:Deprecated("Use [taskContainers] instead.")
-    val firstThumbnailViewDeprecated: TaskThumbnailViewDeprecated
-        /** Returns the first thumbnailView of the TaskView. */
-        get() = taskContainers[0].thumbnailViewDeprecated
-
-    @get:Deprecated("Use [taskContainers] instead.")
     val firstSnapshotView: View
         /** Returns the first snapshotView of the TaskView. */
         get() = taskContainers[0].snapshotView
@@ -1398,10 +1393,8 @@
     }
 
     open fun setOverlayEnabled(overlayEnabled: Boolean) {
-        // TODO(b/335606129) Investigate the usage of [TaskOverlay] in the new TaskThumbnailView.
-        //  and if it's still necessary we should support that in the new TTV class.
         if (!enableRefactorTaskThumbnail()) {
-            taskContainers.forEach { it.thumbnailViewDeprecated.setOverlayEnabled(overlayEnabled) }
+            taskContainers.forEach { it.setOverlayEnabled(overlayEnabled) }
         }
     }
 
@@ -1476,7 +1469,9 @@
     protected open fun updateSnapshotRadius() {
         updateCurrentFullscreenParams()
         taskContainers.forEach {
-            it.thumbnailViewDeprecated.setFullscreenParams(getThumbnailFullscreenParams())
+            if (!enableRefactorTaskThumbnail()) {
+                it.thumbnailViewDeprecated.setFullscreenParams(getThumbnailFullscreenParams())
+            }
             it.overlay.setFullscreenParams(getThumbnailFullscreenParams())
         }
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
index bfad697..2f0b446 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsControllerTest.kt
@@ -20,19 +20,20 @@
 import android.content.ComponentName
 import android.content.Intent
 import android.os.Process
-import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.BubbleTextView
 import com.android.launcher3.appprediction.PredictionRowView
 import com.android.launcher3.model.data.AppInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.notification.NotificationKeyData
-import com.android.launcher3.taskbar.TaskbarUnitTestRule
-import com.android.launcher3.taskbar.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayController
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
 import org.junit.Test
@@ -43,84 +44,95 @@
 class TaskbarAllAppsControllerTest {
 
     @get:Rule
-    val taskbarUnitTestRule = TaskbarUnitTestRule(this, getInstrumentation().targetContext)
+    val taskbarUnitTestRule =
+        TaskbarUnitTestRule(
+            this,
+            TaskbarWindowSandboxContext.create(getInstrumentation().targetContext),
+        )
     @get:Rule val animatorTestRule = AnimatorTestRule(this)
 
     @InjectController lateinit var allAppsController: TaskbarAllAppsController
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     @Test
-    @UiThreadTest
     fun testToggle_once_showsAllApps() {
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync { allAppsController.toggle() }
         assertThat(allAppsController.isOpen).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testToggle_twice_closesAllApps() {
-        allAppsController.toggle()
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync {
+            allAppsController.toggle()
+            allAppsController.toggle()
+        }
         assertThat(allAppsController.isOpen).isFalse()
     }
 
     @Test
-    @UiThreadTest
     fun testToggle_taskbarRecreated_allAppsReopened() {
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync { allAppsController.toggle() }
         taskbarUnitTestRule.recreateTaskbar()
         assertThat(allAppsController.isOpen).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testSetApps_beforeOpened_cachesInfo() {
-        allAppsController.setApps(TEST_APPS, 0, emptyMap())
-        allAppsController.toggle()
+        val overlayContext =
+            TestUtil.getOnUiThread {
+                allAppsController.setApps(TEST_APPS, 0, emptyMap())
+                allAppsController.toggle()
+                overlayController.requestWindow()
+            }
 
-        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetApps_afterOpened_updatesStore() {
-        allAppsController.toggle()
-        allAppsController.setApps(TEST_APPS, 0, emptyMap())
+        val overlayContext =
+            TestUtil.getOnUiThread {
+                allAppsController.toggle()
+                allAppsController.setApps(TEST_APPS, 0, emptyMap())
+                overlayController.requestWindow()
+            }
 
-        val overlayContext = overlayController.requestWindow()
         assertThat(overlayContext.appsView.appsStore.apps).isEqualTo(TEST_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetPredictedApps_beforeOpened_cachesInfo() {
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-        allAppsController.toggle()
-
         val predictedApps =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-                .predictedApps
+            TestUtil.getOnUiThread {
+                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+                allAppsController.toggle()
+
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .predictedApps
+            }
+
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
     @Test
-    @UiThreadTest
     fun testSetPredictedApps_afterOpened_cachesInfo() {
-        allAppsController.toggle()
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-
         val predictedApps =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-                .predictedApps
+            TestUtil.getOnUiThread {
+                allAppsController.toggle()
+                allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .predictedApps
+            }
+
         assertThat(predictedApps).isEqualTo(TEST_PREDICTED_APPS)
     }
 
@@ -136,36 +148,38 @@
         }
 
         // Ensure the recycler view fully inflates before trying to grab an icon.
-        getInstrumentation().runOnMainSync {
-            val btv =
+        val btv =
+            TestUtil.getOnUiThread {
                 overlayController
                     .requestWindow()
                     .appsView
                     .activeRecyclerView
                     .findViewHolderForAdapterPosition(0)
                     ?.itemView as? BubbleTextView
-            assertThat(btv?.hasDot()).isTrue()
-        }
+            }
+        assertThat(btv?.hasDot()).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateNotificationDots_predictedApp_hasDot() {
-        allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
-        allAppsController.toggle()
+        getInstrumentation().runOnMainSync {
+            allAppsController.setPredictedApps(TEST_PREDICTED_APPS)
+            allAppsController.toggle()
+            taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
+                PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
+                NotificationKeyData("key"),
+            )
+        }
 
-        taskbarUnitTestRule.activityContext.popupDataProvider.onNotificationPosted(
-            PackageUserKey.fromItemInfo(TEST_PREDICTED_APPS[0]),
-            NotificationKeyData("key"),
-        )
-
-        val predictionRowView =
-            overlayController
-                .requestWindow()
-                .appsView
-                .floatingHeaderView
-                .findFixedRowByType(PredictionRowView::class.java)
-        val btv = predictionRowView.getChildAt(0) as BubbleTextView
+        val btv =
+            TestUtil.getOnUiThread {
+                overlayController
+                    .requestWindow()
+                    .appsView
+                    .floatingHeaderView
+                    .findFixedRowByType(PredictionRowView::class.java)
+                    .getChildAt(0) as BubbleTextView
+            }
         assertThat(btv.hasDot()).isTrue()
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
index 72bdc16..f946d4d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayControllerTest.kt
@@ -18,7 +18,6 @@
 
 import android.app.ActivityManager.RunningTaskInfo
 import android.view.MotionEvent
-import androidx.test.annotation.UiThreadTest
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingView.TYPE_OPTIONS_POPUP
@@ -26,11 +25,12 @@
 import com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_OVERLAY_PROXY
 import com.android.launcher3.AbstractFloatingView.hasOpenView
 import com.android.launcher3.taskbar.TaskbarActivityContext
-import com.android.launcher3.taskbar.TaskbarUnitTestRule
-import com.android.launcher3.taskbar.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarWindowSandboxContext
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
-import com.android.launcher3.views.BaseDragLayer
+import com.android.launcher3.util.TestUtil.getOnUiThread
 import com.android.systemui.shared.system.TaskStackChangeListeners
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
@@ -42,81 +42,80 @@
 class TaskbarOverlayControllerTest {
 
     @get:Rule
-    val taskbarUnitTestRule = TaskbarUnitTestRule(this, getInstrumentation().targetContext)
+    val taskbarUnitTestRule =
+        TaskbarUnitTestRule(
+            this,
+            TaskbarWindowSandboxContext.create(getInstrumentation().targetContext),
+        )
     @InjectController lateinit var overlayController: TaskbarOverlayController
 
     private val taskbarContext: TaskbarActivityContext
         get() = taskbarUnitTestRule.activityContext
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_twice_reusesWindow() {
-        val context1 = overlayController.requestWindow()
-        val context2 = overlayController.requestWindow()
+        val (context1, context2) =
+            getOnUiThread {
+                Pair(overlayController.requestWindow(), overlayController.requestWindow())
+            }
         assertThat(context1).isSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_afterHidingExistingWindow_createsNewWindow() {
-        val context1 = overlayController.requestWindow()
-        overlayController.hideWindow()
+        val context1 = getOnUiThread { overlayController.requestWindow() }
+        getInstrumentation().runOnMainSync { overlayController.hideWindow() }
 
-        val context2 = overlayController.requestWindow()
+        val context2 = getOnUiThread { overlayController.requestWindow() }
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_afterHidingOverlay_createsNewWindow() {
-        val context1 = overlayController.requestWindow()
-        TestOverlayView.show(context1)
-        overlayController.hideWindow()
+        val context1 = getOnUiThread { overlayController.requestWindow() }
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(context1)
+            overlayController.hideWindow()
+        }
 
-        val context2 = overlayController.requestWindow()
+        val context2 = getOnUiThread { overlayController.requestWindow() }
         assertThat(context1).isNotSameInstanceAs(context2)
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_addsProxyView() {
-        TestOverlayView.show(overlayController.requestWindow())
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(overlayController.requestWindow())
+        }
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testRequestWindow_closeProxyView_closesOverlay() {
-        val overlay = TestOverlayView.show(overlayController.requestWindow())
-        AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        getInstrumentation().runOnMainSync {
+            AbstractFloatingView.closeOpenContainer(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)
+        }
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testRequestWindow_attachesDragLayer() {
-        lateinit var dragLayer: BaseDragLayer<*>
-        getInstrumentation().runOnMainSync {
-            dragLayer = overlayController.requestWindow().dragLayer
-        }
-
+        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
         // Allow drag layer to attach before checking.
         getInstrumentation().runOnMainSync { assertThat(dragLayer.isAttachedToWindow).isTrue() }
     }
 
     @Test
-    @UiThreadTest
     fun testHideWindow_closesOverlay() {
-        val overlay = TestOverlayView.show(overlayController.requestWindow())
-        overlayController.hideWindow()
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
+        getInstrumentation().runOnMainSync { overlayController.hideWindow() }
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
     fun testHideWindow_detachesDragLayer() {
-        lateinit var dragLayer: BaseDragLayer<*>
-        getInstrumentation().runOnMainSync {
-            dragLayer = overlayController.requestWindow().dragLayer
-        }
+        val dragLayer = getOnUiThread { overlayController.requestWindow().dragLayer }
 
         // Wait for drag layer to be attached to window before hiding.
         getInstrumentation().runOnMainSync {
@@ -126,44 +125,45 @@
     }
 
     @Test
-    @UiThreadTest
     fun testTwoOverlays_closeOne_windowStaysOpen() {
-        val context = overlayController.requestWindow()
-        val overlay1 = TestOverlayView.show(context)
-        val overlay2 = TestOverlayView.show(context)
+        val (overlay1, overlay2) =
+            getOnUiThread {
+                val context = overlayController.requestWindow()
+                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
+            }
 
-        overlay1.close(false)
+        getInstrumentation().runOnMainSync { overlay1.close(false) }
         assertThat(overlay2.isOpen).isTrue()
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isTrue()
     }
 
     @Test
-    @UiThreadTest
     fun testTwoOverlays_closeAll_closesWindow() {
-        val context = overlayController.requestWindow()
-        val overlay1 = TestOverlayView.show(context)
-        val overlay2 = TestOverlayView.show(context)
+        val (overlay1, overlay2) =
+            getOnUiThread {
+                val context = overlayController.requestWindow()
+                Pair(TestOverlayView.show(context), TestOverlayView.show(context))
+            }
 
-        overlay1.close(false)
-        overlay2.close(false)
+        getInstrumentation().runOnMainSync {
+            overlay1.close(false)
+            overlay2.close(false)
+        }
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isFalse()
     }
 
     @Test
-    @UiThreadTest
     fun testRecreateTaskbar_closesWindow() {
-        TestOverlayView.show(overlayController.requestWindow())
+        getInstrumentation().runOnMainSync {
+            TestOverlayView.show(overlayController.requestWindow())
+        }
         taskbarUnitTestRule.recreateTaskbar()
         assertThat(hasOpenView(taskbarContext, TYPE_TASKBAR_OVERLAY_PROXY)).isFalse()
     }
 
     @Test
     fun testTaskMovedToFront_closesOverlay() {
-        lateinit var overlay: TestOverlayView
-        getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
-        }
-
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         TaskStackChangeListeners.getInstance().listenerImpl.onTaskMovedToFront(RunningTaskInfo())
         // Make sure TaskStackChangeListeners' Handler posts the callback before checking state.
         getInstrumentation().runOnMainSync { assertThat(overlay.isOpen).isFalse() }
@@ -171,9 +171,8 @@
 
     @Test
     fun testTaskStackChanged_allAppsClosed_overlayStaysOpen() {
-        lateinit var overlay: TestOverlayView
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
             taskbarContext.controllers.sharedState?.allAppsVisible = false
         }
 
@@ -183,9 +182,8 @@
 
     @Test
     fun testTaskStackChanged_allAppsOpen_closesOverlay() {
-        lateinit var overlay: TestOverlayView
+        val overlay = getOnUiThread { TestOverlayView.show(overlayController.requestWindow()) }
         getInstrumentation().runOnMainSync {
-            overlay = TestOverlayView.show(overlayController.requestWindow())
             taskbarContext.controllers.sharedState?.allAppsVisible = true
         }
 
@@ -194,33 +192,39 @@
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayNotRebindSafe_closesOverlay() {
-        val overlayContext = overlayController.requestWindow()
-        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_OPTIONS_POPUP }
+        val context = getOnUiThread { overlayController.requestWindow() }
+        val overlay = getOnUiThread {
+            TestOverlayView.show(context).apply { type = TYPE_OPTIONS_POPUP }
+        }
 
-        overlayController.updateLauncherDeviceProfile(
-            overlayController.launcherDeviceProfile
-                .toBuilder(overlayContext)
-                .setGestureMode(false)
-                .build()
-        )
+        getInstrumentation().runOnMainSync {
+            overlayController.updateLauncherDeviceProfile(
+                overlayController.launcherDeviceProfile
+                    .toBuilder(context)
+                    .setGestureMode(false)
+                    .build()
+            )
+        }
 
         assertThat(overlay.isOpen).isFalse()
     }
 
     @Test
-    @UiThreadTest
     fun testUpdateLauncherDeviceProfile_overlayRebindSafe_overlayStaysOpen() {
-        val overlayContext = overlayController.requestWindow()
-        val overlay = TestOverlayView.show(overlayContext).apply { type = TYPE_TASKBAR_ALL_APPS }
+        val context = getOnUiThread { overlayController.requestWindow() }
+        val overlay = getOnUiThread {
+            TestOverlayView.show(context).apply { type = TYPE_TASKBAR_ALL_APPS }
+        }
 
-        overlayController.updateLauncherDeviceProfile(
-            overlayController.launcherDeviceProfile
-                .toBuilder(overlayContext)
-                .setGestureMode(false)
-                .build()
-        )
+        getInstrumentation().runOnMainSync {
+            overlayController.updateLauncherDeviceProfile(
+                overlayController.launcherDeviceProfile
+                    .toBuilder(context)
+                    .setGestureMode(false)
+                    .build()
+            )
+        }
 
         assertThat(overlay.isOpen).isTrue()
     }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarModeRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
similarity index 60%
rename from quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarModeRule.kt
rename to quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
index 3b53cdc..c48947e 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarModeRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRule.kt
@@ -14,13 +14,13 @@
  * limitations under the License.
  */
 
-package com.android.launcher3.taskbar
+package com.android.launcher3.taskbar.rules
 
-import com.android.launcher3.taskbar.TaskbarModeRule.Mode
-import com.android.launcher3.taskbar.TaskbarModeRule.TaskbarMode
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode
+import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.MainThreadInitializedObject
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.android.launcher3.util.NavigationMode
 import org.junit.rules.TestRule
 import org.junit.runner.Description
@@ -40,7 +40,7 @@
  * Make sure this rule precedes any rules that depend on [DisplayController], or else the instance
  * might be inconsistent across the test lifecycle.
  */
-class TaskbarModeRule(private val context: SandboxContext) : TestRule {
+class TaskbarModeRule(private val context: TaskbarWindowSandboxContext) : TestRule {
     /** The selected Taskbar mode. */
     enum class Mode {
         TRANSIENT,
@@ -60,23 +60,25 @@
             override fun evaluate() {
                 val mode = taskbarMode.mode
 
-                context.putObject(
-                    DisplayController.INSTANCE,
-                    object : DisplayController(context) {
-                        override fun getInfo(): Info {
-                            return spy(super.getInfo()) {
-                                on { isTransientTaskbar } doReturn (mode == Mode.TRANSIENT)
-                                on { isPinnedTaskbar } doReturn (mode == Mode.PINNED)
-                                on { navigationMode } doReturn
-                                    when (mode) {
-                                        Mode.TRANSIENT,
-                                        Mode.PINNED -> NavigationMode.NO_BUTTON
-                                        Mode.THREE_BUTTONS -> NavigationMode.THREE_BUTTONS
-                                    }
+                getInstrumentation().runOnMainSync {
+                    context.applicationContext.putObject(
+                        DisplayController.INSTANCE,
+                        object : DisplayController(context) {
+                            override fun getInfo(): Info {
+                                return spy(super.getInfo()) {
+                                    on { isTransientTaskbar } doReturn (mode == Mode.TRANSIENT)
+                                    on { isPinnedTaskbar } doReturn (mode == Mode.PINNED)
+                                    on { navigationMode } doReturn
+                                        when (mode) {
+                                            Mode.TRANSIENT,
+                                            Mode.PINNED -> NavigationMode.NO_BUTTON
+                                            Mode.THREE_BUTTONS -> NavigationMode.THREE_BUTTONS
+                                        }
+                                }
                             }
-                        }
-                    },
-                )
+                        },
+                    )
+                }
 
                 base.evaluate()
             }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarModeRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
similarity index 87%
rename from quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarModeRuleTest.kt
rename to quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
index 7dfbb9a..f75e542 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarModeRuleTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarModeRuleTest.kt
@@ -14,17 +14,16 @@
  * limitations under the License.
  */
 
-package com.android.launcher3.taskbar
+package com.android.launcher3.taskbar.rules
 
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
 import com.android.launcher3.InvariantDeviceProfile
-import com.android.launcher3.taskbar.TaskbarModeRule.Mode.PINNED
-import com.android.launcher3.taskbar.TaskbarModeRule.Mode.THREE_BUTTONS
-import com.android.launcher3.taskbar.TaskbarModeRule.Mode.TRANSIENT
-import com.android.launcher3.taskbar.TaskbarModeRule.TaskbarMode
+import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.PINNED
+import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.THREE_BUTTONS
+import com.android.launcher3.taskbar.rules.TaskbarModeRule.Mode.TRANSIENT
+import com.android.launcher3.taskbar.rules.TaskbarModeRule.TaskbarMode
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.LauncherMultivalentJUnit
-import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
 import com.android.launcher3.util.NavigationMode
 import com.google.common.truth.Truth.assertThat
 import org.junit.Rule
@@ -34,7 +33,7 @@
 @RunWith(LauncherMultivalentJUnit::class)
 class TaskbarModeRuleTest {
 
-    private val context = SandboxContext(getInstrumentation().targetContext)
+    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
 
     @get:Rule val taskbarModeRule = TaskbarModeRule(context)
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
similarity index 61%
rename from quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt
rename to quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index bbf738e..8a64949 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -14,23 +14,31 @@
  * limitations under the License.
  */
 
-package com.android.launcher3.taskbar
+package com.android.launcher3.taskbar.rules
 
 import android.app.Instrumentation
 import android.app.PendingIntent
-import android.content.Context
 import android.content.IIntentSender
 import android.content.Intent
+import android.provider.Settings
+import android.provider.Settings.Secure.NAV_BAR_KIDS_MODE
+import android.provider.Settings.Secure.USER_SETUP_COMPLETE
 import androidx.test.platform.app.InstrumentationRegistry
 import androidx.test.rule.ServiceTestRule
 import com.android.launcher3.LauncherAppState
+import com.android.launcher3.taskbar.TaskbarActivityContext
+import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
 import com.android.launcher3.util.LauncherMultivalentJUnit.Companion.isRunningInRobolectric
+import com.android.launcher3.util.ModelTestExtensions.loadModelSync
+import com.android.launcher3.util.TestUtil
 import com.android.quickstep.AllAppsActionManager
 import com.android.quickstep.TouchInteractionService
 import com.android.quickstep.TouchInteractionService.TISBinder
 import org.junit.Assume.assumeTrue
+import org.junit.rules.RuleChain
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
@@ -47,12 +55,11 @@
  * that code that is executed on the main thread in production should also happen on that thread
  * when tested.
  *
- * `@UiThreadTest` is a simple way to run an entire test body on the main thread. But if a test
- * executes code that appends message(s) to the main thread's `MessageQueue`, the annotation will
- * prevent those messages from being processed until after the test body finishes.
+ * `@UiThreadTest` is incompatible with this rule. The annotation causes this rule to run on the
+ * main thread, but it needs to be run on the test thread for it to work properly. Instead, only run
+ * code that requires the main thread using something like [Instrumentation.runOnMainSync] or
+ * [TestUtil.getOnUiThread].
  *
- * To test pending messages, instead use something like [Instrumentation.runOnMainSync] to perform
- * only sections of the test body on the main thread synchronously:
  * ```
  * @Test
  * fun example() {
@@ -62,10 +69,18 @@
  * }
  * ```
  */
-class TaskbarUnitTestRule(private val testInstance: Any, private val context: Context) : TestRule {
+class TaskbarUnitTestRule(
+    private val testInstance: Any,
+    private val context: TaskbarWindowSandboxContext,
+) : TestRule {
+
     private val instrumentation = InstrumentationRegistry.getInstrumentation()
     private val serviceTestRule = ServiceTestRule()
 
+    private val userSetupCompleteRule = TaskbarSecureSettingRule(USER_SETUP_COMPLETE)
+    private val kidsModeRule = TaskbarSecureSettingRule(NAV_BAR_KIDS_MODE)
+    private val settingRules = RuleChain.outerRule(userSetupCompleteRule).around(kidsModeRule)
+
     private lateinit var taskbarManager: TaskbarManager
 
     val activityContext: TaskbarActivityContext
@@ -75,15 +90,34 @@
         }
 
     override fun apply(base: Statement, description: Description): Statement {
+        return settingRules.apply(createStatement(base, description), description)
+    }
+
+    private fun createStatement(base: Statement, description: Description): Statement {
         return object : Statement() {
             override fun evaluate() {
 
+                // Only run test when Taskbar is enabled.
                 instrumentation.runOnMainSync {
                     assumeTrue(
                         LauncherAppState.getIDP(context).getDeviceProfile(context).isTaskbarPresent
                     )
                 }
 
+                // Process secure setting annotations.
+                instrumentation.runOnMainSync {
+                    userSetupCompleteRule.putInt(
+                        if (description.getAnnotation(UserSetupMode::class.java) != null) {
+                            0
+                        } else {
+                            1
+                        }
+                    )
+                    kidsModeRule.putInt(
+                        if (description.getAnnotation(NavBarKidsMode::class.java) != null) 1 else 0
+                    )
+                }
+
                 // Check for existing Taskbar instance from Launcher process.
                 val launcherTaskbarManager: TaskbarManager? =
                     if (!isRunningInRobolectric) {
@@ -100,8 +134,8 @@
                         null
                     }
 
-                instrumentation.runOnMainSync {
-                    taskbarManager =
+                taskbarManager =
+                    TestUtil.getOnUiThread {
                         TaskbarManager(
                             context,
                             AllAppsActionManager(context, UI_HELPER_EXECUTOR) {
@@ -109,12 +143,14 @@
                             },
                             object : TaskbarNavButtonCallbacks {},
                         )
-                }
+                    }
 
                 try {
+                    LauncherAppState.getInstance(context).model.loadModelSync()
+
                     // Replace Launcher Taskbar window with test instance.
                     instrumentation.runOnMainSync {
-                        launcherTaskbarManager?.removeTaskbarRootViewFromWindow()
+                        launcherTaskbarManager?.setSuspended(true)
                         taskbarManager.onUserUnlocked() // Required to complete initialization.
                     }
 
@@ -124,7 +160,7 @@
                     // Revert Taskbar window.
                     instrumentation.runOnMainSync {
                         taskbarManager.destroy()
-                        launcherTaskbarManager?.addTaskbarRootViewToWindow()
+                        launcherTaskbarManager?.setSuspended(false)
                     }
                 }
             }
@@ -133,7 +169,7 @@
 
     /** Simulates Taskbar recreation lifecycle. */
     fun recreateTaskbar() {
-        taskbarManager.recreateTaskbar()
+        instrumentation.runOnMainSync { taskbarManager.recreateTaskbar() }
         injectControllers()
     }
 
@@ -162,4 +198,35 @@
     @Retention(AnnotationRetention.RUNTIME)
     @Target(AnnotationTarget.FIELD)
     annotation class InjectController
+
+    /** Overrides [USER_SETUP_COMPLETE] to be `false` for tests. */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+    annotation class UserSetupMode
+
+    /** Overrides [NAV_BAR_KIDS_MODE] to be `true` for tests. */
+    @Retention(AnnotationRetention.RUNTIME)
+    @Target(AnnotationTarget.CLASS, AnnotationTarget.FUNCTION)
+    annotation class NavBarKidsMode
+
+    /** Rule for Taskbar integer-based secure settings. */
+    private inner class TaskbarSecureSettingRule(private val settingName: String) : TestRule {
+
+        override fun apply(base: Statement, description: Description): Statement {
+            return object : Statement() {
+                override fun evaluate() {
+                    val originalValue =
+                        Settings.Secure.getInt(context.contentResolver, settingName, /* def= */ 0)
+                    try {
+                        base.evaluate()
+                    } finally {
+                        instrumentation.runOnMainSync { putInt(originalValue) }
+                    }
+                }
+            }
+        }
+
+        /** Puts [value] into secure settings under [settingName]. */
+        fun putInt(value: Int) = Settings.Secure.putInt(context.contentResolver, settingName, value)
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
new file mode 100644
index 0000000..234e499
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRuleTest.kt
@@ -0,0 +1,184 @@
+/*
+ * 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.rules
+
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.taskbar.TaskbarActivityContext
+import com.android.launcher3.taskbar.TaskbarKeyguardController
+import com.android.launcher3.taskbar.TaskbarManager
+import com.android.launcher3.taskbar.TaskbarStashController
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.InjectController
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.NavBarKidsMode
+import com.android.launcher3.taskbar.rules.TaskbarUnitTestRule.UserSetupMode
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
+import com.google.common.truth.Truth.assertThat
+import org.junit.Assert.assertThrows
+import org.junit.Test
+import org.junit.runner.Description
+import org.junit.runner.RunWith
+import org.junit.runners.model.Statement
+
+@RunWith(LauncherMultivalentJUnit::class)
+@EmulatedDevices(["pixelFoldable2023"])
+class TaskbarUnitTestRuleTest {
+
+    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
+
+    @Test
+    fun testSetup_taskbarInitialized() {
+        onSetup { assertThat(activityContext).isInstanceOf(TaskbarActivityContext::class.java) }
+    }
+
+    @Test
+    fun testRecreateTaskbar_activityContextChanged() {
+        onSetup {
+            val context1 = activityContext
+            recreateTaskbar()
+            val context2 = activityContext
+            assertThat(context1).isNotSameInstanceAs(context2)
+        }
+    }
+
+    @Test
+    fun testTeardown_taskbarDestroyed() {
+        val testRule = TaskbarUnitTestRule(this, context)
+        testRule.apply(EMPTY_STATEMENT, DESCRIPTION).evaluate()
+        assertThrows(RuntimeException::class.java) { testRule.activityContext }
+    }
+
+    @Test
+    fun testInjectController_validControllerType_isInjected() {
+        val testClass =
+            object {
+                @InjectController lateinit var controller: TaskbarStashController
+                val isInjected: Boolean
+                    get() = ::controller.isInitialized
+            }
+
+        TaskbarUnitTestRule(testClass, context).apply(EMPTY_STATEMENT, DESCRIPTION).evaluate()
+
+        onSetup(TaskbarUnitTestRule(testClass, context)) {
+            assertThat(testClass.isInjected).isTrue()
+        }
+    }
+
+    @Test
+    fun testInjectController_multipleControllers_areInjected() {
+        val testClass =
+            object {
+                @InjectController lateinit var controller1: TaskbarStashController
+                @InjectController lateinit var controller2: TaskbarKeyguardController
+                val areInjected: Boolean
+                    get() = ::controller1.isInitialized && ::controller2.isInitialized
+            }
+
+        onSetup(TaskbarUnitTestRule(testClass, context)) {
+            assertThat(testClass.areInjected).isTrue()
+        }
+    }
+
+    @Test
+    fun testInjectController_invalidControllerType_exceptionThrown() {
+        val testClass =
+            object {
+                @InjectController lateinit var manager: TaskbarManager // Not a controller.
+            }
+
+        // We cannot use #assertThrows because we also catch an assumption violated exception when
+        // running #evaluate on devices that do not support Taskbar.
+        val result =
+            try {
+                TaskbarUnitTestRule(testClass, context)
+                    .apply(EMPTY_STATEMENT, DESCRIPTION)
+                    .evaluate()
+            } catch (e: NoSuchElementException) {
+                e
+            }
+        assertThat(result).isInstanceOf(NoSuchElementException::class.java)
+    }
+
+    @Test
+    fun testInjectController_recreateTaskbar_controllerChanged() {
+        val testClass =
+            object {
+                @InjectController lateinit var controller: TaskbarStashController
+            }
+
+        onSetup(TaskbarUnitTestRule(testClass, context)) {
+            val controller1 = testClass.controller
+            recreateTaskbar()
+            val controller2 = testClass.controller
+            assertThat(controller1).isNotSameInstanceAs(controller2)
+        }
+    }
+
+    @Test
+    fun testUserSetupMode_default_isComplete() {
+        onSetup { assertThat(activityContext.isUserSetupComplete).isTrue() }
+    }
+
+    @Test
+    fun testUserSetupMode_withAnnotation_isIncomplete() {
+        @UserSetupMode class Mode
+        onSetup(description = Description.createSuiteDescription(Mode::class.java)) {
+            assertThat(activityContext.isUserSetupComplete).isFalse()
+        }
+    }
+
+    @Test
+    fun testNavBarKidsMode_default_navBarNotForcedVisible() {
+        onSetup { assertThat(activityContext.isNavBarForceVisible).isFalse() }
+    }
+
+    @Test
+    fun testNavBarKidsMode_withAnnotation_navBarForcedVisible() {
+        @NavBarKidsMode class Mode
+        onSetup(description = Description.createSuiteDescription(Mode::class.java)) {
+            assertThat(activityContext.isNavBarForceVisible).isTrue()
+        }
+    }
+
+    /**
+     * Executes [runTest] after the [testRule] setup phase completes.
+     *
+     * A [description] can also be provided to mimic annotating a test or test class.
+     */
+    private fun onSetup(
+        testRule: TaskbarUnitTestRule = TaskbarUnitTestRule(this, context),
+        description: Description = DESCRIPTION,
+        runTest: TaskbarUnitTestRule.() -> Unit,
+    ) {
+        testRule
+            .apply(
+                object : Statement() {
+                    override fun evaluate() = runTest(testRule)
+                },
+                description,
+            )
+            .evaluate()
+    }
+
+    private companion object {
+        private val EMPTY_STATEMENT =
+            object : Statement() {
+                override fun evaluate() = Unit
+            }
+        private val DESCRIPTION =
+            Description.createSuiteDescription(TaskbarUnitTestRuleTest::class.java)
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
new file mode 100644
index 0000000..321e7a9
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContext.kt
@@ -0,0 +1,61 @@
+/*
+ * 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.rules
+
+import android.content.Context
+import android.content.ContextWrapper
+import android.os.Bundle
+import android.view.Display
+import com.android.launcher3.util.MainThreadInitializedObject
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
+
+/**
+ * Sandbox wrapper where [createWindowContext] provides contexts that are still sandboxed within
+ * [application].
+ *
+ * Taskbar can create window contexts, which need to operate under the same sandbox application, but
+ * [Context.getApplicationContext] by default returns the actual application. For this reason,
+ * [SandboxContext] overrides [getApplicationContext] to return itself, which prevents leaving the
+ * sandbox. [SandboxContext] and the real application have different sets of
+ * [MainThreadInitializedObject] instances, so overriding the application prevents the latter set
+ * from leaking into the sandbox. Similarly, this implementation overrides [getApplicationContext]
+ * to return the original sandboxed [application], and it wraps created windowed contexts to
+ * propagate this [application].
+ */
+class TaskbarWindowSandboxContext
+private constructor(private val application: SandboxContext, base: Context) : ContextWrapper(base) {
+
+    override fun createWindowContext(type: Int, options: Bundle?): Context {
+        return TaskbarWindowSandboxContext(application, super.createWindowContext(type, options))
+    }
+
+    override fun createWindowContext(display: Display, type: Int, options: Bundle?): Context {
+        return TaskbarWindowSandboxContext(
+            application,
+            super.createWindowContext(display, type, options),
+        )
+    }
+
+    override fun getApplicationContext() = application
+
+    companion object {
+        /** Creates a [TaskbarWindowSandboxContext] to sandbox [base] for Taskbar tests. */
+        fun create(base: Context): TaskbarWindowSandboxContext {
+            return SandboxContext(base).let { TaskbarWindowSandboxContext(it, it) }
+        }
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
new file mode 100644
index 0000000..ad4b4de
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarWindowSandboxContextTest.kt
@@ -0,0 +1,44 @@
+/*
+ * 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.rules
+
+import android.view.WindowManager.LayoutParams.TYPE_APPLICATION_OVERLAY
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.MainThreadInitializedObject.SandboxContext
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(LauncherMultivalentJUnit::class)
+class TaskbarWindowSandboxContextTest {
+
+    private val context = TaskbarWindowSandboxContext.create(getInstrumentation().targetContext)
+
+    @Test
+    fun testCreateWindowContext_applicationContextSandboxed() {
+        val windowContext = context.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        assertThat(windowContext.applicationContext).isInstanceOf(SandboxContext::class.java)
+    }
+
+    @Test
+    fun testCreateWindowContext_nested_applicationContextSandboxed() {
+        val windowContext = context.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        val nestedContext = windowContext.createWindowContext(TYPE_APPLICATION_OVERLAY, null)
+        assertThat(nestedContext.applicationContext).isInstanceOf(SandboxContext::class.java)
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
new file mode 100644
index 0000000..40482c4
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
@@ -0,0 +1,160 @@
+/*
+ * 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.quickstep.task.viewmodel
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Bitmap
+import android.graphics.Color
+import android.graphics.Matrix
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
+import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [TaskOverlayViewModel] */
+@RunWith(AndroidJUnit4::class)
+class TaskOverlayViewModelTest {
+    private val task =
+        Task(Task.TaskKey(TASK_ID, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
+            colorBackground = Color.BLACK
+        }
+    private val thumbnailData =
+        ThumbnailData(
+            thumbnail =
+                mock<Bitmap>().apply {
+                    whenever(width).thenReturn(THUMBNAIL_WIDTH)
+                    whenever(height).thenReturn(THUMBNAIL_HEIGHT)
+                }
+        )
+    private val recentsViewData = RecentsViewData()
+    private val tasksRepository = FakeTasksRepository()
+    private val systemUnderTest = TaskOverlayViewModel(task, recentsViewData, tasksRepository)
+
+    @Test
+    fun initialStateIsDisabled() = runTest {
+        assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
+    }
+
+    @Test
+    fun recentsViewOverlayDisabled_Disabled() = runTest {
+        recentsViewData.overlayEnabled.value = false
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+
+        assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
+    }
+
+    @Test
+    fun taskNotFullyVisible_Disabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf()
+
+        assertThat(systemUnderTest.overlayState.first()).isEqualTo(Disabled)
+    }
+
+    @Test
+    fun noThumbnail_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        task.isLocked = false
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = false,
+                    thumbnail = null,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    @Test
+    fun withThumbnail_RealSnapshot_NotLocked_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+        thumbnailData.isRealSnapshot = true
+        task.isLocked = false
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = true,
+                    thumbnail = thumbnailData.thumbnail,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    @Test
+    fun withThumbnail_RealSnapshot_Locked_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+        thumbnailData.isRealSnapshot = true
+        task.isLocked = true
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = false,
+                    thumbnail = thumbnailData.thumbnail,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    @Test
+    fun withThumbnail_FakeSnapshot_Enabled() = runTest {
+        recentsViewData.overlayEnabled.value = true
+        recentsViewData.settledFullyVisibleTaskIds.value = setOf(TASK_ID)
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+        thumbnailData.isRealSnapshot = false
+        task.isLocked = false
+
+        assertThat(systemUnderTest.overlayState.first())
+            .isEqualTo(
+                Enabled(
+                    isRealSnapshot = false,
+                    thumbnail = thumbnailData.thumbnail,
+                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
+                )
+            )
+    }
+
+    companion object {
+        const val TASK_ID = 0
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/DesktopTaskTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/DesktopTaskTest.kt
new file mode 100644
index 0000000..7aed579
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/DesktopTaskTest.kt
@@ -0,0 +1,67 @@
+/*
+ * 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.quickstep.util
+
+import android.content.ComponentName
+import android.content.Intent
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.systemui.shared.recents.model.Task
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(LauncherMultivalentJUnit::class)
+class DesktopTaskTest {
+
+    @Test
+    fun testDesktopTask_sameInstance_isEqual() {
+        val task = DesktopTask(createTasks(1))
+        assertThat(task).isEqualTo(task)
+    }
+
+    @Test
+    fun testDesktopTask_identicalConstructor_isEqual() {
+        val task1 = DesktopTask(createTasks(1))
+        val task2 = DesktopTask(createTasks(1))
+        assertThat(task1).isEqualTo(task2)
+    }
+
+    @Test
+    fun testDesktopTask_copy_isEqual() {
+        val task1 = DesktopTask(createTasks(1))
+        val task2 = task1.copy()
+        assertThat(task1).isEqualTo(task2)
+    }
+
+    @Test
+    fun testDesktopTask_differentId_isNotEqual() {
+        val task1 = DesktopTask(createTasks(1))
+        val task2 = DesktopTask(createTasks(2))
+        assertThat(task1).isNotEqualTo(task2)
+    }
+
+    @Test
+    fun testDesktopTask_differentLength_isNotEqual() {
+        val task1 = DesktopTask(createTasks(1))
+        val task2 = DesktopTask(createTasks(1, 2))
+        assertThat(task1).isNotEqualTo(task2)
+    }
+
+    private fun createTasks(vararg ids: Int): List<Task> {
+        return ids.map { Task(Task.TaskKey(it, 0, Intent(), ComponentName("", ""), 0, 0)) }
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
new file mode 100644
index 0000000..a6d3887
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/util/GroupTaskTest.kt
@@ -0,0 +1,109 @@
+/*
+ * 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.quickstep.util
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Rect
+import com.android.launcher3.util.LauncherMultivalentJUnit
+import com.android.launcher3.util.SplitConfigurationOptions
+import com.android.quickstep.views.TaskView
+import com.android.systemui.shared.recents.model.Task
+import com.android.wm.shell.common.split.SplitScreenConstants
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(LauncherMultivalentJUnit::class)
+class GroupTaskTest {
+
+    @Test
+    fun testGroupTask_sameInstance_isEqual() {
+        val task = GroupTask(createTask(1))
+        assertThat(task).isEqualTo(task)
+    }
+
+    @Test
+    fun testGroupTask_identicalConstructor_isEqual() {
+        val task1 = GroupTask(createTask(1))
+        val task2 = GroupTask(createTask(1))
+        assertThat(task1).isEqualTo(task2)
+    }
+
+    @Test
+    fun testGroupTask_copy_isEqual() {
+        val task1 = GroupTask(createTask(1))
+        val task2 = task1.copy()
+        assertThat(task1).isEqualTo(task2)
+    }
+
+    @Test
+    fun testGroupTask_differentId_isNotEqual() {
+        val task1 = GroupTask(createTask(1))
+        val task2 = GroupTask(createTask(2))
+        assertThat(task1).isNotEqualTo(task2)
+    }
+
+    @Test
+    fun testGroupTask_equalSplitTasks_isEqual() {
+        val splitBounds =
+            SplitConfigurationOptions.SplitBounds(
+                Rect(),
+                Rect(),
+                1,
+                2,
+                SplitScreenConstants.SNAP_TO_50_50
+            )
+        val task1 = GroupTask(createTask(1), createTask(2), splitBounds, TaskView.Type.GROUPED)
+        val task2 = GroupTask(createTask(1), createTask(2), splitBounds, TaskView.Type.GROUPED)
+        assertThat(task1).isEqualTo(task2)
+    }
+
+    @Test
+    fun testGroupTask_differentSplitTasks_isNotEqual() {
+        val splitBounds1 =
+            SplitConfigurationOptions.SplitBounds(
+                Rect(),
+                Rect(),
+                1,
+                2,
+                SplitScreenConstants.SNAP_TO_50_50
+            )
+        val splitBounds2 =
+            SplitConfigurationOptions.SplitBounds(
+                Rect(),
+                Rect(),
+                1,
+                2,
+                SplitScreenConstants.SNAP_TO_30_70
+            )
+        val task1 = GroupTask(createTask(1), createTask(2), splitBounds1, TaskView.Type.GROUPED)
+        val task2 = GroupTask(createTask(1), createTask(2), splitBounds2, TaskView.Type.GROUPED)
+        assertThat(task1).isNotEqualTo(task2)
+    }
+
+    @Test
+    fun testGroupTask_differentType_isNotEqual() {
+        val task1 = GroupTask(createTask(1), null, null, TaskView.Type.SINGLE)
+        val task2 = GroupTask(createTask(1), null, null, TaskView.Type.DESKTOP)
+        assertThat(task1).isNotEqualTo(task2)
+    }
+
+    private fun createTask(id: Int): Task {
+        return Task(Task.TaskKey(id, 0, Intent(), ComponentName("", ""), 0, 0))
+    }
+}
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarHoverToolTipControllerTest.java b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarHoverToolTipControllerTest.java
index 9ed3906..ef3a833 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarHoverToolTipControllerTest.java
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarHoverToolTipControllerTest.java
@@ -39,6 +39,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.BubbleTextView;
+import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.data.FolderInfo;
@@ -66,6 +67,7 @@
     @Mock private MotionEvent mMotionEvent;
     @Mock private BubbleTextView mHoverBubbleTextView;
     @Mock private FolderIcon mHoverFolderIcon;
+    @Mock private AppPairIcon mAppPairIcon;
     @Mock private Display mDisplay;
     @Mock private TaskbarDragLayer mTaskbarDragLayer;
     private Folder mSpyFolderView;
@@ -85,6 +87,7 @@
         when(taskbarActivityContext.getDragLayer()).thenReturn(mTaskbarDragLayer);
         when(taskbarActivityContext.getMainLooper()).thenReturn(context.getMainLooper());
         when(taskbarActivityContext.getDisplay()).thenReturn(mDisplay);
+        when(taskbarActivityContext.isIconAlignedWithHotseat()).thenReturn(false);
 
         when(mTaskbarDragLayer.getChildCount()).thenReturn(1);
         mSpyFolderView = spy(new Folder(new ActivityContextWrapper(context), null));
@@ -213,6 +216,49 @@
         assertThat(hoverHandled).isFalse();
     }
 
+    @Test
+    public void onHover_hoverEnterAppPair_revealToolTip() {
+        when(mMotionEvent.getAction()).thenReturn(MotionEvent.ACTION_HOVER_ENTER);
+        when(mMotionEvent.getActionMasked()).thenReturn(MotionEvent.ACTION_HOVER_ENTER);
+
+        boolean hoverHandled =
+                mTaskbarHoverToolTipController.onHover(mAppPairIcon, mMotionEvent);
+        waitForIdleSync();
+
+        assertThat(hoverHandled).isTrue();
+        verify(taskbarActivityContext).setAutohideSuspendFlag(FLAG_AUTOHIDE_SUSPEND_HOVERING_ICONS,
+                true);
+    }
+
+    @Test
+    public void onHover_hoverExitAppPair_closeToolTip() {
+        when(mMotionEvent.getAction()).thenReturn(MotionEvent.ACTION_HOVER_EXIT);
+        when(mMotionEvent.getActionMasked()).thenReturn(MotionEvent.ACTION_HOVER_EXIT);
+
+        boolean hoverHandled =
+                mTaskbarHoverToolTipController.onHover(mAppPairIcon, mMotionEvent);
+        waitForIdleSync();
+
+        assertThat(hoverHandled).isTrue();
+        verify(taskbarActivityContext).setAutohideSuspendFlag(FLAG_AUTOHIDE_SUSPEND_HOVERING_ICONS,
+                false);
+    }
+
+    @Test
+    public void onHover_hoverEnterIconAlignedWithHotseat_noReveal() {
+        when(mMotionEvent.getAction()).thenReturn(MotionEvent.ACTION_HOVER_ENTER);
+        when(mMotionEvent.getActionMasked()).thenReturn(MotionEvent.ACTION_HOVER_ENTER);
+        when(taskbarActivityContext.isIconAlignedWithHotseat()).thenReturn(true);
+
+        boolean hoverHandled =
+                mTaskbarHoverToolTipController.onHover(mHoverBubbleTextView, mMotionEvent);
+        waitForIdleSync();
+
+        assertThat(hoverHandled).isTrue();
+        verify(taskbarActivityContext).setAutohideSuspendFlag(FLAG_AUTOHIDE_SUSPEND_HOVERING_ICONS,
+                true);
+    }
+
     private void waitForIdleSync() {
         mTestableLooper.processAllMessages();
     }
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
index 486dc68..13c4f72 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
@@ -44,6 +44,7 @@
 import org.mockito.junit.MockitoJUnit
 import org.mockito.kotlin.any
 import org.mockito.kotlin.doAnswer
+import org.mockito.kotlin.times
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 
@@ -56,7 +57,6 @@
     @Mock private lateinit var mockRecentsModel: RecentsModel
     @Mock private lateinit var mockDesktopVisibilityController: DesktopVisibilityController
 
-    private var nextTaskId: Int = 500
     private var taskListChangeId: Int = 1
 
     private lateinit var recentAppsController: TaskbarRecentAppsController
@@ -478,6 +478,82 @@
         assertThat(shownPackages).containsExactlyElementsIn(expectedPackages)
     }
 
+    @Test
+    fun onRecentTasksChanged_notInDesktopMode_noActualChangeToRecents_commitRunningAppsToUI_notCalled() {
+        setInDesktopMode(false)
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = emptyList(),
+            runningTaskPackages = emptyList(),
+            recentTaskPackages = listOf(RECENT_PACKAGE_1, RECENT_PACKAGE_2)
+        )
+        verify(taskbarViewController, times(1)).commitRunningAppsToUI()
+        // Call onRecentTasksChanged() again with the same tasks, verify it's a no-op.
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = emptyList(),
+            runningTaskPackages = emptyList(),
+            recentTaskPackages = listOf(RECENT_PACKAGE_1, RECENT_PACKAGE_2)
+        )
+        verify(taskbarViewController, times(1)).commitRunningAppsToUI()
+    }
+
+    @Test
+    fun onRecentTasksChanged_inDesktopMode_noActualChangeToRunning_commitRunningAppsToUI_notCalled() {
+        setInDesktopMode(true)
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = emptyList(),
+            runningTaskPackages = listOf(RUNNING_APP_PACKAGE_1, RUNNING_APP_PACKAGE_2),
+            recentTaskPackages = emptyList()
+        )
+        verify(taskbarViewController, times(1)).commitRunningAppsToUI()
+        // Call onRecentTasksChanged() again with the same tasks, verify it's a no-op.
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = emptyList(),
+            runningTaskPackages = listOf(RUNNING_APP_PACKAGE_1, RUNNING_APP_PACKAGE_2),
+            recentTaskPackages = emptyList()
+        )
+        verify(taskbarViewController, times(1)).commitRunningAppsToUI()
+    }
+
+    @Test
+    fun onRecentTasksChanged_onlyMinimizedChanges_commitRunningAppsToUI_isCalled() {
+        setInDesktopMode(true)
+        val runningTasks = listOf(RUNNING_APP_PACKAGE_1, RUNNING_APP_PACKAGE_2)
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = emptyList(),
+            runningTaskPackages = runningTasks,
+            minimizedTaskIndices = setOf(0),
+            recentTaskPackages = emptyList()
+        )
+        verify(taskbarViewController, times(1)).commitRunningAppsToUI()
+        // Call onRecentTasksChanged() again with a new minimized app, verify we update UI.
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = emptyList(),
+            runningTaskPackages = runningTasks,
+            minimizedTaskIndices = setOf(0, 1),
+            recentTaskPackages = emptyList()
+        )
+        verify(taskbarViewController, times(2)).commitRunningAppsToUI()
+    }
+
+    @Test
+    fun onRecentTasksChanged_hotseatAppStartsRunning_commitRunningAppsToUI_isCalled() {
+        setInDesktopMode(true)
+        val hotseatPackages = listOf(HOTSEAT_PACKAGE_1, HOTSEAT_PACKAGE_2)
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = hotseatPackages,
+            runningTaskPackages = listOf(RUNNING_APP_PACKAGE_1),
+            recentTaskPackages = emptyList()
+        )
+        verify(taskbarViewController, times(1)).commitRunningAppsToUI()
+        // Call onRecentTasksChanged() again with a new running app, verify we update UI.
+        prepareHotseatAndRunningAndRecentApps(
+            hotseatPackages = hotseatPackages,
+            runningTaskPackages = listOf(RUNNING_APP_PACKAGE_1, HOTSEAT_PACKAGE_1),
+            recentTaskPackages = emptyList()
+        )
+        verify(taskbarViewController, times(2)).commitRunningAppsToUI()
+    }
+
     private fun prepareHotseatAndRunningAndRecentApps(
         hotseatPackages: List<String>,
         runningTaskPackages: List<String>,
@@ -556,9 +632,11 @@
     }
 
     private fun createTask(packageName: String, isVisible: Boolean = true): Task {
+        // Use the number at the end of the test packageName as the id.
+        val id = packageName[packageName.length - 1].code
         return Task(
                 Task.TaskKey(
-                    nextTaskId++,
+                    id,
                     WINDOWING_MODE_FREEFORM,
                     Intent().apply { `package` = packageName },
                     ComponentName(packageName, "TestActivity"),
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index d905801..5c052b2 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -771,6 +771,7 @@
             // initialized properly.
             onSaveInstanceState(new Bundle());
             mModel.rebindCallbacks();
+            updateDisallowBack();
         } finally {
             Trace.endSection();
         }
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 19a3002..a296f46 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -381,6 +381,28 @@
     }
 
     /**
+     * Scales a {@code RectF} in place about a specified pivot point.
+     *
+     * <p>This method modifies the given {@code RectF} directly to scale it proportionally
+     * by the given {@code scale}, while preserving its center at the specified
+     * {@code (pivotX, pivotY)} coordinates.
+     *
+     * @param rectF the {@code RectF} to scale, modified directly.
+     * @param pivotX the x-coordinate of the pivot point about which to scale.
+     * @param pivotY the y-coordinate of the pivot point about which to scale.
+     * @param scale the factor by which to scale the rectangle. Values less than 1 will
+     *                    shrink the rectangle, while values greater than 1 will enlarge it.
+     */
+    public static void scaleRectFAboutPivot(RectF rectF, float pivotX, float pivotY, float scale) {
+        rectF.offset(-pivotX, -pivotY);
+        rectF.left *= scale;
+        rectF.top *= scale;
+        rectF.right *= scale;
+        rectF.bottom *= scale;
+        rectF.offset(pivotX, pivotY);
+    }
+
+    /**
      * Maps t from one range to another range.
      * @param t The value to map.
      * @param fromMin The lower bound of the range that t is being mapped from.
diff --git a/src/com/android/launcher3/apppairs/AppPairIcon.java b/src/com/android/launcher3/apppairs/AppPairIcon.java
index 32445ec..870c876 100644
--- a/src/com/android/launcher3/apppairs/AppPairIcon.java
+++ b/src/com/android/launcher3/apppairs/AppPairIcon.java
@@ -18,10 +18,12 @@
 
 import static com.android.launcher3.BubbleTextView.DISPLAY_FOLDER;
 
+import android.animation.ObjectAnimator;
 import android.content.Context;
 import android.graphics.Paint;
 import android.graphics.Rect;
 import android.util.AttributeSet;
+import android.util.FloatProperty;
 import android.view.LayoutInflater;
 import android.view.ViewGroup;
 import android.widget.FrameLayout;
@@ -54,6 +56,26 @@
 public class AppPairIcon extends FrameLayout implements DraggableView, Reorderable {
     private static final String TAG = "AppPairIcon";
 
+    // The duration of the scaling animation on hover enter/exit.
+    private static final int HOVER_SCALE_DURATION = 150;
+    // The default scale of the icon when not hovered.
+    private static final Float HOVER_SCALE_DEFAULT = 1f;
+    // The max scale of the icon when hovered.
+    private static final Float HOVER_SCALE_MAX = 1.1f;
+    // Animates the scale of the icon background on hover.
+    private static final FloatProperty<AppPairIcon> HOVER_SCALE_PROPERTY =
+            new FloatProperty<>("hoverScale") {
+                @Override
+                public void setValue(AppPairIcon view, float scale) {
+                    view.mIconGraphic.setHoverScale(scale);
+                }
+
+                @Override
+                public Float get(AppPairIcon view) {
+                    return view.mIconGraphic.getHoverScale();
+                }
+            };
+
     // A view that holds the app pair icon graphic.
     private AppPairIconGraphic mIconGraphic;
     // A view that holds the app pair's title.
@@ -250,4 +272,14 @@
         }
         super.onMeasure(widthMeasureSpec, heightMeasureSpec);
     }
+
+    @Override
+    public void onHoverChanged(boolean hovered) {
+        super.onHoverChanged(hovered);
+        ObjectAnimator
+                .ofFloat(this, HOVER_SCALE_PROPERTY,
+                        hovered ? HOVER_SCALE_MAX : HOVER_SCALE_DEFAULT)
+                .setDuration(HOVER_SCALE_DURATION)
+                .start();
+    }
 }
diff --git a/src/com/android/launcher3/apppairs/AppPairIconDrawable.java b/src/com/android/launcher3/apppairs/AppPairIconDrawable.java
index db83d91..114ed2e 100644
--- a/src/com/android/launcher3/apppairs/AppPairIconDrawable.java
+++ b/src/com/android/launcher3/apppairs/AppPairIconDrawable.java
@@ -26,6 +26,7 @@
 
 import androidx.annotation.NonNull;
 
+import com.android.launcher3.Utilities;
 import com.android.launcher3.icons.FastBitmapDrawable;
 
 /**
@@ -128,6 +129,18 @@
                 height - (mP.getStandardIconPadding() + mP.getOuterPadding())
         );
 
+        // Scale each background from its center edge closest to the center channel.
+        Utilities.scaleRectFAboutPivot(
+                leftSide,
+                leftSide.left + leftSide.width(),
+                leftSide.top + leftSide.centerY(),
+                mP.getHoverScale());
+        Utilities.scaleRectFAboutPivot(
+                rightSide,
+                rightSide.left,
+                rightSide.top + rightSide.centerY(),
+                mP.getHoverScale());
+
         drawCustomRoundedRect(canvas, leftSide, new float[]{
                 mP.getBigRadius(), mP.getBigRadius(),
                 mP.getSmallRadius(), mP.getSmallRadius(),
@@ -163,6 +176,18 @@
                 height - (mP.getStandardIconPadding() + mP.getOuterPadding())
         );
 
+        // Scale each background from its center edge closest to the center channel.
+        Utilities.scaleRectFAboutPivot(
+                topSide,
+                topSide.left + topSide.centerX(),
+                topSide.top + topSide.height(),
+                mP.getHoverScale());
+        Utilities.scaleRectFAboutPivot(
+                bottomSide,
+                bottomSide.left + bottomSide.centerX(),
+                bottomSide.top,
+                mP.getHoverScale());
+
         drawCustomRoundedRect(canvas, topSide, new float[]{
                 mP.getBigRadius(), mP.getBigRadius(),
                 mP.getBigRadius(), mP.getBigRadius(),
diff --git a/src/com/android/launcher3/apppairs/AppPairIconDrawingParams.kt b/src/com/android/launcher3/apppairs/AppPairIconDrawingParams.kt
index 45dc013..5b546d6 100644
--- a/src/com/android/launcher3/apppairs/AppPairIconDrawingParams.kt
+++ b/src/com/android/launcher3/apppairs/AppPairIconDrawingParams.kt
@@ -64,6 +64,8 @@
     var isLeftRightSplit: Boolean = true
     // The background paint color (based on container).
     var bgColor: Int = 0
+    // The scale of the icon background while hovered.
+    var hoverScale: Float = 1f
 
     init {
         val activity: ActivityContext = ActivityContext.lookupContext(context)
diff --git a/src/com/android/launcher3/apppairs/AppPairIconGraphic.kt b/src/com/android/launcher3/apppairs/AppPairIconGraphic.kt
index dce97eb..034b686 100644
--- a/src/com/android/launcher3/apppairs/AppPairIconGraphic.kt
+++ b/src/com/android/launcher3/apppairs/AppPairIconGraphic.kt
@@ -139,4 +139,19 @@
         super.dispatchDraw(canvas)
         drawable.draw(canvas)
     }
+
+    /**
+     * Sets the scale of the icon background while hovered.
+     */
+    fun setHoverScale(scale: Float) {
+        drawParams.hoverScale = scale
+        redraw()
+    }
+
+    /**
+     * Gets the scale of the icon background while hovered.
+     */
+    fun getHoverScale(): Float {
+        return drawParams.hoverScale
+    }
 }
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index dcc55e6..d3c1a02 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -65,6 +65,7 @@
 
 import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import androidx.core.content.res.ResourcesCompat;
 
 import com.android.launcher3.AbstractFloatingView;
@@ -1693,6 +1694,11 @@
         return windowBottomPx - folderBottomPx;
     }
 
+    @VisibleForTesting
+    public boolean getDeleteFolderOnDropCompleted() {
+        return mDeleteFolderOnDropCompleted;
+    }
+
     /**
      * Save this listener for the special case of when we update the state and concurrently
      * add another listener to {@link #mOnFolderStateChangedListeners} to avoid a
diff --git a/src/com/android/launcher3/widget/WidgetCell.java b/src/com/android/launcher3/widget/WidgetCell.java
index 35372d3..b7ad95e 100644
--- a/src/com/android/launcher3/widget/WidgetCell.java
+++ b/src/com/android/launcher3/widget/WidgetCell.java
@@ -40,7 +40,6 @@
 import android.view.View;
 import android.view.ViewGroup;
 import android.view.ViewPropertyAnimator;
-import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.Button;
 import android.widget.FrameLayout;
 import android.widget.LinearLayout;
@@ -500,12 +499,6 @@
     }
 
     @Override
-    public void onInitializeAccessibilityNodeInfo(AccessibilityNodeInfo info) {
-        super.onInitializeAccessibilityNodeInfo(info);
-        info.removeAction(AccessibilityNodeInfo.AccessibilityAction.ACTION_CLICK);
-    }
-
-    @Override
     protected void onMeasure(int widthMeasureSpec, int heightMeasureSpec) {
         ViewGroup.LayoutParams containerLp = mWidgetImageContainer.getLayoutParams();
         int maxWidth = MeasureSpec.getSize(widthMeasureSpec);
diff --git a/tests/src/com/android/launcher3/folder/FolderTest.kt b/tests/src/com/android/launcher3/folder/FolderTest.kt
new file mode 100644
index 0000000..e1daa74
--- /dev/null
+++ b/tests/src/com/android/launcher3/folder/FolderTest.kt
@@ -0,0 +1,86 @@
+/*
+ * 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.folder
+
+import android.content.Context
+import android.graphics.Point
+import android.view.View
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.DropTarget.DragObject
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.celllayout.board.FolderPoint
+import com.android.launcher3.celllayout.board.TestWorkspaceBuilder
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.ModelTestExtensions.clearModelDb
+import junit.framework.TestCase.assertEquals
+import org.junit.After
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
+
+/** Tests for [Folder] */
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class FolderTest {
+
+    private val context: Context =
+        ActivityContextWrapper(ApplicationProvider.getApplicationContext())
+    private val workspaceBuilder = TestWorkspaceBuilder(context)
+    private val folder: Folder = Mockito.spy(Folder(context, null))
+
+    @After
+    fun tearDown() {
+        LauncherAppState.getInstance(context).model.clearModelDb()
+    }
+
+    @Test
+    fun `Undo a folder with 1 icon when onDropCompleted is called`() {
+        val folderInfo =
+            workspaceBuilder.createFolderInCell(FolderPoint(Point(1, 0), TWO_ICON_FOLDER_TYPE), 0)
+        folder.mInfo = folderInfo
+        folder.mInfo.getContents().removeAt(0)
+        folder.mContent = Mockito.mock(FolderPagedView::class.java)
+        val dragLayout = Mockito.mock(View::class.java)
+        val dragObject = Mockito.mock(DragObject::class.java)
+        assertEquals(folder.deleteFolderOnDropCompleted, false)
+        folder.onDropCompleted(dragLayout, dragObject, true)
+        verify(folder, times(1)).replaceFolderWithFinalItem()
+        assertEquals(folder.deleteFolderOnDropCompleted, false)
+    }
+
+    @Test
+    fun `Do not undo a folder with 2 icons when onDropCompleted is called`() {
+        val folderInfo =
+            workspaceBuilder.createFolderInCell(FolderPoint(Point(1, 0), TWO_ICON_FOLDER_TYPE), 0)
+        folder.mInfo = folderInfo
+        folder.mContent = Mockito.mock(FolderPagedView::class.java)
+        val dragLayout = Mockito.mock(View::class.java)
+        val dragObject = Mockito.mock(DragObject::class.java)
+        assertEquals(folder.deleteFolderOnDropCompleted, false)
+        folder.onDropCompleted(dragLayout, dragObject, true)
+        verify(folder, times(0)).replaceFolderWithFinalItem()
+        assertEquals(folder.deleteFolderOnDropCompleted, false)
+    }
+
+    companion object {
+        const val TWO_ICON_FOLDER_TYPE = 'A'
+    }
+}
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
index 28d1faa..d40d3bc 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
@@ -23,8 +23,6 @@
 import static com.android.launcher3.provider.LauncherDbUtils.itemIdMatch;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.WidgetUtils.createWidgetInfo;
-import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
-import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -56,7 +54,6 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
 import com.android.launcher3.util.rule.ShellCommandRule;
-import com.android.launcher3.util.rule.TestStabilityRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
 
@@ -143,7 +140,6 @@
     }
 
     @Test
-    @TestStabilityRule.Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/310242894
     public void testPendingWidget_withConfigScreen() {
         // A non-restored widget with config screen get bound and shows a 'Click to setup' UI.
         // Do not bind the widget
@@ -193,7 +189,6 @@
     }
 
     @Test
-    @TestStabilityRule.Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/310242894
     public void testPendingWidget_notRestored_brokenInstall() {
         // A widget which is was being installed once, even if its not being
         // installed at the moment is not removed.
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
index d653317..5dee322 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplThemeIconsTest.java
@@ -16,6 +16,8 @@
 package com.android.launcher3.ui.workspace;
 
 import static com.android.launcher3.util.TestConstants.AppNames.TEST_APP_NAME;
+import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
+import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -38,6 +40,7 @@
 import com.android.launcher3.tapl.HomeAppIconMenuItem;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.util.Executors;
+import com.android.launcher3.util.rule.TestStabilityRule;
 
 import org.junit.Test;
 
@@ -111,6 +114,7 @@
     }
 
     @Test
+    @TestStabilityRule.Stability(flavors = LOCAL | PLATFORM_POSTSUBMIT) // b/350557998
     public void testShortcutIconWithTheme() throws Exception {
         setThemeEnabled(true);
         initialize(this);
diff --git a/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
index ae24a57..bc26c00 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
@@ -352,8 +352,11 @@
     }
 
     private void assertPagesExist(Launcher launcher, int... pageIds) {
+        waitForLauncherCondition("Existing page count does NOT match. "
+                + "Expected: " + pageIds.length
+                + ". Actual: " + launcher.getWorkspace().getPageCount(),
+                l -> pageIds.length == l.getWorkspace().getPageCount());
         int pageCount = launcher.getWorkspace().getPageCount();
-        assertEquals("Existing page count does NOT match.", pageIds.length, pageCount);
         for (int i = 0; i < pageCount; i++) {
             CellLayout page = (CellLayout) launcher.getWorkspace().getPageAt(i);
             int pageId = launcher.getWorkspace().getCellLayoutId(page);