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 0fa3fbc..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;
@@ -80,6 +79,7 @@
 import android.view.ViewGroup;
 import android.view.ViewTreeObserver;
 import android.view.WindowManager;
+import android.view.inputmethod.Flags;
 import android.widget.FrameLayout;
 import android.widget.ImageView;
 import android.widget.LinearLayout;
@@ -247,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);
         }
     }
@@ -274,7 +274,10 @@
                 InputMethodService.canImeRenderGesturalNavButtons() && mContext.imeDrawsImeNavBar();
         if (!mIsImeRenderingNavButtons) {
             // IME switcher
-            mImeSwitcherButton = addButton(R.drawable.ic_ime_switcher, BUTTON_IME_SWITCH,
+            final int switcherResId = Flags.imeSwitcherRevamp()
+                    ? com.android.internal.R.drawable.ic_ime_switcher_new
+                    : R.drawable.ic_ime_switcher;
+            mImeSwitcherButton = addButton(switcherResId, BUTTON_IME_SWITCH,
                     isThreeButtonNav ? mStartContextualContainer : mEndContextualContainer,
                     mControllers.navButtonController, R.id.ime_switcher);
             mPropertyHolders.add(new StatePropertyHolder(mImeSwitcherButton,
@@ -361,7 +364,7 @@
                 R.bool.floating_rotation_button_position_left);
         mControllers.rotationButtonController.setRotationButton(mFloatingRotationButton,
                 mRotationButtonListener);
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.init();
         }
 
@@ -369,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];
@@ -625,7 +628,7 @@
     }
 
     public void setWallpaperVisible(boolean isVisible) {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.setWallpaperVisibility(isVisible);
         }
     }
@@ -638,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) {
@@ -661,7 +664,7 @@
     }
 
     public void transitionTo(@BarTransitions.TransitionMode int barMode, boolean animate) {
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.transitionTo(barMode, animate);
         }
     }
@@ -765,7 +768,7 @@
 
     private void onDarkIntensityChanged() {
         updateNavButtonColor();
-        if (enableTaskbarOnPhones() && mContext.isPhoneButtonNavMode()) {
+        if (mContext.isPhoneMode()) {
             mTaskbarTransitions.onDarkIntensityChanged(mTaskbarNavButtonDarkIntensity.value);
         }
     }
@@ -1115,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/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/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/data/RecentsDeviceProfileRepository.kt b/quickstep/src/com/android/quickstep/recents/data/RecentsDeviceProfileRepository.kt
new file mode 100644
index 0000000..adf904c
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentsDeviceProfileRepository.kt
@@ -0,0 +1,36 @@
+/*
+ * 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.data
+
+import com.android.quickstep.views.RecentsViewContainer
+
+/**
+ * Repository for shrink down version of [com.android.launcher3.DeviceProfile] that only contains
+ * data related to Recents.
+ */
+class RecentsDeviceProfileRepository(private val container: RecentsViewContainer) {
+
+    fun getRecentsDeviceProfile() =
+        with(container.deviceProfile) { RecentsDeviceProfile(isLargeScreen = isTablet) }
+
+    /**
+     * Container to hold [com.android.launcher3.DeviceProfile] related to Recents.
+     *
+     * @property isLargeScreen whether the current device posture has a large screen
+     */
+    data class RecentsDeviceProfile(val isLargeScreen: Boolean)
+}
diff --git a/quickstep/src/com/android/quickstep/recents/data/RecentsRotationStateRepository.kt b/quickstep/src/com/android/quickstep/recents/data/RecentsRotationStateRepository.kt
new file mode 100644
index 0000000..6ead704
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/data/RecentsRotationStateRepository.kt
@@ -0,0 +1,35 @@
+/*
+ * 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.data
+
+import com.android.quickstep.util.RecentsOrientedState
+
+/**
+ * Repository for [RecentsRotationState] which holds orientation/rotation related information
+ * related to Recents
+ */
+class RecentsRotationStateRepository(private val state: RecentsOrientedState) {
+    fun getRecentsRotationState() =
+        with(state) { RecentsRotationState(activityRotation = recentsActivityRotation) }
+
+    /**
+     * Container to hold orientation/rotation related information related to Recents.
+     *
+     * @property activityRotation rotation of the activity hosting RecentsView
+     */
+    data class RecentsRotationState(val activityRotation: Int)
+}
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 5e55e2e..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
@@ -25,63 +26,90 @@
 import com.android.quickstep.views.RecentsView
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.systemui.shared.recents.model.Task
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.MainScope
-import kotlinx.coroutines.launch
+import kotlinx.coroutines.CoroutineName
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 
 /**
  * Helper for [TaskOverlayFactory.TaskOverlay] to interact with [TaskOverlayViewModel], this helper
  * should merge with [TaskOverlayFactory.TaskOverlay] when it's migrated to MVVM.
  */
 class TaskOverlayHelper(val task: Task, val overlay: TaskOverlayFactory.TaskOverlay<*>) {
-    private lateinit var job: Job
+    private lateinit var overlayInitializedScope: CoroutineScope
     private var uiState: TaskOverlayUiState = Disabled
 
     // TODO(b/335649589): Ideally create and obtain this from DI. This ViewModel should be scoped
     //  to [TaskView], and also shared between [TaskView] and [TaskThumbnailView]
     //  This is using a lazy for now because the dependencies cannot be obtained without DI.
-    private val taskOverlayViewModel by lazy {
+    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() {
-        // TODO(b/335396935): This should be changed to TaskView's scope.
-        job =
-            MainScope().launch {
-                taskOverlayViewModel.overlayState.collect {
-                    uiState = it
-                    if (it is Enabled) {
-                        Log.d(
-                            TAG,
-                            "initOverlay - taskId: ${task.key.id}, thumbnail: ${it.thumbnail}"
-                        )
-                        overlay.initOverlay(
-                            task,
-                            it.thumbnail,
-                            it.thumbnailMatrix,
-                            /* rotated= */ false
-                        )
-                    } else {
-                        Log.d(TAG, "reset - taskId: ${task.key.id}")
-                        overlay.reset()
-                    }
+        overlayInitializedScope =
+            CoroutineScope(SupervisorJob() + Dispatchers.Main + CoroutineName("TaskOverlayHelper"))
+        viewModel.overlayState
+            .onEach {
+                uiState = it
+                if (it is Enabled) {
+                    initOverlay(it)
+                } else {
+                    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() {
-        job.cancel()
+        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 3273809..1db1447 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -191,6 +191,8 @@
 import com.android.quickstep.TopTaskTracker;
 import com.android.quickstep.ViewUtils;
 import com.android.quickstep.orientation.RecentsPagedOrientationHandler;
+import com.android.quickstep.recents.data.RecentsDeviceProfileRepository;
+import com.android.quickstep.recents.data.RecentsRotationStateRepository;
 import com.android.quickstep.recents.data.TasksRepository;
 import com.android.quickstep.recents.viewmodel.RecentsViewData;
 import com.android.quickstep.util.ActiveGestureErrorDetector;
@@ -465,6 +467,10 @@
     public final RecentsViewData mRecentsViewData = new RecentsViewData();
     @Nullable
     public final TasksRepository mTasksRepository;
+    @Nullable
+    public final RecentsRotationStateRepository mOrientedStateRepository;
+    @Nullable
+    public final RecentsDeviceProfileRepository mDeviceProfileRepository;
 
     protected final RecentsOrientedState mOrientationState;
     protected final BaseContainerInterface<STATE_TYPE, CONTAINER_TYPE> mSizeStrategy;
@@ -822,8 +828,12 @@
         if (enableRefactorTaskThumbnail()) {
             mTasksRepository = new TasksRepository(
                     mModel, mModel.getThumbnailCache(), mModel.getIconCache());
+            mOrientedStateRepository = new RecentsRotationStateRepository(mOrientationState);
+            mDeviceProfileRepository = new RecentsDeviceProfileRepository(mContainer);
         } else {
             mTasksRepository = null;
+            mOrientedStateRepository = null;
+            mDeviceProfileRepository = null;
         }
 
         mClearAllButton = (ClearAllButton) LayoutInflater.from(context)
@@ -2810,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/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/recents/data/RecentsDeviceProfileRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/RecentsDeviceProfileRepositoryTest.kt
new file mode 100644
index 0000000..eff926d
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/RecentsDeviceProfileRepositoryTest.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.data
+
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.FakeInvariantDeviceProfileTest
+import com.android.quickstep.views.RecentsViewContainer
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [RecentsDeviceProfileRepository] */
+@RunWith(AndroidJUnit4::class)
+class RecentsDeviceProfileRepositoryTest : FakeInvariantDeviceProfileTest() {
+    private val recentsViewContainer = mock<RecentsViewContainer>()
+
+    private val systemUnderTest = RecentsDeviceProfileRepository(recentsViewContainer)
+
+    @Test
+    fun deviceProfileMappedCorrectly() {
+        initializeVarsForTablet()
+        val tabletDeviceProfile = newDP()
+        whenever(recentsViewContainer.deviceProfile).thenReturn(tabletDeviceProfile)
+
+        assertThat(systemUnderTest.getRecentsDeviceProfile())
+            .isEqualTo(RecentsDeviceProfileRepository.RecentsDeviceProfile(isLargeScreen = true))
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/RecentsRotationStateRepositoryTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/RecentsRotationStateRepositoryTest.kt
new file mode 100644
index 0000000..1f4da26
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/data/RecentsRotationStateRepositoryTest.kt
@@ -0,0 +1,41 @@
+/*
+ * 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.data
+
+import android.view.Surface.ROTATION_90
+import com.android.quickstep.util.RecentsOrientedState
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+/** Test for [RecentsRotationStateRepository] */
+class RecentsRotationStateRepositoryTest {
+    private val recentsOrientedState = mock<RecentsOrientedState>()
+
+    private val systemUnderTest = RecentsRotationStateRepository(recentsOrientedState)
+
+    @Test
+    fun orientedStateMappedCorrectly() {
+        whenever(recentsOrientedState.recentsActivityRotation).thenReturn(ROTATION_90)
+
+        assertThat(systemUnderTest.getRecentsRotationState())
+            .isEqualTo(
+                RecentsRotationStateRepository.RecentsRotationState(activityRotation = ROTATION_90)
+            )
+    }
+}
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/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
index ec245ee..c24e974 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsTaskbar.java
@@ -24,6 +24,7 @@
 import androidx.test.filters.LargeTest;
 
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
+import com.android.launcher3.util.rule.ScreenRecordRule;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -133,6 +134,7 @@
 
     @Test
     @PortraitLandscape
+    @ScreenRecordRule.ScreenRecord // b/349439239
     public void testLaunchAppInSplitscreen_fromTaskbarAllApps() {
         getTaskbar().openAllApps()
                 .getAppIcon(TEST_APP_NAME)
diff --git a/res/drawable/cloud_download_24px.xml b/res/drawable/cloud_download_24px.xml
new file mode 100644
index 0000000..6f7c95a
--- /dev/null
+++ b/res/drawable/cloud_download_24px.xml
@@ -0,0 +1,11 @@
+<!-- GM3 icon cloud_download:vd_theme_24 -->
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:viewportWidth="960"
+    android:viewportHeight="960"
+    android:tint="?attr/colorControlNormal">
+  <path
+      android:fillColor="@android:color/white"
+      android:pathData="M260,800Q169,800 104.5,737Q40,674 40,583Q40,505 87,444Q134,383 210,366Q227,294 295,229Q363,164 440,164Q473,164 496.5,187.5Q520,211 520,244L520,486L584,424L640,480L480,640L320,480L376,424L440,486L440,244Q364,258 322,317.5Q280,377 280,440L260,440Q202,440 161,481Q120,522 120,580Q120,638 161,679Q202,720 260,720L740,720Q782,720 811,691Q840,662 840,620Q840,578 811,549Q782,520 740,520L680,520L680,440Q680,392 658,350.5Q636,309 600,280L600,187Q674,222 717,290.5Q760,359 760,440L760,440L760,440Q829,448 874.5,499.5Q920,551 920,620Q920,695 867.5,747.5Q815,800 740,800L260,800ZM480,442Q480,442 480,442Q480,442 480,442L480,442Q480,442 480,442Q480,442 480,442L480,442Q480,442 480,442Q480,442 480,442L480,442Q480,442 480,442Q480,442 480,442Q480,442 480,442Q480,442 480,442L480,442Q480,442 480,442Q480,442 480,442Q480,442 480,442Q480,442 480,442L480,442L480,442Q480,442 480,442Q480,442 480,442Z"/>
+</vector>
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/BubbleTextView.java b/src/com/android/launcher3/BubbleTextView.java
index 83427a0..0cb2137 100644
--- a/src/com/android/launcher3/BubbleTextView.java
+++ b/src/com/android/launcher3/BubbleTextView.java
@@ -40,11 +40,15 @@
 import android.graphics.drawable.ColorDrawable;
 import android.graphics.drawable.Drawable;
 import android.icu.text.MessageFormat;
+import android.text.Spannable;
+import android.text.SpannableString;
 import android.text.StaticLayout;
 import android.text.TextPaint;
 import android.text.TextUtils;
 import android.text.TextUtils.TruncateAt;
+import android.text.style.ImageSpan;
 import android.util.AttributeSet;
+import android.util.Log;
 import android.util.Property;
 import android.util.Size;
 import android.util.TypedValue;
@@ -54,6 +58,7 @@
 import android.view.ViewDebug;
 import android.widget.TextView;
 
+import androidx.annotation.DrawableRes;
 import androidx.annotation.Nullable;
 import androidx.annotation.UiThread;
 import androidx.annotation.VisibleForTesting;
@@ -96,6 +101,8 @@
 public class BubbleTextView extends TextView implements ItemInfoUpdateReceiver,
         IconLabelDotView, DraggableView, Reorderable {
 
+    public static final String TAG = "BubbleTextView";
+
     public static final int DISPLAY_WORKSPACE = 0;
     public static final int DISPLAY_ALL_APPS = 1;
     public static final int DISPLAY_FOLDER = 2;
@@ -494,7 +501,13 @@
             mLastOriginalText = label;
             mLastModifiedText = mLastOriginalText;
             mBreakPointsIntArray = StringMatcherUtility.getListOfBreakpoints(label, MATCHER);
-            setText(label);
+            if (Flags.enableNewArchivingIcon()
+                    && info instanceof ItemInfoWithIcon infoWithIcon
+                    && infoWithIcon.isInactiveArchive()) {
+                setTextWithStartIcon(label, R.drawable.cloud_download_24px);
+            } else {
+                setText(label);
+            }
         }
         if (info.contentDescription != null) {
             setContentDescription(info.isDisabled()
@@ -804,7 +817,13 @@
                     getLineSpacingExtra());
             if (!TextUtils.equals(modifiedString, mLastModifiedText)) {
                 mLastModifiedText = modifiedString;
-                setText(modifiedString);
+                if (Flags.enableNewArchivingIcon()
+                        && getTag() instanceof ItemInfoWithIcon infoWithIcon
+                        && infoWithIcon.isInactiveArchive()) {
+                    setTextWithStartIcon(modifiedString, R.drawable.cloud_download_24px);
+                } else {
+                    setText(modifiedString);
+                }
                 // if text contains NEW_LINE, set max lines to 2
                 if (TextUtils.indexOf(modifiedString, NEW_LINE) != -1) {
                     setSingleLine(false);
@@ -825,6 +844,28 @@
         super.setTextColor(getModifiedColor());
     }
 
+    /**
+     * Uses a SpannableString to set text with a Drawable at the start of the TextView
+     * @param text text to use for TextView
+     * @param drawableRes Drawable Resource to use for drawing image at start of text
+     */
+    private void setTextWithStartIcon(CharSequence text, @DrawableRes int drawableRes) {
+        Drawable drawable = getContext().getDrawable(drawableRes);
+        if (drawable == null) {
+            setText(text);
+            Log.w(TAG, "setTextWithStartIcon: start icon Drawable not found from resources"
+                    + ", will just set text instead. text=" + text);
+            return;
+        }
+        drawable.setTint(getCurrentTextColor());
+        drawable.setBounds(0, 0, Math.round(getTextSize()), Math.round(getTextSize()));
+        ImageSpan imageSpan = new ImageSpan(drawable, ImageSpan.ALIGN_CENTER);
+        // First space will be replaced with Drawable, second space is for space before text.
+        SpannableString spannable = new SpannableString("  " + text);
+        spannable.setSpan(imageSpan, 0, 1, Spannable.SPAN_INCLUSIVE_EXCLUSIVE);
+        setText(spannable);
+    }
+
     @Override
     public void setTextColor(ColorStateList colors) {
         mTextColor = colors.getDefaultColor();
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index 85c8b57..b41da0f 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -115,6 +115,9 @@
         if (BuildCompat.isAtLeastV() && Flags.enableSupportForArchiving()) {
             ArchiveCompatibilityParams params = new ArchiveCompatibilityParams();
             params.setEnableUnarchivalConfirmation(false);
+            if (Flags.enableNewArchivingIcon()) {
+                params.setEnableIconOverlay(false);
+            }
             launcherApps.setArchiveCompatibility(params);
         }
 
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..aefb7e9 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,11 @@
     /** Add Private Space Header view elements based upon {@link UserProfileState} */
     public void bindPrivateSpaceHeaderViewElements(RelativeLayout parent) {
         mPSHeader = parent;
+        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. */
@@ -378,9 +379,10 @@
         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 +393,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 +414,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 +595,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 +651,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 +671,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 +689,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 +706,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 +749,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();
@@ -720,36 +767,10 @@
         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 +781,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 +806,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;
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/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/WidgetManagerHelper.java b/src/com/android/launcher3/widget/WidgetManagerHelper.java
index 9132b4f..23d0585 100644
--- a/src/com/android/launcher3/widget/WidgetManagerHelper.java
+++ b/src/com/android/launcher3/widget/WidgetManagerHelper.java
@@ -32,6 +32,7 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.annotation.RequiresApi;
+import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
@@ -57,8 +58,13 @@
     final Context mContext;
 
     public WidgetManagerHelper(Context context) {
+        this(context, AppWidgetManager.getInstance(context));
+    }
+
+    @VisibleForTesting
+    public WidgetManagerHelper(Context context, AppWidgetManager appWidgetManager) {
         mContext = context;
-        mAppWidgetManager = AppWidgetManager.getInstance(context);
+        mAppWidgetManager = appWidgetManager;
     }
 
     /**
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/FakeInvariantDeviceProfileTest.kt b/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
index 0538870..954dc8f 100644
--- a/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/FakeInvariantDeviceProfileTest.kt
@@ -47,14 +47,13 @@
 abstract class FakeInvariantDeviceProfileTest {
 
     protected lateinit var context: Context
-    protected var inv: InvariantDeviceProfile? = null
-    protected val info: Info = mock()
-    protected var windowBounds: WindowBounds? = null
-    protected var isMultiWindowMode: Boolean = false
-    protected var transposeLayoutWithOrientation: Boolean = false
-    protected var useTwoPanels: Boolean = false
-    protected var isGestureMode: Boolean = true
-    protected var isTransientTaskbar: Boolean = true
+    protected lateinit var inv: InvariantDeviceProfile
+    protected val info = mock<Info>()
+    protected lateinit var windowBounds: WindowBounds
+    private var transposeLayoutWithOrientation = false
+    private var useTwoPanels = false
+    private var isGestureMode = true
+    private var isTransientTaskbar = true
 
     @Rule @JvmField val limitDevicesRule = LimitDevicesRule()
 
@@ -73,7 +72,7 @@
             info,
             windowBounds,
             SparseArray(),
-            isMultiWindowMode,
+            /*isMultiWindowMode=*/ false,
             transposeLayoutWithOrientation,
             useTwoPanels,
             isGestureMode,
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/RunnableListTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/RunnableListTest.kt
new file mode 100644
index 0000000..093afc9
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/util/RunnableListTest.kt
@@ -0,0 +1,113 @@
+/*
+ * 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 androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.Before
+import org.junit.Test
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.reset
+import org.mockito.kotlin.verify
+import org.mockito.kotlin.verifyZeroInteractions
+
+@SmallTest
+class RunnableListTest {
+
+    @Mock private lateinit var runnable1: Runnable
+    @Mock private lateinit var runnable2: Runnable
+
+    private val underTest = RunnableList()
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+    }
+
+    @Test
+    fun not_destroyedByDefault() {
+        assertThat(underTest.isDestroyed).isFalse()
+    }
+
+    @Test
+    fun add_and_run() {
+        underTest.add(runnable1)
+        underTest.add(runnable2)
+
+        underTest.executeAllAndDestroy()
+
+        verify(runnable1).run()
+        verify(runnable2).run()
+        assertThat(underTest.isDestroyed).isTrue()
+    }
+
+    @Test
+    fun add_to_destroyed_runnableList_run_immediately() {
+        underTest.executeAllAndDestroy()
+
+        underTest.add(runnable1)
+
+        verify(runnable1).run()
+    }
+
+    @Test
+    fun second_executeAllAndDestroy_noOp() {
+        underTest.executeAllAndDestroy()
+        underTest.add(runnable1)
+        reset(runnable1)
+
+        underTest.executeAllAndDestroy()
+
+        verifyZeroInteractions(runnable1)
+    }
+
+    @Test
+    fun executeAllAndClear_run_not_destroy() {
+        underTest.add(runnable1)
+        underTest.add(runnable2)
+
+        underTest.executeAllAndClear()
+
+        verify(runnable1).run()
+        verify(runnable2).run()
+        assertThat(underTest.isDestroyed).isFalse()
+    }
+
+    @Test
+    fun executeAllAndClear_not_destroy() {
+        underTest.executeAllAndClear()
+        underTest.add(runnable1)
+        reset(runnable1)
+
+        underTest.executeAllAndClear()
+
+        verify(runnable1).run()
+    }
+
+    @Test
+    fun remove_and_run_not_executed() {
+        underTest.add(runnable1)
+        underTest.add(runnable2)
+
+        underTest.remove(runnable1)
+        underTest.executeAllAndClear()
+
+        verifyZeroInteractions(runnable1)
+        verify(runnable2).run()
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/ViewCacheTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/ViewCacheTest.kt
similarity index 100%
rename from tests/src/com/android/launcher3/util/ViewCacheTest.kt
rename to tests/multivalentTests/src/com/android/launcher3/util/ViewCacheTest.kt
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/ListenableHostViewTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/ListenableHostViewTest.kt
new file mode 100644
index 0000000..6c71f36
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/ListenableHostViewTest.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.launcher3.widget
+
+import android.content.Context
+import android.view.accessibility.AccessibilityNodeInfo
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.util.ActivityContextWrapper
+import com.google.common.truth.Truth
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class ListenableHostViewTest {
+
+    private val context: Context
+        get() = ActivityContextWrapper(InstrumentationRegistry.getInstrumentation().targetContext)
+
+    @Test
+    fun updateAppWidget_notifiesListeners() {
+        val hostView = ListenableHostView(context)
+        var wasNotifiedOfUpdate = false
+        val updateListener = Runnable { wasNotifiedOfUpdate = true }
+        hostView.addUpdateListener(updateListener)
+        hostView.beginDeferringUpdates()
+        hostView.updateAppWidget(null)
+        Truth.assertThat(wasNotifiedOfUpdate).isTrue()
+    }
+
+    @Test
+    fun onInitializeAccessibilityNodeInfo_correctlySetsClassName() {
+        val hostView = ListenableHostView(context)
+        val nodeInfo = AccessibilityNodeInfo()
+        hostView.onInitializeAccessibilityNodeInfo(nodeInfo)
+        Truth.assertThat(nodeInfo.className).isEqualTo(LauncherAppWidgetHostView::class.java.name)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/WidgetAddFlowHandlerTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/WidgetAddFlowHandlerTest.kt
new file mode 100644
index 0000000..242ec7c
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/WidgetAddFlowHandlerTest.kt
@@ -0,0 +1,107 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget
+
+import android.content.Context
+import android.os.Bundle
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.Launcher
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.LauncherAppWidgetInfo
+import com.android.launcher3.util.ActivityContextWrapper
+import com.google.common.truth.Truth
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.any
+import org.mockito.Mockito.verify
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class WidgetAddFlowHandlerTest {
+
+    private val context: Context
+        get() = ActivityContextWrapper(InstrumentationRegistry.getInstrumentation().targetContext)
+
+    private val providerInfo =
+        LauncherAppWidgetProviderInfo().apply {
+            configure = InstrumentationRegistry.getInstrumentation().componentName
+        }
+    private val appWidgetHolder: LauncherWidgetHolder = mock<LauncherWidgetHolder>()
+    private val launcher: Launcher =
+        mock<Launcher>().also { whenever(it.appWidgetHolder).thenReturn(appWidgetHolder) }
+    private val appWidgetInfo = LauncherAppWidgetInfo().apply { appWidgetId = 123 }
+    private val requestCode = 123
+    private val flowHandler = WidgetAddFlowHandler(providerInfo)
+
+    @Test
+    fun valuesShouldRemainTheSame_beforeAndAfter_parcelization() {
+        with(Bundle()) {
+            val testKey = "testKey"
+            putParcelable(testKey, flowHandler)
+            Truth.assertThat(getParcelable(testKey, WidgetAddFlowHandler::class.java))
+                .isEqualTo(flowHandler)
+        }
+    }
+
+    @Test
+    fun describeContents_shouldReturn_0() {
+        Truth.assertThat(flowHandler.describeContents()).isEqualTo(0)
+    }
+
+    @Test
+    fun startBindFlow_shouldCorrectly_startLauncherFlowBinding() {
+        flowHandler.startBindFlow(launcher, appWidgetInfo.appWidgetId, appWidgetInfo, requestCode)
+        verify(launcher).setWaitingForResult(any())
+        verify(appWidgetHolder)
+            .startBindFlow(launcher, appWidgetInfo.appWidgetId, providerInfo, requestCode)
+    }
+
+    @Test
+    fun startConfigActivityWithCustomAppWidgetId_shouldAskLauncherToStartConfigActivity() {
+        flowHandler.startConfigActivity(
+            launcher,
+            appWidgetInfo.appWidgetId,
+            ItemInfo(),
+            requestCode
+        )
+        verify(launcher).setWaitingForResult(any())
+        verify(appWidgetHolder)
+            .startConfigActivity(launcher, appWidgetInfo.appWidgetId, requestCode)
+    }
+
+    @Test
+    fun startConfigActivity_shouldAskLauncherToStartConfigActivity() {
+        flowHandler.startConfigActivity(launcher, appWidgetInfo, requestCode)
+        verify(launcher).setWaitingForResult(any())
+        verify(appWidgetHolder)
+            .startConfigActivity(launcher, appWidgetInfo.appWidgetId, requestCode)
+    }
+
+    @Test
+    fun needsConfigure_returnsTrue_ifFlagsAndProviderInfoDetermineSo() {
+        Truth.assertThat(flowHandler.needsConfigure()).isTrue()
+    }
+
+    @Test
+    fun getProviderInfo_returnCorrectProviderInfo() {
+        Truth.assertThat(flowHandler.getProviderInfo(context)).isSameInstanceAs(providerInfo)
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/widget/WidgetManagerHelperTest.kt b/tests/multivalentTests/src/com/android/launcher3/widget/WidgetManagerHelperTest.kt
new file mode 100644
index 0000000..f1cfb79
--- /dev/null
+++ b/tests/multivalentTests/src/com/android/launcher3/widget/WidgetManagerHelperTest.kt
@@ -0,0 +1,128 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget
+
+import android.appwidget.AppWidgetManager
+import android.content.Context
+import android.content.pm.ActivityInfo
+import android.os.Bundle
+import android.os.Process
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.PackageUserKey
+import com.google.common.truth.Truth
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.whenever
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class WidgetManagerHelperTest {
+
+    private val context: Context
+        get() = ActivityContextWrapper(InstrumentationRegistry.getInstrumentation().targetContext)
+
+    private val info =
+        LauncherAppWidgetProviderInfo().apply {
+            provider = InstrumentationRegistry.getInstrumentation().componentName
+            providerInfo =
+                mock(ActivityInfo::class.java).apply { applicationInfo = context.applicationInfo }
+        }
+
+    @Mock private lateinit var appWidgetManager: AppWidgetManager
+
+    private lateinit var underTest: WidgetManagerHelper
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+        underTest = WidgetManagerHelper(context, appWidgetManager)
+    }
+
+    @Test
+    fun getAllProviders_returnsCorrectWidgetProviderInfo() {
+        val packageUserKey =
+            mock(PackageUserKey::class.java).apply {
+                mPackageName = context.packageName
+                mUser = Process.myUserHandle()
+            }
+        val desiredResult = listOf(info)
+        whenever(
+                appWidgetManager.getInstalledProvidersForPackage(
+                    packageUserKey.mPackageName,
+                    packageUserKey.mUser
+                )
+            )
+            .thenReturn(desiredResult)
+        Truth.assertThat(underTest.getAllProviders(packageUserKey)).isSameInstanceAs(desiredResult)
+    }
+
+    @Test
+    fun getLauncherAppWidgetInfo_returnsCorrectInfo_ifWidgetExists() {
+        val id = 123
+        whenever(appWidgetManager.getAppWidgetInfo(id)).thenReturn(info)
+        val componentName = InstrumentationRegistry.getInstrumentation().componentName
+        Truth.assertThat(underTest.getLauncherAppWidgetInfo(id, componentName))
+            .isSameInstanceAs(info)
+    }
+
+    @Test
+    fun bindAppWidgetIdIfAllowed_correctly_forwardsBindCommandToAppWidgetManager() {
+        val id = 124
+        val options = Bundle()
+        underTest.bindAppWidgetIdIfAllowed(id, info, options)
+        verify(appWidgetManager).bindAppWidgetIdIfAllowed(id, info.profile, info.provider, options)
+    }
+
+    @Test
+    fun findProvider_returnsNull_ifNoProviderExists() {
+        val info =
+            underTest.getLauncherAppWidgetInfo(
+                1,
+                InstrumentationRegistry.getInstrumentation().componentName
+            )
+        Truth.assertThat(info).isNull()
+    }
+
+    @Test
+    fun isAppWidgetRestored_returnsTrue_ifWidgetIsRestored() {
+        val id = 126
+        whenever(appWidgetManager.getAppWidgetOptions(id))
+            .thenReturn(
+                Bundle().apply {
+                    putBoolean(WidgetManagerHelper.WIDGET_OPTION_RESTORE_COMPLETED, true)
+                }
+            )
+        Truth.assertThat(underTest.isAppWidgetRestored(id)).isTrue()
+    }
+
+    @Test
+    fun loadGeneratedPreview_returnsWidgetPreview_fromAppWidgetManager() {
+        val widgetCategory = 130
+        with(info) {
+            underTest.loadGeneratedPreview(this, widgetCategory)
+            verify(appWidgetManager).getWidgetPreview(provider, profile, widgetCategory)
+        }
+    }
+}
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/dragging/TaplDragTest.java b/tests/src/com/android/launcher3/dragging/TaplDragTest.java
index 41abcf8..76c1948 100644
--- a/tests/src/com/android/launcher3/dragging/TaplDragTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplDragTest.java
@@ -65,6 +65,7 @@
     @Test
     @PortraitLandscape
     @PlatinumTest(focusArea = "launcher")
+    @ScreenRecordRule.ScreenRecord // b/353600888
     public void testDragToFolder() {
         // TODO: add the use case to drag an icon to an existing folder. Currently it either fails
         // on tablets or phones due to difference in resolution.
@@ -97,6 +98,7 @@
      * icon left.
      */
     @Test
+    @ScreenRecordRule.ScreenRecord // b/353600888
     public void testDragOutOfFolder() {
         final HomeAppIcon playStoreIcon = createShortcutIfNotExist(STORE_APP_NAME, 0, 1);
         final HomeAppIcon photosIcon = createShortcutInCenterIfNotExist(PHOTOS_APP_NAME);
@@ -174,13 +176,13 @@
     public void testDragAndCancelAppIcon() {
         final HomeAppIcon homeAppIcon = createShortcutInCenterIfNotExist(GMAIL_APP_NAME);
         Point positionBeforeDrag =
-                mLauncher.getWorkspace().getWorkspaceIconsPositions().get(GMAIL_APP_NAME);
+                mLauncher.getWorkspace().getWorkspaceIconPosition(GMAIL_APP_NAME);
         assertNotNull("App not found in Workspace before dragging.", positionBeforeDrag);
 
         mLauncher.getWorkspace().dragAndCancelAppIcon(homeAppIcon);
 
         Point positionAfterDrag =
-                mLauncher.getWorkspace().getWorkspaceIconsPositions().get(GMAIL_APP_NAME);
+                mLauncher.getWorkspace().getWorkspaceIconPosition(GMAIL_APP_NAME);
         assertNotNull("App not found in Workspace after dragging.", positionAfterDrag);
         assertEquals("App not returned to same position in Workspace after drag & cancel",
                 positionBeforeDrag, positionAfterDrag);
diff --git a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
index 46cafa7..907aa50 100644
--- a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
@@ -46,7 +46,6 @@
 
 import java.io.IOException;
 import java.util.Arrays;
-import java.util.Map;
 
 /**
  * Test runs in Out of process (Oop) and In process (Ipc)
@@ -60,6 +59,7 @@
      */
     @Test
     @PortraitLandscape
+    @ScreenRecordRule.ScreenRecord // b/349439239
     public void testDeleteFromWorkspace() {
         for (String appName : new String[]{GMAIL_APP_NAME, STORE_APP_NAME, TEST_APP_NAME}) {
             final HomeAppIcon homeAppIcon = createShortcutInCenterIfNotExist(appName);
@@ -155,18 +155,14 @@
                 createShortcutIfNotExist(appNames[i], gridPositions[i]);
             }
 
-            Map<String, Point> initialPositions =
-                    mLauncher.getWorkspace().getWorkspaceIconsPositions();
-            assertThat(initialPositions.keySet()).containsAtLeastElementsIn(appNames);
+            Point initialPosition =
+                    mLauncher.getWorkspace().getWorkspaceIconPosition(DUMMY_APP_NAME);
+            assertThat(initialPosition).isNotNull();
 
             final Workspace workspace = mLauncher.getWorkspace().getWorkspaceAppIcon(
                     DUMMY_APP_NAME).uninstall();
             workspace.verifyWorkspaceAppIconIsGone(
                     DUMMY_APP_NAME + " was expected to disappear after uninstall.", DUMMY_APP_NAME);
-
-            Log.d(UIOBJECT_STALE_ELEMENT, "second getWorkspaceIconsPositions()");
-            Map<String, Point> finalPositions = workspace.getWorkspaceIconsPositions();
-            assertThat(finalPositions).doesNotContainKey(DUMMY_APP_NAME);
         } finally {
             TestUtil.uninstallDummyApp();
         }
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/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java b/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
index eb05000..20c5a25 100644
--- a/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
+++ b/tests/src/com/android/launcher3/ui/workspace/TaplTwoPanelWorkspaceTest.java
@@ -20,8 +20,6 @@
 import static com.android.launcher3.util.TestConstants.AppNames.MAPS_APP_NAME;
 import static com.android.launcher3.util.TestConstants.AppNames.MESSAGES_APP_NAME;
 import static com.android.launcher3.util.TestConstants.AppNames.STORE_APP_NAME;
-import static com.android.launcher3.util.rule.TestStabilityRule.LOCAL;
-import static com.android.launcher3.util.rule.TestStabilityRule.PLATFORM_POSTSUBMIT;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
@@ -40,7 +38,6 @@
 import com.android.launcher3.util.LauncherLayoutBuilder;
 import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ScreenRecordRule;
-import com.android.launcher3.util.rule.TestStabilityRule;
 
 import org.junit.After;
 import org.junit.Before;
@@ -288,6 +285,7 @@
 
     @Test
     @PortraitLandscape
+    @ScreenRecordRule.ScreenRecord // b/330232490
     public void testEmptyPagesGetRemovedIfBothPagesAreEmpty() {
         Workspace workspace = mLauncher.getWorkspace();
 
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)) {
diff --git a/tests/tapl/com/android/launcher3/tapl/Workspace.java b/tests/tapl/com/android/launcher3/tapl/Workspace.java
index 9ac6768..748d576 100644
--- a/tests/tapl/com/android/launcher3/tapl/Workspace.java
+++ b/tests/tapl/com/android/launcher3/tapl/Workspace.java
@@ -345,17 +345,34 @@
      * @return map of text -> center of the view. In case of icons with the same name, the one with
      * lower x coordinate is selected.
      */
-    public Map<String, Point> getWorkspaceIconsPositions() {
+    public Map<String, Point> getAllWorkspaceIconsPositions() {
         final UiObject2 workspace = verifyActiveContainer();
-        mLauncher.waitForLauncherInitialized(); // b/319501259
         List<UiObject2> workspaceIcons =
                 mLauncher.waitForObjectsInContainer(workspace, AppIcon.getAnyAppIconSelector());
-        return workspaceIcons.stream()
+        return getIconPositionMap(workspaceIcons);
+    }
+
+    /**
+     * @return point where icon is found for given the app name,
+     * point is visible center of the icon.
+     */
+    @NonNull
+    public Point getWorkspaceIconPosition(String appName) {
+        final UiObject2 workspace = verifyActiveContainer();
+
+        UiObject2 workspaceIcon =
+                mLauncher.waitForObjectInContainer(workspace,
+                        AppIcon.getAppIconSelector(appName, mLauncher));
+        return workspaceIcon.getVisibleCenter();
+    }
+
+    private Map<String, Point> getIconPositionMap(List<UiObject2> icons) {
+        return icons.stream()
                 .collect(
                         Collectors.toMap(
                                 /* keyMapper= */ uiObject21 -> {
-                                    Log.d(UIOBJECT_STALE_ELEMENT, "keyText: " +
-                                            uiObject21.getText());
+                                    Log.d(UIOBJECT_STALE_ELEMENT, "keyText: "
+                                            + uiObject21.getText());
                                     return uiObject21.getText();
                                 },
                                 /* valueMapper= */ uiObject2 -> {