diff --git a/quickstep/res/layout/task_thumbnail.xml b/quickstep/res/layout/task_thumbnail.xml
index 34640e6..b1fe89e 100644
--- a/quickstep/res/layout/task_thumbnail.xml
+++ b/quickstep/res/layout/task_thumbnail.xml
@@ -23,6 +23,7 @@
         android:layout_width="match_parent"
         android:layout_height="match_parent"
         android:importantForAccessibility="no"
+        android:scaleType="matrix"
         android:visibility="gone"/>
 
     <com.android.quickstep.task.thumbnail.LiveTileView
diff --git a/quickstep/res/values/styles.xml b/quickstep/res/values/styles.xml
index 952505a..153bc99 100644
--- a/quickstep/res/values/styles.xml
+++ b/quickstep/res/values/styles.xml
@@ -312,9 +312,15 @@
         <item name="android:lineHeight">20sp</item>
     </style>
 
-    <style name="WidgetPickerActivityTheme" parent="@android:style/Theme.Translucent.NoTitleBar">
-        <item name="widgetsTheme">@style/WidgetContainerTheme</item>
+    <style name="WidgetPickerActivityTheme" parent="@android:style/Theme.DeviceDefault.DayNight">
         <item name="android:windowBackground">@android:color/transparent</item>
+        <item name="android:windowNoTitle">true</item>
+        <item name="android:windowContentOverlay">@null</item>
+        <item name="android:colorBackgroundCacheHint">@null</item>
+        <item name="android:windowIsTranslucent">true</item>
+        <item name="android:windowAnimationStyle">@android:style/Animation</item>
+
+        <item name="widgetsTheme">@style/WidgetContainerTheme</item>
         <item name="pageIndicatorDotColor">@color/page_indicator_dot_color_light</item>
     </style>
 </resources>
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index 5a74f4a..5e5487b 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -1644,7 +1644,15 @@
     }
 
     private void addCujInstrumentation(Animator anim, int cuj) {
-        anim.addListener(new AnimationSuccessListener() {
+        anim.addListener(getCujAnimationSuccessListener(cuj));
+    }
+
+    private void addCujInstrumentation(RectFSpringAnim anim, int cuj) {
+        anim.addAnimatorListener(getCujAnimationSuccessListener(cuj));
+    }
+
+    private AnimationSuccessListener getCujAnimationSuccessListener(int cuj) {
+        return new AnimationSuccessListener() {
             @Override
             public void onAnimationStart(Animator animation) {
                 mDragLayer.getViewTreeObserver().addOnDrawListener(
@@ -1678,7 +1686,7 @@
             public void onAnimationSuccess(Animator animator) {
                 InteractionJankMonitorWrapper.end(cuj);
             }
-        });
+        };
     }
 
     /**
@@ -1759,10 +1767,6 @@
         // invisibility on touch down, and only reset it after the animation to home
         // is initialized.
         if (launcherIsForceInvisibleOrOpening || fromPredictiveBack) {
-            addCujInstrumentation(anim, playFallBackAnimation
-                    ? Cuj.CUJ_LAUNCHER_APP_CLOSE_TO_HOME_FALLBACK
-                    : Cuj.CUJ_LAUNCHER_APP_CLOSE_TO_HOME);
-
             AnimatorListenerAdapter endListener = new AnimatorListenerAdapter() {
                 @Override
                 public void onAnimationEnd(Animator animation) {
@@ -1772,6 +1776,14 @@
                 }
             };
 
+            if (rectFSpringAnim != null && anim.getChildAnimations().isEmpty()) {
+                addCujInstrumentation(rectFSpringAnim, Cuj.CUJ_LAUNCHER_APP_CLOSE_TO_HOME);
+            } else {
+                addCujInstrumentation(anim, playFallBackAnimation
+                        ? Cuj.CUJ_LAUNCHER_APP_CLOSE_TO_HOME_FALLBACK
+                        : Cuj.CUJ_LAUNCHER_APP_CLOSE_TO_HOME);
+            }
+
             if (fromPredictiveBack && rectFSpringAnim != null) {
                 rectFSpringAnim.addAnimatorListener(endListener);
             } else {
diff --git a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
index 44d8a5c..d0be2f3 100644
--- a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
+++ b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
@@ -28,6 +28,7 @@
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ClipData;
 import android.content.ClipDescription;
+import android.content.Context;
 import android.content.Intent;
 import android.os.Bundle;
 import android.util.Log;
@@ -60,6 +61,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
@@ -128,9 +130,23 @@
     /** A set of user ids that should be filtered out from the selected widgets. */
     @NonNull
     Set<Integer> mFilteredUserIds = new HashSet<>();
+
     @Nullable
     private WidgetsFullSheet mWidgetSheet;
 
+    private final Predicate<WidgetItem> mWidgetsFilter = widget -> {
+        final WidgetAcceptabilityVerdict verdict =
+                isWidgetAcceptable(widget, /* applySizeFilter=*/ false);
+        verdict.maybeLogVerdict();
+        return verdict.isAcceptable;
+    };
+    private final Predicate<WidgetItem> mDefaultWidgetsFilter = widget -> {
+        final WidgetAcceptabilityVerdict verdict =
+                isWidgetAcceptable(widget, /* applySizeFilter=*/ true);
+        verdict.maybeLogVerdict();
+        return verdict.isAcceptable;
+    };
+
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
@@ -276,18 +292,15 @@
     private void refreshAndBindWidgets() {
         MODEL_EXECUTOR.execute(() -> {
             LauncherAppState app = LauncherAppState.getInstance(this);
+            Context context = app.getContext();
+
             mModel.update(app, null);
             final List<WidgetsListBaseEntry> allWidgets =
-                    mModel.getFilteredWidgetsListForPicker(
-                            app.getContext(),
-                            /*widgetItemFilter=*/ widget -> {
-                                final WidgetAcceptabilityVerdict verdict =
-                                        isWidgetAcceptable(widget);
-                                verdict.maybeLogVerdict();
-                                return verdict.isAcceptable;
-                            }
-                    );
-            bindWidgets(allWidgets);
+                    mModel.getFilteredWidgetsListForPicker(context, mWidgetsFilter);
+            final List<WidgetsListBaseEntry> defaultWidgets =
+                    shouldShowDefaultWidgets() ? mModel.getFilteredWidgetsListForPicker(context,
+                            mDefaultWidgetsFilter) : List.of();
+            bindWidgets(allWidgets, defaultWidgets);
             // Open sheet once widgets are available, so that it doesn't interrupt the open
             // animation.
             openWidgetsSheet();
@@ -307,8 +320,9 @@
         });
     }
 
-    private void bindWidgets(List<WidgetsListBaseEntry> widgets) {
-        MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setAllWidgets(widgets));
+    private void bindWidgets(List<WidgetsListBaseEntry> allWidgets,
+            List<WidgetsListBaseEntry> defaultWidgets) {
+        MAIN_EXECUTOR.execute(() -> mPopupDataProvider.setAllWidgets(allWidgets, defaultWidgets));
     }
 
    private void openWidgetsSheet() {
@@ -380,7 +394,13 @@
         }
     };
 
-    private WidgetAcceptabilityVerdict isWidgetAcceptable(WidgetItem widget) {
+    private boolean shouldShowDefaultWidgets() {
+        // If optional filters such as size filter are present, we display them as default widgets.
+        return mDesiredWidgetWidth != 0 || mDesiredWidgetHeight != 0;
+    }
+
+    private WidgetAcceptabilityVerdict isWidgetAcceptable(WidgetItem widget,
+            boolean applySizeFilter) {
         final AppWidgetProviderInfo info = widget.widgetInfo;
         if (info == null) {
             return rejectWidget(widget, "shortcut");
@@ -401,61 +421,63 @@
                     info.widgetCategory);
         }
 
-        if (mDesiredWidgetWidth == 0 && mDesiredWidgetHeight == 0) {
-            // Accept the widget if the desired dimensions are unspecified.
-            return acceptWidget(widget);
-        }
-
-        final boolean isHorizontallyResizable =
-                (info.resizeMode & AppWidgetProviderInfo.RESIZE_HORIZONTAL) != 0;
-        if (mDesiredWidgetWidth > 0 && isHorizontallyResizable) {
-            if (info.maxResizeWidth > 0
-                    && info.maxResizeWidth >= info.minWidth
-                    && info.maxResizeWidth < mDesiredWidgetWidth) {
-                return rejectWidget(
-                        widget,
-                        "maxResizeWidth[%d] < mDesiredWidgetWidth[%d]",
-                        info.maxResizeWidth,
-                        mDesiredWidgetWidth);
+        if (applySizeFilter) {
+            if (mDesiredWidgetWidth == 0 && mDesiredWidgetHeight == 0) {
+                // Accept the widget if the desired dimensions are unspecified.
+                return acceptWidget(widget);
             }
 
-            final int minWidth = Math.min(info.minResizeWidth, info.minWidth);
-            if (minWidth > mDesiredWidgetWidth) {
-                return rejectWidget(
-                        widget,
-                        "min(minWidth[%d], minResizeWidth[%d]) > mDesiredWidgetWidth[%d]",
-                        info.minWidth,
-                        info.minResizeWidth,
-                        mDesiredWidgetWidth);
-            }
-        }
+            final boolean isHorizontallyResizable =
+                    (info.resizeMode & AppWidgetProviderInfo.RESIZE_HORIZONTAL) != 0;
+            if (mDesiredWidgetWidth > 0 && isHorizontallyResizable) {
+                if (info.maxResizeWidth > 0
+                        && info.maxResizeWidth >= info.minWidth
+                        && info.maxResizeWidth < mDesiredWidgetWidth) {
+                    return rejectWidget(
+                            widget,
+                            "maxResizeWidth[%d] < mDesiredWidgetWidth[%d]",
+                            info.maxResizeWidth,
+                            mDesiredWidgetWidth);
+                }
 
-        final boolean isVerticallyResizable =
-                (info.resizeMode & AppWidgetProviderInfo.RESIZE_VERTICAL) != 0;
-        if (mDesiredWidgetHeight > 0 && isVerticallyResizable) {
-            if (info.maxResizeHeight > 0
-                    && info.maxResizeHeight >= info.minHeight
-                    && info.maxResizeHeight < mDesiredWidgetHeight) {
-                return rejectWidget(
-                        widget,
-                        "maxResizeHeight[%d] < mDesiredWidgetHeight[%d]",
-                        info.maxResizeHeight,
-                        mDesiredWidgetHeight);
+                final int minWidth = Math.min(info.minResizeWidth, info.minWidth);
+                if (minWidth > mDesiredWidgetWidth) {
+                    return rejectWidget(
+                            widget,
+                            "min(minWidth[%d], minResizeWidth[%d]) > mDesiredWidgetWidth[%d]",
+                            info.minWidth,
+                            info.minResizeWidth,
+                            mDesiredWidgetWidth);
+                }
             }
 
-            final int minHeight = Math.min(info.minResizeHeight, info.minHeight);
-            if (minHeight > mDesiredWidgetHeight) {
-                return rejectWidget(
-                        widget,
-                        "min(minHeight[%d], minResizeHeight[%d]) > mDesiredWidgetHeight[%d]",
-                        info.minHeight,
-                        info.minResizeHeight,
-                        mDesiredWidgetHeight);
-            }
-        }
+            final boolean isVerticallyResizable =
+                    (info.resizeMode & AppWidgetProviderInfo.RESIZE_VERTICAL) != 0;
+            if (mDesiredWidgetHeight > 0 && isVerticallyResizable) {
+                if (info.maxResizeHeight > 0
+                        && info.maxResizeHeight >= info.minHeight
+                        && info.maxResizeHeight < mDesiredWidgetHeight) {
+                    return rejectWidget(
+                            widget,
+                            "maxResizeHeight[%d] < mDesiredWidgetHeight[%d]",
+                            info.maxResizeHeight,
+                            mDesiredWidgetHeight);
+                }
 
-        if (!isHorizontallyResizable || !isVerticallyResizable) {
-            return rejectWidget(widget, "not resizeable");
+                final int minHeight = Math.min(info.minResizeHeight, info.minHeight);
+                if (minHeight > mDesiredWidgetHeight) {
+                    return rejectWidget(
+                            widget,
+                            "min(minHeight[%d], minResizeHeight[%d]) > mDesiredWidgetHeight[%d]",
+                            info.minHeight,
+                            info.minResizeHeight,
+                            mDesiredWidgetHeight);
+                }
+            }
+
+            if (!isHorizontallyResizable || !isVerticallyResizable) {
+                return rejectWidget(widget, "not resizeable");
+            }
         }
 
         return acceptWidget(widget);
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 64bb05e..0395d32 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -65,7 +65,7 @@
                 Collectors.toSet());
         Predicate<WidgetItem> notOnWorkspace = w -> !widgetsInWorkspace.contains(w);
         Map<ComponentKey, WidgetItem> allWidgets =
-                dataModel.widgetsModel.getAllWidgetComponentsWithoutShortcuts();
+                dataModel.widgetsModel.getWidgetsByComponentKey();
 
         List<WidgetItem> servicePredictedItems = new ArrayList<>();
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 5c1e638..ea2adcf 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -48,7 +48,6 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SHORTCUT_HELPER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
-import static com.android.wm.shell.Flags.enableTaskbarOnPhones;
 
 import android.animation.ArgbEvaluator;
 import android.animation.ObjectAnimator;
@@ -248,7 +247,7 @@
                 ? context.getColor(R.color.taskbar_nav_icon_light_color)
                 : context.getColor(R.color.taskbar_nav_icon_dark_color);
 
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions = new TaskbarTransitions(mContext, mNavButtonsView);
         }
     }
@@ -365,7 +364,7 @@
                 R.bool.floating_rotation_button_position_left);
         mControllers.rotationButtonController.setRotationButton(mFloatingRotationButton,
                 mRotationButtonListener);
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.init();
         }
 
@@ -373,7 +372,7 @@
         mPropertyHolders.forEach(StatePropertyHolder::endAnimation);
 
         // Initialize things needed to move nav buttons to separate window.
-        mSeparateWindowParent = new BaseDragLayer<TaskbarActivityContext>(mContext, null, 0) {
+        mSeparateWindowParent = new BaseDragLayer<>(mContext, null, 0) {
             @Override
             public void recreateControllers() {
                 mControllers = new TouchController[0];
@@ -629,7 +628,7 @@
     }
 
     public void setWallpaperVisible(boolean isVisible) {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.setWallpaperVisibility(isVisible);
         }
     }
@@ -642,20 +641,20 @@
     }
 
     public void checkNavBarModes() {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             boolean isBarHidden = (mSysuiStateFlags & SYSUI_STATE_NAV_BAR_HIDDEN) != 0;
             mTaskbarTransitions.transitionTo(mTransitionMode, !isBarHidden);
         }
     }
 
     public void finishBarAnimations() {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.finishAnimations();
         }
     }
 
     public void touchAutoDim(boolean reset) {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.setAutoDim(false);
             mHandler.removeCallbacks(mAutoDim);
             if (reset) {
@@ -665,7 +664,7 @@
     }
 
     public void transitionTo(@BarTransitions.TransitionMode int barMode, boolean animate) {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.transitionTo(barMode, animate);
         }
     }
@@ -769,7 +768,7 @@
 
     private void onDarkIntensityChanged() {
         updateNavButtonColor();
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.onDarkIntensityChanged(mTaskbarNavButtonDarkIntensity.value);
         }
     }
@@ -1119,7 +1118,7 @@
                 + mOnBackgroundNavButtonColorOverrideMultiplier.value);
 
         mNavButtonsView.dumpLogs(prefix + "\t", pw);
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.dumpLogs(prefix + "\t", pw);
         }
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 267e19c..430c003 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -331,9 +331,10 @@
         applyState(/* duration = */ 0);
 
         // Hide the background while stashed so it doesn't show on fast swipes home
-        boolean shouldHideTaskbarBackground = enableScalingRevealHomeAnimation()
-                && DisplayController.isTransientTaskbar(mActivity)
-                && isStashed();
+        boolean shouldHideTaskbarBackground = mActivity.isPhoneMode() ||
+                (enableScalingRevealHomeAnimation()
+                        && DisplayController.isTransientTaskbar(mActivity)
+                        && isStashed());
 
         mTaskbarBackgroundAlphaForStash.setValue(shouldHideTaskbarBackground ? 0 : 1);
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index 527e3a3..b21c414 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -46,6 +46,7 @@
 import android.view.animation.Interpolator;
 
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 import androidx.core.view.OneShotPreDrawListener;
 
 import com.android.app.animation.Interpolators;
@@ -96,6 +97,8 @@
     public static final int ALPHA_INDEX_SMALL_SCREEN = 6;
     private static final int NUM_ALPHA_CHANNELS = 7;
 
+    private static boolean sEnableModelLoadingForTests = true;
+
     private final TaskbarActivityContext mActivity;
     private final TaskbarView mTaskbarView;
     private final MultiValueAlpha mTaskbarIconAlpha;
@@ -192,7 +195,7 @@
         mTaskbarIconTranslationXForPinning.updateValue(pinningValue);
 
         mModelCallbacks.init(controllers);
-        if (mActivity.isUserSetupComplete()) {
+        if (mActivity.isUserSetupComplete() && sEnableModelLoadingForTests) {
             // Only load the callbacks if user setup is completed
             LauncherAppState.getInstance(mActivity).getModel().addCallbacksAndLoad(mModelCallbacks);
         }
@@ -924,4 +927,10 @@
 
         mModelCallbacks.dumpLogs(prefix + "\t", pw);
     }
+
+    /** Enables model loading for tests. */
+    @VisibleForTesting
+    public static void enableModelLoadingForTests(boolean enable) {
+        sEnableModelLoadingForTests = enable;
+    }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index 7426dc7..83d4df2 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -33,7 +33,6 @@
 
 import android.annotation.BinderThread;
 import android.annotation.Nullable;
-import android.app.Notification;
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherApps;
@@ -463,15 +462,6 @@
     /** Tells WMShell to show the currently selected bubble. */
     public void showSelectedBubble() {
         if (getSelectedBubbleKey() != null) {
-            if (mSelectedBubble instanceof BubbleBarBubble) {
-                // Because we've visited this bubble, we should suppress the notification.
-                // This is updated on WMShell side when we show the bubble, but that update isn't
-                // passed to launcher, instead we apply it directly here.
-                BubbleInfo info = ((BubbleBarBubble) mSelectedBubble).getInfo();
-                info.setFlags(
-                        info.getFlags() | Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION);
-                mSelectedBubble.getView().updateDotVisibility(true /* animate */);
-            }
             mLastSentBubbleBarTop = mBarView.getRestingTopPositionOnScreen();
             mSystemUiProxy.showBubble(getSelectedBubbleKey(), mLastSentBubbleBarTop);
         } else {
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index fd989b1..4794dfd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -240,6 +240,10 @@
                         BubbleView firstBubble = (BubbleView) getChildAt(0);
                         mUpdateSelectedBubbleAfterCollapse.accept(firstBubble.getBubble().getKey());
                     }
+                    // If the bar was just expanded, remove the dot from the selected bubble.
+                    if (mIsBarExpanded && mSelectedBubbleView != null) {
+                        mSelectedBubbleView.markSeen();
+                    }
                     updateWidth();
                 },
                 /* onUpdate= */ animator -> {
@@ -665,7 +669,7 @@
     }
 
     /** Add a new bubble to the bubble bar. */
-    public void addBubble(View bubble) {
+    public void addBubble(BubbleView bubble) {
         FrameLayout.LayoutParams lp = new FrameLayout.LayoutParams((int) mIconSize, (int) mIconSize,
                 Gravity.LEFT);
         if (isExpanded()) {
@@ -673,6 +677,7 @@
             bubble.setScaleX(0f);
             bubble.setScaleY(0f);
             addView(bubble, 0, lp);
+            bubble.showDotIfNeeded(/* animate= */ false);
 
             mBubbleAnimator = new BubbleAnimator(mIconSize, mExpandedBarIconsSpacing,
                     getChildCount(), mBubbleBarLocation.isOnLeft(isLayoutRtl()));
@@ -825,6 +830,23 @@
         updateBubbleAccessibilityStates();
         updateContentDescription();
         mDismissedByDragBubbleView = null;
+        updateNotificationDotsIfCollapsed();
+    }
+
+    private void updateNotificationDotsIfCollapsed() {
+        if (isExpanded()) {
+            return;
+        }
+        for (int i = 0; i < getChildCount(); i++) {
+            BubbleView bubbleView = (BubbleView) getChildAt(i);
+            // when we're collapsed, the first bubble should show the dot if it has it. the rest of
+            // the bubbles should hide their dots.
+            if (i == 0 && bubbleView.hasUnseenContent()) {
+                bubbleView.showDotIfNeeded(/* animate= */ true);
+            } else {
+                bubbleView.hideDot();
+            }
+        }
     }
 
     private void updateWidth() {
@@ -865,7 +887,6 @@
         float bubbleBarAnimatedTop = viewBottom - getBubbleBarHeight();
         // When translating X & Y the scale is ignored, so need to deduct it from the translations
         final float ty = bubbleBarAnimatedTop + mBubbleBarPadding - getScaleIconShift();
-        final boolean animate = getVisibility() == VISIBLE;
         final boolean onLeft = bubbleBarLocation.isOnLeft(isLayoutRtl());
         // elevation state is opposite to widthState - when expanded all icons are flat
         float elevationState = (1 - widthState);
@@ -897,10 +918,10 @@
             bv.setZ(fullElevationForChild * elevationState);
 
             // only update the dot scale if we're expanding or collapsing
-            // TODO b/351904597: update the dot for the first bubble after removal and reorder
-            // since those might happen when the bar is collapsed and will need their dot back
             if (mWidthAnimator.isRunning()) {
-                bv.setDotScale(widthState);
+                // The dot for the selected bubble scales in the opposite direction of the expansion
+                // animation.
+                bv.showDotIfNeeded(bv == mSelectedBubbleView ? 1 - widthState : widthState);
             }
 
             if (mIsBarExpanded) {
@@ -1025,6 +1046,7 @@
             }
             updateBubblesLayoutProperties(mBubbleBarLocation);
             updateContentDescription();
+            updateNotificationDotsIfCollapsed();
         }
     }
 
@@ -1049,6 +1071,14 @@
         if (mBubbleAnimator == null) {
             updateArrowForSelected(previouslySelectedBubble != null);
         }
+        if (view != null) {
+            if (isExpanded()) {
+                view.markSeen();
+            } else {
+                // when collapsed, the selected bubble should show the dot if it has it
+                view.showDotIfNeeded(/* animate= */ true);
+            }
+        }
     }
 
     /**
@@ -1316,6 +1346,7 @@
     public void dump(PrintWriter pw) {
         pw.println("BubbleBarView state:");
         pw.println("  visibility: " + getVisibility());
+        pw.println("  alpha: " + getAlpha());
         pw.println("  translation Y: " + getTranslationY());
         pw.println("  bubbles in bar (childCount = " + getChildCount() + ")");
         for (BubbleView bubbleView: getBubbles()) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 24b9139..1c33be6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -115,7 +115,7 @@
                 dp -> updateBubbleBarIconSize(dp.taskbarIconSize, /* animate= */ true));
         updateBubbleBarIconSize(mActivity.getDeviceProfile().taskbarIconSize, /* animate= */ false);
         mBubbleBarScale.updateValue(1f);
-        mBubbleClickListener = v -> onBubbleClicked(v);
+        mBubbleClickListener = v -> onBubbleClicked((BubbleView) v);
         mBubbleBarClickListener = v -> onBubbleBarClicked();
         mBubbleDragController.setupBubbleBarView(mBarView);
         mBarView.setOnClickListener(mBubbleBarClickListener);
@@ -139,8 +139,9 @@
         });
     }
 
-    private void onBubbleClicked(View v) {
-        BubbleBarItem bubble = ((BubbleView) v).getBubble();
+    private void onBubbleClicked(BubbleView bubbleView) {
+        bubbleView.markSeen();
+        BubbleBarItem bubble = bubbleView.getBubble();
         if (bubble == null) {
             Log.e(TAG, "bubble click listener, bubble was null");
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
index 4c468bb..acb6b4e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.taskbar.bubbles;
 
 import android.annotation.Nullable;
+import android.app.Notification;
 import android.content.Context;
 import android.graphics.Bitmap;
 import android.graphics.Canvas;
@@ -35,6 +36,7 @@
 import com.android.launcher3.icons.DotRenderer;
 import com.android.launcher3.icons.IconNormalizer;
 import com.android.wm.shell.animation.Interpolators;
+import com.android.wm.shell.common.bubbles.BubbleInfo;
 
 // TODO: (b/276978250) This is will be similar to WMShell's BadgedImageView, it'd be nice to share.
 
@@ -217,9 +219,9 @@
     }
 
     void updateDotVisibility(boolean animate) {
-        final float targetScale = shouldDrawDot() ? 1f : 0f;
+        final float targetScale = hasUnseenContent() ? 1f : 0f;
         if (animate) {
-            animateDotScale();
+            animateDotScale(targetScale);
         } else {
             mDotScale = targetScale;
             mAnimatingToDotScale = targetScale;
@@ -241,18 +243,27 @@
         mAppIcon.setVisibility(show ? VISIBLE : GONE);
     }
 
-    /** Whether the dot indicating unseen content in a bubble should be shown. */
-    private boolean shouldDrawDot() {
-        boolean bubbleHasUnseenContent = mBubble != null
+    boolean hasUnseenContent() {
+        return mBubble != null
                 && mBubble instanceof BubbleBarBubble
                 && !((BubbleBarBubble) mBubble).getInfo().isNotificationSuppressed();
-        // Always render the dot if it's animating, since it could be animating out. Otherwise, show
-        // it if the bubble wants to show it, and we aren't suppressing it.
-        return bubbleHasUnseenContent || mDotIsAnimating;
     }
 
-    /** How big the dot should be, fraction from 0 to 1. */
-    void setDotScale(float fraction) {
+    /**
+     * Used to determine if we can skip drawing frames.
+     *
+     * <p>Generally we should draw the dot when it is requested to be shown and there is unseen
+     * content. But when the dot is removed, we still want to draw frames so that it can be scaled
+     * out.
+     */
+    private boolean shouldDrawDot() {
+        // if there's no dot there's nothing to draw, unless the dot was removed and we're in the
+        // middle of removing it
+        return hasUnseenContent() || mDotIsAnimating;
+    }
+
+    /** Updates the dot scale to the specified fraction from 0 to 1. */
+    private void setDotScale(float fraction) {
         if (!shouldDrawDot()) {
             return;
         }
@@ -260,11 +271,41 @@
         invalidate();
     }
 
-    /**
-     * Animates the dot to the given scale.
-     */
-    private void animateDotScale() {
-        float toScale = shouldDrawDot() ? 1f : 0f;
+    void showDotIfNeeded(float fraction) {
+        if (!hasUnseenContent()) {
+            return;
+        }
+        setDotScale(fraction);
+    }
+
+    void showDotIfNeeded(boolean animate) {
+        // only show the dot if we have unseen content
+        if (!hasUnseenContent()) {
+            return;
+        }
+        if (animate) {
+            animateDotScale(1f);
+        } else {
+            setDotScale(1f);
+        }
+    }
+
+    void hideDot() {
+        animateDotScale(0f);
+    }
+
+    /** Marks this bubble such that it no longer has unseen content, and hides the dot. */
+    void markSeen() {
+        if (mBubble instanceof BubbleBarBubble bubble) {
+            BubbleInfo info = bubble.getInfo();
+            info.setFlags(
+                    info.getFlags() | Notification.BubbleMetadata.FLAG_SUPPRESS_NOTIFICATION);
+            hideDot();
+        }
+    }
+
+    /** Animates the dot to the given scale. */
+    private void animateDotScale(float toScale) {
         boolean isDotScaleChanging = Float.compare(mDotScale, toScale) != 0;
 
         // Don't restart the animation if we're already animating to the given value or if the dot
@@ -277,8 +318,6 @@
 
         final boolean showDot = toScale > 0f;
 
-        // Do NOT wait until after animation ends to setShowDot
-        // to avoid overriding more recent showDot states.
         clearAnimation();
         animate()
                 .setDuration(200)
@@ -293,7 +332,6 @@
                 }).start();
     }
 
-
     @Override
     public String toString() {
         String toString = mBubble != null ? mBubble.getKey() : "null";
diff --git a/quickstep/src/com/android/quickstep/DeviceConfigWrapper.kt b/quickstep/src/com/android/quickstep/DeviceConfigWrapper.kt
index 3549a12..904ed69 100644
--- a/quickstep/src/com/android/quickstep/DeviceConfigWrapper.kt
+++ b/quickstep/src/com/android/quickstep/DeviceConfigWrapper.kt
@@ -64,11 +64,19 @@
             "Enable two stage for LPNH duration and touch slop"
         )
 
-    val twoStageMultiplier =
+    val twoStageDurationPercentage =
         propReader.get(
-            "TWO_STAGE_MULTIPLIER",
-            2,
-            "Extends the duration and touch slop if the initial slop is passed"
+            "TWO_STAGE_DURATION_PERCENTAGE",
+            200,
+            "Extends the duration to trigger a long press after a fraction of the gesture " +
+                "slop is passed, expressed as a percentage (i.e. 200 = 2x)."
+        )
+
+    val twoStageSlopPercentage =
+        propReader.get(
+            "TWO_STAGE_SLOP_PERCENTAGE",
+            50,
+            "Percentage of gesture slop region to trigger the extended long press duration."
         )
 
     val animateLpnh = propReader.get("ANIMATE_LPNH", false, "Animates navbar when long pressing")
diff --git a/quickstep/src/com/android/quickstep/LauncherRestoreEventLoggerImpl.kt b/quickstep/src/com/android/quickstep/LauncherRestoreEventLoggerImpl.kt
index 27bd03d..7d22c52 100644
--- a/quickstep/src/com/android/quickstep/LauncherRestoreEventLoggerImpl.kt
+++ b/quickstep/src/com/android/quickstep/LauncherRestoreEventLoggerImpl.kt
@@ -5,6 +5,7 @@
 import android.app.backup.BackupRestoreEventLogger.BackupRestoreDataType
 import android.app.backup.BackupRestoreEventLogger.BackupRestoreError
 import android.content.Context
+import androidx.annotation.VisibleForTesting
 import com.android.launcher3.Flags.enableLauncherBrMetricsFixed
 import com.android.launcher3.LauncherSettings.Favorites
 import com.android.launcher3.backuprestore.LauncherRestoreEventLogger
@@ -29,8 +30,8 @@
         @BackupRestoreDataType private const val DATA_TYPE_APP_PAIR = "app_pair"
     }
 
-    private val restoreEventLogger: BackupRestoreEventLogger =
-        BackupManager(context).delayedRestoreLogger
+    @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
+    val restoreEventLogger: BackupRestoreEventLogger = BackupManager(context).delayedRestoreLogger
 
     /**
      * For logging when multiple items of a given data type failed to restore.
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index b796951..f2b6005 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -18,6 +18,7 @@
 import static android.os.Process.THREAD_PRIORITY_BACKGROUND;
 
 import static com.android.launcher3.Flags.enableGridOnlyOverview;
+import static com.android.launcher3.Flags.enableRefactorTaskThumbnail;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.quickstep.TaskUtils.checkCurrentOrManagedUserId;
 
@@ -108,7 +109,7 @@
         mIconCache = iconCache;
         mIconCache.registerTaskVisualsChangeListener(this);
         mThumbnailCache = thumbnailCache;
-        if (enableGridOnlyOverview()) {
+        if (isCachePreloadingEnabled()) {
             mCallbacks = new ComponentCallbacks() {
                 @Override
                 public void onConfigurationChanged(Configuration configuration) {
@@ -342,7 +343,7 @@
      * highResLoadingState is enabled
      */
     public void preloadCacheIfNeeded() {
-        if (!enableGridOnlyOverview()) {
+        if (!isCachePreloadingEnabled()) {
             return;
         }
 
@@ -368,7 +369,7 @@
      * Updates cache size and preloads more tasks if cache size increases
      */
     public void updateCacheSizeAndPreloadIfNeeded() {
-        if (!enableGridOnlyOverview()) {
+        if (!isCachePreloadingEnabled()) {
             return;
         }
 
@@ -387,6 +388,10 @@
         mTaskStackChangeListeners.unregisterTaskStackListener(this);
     }
 
+    private boolean isCachePreloadingEnabled() {
+        return enableGridOnlyOverview() || enableRefactorTaskThumbnail();
+    }
+
     /**
      * Listener for receiving running tasks changes
      */
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 08bb6cd..723aa03 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -332,7 +332,8 @@
 
         if (ENABLE_SHELL_TRANSITIONS) {
             final ActivityOptions options = ActivityOptions.makeBasic();
-            options.setPendingIntentBackgroundActivityLaunchAllowedByPermission(true);
+            options.setPendingIntentBackgroundActivityStartMode(
+                    ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOW_ALWAYS);
             // Use regular (non-transient) launch for all apps page to control IME.
             if (!containerInterface.allowAllAppsFromOverview()) {
                 options.setTransientLaunch();
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index 80902e3..f4ff4b2 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -184,6 +184,10 @@
             return mTaskContainer.getTaskView();
         }
 
+        public View getSnapshotView() {
+            return mTaskContainer.getSnapshotView();
+        }
+
         /**
          * Called when the current task is interactive for the user
          */
@@ -312,7 +316,7 @@
             // the difference between the bitmap bounds and the projected view bounds.
             Matrix boundsToBitmapSpace = new Matrix();
             Matrix thumbnailMatrix = enableRefactorTaskThumbnail()
-                    ? mHelper.getEnabledState().getThumbnailMatrix()
+                    ? mHelper.getThumbnailMatrix()
                     : mTaskContainer.getThumbnailViewDeprecated().getThumbnailMatrix();
             thumbnailMatrix.invert(boundsToBitmapSpace);
             RectF boundsInBitmapSpace = new RectF();
diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index 77124bf..dad34ac 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -23,7 +23,6 @@
 import static com.android.launcher3.Flags.enableRefactorTaskThumbnail;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_FREE_FORM_TAP;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_BOTTOM_OR_RIGHT;
-import static com.android.window.flags.Flags.enableDesktopWindowingMode;
 
 import android.app.ActivityOptions;
 import android.graphics.Bitmap;
@@ -63,6 +62,7 @@
 import com.android.systemui.shared.recents.view.AppTransitionAnimationSpecsFuture;
 import com.android.systemui.shared.recents.view.RecentsTransition;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.wm.shell.shared.desktopmode.DesktopModeStatus;
 
 import java.util.Collections;
 import java.util.List;
@@ -394,7 +394,7 @@
             return Settings.Global.getInt(
                     container.asContext().getContentResolver(),
                     Settings.Global.DEVELOPMENT_ENABLE_FREEFORM_WINDOWS_SUPPORT, 0) != 0
-                    && !enableDesktopWindowingMode();
+                    && !DesktopModeStatus.canEnterDesktopMode(container.asContext());
         }
     };
 
diff --git a/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java b/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java
index 848a43a..186c453 100644
--- a/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java
+++ b/quickstep/src/com/android/quickstep/inputconsumers/NavHandleLongPressInputConsumer.java
@@ -73,17 +73,32 @@
         super(delegate, inputMonitor);
         mScreenWidth = DisplayController.INSTANCE.get(context).getInfo().currentSize.x;
         mDeepPressEnabled = DeviceConfigWrapper.get().getEnableLpnhDeepPress();
-        int twoStageMultiplier = DeviceConfigWrapper.get().getTwoStageMultiplier();
         AssistStateManager assistStateManager = AssistStateManager.INSTANCE.get(context);
         if (assistStateManager.getLPNHDurationMillis().isPresent()) {
             mLongPressTimeout = assistStateManager.getLPNHDurationMillis().get().intValue();
         } else {
             mLongPressTimeout = ViewConfiguration.getLongPressTimeout();
         }
-        mOuterLongPressTimeout = mLongPressTimeout * twoStageMultiplier;
-        mTouchSlopSquaredOriginal = deviceState.getSquaredTouchSlop();
-        mTouchSlopSquared = mTouchSlopSquaredOriginal;
-        mOuterTouchSlopSquared = mTouchSlopSquared * (twoStageMultiplier * twoStageMultiplier);
+        float twoStageDurationMultiplier =
+                (DeviceConfigWrapper.get().getTwoStageDurationPercentage() / 100f);
+        mOuterLongPressTimeout = (int) (mLongPressTimeout * twoStageDurationMultiplier);
+
+        float gestureNavTouchSlopSquared = deviceState.getSquaredTouchSlop();
+        float twoStageSlopMultiplier =
+                (DeviceConfigWrapper.get().getTwoStageSlopPercentage() / 100f);
+        float twoStageSlopMultiplierSquared = twoStageSlopMultiplier * twoStageSlopMultiplier;
+        if (DeviceConfigWrapper.get().getEnableLpnhTwoStages()) {
+            // For 2 stages, the outer touch slop should match gesture nav.
+            mTouchSlopSquared = gestureNavTouchSlopSquared * twoStageSlopMultiplierSquared;
+            mOuterTouchSlopSquared = gestureNavTouchSlopSquared;
+        } else {
+            // For single stage, the touch slop should match gesture nav.
+            mTouchSlopSquared = gestureNavTouchSlopSquared;
+            // Note: This outer slop is not actually used for single-stage (flag disabled).
+            mOuterTouchSlopSquared = gestureNavTouchSlopSquared;
+        }
+        mTouchSlopSquaredOriginal = mTouchSlopSquared;
+
         mGestureState = gestureState;
         mGestureState.setIsInExtendedSlopRegion(false);
         if (DEBUG_NAV_HANDLE) {
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
new file mode 100644
index 0000000..f060d7d
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCase.kt
@@ -0,0 +1,54 @@
+/*
+ * 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.recents.usecase
+
+import android.graphics.Matrix
+import android.graphics.Rect
+import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.quickstep.recents.data.RecentsDeviceProfileRepository
+import com.android.quickstep.recents.data.RecentsRotationStateRepository
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
+import com.android.systemui.shared.recents.utilities.PreviewPositionHelper
+import kotlinx.coroutines.flow.firstOrNull
+
+/** Use case for retrieving [Matrix] for positioning Thumbnail in a View */
+class GetThumbnailPositionUseCase(
+    private val deviceProfileRepository: RecentsDeviceProfileRepository,
+    private val rotationStateRepository: RecentsRotationStateRepository,
+    private val tasksRepository: RecentTasksRepository,
+    private val previewPositionHelper: PreviewPositionHelper = PreviewPositionHelper()
+) {
+    suspend fun run(taskId: Int, width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
+        val thumbnailData =
+            tasksRepository.getThumbnailById(taskId).firstOrNull() ?: return MissingThumbnail
+        val thumbnail = thumbnailData.thumbnail ?: return MissingThumbnail
+        previewPositionHelper.updateThumbnailMatrix(
+            Rect(0, 0, thumbnail.width, thumbnail.height),
+            thumbnailData,
+            width,
+            height,
+            deviceProfileRepository.getRecentsDeviceProfile().isLargeScreen,
+            rotationStateRepository.getRecentsRotationState().activityRotation,
+            isRtl
+        )
+        return MatrixScaling(
+            previewPositionHelper.matrix,
+            previewPositionHelper.isOrientationChanged
+        )
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/task/util/GetThumbnailUseCase.kt b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt
similarity index 95%
rename from quickstep/src/com/android/quickstep/task/util/GetThumbnailUseCase.kt
rename to quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt
index e8dd04c3..3aa808e 100644
--- a/quickstep/src/com/android/quickstep/task/util/GetThumbnailUseCase.kt
+++ b/quickstep/src/com/android/quickstep/recents/usecase/GetThumbnailUseCase.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.quickstep.task.util
+package com.android.quickstep.recents.usecase
 
 import android.graphics.Bitmap
 import com.android.quickstep.recents.data.RecentTasksRepository
diff --git a/quickstep/src/com/android/quickstep/recents/usecase/ThumbnailPositionState.kt b/quickstep/src/com/android/quickstep/recents/usecase/ThumbnailPositionState.kt
new file mode 100644
index 0000000..1a1bef7
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/usecase/ThumbnailPositionState.kt
@@ -0,0 +1,26 @@
+/*
+ * 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.recents.usecase
+
+import android.graphics.Matrix
+
+/** State on how a task Thumbnail can fit on given canvas */
+sealed class ThumbnailPositionState {
+    data object MissingThumbnail : ThumbnailPositionState()
+
+    data class MatrixScaling(val matrix: Matrix, val isRotated: Boolean) : ThumbnailPositionState()
+}
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
index feee11f..5fb5b90 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskOverlayUiState.kt
@@ -17,15 +17,10 @@
 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()
+    data class Enabled(val isRealSnapshot: Boolean, val thumbnail: Bitmap?) : TaskOverlayUiState()
 }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt
index 40f9b28..3b3a811 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailUiState.kt
@@ -17,18 +17,17 @@
 package com.android.quickstep.task.thumbnail
 
 import android.graphics.Bitmap
-import android.graphics.Rect
 import androidx.annotation.ColorInt
 
 sealed class TaskThumbnailUiState {
     data object Uninitialized : TaskThumbnailUiState()
+
     data object LiveTile : TaskThumbnailUiState()
+
     data class BackgroundOnly(@ColorInt val backgroundColor: Int) : TaskThumbnailUiState()
-    data class Snapshot(
-        val bitmap: Bitmap,
-        val drawnRect: Rect,
-        @ColorInt val backgroundColor: Int
-    ) : TaskThumbnailUiState()
+
+    data class Snapshot(val bitmap: Bitmap, @ColorInt val backgroundColor: Int) :
+        TaskThumbnailUiState()
 }
 
 data class TaskThumbnail(val taskId: Int, val isRunning: Boolean)
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index d22fc94..d1be3fb 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -31,6 +31,7 @@
 import com.android.launcher3.R
 import com.android.launcher3.Utilities
 import com.android.launcher3.util.ViewPool
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
@@ -61,14 +62,21 @@
             (parent as TaskView).taskViewData,
             (parent as TaskView).getTaskContainerForTaskThumbnailView(this)!!.taskContainerData,
             recentsView.mTasksRepository!!,
+            GetThumbnailPositionUseCase(
+                recentsView.mDeviceProfileRepository!!,
+                recentsView.mOrientedStateRepository!!,
+                recentsView.mTasksRepository!!
+            )
         )
     }
+
     private lateinit var viewAttachedScope: CoroutineScope
 
     private val scrimView: View by lazy { findViewById(R.id.task_thumbnail_scrim) }
     private val liveTileView: LiveTileView by lazy { findViewById(R.id.task_thumbnail_live_tile) }
-    private val thumbnail: ImageView by lazy { findViewById(R.id.task_thumbnail) }
+    private val thumbnailView: ImageView by lazy { findViewById(R.id.task_thumbnail) }
 
+    private var uiState: TaskThumbnailUiState = Uninitialized
     private var inheritedScale: Float = 1f
 
     private val _measuredBounds = Rect()
@@ -97,6 +105,7 @@
             CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("TaskThumbnailView"))
         viewModel.uiState
             .onEach { viewModelUiState ->
+                uiState = viewModelUiState
                 resetViews()
                 when (viewModelUiState) {
                     is Uninitialized -> {}
@@ -135,7 +144,14 @@
     }
 
     override fun onRecycle() {
-        // Do nothing
+        uiState = Uninitialized
+    }
+
+    override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
+        super.onSizeChanged(w, h, oldw, oldh)
+        if (uiState is Snapshot) {
+            setImageMatrix()
+        }
     }
 
     override fun onConfigurationChanged(newConfig: Configuration?) {
@@ -148,7 +164,7 @@
 
     private fun resetViews() {
         liveTileView.isVisible = false
-        thumbnail.isVisible = false
+        thumbnailView.isVisible = false
         scrimView.alpha = 0f
         setBackgroundColor(Color.BLACK)
     }
@@ -163,8 +179,13 @@
 
     private fun drawSnapshot(snapshot: Snapshot) {
         drawBackground(snapshot.backgroundColor)
-        thumbnail.setImageBitmap(snapshot.bitmap)
-        thumbnail.isVisible = true
+        thumbnailView.setImageBitmap(snapshot.bitmap)
+        thumbnailView.isVisible = true
+        setImageMatrix()
+    }
+
+    private fun setImageMatrix() {
+        thumbnailView.imageMatrix = viewModel.getThumbnailPositionState(width, height, isLayoutRtl)
     }
 
     private fun getCurrentCornerRadius() =
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
index d8729a6..e03e114 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModel.kt
@@ -17,9 +17,11 @@
 package com.android.quickstep.task.thumbnail
 
 import android.annotation.ColorInt
-import android.graphics.Rect
+import android.graphics.Matrix
 import androidx.core.graphics.ColorUtils
 import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
+import com.android.quickstep.recents.usecase.ThumbnailPositionState
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
@@ -37,6 +39,7 @@
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.runBlocking
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class TaskThumbnailViewModel(
@@ -44,9 +47,10 @@
     taskViewData: TaskViewData,
     taskContainerData: TaskContainerData,
     private val tasksRepository: RecentTasksRepository,
+    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase
 ) {
     private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
-    private var boundTaskIsRunning = false
+    private lateinit var taskThumbnail: TaskThumbnail
 
     /**
      * Progress for changes in corner radius. progress: 0 = overview corner radius; 1 = fullscreen
@@ -66,16 +70,12 @@
                 taskFlow.map { taskVal ->
                     when {
                         taskVal == null -> Uninitialized
-                        boundTaskIsRunning -> LiveTile
+                        taskThumbnail.isRunning -> LiveTile
                         isBackgroundOnly(taskVal) ->
                             BackgroundOnly(taskVal.colorBackground.removeAlpha())
                         isSnapshotState(taskVal) -> {
                             val bitmap = taskVal.thumbnail?.thumbnail!!
-                            Snapshot(
-                                bitmap,
-                                Rect(0, 0, bitmap.width, bitmap.height),
-                                taskVal.colorBackground.removeAlpha()
-                            )
+                            Snapshot(bitmap, taskVal.colorBackground.removeAlpha())
                         }
                         else -> Uninitialized
                     }
@@ -84,10 +84,22 @@
             .distinctUntilChanged()
 
     fun bind(taskThumbnail: TaskThumbnail) {
-        boundTaskIsRunning = taskThumbnail.isRunning
+        this.taskThumbnail = taskThumbnail
         task.value = tasksRepository.getTaskDataById(taskThumbnail.taskId)
     }
 
+    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): Matrix {
+        return runBlocking {
+            when (
+                val thumbnailPositionState =
+                    getThumbnailPositionUseCase.run(taskThumbnail.taskId, width, height, isRtl)
+            ) {
+                is ThumbnailPositionState.MatrixScaling -> thumbnailPositionState.matrix
+                is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
+            }
+        }
+    }
+
     private fun isBackgroundOnly(task: Task): Boolean = task.isLocked || task.thumbnail == null
 
     private fun isSnapshotState(task: Task): Boolean {
diff --git a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
index a9f7041..cc53be9 100644
--- a/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
+++ b/quickstep/src/com/android/quickstep/task/util/TaskOverlayHelper.kt
@@ -18,6 +18,7 @@
 
 import android.util.Log
 import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
@@ -44,45 +45,71 @@
     // 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 {
+    private val viewModel by lazy {
         val recentsView =
             RecentsViewContainer.containerFromContext<RecentsViewContainer>(
                     overlay.taskView.context
                 )
                 .getOverviewPanel<RecentsView<*, *>>()
-        TaskOverlayViewModel(task, recentsView.mRecentsViewData!!, recentsView.mTasksRepository!!)
+        TaskOverlayViewModel(
+            task,
+            recentsView.mRecentsViewData!!,
+            recentsView.mTasksRepository!!,
+            GetThumbnailPositionUseCase(
+                recentsView.mDeviceProfileRepository!!,
+                recentsView.mOrientedStateRepository!!,
+                recentsView.mTasksRepository
+            )
+        )
     }
 
     // TODO(b/331753115): TaskOverlay should listen for state changes and react.
     val enabledState: Enabled
         get() = uiState as Enabled
 
+    fun getThumbnailMatrix() = getThumbnailPositionState().matrix
+
+    private fun getThumbnailPositionState() =
+        viewModel.getThumbnailPositionState(
+            overlay.snapshotView.width,
+            overlay.snapshotView.height,
+            overlay.snapshotView.isLayoutRtl
+        )
+
     fun init() {
         overlayInitializedScope =
             CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("TaskOverlayHelper"))
-        taskOverlayViewModel.overlayState
+        viewModel.overlayState
             .onEach {
                 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
-                    )
+                    initOverlay(it)
                 } else {
-                    Log.d(TAG, "reset - taskId: ${task.key.id}")
-                    overlay.reset()
+                    reset()
                 }
             }
             .launchIn(overlayInitializedScope)
+        overlay.snapshotView.addOnLayoutChangeListener { _, _, _, _, _, _, _, _, _ ->
+            (uiState as? Enabled)?.let { initOverlay(it) }
+        }
+    }
+
+    private fun initOverlay(enabledState: Enabled) {
+        Log.d(TAG, "initOverlay - taskId: ${task.key.id}, thumbnail: ${enabledState.thumbnail}")
+        with(getThumbnailPositionState()) {
+            overlay.initOverlay(task, enabledState.thumbnail, matrix, isRotated)
+        }
+    }
+
+    private fun reset() {
+        Log.d(TAG, "reset - taskId: ${task.key.id}")
+        overlay.reset()
     }
 
     fun destroy() {
         overlayInitializedScope.cancel()
         uiState = Disabled
-        overlay.reset()
+        reset()
     }
 
     companion object {
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
index 47f32fb..ca3bbb4 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModel.kt
@@ -18,6 +18,9 @@
 
 import android.graphics.Matrix
 import com.android.quickstep.recents.data.RecentTasksRepository
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Disabled
 import com.android.quickstep.task.thumbnail.TaskOverlayUiState.Enabled
@@ -25,12 +28,14 @@
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.runBlocking
 
 /** View model for TaskOverlay */
 class TaskOverlayViewModel(
-    task: Task,
+    private val task: Task,
     recentsViewData: RecentsViewData,
     tasksRepository: RecentTasksRepository,
+    private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase
 ) {
     val overlayState =
         combine(
@@ -42,14 +47,33 @@
                     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()
+
+    fun getThumbnailPositionState(width: Int, height: Int, isRtl: Boolean): ThumbnailPositionState {
+        return runBlocking {
+            val matrix: Matrix
+            val isRotated: Boolean
+            when (
+                val thumbnailPositionState =
+                    getThumbnailPositionUseCase.run(task.key.id, width, height, isRtl)
+            ) {
+                is MatrixScaling -> {
+                    matrix = thumbnailPositionState.matrix
+                    isRotated = thumbnailPositionState.isRotated
+                }
+                is MissingThumbnail -> {
+                    matrix = Matrix.IDENTITY_MATRIX
+                    isRotated = false
+                }
+            }
+            ThumbnailPositionState(matrix, isRotated)
+        }
+    }
+
+    data class ThumbnailPositionState(val matrix: Matrix, val isRotated: Boolean)
 }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index d888eb9..1db1447 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -2820,7 +2820,7 @@
         int[] runningTaskIds = Arrays.stream(runningTasks).mapToInt(task -> task.key.id).toArray();
         TaskView matchingTaskView = null;
         if (hasDesktopTask(runningTasks) && runningTaskIds.length == 1) {
-            // TODO(b/249371338): Unsure if it's expected, desktop runningTasks only have a single
+            // TODO(b/342635213): Unsure if it's expected, desktop runningTasks only have a single
             // taskId, therefore we match any DesktopTaskView that contains the runningTaskId.
             TaskView taskview = getTaskViewByTaskId(runningTaskIds[0]);
             if (taskview instanceof DesktopTaskView) {
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 74d120f..12d1753 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -29,9 +29,9 @@
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory
 import com.android.quickstep.TaskUtils
+import com.android.quickstep.recents.usecase.GetThumbnailUseCase
 import com.android.quickstep.task.thumbnail.TaskThumbnail
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
-import com.android.quickstep.task.util.GetThumbnailUseCase
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.systemui.shared.recents.model.Task
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
new file mode 100644
index 0000000..8bad3b9
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/BubbleViewTest.kt
@@ -0,0 +1,76 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.taskbar.bubbles
+
+import android.content.Context
+import android.graphics.Color
+import android.graphics.Path
+import android.graphics.drawable.ColorDrawable
+import android.view.LayoutInflater
+import androidx.core.graphics.drawable.toBitmap
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.R
+import com.android.wm.shell.common.bubbles.BubbleInfo
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class BubbleViewTest {
+
+    private val context = ApplicationProvider.getApplicationContext<Context>()
+    private lateinit var bubbleView: BubbleView
+    private lateinit var overflowView: BubbleView
+    private lateinit var bubble: BubbleBarBubble
+
+    @Test
+    fun hasUnseenContent_bubble() {
+        setupBubbleViews()
+        assertThat(bubbleView.hasUnseenContent()).isTrue()
+
+        bubbleView.markSeen()
+        assertThat(bubbleView.hasUnseenContent()).isFalse()
+    }
+
+    @Test
+    fun hasUnseenContent_overflow() {
+        setupBubbleViews()
+        assertThat(overflowView.hasUnseenContent()).isFalse()
+    }
+
+    private fun setupBubbleViews() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync {
+            val inflater = LayoutInflater.from(context)
+
+            val bitmap = ColorDrawable(Color.WHITE).toBitmap(width = 20, height = 20)
+            overflowView = inflater.inflate(R.layout.bubblebar_item_view, null, false) as BubbleView
+            overflowView.setOverflow(BubbleBarOverflow(overflowView), bitmap)
+
+            val bubbleInfo =
+                BubbleInfo("key", 0, null, null, 0, context.packageName, null, null, false)
+            bubbleView = inflater.inflate(R.layout.bubblebar_item_view, null, false) as BubbleView
+            bubble =
+                BubbleBarBubble(bubbleInfo, bubbleView, bitmap, bitmap, Color.WHITE, Path(), "")
+            bubbleView.setBubble(bubble)
+        }
+        InstrumentationRegistry.getInstrumentation().waitForIdleSync()
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index a966d2a..bbcf566 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -29,10 +29,10 @@
 import com.android.launcher3.taskbar.TaskbarActivityContext
 import com.android.launcher3.taskbar.TaskbarManager
 import com.android.launcher3.taskbar.TaskbarNavButtonController.TaskbarNavButtonCallbacks
+import com.android.launcher3.taskbar.TaskbarViewController
 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
@@ -152,7 +152,7 @@
                     }
 
                 try {
-                    LauncherAppState.getInstance(context).model.loadModelSync()
+                    TaskbarViewController.enableModelLoadingForTests(false)
 
                     // Replace Launcher Taskbar window with test instance.
                     instrumentation.runOnMainSync {
@@ -167,6 +167,8 @@
                         taskbarManager.destroy()
                         launcherTaskbarManager?.setSuspended(false)
                     }
+
+                    TaskbarViewController.enableModelLoadingForTests(true)
                 }
             }
         }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherRestoreEventLoggerImplTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherRestoreEventLoggerImplTest.kt
new file mode 100644
index 0000000..24f9696
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/LauncherRestoreEventLoggerImplTest.kt
@@ -0,0 +1,139 @@
+package com.android.quickstep
+
+/*
+ * 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.
+ */
+
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.Flags
+import com.android.launcher3.LauncherSettings.Favorites
+import com.android.launcher3.util.LauncherModelHelper
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+@EnableFlags(Flags.FLAG_ENABLE_LAUNCHER_BR_METRICS_FIXED)
+class LauncherRestoreEventLoggerImplTest {
+
+    @get:Rule val setFlagsRule = SetFlagsRule()
+
+    private val mLauncherModelHelper = LauncherModelHelper()
+    private val mSandboxContext: SandboxModelContext = mLauncherModelHelper.sandboxContext
+    private lateinit var loggerUnderTest: LauncherRestoreEventLoggerImpl
+
+    @Before
+    fun setup() {
+        loggerUnderTest = LauncherRestoreEventLoggerImpl(mSandboxContext)
+    }
+
+    @After
+    fun teardown() {
+        loggerUnderTest.restoreEventLogger.clearData()
+        mLauncherModelHelper.destroy()
+    }
+
+    @Test
+    fun `logLauncherItemsRestoreFailed logs multiple items as failing restore`() {
+        // Given
+        val expectedDataType = "application"
+        val expectedError = "test_failure"
+        // When
+        loggerUnderTest.logLauncherItemsRestoreFailed(
+            dataType = expectedDataType,
+            count = 5,
+            error = expectedError
+        )
+        // Then
+        val actualResult = loggerUnderTest.restoreEventLogger.loggingResults.first()
+        assertThat(actualResult.dataType).isEqualTo(expectedDataType)
+        assertThat(actualResult.successCount).isEqualTo(0)
+        assertThat(actualResult.failCount).isEqualTo(5)
+        assertThat(actualResult.errors.keys).containsExactly(expectedError)
+    }
+
+    @Test
+    fun `logLauncherItemsRestored logs multiple items as restored`() {
+        // Given
+        val expectedDataType = "application"
+        // When
+        loggerUnderTest.logLauncherItemsRestored(dataType = expectedDataType, count = 5)
+        // Then
+        val actualResult = loggerUnderTest.restoreEventLogger.loggingResults.first()
+        assertThat(actualResult.dataType).isEqualTo(expectedDataType)
+        assertThat(actualResult.successCount).isEqualTo(5)
+        assertThat(actualResult.failCount).isEqualTo(0)
+        assertThat(actualResult.errors.keys).isEmpty()
+    }
+
+    @Test
+    fun `logSingleFavoritesItemRestored logs a single Favorites Item as restored`() {
+        // Given
+        val expectedDataType = "widget"
+        // When
+        loggerUnderTest.logSingleFavoritesItemRestored(favoritesId = Favorites.ITEM_TYPE_APPWIDGET)
+        // Then
+        val actualResult = loggerUnderTest.restoreEventLogger.loggingResults.first()
+        assertThat(actualResult.dataType).isEqualTo(expectedDataType)
+        assertThat(actualResult.successCount).isEqualTo(1)
+        assertThat(actualResult.failCount).isEqualTo(0)
+        assertThat(actualResult.errors.keys).isEmpty()
+    }
+
+    @Test
+    fun `logSingleFavoritesItemRestoreFailed logs a single Favorites Item as failing restore`() {
+        // Given
+        val expectedDataType = "widget"
+        val expectedError = "test_failure"
+        // When
+        loggerUnderTest.logSingleFavoritesItemRestoreFailed(
+            favoritesId = Favorites.ITEM_TYPE_APPWIDGET,
+            error = expectedError
+        )
+        // Then
+        val actualResult = loggerUnderTest.restoreEventLogger.loggingResults.first()
+        assertThat(actualResult.dataType).isEqualTo(expectedDataType)
+        assertThat(actualResult.successCount).isEqualTo(0)
+        assertThat(actualResult.failCount).isEqualTo(1)
+        assertThat(actualResult.errors.keys).containsExactly(expectedError)
+    }
+
+    @Test
+    fun `logFavoritesItemsRestoreFailed logs multiple Favorites Items as failing restore`() {
+        // Given
+        val expectedDataType = "deep_shortcut"
+        val expectedError = "test_failure"
+        // When
+        loggerUnderTest.logFavoritesItemsRestoreFailed(
+            favoritesId = Favorites.ITEM_TYPE_DEEP_SHORTCUT,
+            count = 5,
+            error = expectedError
+        )
+        // Then
+        val actualResult = loggerUnderTest.restoreEventLogger.loggingResults.first()
+        assertThat(actualResult.dataType).isEqualTo(expectedDataType)
+        assertThat(actualResult.successCount).isEqualTo(0)
+        assertThat(actualResult.failCount).isEqualTo(5)
+        assertThat(actualResult.errors.keys).containsExactly(expectedError)
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCaseTest.kt
new file mode 100644
index 0000000..e657d59
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailPositionUseCaseTest.kt
@@ -0,0 +1,134 @@
+/*
+ * 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.recents.usecase
+
+import android.content.ComponentName
+import android.content.Intent
+import android.graphics.Bitmap
+import android.graphics.Color
+import android.graphics.Matrix
+import android.graphics.Rect
+import android.view.Surface.ROTATION_90
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.recents.data.RecentsDeviceProfileRepository
+import com.android.quickstep.recents.data.RecentsRotationStateRepository
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
+import com.android.systemui.shared.recents.model.Task
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.android.systemui.shared.recents.utilities.PreviewPositionHelper
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+/** Test for [GetThumbnailPositionUseCase] */
+@RunWith(AndroidJUnit4::class)
+class GetThumbnailPositionUseCaseTest {
+    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 deviceProfileRepository = mock<RecentsDeviceProfileRepository>()
+    private val rotationStateRepository = mock<RecentsRotationStateRepository>()
+    private val tasksRepository = FakeTasksRepository()
+    private val previewPositionHelper = mock<PreviewPositionHelper>()
+
+    private val systemUnderTest =
+        GetThumbnailPositionUseCase(
+            deviceProfileRepository,
+            rotationStateRepository,
+            tasksRepository,
+            previewPositionHelper
+        )
+
+    @Test
+    fun invisibleTask_returnsIdentityMatrix() = runTest {
+        tasksRepository.seedTasks(listOf(task))
+
+        assertThat(systemUnderTest.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl = true))
+            .isInstanceOf(MissingThumbnail::class.java)
+    }
+
+    @Test
+    fun visibleTaskWithoutThumbnailData_returnsIdentityMatrix() = runTest {
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+
+        assertThat(systemUnderTest.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl = true))
+            .isInstanceOf(MissingThumbnail::class.java)
+    }
+
+    @Test
+    fun visibleTaskWithThumbnailData_returnsTransformedMatrix() = runTest {
+        tasksRepository.seedThumbnailData(mapOf(TASK_ID to thumbnailData))
+        tasksRepository.seedTasks(listOf(task))
+        tasksRepository.setVisibleTasks(listOf(TASK_ID))
+
+        val isLargeScreen = true
+        val activityRotation = ROTATION_90
+        val isRtl = true
+        val isRotated = true
+
+        whenever(deviceProfileRepository.getRecentsDeviceProfile())
+            .thenReturn(RecentsDeviceProfileRepository.RecentsDeviceProfile(isLargeScreen))
+        whenever(rotationStateRepository.getRecentsRotationState())
+            .thenReturn(RecentsRotationStateRepository.RecentsRotationState(activityRotation))
+
+        whenever(previewPositionHelper.matrix).thenReturn(MATRIX)
+        whenever(previewPositionHelper.isOrientationChanged).thenReturn(isRotated)
+
+        assertThat(systemUnderTest.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .isEqualTo(MatrixScaling(MATRIX, isRotated))
+
+        verify(previewPositionHelper)
+            .updateThumbnailMatrix(
+                Rect(0, 0, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT),
+                thumbnailData,
+                CANVAS_WIDTH,
+                CANVAS_HEIGHT,
+                isLargeScreen,
+                activityRotation,
+                isRtl
+            )
+    }
+
+    companion object {
+        const val TASK_ID = 2
+        const val THUMBNAIL_WIDTH = 100
+        const val THUMBNAIL_HEIGHT = 200
+        const val CANVAS_WIDTH = 300
+        const val CANVAS_HEIGHT = 600
+        val MATRIX =
+            Matrix().apply {
+                setValues(floatArrayOf(2.3f, 4.5f, 2.6f, 7.4f, 3.4f, 2.3f, 2.5f, 6.0f, 3.4f))
+            }
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/util/GetThumbnailUseCaseTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt
similarity index 98%
rename from quickstep/tests/multivalentTests/src/com/android/quickstep/task/util/GetThumbnailUseCaseTest.kt
rename to quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt
index 414f8ca..12a94cf 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/util/GetThumbnailUseCaseTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/usecase/GetThumbnailUseCaseTest.kt
@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package com.android.quickstep.task.util
+package com.android.quickstep.recents.usecase
 
 import android.content.ComponentName
 import android.content.Intent
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
index b78f871..fddc740 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelTest.kt
@@ -20,9 +20,12 @@
 import android.content.Intent
 import android.graphics.Bitmap
 import android.graphics.Color
-import android.graphics.Rect
+import android.graphics.Matrix
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
 import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
@@ -48,8 +51,15 @@
     private val taskViewData by lazy { TaskViewData(taskViewType) }
     private val taskContainerData = TaskContainerData()
     private val tasksRepository = FakeTasksRepository()
+    private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
     private val systemUnderTest by lazy {
-        TaskThumbnailViewModel(recentsViewData, taskViewData, taskContainerData, tasksRepository)
+        TaskThumbnailViewModel(
+            recentsViewData,
+            taskViewData,
+            taskContainerData,
+            tasksRepository,
+            mGetThumbnailPositionUseCase
+        )
     }
 
     private val tasks = (0..5).map(::createTaskWithId)
@@ -149,7 +159,6 @@
                 Snapshot(
                     backgroundColor = Color.rgb(2, 2, 2),
                     bitmap = expectedThumbnailData.thumbnail!!,
-                    drawnRect = Rect(0, 0, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
                 )
             )
     }
@@ -170,11 +179,38 @@
                 Snapshot(
                     backgroundColor = Color.rgb(2, 2, 2),
                     bitmap = expectedThumbnailData.thumbnail!!,
-                    drawnRect = Rect(0, 0, THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
                 )
             )
     }
 
+    @Test
+    fun getSnapshotMatrix_MissingThumbnail() = runTest {
+        val taskId = 2
+        val recentTask = TaskThumbnail(taskId = taskId, isRunning = false)
+        val isRtl = true
+
+        whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .thenReturn(MissingThumbnail)
+
+        systemUnderTest.bind(recentTask)
+        assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .isEqualTo(Matrix.IDENTITY_MATRIX)
+    }
+
+    @Test
+    fun getSnapshotMatrix_MatrixScaling() = runTest {
+        val taskId = 2
+        val recentTask = TaskThumbnail(taskId = taskId, isRunning = false)
+        val isRtl = true
+
+        whenever(mGetThumbnailPositionUseCase.run(taskId, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .thenReturn(MatrixScaling(MATRIX, isRotated = false))
+
+        systemUnderTest.bind(recentTask)
+        assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .isEqualTo(MATRIX)
+    }
+
     private fun createTaskWithId(taskId: Int) =
         Task(Task.TaskKey(taskId, 0, Intent(), ComponentName("", ""), 0, 2000)).apply {
             colorBackground = Color.argb(taskId, taskId, taskId, taskId)
@@ -191,5 +227,11 @@
     companion object {
         const val THUMBNAIL_WIDTH = 100
         const val THUMBNAIL_HEIGHT = 200
+        const val CANVAS_WIDTH = 300
+        const val CANVAS_HEIGHT = 600
+        val MATRIX =
+            Matrix().apply {
+                setValues(floatArrayOf(2.3f, 4.5f, 2.6f, 7.4f, 3.4f, 2.3f, 2.5f, 6.0f, 3.4f))
+            }
     }
 }
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
index 40482c4..64937b4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/viewmodel/TaskOverlayViewModelTest.kt
@@ -23,9 +23,13 @@
 import android.graphics.Matrix
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.quickstep.recents.data.FakeTasksRepository
+import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
+import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
 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.quickstep.task.viewmodel.TaskOverlayViewModel.ThumbnailPositionState
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
@@ -53,7 +57,9 @@
         )
     private val recentsViewData = RecentsViewData()
     private val tasksRepository = FakeTasksRepository()
-    private val systemUnderTest = TaskOverlayViewModel(task, recentsViewData, tasksRepository)
+    private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
+    private val systemUnderTest =
+        TaskOverlayViewModel(task, recentsViewData, tasksRepository, mGetThumbnailPositionUseCase)
 
     @Test
     fun initialStateIsDisabled() = runTest {
@@ -87,7 +93,6 @@
                 Enabled(
                     isRealSnapshot = false,
                     thumbnail = null,
-                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
                 )
             )
     }
@@ -107,7 +112,6 @@
                 Enabled(
                     isRealSnapshot = true,
                     thumbnail = thumbnailData.thumbnail,
-                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
                 )
             )
     }
@@ -127,7 +131,6 @@
                 Enabled(
                     isRealSnapshot = false,
                     thumbnail = thumbnailData.thumbnail,
-                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
                 )
             )
     }
@@ -147,14 +150,42 @@
                 Enabled(
                     isRealSnapshot = false,
                     thumbnail = thumbnailData.thumbnail,
-                    thumbnailMatrix = Matrix.IDENTITY_MATRIX
                 )
             )
     }
 
+    @Test
+    fun getThumbnailMatrix_MissingThumbnail() = runTest {
+        val isRtl = true
+
+        whenever(mGetThumbnailPositionUseCase.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .thenReturn(MissingThumbnail)
+
+        assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .isEqualTo(ThumbnailPositionState(Matrix.IDENTITY_MATRIX, isRotated = false))
+    }
+
+    @Test
+    fun getThumbnailMatrix_MatrixScaling() = runTest {
+        val isRtl = true
+        val isRotated = true
+
+        whenever(mGetThumbnailPositionUseCase.run(TASK_ID, CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .thenReturn(MatrixScaling(MATRIX, isRotated))
+
+        assertThat(systemUnderTest.getThumbnailPositionState(CANVAS_WIDTH, CANVAS_HEIGHT, isRtl))
+            .isEqualTo(ThumbnailPositionState(MATRIX, isRotated))
+    }
+
     companion object {
         const val TASK_ID = 0
         const val THUMBNAIL_WIDTH = 100
         const val THUMBNAIL_HEIGHT = 200
+        const val CANVAS_WIDTH = 300
+        const val CANVAS_HEIGHT = 600
+        val MATRIX =
+            Matrix().apply {
+                setValues(floatArrayOf(2.3f, 4.5f, 2.6f, 7.4f, 3.4f, 2.3f, 2.5f, 6.0f, 3.4f))
+            }
     }
 }
diff --git a/res/drawable/ic_more_vert_dots.xml b/res/drawable/ic_more_vert_dots.xml
new file mode 100644
index 0000000..c4659f8
--- /dev/null
+++ b/res/drawable/ic_more_vert_dots.xml
@@ -0,0 +1,26 @@
+<!--
+  ~ 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.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="24"
+    android:viewportHeight="24"
+    android:tint="?attr/colorControlNormal">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M12,8c1.1,0 2,-0.9 2,-2s-0.9,-2 -2,-2 -2,0.9 -2,2 0.9,2 2,2zM12,10c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2zM12,16c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2z" />
+</vector>
\ No newline at end of file
diff --git a/res/layout/private_space_header.xml b/res/layout/private_space_header.xml
index 9c0f129..52180cf 100644
--- a/res/layout/private_space_header.xml
+++ b/res/layout/private_space_header.xml
@@ -43,6 +43,7 @@
             android:layout_height="@dimen/ps_header_image_height"
             android:background="@drawable/ps_settings_background"
             android:src="@drawable/ic_ps_settings"
+            android:visibility="gone"
             android:contentDescription="@string/ps_container_settings" />
         <LinearLayout
             android:id="@+id/ps_lock_unlock_button"
@@ -71,7 +72,9 @@
                 android:textColor="@color/material_color_on_primary_fixed"
                 android:textSize="14sp"
                 android:text="@string/ps_container_lock_title"
+                android:maxLines="1"
                 android:visibility="gone"
+                android:alpha="0"
                 style="@style/TextHeadline"/>
         </LinearLayout>
     </LinearLayout>
diff --git a/res/layout/widgets_two_pane_sheet_paged_view.xml b/res/layout/widgets_two_pane_sheet_paged_view.xml
index 887efb8..1f41680 100644
--- a/res/layout/widgets_two_pane_sheet_paged_view.xml
+++ b/res/layout/widgets_two_pane_sheet_paged_view.xml
@@ -55,18 +55,39 @@
             android:clipToOutline="true"
             android:orientation="vertical">
 
-            <FrameLayout
+            <LinearLayout
                 android:id="@+id/search_bar_container"
                 android:layout_width="match_parent"
                 android:layout_height="wrap_content"
+                android:orientation="horizontal"
                 android:background="?attr/widgetPickerPrimarySurfaceColor"
-                android:clipToPadding="false"
-                android:elevation="0.1dp"
-                android:paddingBottom="8dp"
+                android:gravity="center_vertical"
                 launcher:layout_sticky="true">
+                <FrameLayout
+                    android:layout_width="0dp"
+                    android:layout_height="wrap_content"
+                    android:layout_weight="1"
+                    android:clipToPadding="false"
+                    android:elevation="0.1dp"
+                    android:paddingBottom="8dp">
 
-                <include layout="@layout/widgets_search_bar" />
-            </FrameLayout>
+                    <include layout="@layout/widgets_search_bar" />
+                </FrameLayout>
+
+                <ImageButton
+                    android:id="@+id/widget_picker_widget_options_menu"
+                    android:layout_width="48dp"
+                    android:layout_height="48dp"
+                    android:layout_marginStart="8dp"
+                    android:layout_marginBottom="8dp"
+                    android:layout_gravity="bottom"
+                    android:background="@drawable/full_rounded_transparent_ripple"
+                    android:contentDescription="@string/widget_picker_widget_options_button_description"
+                    android:padding="12dp"
+                    android:src="@drawable/ic_more_vert_dots"
+                    android:visibility="gone"
+                    android:tint="?attr/widgetPickerWidgetOptionsMenuColor" />
+            </LinearLayout>
 
             <FrameLayout
                 android:layout_width="match_parent"
diff --git a/res/layout/widgets_two_pane_sheet_recyclerview.xml b/res/layout/widgets_two_pane_sheet_recyclerview.xml
index f3d3b16..c6b3b74 100644
--- a/res/layout/widgets_two_pane_sheet_recyclerview.xml
+++ b/res/layout/widgets_two_pane_sheet_recyclerview.xml
@@ -39,19 +39,40 @@
             android:clipToOutline="true"
             android:orientation="vertical">
 
-            <FrameLayout
+            <LinearLayout
                 android:id="@+id/search_bar_container"
                 android:layout_width="match_parent"
                 android:layout_height="wrap_content"
+                android:gravity="center_vertical"
+                android:orientation="horizontal"
                 android:background="?attr/widgetPickerPrimarySurfaceColor"
-                android:clipToPadding="false"
-                android:elevation="0.1dp"
-                android:paddingBottom="16dp"
                 android:paddingHorizontal="@dimen/widget_list_horizontal_margin_two_pane"
                 launcher:layout_sticky="true">
+                <FrameLayout
+                    android:layout_width="0dp"
+                    android:layout_height="wrap_content"
+                    android:layout_weight="1"
+                    android:clipToPadding="false"
+                    android:elevation="0.1dp"
+                    android:paddingBottom="16dp">
 
-                <include layout="@layout/widgets_search_bar" />
-            </FrameLayout>
+                    <include layout="@layout/widgets_search_bar" />
+                </FrameLayout>
+
+                <ImageButton
+                    android:id="@+id/widget_picker_widget_options_menu"
+                    android:layout_width="48dp"
+                    android:layout_height="48dp"
+                    android:layout_marginStart="8dp"
+                    android:layout_marginBottom="16dp"
+                    android:layout_gravity="bottom"
+                    android:background="@drawable/full_rounded_transparent_ripple"
+                    android:contentDescription="@string/widget_picker_widget_options_button_description"
+                    android:padding="12dp"
+                    android:src="@drawable/ic_more_vert_dots"
+                    android:visibility="gone"
+                    android:tint="?attr/widgetPickerWidgetOptionsMenuColor" />
+            </LinearLayout>
 
             <FrameLayout
                 android:layout_width="match_parent"
diff --git a/res/values-night-v34/colors.xml b/res/values-night-v34/colors.xml
index af28119..abce763 100644
--- a/res/values-night-v34/colors.xml
+++ b/res/values-night-v34/colors.xml
@@ -27,4 +27,7 @@
         @android:color/system_on_surface_dark</color>
     <color name="widget_cell_subtitle_color_dark">
         @android:color/system_on_surface_variant_dark</color>
+    <color name="widget_picker_menu_options_color_dark">
+        @android:color/system_on_surface_variant_dark
+    </color>
 </resources>
diff --git a/res/values-night/styles.xml b/res/values-night/styles.xml
index c95722f..9d50d07 100644
--- a/res/values-night/styles.xml
+++ b/res/values-night/styles.xml
@@ -26,9 +26,15 @@
         <item name="android:backgroundDimEnabled">true</item>
     </style>
 
-    <style name="WidgetPickerActivityTheme" parent="@android:style/Theme.Translucent.NoTitleBar">
-        <item name="widgetsTheme">@style/WidgetContainerTheme.Dark</item>
+    <style name="WidgetPickerActivityTheme" parent="@android:style/Theme.DeviceDefault.DayNight">
         <item name="android:windowBackground">@android:color/transparent</item>
+        <item name="android:windowNoTitle">true</item>
+        <item name="android:windowContentOverlay">@null</item>
+        <item name="android:colorBackgroundCacheHint">@null</item>
+        <item name="android:windowIsTranslucent">true</item>
+        <item name="android:windowAnimationStyle">@android:style/Animation</item>
+
+        <item name="widgetsTheme">@style/WidgetContainerTheme.Dark</item>
         <item name="pageIndicatorDotColor">@color/page_indicator_dot_color_dark</item>
     </style>
 </resources>
diff --git a/res/values-v34/colors.xml b/res/values-v34/colors.xml
index 26d3712..d19d4a6 100644
--- a/res/values-v34/colors.xml
+++ b/res/values-v34/colors.xml
@@ -27,4 +27,7 @@
         @android:color/system_on_surface_light</color>
     <color name="widget_cell_subtitle_color_light">
         @android:color/system_on_surface_variant_light</color>
+    <color name="widget_picker_menu_options_color_light">
+        @android:color/system_on_surface_variant_light
+    </color>
 </resources>
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index e4e047e..eda3647 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -62,6 +62,7 @@
     <attr name="preloadIconBackgroundColor" format="color" />
     <attr name="widgetPickerTitleColor" format="color"/>
     <attr name="widgetPickerDescriptionColor" format="color"/>
+    <attr name="widgetPickerWidgetOptionsMenuColor" format="color"/>
     <attr name="widgetPickerPrimarySurfaceColor" format="color"/>
     <attr name="widgetPickerSecondarySurfaceColor" format="color"/>
     <attr name="widgetPickerHeaderAppTitleColor" format="color"/>
diff --git a/res/values/colors.xml b/res/values/colors.xml
index 8fa1992..02f69f6 100644
--- a/res/values/colors.xml
+++ b/res/values/colors.xml
@@ -105,6 +105,7 @@
     <color name="widget_picker_secondary_surface_color_light">#FAF9F8</color>
     <color name="widget_picker_title_color_light">#1F1F1F</color>
     <color name="widget_picker_description_color_light">#4C4D50</color>
+    <color name="widget_picker_menu_options_color_light">@color/material_color_on_surface_variant</color>
     <color name="widget_picker_header_app_title_color_light">#1F1F1F</color>
     <color name="widget_picker_header_app_subtitle_color_light">#444746</color>
     <color name="widget_picker_header_background_color_light">#C2E7FF</color>
@@ -125,6 +126,7 @@
     <color name="widget_picker_secondary_surface_color_dark">#393939</color>
     <color name="widget_picker_title_color_dark">#E3E3E3</color>
     <color name="widget_picker_description_color_dark">#CCCDCF</color>
+    <color name="widget_picker_menu_options_color_dark">@color/material_color_on_surface_variant</color>
     <color name="widget_picker_header_app_title_color_dark">#E3E3E3</color>
     <color name="widget_picker_header_app_subtitle_color_dark">#C4C7C5</color>
     <color name="widget_picker_header_background_color_dark">#004A77</color>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index d33adc4..3b458c2 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -61,6 +61,12 @@
     <string name="long_press_widget_to_add">Touch &amp; hold to move a widget.</string>
     <!-- Accessibility spoken hint message in widget picker, which allows user to add a widget. Custom action is the label for additional accessibility actions available in this mode [CHAR_LIMIT=100] -->
     <string name="long_accessible_way_to_add">Double-tap &amp; hold to move a widget or use custom actions.</string>
+    <!-- Accessibility label for the icon button shown in the widget picker that opens a overflow
+         menu with widgets viewing options. [CHAR_LIMIT=25] -->
+    <string name="widget_picker_widget_options_button_description">More options</string>
+    <!-- Label for the checkbox shown in the widget picker toggles whether to show all widgets or
+     the default set. [CHAR_LIMIT=25] -->
+    <string name="widget_picker_show_all_widgets_menu_item_title">Show all widgets</string>
     <!-- The format string for the dimensions of a widget in the drawer -->
     <!-- There is a special version of this format string for Farsi -->
     <string name="widget_dims_format">%1$d \u00d7 %2$d</string>
diff --git a/res/values/styles.xml b/res/values/styles.xml
index f7273a0..7b43a3b 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -244,6 +244,7 @@
             @color/widget_picker_secondary_surface_color_light</item>
         <item name="widgetPickerTitleColor">@color/widget_picker_title_color_light</item>
         <item name="widgetPickerDescriptionColor">@color/widget_picker_description_color_light</item>
+        <item name="widgetPickerWidgetOptionsMenuColor">@color/widget_picker_menu_options_color_light</item>
         <item name="widgetPickerHeaderAppTitleColor">
             @color/widget_picker_header_app_title_color_light</item>
         <item name="widgetPickerHeaderAppSubtitleColor">
@@ -284,6 +285,7 @@
         <item name="widgetPickerTitleColor">
             @color/widget_picker_title_color_dark</item>
         <item name="widgetPickerDescriptionColor">@color/widget_picker_description_color_dark</item>
+        <item name="widgetPickerWidgetOptionsMenuColor">@color/widget_picker_menu_options_color_dark</item>
         <item name="widgetPickerHeaderAppTitleColor">
             @color/widget_picker_header_app_title_color_dark</item>
         <item name="widgetPickerHeaderAppSubtitleColor">
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index cb897dc..5949732 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -70,6 +70,7 @@
 import static com.android.launcher3.Utilities.postAsyncCallback;
 import static com.android.launcher3.config.FeatureFlags.FOLDABLE_SINGLE_PAGE;
 import static com.android.launcher3.config.FeatureFlags.MULTI_SELECT_EDIT_MODE;
+import static com.android.launcher3.folder.FolderGridOrganizer.createFolderGridOrganizer;
 import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.HIDE;
 import static com.android.launcher3.logging.KeyboardStateManager.KeyboardState.SHOW;
 import static com.android.launcher3.logging.StatsLogManager.EventEnum;
@@ -188,7 +189,6 @@
 import com.android.launcher3.dragndrop.DragView;
 import com.android.launcher3.dragndrop.LauncherDragController;
 import com.android.launcher3.folder.Folder;
-import com.android.launcher3.folder.FolderGridOrganizer;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.keyboard.ViewGroupFocusHelper;
@@ -816,7 +816,7 @@
             View collectionIcon = getWorkspace().getHomescreenIconByItemId(info.container);
             if (collectionIcon instanceof FolderIcon folderIcon
                     && collectionIcon.getTag() instanceof FolderInfo) {
-                if (new FolderGridOrganizer(getDeviceProfile())
+                if (createFolderGridOrganizer(getDeviceProfile())
                         .setFolderInfo((FolderInfo) folderIcon.getTag())
                         .isItemInPreview(info.rank)) {
                     folderIcon.invalidate();
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index a296f46..a448228 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -106,8 +106,6 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.function.Predicate;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 /**
  * Various utilities shared amongst the Launcher's classes.
@@ -116,8 +114,7 @@
 
     private static final String TAG = "Launcher.Utilities";
 
-    private static final Pattern sTrimPattern =
-            Pattern.compile("^[\\s|\\p{javaSpaceChar}]*(.*)[\\s|\\p{javaSpaceChar}]*$");
+    private static final String TRIM_PATTERN = "(^\\h+|\\h+$)";
 
     private static final Matrix sMatrix = new Matrix();
     private static final Matrix sInverseMatrix = new Matrix();
@@ -445,10 +442,7 @@
         if (s == null) {
             return "";
         }
-
-        // Just strip any sequence of whitespace or java space characters from the beginning and end
-        Matcher m = sTrimPattern.matcher(s);
-        return m.replaceAll("$1");
+        return s.toString().replaceAll(TRIM_PATTERN, "").trim();
     }
 
     /**
@@ -722,14 +716,59 @@
     }
 
     /**
-     * Rotates `inOutBounds` by `delta` 90-degree increments. Rotation is visually CCW. Parent
+     * Rotates `inOutBounds` by `delta` 90-degree increments. Rotation is visually CW. Parent
      * sizes represent the "space" that will rotate carrying inOutBounds along with it to determine
      * the final bounds.
+     *
+     * As an example if this is the input:
+     * +-------------+
+     * |   +-----+   |
+     * |   |     |   |
+     * |   +-----+   |
+     * |             |
+     * |             |
+     * |             |
+     * +-------------+
+     * This would be case delta % 4 == 0:
+     * +-------------+
+     * |   +-----+   |
+     * |   |     |   |
+     * |   +-----+   |
+     * |             |
+     * |             |
+     * |             |
+     * +-------------+
+     * This would be case delta % 4 == 1:
+     * +----------------+
+     * |          +--+  |
+     * |          |  |  |
+     * |          |  |  |
+     * |          +--+  |
+     * |                |
+     * +----------------+
+     * This would be case delta % 4 == 2:
+     * +-------------+
+     * |             |
+     * |             |
+     * |             |
+     * |   +-----+   |
+     * |   |     |   |
+     * |   +-----+   |
+     * +-------------+
+     * This would be case delta % 4 == 3:
+     * +----------------+
+     * |  +--+          |
+     * |  |  |          |
+     * |  |  |          |
+     * |  +--+          |
+     * |                |
+     * +----------------+
      */
     public static void rotateBounds(Rect inOutBounds, int parentWidth, int parentHeight,
             int delta) {
         int rdelta = ((delta % 4) + 4) % 4;
         int origLeft = inOutBounds.left;
+        int origTop = inOutBounds.top;
         switch (rdelta) {
             case 0:
                 return;
@@ -741,6 +780,8 @@
                 return;
             case 2:
                 inOutBounds.left = parentWidth - inOutBounds.right;
+                inOutBounds.top = parentHeight - inOutBounds.bottom;
+                inOutBounds.bottom = parentHeight - origTop;
                 inOutBounds.right = parentWidth - origLeft;
                 return;
             case 3:
@@ -830,6 +871,9 @@
             @NonNull Rect inclusionBounds,
             @NonNull Rect exclusionBounds,
             @AdjustmentDirection int adjustmentDirection) {
+        if (!Rect.intersects(targetViewBounds, exclusionBounds)) {
+            return;
+        }
         switch (adjustmentDirection) {
             case TRANSLATE_RIGHT:
                 targetView.setTranslationX(Math.min(
diff --git a/src/com/android/launcher3/accessibility/DragAndDropAccessibilityDelegate.java b/src/com/android/launcher3/accessibility/DragAndDropAccessibilityDelegate.java
index d0fc175..6f73e07 100644
--- a/src/com/android/launcher3/accessibility/DragAndDropAccessibilityDelegate.java
+++ b/src/com/android/launcher3/accessibility/DragAndDropAccessibilityDelegate.java
@@ -29,9 +29,9 @@
 import androidx.customview.widget.ExploreByTouchHelper;
 
 import com.android.launcher3.CellLayout;
-import com.android.launcher3.Launcher;
 import com.android.launcher3.R;
-import com.android.launcher3.dragndrop.DragLayer;
+import com.android.launcher3.views.ActivityContext;
+import com.android.launcher3.views.BaseDragLayer;
 
 import java.util.List;
 
@@ -47,16 +47,17 @@
 
     protected final CellLayout mView;
     protected final Context mContext;
+    protected final ActivityContext mActivityContext;
     protected final LauncherAccessibilityDelegate mDelegate;
-    protected final DragLayer mDragLayer;
+    protected final BaseDragLayer<?> mDragLayer;
 
     public DragAndDropAccessibilityDelegate(CellLayout forView) {
         super(forView);
         mView = forView;
         mContext = mView.getContext();
-        Launcher launcher = Launcher.getLauncher(mContext);
-        mDelegate = launcher.getAccessibilityDelegate();
-        mDragLayer = launcher.getDragLayer();
+        mActivityContext = ActivityContext.lookupContext(mContext);
+        mDelegate = (LauncherAccessibilityDelegate) mActivityContext.getAccessibilityDelegate();
+        mDragLayer = mActivityContext.getDragLayer();
     }
 
     @Override
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 0f4204f..c1264d6 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -41,7 +41,6 @@
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.AnimatorSet;
-import android.animation.LayoutTransition;
 import android.animation.ObjectAnimator;
 import android.animation.ValueAnimator;
 import android.content.Context;
@@ -92,14 +91,14 @@
 public class PrivateProfileManager extends UserProfileManager {
 
     private static final String TAG = "PrivateProfileManager";
-    private static final int EXPAND_COLLAPSE_DURATION = 800;
+    private static final int EXPAND_COLLAPSE_DURATION = 400;
     private static final int SETTINGS_OPACITY_DURATION = 400;
     private static final int TEXT_UNLOCK_OPACITY_DURATION = 300;
     private static final int TEXT_LOCK_OPACITY_DURATION = 50;
     private static final int APP_OPACITY_DURATION = 400;
     private static final int MASK_VIEW_DURATION = 200;
     private static final int APP_OPACITY_DELAY = 400;
-    private static final int SETTINGS_AND_LOCK_GROUP_TRANSITION_DELAY = 400;
+    private static final int PILL_TRANSITION_DELAY = 400;
     private static final int SETTINGS_OPACITY_DELAY = 400;
     private static final int LOCK_TEXT_OPACITY_DELAY = 500;
     private static final int MASK_VIEW_DELAY = 400;
@@ -109,6 +108,8 @@
     private final Predicate<UserHandle> mPrivateProfileMatcher;
     private final int mPsHeaderHeight;
     private final int mFloatingMaskViewCornerRadius;
+    private final int mLockTextMarginStart;
+    private final int mLockTextMarginEnd;
     private final RecyclerView.OnScrollListener mOnIdleScrollListener =
             new RecyclerView.OnScrollListener() {
         @Override
@@ -133,6 +134,11 @@
     private Runnable mOnPSHeaderAdded;
     @Nullable
     private RelativeLayout mPSHeader;
+    @Nullable
+    private TextView mLockText;
+    @Nullable
+    private PrivateSpaceSettingsButton mPrivateSpaceSettingsButton;
+    @Nullable
     private ConstraintLayout mFloatingMaskView;
     private final String mLockedStateContentDesc;
     private final String mUnLockedStateContentDesc;
@@ -155,6 +161,10 @@
                 .getString(R.string.ps_container_unlock_button_content_description);
         mFloatingMaskViewCornerRadius = mAllApps.getContext().getResources().getDimensionPixelSize(
                 R.dimen.ps_floating_mask_corner_radius);
+        mLockTextMarginStart = mAllApps.getContext().getResources().getDimensionPixelSize(
+                R.dimen.ps_lock_icon_text_margin_start_expanded);
+        mLockTextMarginEnd = mAllApps.getContext().getResources().getDimensionPixelSize(
+                R.dimen.ps_lock_icon_text_margin_end_expanded);
     }
 
     /** Adds Private Space Header to the layout. */
@@ -354,20 +364,12 @@
     /** Add Private Space Header view elements based upon {@link UserProfileState} */
     public void bindPrivateSpaceHeaderViewElements(RelativeLayout parent) {
         mPSHeader = parent;
+        Log.d(TAG, "bindPrivateSpaceHeaderViewElements: " + "Binding private space.");
+        updateView();
         if (mOnPSHeaderAdded != null) {
             MAIN_EXECUTOR.execute(mOnPSHeaderAdded);
             mOnPSHeaderAdded = null;
         }
-        // Set the transition duration for the settings and lock button to animate.
-        ViewGroup settingAndLockGroup = mPSHeader.findViewById(R.id.settingsAndLockGroup);
-        if (mReadyToAnimate) {
-            enableLayoutTransition(settingAndLockGroup);
-        } else {
-            // Ensure any unwanted animations to not happen.
-            settingAndLockGroup.setLayoutTransition(null);
-            Log.d(TAG, "bindPrivateSpaceHeaderViewElements: removing transitions ");
-        }
-        updateView();
     }
 
     /** Update the states of the views that make up the header at the state it is called in. */
@@ -375,12 +377,15 @@
         if (mPSHeader == null) {
             return;
         }
+        Log.d(TAG, "bindPrivateSpaceHeaderViewElements: " + "Updating view with state: "
+                + getCurrentState());
         mPSHeader.setAlpha(1);
         ViewGroup lockPill = mPSHeader.findViewById(R.id.ps_lock_unlock_button);
         assert lockPill != null;
-        TextView lockText = lockPill.findViewById(R.id.lock_text);
-        PrivateSpaceSettingsButton settingsButton = mPSHeader.findViewById(R.id.ps_settings_button);
-        assert settingsButton != null;
+        mLockText = lockPill.findViewById(R.id.lock_text);
+        assert mLockText != null;
+        mPrivateSpaceSettingsButton = mPSHeader.findViewById(R.id.ps_settings_button);
+        assert mPrivateSpaceSettingsButton != null;
         //Add image for private space transitioning view
         ImageView transitionView = mPSHeader.findViewById(R.id.ps_transition_image);
         assert transitionView != null;
@@ -391,12 +396,18 @@
                 // Remove header from accessibility target when enabled.
                 mPSHeader.setImportantForAccessibility(View.IMPORTANT_FOR_ACCESSIBILITY_NO);
 
-                lockText.setVisibility(VISIBLE);
+                if (!mReadyToAnimate) {
+                    // Don't set visibilities when animating as the animation will handle it.
+                    mLockText.setVisibility(VISIBLE);
+                    mLockText.setAlpha(1);
+                    mLockText.setHorizontallyScrolling(false);
+                    mPrivateSpaceSettingsButton.setVisibility(
+                            isPrivateSpaceSettingsAvailable() ? VISIBLE : GONE);
+                }
                 lockPill.setVisibility(VISIBLE);
                 lockPill.setOnClickListener(view -> lockingAction(/* lock */ true));
                 lockPill.setContentDescription(mUnLockedStateContentDesc);
 
-                settingsButton.setVisibility(isPrivateSpaceSettingsAvailable() ? VISIBLE : GONE);
                 transitionView.setVisibility(GONE);
             }
             case STATE_DISABLED -> {
@@ -406,12 +417,14 @@
                 mPSHeader.setImportantForAccessibility(View.IMPORTANT_FOR_ACCESSIBILITY_YES);
                 mPSHeader.setContentDescription(mLockedStateContentDesc);
 
-                lockText.setVisibility(GONE);
+                mLockText.setVisibility(GONE);
+                mLockText.setAlpha(0);
+                mLockText.setHorizontallyScrolling(false);
                 lockPill.setVisibility(VISIBLE);
                 lockPill.setOnClickListener(view -> lockingAction(/* lock */ false));
                 lockPill.setContentDescription(mLockedStateContentDesc);
 
-                settingsButton.setVisibility(GONE);
+                mPrivateSpaceSettingsButton.setVisibility(GONE);
                 transitionView.setVisibility(GONE);
             }
             case STATE_TRANSITION -> {
@@ -585,6 +598,51 @@
         return alphaAnim;
     }
 
+    private ValueAnimator animatePillTransition(boolean isExpanding) {
+        if (mLockText == null) {
+            return new ValueAnimator().setDuration(0);
+        }
+        mLockText.measure(0,0);
+        int currentWidth = mLockText.getWidth();
+        int fullWidth = mLockText.getMeasuredWidth();
+        float from = isExpanding ? 0 : currentWidth;
+        float to = isExpanding ? fullWidth : 0;
+        ValueAnimator pillAnim = ObjectAnimator.ofFloat(from, to);
+        pillAnim.setStartDelay(isExpanding ? PILL_TRANSITION_DELAY : 0);
+        pillAnim.setDuration(EXPAND_COLLAPSE_DURATION);
+        pillAnim.setInterpolator(Interpolators.STANDARD);
+        pillAnim.addUpdateListener(new ValueAnimator.AnimatorUpdateListener() {
+            @Override
+            public void onAnimationUpdate(ValueAnimator valueAnimator) {
+                float translation = (float) valueAnimator.getAnimatedValue();
+                float translationFraction = translation / fullWidth;
+                ViewGroup.MarginLayoutParams layoutParams =
+                        (ViewGroup.MarginLayoutParams) mLockText.getLayoutParams();
+                layoutParams.width = (int) translation;
+                layoutParams.setMarginStart((int) (mLockTextMarginStart * translationFraction));
+                layoutParams.setMarginEnd((int) (mLockTextMarginEnd * translationFraction));
+                mLockText.setLayoutParams(layoutParams);
+                mLockText.requestLayout();
+            }
+        });
+        pillAnim.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationEnd(Animator animator) {
+                if (!isExpanding) {
+                    mLockText.setVisibility(GONE);
+                }
+                mLockText.setHorizontallyScrolling(false);
+            }
+
+            @Override
+            public void onAnimationStart(Animator animator) {
+                mLockText.setHorizontallyScrolling(true);
+                mLockText.setVisibility(VISIBLE);
+            }
+        });
+        return pillAnim;
+    }
+
     /**
      * Using PropertySetter{@link PropertySetter}, we can update the view's attributes within an
      * animation. At the moment, collapsing, setting alpha changes, and animating the text is done
@@ -596,22 +654,12 @@
         }
         if (mPSHeader == null) {
             mOnPSHeaderAdded = () -> updatePrivateStateAnimator(expand);
-            setAnimationRunning(false);
+            // Set animation to true, because onBind will be called after this return where we want
+            // the views to be updated accordingly so animation can happen.
+            setAnimationRunning(true);
             return;
         }
         attachFloatingMaskView(expand);
-        ViewGroup settingsAndLockGroup = mPSHeader.findViewById(R.id.settingsAndLockGroup);
-        TextView lockText = mPSHeader.findViewById(R.id.lock_text);
-        PrivateSpaceSettingsButton privateSpaceSettingsButton =
-                mPSHeader.findViewById(R.id.ps_settings_button);
-        if (settingsAndLockGroup.getLayoutTransition() == null) {
-            // Set a new transition if the current ViewGroup does not already contain one as each
-            // transition should only happen once when applied.
-            enableLayoutTransition(settingsAndLockGroup);
-        }
-        settingsAndLockGroup.getLayoutTransition().setStartDelay(
-                LayoutTransition.CHANGING,
-                expand ? SETTINGS_AND_LOCK_GROUP_TRANSITION_DELAY : NO_DELAY);
         PropertySetter headerSetter = new AnimatedPropertySetter();
         headerSetter.add(updateSettingsGearAlpha(expand));
         headerSetter.add(updateLockTextAlpha(expand));
@@ -626,8 +674,6 @@
                                 ? LAUNCHER_PRIVATE_SPACE_UNLOCK_ANIMATION_BEGIN
                                 : LAUNCHER_PRIVATE_SPACE_LOCK_ANIMATION_BEGIN,
                         mAllApps.getActiveRecyclerView());
-                // Animate the collapsing of the text at the same time while updating lock button.
-                lockText.setVisibility(expand ? VISIBLE : GONE);
                 setAnimationRunning(true);
             }
 
@@ -646,10 +692,10 @@
                             : LAUNCHER_PRIVATE_SPACE_LOCK_ANIMATION_END,
                     mAllApps.getActiveRecyclerView());
             Log.d(TAG, "updatePrivateStateAnimator: lockText visibility: "
-                    + lockText.getVisibility() + " lockTextAlpha: " + lockText.getAlpha());
+                    + mLockText.getVisibility() + " lockTextAlpha: " + mLockText.getAlpha());
             Log.d(TAG, "updatePrivateStateAnimator: settingsCog visibility: "
-                    + privateSpaceSettingsButton.getVisibility()
-                    + " settingsCogAlpha: " + privateSpaceSettingsButton.getAlpha());
+                    + mPrivateSpaceSettingsButton.getVisibility()
+                    + " settingsCogAlpha: " + mPrivateSpaceSettingsButton.getAlpha());
             if (!expand) {
                 mAllApps.mAH.get(MAIN).mRecyclerView.removeItemDecoration(
                         mPrivateAppsSectionDecorator);
@@ -663,15 +709,19 @@
         }));
         if (expand) {
             animatorSet.playTogether(animateAlphaOfIcons(true),
+                    animatePillTransition(true),
                     translateFloatingMaskView(false));
         } else {
+            AnimatorSet parallelSet = new AnimatorSet();
+            parallelSet.playTogether(animateAlphaOfIcons(false),
+                    animatePillTransition(false));
             if (isPrivateSpaceHidden()) {
-                animatorSet.playSequentially(animateAlphaOfIcons(false),
+                animatorSet.playSequentially(parallelSet,
                         animateAlphaOfPrivateSpaceContainer(),
                         animateCollapseAnimation());
             } else {
                 animatorSet.playSequentially(translateFloatingMaskView(true),
-                        animateAlphaOfIcons(false),
+                        parallelSet,
                         animateCollapseAnimation());
             }
         }
@@ -702,7 +752,7 @@
     /** Fades out the private space container. */
     private ValueAnimator translateFloatingMaskView(boolean animateIn) {
         if (!Flags.privateSpaceAddFloatingMaskView() || mFloatingMaskView == null) {
-            return new ValueAnimator();
+            return new ValueAnimator().setDuration(0);
         }
         // Translate base on the height amount. Translates out on expand and in on collapse.
         float floatingMaskViewHeight = getFloatingMaskViewHeight();
@@ -714,42 +764,19 @@
         alphaAnim.addUpdateListener(new ValueAnimator.AnimatorUpdateListener() {
             @Override
             public void onAnimationUpdate(ValueAnimator valueAnimator) {
+                if (mFloatingMaskView == null) {
+                    return;
+                }
                 mFloatingMaskView.setTranslationY((float) valueAnimator.getAnimatedValue());
             }
         });
         return alphaAnim;
     }
 
-    /** Animates the layout changes when the text of the button becomes visible/gone. */
-    private void enableLayoutTransition(ViewGroup settingsAndLockGroup) {
-        LayoutTransition settingsAndLockTransition = new LayoutTransition();
-        settingsAndLockTransition.enableTransitionType(LayoutTransition.CHANGING);
-        settingsAndLockTransition.setDuration(EXPAND_COLLAPSE_DURATION);
-        settingsAndLockTransition.setInterpolator(LayoutTransition.CHANGING,
-                Interpolators.STANDARD);
-        settingsAndLockTransition.addTransitionListener(new LayoutTransition.TransitionListener() {
-            @Override
-            public void startTransition(LayoutTransition transition, ViewGroup viewGroup,
-                    View view, int i) {
-                Log.d(TAG, "updatePrivateStateAnimator: transition started: " + transition);
-            }
-            @Override
-            public void endTransition(LayoutTransition transition, ViewGroup viewGroup,
-                    View view, int i) {
-                settingsAndLockGroup.setLayoutTransition(null);
-                mReadyToAnimate = false;
-                Log.d(TAG, "updatePrivateStateAnimator: transition finished: " + transition);
-            }
-        });
-        settingsAndLockGroup.setLayoutTransition(settingsAndLockTransition);
-        Log.d(TAG, "updatePrivateStateAnimator: setting transition: "
-                + settingsAndLockTransition);
-    }
-
     /** Change the settings gear alpha when expanded or collapsed. */
     private ValueAnimator updateSettingsGearAlpha(boolean expand) {
-        if (mPSHeader == null) {
-            return new ValueAnimator();
+        if (mPrivateSpaceSettingsButton == null || !isPrivateSpaceSettingsAvailable()) {
+            return new ValueAnimator().setDuration(0);
         }
         float from = expand ? 0 : 1;
         float to = expand ? 1 : 0;
@@ -760,16 +787,21 @@
         settingsAlphaAnim.addUpdateListener(new ValueAnimator.AnimatorUpdateListener() {
             @Override
             public void onAnimationUpdate(ValueAnimator valueAnimator) {
-                mPSHeader.findViewById(R.id.ps_settings_button)
-                        .setAlpha((float) valueAnimator.getAnimatedValue());
+                mPrivateSpaceSettingsButton.setAlpha((float) valueAnimator.getAnimatedValue());
+            }
+        });
+        settingsAlphaAnim.addListener(new AnimatorListenerAdapter() {
+            @Override
+            public void onAnimationStart(Animator animator) {
+                mPrivateSpaceSettingsButton.setVisibility(VISIBLE);
             }
         });
         return settingsAlphaAnim;
     }
 
     private ValueAnimator updateLockTextAlpha(boolean expand) {
-        if (mPSHeader == null) {
-            return new ValueAnimator();
+        if (mLockText == null) {
+            return new ValueAnimator().setDuration(0);
         }
         float from = expand ? 0 : 1;
         float to = expand ? 1 : 0;
@@ -780,8 +812,7 @@
         alphaAnim.addUpdateListener(new ValueAnimator.AnimatorUpdateListener() {
             @Override
             public void onAnimationUpdate(ValueAnimator valueAnimator) {
-                mPSHeader.findViewById(R.id.lock_text).setAlpha(
-                        (float) valueAnimator.getAnimatedValue());
+                mLockText.setAlpha((float) valueAnimator.getAnimatedValue());
             }
         });
         return alphaAnim;
@@ -819,8 +850,22 @@
         if (!Flags.privateSpaceAddFloatingMaskView()) {
             return;
         }
+        // Use getLocationOnScreen() as simply checking for mPSHeader.getBottom() is only relative
+        // to its parent.
+        int[] psHeaderLocation = new int[2];
+        mPSHeader.getLocationOnScreen(psHeaderLocation);
+        int psHeaderBottomY = psHeaderLocation[1] + mPsHeaderHeight;
+        // Calculate the topY of the floatingMaskView as if it was added.
+        int floatingMaskViewBottomBoxTopY =
+                (int) (mAllApps.getBottom() - getMainRecyclerView().getPaddingBottom());
+        // Don't attach if the header will be clipped by the floating mask view.
+        if (psHeaderBottomY > floatingMaskViewBottomBoxTopY) {
+            mFloatingMaskView = null;
+            return;
+        }
         mFloatingMaskView = (FloatingMaskView) mAllApps.getLayoutInflater().inflate(
                 R.layout.private_space_mask_view, mAllApps, false);
+        assert mFloatingMaskView != null;
         mAllApps.addView(mFloatingMaskView);
         // Translate off the screen first if its collapsing so this header view isn't visible to
         // user when animation starts.
diff --git a/src/com/android/launcher3/folder/Folder.java b/src/com/android/launcher3/folder/Folder.java
index d3c1a02..175ab4e 100644
--- a/src/com/android/launcher3/folder/Folder.java
+++ b/src/com/android/launcher3/folder/Folder.java
@@ -26,6 +26,7 @@
 import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.compat.AccessibilityManagerCompat.sendCustomAccessibilityEvent;
 import static com.android.launcher3.config.FeatureFlags.ALWAYS_USE_HARDWARE_OPTIMIZATION_FOR_FOLDER_ANIMATIONS;
+import static com.android.launcher3.folder.FolderGridOrganizer.createFolderGridOrganizer;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_LABEL_UPDATED;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ITEM_DROP_COMPLETED;
 import static com.android.launcher3.testing.shared.TestProtocol.FOLDER_OPENED_MESSAGE;
@@ -1106,7 +1107,7 @@
     }
 
     private void updateItemLocationsInDatabaseBatch(boolean isBind) {
-        FolderGridOrganizer verifier = new FolderGridOrganizer(
+        FolderGridOrganizer verifier = createFolderGridOrganizer(
                 mActivityContext.getDeviceProfile()).setFolderInfo(mInfo);
 
         ArrayList<ItemInfo> items = new ArrayList<>();
@@ -1404,7 +1405,7 @@
 
     @Override
     public void onAdd(ItemInfo item, int rank) {
-        FolderGridOrganizer verifier = new FolderGridOrganizer(
+        FolderGridOrganizer verifier = createFolderGridOrganizer(
                 mActivityContext.getDeviceProfile()).setFolderInfo(mInfo);
         verifier.updateRankAndPos(item, rank);
         mLauncherDelegate.getModelWriter().addOrMoveItemInDatabase(item, mInfo.id, 0, item.cellX,
diff --git a/src/com/android/launcher3/folder/FolderAnimationManager.java b/src/com/android/launcher3/folder/FolderAnimationManager.java
index 37a8d9b..3c4cf5a 100644
--- a/src/com/android/launcher3/folder/FolderAnimationManager.java
+++ b/src/com/android/launcher3/folder/FolderAnimationManager.java
@@ -21,6 +21,7 @@
 import static com.android.launcher3.BubbleTextView.TEXT_ALPHA_PROPERTY;
 import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
 import static com.android.launcher3.folder.ClippedFolderIconLayoutRule.MAX_NUM_ITEMS_IN_PREVIEW;
+import static com.android.launcher3.folder.FolderGridOrganizer.createFolderGridOrganizer;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -99,7 +100,7 @@
 
         mContext = folder.getContext();
         mDeviceProfile = folder.mActivityContext.getDeviceProfile();
-        mPreviewVerifier = new FolderGridOrganizer(mDeviceProfile);
+        mPreviewVerifier = createFolderGridOrganizer(mDeviceProfile);
 
         mIsOpening = isOpening;
 
diff --git a/src/com/android/launcher3/folder/FolderGridOrganizer.java b/src/com/android/launcher3/folder/FolderGridOrganizer.java
index 593673d..3669d8b 100644
--- a/src/com/android/launcher3/folder/FolderGridOrganizer.java
+++ b/src/com/android/launcher3/folder/FolderGridOrganizer.java
@@ -47,13 +47,20 @@
     /**
      * Note: must call {@link #setFolderInfo(FolderInfo)} manually for verifier to work.
      */
-    public FolderGridOrganizer(DeviceProfile profile) {
-        mMaxCountX = profile.numFolderColumns;
-        mMaxCountY = profile.numFolderRows;
+    public FolderGridOrganizer(int maxCountX, int maxCountY) {
+        mMaxCountX = maxCountX;
+        mMaxCountY = maxCountY;
         mMaxItemsPerPage = mMaxCountX * mMaxCountY;
     }
 
     /**
+     * Creates a FolderGridOrganizer for the given DeviceProfile
+     */
+    public static FolderGridOrganizer createFolderGridOrganizer(DeviceProfile profile) {
+        return new FolderGridOrganizer(profile.numFolderColumns, profile.numFolderRows);
+    }
+
+    /**
      * Updates the organizer with the provided folder info
      */
     public FolderGridOrganizer setFolderInfo(FolderInfo info) {
diff --git a/src/com/android/launcher3/folder/FolderIcon.java b/src/com/android/launcher3/folder/FolderIcon.java
index 00636a3..e9859cf 100644
--- a/src/com/android/launcher3/folder/FolderIcon.java
+++ b/src/com/android/launcher3/folder/FolderIcon.java
@@ -19,6 +19,7 @@
 import static com.android.launcher3.Flags.enableCursorHoverStates;
 import static com.android.launcher3.folder.ClippedFolderIconLayoutRule.ICON_OVERLAP_FACTOR;
 import static com.android.launcher3.folder.ClippedFolderIconLayoutRule.MAX_NUM_ITEMS_IN_PREVIEW;
+import static com.android.launcher3.folder.FolderGridOrganizer.createFolderGridOrganizer;
 import static com.android.launcher3.folder.PreviewItemManager.INITIAL_ITEM_ANIMATION_DURATION;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_AUTO_LABELED;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_AUTO_LABELING_SKIPPED_EMPTY_PRIMARY;
@@ -223,7 +224,7 @@
 
         icon.setAccessibilityDelegate(activity.getAccessibilityDelegate());
 
-        icon.mPreviewVerifier = new FolderGridOrganizer(activity.getDeviceProfile());
+        icon.mPreviewVerifier = createFolderGridOrganizer(activity.getDeviceProfile());
         icon.mPreviewVerifier.setFolderInfo(folderInfo);
         icon.updatePreviewItems(false);
 
diff --git a/src/com/android/launcher3/folder/FolderPagedView.java b/src/com/android/launcher3/folder/FolderPagedView.java
index 8eaa0dc..1b5ef42 100644
--- a/src/com/android/launcher3/folder/FolderPagedView.java
+++ b/src/com/android/launcher3/folder/FolderPagedView.java
@@ -18,6 +18,7 @@
 
 import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
 import static com.android.launcher3.AbstractFloatingView.TYPE_FOLDER;
+import static com.android.launcher3.folder.FolderGridOrganizer.createFolderGridOrganizer;
 
 import android.annotation.SuppressLint;
 import android.content.Context;
@@ -100,10 +101,21 @@
     private boolean mViewsBound = false;
 
     public FolderPagedView(Context context, AttributeSet attrs) {
+        this(
+                context,
+                attrs,
+                createFolderGridOrganizer(ActivityContext.lookupContext(context).getDeviceProfile())
+        );
+    }
+
+    public FolderPagedView(
+            Context context,
+            AttributeSet attrs,
+            FolderGridOrganizer folderGridOrganizer
+    ) {
         super(context, attrs);
         ActivityContext activityContext = ActivityContext.lookupContext(context);
-        DeviceProfile profile = activityContext.getDeviceProfile();
-        mOrganizer = new FolderGridOrganizer(profile);
+        mOrganizer = folderGridOrganizer;
 
         mIsRtl = Utilities.isRtl(getResources());
         setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_YES);
diff --git a/src/com/android/launcher3/graphics/IconPalette.java b/src/com/android/launcher3/graphics/IconPalette.java
index 778b32a..00f1c67 100644
--- a/src/com/android/launcher3/graphics/IconPalette.java
+++ b/src/com/android/launcher3/graphics/IconPalette.java
@@ -16,22 +16,15 @@
 
 package com.android.launcher3.graphics;
 
-import android.app.Notification;
 import android.content.Context;
 import android.graphics.Color;
-import android.util.Log;
 
-import androidx.core.graphics.ColorUtils;
-
-import com.android.launcher3.R;
 import com.android.launcher3.util.Themes;
 
 /**
  * Contains colors based on the dominant color of an icon.
  */
 public class IconPalette {
-
-    private static final boolean DEBUG = false;
     private static final String TAG = "IconPalette";
 
     private static final float MIN_PRELOAD_COLOR_SATURATION = 0.2f;
@@ -54,95 +47,4 @@
         }
         return result;
     }
-
-    /**
-     * Resolves a color such that it has enough contrast to be used as the
-     * color of an icon or text on the given background color.
-     *
-     * @return a color of the same hue with enough contrast against the background.
-     *
-     * This was copied from com.android.internal.util.NotificationColorUtil.
-     */
-    public static int resolveContrastColor(Context context, int color, int background) {
-        final int resolvedColor = resolveColor(context, color);
-
-        int contrastingColor = ensureTextContrast(resolvedColor, background);
-
-        if (contrastingColor != resolvedColor) {
-            if (DEBUG){
-                Log.w(TAG, String.format(
-                        "Enhanced contrast of notification for %s " +
-                                "%s (over background) by changing #%s to %s",
-                        context.getPackageName(),
-                        contrastChange(resolvedColor, contrastingColor, background),
-                        Integer.toHexString(resolvedColor), Integer.toHexString(contrastingColor)));
-            }
-        }
-        return contrastingColor;
-    }
-
-    /**
-     * Resolves {@param color} to an actual color if it is {@link Notification#COLOR_DEFAULT}
-     *
-     * This was copied from com.android.internal.util.NotificationColorUtil.
-     */
-    private static int resolveColor(Context context, int color) {
-        if (color == Notification.COLOR_DEFAULT) {
-            return context.getColor(R.color.notification_icon_default_color);
-        }
-        return color;
-    }
-
-    /** For debugging. This was copied from com.android.internal.util.NotificationColorUtil. */
-    private static String contrastChange(int colorOld, int colorNew, int bg) {
-        return String.format("from %.2f:1 to %.2f:1",
-                ColorUtils.calculateContrast(colorOld, bg),
-                ColorUtils.calculateContrast(colorNew, bg));
-    }
-
-    /**
-     * Finds a text color with sufficient contrast over bg that has the same hue as the original
-     * color.
-     *
-     * This was copied from com.android.internal.util.NotificationColorUtil.
-     */
-    private static int ensureTextContrast(int color, int bg) {
-        return findContrastColor(color, bg, 4.5);
-    }
-    /**
-     * Finds a suitable color such that there's enough contrast.
-     *
-     * @param fg the color to start searching from.
-     * @param bg the color to ensure contrast against.
-     * @param minRatio the minimum contrast ratio required.
-     * @return a color with the same hue as {@param color}, potentially darkened to meet the
-     *          contrast ratio.
-     *
-     * This was copied from com.android.internal.util.NotificationColorUtil.
-     */
-    private static int findContrastColor(int fg, int bg, double minRatio) {
-        if (ColorUtils.calculateContrast(fg, bg) >= minRatio) {
-            return fg;
-        }
-
-        double[] lab = new double[3];
-        ColorUtils.colorToLAB(bg, lab);
-        double bgL = lab[0];
-        ColorUtils.colorToLAB(fg, lab);
-        double fgL = lab[0];
-        boolean isBgDark = bgL < 50;
-
-        double low = isBgDark ? fgL : 0, high = isBgDark ? 100 : fgL;
-        final double a = lab[1], b = lab[2];
-        for (int i = 0; i < 15 && high - low > 0.00001; i++) {
-            final double l = (low + high) / 2;
-            fg = ColorUtils.LABToColor(l, a, b);
-            if (ColorUtils.calculateContrast(fg, bg) > minRatio) {
-                if (isBgDark) high = l; else low = l;
-            } else {
-                if (isBgDark) low = l; else high = l;
-            }
-        }
-        return ColorUtils.LABToColor(low, a, b);
-    }
 }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index 6088941..2408955 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -78,8 +78,6 @@
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
-import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.AppPairInfo;
 import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.FolderInfo;
@@ -106,6 +104,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 
 /**
  * Utility class for generating the preview of Launcher for a given InvariantDeviceProfile.
@@ -376,15 +375,6 @@
                 getApplicationContext(), providerInfo));
     }
 
-    private void inflateAndAddWidgets(LauncherAppWidgetInfo info, WidgetsModel widgetsModel) {
-        WidgetItem widgetItem = widgetsModel.getWidgetProviderInfoByProviderName(
-                info.providerName, info.user, mContext);
-        if (widgetItem == null) {
-            return;
-        }
-        inflateAndAddWidgets(info, widgetItem.widgetInfo);
-    }
-
     private void inflateAndAddWidgets(
             LauncherAppWidgetInfo info, LauncherAppWidgetProviderInfo providerInfo) {
         AppWidgetHostView view = mAppWidgetHost.createView(
@@ -468,17 +458,22 @@
                     break;
             }
         }
+        Map<ComponentKey, AppWidgetProviderInfo> widgetsMap = widgetProviderInfoMap;
         for (ItemInfo itemInfo : currentAppWidgets) {
             switch (itemInfo.itemType) {
                 case Favorites.ITEM_TYPE_APPWIDGET:
                 case Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                    if (widgetProviderInfoMap != null) {
-                        inflateAndAddWidgets(
-                                (LauncherAppWidgetInfo) itemInfo, widgetProviderInfoMap);
-                    } else {
-                        inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo,
-                                dataModel.widgetsModel);
+                    if (widgetsMap == null) {
+                        widgetsMap = dataModel.widgetsModel.getWidgetsByComponentKey()
+                                .entrySet()
+                                .stream()
+                                .filter(entry -> entry.getValue().widgetInfo != null)
+                                .collect(Collectors.toMap(
+                                        Map.Entry::getKey,
+                                        entry -> entry.getValue().widgetInfo
+                                ));
                     }
+                    inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo, widgetsMap);
                     break;
                 default:
                     break;
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 312c6f4..269cb9f 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -568,7 +568,7 @@
     private void processFolderItems() {
         // Sort the folder items, update ranks, and make sure all preview items are high res.
         List<FolderGridOrganizer> verifiers = mApp.getInvariantDeviceProfile().supportedProfiles
-                .stream().map(FolderGridOrganizer::new).toList();
+                .stream().map(FolderGridOrganizer::createFolderGridOrganizer).toList();
         for (CollectionInfo collection : mBgDataModel.collections) {
             if (!(collection instanceof FolderInfo folder)) {
                 continue;
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index 079987b..2febb22 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -109,7 +109,7 @@
         final IconCache iconCache = app.getIconCache();
 
         final String[] packages = mPackages;
-        final int N = packages.length;
+        final int packageCount = packages.length;
         final FlagOp flagOp;
         final HashSet<String> packageSet = new HashSet<>(Arrays.asList(packages));
         final Predicate<ItemInfo> matcher = mOp == OP_USER_AVAILABILITY_CHANGE
@@ -123,7 +123,7 @@
         }
         switch (mOp) {
             case OP_ADD: {
-                for (int i = 0; i < N; i++) {
+                for (int i = 0; i < packageCount; i++) {
                     iconCache.updateIconsForPkg(packages[i], mUser);
                     if (FeatureFlags.PROMISE_APPS_IN_ALL_APPS.get()) {
                         if (DEBUG) {
@@ -146,7 +146,7 @@
                             + " Look for earlier AllAppsList logs to find more information.");
                     removedComponents.add(a.componentName);
                 })) {
-                    for (int i = 0; i < N; i++) {
+                    for (int i = 0; i < packageCount; i++) {
                         iconCache.updateIconsForPkg(packages[i], mUser);
                         activitiesLists.put(packages[i],
                                 appsList.updatePackage(context, packages[i], mUser));
@@ -156,13 +156,13 @@
                 flagOp = FlagOp.NO_OP.removeFlag(WorkspaceItemInfo.FLAG_DISABLED_NOT_AVAILABLE);
                 break;
             case OP_REMOVE: {
-                for (int i = 0; i < N; i++) {
+                for (int i = 0; i < packageCount; i++) {
                     iconCache.removeIconsForPkg(packages[i], mUser);
                 }
                 // Fall through
             }
             case OP_UNAVAILABLE:
-                for (int i = 0; i < N; i++) {
+                for (int i = 0; i < packageCount; i++) {
                     if (DEBUG) {
                         Log.d(TAG, getOpString() + ": removing package=" + packages[i]);
                     }
@@ -217,44 +217,44 @@
             // For system apps, package manager send OP_UPDATE when an app is enabled.
             final boolean isNewApkAvailable = mOp == OP_ADD || mOp == OP_UPDATE;
             synchronized (dataModel) {
-                dataModel.forAllWorkspaceItemInfos(mUser, si -> {
+                dataModel.forAllWorkspaceItemInfos(mUser, itemInfo -> {
 
                     boolean infoUpdated = false;
                     boolean shortcutUpdated = false;
 
-                    ComponentName cn = si.getTargetComponent();
-                    if (cn != null && matcher.test(si)) {
+                    ComponentName cn = itemInfo.getTargetComponent();
+                    if (cn != null && matcher.test(itemInfo)) {
                         String packageName = cn.getPackageName();
 
-                        if (si.hasStatusFlag(WorkspaceItemInfo.FLAG_SUPPORTS_WEB_UI)) {
-                            forceKeepShortcuts.add(si.id);
+                        if (itemInfo.hasStatusFlag(WorkspaceItemInfo.FLAG_SUPPORTS_WEB_UI)) {
+                            forceKeepShortcuts.add(itemInfo.id);
                             if (mOp == OP_REMOVE) {
                                 return;
                             }
                         }
 
-                        if (si.isPromise() && isNewApkAvailable) {
+                        if (itemInfo.isPromise() && isNewApkAvailable) {
                             boolean isTargetValid = !cn.getClassName().equals(
                                     IconCache.EMPTY_CLASS_NAME);
-                            if (si.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
+                            if (itemInfo.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
                                 List<ShortcutInfo> shortcut =
                                         new ShortcutRequest(context, mUser)
                                                 .forPackage(cn.getPackageName(),
-                                                        si.getDeepShortcutId())
+                                                        itemInfo.getDeepShortcutId())
                                                 .query(ShortcutRequest.PINNED);
                                 if (shortcut.isEmpty()) {
                                     isTargetValid = false;
                                     if (DEBUG) {
                                         Log.d(TAG, "Pinned Shortcut not found for updated"
-                                                + " package=" + si.getTargetPackage());
+                                                + " package=" + itemInfo.getTargetPackage());
                                     }
                                 } else {
                                     if (DEBUG) {
                                         Log.d(TAG, "Found pinned shortcut for updated"
-                                                + " package=" + si.getTargetPackage()
+                                                + " package=" + itemInfo.getTargetPackage()
                                                 + ", isTargetValid=" + isTargetValid);
                                     }
-                                    si.updateFromDeepShortcutInfo(shortcut.get(0), context);
+                                    itemInfo.updateFromDeepShortcutInfo(shortcut.get(0), context);
                                     infoUpdated = true;
                                 }
                             } else if (isTargetValid) {
@@ -262,39 +262,39 @@
                                         .isActivityEnabled(cn, mUser);
                             }
 
-                            if (!isTargetValid && (si.hasStatusFlag(
+                            if (!isTargetValid && (itemInfo.hasStatusFlag(
                                     FLAG_RESTORED_ICON | FLAG_AUTOINSTALL_ICON)
-                                    || si.isArchived())) {
-                                if (updateWorkspaceItemIntent(context, si, packageName)) {
+                                    || itemInfo.isArchived())) {
+                                if (updateWorkspaceItemIntent(context, itemInfo, packageName)) {
                                     infoUpdated = true;
-                                } else if (si.hasPromiseIconUi()) {
-                                    removedShortcuts.add(si.id);
+                                } else if (itemInfo.hasPromiseIconUi()) {
+                                    removedShortcuts.add(itemInfo.id);
                                     if (DEBUG) {
                                         FileLog.w(TAG, "Removing restored shortcut promise icon"
                                                 + " that no longer points to valid component."
-                                                + " id=" + si.id
-                                                + ", package=" + si.getTargetPackage()
-                                                + ", status=" + si.status
-                                                + ", isArchived=" + si.isArchived());
+                                                + " id=" + itemInfo.id
+                                                + ", package=" + itemInfo.getTargetPackage()
+                                                + ", status=" + itemInfo.status
+                                                + ", isArchived=" + itemInfo.isArchived());
                                     }
                                     return;
                                 }
                             } else if (!isTargetValid) {
-                                removedShortcuts.add(si.id);
+                                removedShortcuts.add(itemInfo.id);
                                 if (DEBUG) {
                                     FileLog.w(TAG, "Removing shortcut that no longer points to"
                                             + " valid component."
-                                            + " id=" + si.id
-                                            + " package=" + si.getTargetPackage()
-                                            + " status=" + si.status);
+                                            + " id=" + itemInfo.id
+                                            + " package=" + itemInfo.getTargetPackage()
+                                            + " status=" + itemInfo.status);
                                 }
                                 return;
                             } else {
-                                si.status = WorkspaceItemInfo.DEFAULT;
+                                itemInfo.status = WorkspaceItemInfo.DEFAULT;
                                 infoUpdated = true;
                             }
                         } else if (isNewApkAvailable && removedComponents.contains(cn)) {
-                            if (updateWorkspaceItemIntent(context, si, packageName)) {
+                            if (updateWorkspaceItemIntent(context, itemInfo, packageName)) {
                                 infoUpdated = true;
                             }
                         }
@@ -304,7 +304,7 @@
                                     packageName);
                             // TODO: See if we can migrate this to
                             //  AppInfo#updateRuntimeFlagsForActivityTarget
-                            si.setProgressLevel(
+                            itemInfo.setProgressLevel(
                                     activities == null || activities.isEmpty()
                                             ? 100
                                             : PackageManagerHelper.getLoadingProgress(
@@ -313,42 +313,42 @@
                             // In case an app is archived, we need to make sure that archived state
                             // in WorkspaceItemInfo is refreshed.
                             if (Flags.enableSupportForArchiving() && !activities.isEmpty()) {
-                                boolean newArchivalState = activities.get(
-                                        0).getActivityInfo().isArchived;
-                                if (newArchivalState != si.isArchived()) {
-                                    si.runtimeStatusFlags ^= FLAG_ARCHIVED;
+                                boolean newArchivalState = activities.get(0)
+                                        .getActivityInfo().isArchived;
+                                if (newArchivalState != itemInfo.isArchived()) {
+                                    itemInfo.runtimeStatusFlags ^= FLAG_ARCHIVED;
                                     infoUpdated = true;
                                 }
                             }
-                            if (si.itemType == Favorites.ITEM_TYPE_APPLICATION) {
+                            if (itemInfo.itemType == Favorites.ITEM_TYPE_APPLICATION) {
                                 if (activities != null && !activities.isEmpty()) {
-                                    si.setNonResizeable(ApiWrapper.INSTANCE.get(context)
+                                    itemInfo.setNonResizeable(ApiWrapper.INSTANCE.get(context)
                                             .isNonResizeableActivity(activities.get(0)));
                                 }
-                                iconCache.getTitleAndIcon(si, si.usingLowResIcon());
+                                iconCache.getTitleAndIcon(itemInfo, itemInfo.usingLowResIcon());
                                 infoUpdated = true;
                             }
                         }
 
-                        int oldRuntimeFlags = si.runtimeStatusFlags;
-                        si.runtimeStatusFlags = flagOp.apply(si.runtimeStatusFlags);
-                        if (si.runtimeStatusFlags != oldRuntimeFlags) {
+                        int oldRuntimeFlags = itemInfo.runtimeStatusFlags;
+                        itemInfo.runtimeStatusFlags = flagOp.apply(itemInfo.runtimeStatusFlags);
+                        if (itemInfo.runtimeStatusFlags != oldRuntimeFlags) {
                             shortcutUpdated = true;
                         }
                     }
 
                     if (infoUpdated || shortcutUpdated) {
-                        updatedWorkspaceItems.add(si);
+                        updatedWorkspaceItems.add(itemInfo);
                     }
-                    if (infoUpdated && si.id != ItemInfo.NO_ID) {
-                        taskController.getModelWriter().updateItemInDatabase(si);
+                    if (infoUpdated && itemInfo.id != ItemInfo.NO_ID) {
+                        taskController.getModelWriter().updateItemInDatabase(itemInfo);
                     }
                 });
 
                 for (LauncherAppWidgetInfo widgetInfo : dataModel.appWidgets) {
                     if (mUser.equals(widgetInfo.user)
                             && widgetInfo.hasRestoreFlag(
-                                    LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY)
+                            LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY)
                             && packageSet.contains(widgetInfo.providerName.getPackageName())) {
                         widgetInfo.restoreStatus &=
                                 ~LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY
@@ -391,7 +391,7 @@
         } else if (mOp == OP_UPDATE) {
             // Mark disabled packages in the broadcast to be removed
             final LauncherApps launcherApps = context.getSystemService(LauncherApps.class);
-            for (int i=0; i<N; i++) {
+            for (int i = 0; i < packageCount; i++) {
                 if (!launcherApps.isPackageEnabled(packages[i], mUser)) {
                     if (DEBUG) {
                         Log.d(TAG, "OP_UPDATE:"
@@ -423,7 +423,7 @@
         if (mOp == OP_ADD) {
             // Load widgets for the new package. Changes due to app updates are handled through
             // AppWidgetHost events, this is just to initialize the long-press options.
-            for (int i = 0; i < N; i++) {
+            for (int i = 0; i < packageCount; i++) {
                 dataModel.widgetsModel.update(app, new PackageUserKey(packages[i], mUser));
             }
             taskController.bindUpdatedWidgets(dataModel);
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index 454ae96..58ebf0f 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -54,7 +54,9 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 /**
  * Widgets data model that is used by the adapters of the widget views and controllers.
@@ -67,7 +69,26 @@
     private static final boolean DEBUG = false;
 
     /* Map of widgets and shortcuts that are tracked per package. */
-    private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsList = new HashMap<>();
+    private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsByPackageItem = new HashMap<>();
+
+    /**
+     * Returns all widgets keyed by their component key.
+     */
+    public synchronized Map<ComponentKey, WidgetItem> getWidgetsByComponentKey() {
+        return mWidgetsByPackageItem.values().stream()
+                .flatMap(Collection::stream).distinct()
+                .collect(Collectors.toMap(
+                        widget -> new ComponentKey(widget.componentName, widget.user),
+                        Function.identity()
+                ));
+    }
+
+    /**
+     * Returns widgets grouped by the package item that they should belong to.
+     */
+    public synchronized Map<PackageItemInfo, List<WidgetItem>> getWidgetsByPackageItem() {
+        return mWidgetsByPackageItem;
+    }
 
     /**
      * Returns a list of {@link WidgetsListBaseEntry} filtered using given widget item filter. All
@@ -85,7 +106,8 @@
         ArrayList<WidgetsListBaseEntry> result = new ArrayList<>();
         AlphabeticIndexCompat indexer = new AlphabeticIndexCompat(context);
 
-        for (Map.Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
+        for (Map.Entry<PackageItemInfo, List<WidgetItem>> entry :
+                mWidgetsByPackageItem.entrySet()) {
             PackageItemInfo pkgItem = entry.getKey();
             List<WidgetItem> widgetItems = entry.getValue()
                     .stream()
@@ -112,41 +134,6 @@
         return getFilteredWidgetsListForPicker(context, /*widgetItemFilter=*/ item -> true);
     }
 
-    /** Returns a mapping of packages to their widgets without static shortcuts. */
-    public synchronized Map<PackageUserKey, List<WidgetItem>> getAllWidgetsWithoutShortcuts() {
-        if (!WIDGETS_ENABLED) {
-            return Collections.emptyMap();
-        }
-        Map<PackageUserKey, List<WidgetItem>> packagesToWidgets = new HashMap<>();
-        mWidgetsList.forEach((packageItemInfo, widgetsAndShortcuts) -> {
-            List<WidgetItem> widgets = widgetsAndShortcuts.stream()
-                    .filter(item -> item.widgetInfo != null)
-                    .collect(toList());
-            if (widgets.size() > 0) {
-                packagesToWidgets.put(
-                        new PackageUserKey(packageItemInfo.packageName, packageItemInfo.user),
-                        widgets);
-            }
-        });
-        return packagesToWidgets;
-    }
-
-    /**
-     * Returns a map of widget component keys to corresponding widget items. Excludes the
-     * shortcuts.
-     */
-    public synchronized Map<ComponentKey, WidgetItem> getAllWidgetComponentsWithoutShortcuts() {
-        if (!WIDGETS_ENABLED) {
-            return Collections.emptyMap();
-        }
-        Map<ComponentKey, WidgetItem> widgetsMap = new HashMap<>();
-        mWidgetsList.forEach((packageItemInfo, widgetsAndShortcuts) ->
-                widgetsAndShortcuts.stream().filter(item -> item.widgetInfo != null).forEach(
-                        item -> widgetsMap.put(new ComponentKey(item.componentName, item.user),
-                                item)));
-        return widgetsMap;
-    }
-
     /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
@@ -210,14 +197,14 @@
 
         if (packageUser == null) {
             // Clear the list if this is an update on all widgets and shortcuts.
-            mWidgetsList.clear();
+            mWidgetsByPackageItem.clear();
         } else {
             // Otherwise, only clear the widgets and shortcuts for the changed package.
-            mWidgetsList.remove(packageItemInfoCache.getOrCreate(packageUser));
+            mWidgetsByPackageItem.remove(packageItemInfoCache.getOrCreate(packageUser));
         }
 
         // add and update.
-        mWidgetsList.putAll(rawWidgetsShortcuts.stream()
+        mWidgetsByPackageItem.putAll(rawWidgetsShortcuts.stream()
                 .filter(new WidgetValidityCheck(app))
                 .filter(new WidgetFlagCheck())
                 .flatMap(widgetItem -> getPackageUserKeys(app.getContext(), widgetItem).stream()
@@ -237,7 +224,7 @@
             return;
         }
         WidgetManagerHelper widgetManager = new WidgetManagerHelper(app.getContext());
-        for (Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsList.entrySet()) {
+        for (Entry<PackageItemInfo, List<WidgetItem>> entry : mWidgetsByPackageItem.entrySet()) {
             if (packageNames.contains(entry.getKey().packageName)) {
                 List<WidgetItem> items = entry.getValue();
                 int count = items.size();
@@ -258,50 +245,6 @@
         }
     }
 
-    private PackageItemInfo createPackageItemInfo(
-            ComponentName providerName,
-            UserHandle user,
-            int category
-    ) {
-        if (category == NO_CATEGORY) {
-            return new PackageItemInfo(providerName.getPackageName(), user);
-        } else {
-            return new PackageItemInfo("" , category, user);
-        }
-    }
-
-    private IntSet getCategories(ComponentName providerName, Context context) {
-        IntSet categories = WidgetSections.getWidgetsToCategory(context).get(providerName);
-        if (categories != null) {
-            return categories;
-        }
-        categories = new IntSet();
-        categories.add(NO_CATEGORY);
-        return categories;
-    }
-
-    public WidgetItem getWidgetProviderInfoByProviderName(
-            ComponentName providerName, UserHandle user, Context context) {
-        if (!WIDGETS_ENABLED) {
-            return null;
-        }
-        IntSet categories = getCategories(providerName, context);
-
-        // Checking if we have a provider in any of the categories.
-        for (Integer category: categories) {
-            PackageItemInfo key = createPackageItemInfo(providerName, user, category);
-            List<WidgetItem> widgets = mWidgetsList.get(key);
-            if (widgets != null) {
-                return widgets.stream().filter(
-                                item -> item.componentName.equals(providerName)
-                        )
-                        .findFirst()
-                        .orElse(null);
-            }
-        }
-        return null;
-    }
-
     /** Returns {@link PackageItemInfo} of a pending widget. */
     public static PackageItemInfo newPendingItemInfo(Context context, ComponentName provider,
             UserHandle user) {
diff --git a/src/com/android/launcher3/popup/PopupDataProvider.java b/src/com/android/launcher3/popup/PopupDataProvider.java
index fb463f7..7e139c3 100644
--- a/src/com/android/launcher3/popup/PopupDataProvider.java
+++ b/src/com/android/launcher3/popup/PopupDataProvider.java
@@ -64,6 +64,11 @@
 
     /** All installed widgets. */
     private List<WidgetsListBaseEntry> mAllWidgets = List.of();
+    /**
+     * Selectively chosen installed widgets which may be preferred for default display over the list
+     * of all widgets.
+     */
+    private List<WidgetsListBaseEntry> mDefaultWidgets = List.of();
     /** Widgets that can be recommended to the users. */
     private List<ItemInfo> mRecommendedWidgets = List.of();
 
@@ -194,6 +199,18 @@
 
     public void setAllWidgets(List<WidgetsListBaseEntry> allWidgets) {
         mAllWidgets = allWidgets;
+        mDefaultWidgets = List.of();
+        mChangeListener.onWidgetsBound();
+    }
+
+    /**
+     * Sets the list of widgets to be displayed by default and a complete list that can be displayed
+     * when user chooses to show all widgets.
+     */
+    public void setAllWidgets(List<WidgetsListBaseEntry> allWidgets,
+            List<WidgetsListBaseEntry> defaultWidgets) {
+        mAllWidgets = allWidgets;
+        mDefaultWidgets = defaultWidgets;
         mChangeListener.onWidgetsBound();
     }
 
@@ -205,6 +222,14 @@
         return mAllWidgets;
     }
 
+    /**
+     * Returns a "selectively" chosen list of widgets that may be preferred to be shown by default
+     * instead of a complete list.
+     */
+    public List<WidgetsListBaseEntry> getDefaultWidgets() {
+        return mDefaultWidgets;
+    }
+
     /** Returns a list of recommended widgets. */
     public List<WidgetItem> getRecommendedWidgets() {
         HashMap<ComponentKey, WidgetItem> allWidgetItems = new HashMap<>();
@@ -259,8 +284,10 @@
     }
 
     /** Gets the WidgetsListContentEntry for the currently selected header. */
-    public WidgetsListContentEntry getSelectedAppWidgets(PackageUserKey packageUserKey) {
-        return (WidgetsListContentEntry) mAllWidgets.stream()
+    public WidgetsListContentEntry getSelectedAppWidgets(PackageUserKey packageUserKey,
+            boolean useDefault) {
+        List<WidgetsListBaseEntry> widgets = useDefault ? mDefaultWidgets : mAllWidgets;
+        return (WidgetsListContentEntry) widgets.stream()
                 .filter(row -> row instanceof WidgetsListContentEntry
                         && PackageUserKey.fromPackageItemInfo(row.mPkgItem).equals(packageUserKey))
                 .findAny()
diff --git a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
index 398b1df..5ad9222 100644
--- a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
@@ -22,6 +22,8 @@
 import android.os.Parcel;
 import android.os.Parcelable;
 
+import androidx.annotation.VisibleForTesting;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
@@ -52,6 +54,9 @@
         }
     }
 
+    @VisibleForTesting
+    CustomAppWidgetProviderInfo() {}
+
     @Override
     public void initSpans(Context context, InvariantDeviceProfile idp) {
         mIsMinSizeFulfilled = Math.min(spanX, minSpanX) <= idp.numColumns
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 50012b3..faa5d12 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -30,6 +30,7 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.R;
 import com.android.launcher3.util.MainThreadInitializedObject;
@@ -45,6 +46,7 @@
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
@@ -62,9 +64,16 @@
     private final HashMap<ComponentName, CustomWidgetPlugin> mPlugins;
     private final List<CustomAppWidgetProviderInfo> mCustomWidgets;
     private Consumer<PackageUserKey> mWidgetRefreshCallback;
+    private final @NonNull AppWidgetManager mAppWidgetManager;
 
     private CustomWidgetManager(Context context) {
+        this(context, AppWidgetManager.getInstance(context));
+    }
+
+    @VisibleForTesting
+    CustomWidgetManager(Context context, @NonNull AppWidgetManager widgetManager) {
         mContext = context;
+        mAppWidgetManager = widgetManager;
         mPlugins = new HashMap<>();
         mCustomWidgets = new ArrayList<>();
         PluginManagerWrapper.INSTANCE.get(context)
@@ -94,7 +103,7 @@
 
     @Override
     public void onPluginConnected(CustomWidgetPlugin plugin, Context context) {
-        List<AppWidgetProviderInfo> providers = AppWidgetManager.getInstance(context)
+        List<AppWidgetProviderInfo> providers = mAppWidgetManager
                 .getInstalledProvidersForProfile(Process.myUserHandle());
         if (providers.isEmpty()) return;
         Parcel parcel = Parcel.obtain();
@@ -113,6 +122,12 @@
         mCustomWidgets.removeIf(w -> w.getComponent().equals(cn));
     }
 
+    @VisibleForTesting
+    @NonNull
+    Map<ComponentName, CustomWidgetPlugin> getPlugins() {
+        return mPlugins;
+    }
+
     /**
      * Inject a callback function to refresh the widgets.
      */
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 2e36583..21b7be4 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -69,6 +69,7 @@
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.picker.search.SearchModeListener;
 import com.android.launcher3.widget.picker.search.WidgetsSearchBar;
+import com.android.launcher3.widget.picker.search.WidgetsSearchBar.WidgetsSearchDataProvider;
 import com.android.launcher3.workprofile.PersonalWorkPagedView;
 import com.android.launcher3.workprofile.PersonalWorkSlidingTabStrip.OnActivePageChangedListener;
 
@@ -135,7 +136,7 @@
     private WidgetsRecyclerView mCurrentTouchEventRecyclerView;
     @Nullable
     PersonalWorkPagedView mViewPager;
-    private boolean mIsInSearchMode;
+    protected boolean mIsInSearchMode;
     private boolean mIsNoWidgetsViewNeeded;
     @Px
     protected int mMaxSpanPerRow;
@@ -245,8 +246,12 @@
         mSearchBarContainer = mSearchScrollView.findViewById(R.id.search_bar_container);
         mSearchBar = mSearchScrollView.findViewById(R.id.widgets_search_bar);
 
-        mSearchBar.initialize(
-                mActivityContext.getPopupDataProvider(), /* searchModeListener= */ this);
+        mSearchBar.initialize(new WidgetsSearchDataProvider() {
+            @Override
+            public List<WidgetsListBaseEntry> getWidgets() {
+                return getWidgetsToDisplay();
+            }
+        }, /* searchModeListener= */ this);
     }
 
     private void setDeviceManagementResources() {
@@ -462,22 +467,28 @@
         setTranslationShift(mTranslationShift);
     }
 
+    /**
+     * Returns all displayable widgets.
+     */
+    protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
+        return mActivityContext.getPopupDataProvider().getAllWidgets();
+    }
+
     @Override
     public void onWidgetsBound() {
         if (mIsInSearchMode) {
             return;
         }
-        List<WidgetsListBaseEntry> allWidgets =
-                mActivityContext.getPopupDataProvider().getAllWidgets();
+        List<WidgetsListBaseEntry> widgets = getWidgetsToDisplay();
 
         AdapterHolder primaryUserAdapterHolder = mAdapters.get(AdapterHolder.PRIMARY);
-        primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(allWidgets);
+        primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets);
 
         if (mHasWorkProfile) {
             mViewPager.setVisibility(VISIBLE);
             mTabBar.setVisibility(VISIBLE);
             AdapterHolder workUserAdapterHolder = mAdapters.get(AdapterHolder.WORK);
-            workUserAdapterHolder.mWidgetsListAdapter.setWidgets(allWidgets);
+            workUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets);
             onActivePageChanged(mViewPager.getCurrentPage());
         } else {
             onActivePageChanged(0);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index c84680d..c4c755a 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -27,7 +27,9 @@
 import android.graphics.Rect;
 import android.os.Process;
 import android.util.AttributeSet;
+import android.view.Gravity;
 import android.view.LayoutInflater;
+import android.view.MenuItem;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.ViewGroup;
@@ -35,6 +37,7 @@
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.FrameLayout;
 import android.widget.LinearLayout;
+import android.widget.PopupMenu;
 import android.widget.ScrollView;
 import android.widget.TextView;
 
@@ -83,6 +86,17 @@
     private PackageUserKey mSelectedHeader;
     private TextView mHeaderDescription;
 
+    /**
+     * A menu displayed for options (e.g. "show all widgets" filter) around widget lists in the
+     * picker.
+     */
+    protected View mWidgetOptionsMenu;
+    /**
+     * State of the options in the menu (if displayed to the user).
+     */
+    @Nullable
+    protected WidgetOptionsMenuState mWidgetOptionsMenuState = null;
+
     public WidgetsTwoPaneSheet(Context context, AttributeSet attrs, int defStyleAttr) {
         super(context, attrs, defStyleAttr);
     }
@@ -130,6 +144,9 @@
         mHeaderTitle = mContent.findViewById(R.id.title);
         mHeaderDescription = mContent.findViewById(R.id.widget_picker_description);
 
+        mWidgetOptionsMenu = mContent.findViewById(R.id.widget_picker_widget_options_menu);
+        setupWidgetOptionsMenu();
+
         mRightPane = mContent.findViewById(R.id.right_pane);
         mRightPaneScrollView = mContent.findViewById(R.id.right_pane_scroll_view);
         mRightPaneScrollView.setOverScrollMode(View.OVER_SCROLL_NEVER);
@@ -155,6 +172,40 @@
         }
     }
 
+    protected void setupWidgetOptionsMenu() {
+        mWidgetOptionsMenu.setOnClickListener(new OnClickListener() {
+            @Override
+            public void onClick(View v) {
+                if (mWidgetOptionsMenuState != null) {
+                    PopupMenu popupMenu = new PopupMenu(mActivityContext, /*anchor=*/ v,
+                            Gravity.END);
+                    MenuItem menuItem = popupMenu.getMenu().add(
+                            R.string.widget_picker_show_all_widgets_menu_item_title);
+                    menuItem.setCheckable(true);
+                    menuItem.setChecked(mWidgetOptionsMenuState.showAllWidgets);
+                    menuItem.setOnMenuItemClickListener(
+                            item -> onShowAllWidgetsMenuItemClick(item));
+                    popupMenu.show();
+                }
+            }
+        });
+    }
+
+    private boolean onShowAllWidgetsMenuItemClick(MenuItem menuItem) {
+        mWidgetOptionsMenuState.showAllWidgets = !mWidgetOptionsMenuState.showAllWidgets;
+        menuItem.setChecked(mWidgetOptionsMenuState.showAllWidgets);
+
+        // Refresh widgets
+        onWidgetsBound();
+        if (mIsInSearchMode) {
+            mSearchBar.reset();
+        } else if (!mSuggestedWidgetsPackageUserKey.equals(mSelectedHeader)) {
+            mAdapters.get(mActivePage).mWidgetsListAdapter.selectFirstHeaderEntry();
+            mAdapters.get(mActivePage).mWidgetsRecyclerView.scrollToTop();
+        }
+        return true;
+    }
+
     @Override
     protected int getTabletHorizontalMargin(DeviceProfile deviceProfile) {
         if (enableCategorizedWidgetSuggestions()) {
@@ -234,6 +285,29 @@
     }
 
     @Override
+    protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
+        List<WidgetsListBaseEntry> allWidgets =
+                mActivityContext.getPopupDataProvider().getAllWidgets();
+        List<WidgetsListBaseEntry> defaultWidgets =
+                mActivityContext.getPopupDataProvider().getDefaultWidgets();
+
+        if (allWidgets.isEmpty() || defaultWidgets.isEmpty()) {
+            // no menu if there are no default widgets to show
+            mWidgetOptionsMenuState = null;
+            mWidgetOptionsMenu.setVisibility(GONE);
+        } else {
+            if (mWidgetOptionsMenuState == null) {
+                mWidgetOptionsMenuState = new WidgetOptionsMenuState();
+            }
+
+            mWidgetOptionsMenu.setVisibility(VISIBLE);
+            return mWidgetOptionsMenuState.showAllWidgets ? allWidgets : defaultWidgets;
+        }
+
+        return allWidgets;
+    }
+
+    @Override
     public void onWidgetsBound() {
         super.onWidgetsBound();
         if (mRecommendedWidgetsCount == 0 && mSelectedHeader == null) {
@@ -435,8 +509,11 @@
                 final boolean isUserClick = mSelectedHeader != null
                         && !getAccessibilityInitialFocusView().isAccessibilityFocused();
                 mSelectedHeader = selectedHeader;
-                WidgetsListContentEntry contentEntry = mActivityContext.getPopupDataProvider()
-                        .getSelectedAppWidgets(selectedHeader);
+                WidgetsListContentEntry contentEntry =
+                        mActivityContext.getPopupDataProvider().getSelectedAppWidgets(
+                                selectedHeader, /*useDefault=*/
+                                (mWidgetOptionsMenuState != null
+                                        && !mWidgetOptionsMenuState.showAllWidgets));
 
                 if (contentEntry == null || mRightPane == null) {
                     return;
@@ -570,4 +647,15 @@
          */
         void onHeaderChanged(@NonNull PackageUserKey selectedHeader);
     }
+
+    /**
+     * Holds the selection state of the options menu (if presented to the user).
+     */
+    protected static class WidgetOptionsMenuState {
+        /**
+         * UI state indicating whether to show default or all widgets.
+         * <p>If true, shows all widgets; else shows the default widgets.</p>
+         */
+        public boolean showAllWidgets = false;
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java b/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java
index 65937b6..92caf3e 100644
--- a/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java
+++ b/src/com/android/launcher3/widget/picker/search/LauncherWidgetsSearchBar.java
@@ -26,7 +26,6 @@
 
 import com.android.launcher3.ExtendedEditText;
 import com.android.launcher3.R;
-import com.android.launcher3.popup.PopupDataProvider;
 
 /**
  * View for a search bar with an edit text with a cancel button.
@@ -51,7 +50,8 @@
     }
 
     @Override
-    public void initialize(PopupDataProvider dataProvider, SearchModeListener searchModeListener) {
+    public void initialize(WidgetsSearchDataProvider dataProvider,
+            SearchModeListener searchModeListener) {
         mController = new WidgetsSearchBarController(
                 new SimpleWidgetsSearchAlgorithm(dataProvider),
                 mEditText, mCancelButton, searchModeListener);
diff --git a/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithm.java b/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithm.java
index 613066a..0e88787 100644
--- a/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithm.java
+++ b/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithm.java
@@ -21,13 +21,13 @@
 import android.os.Handler;
 
 import com.android.launcher3.model.WidgetItem;
-import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.search.SearchAlgorithm;
 import com.android.launcher3.search.SearchCallback;
 import com.android.launcher3.search.StringMatcherUtility.StringMatcher;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.search.WidgetsSearchBar.WidgetsSearchDataProvider;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -39,9 +39,9 @@
 public final class SimpleWidgetsSearchAlgorithm implements SearchAlgorithm<WidgetsListBaseEntry> {
 
     private final Handler mResultHandler;
-    private final PopupDataProvider mDataProvider;
+    private final WidgetsSearchDataProvider mDataProvider;
 
-    public SimpleWidgetsSearchAlgorithm(PopupDataProvider dataProvider) {
+    public SimpleWidgetsSearchAlgorithm(WidgetsSearchDataProvider dataProvider) {
         mResultHandler = new Handler();
         mDataProvider = dataProvider;
     }
@@ -63,9 +63,9 @@
      * Returns entries for all matched widgets
      */
     public static ArrayList<WidgetsListBaseEntry> getFilteredWidgets(
-            PopupDataProvider dataProvider, String input) {
+            WidgetsSearchDataProvider dataProvider, String input) {
         ArrayList<WidgetsListBaseEntry> results = new ArrayList<>();
-        dataProvider.getAllWidgets().stream()
+        dataProvider.getWidgets().stream()
                 .filter(entry -> entry instanceof WidgetsListHeaderEntry)
                 .forEach(headerEntry -> {
                     List<WidgetItem> matchedWidgetItems = filterWidgetItems(
diff --git a/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
index 44a5e80..ab504e7 100644
--- a/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
+++ b/src/com/android/launcher3/widget/picker/search/WidgetsSearchBar.java
@@ -16,7 +16,9 @@
 
 package com.android.launcher3.widget.picker.search;
 
-import com.android.launcher3.popup.PopupDataProvider;
+import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+
+import java.util.List;
 
 /**
  * Interface for a widgets picker search bar.
@@ -25,7 +27,7 @@
     /**
      * Attaches a controller to the search bar which interacts with {@code searchModeListener}.
      */
-    void initialize(PopupDataProvider dataProvider, SearchModeListener searchModeListener);
+    void initialize(WidgetsSearchDataProvider dataProvider, SearchModeListener searchModeListener);
 
     /**
      * Clears search bar.
@@ -44,4 +46,15 @@
      * Sets the vertical location, in pixels, of this search bar relative to its top position.
      */
     void setTranslationY(float translationY);
+
+
+    /**
+     * Provides corpus from which search results must be returned.
+     */
+    interface WidgetsSearchDataProvider {
+        /**
+         * Returns the widgets from which the search should return the results.
+         */
+        List<WidgetsListBaseEntry> getWidgets();
+    }
 }
diff --git a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java b/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java
index fab3015..dc3b321 100644
--- a/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java
+++ b/tests/multivalentTests/shared/com/android/launcher3/testing/shared/TestProtocol.java
@@ -157,7 +157,6 @@
             "get-overview-current-page-index";
     public static final String REQUEST_GET_SPLIT_SELECTION_ACTIVE = "get-split-selection-active";
     public static final String REQUEST_ENABLE_ROTATION = "enable_rotation";
-    public static final String REQUEST_ENABLE_SUGGESTION = "enable-suggestion";
     public static final String REQUEST_MODEL_QUEUE_CLEARED = "model-queue-cleared";
 
     public static boolean sDebugTracing = false;
diff --git a/tests/multivalentTests/src/com/android/launcher3/UtilitiesTest.kt b/tests/multivalentTests/src/com/android/launcher3/UtilitiesTest.kt
index 60a4197..5a26087 100644
--- a/tests/multivalentTests/src/com/android/launcher3/UtilitiesTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/UtilitiesTest.kt
@@ -17,12 +17,19 @@
 package com.android.launcher3
 
 import android.content.Context
+import android.content.ContextWrapper
+import android.graphics.Rect
+import android.graphics.RectF
 import android.view.View
 import android.view.ViewGroup
 import androidx.test.core.app.ApplicationProvider.getApplicationContext
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import com.android.launcher3.util.ActivityContextWrapper
-import org.junit.Assert.*
+import kotlin.random.Random
+import org.junit.Assert.assertArrayEquals
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertFalse
+import org.junit.Assert.assertTrue
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -30,6 +37,10 @@
 @RunWith(AndroidJUnit4::class)
 class UtilitiesTest {
 
+    companion object {
+        const val SEED = 827
+    }
+
     private lateinit var mContext: Context
 
     @Before
@@ -94,4 +105,283 @@
         assertTrue(Utilities.pointInView(view, -5f, -5f, 10f)) // Inside slop
         assertFalse(Utilities.pointInView(view, 115f, 115f, 10f)) // Outside slop
     }
+
+    @Test
+    fun testNumberBounding() {
+        assertEquals(887.99f, Utilities.boundToRange(887.99f, 0f, 1000f))
+        assertEquals(2.777f, Utilities.boundToRange(887.99f, 0f, 2.777f))
+        assertEquals(900f, Utilities.boundToRange(887.99f, 900f, 1000f))
+
+        assertEquals(9383667L, Utilities.boundToRange(9383667L, -999L, 9999999L))
+        assertEquals(9383668L, Utilities.boundToRange(9383667L, 9383668L, 9999999L))
+        assertEquals(42L, Utilities.boundToRange(9383667L, -999L, 42L))
+
+        assertEquals(345, Utilities.boundToRange(345, 2, 500))
+        assertEquals(400, Utilities.boundToRange(345, 400, 500))
+        assertEquals(300, Utilities.boundToRange(345, 2, 300))
+
+        val random = Random(SEED)
+        for (i in 1..300) {
+            val value = random.nextFloat()
+            val lowerBound = random.nextFloat()
+            val higherBound = lowerBound + random.nextFloat()
+
+            assertEquals(
+                "Utilities.boundToRange doesn't match Kotlin coerceIn",
+                value.coerceIn(lowerBound, higherBound),
+                Utilities.boundToRange(value, lowerBound, higherBound)
+            )
+            assertEquals(
+                "Utilities.boundToRange doesn't match Kotlin coerceIn",
+                value.toInt().coerceIn(lowerBound.toInt(), higherBound.toInt()),
+                Utilities.boundToRange(value.toInt(), lowerBound.toInt(), higherBound.toInt())
+            )
+            assertEquals(
+                "Utilities.boundToRange doesn't match Kotlin coerceIn",
+                value.toLong().coerceIn(lowerBound.toLong(), higherBound.toLong()),
+                Utilities.boundToRange(value.toLong(), lowerBound.toLong(), higherBound.toLong())
+            )
+            assertEquals(
+                "If the lower bound is higher than lower bound, it should return the lower bound",
+                higherBound,
+                Utilities.boundToRange(value, higherBound, lowerBound)
+            )
+        }
+    }
+
+    @Test
+    fun testTranslateOverlappingView() {
+        testConcentricOverlap()
+        leftDownCornerOverlap()
+        noOverlap()
+    }
+
+    /*
+        Test Case: Rectangle Contained Within Another Rectangle
+
+           +-------------+  <-- exclusionBounds
+           |             |
+           |   +-----+   |
+           |   |     |   |  <-- targetViewBounds
+           |   |     |   |
+           |   +-----+   |
+           |             |
+           +-------------+
+    */
+    private fun testConcentricOverlap() {
+        val targetView = View(ContextWrapper(getApplicationContext()))
+        val targetViewBounds = Rect(40, 40, 60, 60)
+        val inclusionBounds = Rect(0, 0, 100, 100)
+        val exclusionBounds = Rect(30, 30, 70, 70)
+
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_RIGHT
+        )
+        assertEquals(30f, targetView.translationX)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_LEFT
+        )
+        assertEquals(-30f, targetView.translationX)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_DOWN
+        )
+        assertEquals(30f, targetView.translationY)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_UP
+        )
+        assertEquals(-30f, targetView.translationY)
+    }
+
+    /*
+    Test Case: Non-Overlapping Rectangles
+
+        +-----------------+      <-- targetViewBounds
+        |                 |
+        |                 |
+        +-----------------+
+
+                 +-----------+     <-- exclusionBounds
+                 |           |
+                 |           |
+                 +-----------+
+    */
+    private fun noOverlap() {
+        val targetView = View(ContextWrapper(getApplicationContext()))
+        val targetViewBounds = Rect(10, 10, 20, 20)
+
+        val inclusionBounds = Rect(0, 0, 100, 100)
+        val exclusionBounds = Rect(30, 30, 40, 40)
+
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_RIGHT
+        )
+        assertEquals(0f, targetView.translationX)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_LEFT
+        )
+        assertEquals(0f, targetView.translationX)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_DOWN
+        )
+        assertEquals(0f, targetView.translationY)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_UP
+        )
+        assertEquals(0f, targetView.translationY)
+    }
+
+    /*
+    Test Case: Rectangles Overlapping at Corners
+
+       +------------+         <-- exclusionBounds
+       |            |
+    +-------+       |
+    |  |    |       |          <-- targetViewBounds
+    |  +------------+
+    |       |
+    +-------+
+    */
+    private fun leftDownCornerOverlap() {
+        val targetView = View(ContextWrapper(getApplicationContext()))
+        val targetViewBounds = Rect(20, 20, 30, 30)
+
+        val inclusionBounds = Rect(0, 0, 100, 100)
+        val exclusionBounds = Rect(25, 25, 35, 35)
+
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_RIGHT
+        )
+        assertEquals(15f, targetView.translationX)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_LEFT
+        )
+        assertEquals(-5f, targetView.translationX)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_DOWN
+        )
+        assertEquals(15f, targetView.translationY)
+        Utilities.translateOverlappingView(
+            targetView,
+            targetViewBounds,
+            inclusionBounds,
+            exclusionBounds,
+            Utilities.TRANSLATE_UP
+        )
+        assertEquals(-5f, targetView.translationY)
+    }
+
+    @Test
+    fun trim() {
+        val expectedString = "Hello World"
+        assertEquals(expectedString, Utilities.trim("Hello World   "))
+        // Basic trimming
+        assertEquals(expectedString, Utilities.trim("   Hello World  "))
+        assertEquals(expectedString, Utilities.trim("   Hello World"))
+
+        // Non-breaking whitespace
+        assertEquals("Hello World", Utilities.trim("\u00A0\u00A0Hello World\u00A0\u00A0"))
+
+        // Whitespace combinations
+        assertEquals(expectedString, Utilities.trim("\t \r\n Hello World \n\r"))
+        assertEquals(expectedString, Utilities.trim("\nHello World   "))
+
+        // Null input
+        assertEquals("", Utilities.trim(null))
+
+        // Empty String
+        assertEquals("", Utilities.trim(""))
+    }
+
+    @Test
+    fun getProgress() {
+        // Basic test
+        assertEquals(0.5f, Utilities.getProgress(50f, 0f, 100f), 0.001f)
+
+        // Negative values
+        assertEquals(0.5f, Utilities.getProgress(-20f, -50f, 10f), 0.001f)
+
+        // Outside of range
+        assertEquals(1.2f, Utilities.getProgress(120f, 0f, 100f), 0.001f)
+    }
+
+    @Test
+    fun scaleRectFAboutPivot() {
+        // Enlarge
+        var rectF = RectF(10f, 20f, 50f, 80f)
+        Utilities.scaleRectFAboutPivot(rectF, 30f, 50f, 1.5f)
+        assertEquals(RectF(0f, 5f, 60f, 95f), rectF)
+
+        // Shrink
+        rectF = RectF(10f, 20f, 50f, 80f)
+        Utilities.scaleRectFAboutPivot(rectF, 30f, 50f, 0.5f)
+        assertEquals(RectF(20f, 35f, 40f, 65f), rectF)
+
+        // No scale
+        rectF = RectF(10f, 20f, 50f, 80f)
+        Utilities.scaleRectFAboutPivot(rectF, 30f, 50f, 1.0f)
+        assertEquals(RectF(10f, 20f, 50f, 80f), rectF)
+    }
+
+    @Test
+    fun rotateBounds() {
+        var rect = Rect(20, 70, 60, 80)
+        Utilities.rotateBounds(rect, 100, 100, 0)
+        assertEquals(Rect(20, 70, 60, 80), rect)
+
+        rect = Rect(20, 70, 60, 80)
+        Utilities.rotateBounds(rect, 100, 100, 1)
+        assertEquals(Rect(70, 40, 80, 80), rect)
+
+        rect = Rect(20, 70, 60, 80)
+        Utilities.rotateBounds(rect, 100, 100, 2)
+        assertEquals(Rect(40, 20, 80, 30), rect)
+
+        rect = Rect(20, 70, 60, 80)
+        Utilities.rotateBounds(rect, 100, 100, 3)
+        assertEquals(Rect(20, 20, 30, 60), rect)
+    }
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/accessibility/FolderAccessibilityHelperTest.kt b/tests/multivalentTests/src/com/android/launcher3/accessibility/FolderAccessibilityHelperTest.kt
new file mode 100644
index 0000000..1cbe1df
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/accessibility/FolderAccessibilityHelperTest.kt
@@ -0,0 +1,114 @@
+/*
+ * 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.accessibility // Use the original package
+
+// Imports
+import android.content.Context
+import androidx.test.core.app.ApplicationProvider.getApplicationContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.CellLayout
+import com.android.launcher3.folder.FolderPagedView
+import com.android.launcher3.util.ActivityContextWrapper
+import kotlin.math.min
+import org.junit.Assert.assertEquals
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.`when`
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class FolderAccessibilityHelperTest {
+
+    // Context
+    private lateinit var mContext: Context
+    // Mocks
+    @Mock private lateinit var mockParent: FolderPagedView
+    @Mock private lateinit var mockLayout: CellLayout
+
+    private var countX = 4
+    private var countY = 3
+    private var index = 1
+
+    // System under test
+    private lateinit var folderAccessibilityHelper: FolderAccessibilityHelper
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        mContext = ActivityContextWrapper(getApplicationContext())
+        `when`(mockLayout.parent).thenReturn(mockParent)
+        `when`(mockLayout.context).thenReturn(mContext)
+
+        // mStartPosition isn't recalculated after the constructor
+        // If you want to create new tests with different starting params,
+        // rebuild the folderAccessibilityHelper object
+        val countX = 4
+        val countY = 3
+        val index = 1
+        `when`(mockParent.indexOfChild(mockLayout)).thenReturn(index)
+        `when`(mockLayout.countX).thenReturn(countX)
+        `when`(mockLayout.countY).thenReturn(countY)
+
+        folderAccessibilityHelper = FolderAccessibilityHelper(mockLayout)
+    }
+
+    // Test for intersectsValidDropTarget()
+    @Test
+    fun testIntersectsValidDropTarget() {
+        // Setup
+        val id = 5
+        val allocatedContentSize = 20
+        // Make layout function public @VisibleForTesting
+        `when`(mockParent.allocatedContentSize).thenReturn(allocatedContentSize)
+
+        // Execute
+        val result = folderAccessibilityHelper.intersectsValidDropTarget(id)
+
+        // Verify
+        val expectedResult = min(id, allocatedContentSize - (index * countX * countY) - 1)
+        assertEquals(expectedResult, result)
+    }
+
+    // Test for getLocationDescriptionForIconDrop()
+    @Test
+    fun testGetLocationDescriptionForIconDrop() {
+        // Setup
+        val id = 5
+
+        // Execute
+        val result = folderAccessibilityHelper.getLocationDescriptionForIconDrop(id)
+
+        // Verify
+        val expectedResult = "Move to position ${id + (index * countX * countY) + 1}"
+        assertEquals(expectedResult, result)
+    }
+
+    // Test for getConfirmationForIconDrop()
+    @Test
+    fun testGetConfirmationForIconDrop() {
+        // Execute
+        val result =
+            folderAccessibilityHelper.getConfirmationForIconDrop(0) // Id doesn't matter here
+
+        // Verify
+        assertEquals("Item moved", result)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/celllayout/CellLayoutMethodsTest.kt b/tests/multivalentTests/src/com/android/launcher3/celllayout/CellLayoutMethodsTest.kt
new file mode 100644
index 0000000..e8459d6
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/celllayout/CellLayoutMethodsTest.kt
@@ -0,0 +1,69 @@
+/*
+ * 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.celllayout
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class CellLayoutMethodsTest {
+
+    @JvmField @Rule var cellLayoutBuilder = UnitTestCellLayoutBuilderRule()
+
+    @Test
+    fun pointToCellExact() {
+        val width = 1000
+        val height = 1000
+        val columns = 30
+        val rows = 30
+        val cl = cellLayoutBuilder.createCellLayout(columns, rows, false, width, height)
+
+        val res = intArrayOf(0, 0)
+        for (col in 0..<columns) {
+            for (row in 0..<rows) {
+                val x = (width / columns) * col
+                val y = (height / rows) * row
+                cl.pointToCellExact(x, y, res)
+                cl.pointToCellExact(x, y, res)
+                assertValues(col, res, row, columns, rows, width, height, x, y)
+            }
+        }
+
+        cl.pointToCellExact(-10, -10, res)
+        assertValues(0, res, 0, columns, rows, width, height, -10, -10)
+        cl.pointToCellExact(width + 10, height + 10, res)
+        assertValues(columns - 1, res, rows - 1, columns, rows, width, height, -10, -10)
+    }
+
+    private fun assertValues(
+        col: Int,
+        res: IntArray,
+        row: Int,
+        columns: Int,
+        rows: Int,
+        width: Int,
+        height: Int,
+        x: Int,
+        y: Int
+    ) {
+        assert(col == res[0] && row == res[1]) {
+            "Cell Layout with values (c= $columns, r= $rows, w= $width, h= $height) didn't mapped correctly the pixels ($x, $y) to the cells ($col, $row) with result (${res[0]}, ${res[1]})"
+        }
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/celllayout/ReorderAlgorithmUnitTest.java b/tests/multivalentTests/src/com/android/launcher3/celllayout/ReorderAlgorithmUnitTest.java
index 8a9711d..30953cc 100644
--- a/tests/multivalentTests/src/com/android/launcher3/celllayout/ReorderAlgorithmUnitTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/celllayout/ReorderAlgorithmUnitTest.java
@@ -144,8 +144,8 @@
 
     public ItemConfiguration solve(CellLayoutBoard board, int x, int y, int spanX,
             int spanY, int minSpanX, int minSpanY, boolean isMulti) {
-        CellLayout cl = mCellLayoutBuilder.createCellLayout(board.getWidth(), board.getHeight(),
-                isMulti);
+        CellLayout cl = mCellLayoutBuilder.createCellLayoutDefaultSize(board.getWidth(),
+                board.getHeight(), isMulti);
 
         // The views have to be sorted or the result can vary
         board.getIcons()
diff --git a/tests/multivalentTests/src/com/android/launcher3/celllayout/UnitTestCellLayoutBuilderRule.kt b/tests/multivalentTests/src/com/android/launcher3/celllayout/UnitTestCellLayoutBuilderRule.kt
index b63966d..f624be1 100644
--- a/tests/multivalentTests/src/com/android/launcher3/celllayout/UnitTestCellLayoutBuilderRule.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/celllayout/UnitTestCellLayoutBuilderRule.kt
@@ -64,11 +64,21 @@
         dp.inv.numRows = prevNumRows
     }
 
-    fun createCellLayout(width: Int, height: Int, isMulti: Boolean): CellLayout {
+    fun createCellLayoutDefaultSize(columns: Int, rows: Int, isMulti: Boolean): CellLayout {
+        return createCellLayout(columns, rows, isMulti)
+    }
+
+    fun createCellLayout(
+        columns: Int,
+        rows: Int,
+        isMulti: Boolean,
+        width: Int = 1000,
+        height: Int = 1000
+    ): CellLayout {
         val dp = getDeviceProfile()
         // modify the device profile.
-        dp.inv.numColumns = if (isMulti) width / 2 else width
-        dp.inv.numRows = height
+        dp.inv.numColumns = if (isMulti) columns / 2 else columns
+        dp.inv.numRows = rows
         dp.cellLayoutBorderSpacePx = Point(0, 0)
         val cl =
             if (isMulti) MultipageCellLayout(getWrappedContext(applicationContext, dp))
@@ -76,8 +86,8 @@
         // I put a very large number for width and height so that all the items can fit, it doesn't
         // need to be exact, just bigger than the sum of cell border
         cl.measure(
-            View.MeasureSpec.makeMeasureSpec(10000, View.MeasureSpec.EXACTLY),
-            View.MeasureSpec.makeMeasureSpec(10000, View.MeasureSpec.EXACTLY)
+            View.MeasureSpec.makeMeasureSpec(width, View.MeasureSpec.EXACTLY),
+            View.MeasureSpec.makeMeasureSpec(height, View.MeasureSpec.EXACTLY)
         )
         return cl
     }
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
new file mode 100644
index 0000000..71f7d47
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
@@ -0,0 +1,209 @@
+/*
+ * 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.model
+
+import android.appwidget.AppWidgetManager
+import android.content.ComponentName
+import android.content.Context
+import android.os.UserHandle
+import android.platform.test.rule.AllowedDevices
+import android.platform.test.rule.DeviceProduct
+import android.platform.test.rule.LimitDevicesRule
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.data.PackageItemInfo
+import com.android.launcher3.pm.UserCache
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.ComponentKey
+import com.android.launcher3.util.Executors
+import com.android.launcher3.util.IntSet
+import com.android.launcher3.util.PackageUserKey
+import com.android.launcher3.util.WidgetUtils.createAppWidgetProviderInfo
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.android.launcher3.widget.WidgetSections
+import com.android.launcher3.widget.WidgetSections.NO_CATEGORY
+import com.google.common.truth.Truth.assertThat
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import org.junit.Assert.fail
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.spy
+import org.mockito.junit.MockitoJUnit
+import org.mockito.junit.MockitoRule
+import org.mockito.kotlin.any
+import org.mockito.kotlin.whenever
+
+@AllowedDevices(allowed = [DeviceProduct.ROBOLECTRIC])
+@RunWith(AndroidJUnit4::class)
+class WidgetsModelTest {
+    @Rule @JvmField val limitDevicesRule = LimitDevicesRule()
+    @Rule @JvmField val mockitoRule: MockitoRule = MockitoJUnit.rule()
+
+    @Mock private lateinit var appWidgetManager: AppWidgetManager
+    @Mock private lateinit var app: LauncherAppState
+    @Mock private lateinit var iconCacheMock: IconCache
+
+    private lateinit var context: Context
+    private lateinit var idp: InvariantDeviceProfile
+    private lateinit var underTest: WidgetsModel
+
+    private var widgetSectionCategory: Int = 0
+    private lateinit var appAPackage: String
+
+    @Before
+    fun setUp() {
+        val appContext: Context = ApplicationProvider.getApplicationContext()
+        idp = InvariantDeviceProfile.INSTANCE[appContext]
+
+        context =
+            object : ActivityContextWrapper(ApplicationProvider.getApplicationContext()) {
+                override fun getSystemService(name: String): Any? {
+                    if (name == "appwidget") {
+                        return appWidgetManager
+                    }
+                    return super.getSystemService(name)
+                }
+
+                override fun getDeviceProfile(): DeviceProfile {
+                    return idp.getDeviceProfile(applicationContext).copy(applicationContext)
+                }
+            }
+
+        whenever(iconCacheMock.getTitleNoCache(any<LauncherAppWidgetProviderInfo>()))
+            .thenReturn("title")
+        whenever(app.iconCache).thenReturn(iconCacheMock)
+        whenever(app.context).thenReturn(context)
+        whenever(app.invariantDeviceProfile).thenReturn(idp)
+
+        val widgetToCategoryEntry: Map.Entry<ComponentName, IntSet> =
+            WidgetSections.getWidgetsToCategory(context).entries.first()
+        widgetSectionCategory = widgetToCategoryEntry.value.first()
+        val appAWidgetComponent = widgetToCategoryEntry.key
+        appAPackage = appAWidgetComponent.packageName
+
+        whenever(appWidgetManager.getInstalledProvidersForProfile(any()))
+            .thenReturn(
+                listOf(
+                    // First widget from widget sections xml
+                    createAppWidgetProviderInfo(appAWidgetComponent),
+                    // A widget that belongs to same package as the widget from widget sections
+                    // xml, but, because it's not mentioned in xml, it would be included in its
+                    // own package section.
+                    createAppWidgetProviderInfo(
+                        ComponentName.createRelative(appAPackage, APP_A_TEST_WIDGET_NAME)
+                    ),
+                    // A widget in different package (none of that app's widgets are in widget
+                    // sections xml)
+                    createAppWidgetProviderInfo(AppBTestWidgetComponent),
+                )
+            )
+
+        val userCache = spy(UserCache.INSTANCE.get(context))
+        whenever(userCache.userProfiles).thenReturn(listOf(UserHandle.CURRENT))
+
+        underTest = WidgetsModel()
+    }
+
+    @Test
+    fun widgetsByPackage_treatsWidgetSectionsAsSeparatePackageItems() {
+        loadWidgets()
+
+        val packages: Map<PackageItemInfo, List<WidgetItem>> = underTest.widgetsByPackageItem
+
+        // expect 3 package items
+        // one for the custom section with widget from appA
+        // one for package section for second widget from appA (that wasn't listed in xml)
+        // and one for package section for appB
+        assertThat(packages).hasSize(3)
+
+        // Each package item when used as a key is distinct (i.e. even if appA is split into custom
+        // package and owner package section, each of them is a distinct key). This ensures that
+        // clicking on a custom widget section doesn't take user to app package section.
+        val distinctPackageUserKeys =
+            packages.map { PackageUserKey.fromPackageItemInfo(it.key) }.distinct()
+        assertThat(distinctPackageUserKeys).hasSize(3)
+
+        val customSections = packages.filter { it.key.widgetCategory == widgetSectionCategory }
+        assertThat(customSections).hasSize(1)
+        val widgetsInCustomSection = customSections.entries.first().value
+        assertThat(widgetsInCustomSection).hasSize(1)
+
+        val packageSections = packages.filter { it.key.widgetCategory == NO_CATEGORY }
+        assertThat(packageSections).hasSize(2)
+
+        // App A's package section
+        val appAPackageSection = packageSections.filter { it.key.packageName == appAPackage }
+        assertThat(appAPackageSection).hasSize(1)
+        val widgetsInAppASection = appAPackageSection.entries.first().value
+        assertThat(widgetsInAppASection).hasSize(1)
+
+        // App B's package section
+        val appBPackageSection =
+            packageSections.filter { it.key.packageName == AppBTestWidgetComponent.packageName }
+        assertThat(appBPackageSection).hasSize(1)
+        val widgetsInAppBSection = appBPackageSection.entries.first().value
+        assertThat(widgetsInAppBSection).hasSize(1)
+    }
+
+    @Test
+    fun widgetComponentMap_returnsWidgets() {
+        loadWidgets()
+
+        val widgetsByComponentKey: Map<ComponentKey, WidgetItem> = underTest.widgetsByComponentKey
+
+        assertThat(widgetsByComponentKey).hasSize(3)
+        widgetsByComponentKey.forEach { entry ->
+            assertThat(entry.key).isEqualTo(entry.value as ComponentKey)
+        }
+    }
+
+    @Test
+    fun widgets_noData_returnsEmpty() {
+        // no loadWidgets()
+
+        assertThat(underTest.widgetsByComponentKey).isEmpty()
+    }
+
+    private fun loadWidgets() {
+        val latch = CountDownLatch(1)
+        Executors.MODEL_EXECUTOR.execute {
+            underTest.update(app, /* packageUser= */ null)
+            latch.countDown()
+        }
+        if (!latch.await(LOAD_WIDGETS_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+            fail("Timed out waiting widgets to load")
+        }
+    }
+
+    companion object {
+        // Another widget within app A
+        private const val APP_A_TEST_WIDGET_NAME = "MyProvider"
+
+        private val AppBTestWidgetComponent: ComponentName =
+            ComponentName.createRelative("com.test.package", "TestProvider")
+
+        private const val LOAD_WIDGETS_TIMEOUT_SECONDS = 2L
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/ItemInfoMatcherTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/ItemInfoMatcherTest.kt
new file mode 100644
index 0000000..daba024
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/ItemInfoMatcherTest.kt
@@ -0,0 +1,213 @@
+/*
+ * 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.util
+
+import android.content.ComponentName
+import android.content.Intent
+import android.os.UserHandle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
+import com.android.launcher3.model.data.FolderInfo
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.shortcuts.ShortcutKey
+import com.android.launcher3.util.ItemInfoMatcher.ofShortcutKeys
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ItemInfoMatcherTest {
+
+    @Test
+    fun `ofUser returns Predicate for ItemInfo containing given UserHandle`() {
+        val expectedItemInfo = ItemInfo().apply { user = UserHandle(11) }
+        val unexpectedItemInfo = ItemInfo().apply { user = UserHandle(0) }
+        val itemInfoStream = listOf(expectedItemInfo, unexpectedItemInfo).stream()
+
+        val predicate = ItemInfoMatcher.ofUser(UserHandle(11))
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+
+    @Test
+    fun `ofComponents returns Predicate for ItemInfo containing target Component and UserHandle`() {
+        // Given
+        val expectedUserHandle = UserHandle(0)
+        val expectedComponentName = ComponentName("expectedPackage", "expectedClass")
+        val expectedItemInfo = spy(ItemInfo())
+        expectedItemInfo.user = expectedUserHandle
+        whenever(expectedItemInfo.targetComponent).thenReturn(expectedComponentName)
+
+        val unexpectedComponentName = ComponentName("unexpectedPackage", "unexpectedClass")
+        val unexpectedItemInfo1 = spy(ItemInfo())
+        unexpectedItemInfo1.user = expectedUserHandle
+        whenever(unexpectedItemInfo1.targetComponent).thenReturn(unexpectedComponentName)
+
+        val unexpectedItemInfo2 = spy(ItemInfo())
+        unexpectedItemInfo2.user = UserHandle(10)
+        whenever(unexpectedItemInfo2.targetComponent).thenReturn(expectedComponentName)
+
+        val itemInfoStream =
+            listOf(expectedItemInfo, unexpectedItemInfo1, unexpectedItemInfo2).stream()
+
+        // When
+        val predicate =
+            ItemInfoMatcher.ofComponents(hashSetOf(expectedComponentName), expectedUserHandle)
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        // Then
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+
+    @Test
+    fun `ofPackages returns Predicate for ItemInfo containing UserHandle and target package`() {
+        // Given
+        val expectedUserHandle = UserHandle(0)
+        val expectedPackage = "expectedPackage"
+        val expectedComponentName = ComponentName(expectedPackage, "expectedClass")
+        val expectedItemInfo = spy(ItemInfo())
+        expectedItemInfo.user = expectedUserHandle
+        whenever(expectedItemInfo.targetComponent).thenReturn(expectedComponentName)
+
+        val unexpectedPackage = "unexpectedPackage"
+        val unexpectedComponentName = ComponentName(unexpectedPackage, "unexpectedClass")
+        val unexpectedItemInfo1 = spy(ItemInfo())
+        unexpectedItemInfo1.user = expectedUserHandle
+        whenever(unexpectedItemInfo1.targetComponent).thenReturn(unexpectedComponentName)
+
+        val unexpectedItemInfo2 = spy(ItemInfo())
+        unexpectedItemInfo2.user = UserHandle(10)
+        whenever(unexpectedItemInfo2.targetComponent).thenReturn(expectedComponentName)
+
+        val itemInfoStream =
+            listOf(expectedItemInfo, unexpectedItemInfo1, unexpectedItemInfo2).stream()
+
+        // When
+        val predicate = ItemInfoMatcher.ofPackages(setOf(expectedPackage), expectedUserHandle)
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        // Then
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+
+    @Test
+    fun `ofShortcutKeys returns Predicate for Deep Shortcut Info containing given ShortcutKey`() {
+        // Given
+        val expectedItemInfo = spy(ItemInfo())
+        expectedItemInfo.itemType = ITEM_TYPE_DEEP_SHORTCUT
+        val expectedIntent =
+            Intent().apply {
+                putExtra("shortcut_id", "expectedShortcut")
+                `package` = "expectedPackage"
+            }
+        whenever(expectedItemInfo.intent).thenReturn(expectedIntent)
+
+        val unexpectedIntent =
+            Intent().apply {
+                putExtra("shortcut_id", "unexpectedShortcut")
+                `package` = "unexpectedPackage"
+            }
+        val unexpectedItemInfo = spy(ItemInfo())
+        unexpectedItemInfo.itemType = ITEM_TYPE_DEEP_SHORTCUT
+        whenever(unexpectedItemInfo.intent).thenReturn(unexpectedIntent)
+
+        val itemInfoStream = listOf(expectedItemInfo, unexpectedItemInfo).stream()
+        val expectedShortcutKey = ShortcutKey.fromItemInfo(expectedItemInfo)
+
+        // When
+        val predicate = ItemInfoMatcher.ofShortcutKeys(setOf(expectedShortcutKey))
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        // Then
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+
+    @Test
+    fun `forFolderMatch returns Predicate to match against children within Folder ItemInfo`() {
+        // Given
+        val expectedItemInfo = spy(FolderInfo())
+        expectedItemInfo.itemType = ITEM_TYPE_FOLDER
+        val expectedIntent =
+            Intent().apply {
+                putExtra("shortcut_id", "expectedShortcut")
+                `package` = "expectedPackage"
+            }
+        val expectedChildInfo = spy(ItemInfo())
+        expectedChildInfo.itemType = ITEM_TYPE_DEEP_SHORTCUT
+        whenever(expectedChildInfo.intent).thenReturn(expectedIntent)
+        whenever(expectedItemInfo.getContents()).thenReturn(arrayListOf(expectedChildInfo))
+
+        val unexpectedItemInfo = spy(FolderInfo())
+        unexpectedItemInfo.itemType = ITEM_TYPE_FOLDER
+
+        val itemInfoStream = listOf(expectedItemInfo, unexpectedItemInfo).stream()
+        val expectedShortcutKey = ShortcutKey.fromItemInfo(expectedChildInfo)
+
+        // When
+        val predicate = ItemInfoMatcher.forFolderMatch(ofShortcutKeys(setOf(expectedShortcutKey)))
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        // Then
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+
+    @Test
+    fun `ofItemIds returns Predicate to match ItemInfo that contains given ids`() {
+        // Given
+        val expectedItemInfo = spy(ItemInfo())
+        expectedItemInfo.id = 1
+
+        val unexpectedItemInfo = spy(ItemInfo())
+        unexpectedItemInfo.id = 2
+
+        val itemInfoStream = listOf(expectedItemInfo, unexpectedItemInfo).stream()
+
+        // When
+        val expectedIds = IntSet().apply { add(1) }
+        val predicate = ItemInfoMatcher.ofItemIds(expectedIds)
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        // Then
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+
+    @Test
+    fun `ofItems returns Predicate matching against provided ItemInfo`() {
+        // Given
+        val expectedItemInfo = spy(ItemInfo())
+        expectedItemInfo.id = 1
+
+        val unexpectedItemInfo = spy(ItemInfo())
+        unexpectedItemInfo.id = 2
+
+        val itemInfoStream = listOf(expectedItemInfo, unexpectedItemInfo).stream()
+
+        // When
+        val expectedItems = setOf(expectedItemInfo)
+        val predicate = ItemInfoMatcher.ofItems(expectedItems)
+        val actualResults = itemInfoStream.filter(predicate).toList()
+
+        // Then
+        assertThat(actualResults).containsExactly(expectedItemInfo)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
new file mode 100644
index 0000000..0a3035a
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfoTest.kt
@@ -0,0 +1,58 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget.custom
+
+import android.content.ComponentName
+import android.content.pm.PackageManager
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.MediumTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import org.junit.Assert.assertEquals
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.mock
+
+@MediumTest
+@RunWith(AndroidJUnit4::class)
+class CustomAppWidgetProviderInfoTest {
+
+    private lateinit var underTest: CustomAppWidgetProviderInfo
+
+    @Before
+    fun setup() {
+        underTest = CustomAppWidgetProviderInfo()
+        underTest.provider = PROVIDER_NAME
+    }
+
+    @Test
+    fun info_to_string() {
+        assertEquals("WidgetProviderInfo($PROVIDER_NAME)", underTest.toString())
+    }
+
+    @Test
+    fun get_label() {
+        underTest.label = "  TEST_LABEL"
+        assertEquals(LABEL_NAME, underTest.getLabel(mock(PackageManager::class.java)))
+    }
+
+    companion object {
+        private val PROVIDER_NAME =
+            ComponentName(getInstrumentation().targetContext.packageName, "TEST_PACKAGE")
+        private const val LABEL_NAME = "TEST_LABEL"
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomWidgetManagerTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomWidgetManagerTest.kt
new file mode 100644
index 0000000..4b5710d
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/custom/CustomWidgetManagerTest.kt
@@ -0,0 +1,144 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget.custom
+
+import android.appwidget.AppWidgetManager
+import android.content.ComponentName
+import android.os.Process
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.android.launcher3.util.PluginManagerWrapper
+import com.android.launcher3.util.WidgetUtils
+import com.android.launcher3.widget.LauncherAppWidgetHostView
+import com.android.launcher3.widget.LauncherAppWidgetProviderInfo
+import com.android.systemui.plugins.CustomWidgetPlugin
+import org.junit.After
+import org.junit.Assert.assertEquals
+import org.junit.Assert.assertTrue
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.mock
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.same
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class CustomWidgetManagerTest {
+
+    @get:Rule val setFlagsRule = SetFlagsRule()
+
+    private val context = SandboxModelContext()
+    private lateinit var underTest: CustomWidgetManager
+
+    @Mock private lateinit var pluginManager: PluginManagerWrapper
+    @Mock private lateinit var mockAppWidgetManager: AppWidgetManager
+
+    @Before
+    fun setUp() {
+        MockitoAnnotations.initMocks(this)
+        context.putObject(PluginManagerWrapper.INSTANCE, pluginManager)
+        underTest = CustomWidgetManager(context, mockAppWidgetManager)
+    }
+
+    @After
+    fun tearDown() {
+        underTest.close()
+    }
+
+    @Test
+    fun plugin_manager_added_after_initialization() {
+        verify(pluginManager)
+            .addPluginListener(same(underTest), same(CustomWidgetPlugin::class.java), eq(true))
+    }
+
+    @Test
+    fun close_widget_manager_should_remove_plugin_listener() {
+        underTest.close()
+        verify(pluginManager).removePluginListener(same(underTest))
+    }
+
+    @Test
+    fun on_plugin_connected_no_provider_info() {
+        doReturn(emptyList<LauncherAppWidgetProviderInfo>())
+            .whenever(mockAppWidgetManager)
+            .getInstalledProvidersForProfile(any())
+        val mockPlugin = mock(CustomWidgetPlugin::class.java)
+        underTest.onPluginConnected(mockPlugin, context)
+        assertEquals(0, underTest.plugins.size)
+    }
+
+    @Test
+    fun on_plugin_connected_exist_provider_info() {
+        doReturn(listOf(WidgetUtils.createAppWidgetProviderInfo(TEST_COMPONENT_NAME)))
+            .whenever(mockAppWidgetManager)
+            .getInstalledProvidersForProfile(eq(Process.myUserHandle()))
+        val mockPlugin = mock(CustomWidgetPlugin::class.java)
+        underTest.onPluginConnected(mockPlugin, context)
+        assertEquals(1, underTest.plugins.size)
+    }
+
+    @Test
+    fun on_plugin_disconnected() {
+        doReturn(listOf(WidgetUtils.createAppWidgetProviderInfo(TEST_COMPONENT_NAME)))
+            .whenever(mockAppWidgetManager)
+            .getInstalledProvidersForProfile(eq(Process.myUserHandle()))
+        val mockPlugin = mock(CustomWidgetPlugin::class.java)
+        underTest.onPluginConnected(mockPlugin, context)
+        underTest.onPluginDisconnected(mockPlugin)
+        assertEquals(0, underTest.plugins.size)
+    }
+
+    @Test
+    fun on_view_created() {
+        val mockPlugin = mock(CustomWidgetPlugin::class.java)
+        val mockWidgetView = mock(LauncherAppWidgetHostView::class.java)
+        val mockProviderInfo = mock(CustomAppWidgetProviderInfo::class.java)
+        doReturn(mockProviderInfo).whenever(mockWidgetView).appWidgetInfo
+        mockProviderInfo.provider = TEST_COMPONENT_NAME
+        underTest.plugins.put(TEST_COMPONENT_NAME, mockPlugin)
+        underTest.onViewCreated(mockWidgetView)
+        verify(mockPlugin).onViewCreated(eq(mockWidgetView))
+    }
+
+    @Test
+    fun generate_stream() {
+        assertTrue(underTest.stream().toList().isEmpty())
+        doReturn(listOf(WidgetUtils.createAppWidgetProviderInfo(TEST_COMPONENT_NAME)))
+            .whenever(mockAppWidgetManager)
+            .getInstalledProvidersForProfile(eq(Process.myUserHandle()))
+        val mockPlugin = mock(CustomWidgetPlugin::class.java)
+        underTest.onPluginConnected(mockPlugin, context)
+        assertEquals(1, underTest.stream().toList().size)
+    }
+
+    companion object {
+        private const val TEST_CLASS = "TEST_CLASS"
+        private val TEST_COMPONENT_NAME =
+            ComponentName(getInstrumentation().targetContext.packageName, TEST_CLASS)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
index 0370a6b..24d66a3 100644
--- a/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/picker/search/SimpleWidgetsSearchAlgorithmTest.java
@@ -45,12 +45,12 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.data.PackageItemInfo;
-import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.search.SearchCallback;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
+import com.android.launcher3.widget.picker.search.WidgetsSearchBar.WidgetsSearchDataProvider;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -79,7 +79,7 @@
 
     private SimpleWidgetsSearchAlgorithm mSimpleWidgetsSearchAlgorithm;
     @Mock
-    private PopupDataProvider mDataProvider;
+    private WidgetsSearchDataProvider mDataProvider;
     @Mock
     private SearchCallback<WidgetsListBaseEntry> mSearchCallback;
 
@@ -106,7 +106,7 @@
 
         mSimpleWidgetsSearchAlgorithm = MAIN_EXECUTOR.submit(
                 () -> new SimpleWidgetsSearchAlgorithm(mDataProvider)).get();
-        doReturn(Collections.EMPTY_LIST).when(mDataProvider).getAllWidgets();
+        doReturn(Collections.EMPTY_LIST).when(mDataProvider).getWidgets();
     }
 
     @Test
@@ -114,7 +114,7 @@
         doReturn(List.of(mCalendarHeaderEntry, mCalendarContentEntry, mCameraHeaderEntry,
                 mCameraContentEntry, mClockHeaderEntry, mClockContentEntry))
                 .when(mDataProvider)
-                .getAllWidgets();
+                .getWidgets();
 
         assertEquals(List.of(
                 WidgetsListHeaderEntry.createForSearch(
@@ -135,7 +135,7 @@
         doReturn(List.of(mCalendarHeaderEntry, mCalendarContentEntry, mCameraHeaderEntry,
                 mCameraContentEntry))
                 .when(mDataProvider)
-                .getAllWidgets();
+                .getWidgets();
 
         assertEquals(List.of(
                 WidgetsListHeaderEntry.createForSearch(
@@ -162,7 +162,7 @@
         doReturn(List.of(mCalendarHeaderEntry, mCalendarContentEntry, mCameraHeaderEntry,
                 mCameraContentEntry, mClockHeaderEntry, mClockContentEntry))
                 .when(mDataProvider)
-                .getAllWidgets();
+                .getWidgets();
         mSimpleWidgetsSearchAlgorithm.doSearch("Ca", mSearchCallback);
         getInstrumentation().waitForIdleSync();
         verify(mSearchCallback).onSearchResult(
diff --git a/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt b/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt
new file mode 100644
index 0000000..d9af07a
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/PackageUpdatedTaskTest.kt
@@ -0,0 +1,235 @@
+/*
+ * 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.model
+
+import android.content.ComponentName
+import android.content.pm.ActivityInfo
+import android.content.pm.ApplicationInfo
+import android.content.pm.LauncherActivityInfo
+import android.content.pm.LauncherApps
+import android.os.UserHandle
+import android.platform.test.annotations.EnableFlags
+import android.platform.test.flag.junit.SetFlagsRule
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.AppFilter
+import com.android.launcher3.Flags.FLAG_ENABLE_PRIVATE_SPACE
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings
+import com.android.launcher3.icons.IconCache
+import com.android.launcher3.model.PackageUpdatedTask.OP_ADD
+import com.android.launcher3.model.PackageUpdatedTask.OP_REMOVE
+import com.android.launcher3.model.PackageUpdatedTask.OP_SUSPEND
+import com.android.launcher3.model.PackageUpdatedTask.OP_UNAVAILABLE
+import com.android.launcher3.model.PackageUpdatedTask.OP_UNSUSPEND
+import com.android.launcher3.model.PackageUpdatedTask.OP_UPDATE
+import com.android.launcher3.model.PackageUpdatedTask.OP_USER_AVAILABILITY_CHANGE
+import com.android.launcher3.model.data.AppInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.Executors
+import com.android.launcher3.util.LauncherModelHelper
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
+import com.android.launcher3.util.TestUtil
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
+
+@RunWith(AndroidJUnit4::class)
+class PackageUpdatedTaskTest {
+
+    @get:Rule val setFlagsRule = SetFlagsRule()
+
+    private val mUser = UserHandle(0)
+    private val mDataModel: BgDataModel = BgDataModel()
+    private val mLauncherModelHelper = LauncherModelHelper()
+    private val mContext: SandboxModelContext = spy(mLauncherModelHelper.sandboxContext)
+    private val mAppState: LauncherAppState = spy(LauncherAppState.getInstance(mContext))
+
+    private val expectedPackage = "Test.Package"
+    private val expectedComponent = ComponentName(expectedPackage, "TestClass")
+    private val expectedActivityInfo: LauncherActivityInfo = mock<LauncherActivityInfo>()
+    private val expectedWorkspaceItem = spy(WorkspaceItemInfo())
+
+    private val mockIconCache: IconCache = mock()
+    private val mockTaskController: ModelTaskController = mock<ModelTaskController>()
+    private val mockAppFilter: AppFilter = mock<AppFilter>()
+    private val mockApplicationInfo: ApplicationInfo = mock<ApplicationInfo>()
+    private val mockActivityInfo: ActivityInfo = mock<ActivityInfo>()
+
+    private lateinit var mAllAppsList: AllAppsList
+
+    @Before
+    fun setup() {
+        mAllAppsList = spy(AllAppsList(mockIconCache, mockAppFilter))
+        mLauncherModelHelper.sandboxContext.spyService(LauncherApps::class.java).apply {
+            whenever(getActivityList(expectedPackage, mUser))
+                .thenReturn(listOf(expectedActivityInfo))
+        }
+        whenever(mAppState.iconCache).thenReturn(mockIconCache)
+        whenever(mockTaskController.app).thenReturn(mAppState)
+        whenever(mockAppFilter.shouldShowApp(expectedComponent)).thenReturn(true)
+        mockApplicationInfo.apply {
+            uid = 1
+            isArchived = false
+        }
+        mockActivityInfo.isArchived = false
+        expectedActivityInfo.apply {
+            whenever(applicationInfo).thenReturn(mockApplicationInfo)
+            whenever(activityInfo).thenReturn(mockActivityInfo)
+            whenever(componentName).thenReturn(expectedComponent)
+        }
+        expectedWorkspaceItem.apply {
+            itemType = LauncherSettings.Favorites.ITEM_TYPE_APPLICATION
+            container = LauncherSettings.Favorites.CONTAINER_DESKTOP
+            user = mUser
+            whenever(targetPackage).thenReturn(expectedPackage)
+            whenever(targetComponent).thenReturn(expectedComponent)
+        }
+    }
+
+    @After
+    fun tearDown() {
+        mLauncherModelHelper.destroy()
+    }
+
+    @Test
+    fun `OP_ADD triggers model callbacks and adds new items to AllAppsList`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_ADD, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mockIconCache).updateIconsForPkg(expectedPackage, mUser)
+        verify(mAllAppsList).addPackage(mContext, expectedPackage, mUser)
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWorkspaceItem))
+        verify(mockTaskController).bindUpdatedWidgets(mDataModel)
+        assertThat(mAllAppsList.data.firstOrNull()?.componentName)
+            .isEqualTo(AppInfo(mContext, expectedActivityInfo, mUser).componentName)
+    }
+
+    @Test
+    fun `OP_UPDATE triggers model callbacks and updates items in AllAppsList`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_UPDATE, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mockIconCache).updateIconsForPkg(expectedPackage, mUser)
+        verify(mAllAppsList).updatePackage(mContext, expectedPackage, mUser)
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWorkspaceItem))
+        assertThat(mAllAppsList.data.firstOrNull()?.componentName)
+            .isEqualTo(AppInfo(mContext, expectedActivityInfo, mUser).componentName)
+    }
+
+    @Test
+    fun `OP_REMOVE triggers model callbacks and removes packages and icons`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_REMOVE, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mockIconCache).removeIconsForPkg(expectedPackage, mUser)
+        verify(mAllAppsList).removePackage(expectedPackage, mUser)
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWorkspaceItem))
+        assertThat(mAllAppsList.data).isEmpty()
+    }
+
+    @Test
+    fun `OP_UNAVAILABLE triggers model callbacks and removes package from AllAppsList`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_UNAVAILABLE, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mAllAppsList).removePackage(expectedPackage, mUser)
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWorkspaceItem))
+        assertThat(mAllAppsList.data).isEmpty()
+    }
+
+    @Test
+    fun `OP_SUSPEND triggers model callbacks and updates flags in AllAppsList`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_SUSPEND, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        mAllAppsList.add(AppInfo(mContext, expectedActivityInfo, mUser), expectedActivityInfo)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mAllAppsList).updateDisabledFlags(any(), any())
+        verify(mockTaskController).bindUpdatedWorkspaceItems(listOf(expectedWorkspaceItem))
+        assertThat(mAllAppsList.getAndResetChangeFlag()).isTrue()
+    }
+
+    @Test
+    fun `OP_UNSUSPEND triggers no callbacks when app not suspended`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_UNSUSPEND, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mAllAppsList).updateDisabledFlags(any(), any())
+        verify(mockTaskController).bindUpdatedWorkspaceItems(emptyList())
+        assertThat(mAllAppsList.getAndResetChangeFlag()).isFalse()
+    }
+
+    @EnableFlags(FLAG_ENABLE_PRIVATE_SPACE)
+    @Test
+    fun `OP_USER_AVAILABILITY_CHANGE triggers no callbacks if current user not work or private`() {
+        // Given
+        val taskUnderTest = PackageUpdatedTask(OP_USER_AVAILABILITY_CHANGE, mUser, expectedPackage)
+        // When
+        mDataModel.addItem(mContext, expectedWorkspaceItem, true)
+        TestUtil.runOnExecutorSync(Executors.MODEL_EXECUTOR) {
+            taskUnderTest.execute(mockTaskController, mDataModel, mAllAppsList)
+        }
+        mLauncherModelHelper.loadModelSync()
+        // Then
+        verify(mAllAppsList).updateDisabledFlags(any(), any())
+        verify(mockTaskController).bindUpdatedWorkspaceItems(emptyList())
+        assertThat(mAllAppsList.data).isEmpty()
+    }
+}
diff --git a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
index f02a0c2..25eae44 100644
--- a/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
+++ b/tests/tapl/com/android/launcher3/tapl/LauncherInstrumentation.java
@@ -455,10 +455,6 @@
         getTestInfo(TestProtocol.REQUEST_ENABLE_ROTATION, Boolean.toString(on));
     }
 
-    public void setEnableSuggestion(boolean enableSuggestion) {
-        getTestInfo(TestProtocol.REQUEST_ENABLE_SUGGESTION, Boolean.toString(enableSuggestion));
-    }
-
     public boolean hadNontestEvents() {
         return getTestInfo(TestProtocol.REQUEST_GET_HAD_NONTEST_EVENTS)
                 .getBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD);
@@ -2459,7 +2455,8 @@
     }
 
     float getWindowCornerRadius() {
-        // TODO(b/197326121): Check if the touch is overlapping with the corners by offsetting
+        // Return a larger corner radius to ensure gesture calculated from the radius are offset to
+        // prevent overlapping
         final float tmpBuffer = 100f;
         final Resources resources = getResources();
         if (!supportsRoundedCornersOnWindows(resources)) {