Merge "2b/ Update launcher to use GroupedTaskInfos" into main
diff --git a/quickstep/res/layout/task_desktop.xml b/quickstep/res/layout/task_desktop.xml
index 1564653..0472007 100644
--- a/quickstep/res/layout/task_desktop.xml
+++ b/quickstep/res/layout/task_desktop.xml
@@ -19,16 +19,11 @@
     android:id="@+id/task_view_desktop"
     android:layout_width="match_parent"
     android:layout_height="match_parent"
-    android:clipChildren="true"
-    android:clipToPadding="true"
     android:contentDescription="@string/recent_task_desktop"
     android:defaultFocusHighlightEnabled="false"
     android:focusable="true"
-    android:padding="0.1dp"
     launcher:focusBorderColor="?attr/materialColorOutline"
     launcher:hoverBorderColor="?attr/materialColorPrimary">
-    <!-- Setting a padding of 0.1 dp since android:clipToPadding needs a non-zero value for
-    padding to work-->
     <View
         android:id="@+id/background"
         android:layout_width="match_parent"
@@ -40,4 +35,9 @@
         android:layout_height="wrap_content"
         android:inflatedId="@id/icon" />
 
+    <com.android.quickstep.views.DesktopTaskContentView
+        android:id="@+id/desktop_content"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent" />
+
 </com.android.quickstep.views.DesktopTaskView>
diff --git a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
index 955388d..bd2c7cc 100644
--- a/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
+++ b/quickstep/src/com/android/launcher3/WidgetPickerActivity.java
@@ -23,6 +23,8 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
+import static java.util.Collections.emptyList;
+
 import android.appwidget.AppWidgetManager;
 import android.appwidget.AppWidgetProviderInfo;
 import android.content.ClipData;
@@ -44,6 +46,7 @@
 import com.android.launcher3.model.StringCache;
 import com.android.launcher3.model.WidgetItem;
 import com.android.launcher3.model.WidgetPredictionsRequester;
+import com.android.launcher3.model.WidgetsFilterDataProvider;
 import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
@@ -112,6 +115,7 @@
     private WidgetPredictionsRequester mWidgetPredictionsRequester;
     private final WidgetPickerDataProvider mWidgetPickerDataProvider =
             new WidgetPickerDataProvider();
+    private WidgetsFilterDataProvider mWidgetsFilterDataProvider;
 
     private int mDesiredWidgetWidth;
     private int mDesiredWidgetHeight;
@@ -133,13 +137,13 @@
     @Nullable
     private WidgetsFullSheet mWidgetSheet;
 
-    private final Predicate<WidgetItem> mWidgetsFilter = widget -> {
+    private final Predicate<WidgetItem> mNoShortcutsFilter = widget -> {
         final WidgetAcceptabilityVerdict verdict =
                 isWidgetAcceptable(widget, /* applySizeFilter=*/ false);
         verdict.maybeLogVerdict();
         return verdict.isAcceptable;
     };
-    private final Predicate<WidgetItem> mDefaultWidgetsFilter = widget -> {
+    private final Predicate<WidgetItem> mHostSizeAndNoShortcutsFilter = widget -> {
         final WidgetAcceptabilityVerdict verdict =
                 isWidgetAcceptable(widget, /* applySizeFilter=*/ true);
         verdict.maybeLogVerdict();
@@ -157,6 +161,7 @@
         InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
         mDeviceProfile = idp.getDeviceProfile(this);
         mModel = new WidgetsModel();
+        mWidgetsFilterDataProvider = WidgetsFilterDataProvider.Companion.newInstance(this);
 
         setContentView(R.layout.widget_picker_activity);
         mDragLayer = findViewById(R.id.drag_layer);
@@ -288,13 +293,16 @@
     private void refreshAndBindWidgets() {
         MODEL_EXECUTOR.execute(() -> {
             LauncherAppState app = LauncherAppState.getInstance(this);
+            // Don't have to setup filters - its setup when launcher loads
+            // Just refresh filters with available cached info.
+            mModel.updateWidgetFilters(mWidgetsFilterDataProvider);
             mModel.update(app, null);
 
             StringCache stringCache = new StringCache();
             stringCache.loadStrings(this);
 
             bindStringCache(stringCache);
-            bindWidgets(mModel.getWidgetsByPackageItem());
+            bindWidgets(mModel.getWidgetsByPackageItem(), mModel.getDefaultWidgetsFilter());
             // Open sheet once widgets are available, so that it doesn't interrupt the open
             // animation.
             openWidgetsSheet();
@@ -310,14 +318,23 @@
         MAIN_EXECUTOR.execute(() -> mStringCache = stringCache);
     }
 
-    private void bindWidgets(Map<PackageItemInfo, List<WidgetItem>> widgets) {
+    private void bindWidgets(Map<PackageItemInfo, List<WidgetItem>> widgets,
+            @Nullable Predicate<WidgetItem> defaultWidgetsFilter) {
         WidgetsListBaseEntriesBuilder builder = new WidgetsListBaseEntriesBuilder(
                 mApp.getContext());
 
-        final List<WidgetsListBaseEntry> allWidgets = builder.build(widgets, mWidgetsFilter);
-        final List<WidgetsListBaseEntry> defaultWidgets =
-                shouldShowDefaultWidgets() ? builder.build(widgets,
-                        mDefaultWidgetsFilter) : List.of();
+        final List<WidgetsListBaseEntry> allWidgets = builder.build(widgets, mNoShortcutsFilter);
+
+        // Default list is shown if either defaultWidgetsFilter exists or host has additionally
+        // enforced size filtering.
+        @Nullable Predicate<WidgetItem> defaultListFilter =
+                hasHostSizeFilters() ? mHostSizeAndNoShortcutsFilter : null;
+        if (defaultWidgetsFilter != null) {
+            defaultListFilter = defaultListFilter != null ? defaultListFilter.and(
+                    defaultWidgetsFilter) : defaultWidgetsFilter;
+        }
+        final List<WidgetsListBaseEntry> defaultWidgets = defaultListFilter != null ? builder.build(
+                widgets, defaultListFilter) : emptyList();
 
         MAIN_EXECUTOR.execute(
                 () -> mWidgetPickerDataProvider.setWidgets(allWidgets, defaultWidgets));
@@ -342,6 +359,7 @@
     @Override
     protected void onDestroy() {
         super.onDestroy();
+        MODEL_EXECUTOR.execute(() -> mWidgetsFilterDataProvider.destroy());
         if (mWidgetPredictionsRequester != null) {
             mWidgetPredictionsRequester.clear();
         }
@@ -398,7 +416,7 @@
         }
     }
 
-    private boolean shouldShowDefaultWidgets() {
+    private boolean hasHostSizeFilters() {
         // If optional filters such as size filter are present, we display them as default widgets.
         return mDesiredWidgetWidth != 0 || mDesiredWidgetHeight != 0;
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java
index 929e793..6a908ca 100644
--- a/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/FallbackTaskbarUIController.java
@@ -25,20 +25,24 @@
 
 import com.android.launcher3.popup.SystemShortcut;
 import com.android.launcher3.statemanager.StateManager;
-import com.android.quickstep.RecentsActivity;
+import com.android.launcher3.statemanager.StatefulContainer;
 import com.android.quickstep.TopTaskTracker;
 import com.android.quickstep.fallback.RecentsState;
 import com.android.quickstep.util.TISBindHelper;
 import com.android.quickstep.views.RecentsView;
+import com.android.quickstep.views.RecentsViewContainer;
 
 import java.util.stream.Stream;
 
 /**
  * A data source which integrates with the fallback RecentsActivity instance (for 3P launchers).
+ * @param <T> The type of the RecentsViewContainer that will handle Recents state changes.
  */
-public class FallbackTaskbarUIController extends TaskbarUIController {
+public class FallbackTaskbarUIController
+        <T extends RecentsViewContainer & StatefulContainer<RecentsState>>
+        extends TaskbarUIController {
 
-    private final RecentsActivity mRecentsActivity;
+    private final T mRecentsContainer;
 
     private final StateManager.StateListener<RecentsState> mStateListener =
             new StateManager.StateListener<RecentsState>() {
@@ -46,8 +50,12 @@
                 public void onStateTransitionStart(RecentsState toState) {
                     animateToRecentsState(toState);
 
+                    RecentsView recentsView = getRecentsView();
+                    if (recentsView == null) {
+                        return;
+                    }
                     // Handle tapping on live tile.
-                    getRecentsView().setTaskLaunchListener(toState == RecentsState.DEFAULT
+                    recentsView.setTaskLaunchListener(toState == RecentsState.DEFAULT
                             ? (() -> animateToRecentsState(RecentsState.BACKGROUND_APP)) : null);
                 }
 
@@ -63,23 +71,26 @@
                 }
             };
 
-    public FallbackTaskbarUIController(RecentsActivity recentsActivity) {
-        mRecentsActivity = recentsActivity;
+    public FallbackTaskbarUIController(T recentsContainer) {
+        mRecentsContainer = recentsContainer;
     }
 
     @Override
     protected void init(TaskbarControllers taskbarControllers) {
         super.init(taskbarControllers);
-        mRecentsActivity.setTaskbarUIController(this);
-        mRecentsActivity.getStateManager().addStateListener(mStateListener);
+        mRecentsContainer.setTaskbarUIController(this);
+        mRecentsContainer.getStateManager().addStateListener(mStateListener);
     }
 
     @Override
     protected void onDestroy() {
         super.onDestroy();
-        getRecentsView().setTaskLaunchListener(null);
-        mRecentsActivity.setTaskbarUIController(null);
-        mRecentsActivity.getStateManager().removeStateListener(mStateListener);
+        RecentsView recentsView = getRecentsView();
+        if (recentsView != null) {
+            recentsView.setTaskLaunchListener(null);
+        }
+        mRecentsContainer.setTaskbarUIController(null);
+        mRecentsContainer.getStateManager().removeStateListener(mStateListener);
     }
 
     /**
@@ -108,8 +119,8 @@
     }
 
     @Override
-    public RecentsView getRecentsView() {
-        return mRecentsActivity.getOverviewPanel();
+    public @Nullable RecentsView getRecentsView() {
+        return mRecentsContainer.getOverviewPanel();
     }
 
     @Override
@@ -131,11 +142,11 @@
     @Nullable
     @Override
     protected TISBindHelper getTISBindHelper() {
-        return mRecentsActivity.getTISBindHelper();
+        return mRecentsContainer.getTISBindHelper();
     }
 
     @Override
     protected String getTaskbarUIControllerName() {
-        return "FallbackTaskbarUIController";
+        return "FallbackTaskbarUIController<" + mRecentsContainer.getClass().getSimpleName() + ">";
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
index 4a85acc..5a63ca6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarControllers.java
@@ -169,7 +169,7 @@
         taskbarOverlayController.init(this);
         taskbarAllAppsController.init(this, sharedState.allAppsVisible);
         navButtonController.init(this);
-        bubbleControllers.ifPresent(controllers -> controllers.init(this));
+        bubbleControllers.ifPresent(controllers -> controllers.init(sharedState, this));
         taskbarInsetsController.init(this);
         voiceInteractionWindowController.init(this);
         taskbarRecentAppsController.init(this);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index a0440c1..ab4b1b6 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -71,7 +71,9 @@
 import com.android.quickstep.AllAppsActionManager;
 import com.android.quickstep.RecentsActivity;
 import com.android.quickstep.SystemUiProxy;
+import com.android.quickstep.fallback.window.RecentsWindowManager;
 import com.android.quickstep.util.ContextualSearchInvoker;
+import com.android.quickstep.views.RecentsViewContainer;
 import com.android.systemui.shared.statusbar.phone.BarTransitions;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
@@ -130,6 +132,8 @@
 
     private TaskbarActivityContext mTaskbarActivityContext;
     private StatefulActivity mActivity;
+    private RecentsViewContainer mRecentsViewContainer;
+
     /**
      * Cache a copy here so we can initialize state whenever taskbar is recreated, since
      * this class does not get re-initialized w/ new taskbars.
@@ -403,9 +407,28 @@
         }
         mUnfoldProgressProvider.setSourceProvider(unfoldTransitionProgressProvider);
 
+        if (activity instanceof RecentsViewContainer recentsViewContainer) {
+            setRecentsViewContainer(recentsViewContainer);
+        }
+    }
+
+    /**
+     * Sets the current RecentsViewContainer, from which we create a TaskbarUIController.
+     */
+    public void setRecentsViewContainer(@NonNull RecentsViewContainer recentsViewContainer) {
+        if (mRecentsViewContainer == recentsViewContainer) {
+            return;
+        }
+        if (mRecentsViewContainer == mActivity) {
+            // When switching to RecentsWindowManager (not an Activity), the old mActivity is not
+            // destroyed, nor is there a new Activity to replace it. Thus if we don't clear it here,
+            // it will not get re-set properly if we return to the Activity (e.g. NexusLauncher).
+            mActivityOnDestroyCallback.run();
+        }
+        mRecentsViewContainer = recentsViewContainer;
         if (mTaskbarActivityContext != null) {
             mTaskbarActivityContext.setUIController(
-                    createTaskbarUIControllerForActivity(mActivity));
+                    createTaskbarUIControllerForRecentsViewContainer(mRecentsViewContainer));
         }
     }
 
@@ -428,12 +451,18 @@
     /**
      * Creates a {@link TaskbarUIController} to use while the given StatefulActivity is active.
      */
-    private TaskbarUIController createTaskbarUIControllerForActivity(StatefulActivity activity) {
-        if (activity instanceof QuickstepLauncher) {
-            return new LauncherTaskbarUIController((QuickstepLauncher) activity);
+    private TaskbarUIController createTaskbarUIControllerForRecentsViewContainer(
+            RecentsViewContainer container) {
+        if (container instanceof QuickstepLauncher quickstepLauncher) {
+            return new LauncherTaskbarUIController(quickstepLauncher);
         }
-        if (activity instanceof RecentsActivity) {
-            return new FallbackTaskbarUIController((RecentsActivity) activity);
+        // If a 3P Launcher is default, always use FallbackTaskbarUIController regardless of
+        // whether the recents container is RecentsActivity or RecentsWindowManager.
+        if (container instanceof RecentsActivity recentsActivity) {
+            return new FallbackTaskbarUIController<>(recentsActivity);
+        }
+        if (container instanceof RecentsWindowManager recentsWindowManager) {
+            return new FallbackTaskbarUIController<>(recentsWindowManager);
         }
         return TaskbarUIController.DEFAULT;
     }
@@ -481,9 +510,9 @@
             mSharedState.allAppsVisible = mSharedState.allAppsVisible && isLargeScreenTaskbar;
             mTaskbarActivityContext.init(mSharedState);
 
-            if (mActivity != null) {
+            if (mRecentsViewContainer != null) {
                 mTaskbarActivityContext.setUIController(
-                    createTaskbarUIControllerForActivity(mActivity));
+                        createTaskbarUIControllerForRecentsViewContainer(mRecentsViewContainer));
             }
 
             if (enableTaskbarNoRecreate()) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
index 729cbe9..a64dab1 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
@@ -30,6 +30,10 @@
 import android.view.InsetsFrameProvider;
 
 import com.android.systemui.shared.system.QuickStepContract.SystemUiStateFlags;
+import com.android.wm.shell.shared.bubbles.BubbleBarLocation;
+import com.android.wm.shell.shared.bubbles.BubbleInfo;
+
+import java.util.List;
 
 /**
  * State shared across different taskbar instance
@@ -69,6 +73,15 @@
 
     public boolean allAppsVisible = false;
 
+    public BubbleBarLocation bubbleBarLocation;
+
+    public List<BubbleInfo> bubbleInfoItems;
+
+    /** Returns whether there are a saved bubbles. */
+    public boolean hasSavedBubbles() {
+        return bubbleInfoItems != null && !bubbleInfoItems.isEmpty();
+    }
+
     // LauncherTaskbarUIController#mTaskbarInAppDisplayProgressMultiProp
     public float[] inAppDisplayProgressMultiPropValues = new float[DISPLAY_PROGRESS_COUNT];
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index b22fd6f..30e4e47 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -35,6 +35,7 @@
 import android.util.ArrayMap;
 import android.util.Log;
 
+import com.android.launcher3.taskbar.TaskbarSharedState;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
 import com.android.launcher3.util.Executors.SimpleThreadFactory;
 import com.android.quickstep.SystemUiProxy;
@@ -47,6 +48,7 @@
 import com.android.wm.shell.shared.bubbles.RemovedBubble;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
@@ -112,6 +114,7 @@
 
     private BubbleBarItem mSelectedBubble;
 
+    private TaskbarSharedState mSharedState;
     private ImeVisibilityChecker mImeVisibilityChecker;
     private BubbleBarViewController mBubbleBarViewController;
     private BubbleStashController mBubbleStashController;
@@ -173,12 +176,25 @@
 
     public void onDestroy() {
         mSystemUiProxy.setBubblesListener(null);
+        // Saves bubble bar state
+        BubbleInfo[] bubbleInfoItems = new BubbleInfo[mBubbles.size()];
+        mBubbles.values().forEach(bubbleBarBubble -> {
+            int index = mBubbleBarViewController.bubbleViewIndex(bubbleBarBubble.getView());
+            if (index < 0 || index >= bubbleInfoItems.length) {
+                Log.e(TAG, "Found improper index: " + index + " for " + bubbleBarBubble);
+            } else {
+                bubbleInfoItems[index] = bubbleBarBubble.getInfo();
+            }
+        });
+        mSharedState.bubbleInfoItems = Arrays.asList(bubbleInfoItems);
     }
 
     /** Initializes controllers. */
     public void init(BubbleControllers bubbleControllers,
             BubbleBarLocationListener bubbleBarLocationListener,
-            ImeVisibilityChecker imeVisibilityChecker) {
+            ImeVisibilityChecker imeVisibilityChecker,
+            TaskbarSharedState sharedState) {
+        mSharedState = sharedState;
         mImeVisibilityChecker = imeVisibilityChecker;
         mBubbleBarViewController = bubbleControllers.bubbleBarViewController;
         mBubbleStashController = bubbleControllers.bubbleStashController;
@@ -188,6 +204,7 @@
         mBubbleBarLocationListener = bubbleBarLocationListener;
 
         bubbleControllers.runAfterInit(() -> {
+            restoreSavedState(sharedState);
             mBubbleBarViewController.setHiddenForBubbles(
                     !sBubbleBarEnabled || mBubbles.isEmpty());
             mBubbleStashedHandleViewController.ifPresent(
@@ -266,6 +283,26 @@
         }
     }
 
+    private void restoreSavedState(TaskbarSharedState sharedState) {
+        if (sharedState.bubbleBarLocation != null) {
+            updateBubbleBarLocationInternal(sharedState.bubbleBarLocation);
+        }
+        List<BubbleInfo> bubbleInfos = sharedState.bubbleInfoItems;
+        if (bubbleInfos == null || bubbleInfos.isEmpty()) return;
+        // Iterate in reverse because new bubbles are added in front and the list is in order.
+        for (int i = bubbleInfos.size() - 1; i >= 0; i--) {
+            BubbleBarBubble bubble = mBubbleCreator.populateBubble(mContext,
+                    bubbleInfos.get(i), mBarView, /* existingBubble = */ null);
+            if (bubble == null) {
+                Log.e(TAG, "Could not instantiate BubbleBarBubble for " + bubbleInfos.get(i));
+                continue;
+            }
+            addBubbleInternally(bubble,  /* showAppBadge = */
+                    mBubbleBarViewController.isExpanded() || i == 0,
+                    /* isExpanding = */ false,  /* suppressAnimation = */ true);
+        }
+    }
+
     private void applyViewChanges(BubbleBarViewUpdate update) {
         final boolean isCollapsed = (update.expandedChanged && !update.expanded)
                 || (!update.expandedChanged && !mBubbleBarViewController.isExpanded());
@@ -277,6 +314,12 @@
                 update.initialState || mBubbleBarViewController.isHiddenForSysui()
                         || mImeVisibilityChecker.isImeVisible();
 
+        if (update.initialState && mSharedState.hasSavedBubbles()) {
+            // clear restored state
+            mBubbleBarViewController.removeAllBubbles();
+            mBubbles.clear();
+        }
+
         BubbleBarBubble bubbleToSelect = null;
 
         if (Flags.enableOptionalBubbleOverflow()
@@ -347,8 +390,8 @@
             for (int i = update.currentBubbles.size() - 1; i >= 0; i--) {
                 BubbleBarBubble bubble = update.currentBubbles.get(i);
                 if (bubble != null) {
-                    mBubbles.put(bubble.getKey(), bubble);
-                    mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+                    addBubbleInternally(bubble, /* showAppBadge = */ !isCollapsed || i == 0,
+                            isExpanding, suppressAnimation);
                     if (isCollapsed) {
                         // If we're collapsed, the most recently added bubble will be selected.
                         bubbleToSelect = bubble;
@@ -420,6 +463,7 @@
             }
         }
         if (update.bubbleBarLocation != null) {
+            mSharedState.bubbleBarLocation = update.bubbleBarLocation;
             if (update.bubbleBarLocation != mBubbleBarViewController.getBubbleBarLocation()) {
                 updateBubbleBarLocationInternal(update.bubbleBarLocation);
             }
@@ -519,6 +563,14 @@
         }
     }
 
+    private void addBubbleInternally(BubbleBarBubble bubble, boolean showAppBadge,
+            boolean isExpanding, boolean suppressAnimation) {
+        //TODO(b/360652359): remove setting scale to the app badge once issue is fixed
+        bubble.getView().setBadgeScale(showAppBadge ? 1 : 0);
+        mBubbles.put(bubble.getKey(), bubble);
+        mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
+    }
+
     /** Interface for checking whether the IME is visible. */
     public interface ImeVisibilityChecker {
         /** Whether the IME is visible. */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index fb2c6ea..96fadf7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -328,6 +328,7 @@
     }
 
     private void onBubbleClicked(BubbleView bubbleView) {
+        if (mBubbleBarPinning.isAnimating()) return;
         bubbleView.markSeen();
         BubbleBarItem bubble = bubbleView.getBubble();
         if (bubble == null) {
@@ -876,9 +877,14 @@
     /** Animates the bubble bar to notify the user about a bubble change. */
     public void animateBubbleNotification(BubbleBarBubble bubble, boolean isExpanding,
             boolean isUpdate) {
-        // if we're expanded, don't animate the bubble bar. just show the notification dot.
+        // if we're not already animating another bubble, update the dot visibility. otherwise the
+        // the dot will be handled as part of the animation.
+        if (!mBubbleBarViewAnimator.isAnimating()) {
+            bubble.getView().updateDotVisibility(
+                    /* animate= */ !mBubbleStashController.isStashed());
+        }
+        // if we're expanded, don't animate the bubble bar.
         if (isExpanded()) {
-            bubble.getView().updateDotVisibility(/* animate= */ true);
             return;
         }
         boolean isInApp = mTaskbarStashController.isInApp();
@@ -922,7 +928,7 @@
      * from Launcher.
      */
     public void setExpanded(boolean isExpanded) {
-        if (isExpanded != mBarView.isExpanded()) {
+        if (!mBubbleBarPinning.isAnimating() && isExpanded != mBarView.isExpanded()) {
             mBarView.setExpanded(isExpanded);
             adjustTaskbarAndHotseatToBubbleBarState(isExpanded);
             if (!isExpanded) {
@@ -1064,6 +1070,16 @@
         mSystemUiProxy.removeAllBubbles();
     }
 
+    /** Removes all existing bubble views */
+    public void removeAllBubbles() {
+        mBarView.removeAllViews();
+    }
+
+    /** Returns the view index of the existing bubble */
+    public int bubbleViewIndex(View bubbleView) {
+        return mBarView.indexOfChild(bubbleView);
+    }
+
     /**
      * Set listener to be notified when bubble bar bounds have changed
      */
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
index b5d94bd..d993685 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleControllers.java
@@ -21,6 +21,7 @@
 import android.view.View;
 
 import com.android.launcher3.taskbar.TaskbarControllers;
+import com.android.launcher3.taskbar.TaskbarSharedState;
 import com.android.launcher3.taskbar.bubbles.BubbleBarViewController.TaskbarViewPropertiesProvider;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleBarLocationOnDemandListener;
 import com.android.launcher3.taskbar.bubbles.stashing.BubbleStashController;
@@ -79,7 +80,7 @@
      * BubbleControllers instance, but should be careful to only access things that were created
      * in constructors for now, as some controllers may still be waiting for init().
      */
-    public void init(TaskbarControllers taskbarControllers) {
+    public void init(TaskbarSharedState taskbarSharedState, TaskbarControllers taskbarControllers) {
         BubbleBarLocationCompositeListener bubbleBarLocationListeners =
                 new BubbleBarLocationCompositeListener(
                         taskbarControllers.navbarButtonsViewController,
@@ -88,7 +89,8 @@
                 );
         bubbleBarController.init(this,
                 bubbleBarLocationListeners,
-                taskbarControllers.navbarButtonsViewController::isImeVisible);
+                taskbarControllers.navbarButtonsViewController::isImeVisible,
+                taskbarSharedState);
         bubbleStashedHandleViewController.ifPresent(
                 controller -> controller.init(/* bubbleControllers = */ this));
         bubbleStashController.init(
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
index 4f3e1ae..114edf4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
@@ -299,7 +299,8 @@
         return mBubble;
     }
 
-    void updateDotVisibility(boolean animate) {
+    /** Updates the dot visibility if it's not suppressed based on whether it has unseen content. */
+    public void updateDotVisibility(boolean animate) {
         if (mDotSuppressedForBubbleUpdate) {
             // if the dot is suppressed for an update, there's nothing to do
             return;
@@ -321,16 +322,12 @@
         }
     }
 
-    /**
-     * Suppresses or un-suppresses drawing the dot due to an update for this bubble.
-     *
-     * <p>If the dot is being suppressed and is already visible, it remains visible because it is
-     * used as a starting point for the animation. If the dot is being unsuppressed, it is
-     * redrawn if needed.
-     */
+    /** Suppresses or un-suppresses drawing the dot due to an update for this bubble. */
     public void suppressDotForBubbleUpdate(boolean suppress) {
         mDotSuppressedForBubbleUpdate = suppress;
-        if (!suppress) {
+        if (suppress) {
+            setDotScale(0);
+        } else {
             showDotIfNeeded(/* animate= */ false);
         }
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt
index eca2bc8..6c354f3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/animation/BubbleBarViewAnimator.kt
@@ -59,7 +59,7 @@
 
     private companion object {
         /** The time to show the flyout. */
-        const val FLYOUT_DELAY_MS: Long = 10000
+        const val FLYOUT_DELAY_MS: Long = 3000
         /** The initial scale Y value that the new bubble is set to before the animation starts. */
         const val BUBBLE_ANIMATION_INITIAL_SCALE_Y = 0.3f
         /** The minimum alpha value to make the bubble bar touchable. */
@@ -484,13 +484,14 @@
         val bubble = bubbleView?.bubble as? BubbleBarBubble
         val flyout = bubble?.flyoutMessage
         if (flyout != null) {
-            bubbleView.suppressDotForBubbleUpdate(true)
             bubbleBarFlyoutController.setUpAndShowFlyout(
-                BubbleBarFlyoutMessage(flyout.icon, flyout.title, flyout.message)
-            ) {
-                moveToState(AnimatingBubble.State.IN)
-                bubbleStashController.updateTaskbarTouchRegion()
-            }
+                BubbleBarFlyoutMessage(flyout.icon, flyout.title, flyout.message),
+                onInit = { bubbleView.suppressDotForBubbleUpdate(true) },
+                onEnd = {
+                    moveToState(AnimatingBubble.State.IN)
+                    bubbleStashController.updateTaskbarTouchRegion()
+                },
+            )
         } else {
             moveToState(AnimatingBubble.State.IN)
         }
@@ -563,7 +564,8 @@
         if (!bubbleBarFlyoutController.hasFlyout()) {
             // if the flyout does not yet exist, then we're only animating the bubble bar.
             // the animating bubble has been updated, so the when the flyout expands it will
-            // show the right message.
+            // show the right message. we only need to update the dot visibility.
+            bubbleView.updateDotVisibility(/* animate= */ !bubbleStashController.isStashed)
             return
         }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
index 1452cf6..7b20eea 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutController.kt
@@ -59,7 +59,7 @@
             return rect
         }
 
-    fun setUpAndShowFlyout(message: BubbleBarFlyoutMessage, onEnd: () -> Unit) {
+    fun setUpAndShowFlyout(message: BubbleBarFlyoutMessage, onInit: () -> Unit, onEnd: () -> Unit) {
         flyout?.let(container::removeView)
         val flyout = BubbleBarFlyoutView(container.context, positioner, flyoutScheduler)
 
@@ -76,7 +76,11 @@
         container.addView(flyout, lp)
 
         this.flyout = flyout
-        flyout.showFromCollapsed(message) { showFlyout(AnimationType.MORPH, onEnd) }
+        flyout.showFromCollapsed(message) {
+            flyout.updateExpansionProgress(0f)
+            onInit()
+            showFlyout(AnimationType.MORPH, onEnd)
+        }
     }
 
     private fun showFlyout(animationType: AnimationType, endAction: () -> Unit) {
@@ -160,7 +164,15 @@
                     flyout.updateExpansionProgress(animator.animatedValue as Float)
                 }
         }
-        animator.addListener(onStart = { flyout.setOnClickListener(null) }, onEnd = { endAction() })
+        animator.addListener(
+            onStart = {
+                flyout.setOnClickListener(null)
+                if (animationType == AnimationType.MORPH) {
+                    flyout.updateTranslationToCollapsedPosition()
+                }
+            },
+            onEnd = { endAction() },
+        )
         animator.start()
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
index af8aaf8..418675c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutView.kt
@@ -192,18 +192,8 @@
         title.alpha = 0f
         message.alpha = 0f
         setData(flyoutMessage)
-        val txToCollapsedPosition =
-            if (positioner.isOnLeft) {
-                positioner.distanceToCollapsedPosition.x
-            } else {
-                -positioner.distanceToCollapsedPosition.x
-            }
-        // TODO: b/277815200 - before collapsing, recalculate translationToCollapsedPosition because
-        // the collapsed position may have changed
-        val tyToCollapsedPosition =
-            positioner.distanceToCollapsedPosition.y + triangleHeight - triangleOverlap
-        translationToCollapsedPosition = PointF(txToCollapsedPosition, tyToCollapsedPosition)
 
+        updateTranslationToCollapsedPosition()
         collapsedSize = positioner.collapsedSize
         collapsedCornerRadius = collapsedSize / 2
         collapsedColor = positioner.collapsedColor
@@ -212,7 +202,7 @@
         // calculate the expansion progress required before we start showing the triangle as part of
         // the expansion animation
         minExpansionProgressForTriangle =
-            positioner.distanceToRevealTriangle / tyToCollapsedPosition
+            positioner.distanceToRevealTriangle / translationToCollapsedPosition.y
 
         // post the request to start the expand animation to the looper so the view can measure
         // itself
@@ -259,6 +249,22 @@
         message.text = flyoutMessage.message
     }
 
+    /**
+     * This should be called to update [translationToCollapsedPosition] before we start expanding or
+     * collapsing to make sure that we're animating the flyout to and from the correct position.
+     */
+    fun updateTranslationToCollapsedPosition() {
+        val txToCollapsedPosition =
+            if (positioner.isOnLeft) {
+                positioner.distanceToCollapsedPosition.x
+            } else {
+                -positioner.distanceToCollapsedPosition.x
+            }
+        val tyToCollapsedPosition =
+            positioner.distanceToCollapsedPosition.y + triangleHeight - triangleOverlap
+        translationToCollapsedPosition = PointF(txToCollapsedPosition, tyToCollapsedPosition)
+    }
+
     /** Updates the flyout view with the progress of the animation. */
     fun updateExpansionProgress(fraction: Float) {
         expansionProgress = fraction
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 228dc91..fe68ebc 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -139,6 +139,7 @@
 import com.android.launcher3.statemanager.StateManager.StateHandler;
 import com.android.launcher3.taskbar.LauncherTaskbarUIController;
 import com.android.launcher3.taskbar.TaskbarManager;
+import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.testing.TestLogging;
 import com.android.launcher3.testing.shared.TestProtocol;
 import com.android.launcher3.uioverrides.QuickstepWidgetHolder.QuickstepHolderFactory;
@@ -1089,10 +1090,12 @@
         );
     }
 
-    public void setTaskbarUIController(LauncherTaskbarUIController taskbarUIController) {
-        mTaskbarUIController = taskbarUIController;
+    @Override
+    public void setTaskbarUIController(@Nullable TaskbarUIController taskbarUIController) {
+        mTaskbarUIController = (LauncherTaskbarUIController) taskbarUIController;
     }
 
+    @Override
     public @Nullable LauncherTaskbarUIController getTaskbarUIController() {
         return mTaskbarUIController;
     }
@@ -1399,6 +1402,7 @@
     }
 
     @NonNull
+    @Override
     public TISBindHelper getTISBindHelper() {
         return mTISBindHelper;
     }
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index b19f651..6075294 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -70,6 +70,7 @@
 import com.android.launcher3.statemanager.StatefulActivity;
 import com.android.launcher3.taskbar.FallbackTaskbarUIController;
 import com.android.launcher3.taskbar.TaskbarManager;
+import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.util.ActivityOptionsWrapper;
 import com.android.launcher3.util.ContextTracker;
 import com.android.launcher3.util.RunnableList;
@@ -115,7 +116,7 @@
     private FallbackRecentsView mFallbackRecentsView;
     private OverviewActionsView<?> mActionsView;
     private TISBindHelper mTISBindHelper;
-    private @Nullable FallbackTaskbarUIController mTaskbarUIController;
+    private @Nullable FallbackTaskbarUIController<RecentsActivity> mTaskbarUIController;
 
     private StateManager<RecentsState, RecentsActivity> mStateManager;
 
@@ -174,11 +175,14 @@
         mTISBindHelper.runOnBindToTouchInteractionService(r);
     }
 
-    public void setTaskbarUIController(FallbackTaskbarUIController taskbarUIController) {
-        mTaskbarUIController = taskbarUIController;
+    @Override
+    public void setTaskbarUIController(@Nullable TaskbarUIController taskbarUIController) {
+        mTaskbarUIController = (FallbackTaskbarUIController<RecentsActivity>) taskbarUIController;
     }
 
-    public FallbackTaskbarUIController getTaskbarUIController() {
+    @Nullable
+    @Override
+    public FallbackTaskbarUIController<RecentsActivity> getTaskbarUIController() {
         return mTaskbarUIController;
     }
 
@@ -515,6 +519,7 @@
     }
 
     @NonNull
+    @Override
     public TISBindHelper getTISBindHelper() {
         return mTISBindHelper;
     }
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 032e755..e8f38be 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -776,10 +776,13 @@
         mAllAppsActionManager.setHomeAndOverviewSame(isHomeAndOverviewSame);
         RecentsViewContainer newOverviewContainer =
                 mOverviewComponentObserver.getContainerInterface().getCreatedContainer();
-        if (newOverviewContainer != null
-                && newOverviewContainer instanceof StatefulActivity activity) {
-            //TODO(b/368030750) refactor taskbarManager to accept RecentsViewContainer
-            mTaskbarManager.setActivity(activity);
+        if (newOverviewContainer != null) {
+            if (newOverviewContainer instanceof StatefulActivity activity) {
+                // This will also call setRecentsViewContainer() internally.
+                mTaskbarManager.setActivity(activity);
+            } else {
+                mTaskbarManager.setRecentsViewContainer(newOverviewContainer);
+            }
         }
         mTISBinder.onOverviewTargetChange();
     }
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
index 74f9901..78224ae 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowManager.kt
@@ -40,6 +40,7 @@
 import com.android.launcher3.statemanager.StateManager
 import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory
 import com.android.launcher3.statemanager.StatefulContainer
+import com.android.launcher3.taskbar.TaskbarUIController
 import com.android.launcher3.util.ContextTracker
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.RunnableList
@@ -117,6 +118,7 @@
 
     private var callbacks: RecentsAnimationCallbacks? = null
 
+    private var taskbarUIController: TaskbarUIController? = null
     private var tisBindHelper: TISBindHelper = TISBindHelper(this) {}
 
     // Callback array that corresponds to events defined in @ActivityEvent
@@ -290,6 +292,18 @@
         return tisBindHelper.desktopVisibilityController
     }
 
+    override fun setTaskbarUIController(taskbarUIController: TaskbarUIController?) {
+        this.taskbarUIController = taskbarUIController
+    }
+
+    override fun getTaskbarUIController(): TaskbarUIController? {
+        return taskbarUIController
+    }
+
+    override fun getTISBindHelper(): TISBindHelper {
+        return tisBindHelper
+    }
+
     fun registerInitListener(onInitListener: Predicate<Boolean>) {
         this.onInitListener = onInitListener
     }
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 4f38ec7..275af00 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -73,13 +73,17 @@
         getTaskDataById(taskId).map { it?.thumbnail }.distinctUntilChangedBy { it?.snapshotId }
 
     override fun setVisibleTasks(visibleTaskIdList: Set<Int>) {
-        Log.d(TAG, "setVisibleTasks: $visibleTaskIdList")
-
         // Remove tasks are no longer visible
         val tasksNoLongerVisible = taskRequests.keys.subtract(visibleTaskIdList)
         removeTasks(tasksNoLongerVisible)
         // Add new tasks to be requested
-        visibleTaskIdList.subtract(taskRequests.keys).forEach { taskId -> requestTaskData(taskId) }
+        val newlyVisibleTasks = visibleTaskIdList.subtract(taskRequests.keys)
+        newlyVisibleTasks.forEach { taskId -> requestTaskData(taskId) }
+
+        if (tasksNoLongerVisible.isNotEmpty() || newlyVisibleTasks.isNotEmpty()) {
+            Log.d(TAG, "setVisibleTasks to: $visibleTaskIdList, " +
+                    "removed: $tasksNoLongerVisible, added: $newlyVisibleTasks")
+        }
     }
 
     private fun requestTaskData(taskId: Int) {
diff --git a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
index 4962367..bdfaa48 100644
--- a/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitToWorkspaceController.java
@@ -48,8 +48,11 @@
 import com.android.launcher3.uioverrides.QuickstepLauncher;
 import com.android.quickstep.views.FloatingTaskView;
 import com.android.quickstep.views.RecentsView;
+import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 
+import java.util.Collections;
+
 /** Handles when the stage split lands on the home screen. */
 public class SplitToWorkspaceController {
 
@@ -133,10 +136,20 @@
             // Use Launcher's default click handler
             return false;
         }
-
-        mController.setSecondTask(intent, user, (ItemInfo) tag);
-
-        startWorkspaceAnimation(view, null /*bitmap*/, bitmapInfo.newIcon(mLauncher));
+        // Check for background task matching this tag; if we find one, set second task
+        // via task instead of intent so the bounds and windowing mode will be corrected.
+        mController.findLastActiveTasksAndRunCallback(
+                Collections.singletonList(((ItemInfo) tag).getComponentKey()),
+                false /* findExactPairMatch */,
+                foundTasks -> {
+                    Task foundTask = foundTasks[0];
+                    if (foundTask != null) {
+                        mController.setSecondTask(foundTask, (ItemInfo) tag);
+                    } else {
+                        mController.setSecondTask(intent, user, (ItemInfo) tag);
+                    }
+                    startWorkspaceAnimation(view, null /*bitmap*/, bitmapInfo.newIcon(mLauncher));
+                });
         return true;
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskContentView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskContentView.kt
new file mode 100644
index 0000000..481acac
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskContentView.kt
@@ -0,0 +1,52 @@
+/*
+ * 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.views
+
+import android.content.Context
+import android.graphics.Outline
+import android.graphics.Rect
+import android.util.AttributeSet
+import android.view.View
+import android.view.ViewOutlineProvider
+import android.widget.FrameLayout
+import com.android.quickstep.views.TaskView.FullscreenDrawParams
+
+class DesktopTaskContentView
+@JvmOverloads
+constructor(context: Context, attrs: AttributeSet? = null) : FrameLayout(context, attrs) {
+    private val currentFullscreenParams = FullscreenDrawParams(context)
+    private val taskCornerRadius: Float
+        get() = currentFullscreenParams.cornerRadius
+
+    private val bounds = Rect()
+
+    init {
+        clipToOutline = true
+        outlineProvider =
+            object : ViewOutlineProvider() {
+                override fun getOutline(view: View, outline: Outline) {
+                    outline.setRoundRect(bounds, taskCornerRadius)
+                }
+            }
+    }
+
+    override fun onSizeChanged(w: Int, h: Int, oldw: Int, oldh: Int) {
+        super.onSizeChanged(w, h, oldw, oldh)
+        bounds.set(0, 0, w, h)
+        invalidateOutline()
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
index 15b0a6b..5e842aa 100644
--- a/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/DesktopTaskView.kt
@@ -26,6 +26,7 @@
 import android.util.Log
 import android.view.Gravity
 import android.view.View
+import android.widget.FrameLayout
 import androidx.core.content.res.ResourcesCompat
 import androidx.core.view.updateLayoutParams
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
@@ -81,12 +82,12 @@
     private val tempRect = Rect()
     private lateinit var backgroundView: View
     private lateinit var iconView: TaskViewIcon
-    private var childCountAtInflation = 0
+    private lateinit var contentView: FrameLayout
 
     override fun onFinishInflate() {
         super.onFinishInflate()
         backgroundView =
-            findViewById<View>(R.id.background)!!.apply {
+            findViewById<View>(R.id.background).apply {
                 updateLayoutParams<LayoutParams> {
                     topMargin = container.deviceProfile.overviewTaskThumbnailTopMarginPx
                 }
@@ -113,7 +114,12 @@
                 )
                 setText(resources.getText(R.string.recent_task_desktop))
             }
-        childCountAtInflation = childCount
+        contentView =
+            findViewById<FrameLayout>(R.id.desktop_content).apply {
+                updateLayoutParams<LayoutParams> {
+                    topMargin = container.deviceProfile.overviewTaskThumbnailTopMarginPx
+                }
+            }
     }
 
     /** Updates this desktop task to the gives task list defined in `tasks` */
@@ -137,13 +143,8 @@
                     } else {
                         taskThumbnailViewDeprecatedPool!!.view
                     }
+                contentView.addView(snapshotView, 0)
 
-                addView(
-                    snapshotView,
-                    // Add snapshotView to the front after initial views e.g. icon and
-                    // background.
-                    childCountAtInflation,
-                )
                 TaskContainer(
                     this,
                     task,
@@ -164,7 +165,7 @@
         super.onRecycle()
         visibility = VISIBLE
         taskContainers.forEach {
-            removeView(it.snapshotView)
+            contentView.removeView(it.snapshotView)
             if (enableRefactorTaskThumbnail()) {
                 taskThumbnailViewPool!!.recycle(it.thumbnailView)
             } else {
@@ -227,9 +228,7 @@
                 width = (taskSize.width() * scaleWidth).toInt()
                 height = (taskSize.height() * scaleHeight).toInt()
                 leftMargin = (positionInParent.x * scaleWidth).toInt()
-                topMargin =
-                    (positionInParent.y * scaleHeight).toInt() +
-                        container.deviceProfile.overviewTaskThumbnailTopMarginPx
+                topMargin = (positionInParent.y * scaleHeight).toInt()
             }
             if (DEBUG) {
                 with(it.snapshotView.layoutParams as LayoutParams) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 2d0f15e..9a8041b 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -1245,21 +1245,24 @@
         // - It's the focused task to be moved to the front, we immediately re-add the task
         if (child instanceof TaskView && child != mSplitHiddenTaskView
                 && child != mMovingTaskView) {
-            TaskView taskView = (TaskView) child;
-            for (int i : taskView.getTaskIds()) {
-                mHasVisibleTaskData.delete(i);
-            }
-            if (child instanceof GroupedTaskView) {
-                mGroupedTaskViewPool.recycle((GroupedTaskView) taskView);
-            } else if (child instanceof DesktopTaskView) {
-                mDesktopTaskViewPool.recycle((DesktopTaskView) taskView);
-            } else {
-                mTaskViewPool.recycle(taskView);
-            }
-            mActionsView.updateHiddenFlags(HIDDEN_NO_TASKS, getTaskViewCount() == 0);
+            clearAndRecycleTaskView((TaskView) child);
         }
     }
 
+    private void clearAndRecycleTaskView(TaskView taskView) {
+        for (int taskId : taskView.getTaskIds()) {
+            mHasVisibleTaskData.delete(taskId);
+        }
+        if (taskView instanceof GroupedTaskView) {
+            mGroupedTaskViewPool.recycle((GroupedTaskView) taskView);
+        } else if (taskView instanceof DesktopTaskView) {
+            mDesktopTaskViewPool.recycle((DesktopTaskView) taskView);
+        } else {
+            mTaskViewPool.recycle(taskView);
+        }
+        mActionsView.updateHiddenFlags(HIDDEN_NO_TASKS, getTaskViewCount() == 0);
+    }
+
     @Override
     public void onViewAdded(View child) {
         super.onViewAdded(child);
@@ -5303,6 +5306,13 @@
         mSplitHiddenTaskViewIndex = -1;
         if (mSplitHiddenTaskView != null) {
             mSplitHiddenTaskView.setThumbnailVisibility(VISIBLE, INVALID_TASK_ID);
+            // mSplitHiddenTaskView is set when split select animation starts. The TaskView is only
+            // removed when when the animation finishes. So in the case of overview being dismissed
+            // during the animation, we should not call clearAndRecycleTaskView() because it has
+            // not been removed yet.
+            if (mSplitHiddenTaskView.getParent() == null) {
+                clearAndRecycleTaskView(mSplitHiddenTaskView);
+            }
             mSplitHiddenTaskView = null;
         }
     }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java b/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java
index d8036aa..b04753b 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsViewContainer.java
@@ -16,8 +16,6 @@
 
 package com.android.quickstep.views;
 
-import android.app.Activity;
-import android.content.ComponentName;
 import android.content.Context;
 import android.content.ContextWrapper;
 import android.content.LocusId;
@@ -27,13 +25,16 @@
 import android.view.View;
 import android.view.Window;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.BaseActivity;
 import com.android.launcher3.logger.LauncherAtom;
 import com.android.launcher3.statehandlers.DesktopVisibilityController;
+import com.android.launcher3.taskbar.TaskbarUIController;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.ScrimView;
+import com.android.quickstep.util.TISBindHelper;
 
 /**
  * Interface to be implemented by the parent view of RecentsView
@@ -212,4 +213,10 @@
 
     @Nullable
     DesktopVisibilityController getDesktopVisibilityController();
+
+    void setTaskbarUIController(@Nullable TaskbarUIController taskbarUIController);
+
+    @Nullable TaskbarUIController getTaskbarUIController();
+
+    @NonNull TISBindHelper getTISBindHelper();
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 959516f..25aba39 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -151,7 +151,7 @@
         if (enableRefactorTaskThumbnail()) {
             bindThumbnailView()
         } else {
-            thumbnailViewDeprecated.bind(task, overlay)
+            thumbnailViewDeprecated.bind(task, overlay, taskView)
         }
         overlay.init()
     }
diff --git a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
index 56ca043..5dbc2ef 100644
--- a/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
+++ b/quickstep/src/com/android/quickstep/views/TaskThumbnailViewDeprecated.java
@@ -110,6 +110,7 @@
     private TaskView.FullscreenDrawParams mFullscreenParams;
     private ImageView mSplashView;
     private Drawable mSplashViewDrawable;
+    private TaskView mTaskView;
 
     @Nullable
     private Task mTask;
@@ -153,10 +154,11 @@
     /**
      * Updates the thumbnail to draw the provided task
      */
-    public void bind(Task task, TaskOverlay<?> overlay) {
+    public void bind(Task task, TaskOverlay<?> overlay, TaskView taskView) {
         mOverlay = overlay;
         mOverlay.reset();
         mTask = task;
+        mTaskView = taskView;
         int color = task == null ? Color.BLACK : task.colorBackground | 0xFF000000;
         mPaint.setColor(color);
         mBackgroundPaint.setColor(color);
@@ -292,8 +294,8 @@
 
     public void drawOnCanvas(Canvas canvas, float x, float y, float width, float height,
             float cornerRadius) {
-        if (mTask != null && getTaskView().isRunningTask()
-                && !getTaskView().getShouldShowScreenshot()) {
+        if (mTask != null && mTaskView.isRunningTask()
+                && !mTaskView.getShouldShowScreenshot()) {
             canvas.drawRoundRect(x, y, width, height, cornerRadius, cornerRadius, mClearPaint);
             canvas.drawRoundRect(x, y, width, height, cornerRadius, cornerRadius,
                     mDimmingPaintAfterClearing);
@@ -334,10 +336,6 @@
         }
     }
 
-    public TaskView getTaskView() {
-        return (TaskView) getParent();
-    }
-
     public void setOverlayEnabled(boolean overlayEnabled) {
         if (mOverlayEnabled != overlayEnabled) {
             mOverlayEnabled = overlayEnabled;
@@ -390,9 +388,9 @@
         float viewCenterY = viewHeight / 2f;
         float centeredDrawableLeft = (viewWidth - drawableWidth) / 2f;
         float centeredDrawableTop = (viewHeight - drawableHeight) / 2f;
-        float nonGridScale = getTaskView() == null ? 1 : 1 / getTaskView().getNonGridScale();
-        float recentsMaxScale = getTaskView() == null || getTaskView().getRecentsView() == null
-                ? 1 : 1 / getTaskView().getRecentsView().getMaxScaleForFullScreen();
+        float nonGridScale = mTaskView == null ? 1 : 1 / mTaskView.getNonGridScale();
+        float recentsMaxScale = mTaskView == null || mTaskView.getRecentsView() == null
+                ? 1 : 1 / mTaskView.getRecentsView().getMaxScaleForFullScreen();
         float scaleX = nonGridScale * recentsMaxScale * (1 / getScaleX());
         float scaleY = nonGridScale * recentsMaxScale * (1 / getScaleY());
 
@@ -419,7 +417,7 @@
     }
 
     private boolean isThumbnailRotationDifferentFromTask() {
-        RecentsView recents = getTaskView().getRecentsView();
+        RecentsView recents = mTaskView.getRecentsView();
         if (recents == null || mThumbnailData == null) {
             return false;
         }
@@ -467,7 +465,7 @@
         if (mBitmapShader != null && mThumbnailData != null) {
             mPreviewRect.set(0, 0, mThumbnailData.getThumbnail().getWidth(),
                     mThumbnailData.getThumbnail().getHeight());
-            int currentRotation = getTaskView().getOrientedState().getRecentsActivityRotation();
+            int currentRotation = mTaskView.getOrientedState().getRecentsActivityRotation();
             boolean isRtl = getLayoutDirection() == LAYOUT_DIRECTION_RTL;
             mPreviewPositionHelper.updateThumbnailMatrix(mPreviewRect, mThumbnailData,
                     getMeasuredWidth(), getMeasuredHeight(), dp.isTablet, currentRotation, isRtl);
@@ -475,7 +473,7 @@
             mBitmapShader.setLocalMatrix(mPreviewPositionHelper.getMatrix());
             mPaint.setShader(mBitmapShader);
         }
-        getTaskView().updateCurrentFullscreenParams();
+        mTaskView.updateCurrentFullscreenParams();
         invalidate();
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
index fef82c1..2997ac9 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/bubbles/flyout/BubbleBarFlyoutControllerTest.kt
@@ -76,7 +76,7 @@
     @Test
     fun flyoutPosition_left() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             val flyout = flyoutContainer.getChildAt(0)
             val lp = flyout.layoutParams as FrameLayout.LayoutParams
@@ -89,7 +89,7 @@
     fun flyoutPosition_right() {
         onLeft = false
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             val flyout = flyoutContainer.getChildAt(0)
             val lp = flyout.layoutParams as FrameLayout.LayoutParams
@@ -101,7 +101,7 @@
     @Test
     fun flyoutMessage() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             val flyout = flyoutContainer.getChildAt(0)
             val sender = flyout.findViewById<TextView>(R.id.bubble_flyout_title)
@@ -114,7 +114,7 @@
     @Test
     fun hideFlyout_removedFromContainer() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutController.hasFlyout()).isTrue()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             flyoutController.collapseFlyout {}
@@ -130,7 +130,7 @@
         // boundary
         flyoutTy = -50f
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
         }
         InstrumentationRegistry.getInstrumentation().waitForIdleSync()
@@ -143,7 +143,7 @@
     @Test
     fun showFlyout_withinBoundary() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
         }
         InstrumentationRegistry.getInstrumentation().waitForIdleSync()
@@ -156,7 +156,7 @@
     @Test
     fun collapseFlyout_resetsTopBoundary() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             flyoutController.collapseFlyout {}
             animatorTestRule.advanceTimeBy(300)
@@ -167,7 +167,7 @@
     @Test
     fun cancelFlyout_fadesOutFlyout() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             val flyoutView = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
             assertThat(flyoutView.alpha).isEqualTo(1f)
@@ -181,7 +181,7 @@
     @Test
     fun clickFlyout_notifiesCallback() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutContainer.childCount).isEqualTo(1)
             val flyoutView = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
             assertThat(flyoutView.alpha).isEqualTo(1f)
@@ -194,7 +194,7 @@
     @Test
     fun updateFlyoutWhileExpanding() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             assertThat(flyoutController.hasFlyout()).isTrue()
             val flyout = flyoutContainer.findViewById<View>(R.id.bubble_bar_flyout_view)
             assertThat(flyout.findViewById<TextView>(R.id.bubble_flyout_text).text)
@@ -220,7 +220,7 @@
     @Test
     fun updateFlyoutFullyExpanded() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             animatorTestRule.advanceTimeBy(300)
         }
         assertThat(flyoutController.hasFlyout()).isTrue()
@@ -249,7 +249,7 @@
     @Test
     fun updateFlyoutWhileCollapsing() {
         InstrumentationRegistry.getInstrumentation().runOnMainSync {
-            flyoutController.setUpAndShowFlyout(flyoutMessage) {}
+            setupAndShowFlyout()
             animatorTestRule.advanceTimeBy(300)
         }
         assertThat(flyoutController.hasFlyout()).isTrue()
@@ -280,6 +280,10 @@
         assertThat(flyoutController.hasFlyout()).isTrue()
     }
 
+    private fun setupAndShowFlyout() {
+        flyoutController.setUpAndShowFlyout(flyoutMessage, {}, {})
+    }
+
     class FakeFlyoutCallbacks : FlyoutCallbacks {
 
         var topBoundaryExtendedSpace = 0
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
index 04012c0..df98606 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/FallbackTaskbarUIControllerTest.kt
@@ -33,7 +33,7 @@
 @RunWith(AndroidJUnit4::class)
 class FallbackTaskbarUIControllerTest : TaskbarBaseTestCase() {
 
-    lateinit var fallbackTaskbarUIController: FallbackTaskbarUIController
+    lateinit var fallbackTaskbarUIController: FallbackTaskbarUIController<RecentsActivity>
     lateinit var stateListener: StateManager.StateListener<RecentsState>
 
     private val recentsActivity: RecentsActivity = mock()
diff --git a/res/values-fr-rCA/strings.xml b/res/values-fr-rCA/strings.xml
index cae77dc..a941d88 100644
--- a/res/values-fr-rCA/strings.xml
+++ b/res/values-fr-rCA/strings.xml
@@ -192,7 +192,7 @@
     <string name="private_space_label" msgid="2359721649407947001">"Espace privé"</string>
     <string name="private_space_secondary_label" msgid="9203933341714508907">"Touchez pour configurer ou ouvrir"</string>
     <string name="ps_container_title" msgid="4391796149519594205">"Privé"</string>
-    <string name="ps_container_settings" msgid="6059734123353320479">"Paramètres de l\'Espace privé"</string>
+    <string name="ps_container_settings" msgid="6059734123353320479">"Paramètres de l\'espace privé"</string>
     <string name="ps_container_unlock_button_content_description" msgid="9181551784092204234">"Privé, déverrouillé."</string>
     <string name="ps_container_lock_button_content_description" msgid="5961993384382649530">"Privé, verrouillé."</string>
     <string name="ps_container_lock_title" msgid="2640257399982364682">"Verrouiller"</string>
diff --git a/res/values-it/strings.xml b/res/values-it/strings.xml
index 731f839..9cdb5aa 100644
--- a/res/values-it/strings.xml
+++ b/res/values-it/strings.xml
@@ -192,7 +192,7 @@
     <string name="private_space_label" msgid="2359721649407947001">"Spazio privato"</string>
     <string name="private_space_secondary_label" msgid="9203933341714508907">"Tocca per configurare o aprire"</string>
     <string name="ps_container_title" msgid="4391796149519594205">"Privato"</string>
-    <string name="ps_container_settings" msgid="6059734123353320479">"Impostazioni dello Spazio privato"</string>
+    <string name="ps_container_settings" msgid="6059734123353320479">"Impostazioni dello spazio privato"</string>
     <string name="ps_container_unlock_button_content_description" msgid="9181551784092204234">"Privato, sbloccato."</string>
     <string name="ps_container_lock_button_content_description" msgid="5961993384382649530">"Privato, bloccato."</string>
     <string name="ps_container_lock_title" msgid="2640257399982364682">"Blocca"</string>
diff --git a/res/values/config.xml b/res/values/config.xml
index b0b7aa2..a1ccb67 100644
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -85,6 +85,9 @@
     <string name="local_colors_extraction_class" translatable="false"></string>
     <string name="search_session_manager_class" translatable="false"></string>
 
+    <!-- Filters for widgets displayed in the widget picker  -->
+    <string name="widgets_filter_data_provider_class" translatable="false"></string>
+
     <!-- Scalable Grid configuration -->
     <!-- This is a float because it is converted to dp later in DeviceProfile -->
     <dimen name="hotseat_bar_bottom_space_default">48</dimen>
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 983cf8d..6446f7b 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2629,8 +2629,9 @@
      * See {@code LauncherBindingDelegate}
      */
     @Override
-    public void bindAllWidgets(final List<WidgetsListBaseEntry> allWidgets) {
-        mModelCallbacks.bindAllWidgets(allWidgets);
+    public void bindAllWidgets(@NonNull final List<WidgetsListBaseEntry> allWidgets,
+            @NonNull final List<WidgetsListBaseEntry> defaultWidgets) {
+        mModelCallbacks.bindAllWidgets(allWidgets, defaultWidgets);
     }
 
     @Override
diff --git a/src/com/android/launcher3/LauncherAppState.java b/src/com/android/launcher3/LauncherAppState.java
index b6da164..01d0a74 100644
--- a/src/com/android/launcher3/LauncherAppState.java
+++ b/src/com/android/launcher3/LauncherAppState.java
@@ -48,6 +48,7 @@
 import com.android.launcher3.icons.LauncherIconProvider;
 import com.android.launcher3.icons.LauncherIcons;
 import com.android.launcher3.model.ModelLauncherCallbacks;
+import com.android.launcher3.model.WidgetsFilterDataProvider;
 import com.android.launcher3.notification.NotificationListener;
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.InstallSessionTracker;
@@ -197,7 +198,8 @@
         mIconProvider = new LauncherIconProvider(context);
         mIconCache = new IconCache(mContext, mInvariantDeviceProfile,
                 iconCacheFileName, mIconProvider);
-        mModel = new LauncherModel(context, this, mIconCache, new AppFilter(mContext),
+        mModel = new LauncherModel(context, this, mIconCache,
+                WidgetsFilterDataProvider.Companion.newInstance(context), new AppFilter(mContext),
                 PackageManagerHelper.INSTANCE.get(context), iconCacheFileName != null);
         mOnTerminateCallback.add(mIconCache::close);
         mOnTerminateCallback.add(mModel::destroy);
diff --git a/src/com/android/launcher3/LauncherModel.kt b/src/com/android/launcher3/LauncherModel.kt
index 85ecd58..b56df46 100644
--- a/src/com/android/launcher3/LauncherModel.kt
+++ b/src/com/android/launcher3/LauncherModel.kt
@@ -42,6 +42,7 @@
 import com.android.launcher3.model.ReloadStringCacheTask
 import com.android.launcher3.model.ShortcutsChangedTask
 import com.android.launcher3.model.UserLockStateChangedTask
+import com.android.launcher3.model.WidgetsFilterDataProvider
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.pm.UserCache
@@ -66,6 +67,7 @@
     private val context: Context,
     private val mApp: LauncherAppState,
     private val iconCache: IconCache,
+    private val widgetsFilterDataProvider: WidgetsFilterDataProvider,
     appFilter: AppFilter,
     mPmHelper: PackageManagerHelper,
     isPrimaryInstance: Boolean,
@@ -140,6 +142,11 @@
         owner: BgDataModel.Callbacks?,
     ) = ModelWriter(mApp.context, this, mBgDataModel, verifyChanges, cellPosMapper, owner)
 
+    /** Returns the [WidgetsFilterDataProvider] that manages widget filters. */
+    fun getWidgetsFilterDataProvider(): WidgetsFilterDataProvider {
+        return widgetsFilterDataProvider
+    }
+
     /** Called when the icon for an app changes, outside of package event */
     @WorkerThread
     fun onAppIconChanged(packageName: String, user: UserHandle) {
@@ -160,7 +167,10 @@
     /** Called when the model is destroyed */
     fun destroy() {
         mModelDestroyed = true
-        MODEL_EXECUTOR.execute(modelDelegate::destroy)
+        MODEL_EXECUTOR.execute {
+            modelDelegate.destroy()
+            widgetsFilterDataProvider.destroy()
+        }
     }
 
     fun onBroadcastIntent(intent: Intent) {
@@ -312,6 +322,7 @@
                             mBgDataModel,
                             this.modelDelegate,
                             launcherBinder,
+                            widgetsFilterDataProvider,
                         )
 
                     // Always post the loader task, instead of running directly
@@ -417,6 +428,14 @@
         }
     }
 
+    /** Called when the widget filters are refreshed and available to bind to the model. */
+    fun onWidgetFiltersLoaded() {
+        enqueueModelUpdateTask { taskController, dataModel, _ ->
+            dataModel.widgetsModel.updateWidgetFilters(widgetsFilterDataProvider)
+            taskController.bindUpdatedWidgets(dataModel)
+        }
+    }
+
     fun enqueueModelUpdateTask(task: ModelUpdateTask) {
         if (mModelDestroyed) {
             return
diff --git a/src/com/android/launcher3/ModelCallbacks.kt b/src/com/android/launcher3/ModelCallbacks.kt
index 496d517..5d32525 100644
--- a/src/com/android/launcher3/ModelCallbacks.kt
+++ b/src/com/android/launcher3/ModelCallbacks.kt
@@ -252,8 +252,11 @@
         PopupContainerWithArrow.dismissInvalidPopup(launcher)
     }
 
-    override fun bindAllWidgets(allWidgets: List<WidgetsListBaseEntry>) {
-        launcher.widgetPickerDataProvider.setWidgets(allWidgets, /* defaultWidgets= */ listOf())
+    override fun bindAllWidgets(
+        allWidgets: List<WidgetsListBaseEntry>,
+        defaultWidgets: List<WidgetsListBaseEntry>,
+    ) {
+        launcher.widgetPickerDataProvider.setWidgets(allWidgets, defaultWidgets)
     }
 
     /** Returns the ids of the workspaces to bind. */
diff --git a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
index 0fa275e..e89671e 100644
--- a/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
+++ b/src/com/android/launcher3/dagger/LauncherBaseAppComponent.java
@@ -26,6 +26,7 @@
 import com.android.launcher3.util.PluginManagerWrapper;
 import com.android.launcher3.util.ScreenOnTracker;
 import com.android.launcher3.util.SettingsCache;
+import com.android.launcher3.util.VibratorWrapper;
 import com.android.launcher3.util.window.RefreshRateTracker;
 import com.android.launcher3.widget.custom.CustomWidgetManager;
 
@@ -48,8 +49,9 @@
     RefreshRateTracker getRefreshRateTracker();
     ScreenOnTracker getScreenOnTracker();
     SettingsCache getSettingsCache();
-    PluginManagerWrapper getPluginManagerWrapper();
     PackageManagerHelper getPackageManagerHelper();
+    PluginManagerWrapper getPluginManagerWrapper();
+    VibratorWrapper getVibratorWrapper();
 
     /** Builder for LauncherBaseAppComponent. */
     interface Builder {
diff --git a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
index a5bcd0f..7367f2e 100644
--- a/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
+++ b/src/com/android/launcher3/graphics/GridCustomizationsProvider.java
@@ -36,18 +36,24 @@
 import android.text.TextUtils;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
+
 import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile.GridOption;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel;
 import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.model.BgDataModel;
+import com.android.launcher3.shapes.AppShape;
+import com.android.launcher3.shapes.AppShapesProvider;
 import com.android.launcher3.util.Executors;
 import com.android.launcher3.util.Preconditions;
 import com.android.launcher3.util.RunnableList;
 import com.android.systemui.shared.Flags;
 
 import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
 import java.util.Set;
 import java.util.WeakHashMap;
 import java.util.concurrent.ExecutionException;
@@ -55,18 +61,25 @@
 /**
  * Exposes various launcher grid options and allows the caller to change them.
  * APIs:
- *      /list_options: List the various available grip options, has following columns
- *          name: name of the grid
+ *      /shape_options: List of various available shape options, where each has following fields
+ *          shape_key: key of the shape option
+ *          title: translated title of the shape option
+ *          path: path of the shape, assuming drawn on 100x100 view port
+ *          is_default: true if this shape option is currently set to the system
+ *
+ *      /list_options: List the various available grid options, where each has following fields
+ *          name: key of the grid option
  *          rows: number of rows in the grid
  *          cols: number of columns in the grid
  *          preview_count: number of previews available for this grid option. The preview uri
  *                         looks like /preview/<grid-name>/<preview index starting with 0>
- *          is_default: true if this grid is currently active
+ *          is_default: true if this grid option is currently set to the system
  *
- *     /preview: Opens a file stream for the grid preview
+ *     /get_preview: Open a file stream for the grid preview
  *
- *     /default_grid: Call update to set the current grid, with values
- *          name: name of the grid to apply
+ *     /default_grid: Call update to set the current shape and grid, with values
+ *          shape_key: key of the shape to apply
+ *          name: key of the grid to apply
  */
 public class GridCustomizationsProvider extends ContentProvider {
 
@@ -77,9 +90,16 @@
     private static final String KEY_ROWS = "rows";
     private static final String KEY_COLS = "cols";
     private static final String KEY_PREVIEW_COUNT = "preview_count";
+    // is_default means if a certain option is currently set to the system
     private static final String KEY_IS_DEFAULT = "is_default";
+    private static final String KEY_SHAPE_KEY = "shape_key";
+    private static final String KEY_SHAPE_TITLE = "shape_title";
+    private static final String KEY_PATH = "path";
 
+    // list_options is the key for grid option list
     private static final String KEY_LIST_OPTIONS = "/list_options";
+    private static final String KEY_SHAPE_OPTIONS = "/shape_options";
+    // default_grid is for setting grid and shape to system settings
     private static final String KEY_DEFAULT_GRID = "/default_grid";
 
     private static final String METHOD_GET_PREVIEW = "get_preview";
@@ -95,6 +115,7 @@
     public static final String KEY_GRID_NAME = "grid_name";
 
     private static final int MESSAGE_ID_UPDATE_PREVIEW = 1337;
+    private static final int MESSAGE_ID_UPDATE_SHAPE = 2586;
     private static final int MESSAGE_ID_UPDATE_GRID = 7414;
     private static final int MESSAGE_ID_UPDATE_COLOR = 856;
 
@@ -110,10 +131,37 @@
     @Override
     public Cursor query(Uri uri, String[] projection, String selection,
             String[] selectionArgs, String sortOrder) {
-        switch (uri.getPath()) {
+        Context context = getContext();
+        String path = uri.getPath();
+        if (context == null || path == null) {
+            return null;
+        }
+
+        switch (path) {
+            case KEY_SHAPE_OPTIONS: {
+                if (Flags.newCustomizationPickerUi()) {
+                    MatrixCursor cursor = new MatrixCursor(new String[]{
+                            KEY_SHAPE_KEY, KEY_SHAPE_TITLE, KEY_PATH, KEY_IS_DEFAULT});
+                    List<AppShape> shapes =  AppShapesProvider.INSTANCE.getShapes();
+                    for (int i = 0; i < shapes.size(); i++) {
+                        AppShape shape = shapes.get(i);
+                        cursor.newRow()
+                                .add(KEY_SHAPE_KEY, shape.getKey())
+                                .add(KEY_SHAPE_TITLE, shape.getTitle())
+                                .add(KEY_PATH, shape.getPath())
+                                // TODO (b/348664593): We should fetch the currently-set shape
+                                //  option from the preferences.
+                                .add(KEY_IS_DEFAULT, i == 0);
+                    }
+                    return cursor;
+                } else  {
+                    return null;
+                }
+            }
             case KEY_LIST_OPTIONS: {
                 MatrixCursor cursor = new MatrixCursor(new String[]{
-                        KEY_NAME, KEY_ROWS, KEY_COLS, KEY_PREVIEW_COUNT, KEY_IS_DEFAULT});
+                        KEY_NAME, KEY_GRID_TITLE, KEY_ROWS, KEY_COLS, KEY_PREVIEW_COUNT,
+                        KEY_IS_DEFAULT});
                 InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(getContext());
                 for (GridOption gridOption : idp.parseAllGridOptions(getContext())) {
                     cursor.newRow()
@@ -162,6 +210,14 @@
         }
         switch (path) {
             case KEY_DEFAULT_GRID: {
+                if (Flags.newCustomizationPickerUi()) {
+                    String shapeKey = values.getAsString(KEY_SHAPE_KEY);
+                    Optional<AppShape> optionalShape = AppShapesProvider.INSTANCE.getShapes()
+                            .stream().filter(shape -> shape.getKey().equals(shapeKey)).findFirst();
+                    String pathToSet = optionalShape.map(AppShape::getPath).orElse(null);
+                    // TODO (b/348664593): Apply shapeName to the system. This needs to be a
+                    //  synchronous call.
+                }
                 String gridName = values.getAsString(KEY_NAME);
                 InvariantDeviceProfile idp = InvariantDeviceProfile.INSTANCE.get(context);
                 // Verify that this is a valid grid option
@@ -219,20 +275,30 @@
     }
 
     @Override
-    public Bundle call(String method, String arg, Bundle extras) {
-        if (getContext().checkPermission("android.permission.BIND_WALLPAPER",
+    public Bundle call(@NonNull String method, String arg, Bundle extras) {
+        Context context = getContext();
+        if (context == null) {
+            return null;
+        }
+
+        if (context.checkPermission("android.permission.BIND_WALLPAPER",
                 Binder.getCallingPid(), Binder.getCallingUid())
                 != PackageManager.PERMISSION_GRANTED) {
             return null;
         }
 
-        if (!METHOD_GET_PREVIEW.equals(method)) {
+        if (METHOD_GET_PREVIEW.equals(method)) {
+            return getPreview(extras);
+        } else {
             return null;
         }
-        return getPreview(extras);
     }
 
     private synchronized Bundle getPreview(Bundle request) {
+        Context context = getContext();
+        if (context == null) {
+            return null;
+        }
         RunnableList lifeCycleTracker = new RunnableList();
         try {
             PreviewSurfaceRenderer renderer = new PreviewSurfaceRenderer(
@@ -270,7 +336,9 @@
         public final PreviewSurfaceRenderer renderer;
         public boolean destroyed = false;
 
-        PreviewLifecycleObserver(RunnableList lifeCycleTracker, PreviewSurfaceRenderer renderer) {
+        PreviewLifecycleObserver(
+                RunnableList lifeCycleTracker,
+                PreviewSurfaceRenderer renderer) {
             this.lifeCycleTracker = lifeCycleTracker;
             this.renderer = renderer;
             lifeCycleTracker.add(() -> destroyed = true);
@@ -286,6 +354,17 @@
                 case MESSAGE_ID_UPDATE_PREVIEW:
                     renderer.hideBottomRow(message.getData().getBoolean(KEY_HIDE_BOTTOM_ROW));
                     break;
+                case MESSAGE_ID_UPDATE_SHAPE:
+                    if (Flags.newCustomizationPickerUi()) {
+                        String shapeKey = message.getData().getString(KEY_SHAPE_KEY);
+                        Optional<AppShape> optionalShape = AppShapesProvider.INSTANCE.getShapes()
+                                .stream()
+                                .filter(shape -> shape.getKey().equals(shapeKey))
+                                .findFirst();
+                        String pathToSet = optionalShape.map(AppShape::getPath).orElse(null);
+                        // TODO (b/348664593): Update launcher preview with the given shape
+                    }
+                    break;
                 case MESSAGE_ID_UPDATE_GRID:
                     String gridName = message.getData().getString(KEY_GRID_NAME);
                     if (!TextUtils.isEmpty(gridName)) {
diff --git a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
index 1dd7d45..94c36c0 100644
--- a/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
+++ b/src/com/android/launcher3/graphics/PreviewSurfaceRenderer.java
@@ -305,7 +305,9 @@
                     bgModel,
                     LauncherAppState.getInstance(previewContext).getModel().getModelDelegate(),
                     new BaseLauncherBinder(LauncherAppState.getInstance(previewContext), bgModel,
-                            /* bgAllAppsList= */ null, new Callbacks[0])) {
+                            /* bgAllAppsList= */ null, new Callbacks[0]),
+                    LauncherAppState.getInstance(
+                            previewContext).getModel().getWidgetsFilterDataProvider()) {
 
                 @Override
                 public void run() {
diff --git a/src/com/android/launcher3/model/BaseLauncherBinder.java b/src/com/android/launcher3/model/BaseLauncherBinder.java
index b51f855..c251114 100644
--- a/src/com/android/launcher3/model/BaseLauncherBinder.java
+++ b/src/com/android/launcher3/model/BaseLauncherBinder.java
@@ -24,6 +24,8 @@
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
+import static java.util.Collections.emptyList;
+
 import android.os.Process;
 import android.os.Trace;
 import android.util.Log;
@@ -43,6 +45,7 @@
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
@@ -62,6 +65,7 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.concurrent.Executor;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 /**
@@ -162,9 +166,17 @@
         if (!WIDGETS_ENABLED) {
             return;
         }
+        Map<PackageItemInfo, List<WidgetItem>>
+                widgetsByPackageItem = mBgDataModel.widgetsModel.getWidgetsByPackageItem();
         List<WidgetsListBaseEntry> widgets = new WidgetsListBaseEntriesBuilder(mApp.getContext())
-                .build(mBgDataModel.widgetsModel.getWidgetsByPackageItem());
-        executeCallbacksTask(c -> c.bindAllWidgets(widgets), mUiExecutor);
+                .build(widgetsByPackageItem);
+        Predicate<WidgetItem> filter = mBgDataModel.widgetsModel.getDefaultWidgetsFilter();
+        List<WidgetsListBaseEntry> defaultWidgets =
+                filter != null ? new WidgetsListBaseEntriesBuilder(
+                        mApp.getContext()).build(widgetsByPackageItem,
+                        mBgDataModel.widgetsModel.getDefaultWidgetsFilter()) : emptyList();
+
+        executeCallbacksTask(c -> c.bindAllWidgets(widgets, defaultWidgets), mUiExecutor);
     }
 
     /**
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index 9a9fa5b..b9b1e98 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -537,7 +537,13 @@
         default void bindWidgetsRestored(ArrayList<LauncherAppWidgetInfo> widgets) { }
         default void bindRestoreItemsChange(HashSet<ItemInfo> updates) { }
         default void bindWorkspaceComponentsRemoved(Predicate<ItemInfo> matcher) { }
-        default void bindAllWidgets(List<WidgetsListBaseEntry> widgets) { }
+
+        /**
+         * Binds the app widgets to the providers that share widgets with the UI.
+         */
+        default void bindAllWidgets(@NonNull List<WidgetsListBaseEntry> widgets,
+                @NonNull List<WidgetsListBaseEntry> defaultWidgets) {
+        }
         default void bindSmartspaceWidget() { }
 
         /** Called when workspace has been bound. */
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 06d8b59..a830c96 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -20,6 +20,7 @@
 import static com.android.launcher3.Flags.enableLauncherBrMetricsFixed;
 import static com.android.launcher3.Flags.enableSmartspaceAsAWidget;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.LauncherPrefs.IS_FIRST_LOAD_AFTER_RESTORE;
 import static com.android.launcher3.LauncherPrefs.SHOULD_SHOW_SMARTSPACE;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
@@ -142,6 +143,7 @@
     private final UserManager mUserManager;
     private final UserCache mUserCache;
     private final PackageManagerHelper mPmHelper;
+    private final WidgetsFilterDataProvider mWidgetsFilterDataProvider;
 
     private final InstallSessionHelper mSessionHelper;
     private final IconCache mIconCache;
@@ -158,13 +160,16 @@
     private String mDbName;
 
     public LoaderTask(@NonNull LauncherAppState app, AllAppsList bgAllAppsList, BgDataModel bgModel,
-            ModelDelegate modelDelegate, @NonNull BaseLauncherBinder launcherBinder) {
-        this(app, bgAllAppsList, bgModel, modelDelegate, launcherBinder, new UserManagerState());
+            ModelDelegate modelDelegate, @NonNull BaseLauncherBinder launcherBinder,
+            @NonNull WidgetsFilterDataProvider widgetsFilterDataProvider) {
+        this(app, bgAllAppsList, bgModel, modelDelegate, launcherBinder, widgetsFilterDataProvider,
+                new UserManagerState());
     }
 
     @VisibleForTesting
     LoaderTask(@NonNull LauncherAppState app, AllAppsList bgAllAppsList, BgDataModel bgModel,
             ModelDelegate modelDelegate, @NonNull BaseLauncherBinder launcherBinder,
+            WidgetsFilterDataProvider widgetsFilterDataProvider,
             UserManagerState userManagerState) {
         mApp = app;
         mBgAllAppsList = bgAllAppsList;
@@ -179,6 +184,7 @@
         mIconCache = mApp.getIconCache();
         mUserManagerState = userManagerState;
         mInstallingPkgsCached = null;
+        mWidgetsFilterDataProvider = widgetsFilterDataProvider;
     }
 
     protected synchronized void waitForIdle() {
@@ -330,7 +336,15 @@
             verifyNotStopped();
 
             // fourth step
-            List<CachedObject> allWidgetsList = mBgDataModel.widgetsModel.update(mApp, null);
+            WidgetsModel widgetsModel = mBgDataModel.widgetsModel;
+            if (enableTieredWidgetsByDefaultInPicker()) {
+                // Begin periodic refresh of filters
+                mWidgetsFilterDataProvider.initPeriodicDataRefresh(
+                        mApp.getModel()::onWidgetFiltersLoaded);
+                // And, update model with currently cached data.
+                widgetsModel.updateWidgetFilters(mWidgetsFilterDataProvider);
+            }
+            List<CachedObject> allWidgetsList = widgetsModel.update(mApp, /*packageUser=*/null);
             logASplit("load widgets");
 
             verifyNotStopped();
diff --git a/src/com/android/launcher3/model/ModelTaskController.kt b/src/com/android/launcher3/model/ModelTaskController.kt
index cf2cadc..fc53343 100644
--- a/src/com/android/launcher3/model/ModelTaskController.kt
+++ b/src/com/android/launcher3/model/ModelTaskController.kt
@@ -35,7 +35,7 @@
     val dataModel: BgDataModel,
     val allAppsList: AllAppsList,
     private val model: LauncherModel,
-    private val uiExecutor: Executor
+    private val uiExecutor: Executor,
 ) {
 
     /** Schedules a {@param task} to be executed on the current callbacks. */
@@ -79,10 +79,19 @@
     }
 
     fun bindUpdatedWidgets(dataModel: BgDataModel) {
-        val widgets =
-            WidgetsListBaseEntriesBuilder(app.context)
-                .build(dataModel.widgetsModel.widgetsByPackageItem)
-        scheduleCallbackTask { it.bindAllWidgets(widgets) }
+        val widgetsByPackageItem = dataModel.widgetsModel.widgetsByPackageItem
+        val allWidgets = WidgetsListBaseEntriesBuilder(app.context).build(widgetsByPackageItem)
+
+        val defaultWidgetsFilter = dataModel.widgetsModel.defaultWidgetsFilter
+        val defaultWidgets =
+            if (defaultWidgetsFilter != null) {
+                WidgetsListBaseEntriesBuilder(app.context)
+                    .build(widgetsByPackageItem, defaultWidgetsFilter)
+            } else {
+                emptyList()
+            }
+
+        scheduleCallbackTask { it.bindAllWidgets(allWidgets, defaultWidgets) }
     }
 
     fun deleteAndBindComponentsRemoved(matcher: Predicate<ItemInfo?>, reason: String?) {
@@ -99,7 +108,7 @@
             val packageUserKeyToUidMap =
                 apps.associateBy(
                     keySelector = { PackageUserKey(it.componentName!!.packageName, it.user) },
-                    valueTransform = { it.uid }
+                    valueTransform = { it.uid },
                 )
             scheduleCallbackTask { it.bindAllApplications(apps, flags, packageUserKeyToUidMap) }
         }
diff --git a/src/com/android/launcher3/model/WidgetsFilterDataProvider.kt b/src/com/android/launcher3/model/WidgetsFilterDataProvider.kt
new file mode 100644
index 0000000..0571de3
--- /dev/null
+++ b/src/com/android/launcher3/model/WidgetsFilterDataProvider.kt
@@ -0,0 +1,70 @@
+/*
+ * 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.Context
+import androidx.annotation.WorkerThread
+import com.android.launcher3.R
+import com.android.launcher3.util.ResourceBasedOverride
+import java.util.function.Predicate
+
+/** Helper for the widgets model to load the filters that can be applied to available widgets. */
+open class WidgetsFilterDataProvider(val context: Context) : ResourceBasedOverride {
+    /**
+     * Start regular periodic refresh of widget filtering data starting now (if not started
+     * already).
+     */
+    @WorkerThread
+    open fun initPeriodicDataRefresh(callback: WidgetsFilterLoadedCallback? = null) {
+        // no-op
+    }
+
+    /**
+     * Returns a filter that should be applied to the widget predictions.
+     *
+     * @return null if no filter needs to be applied
+     */
+    @WorkerThread open fun getPredictedWidgetsFilter(): Predicate<WidgetItem>? = null
+
+    /**
+     * Returns a filter that should be applied to the widgets list to see which widgets can be shown
+     * by default.
+     *
+     * @return null if no separate "default" list is supported
+     */
+    @WorkerThread open fun getDefaultWidgetsFilter(): Predicate<WidgetItem>? = null
+
+    /** Called when filter data provider is no longer needed. */
+    open fun destroy() {}
+
+    companion object {
+        /** Returns a new instance of the [WidgetsFilterDataProvider] based on resource override. */
+        fun newInstance(context: Context?): WidgetsFilterDataProvider {
+            return ResourceBasedOverride.Overrides.getObject(
+                WidgetsFilterDataProvider::class.java,
+                context,
+                R.string.widgets_filter_data_provider_class,
+            )
+        }
+    }
+}
+
+/** Interface for the model callback to be invoked when filters are loaded. */
+interface WidgetsFilterLoadedCallback {
+    /** Method called back when widget filters are loaded */
+    fun onWidgetsFilterLoaded()
+}
diff --git a/src/com/android/launcher3/model/WidgetsModel.java b/src/com/android/launcher3/model/WidgetsModel.java
index b450f46..01d4996 100644
--- a/src/com/android/launcher3/model/WidgetsModel.java
+++ b/src/com/android/launcher3/model/WidgetsModel.java
@@ -18,6 +18,8 @@
 import android.util.Log;
 import android.util.Pair;
 
+import androidx.annotation.AnyThread;
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 import androidx.collection.ArrayMap;
 
@@ -65,6 +67,8 @@
 
     /* Map of widgets and shortcuts that are tracked per package. */
     private final Map<PackageItemInfo, List<WidgetItem>> mWidgetsByPackageItem = new HashMap<>();
+    @Nullable private Predicate<WidgetItem> mDefaultWidgetsFilter = null;
+    @Nullable private Predicate<WidgetItem> mPredictedWidgetsFilter = null;
 
     /**
      * Returns all widgets keyed by their component key.
@@ -92,6 +96,37 @@
     }
 
     /**
+     * Returns widget filter that can be applied to {@link WidgetItem}s to check if they can be
+     * shown in the default widgets list.
+     * <p>Returns null if filtering isn't available</p>
+     */
+    @AnyThread
+    public @Nullable Predicate<WidgetItem> getDefaultWidgetsFilter() {
+        return mDefaultWidgetsFilter;
+    }
+
+    /**
+     * Returns widget filter that can be applied to {@link WidgetItem}s to check if they can be
+     * part of widget predictions.
+     * <p>Returns null if filter isn't available</p>
+     */
+    @AnyThread
+    public @Nullable  Predicate<WidgetItem> getPredictedWidgetsFilter() {
+        return mPredictedWidgetsFilter;
+    }
+
+    /**
+     * Updates model with latest filter data in cache.
+     */
+    public void updateWidgetFilters(@NonNull WidgetsFilterDataProvider widgetsFilterDataProvider) {
+        if (!WIDGETS_ENABLED) {
+            return;
+        }
+        mDefaultWidgetsFilter = widgetsFilterDataProvider.getDefaultWidgetsFilter();
+        mPredictedWidgetsFilter = widgetsFilterDataProvider.getPredictedWidgetsFilter();
+    }
+
+    /**
      * @param packageUser If null, all widgets and shortcuts are updated and returned, otherwise
      *                    only widgets and shortcuts associated with the package/user are.
      */
@@ -299,7 +334,7 @@
             if (pInfo == null) {
                 pInfo = new PackageItemInfo(key.mPackageName, key.mWidgetCategory, key.mUser);
                 pInfo.user = key.mUser;
-                mMap.put(key,  pInfo);
+                mMap.put(key, pInfo);
             }
             return pInfo;
         }
diff --git a/src/com/android/launcher3/model/data/WorkspaceItemInfo.java b/src/com/android/launcher3/model/data/WorkspaceItemInfo.java
index f31bf1e..9af61f0 100644
--- a/src/com/android/launcher3/model/data/WorkspaceItemInfo.java
+++ b/src/com/android/launcher3/model/data/WorkspaceItemInfo.java
@@ -189,7 +189,13 @@
         if (TextUtils.isEmpty(label)) {
             label = shortcutInfo.getShortLabel();
         }
-        contentDescription = context.getPackageManager().getUserBadgedLabel(label, user);
+        try {
+            contentDescription = context.getPackageManager().getUserBadgedLabel(label, user);
+        } catch (SecurityException e) {
+            contentDescription = null;
+            Log.e(TAG, "Failed to get content description", e);
+        }
+
         if (shortcutInfo.isEnabled()) {
             runtimeStatusFlags &= ~FLAG_DISABLED_BY_PUBLISHER;
         } else {
diff --git a/src/com/android/launcher3/shapes/AppShape.kt b/src/com/android/launcher3/shapes/AppShape.kt
new file mode 100644
index 0000000..68200a0
--- /dev/null
+++ b/src/com/android/launcher3/shapes/AppShape.kt
@@ -0,0 +1,19 @@
+/*
+ * 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.shapes
+
+class AppShape(val key: String, val title: String, val path: String)
diff --git a/src/com/android/launcher3/shapes/AppShapesProvider.kt b/src/com/android/launcher3/shapes/AppShapesProvider.kt
new file mode 100644
index 0000000..8c2f181
--- /dev/null
+++ b/src/com/android/launcher3/shapes/AppShapesProvider.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.shapes
+
+import com.android.systemui.shared.Flags
+
+object AppShapesProvider {
+
+    val shapes =
+        if (Flags.newCustomizationPickerUi())
+            listOf(
+                AppShape(
+                    "arch",
+                    "arch",
+                    "M100 83.46C100 85.471 100 86.476 99.9 87.321 99.116 93.916 93.916 99.116 87.321 99.9 86.476 100 85.471 100 83.46 100H16.54C14.529 100 13.524 100 12.679 99.9 6.084 99.116.884 93.916.1 87.321 0 86.476 0 85.471 0 83.46L0 50C0 22.386 22.386 0 50 0 77.614 0 100 22.386 100 50V83.46Z",
+                ),
+                AppShape(
+                    "4_sided_cookie",
+                    "4 sided cookie",
+                    "M63.605 3C84.733-6.176 106.176 15.268 97 36.395L95.483 39.888C92.681 46.338 92.681 53.662 95.483 60.112L97 63.605C106.176 84.732 84.733 106.176 63.605 97L60.112 95.483C53.662 92.681 46.338 92.681 39.888 95.483L36.395 97C15.267 106.176-6.176 84.732 3 63.605L4.517 60.112C7.319 53.662 7.319 46.338 4.517 39.888L3 36.395C-6.176 15.268 15.267-6.176 36.395 3L39.888 4.517C46.338 7.319 53.662 7.319 60.112 4.517L63.605 3Z",
+                ),
+                AppShape(
+                    "seven_sided_cookie",
+                    "7 sided cookie",
+                    "M35.209 4.878C36.326 3.895 36.884 3.404 37.397 3.006 44.82-2.742 55.18-2.742 62.603 3.006 63.116 3.404 63.674 3.895 64.791 4.878 65.164 5.207 65.351 5.371 65.539 5.529 68.167 7.734 71.303 9.248 74.663 9.932 74.902 9.981 75.147 10.025 75.637 10.113 77.1 10.375 77.831 10.506 78.461 10.66 87.573 12.893 94.032 21.011 94.176 30.412 94.186 31.062 94.151 31.805 94.08 33.293 94.057 33.791 94.045 34.04 94.039 34.285 93.958 37.72 94.732 41.121 96.293 44.18 96.404 44.399 96.522 44.618 96.759 45.056 97.467 46.366 97.821 47.021 98.093 47.611 102.032 56.143 99.727 66.266 92.484 72.24 91.983 72.653 91.381 73.089 90.177 73.961 89.774 74.254 89.572 74.4 89.377 74.548 86.647 76.626 84.477 79.353 83.063 82.483 82.962 82.707 82.865 82.936 82.671 83.395 82.091 84.766 81.8 85.451 81.51 86.033 77.31 94.44 67.977 98.945 58.801 96.994 58.166 96.859 57.451 96.659 56.019 96.259 55.54 96.125 55.3 96.058 55.063 95.998 51.74 95.154 48.26 95.154 44.937 95.998 44.699 96.058 44.46 96.125 43.981 96.259 42.549 96.659 41.834 96.859 41.199 96.994 32.023 98.945 22.69 94.44 18.49 86.033 18.2 85.451 17.909 84.766 17.329 83.395 17.135 82.936 17.038 82.707 16.937 82.483 15.523 79.353 13.353 76.626 10.623 74.548 10.428 74.4 10.226 74.254 9.823 73.961 8.619 73.089 8.017 72.653 7.516 72.24.273 66.266-2.032 56.143 1.907 47.611 2.179 47.021 2.533 46.366 3.241 45.056 3.478 44.618 3.596 44.399 3.707 44.18 5.268 41.121 6.042 37.72 5.961 34.285 5.955 34.04 5.943 33.791 5.92 33.293 5.849 31.805 5.814 31.062 5.824 30.412 5.968 21.011 12.427 12.893 21.539 10.66 22.169 10.506 22.9 10.375 24.363 10.113 24.853 10.025 25.098 9.981 25.337 9.932 28.697 9.248 31.833 7.734 34.461 5.529 34.649 5.371 34.836 5.207 35.209 4.878Z",
+                ),
+                AppShape(
+                    "sunny",
+                    "sunny",
+                    "M42.846 4.873C46.084-.531 53.916-.531 57.154 4.873L60.796 10.951C62.685 14.103 66.414 15.647 69.978 14.754L76.851 13.032C82.962 11.5 88.5 17.038 86.968 23.149L85.246 30.022C84.353 33.586 85.897 37.315 89.049 39.204L95.127 42.846C100.531 46.084 100.531 53.916 95.127 57.154L89.049 60.796C85.897 62.685 84.353 66.414 85.246 69.978L86.968 76.851C88.5 82.962 82.962 88.5 76.851 86.968L69.978 85.246C66.414 84.353 62.685 85.898 60.796 89.049L57.154 95.127C53.916 100.531 46.084 100.531 42.846 95.127L39.204 89.049C37.315 85.898 33.586 84.353 30.022 85.246L23.149 86.968C17.038 88.5 11.5 82.962 13.032 76.851L14.754 69.978C15.647 66.414 14.103 62.685 10.951 60.796L4.873 57.154C-.531 53.916-.531 46.084 4.873 42.846L10.951 39.204C14.103 37.315 15.647 33.586 14.754 30.022L13.032 23.149C11.5 17.038 17.038 11.5 23.149 13.032L30.022 14.754C33.586 15.647 37.315 14.103 39.204 10.951L42.846 4.873Z",
+                ),
+                AppShape(
+                    "circle",
+                    "circle",
+                    "M99.18 50C99.18 77.162 77.162 99.18 50 99.18 22.838 99.18.82 77.162.82 50 .82 22.839 22.838.82 50 .82 77.162.82 99.18 22.839 99.18 50Z",
+                ),
+                AppShape(
+                    "square",
+                    "square",
+                    "M99.18 53.689C99.18 67.434 99.18 74.306 97.022 79.758 93.897 87.649 87.649 93.897 79.758 97.022 74.306 99.18 67.434 99.18 53.689 99.18H46.311C32.566 99.18 25.694 99.18 20.242 97.022 12.351 93.897 6.103 87.649 2.978 79.758.82 74.306.82 67.434.82 53.689L.82 46.311C.82 32.566.82 25.694 2.978 20.242 6.103 12.351 12.351 6.103 20.242 2.978 25.694.82 32.566.82 46.311.82L53.689.82C67.434.82 74.306.82 79.758 2.978 87.649 6.103 93.897 12.351 97.022 20.242 99.18 25.694 99.18 32.566 99.18 46.311V53.689Z\n",
+                ),
+            )
+        else emptyList()
+}
diff --git a/src/com/android/launcher3/util/SettingsCache.java b/src/com/android/launcher3/util/SettingsCache.java
index 29d5032..8fe6e93 100644
--- a/src/com/android/launcher3/util/SettingsCache.java
+++ b/src/com/android/launcher3/util/SettingsCache.java
@@ -28,6 +28,8 @@
 import android.os.Looper;
 import android.provider.Settings;
 
+import androidx.annotation.UiThread;
+
 import com.android.launcher3.dagger.ApplicationContext;
 import com.android.launcher3.dagger.LauncherAppSingleton;
 import com.android.launcher3.dagger.LauncherBaseAppComponent;
@@ -140,7 +142,9 @@
      * Does not de-dupe if you add same listeners for the same key multiple times.
      * Unregister once complete using {@link #unregister(Uri, OnChangeListener)}
      */
+    @UiThread
     public void register(Uri uri, OnChangeListener changeListener) {
+        Preconditions.assertUIThread();
         if (mListenerMap.containsKey(uri)) {
             mListenerMap.get(uri).add(changeListener);
         } else {
diff --git a/src/com/android/launcher3/util/VibratorWrapper.java b/src/com/android/launcher3/util/VibratorWrapper.java
index adb8f9d..39c9c42 100644
--- a/src/com/android/launcher3/util/VibratorWrapper.java
+++ b/src/com/android/launcher3/util/VibratorWrapper.java
@@ -19,6 +19,7 @@
 import static android.os.VibrationEffect.createPredefined;
 import static android.provider.Settings.System.HAPTIC_FEEDBACK_ENABLED;
 
+import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 
 import android.annotation.SuppressLint;
@@ -31,13 +32,20 @@
 
 import androidx.annotation.VisibleForTesting;
 
+import com.android.launcher3.dagger.ApplicationContext;
+import com.android.launcher3.dagger.LauncherAppSingleton;
+import com.android.launcher3.dagger.LauncherBaseAppComponent;
+
+import javax.inject.Inject;
+
 /**
  * Wrapper around {@link Vibrator} to easily perform haptic feedback where necessary.
  */
-public class VibratorWrapper implements SafeCloseable {
+@LauncherAppSingleton
+public class VibratorWrapper {
 
-    public static final MainThreadInitializedObject<VibratorWrapper> INSTANCE =
-            new MainThreadInitializedObject<>(VibratorWrapper::new);
+    public static final DaggerSingletonObject<VibratorWrapper> INSTANCE =
+            new DaggerSingletonObject<>(LauncherBaseAppComponent::getVibratorWrapper);
 
     public static final AudioAttributes VIBRATION_ATTRS = new AudioAttributes.Builder()
             .setUsage(AudioAttributes.USAGE_ASSISTANCE_SONIFICATION)
@@ -59,38 +67,29 @@
     private final Vibrator mVibrator;
     private final boolean mHasVibrator;
 
-    private final SettingsCache mSettingsCache;
-
     @VisibleForTesting
     final SettingsCache.OnChangeListener mHapticChangeListener =
             isEnabled -> mIsHapticFeedbackEnabled = isEnabled;
 
     private boolean mIsHapticFeedbackEnabled;
 
-    private VibratorWrapper(Context context) {
-        this(context.getSystemService(Vibrator.class), SettingsCache.INSTANCE.get(context));
-    }
+    @Inject
+    public VibratorWrapper(@ApplicationContext Context context, SettingsCache settingsCache,
+            DaggerSingletonTracker tracker) {
 
-    @VisibleForTesting
-    VibratorWrapper(Vibrator vibrator, SettingsCache settingsCache) {
-        mVibrator = vibrator;
+        mVibrator = context.getSystemService(Vibrator.class);
         mHasVibrator = mVibrator.hasVibrator();
-        mSettingsCache = settingsCache;
         if (mHasVibrator) {
-            mSettingsCache.register(HAPTIC_FEEDBACK_URI, mHapticChangeListener);
-            mIsHapticFeedbackEnabled = mSettingsCache.getValue(HAPTIC_FEEDBACK_URI, 0);
+            MAIN_EXECUTOR.execute(
+                    () -> settingsCache.register(HAPTIC_FEEDBACK_URI, mHapticChangeListener));
+            mIsHapticFeedbackEnabled = settingsCache.getValue(HAPTIC_FEEDBACK_URI, 0);
+            tracker.addCloseable(
+                    () -> settingsCache.unregister(HAPTIC_FEEDBACK_URI, mHapticChangeListener));
         } else {
             mIsHapticFeedbackEnabled = false;
         }
     }
 
-    @Override
-    public void close() {
-        if (mHasVibrator) {
-            mSettingsCache.unregister(HAPTIC_FEEDBACK_URI, mHapticChangeListener);
-        }
-    }
-
     /**
      * This should be used to cancel a haptic in case where the haptic shouldn't be vibrating. For
      * example, when no animation is happening but a vibrator happens to be vibrating still.
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
index ff545fe..ae4ff04 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WidgetsModelTest.kt
@@ -43,6 +43,7 @@
 import com.google.common.truth.Truth.assertThat
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
+import java.util.function.Predicate
 import org.junit.Assert.fail
 import org.junit.Before
 import org.junit.Rule
@@ -64,6 +65,7 @@
     @Mock private lateinit var appWidgetManager: AppWidgetManager
     @Mock private lateinit var app: LauncherAppState
     @Mock private lateinit var iconCacheMock: IconCache
+    @Mock private lateinit var widgetsFilterDataProvider: WidgetsFilterDataProvider
 
     private lateinit var context: Context
     private lateinit var idp: InvariantDeviceProfile
@@ -215,6 +217,27 @@
         // No exception
     }
 
+    @Test
+    fun updateWidgetFilters_setsFiltersCorrectly() {
+        val testDefaultWidgetFilter = Predicate<WidgetItem> { w -> w.widgetInfo != null }
+        whenever(widgetsFilterDataProvider.getDefaultWidgetsFilter())
+            .thenReturn(testDefaultWidgetFilter)
+        val testPredicatedWidgetFilter = Predicate<WidgetItem> { w -> w.widgetInfo != null }
+        whenever(widgetsFilterDataProvider.getPredictedWidgetsFilter())
+            .thenReturn(testPredicatedWidgetFilter)
+
+        underTest.updateWidgetFilters(widgetsFilterDataProvider)
+
+        assertThat(underTest.defaultWidgetsFilter).isEqualTo(testDefaultWidgetFilter)
+        assertThat(underTest.predictedWidgetsFilter).isEqualTo(testPredicatedWidgetFilter)
+    }
+
+    @Test
+    fun widgetFilters_nullInitially() {
+        assertThat(underTest.defaultWidgetsFilter).isNull()
+        assertThat(underTest.predictedWidgetsFilter).isNull()
+    }
+
     private fun loadWidgets() {
         val latch = CountDownLatch(1)
         Executors.MODEL_EXECUTOR.execute {
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt
index d321e41..dee98e7 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt
@@ -21,8 +21,8 @@
 import android.os.VibrationEffect.Composition.PRIMITIVE_LOW_TICK
 import android.os.VibrationEffect.Composition.PRIMITIVE_TICK
 import android.os.Vibrator
-import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.launcher3.util.LauncherModelHelper.SandboxModelContext
 import com.android.launcher3.util.VibratorWrapper.HAPTIC_FEEDBACK_URI
 import com.android.launcher3.util.VibratorWrapper.OVERVIEW_HAPTIC
 import com.android.launcher3.util.VibratorWrapper.VIBRATION_ATTRS
@@ -41,25 +41,27 @@
 import org.mockito.kotlin.same
 
 @SmallTest
-@RunWith(AndroidJUnit4::class)
+@RunWith(LauncherMultivalentJUnit::class)
 class VibratorWrapperTest {
 
     @Mock private lateinit var settingsCache: SettingsCache
-    @Mock private lateinit var vibrator: Vibrator
+    private lateinit var vibrator: Vibrator
+    private val context: SandboxModelContext = SandboxModelContext()
     @Captor private lateinit var vibrationEffectCaptor: ArgumentCaptor<VibrationEffect>
-
+    @Mock private lateinit var tracker: DaggerSingletonTracker
     private lateinit var underTest: VibratorWrapper
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
+        vibrator = context.spyService(Vibrator::class.java)
         `when`(settingsCache.getValue(HAPTIC_FEEDBACK_URI, 0)).thenReturn(true)
         `when`(vibrator.hasVibrator()).thenReturn(true)
         `when`(vibrator.areAllPrimitivesSupported(PRIMITIVE_TICK)).thenReturn(true)
         `when`(vibrator.areAllPrimitivesSupported(PRIMITIVE_LOW_TICK)).thenReturn(true)
         `when`(vibrator.getPrimitiveDurations(PRIMITIVE_LOW_TICK)).thenReturn(intArrayOf(10))
 
-        underTest = VibratorWrapper(vibrator, settingsCache)
+        underTest = VibratorWrapper(context, settingsCache, tracker)
     }
 
     @Test
@@ -68,13 +70,6 @@
     }
 
     @Test
-    fun close_unregister_onChangeListener() {
-        underTest.close()
-
-        verify(settingsCache).unregister(HAPTIC_FEEDBACK_URI, underTest.mHapticChangeListener)
-    }
-
-    @Test
     fun vibrate() {
         underTest.vibrate(OVERVIEW_HAPTIC)
 
@@ -117,7 +112,7 @@
     @Test
     fun haptic_feedback_disabled_no_vibrate() {
         `when`(vibrator.hasVibrator()).thenReturn(false)
-        underTest = VibratorWrapper(vibrator, settingsCache)
+        underTest = VibratorWrapper(context, settingsCache, tracker)
 
         underTest.vibrate(OVERVIEW_HAPTIC)
 
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index ef7242f..882061f 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -28,6 +28,7 @@
 import com.android.launcher3.util.UserIconInfo
 import com.google.common.truth.Truth
 import java.util.concurrent.CountDownLatch
+import java.util.function.Predicate
 import junit.framework.Assert.assertEquals
 import org.junit.After
 import org.junit.Before
@@ -76,6 +77,7 @@
     @Mock private lateinit var modelDelegate: ModelDelegate
     @Mock private lateinit var launcherBinder: BaseLauncherBinder
     private lateinit var launcherModel: LauncherModel
+    @Mock private lateinit var widgetsFilterDataProvider: WidgetsFilterDataProvider
     @Mock private lateinit var transaction: LoaderTransaction
     @Mock private lateinit var iconCache: IconCache
     @Mock private lateinit var idleLock: LooperIdleLock
@@ -89,6 +91,7 @@
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
+        setFlagsRule.enableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER)
         launcherModel = mock(LauncherModel::class.java)
         mockitoSession =
             ExtendedMockito.mockitoSession()
@@ -118,6 +121,7 @@
         `when`(launcherBinder.newIdleLock(any())).thenReturn(idleLock)
         `when`(idleLock.awaitLocked(1000)).thenReturn(false)
         `when`(iconCache.updateHandler).thenReturn(iconCacheUpdateHandler)
+        `when`(widgetsFilterDataProvider.getDefaultWidgetsFilter()).thenReturn(Predicate { true })
         context.putObject(UserCache.INSTANCE, userCache)
 
         TestUtil.grantWriteSecurePermission()
@@ -136,17 +140,32 @@
             val mockUserHandles = arrayListOf<UserHandle>(MAIN_HANDLE)
             `when`(userCache.userProfiles).thenReturn(mockUserHandles)
             `when`(userCache.getUserInfo(MAIN_HANDLE)).thenReturn(UserIconInfo(MAIN_HANDLE, 1))
-            LoaderTask(app, bgAllAppsList, this, modelDelegate, launcherBinder)
+            LoaderTask(
+                    app,
+                    bgAllAppsList,
+                    this,
+                    modelDelegate,
+                    launcherBinder,
+                    widgetsFilterDataProvider,
+                )
                 .runSyncOnBackgroundThread()
             Truth.assertThat(workspaceItems.size).isAtLeast(25)
             Truth.assertThat(appWidgets.size).isAtLeast(7)
             Truth.assertThat(collections.size()).isAtLeast(8)
             Truth.assertThat(itemsIdMap.size()).isAtLeast(40)
+            Truth.assertThat(widgetsModel.defaultWidgetsFilter).isNotNull()
         }
 
     @Test
     fun bindsLoadedDataCorrectly() {
-        LoaderTask(app, bgAllAppsList, BgDataModel(), modelDelegate, launcherBinder)
+        LoaderTask(
+                app,
+                bgAllAppsList,
+                BgDataModel(),
+                modelDelegate,
+                launcherBinder,
+                widgetsFilterDataProvider,
+            )
             .runSyncOnBackgroundThread()
 
         verify(launcherBinder).bindWorkspace(true, false)
@@ -155,6 +174,7 @@
         verify(launcherBinder).bindAllApps()
         verify(iconCacheUpdateHandler, times(4)).updateIcons(any(), any<CachingLogic<Any>>(), any())
         verify(launcherBinder).bindDeepShortcuts()
+        verify(widgetsFilterDataProvider).initPeriodicDataRefresh(any())
         verify(launcherBinder).bindWidgets()
         verify(modelDelegate).loadAndBindOtherItems(anyOrNull())
         verify(iconCacheUpdateHandler).finish()
@@ -172,7 +192,15 @@
             `when`(userManagerState?.isUserQuiet(MAIN_HANDLE)).thenReturn(true)
             `when`(userCache.getUserInfo(MAIN_HANDLE)).thenReturn(UserIconInfo(MAIN_HANDLE, 1))
 
-            LoaderTask(app, bgAllAppsList, this, modelDelegate, launcherBinder, userManagerState)
+            LoaderTask(
+                    app,
+                    bgAllAppsList,
+                    this,
+                    modelDelegate,
+                    launcherBinder,
+                    widgetsFilterDataProvider,
+                    userManagerState,
+                )
                 .runSyncOnBackgroundThread()
 
             verify(bgAllAppsList)
@@ -193,7 +221,15 @@
             `when`(userManagerState?.isUserQuiet(MAIN_HANDLE)).thenReturn(true)
             `when`(userCache.getUserInfo(MAIN_HANDLE)).thenReturn(UserIconInfo(MAIN_HANDLE, 3))
 
-            LoaderTask(app, bgAllAppsList, this, modelDelegate, launcherBinder, userManagerState)
+            LoaderTask(
+                    app,
+                    bgAllAppsList,
+                    this,
+                    modelDelegate,
+                    launcherBinder,
+                    widgetsFilterDataProvider,
+                    userManagerState,
+                )
                 .runSyncOnBackgroundThread()
 
             verify(bgAllAppsList)
@@ -232,7 +268,14 @@
         RestoreDbTask.setPending(spyContext)
 
         // When
-        LoaderTask(app, bgAllAppsList, BgDataModel(), modelDelegate, launcherBinder)
+        LoaderTask(
+                app,
+                bgAllAppsList,
+                BgDataModel(),
+                modelDelegate,
+                launcherBinder,
+                widgetsFilterDataProvider,
+            )
             .runSyncOnBackgroundThread()
 
         // Then
@@ -301,7 +344,14 @@
         RestoreDbTask.setPending(spyContext)
 
         // When
-        LoaderTask(app, bgAllAppsList, BgDataModel(), modelDelegate, launcherBinder)
+        LoaderTask(
+                app,
+                bgAllAppsList,
+                BgDataModel(),
+                modelDelegate,
+                launcherBinder,
+                widgetsFilterDataProvider,
+            )
             .runSyncOnBackgroundThread()
 
         // Then
@@ -369,7 +419,14 @@
         Settings.Secure.putInt(spyContext.contentResolver, "launcher_broadcast_installed_apps", 1)
 
         // When
-        LoaderTask(app, bgAllAppsList, BgDataModel(), modelDelegate, launcherBinder)
+        LoaderTask(
+                app,
+                bgAllAppsList,
+                BgDataModel(),
+                modelDelegate,
+                launcherBinder,
+                widgetsFilterDataProvider,
+            )
             .runSyncOnBackgroundThread()
 
         // Then
@@ -404,7 +461,14 @@
         RestoreDbTask.setPending(spyContext)
 
         // When
-        LoaderTask(app, bgAllAppsList, BgDataModel(), modelDelegate, launcherBinder)
+        LoaderTask(
+                app,
+                bgAllAppsList,
+                BgDataModel(),
+                modelDelegate,
+                launcherBinder,
+                widgetsFilterDataProvider,
+            )
             .runSyncOnBackgroundThread()
 
         // Then
diff --git a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
index 1002ca4..b15afc1 100644
--- a/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
+++ b/tests/tapl/com/android/launcher3/tapl/BaseOverview.java
@@ -536,6 +536,10 @@
         int focusedTaskHeight = focusTaskSize.height();
         for (UiObject2 task : taskViews) {
             OverviewTask overviewTask = new OverviewTask(mLauncher, task, this);
+            // Desktop tasks can't be focused tasks, but are the same size.
+            if (overviewTask.isDesktop()) {
+                continue;
+            }
             if (overviewTask.getVisibleHeight() == focusedTaskHeight) {
                 return overviewTask;
             }