diff --git a/quickstep/res/values/config.xml b/quickstep/res/values/config.xml
index 5c80575..f3c9467 100644
--- a/quickstep/res/values/config.xml
+++ b/quickstep/res/values/config.xml
@@ -52,6 +52,11 @@
 
     <integer name="max_depth_blur_radius">23</integer>
 
+    <!-- If predicted widgets from prediction service are less than this number, additional
+    eligible widgets may be added locally by launcher. When set to 0, no widgets will be added
+    locally. -->
+    <integer name="widget_predictions_min_count">6</integer>
+
     <!-- Accessibility actions -->
     <item type="id" name="action_move_to_top_or_left" />
     <item type="id" name="action_move_to_bottom_or_right" />
diff --git a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
index 6916a1d..e160f82 100644
--- a/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
+++ b/quickstep/src/com/android/launcher3/desktop/DesktopAppLaunchTransition.kt
@@ -20,6 +20,7 @@
 import android.animation.AnimatorSet
 import android.animation.ValueAnimator
 import android.content.Context
+import android.graphics.Rect
 import android.os.IBinder
 import android.view.SurfaceControl.Transaction
 import android.view.WindowManager.TRANSIT_OPEN
@@ -31,6 +32,7 @@
 import android.window.TransitionInfo.Change
 import androidx.core.animation.addListener
 import com.android.app.animation.Interpolators
+import com.android.internal.policy.ScreenDecorationsUtils
 import com.android.quickstep.RemoteRunnable
 import com.android.wm.shell.shared.animation.MinimizeAnimator
 import com.android.wm.shell.shared.animation.WindowAnimator
@@ -43,8 +45,19 @@
  * ([android.view.WindowManager.TRANSIT_TO_BACK]) this transition will apply a minimize animation to
  * that window.
  */
-class DesktopAppLaunchTransition(private val context: Context, private val mainExecutor: Executor) :
-    RemoteTransitionStub() {
+class DesktopAppLaunchTransition(
+    private val context: Context,
+    private val mainExecutor: Executor,
+    private val launchType: AppLaunchType,
+) : RemoteTransitionStub() {
+
+    enum class AppLaunchType(
+        val boundsAnimationParams: WindowAnimator.BoundsAnimationParams,
+        val alphaDurationMs: Long,
+    ) {
+        LAUNCH(launchBoundsAnimationDef, /* alphaDurationMs= */ 200L),
+        UNMINIMIZE(unminimizeBoundsAnimationDef, /* alphaDurationMs= */ 100L),
+    }
 
     override fun startAnimation(
         token: IBinder,
@@ -105,18 +118,24 @@
         val boundsAnimator =
             WindowAnimator.createBoundsAnimator(
                 context.resources.displayMetrics,
-                launchBoundsAnimationDef,
+                launchType.boundsAnimationParams,
                 change,
                 transaction,
             )
         val alphaAnimator =
             ValueAnimator.ofFloat(0f, 1f).apply {
-                duration = LAUNCH_ANIM_ALPHA_DURATION_MS
+                duration = launchType.alphaDurationMs
                 interpolator = Interpolators.LINEAR
                 addUpdateListener { animation ->
                     transaction.setAlpha(change.leash, animation.animatedValue as Float).apply()
                 }
             }
+        val clipRect = Rect(change.endAbsBounds).apply { offsetTo(0, 0) }
+        transaction.setCrop(change.leash, clipRect)
+        transaction.setCornerRadius(
+            change.leash,
+            ScreenDecorationsUtils.getWindowCornerRadius(context),
+        )
         return AnimatorSet().apply {
             playTogether(boundsAnimator, alphaAnimator)
             addListener(onEnd = { animation -> onAnimFinish(animation) })
@@ -124,13 +143,18 @@
     }
 
     companion object {
-        private val LAUNCH_CHANGE_MODES = intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT)
-
-        private const val LAUNCH_ANIM_ALPHA_DURATION_MS = 100L
-        private const val MINIMIZE_ANIM_ALPHA_DURATION_MS = 100L
+        val LAUNCH_CHANGE_MODES = intArrayOf(TRANSIT_OPEN, TRANSIT_TO_FRONT)
 
         private val launchBoundsAnimationDef =
             WindowAnimator.BoundsAnimationParams(
+                durationMs = 600,
+                startOffsetYDp = 36f,
+                startScale = 0.95f,
+                interpolator = Interpolators.STANDARD_DECELERATE,
+            )
+
+        private val unminimizeBoundsAnimationDef =
+            WindowAnimator.BoundsAnimationParams(
                 durationMs = 300,
                 startOffsetYDp = 12f,
                 startScale = 0.97f,
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 0395d32..9d9054e 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -16,8 +16,12 @@
 package com.android.launcher3.model;
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 
+import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.toMap;
+
 import android.app.prediction.AppTarget;
 import android.content.ComponentName;
 import android.content.Context;
@@ -26,6 +30,7 @@
 import androidx.annotation.NonNull;
 
 import com.android.launcher3.LauncherModel.ModelUpdateTask;
+import com.android.launcher3.R;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.QuickstepModelDelegate.PredictorState;
 import com.android.launcher3.model.data.ItemInfo;
@@ -34,8 +39,10 @@
 import com.android.launcher3.widget.picker.WidgetRecommendationCategoryProvider;
 
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.Set;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -60,33 +67,72 @@
     @Override
     public void execute(@NonNull ModelTaskController taskController, @NonNull BgDataModel dataModel,
             @NonNull AllAppsList apps) {
+        Predicate<WidgetItem> predictedWidgetsFilter = enableTieredWidgetsByDefaultInPicker()
+                ? dataModel.widgetsModel.getPredictedWidgetsFilter() : null;
         Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
                 widget -> new ComponentKey(widget.providerName, widget.user)).collect(
                 Collectors.toSet());
-        Predicate<WidgetItem> notOnWorkspace = w -> !widgetsInWorkspace.contains(w);
-        Map<ComponentKey, WidgetItem> allWidgets =
-                dataModel.widgetsModel.getWidgetsByComponentKey();
+
+        // Widgets (excluding shortcuts & already added widgets) that belong to apps eligible for
+        // being in predictions.
+        Map<ComponentKey, WidgetItem> allEligibleWidgets =
+                dataModel.widgetsModel.getWidgetsByComponentKey()
+                        .entrySet()
+                        .stream()
+                        .filter(entry -> entry.getValue().widgetInfo != null
+                                && !widgetsInWorkspace.contains(entry.getValue())
+                                && (predictedWidgetsFilter == null
+                                || predictedWidgetsFilter.test(entry.getValue()))
+                        ).collect(toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+        Context context = taskController.getApp().getContext();
 
         List<WidgetItem> servicePredictedItems = new ArrayList<>();
+        List<String> addedWidgetApps = new ArrayList<>();
 
         for (AppTarget app : mTargets) {
             ComponentKey componentKey = new ComponentKey(
                     new ComponentName(app.getPackageName(), app.getClassName()), app.getUser());
-            WidgetItem widget = allWidgets.get(componentKey);
-            if (widget == null) {
+            WidgetItem widget = allEligibleWidgets.get(componentKey);
+            if (widget == null) { // widget not eligible.
                 continue;
             }
             String className = app.getClassName();
             if (!TextUtils.isEmpty(className)) {
-                if (notOnWorkspace.test(widget)) {
-                    servicePredictedItems.add(widget);
-                }
+                servicePredictedItems.add(widget);
+                addedWidgetApps.add(componentKey.componentName.getPackageName());
+            }
+        }
+
+        int minPredictionCount = context.getResources().getInteger(
+                R.integer.widget_predictions_min_count);
+        if (enableTieredWidgetsByDefaultInPicker()
+                && servicePredictedItems.size() < minPredictionCount) {
+            // Eligible apps that aren't already part of predictions.
+            Map<String, List<WidgetItem>> eligibleWidgetsByApp =
+                    allEligibleWidgets.values().stream()
+                            .filter(w -> !addedWidgetApps.contains(
+                                    w.componentName.getPackageName()))
+                            .collect(groupingBy(w -> w.componentName.getPackageName()));
+
+            // Randomize available apps list
+            List<String> appPackages = new ArrayList<>(eligibleWidgetsByApp.keySet());
+            Collections.shuffle(appPackages);
+
+            int widgetsToAdd = minPredictionCount - servicePredictedItems.size();
+            for (String appPackage : appPackages) {
+                if (widgetsToAdd <= 0) break;
+
+                List<WidgetItem> widgetsForApp = eligibleWidgetsByApp.get(appPackage);
+                int index = new Random().nextInt(widgetsForApp.size());
+                // Add a random widget from the app.
+                servicePredictedItems.add(widgetsForApp.get(index));
+                widgetsToAdd--;
             }
         }
 
         List<ItemInfo> items;
         if (enableCategorizedWidgetSuggestions()) {
-            Context context = taskController.getApp().getContext();
             WidgetRecommendationCategoryProvider categoryProvider =
                     WidgetRecommendationCategoryProvider.newInstance(context);
             items = servicePredictedItems.stream()
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
index 7a63f74..390112e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchViewController.java
@@ -36,6 +36,7 @@
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.AnimatorListeners;
 import com.android.launcher3.desktop.DesktopAppLaunchTransition;
+import com.android.launcher3.desktop.DesktopAppLaunchTransition.AppLaunchType;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayContext;
 import com.android.launcher3.taskbar.overlay.TaskbarOverlayDragLayer;
 import com.android.launcher3.util.DisplayController;
@@ -47,7 +48,6 @@
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper;
 import com.android.systemui.shared.system.QuickStepContract;
-import com.android.window.flags.Flags;
 
 import java.io.PrintWriter;
 import java.util.List;
@@ -247,11 +247,13 @@
             return -1;
         }
         RemoteTransition remoteTransition = slideInTransition;
-        if (mOnDesktop && task.task1.isMinimized
-                && Flags.enableDesktopAppLaunchAlttabTransitions()) {
+        if (mOnDesktop
+                && mControllers.taskbarActivityContext.canUnminimizeDesktopTask(task.task1.key.id)
+        ) {
             // This app is being unminimized - use our own transition runner.
             remoteTransition = new RemoteTransition(
-                    new DesktopAppLaunchTransition(context, MAIN_EXECUTOR));
+                    new DesktopAppLaunchTransition(
+                        context, MAIN_EXECUTOR, AppLaunchType.UNMINIMIZE));
         }
         mControllers.taskbarActivityContext.handleGroupTaskLaunch(
                 task,
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 7d8e93c..fbd1b6e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -49,6 +49,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SHORTCUT_HELPER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
+import static com.android.window.flags.Flags.predictiveBackThreeButtonNav;
 import static com.android.wm.shell.Flags.enableBubbleBarInPersistentTaskBar;
 
 import android.animation.Animator;
@@ -71,8 +72,10 @@
 import android.graphics.drawable.RotateDrawable;
 import android.inputmethodservice.InputMethodService;
 import android.os.Handler;
+import android.os.SystemClock;
 import android.util.Property;
 import android.view.Gravity;
+import android.view.KeyEvent;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.View.OnAttachStateChangeListener;
@@ -843,12 +846,43 @@
         buttonView.setImageResource(drawableId);
         buttonView.setContentDescription(parent.getContext().getString(
                 navButtonController.getButtonContentDescription(buttonType)));
-        buttonView.setOnClickListener(view -> navButtonController.onButtonClick(buttonType, view));
-        buttonView.setOnLongClickListener(view ->
-                navButtonController.onButtonLongClick(buttonType, view));
+        if (predictiveBackThreeButtonNav() && buttonType == BUTTON_BACK) {
+            // set up special touch listener for back button to support predictive back
+            setBackButtonTouchListener(buttonView, navButtonController);
+        } else {
+            buttonView.setOnClickListener(view ->
+                    navButtonController.onButtonClick(buttonType, view));
+            buttonView.setOnLongClickListener(view ->
+                    navButtonController.onButtonLongClick(buttonType, view));
+        }
         return buttonView;
     }
 
+    private void setBackButtonTouchListener(View buttonView,
+            TaskbarNavButtonController navButtonController) {
+        buttonView.setOnTouchListener((v, event) -> {
+            if (event.getAction() == MotionEvent.ACTION_MOVE) return false;
+            long time = SystemClock.uptimeMillis();
+            int action = event.getAction();
+            KeyEvent keyEvent = new KeyEvent(time, time,
+                    action == MotionEvent.ACTION_DOWN ? KeyEvent.ACTION_DOWN : KeyEvent.ACTION_UP,
+                    KeyEvent.KEYCODE_BACK, 0);
+            if (event.getAction() == MotionEvent.ACTION_CANCEL) {
+                keyEvent.cancel();
+            }
+            navButtonController.executeBack(keyEvent);
+
+            if (action == MotionEvent.ACTION_UP) {
+                buttonView.performClick();
+            }
+            return false;
+        });
+        buttonView.setOnLongClickListener((view) ->  {
+            navButtonController.onButtonLongClick(BUTTON_BACK, view);
+            return false;
+        });
+    }
+
     private ImageView addButton(ViewGroup parent, @IdRes int id, @LayoutRes int layoutId) {
         ImageView buttonView = (ImageView) mContext.getLayoutInflater()
                 .inflate(layoutId, parent, false);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index e22de06..82acc0c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -59,6 +59,7 @@
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.hardware.display.DisplayManager;
+import android.os.Bundle;
 import android.os.IRemoteCallback;
 import android.os.Process;
 import android.os.Trace;
@@ -72,6 +73,7 @@
 import android.view.WindowManager;
 import android.widget.FrameLayout;
 import android.widget.Toast;
+import android.window.DesktopModeFlags;
 import android.window.RemoteTransition;
 
 import androidx.annotation.NonNull;
@@ -89,6 +91,8 @@
 import com.android.launcher3.anim.AnimatorPlaybackController;
 import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.desktop.DesktopAppLaunchTransition;
+import com.android.launcher3.desktop.DesktopAppLaunchTransition.AppLaunchType;
 import com.android.launcher3.dot.DotInfo;
 import com.android.launcher3.folder.Folder;
 import com.android.launcher3.folder.FolderIcon;
@@ -282,6 +286,7 @@
         BubbleBarController.onTaskbarRecreated();
         if (BubbleBarController.isBubbleBarEnabled()
                 && !mDeviceProfile.isPhone
+                && !mDeviceProfile.isVerticalBarLayout()
                 && bubbleBarView != null
         ) {
             Optional<BubbleStashedHandleViewController> bubbleHandleController = Optional.empty();
@@ -859,6 +864,33 @@
         return makeDefaultActivityOptions(SPLASH_SCREEN_STYLE_UNDEFINED);
     }
 
+    private ActivityOptionsWrapper getActivityLaunchDesktopOptions(ItemInfo info) {
+        if (!DesktopModeFlags.ENABLE_DESKTOP_APP_LAUNCH_TRANSITIONS.isTrue()) {
+            return null;
+        }
+        if (!areDesktopTasksVisible()) {
+            return null;
+        }
+        BubbleTextView.RunningAppState appState =
+                mControllers.taskbarRecentAppsController.getDesktopItemState(info);
+        AppLaunchType launchType = null;
+        switch (appState) {
+            case RUNNING:
+                return null;
+            case MINIMIZED:
+                launchType = AppLaunchType.UNMINIMIZE;
+                break;
+            case NOT_RUNNING:
+                launchType = AppLaunchType.LAUNCH;
+                break;
+        }
+        ActivityOptions options = ActivityOptions.makeRemoteTransition(
+                new RemoteTransition(
+                        new DesktopAppLaunchTransition(
+                                /* context= */ this, getMainExecutor(), launchType)));
+        return new ActivityOptionsWrapper(options, new RunnableList());
+    }
+
     /**
      * Sets a new data-source for this taskbar instance
      */
@@ -1219,10 +1251,10 @@
         mControllers.keyboardQuickSwitchController.closeQuickSwitchView(false);
 
         if (tag instanceof GroupTask groupTask) {
-            handleGroupTaskLaunch(
-                    groupTask,
-                    /* remoteTransition= */ null,
-                    areDesktopTasksVisible());
+            RemoteTransition remoteTransition =
+                    (areDesktopTasksVisible() && canUnminimizeDesktopTask(groupTask.task1.key.id))
+                            ? createUnminimizeRemoteTransition() : null;
+            handleGroupTaskLaunch(groupTask, remoteTransition, areDesktopTasksVisible());
             mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(true);
         } else if (tag instanceof FolderInfo) {
             // Tapping an expandable folder icon on Taskbar
@@ -1240,9 +1272,11 @@
                 mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(true);
             }
         } else if (tag instanceof TaskItemInfo info) {
+            RemoteTransition remoteTransition = canUnminimizeDesktopTask(info.getTaskId())
+                    ? createUnminimizeRemoteTransition() : null;
             UI_HELPER_EXECUTOR.execute(() ->
                     SystemUiProxy.INSTANCE.get(this).showDesktopApp(
-                            info.getTaskId(), /* remoteTransition= */ null));
+                            info.getTaskId(), remoteTransition));
             mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(
                     /* stash= */ true);
         } else if (tag instanceof WorkspaceItemInfo) {
@@ -1361,8 +1395,7 @@
             return;
         }
         if (onDesktop) {
-            boolean useRemoteTransition = task.task1.isMinimized
-                    && com.android.window.flags.Flags.enableDesktopAppLaunchAlttabTransitions();
+            boolean useRemoteTransition = canUnminimizeDesktopTask(task.task1.key.id);
             UI_HELPER_EXECUTOR.execute(() -> {
                 if (onStartCallback != null) {
                     onStartCallback.run();
@@ -1389,6 +1422,20 @@
         mControllers.uiController.launchSplitTasks(task, remoteTransition);
     }
 
+    /** Returns whether the given task is minimized and can be unminimized. */
+    public boolean canUnminimizeDesktopTask(int taskId) {
+        BubbleTextView.RunningAppState runningAppState =
+                mControllers.taskbarRecentAppsController.getRunningAppState(taskId);
+        return runningAppState == BubbleTextView.RunningAppState.MINIMIZED
+                && DesktopModeFlags.ENABLE_DESKTOP_APP_LAUNCH_ALTTAB_TRANSITIONS.isTrue();
+    }
+
+    private RemoteTransition createUnminimizeRemoteTransition() {
+        return new RemoteTransition(
+                new DesktopAppLaunchTransition(
+                        this, getMainExecutor(), AppLaunchType.UNMINIMIZE));
+    }
+
     /**
      * Runs when the user taps a Taskbar icon in TaskbarActivityContext (Overview or inside an app),
      * and calls the appropriate method to animate and launch.
@@ -1487,25 +1534,31 @@
                 .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK);
         try {
             TestLogging.recordEvent(TestProtocol.SEQUENCE_MAIN, "start: taskbarAppIcon");
-            if (info.user.equals(Process.myUserHandle())) {
-                // TODO(b/216683257): Use startActivityForResult for search results that require it.
-                if (taskInRecents != null) {
-                    // Re launch instance from recents
-                    ActivityOptionsWrapper opts = getActivityLaunchOptions(null, info);
-                    opts.options.setLaunchDisplayId(
-                            getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
-                    if (ActivityManagerWrapper.getInstance()
-                            .startActivityFromRecents(taskInRecents.key, opts.options)) {
-                        mControllers.uiController.getRecentsView()
-                                .addSideTaskLaunchCallback(opts.onEndCallback);
-                        return;
-                    }
-                }
-                startActivity(intent);
-            } else {
+            if (!info.user.equals(Process.myUserHandle())) {
+                // TODO b/376819104: support Desktop launch animations for apps in managed profiles
                 getSystemService(LauncherApps.class).startMainActivity(
                         intent.getComponent(), info.user, intent.getSourceBounds(), null);
+                return;
             }
+            // TODO(b/216683257): Use startActivityForResult for search results that require it.
+            if (taskInRecents != null) {
+                // Re launch instance from recents
+                ActivityOptionsWrapper opts = getActivityLaunchOptions(null, info);
+                opts.options.setLaunchDisplayId(
+                        getDisplay() == null ? DEFAULT_DISPLAY : getDisplay().getDisplayId());
+                if (ActivityManagerWrapper.getInstance()
+                        .startActivityFromRecents(taskInRecents.key, opts.options)) {
+                    mControllers.uiController.getRecentsView()
+                            .addSideTaskLaunchCallback(opts.onEndCallback);
+                    return;
+                }
+            }
+            ActivityOptionsWrapper opts = null;
+            if (areDesktopTasksVisible()) {
+                opts = getActivityLaunchDesktopOptions(info);
+            }
+            Bundle optionsBundle = opts == null ? null : opts.options.toBundle();
+            startActivity(intent, optionsBundle);
         } catch (NullPointerException | ActivityNotFoundException | SecurityException e) {
             Toast.makeText(this, R.string.activity_not_found, Toast.LENGTH_SHORT)
                     .show();
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
index bdefea6..c20617d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarModelCallbacks.java
@@ -42,7 +42,6 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.function.Predicate;
 
 /**
@@ -196,26 +195,21 @@
         final TaskbarRecentAppsController recentAppsController =
                 mControllers.taskbarRecentAppsController;
         hotseatItemInfos = recentAppsController.updateHotseatItemInfos(hotseatItemInfos);
-        Set<Integer> runningTaskIds = recentAppsController.getRunningTaskIds();
-        Set<Integer> minimizedTaskIds = recentAppsController.getMinimizedTaskIds();
 
         if (mDeferUpdatesForSUW) {
             ItemInfo[] finalHotseatItemInfos = hotseatItemInfos;
             mDeferredUpdates = () ->
                     commitHotseatItemUpdates(finalHotseatItemInfos,
-                            recentAppsController.getShownTasks(), runningTaskIds,
-                            minimizedTaskIds);
+                            recentAppsController.getShownTasks());
         } else {
-            commitHotseatItemUpdates(hotseatItemInfos,
-                    recentAppsController.getShownTasks(), runningTaskIds, minimizedTaskIds);
+            commitHotseatItemUpdates(hotseatItemInfos, recentAppsController.getShownTasks());
         }
     }
 
-    private void commitHotseatItemUpdates(ItemInfo[] hotseatItemInfos, List<GroupTask> recentTasks,
-            Set<Integer> runningTaskIds, Set<Integer> minimizedTaskIds) {
+    private void commitHotseatItemUpdates(
+            ItemInfo[] hotseatItemInfos, List<GroupTask> recentTasks) {
         mContainer.updateHotseatItems(hotseatItemInfos, recentTasks);
-        mControllers.taskbarViewController.updateIconViewsRunningStates(
-                runningTaskIds, minimizedTaskIds);
+        mControllers.taskbarViewController.updateIconViewsRunningStates();
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
index 8947914..0f9ede9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.taskbar;
 
+import static android.view.MotionEvent.ACTION_UP;
+
 import static com.android.internal.app.AssistUtils.INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS;
 import static com.android.internal.app.AssistUtils.INVOCATION_TYPE_KEY;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASKBAR_A11Y_BUTTON_LONGPRESS;
@@ -31,12 +33,14 @@
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_HOME_KEY;
 import static com.android.systemui.shared.system.ActivityManagerWrapper.CLOSE_SYSTEM_WINDOWS_REASON_RECENTS;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
+import static com.android.window.flags.Flags.predictiveBackThreeButtonNav;
 
 import android.content.Context;
 import android.os.Bundle;
 import android.os.Handler;
 import android.util.Log;
 import android.view.HapticFeedbackConstants;
+import android.view.KeyEvent;
 import android.view.View;
 import android.view.inputmethod.Flags;
 
@@ -141,10 +145,7 @@
         view.performHapticFeedback(HapticFeedbackConstants.VIRTUAL_KEY);
         switch (buttonType) {
             case BUTTON_BACK:
-                logEvent(LAUNCHER_TASKBAR_BACK_BUTTON_TAP);
-                mContextualEduStatsManager.updateEduStats(/* isTrackpadGesture= */ false,
-                        GestureType.BACK);
-                executeBack();
+                executeBack(/* keyEvent */ null);
                 break;
             case BUTTON_HOME:
                 logEvent(LAUNCHER_TASKBAR_HOME_BUTTON_TAP);
@@ -182,7 +183,9 @@
 
         // Provide the same haptic feedback that the system offers for long press.
         // The haptic feedback from long pressing on the home button is handled by circle to search.
-        if (buttonType != BUTTON_HOME) {
+        // There are no haptics for long pressing the back button if predictive back is enabled
+        if (buttonType != BUTTON_HOME
+                && (!predictiveBackThreeButtonNav() || buttonType != BUTTON_BACK)) {
             view.performHapticFeedback(HapticFeedbackConstants.LONG_PRESS);
         }
         switch (buttonType) {
@@ -320,8 +323,13 @@
         mCallbacks.onToggleOverview();
     }
 
-    private void executeBack() {
-        mSystemUiProxy.onBackPressed();
+    void executeBack(@Nullable KeyEvent keyEvent) {
+        if (keyEvent == null || (keyEvent.getAction() == ACTION_UP && !keyEvent.isCanceled())) {
+            logEvent(LAUNCHER_TASKBAR_BACK_BUTTON_TAP);
+            mContextualEduStatsManager.updateEduStats(/* isTrackpadGesture= */ false,
+                    GestureType.BACK);
+        }
+        mSystemUiProxy.onBackEvent(keyEvent);
     }
 
     private void onImeSwitcherPress() {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
index 3e8b615..3d57de4 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarRecentAppsController.kt
@@ -18,6 +18,7 @@
 import android.content.Context
 import android.window.DesktopModeFlags
 import androidx.annotation.VisibleForTesting
+import com.android.launcher3.BubbleTextView.RunningAppState
 import com.android.launcher3.Flags.enableRecentsInTaskbar
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.TaskItemInfo
@@ -72,6 +73,43 @@
     var shownTasks: List<GroupTask> = emptyList()
         private set
 
+    /**
+     * Returns the state of the most active Desktop task represented by the given [ItemInfo].
+     *
+     * If there are several tasks represented by the same [ItemInfo] we return the most active one,
+     * i.e. we return [DesktopAppState.RUNNING] over [DesktopAppState.MINIMIZED], and
+     * [DesktopAppState.MINIMIZED] over [DesktopAppState.NOT_RUNNING].
+     */
+    fun getDesktopItemState(itemInfo: ItemInfo?): RunningAppState {
+        val packageName = itemInfo?.getTargetPackage() ?: return RunningAppState.NOT_RUNNING
+        return getDesktopAppState(packageName, itemInfo.user.identifier)
+    }
+
+    private fun getDesktopAppState(packageName: String, userId: Int): RunningAppState {
+        val tasks = desktopTask?.tasks ?: return RunningAppState.NOT_RUNNING
+        val appTasks =
+            tasks.filter { task ->
+                packageName == task.key.packageName && task.key.userId == userId
+            }
+        if (appTasks.find { getRunningAppState(it.key.id) == RunningAppState.RUNNING } != null) {
+            return RunningAppState.RUNNING
+        }
+        if (appTasks.find { getRunningAppState(it.key.id) == RunningAppState.MINIMIZED } != null) {
+            return RunningAppState.MINIMIZED
+        }
+        return RunningAppState.NOT_RUNNING
+    }
+
+    /** Get the [RunningAppState] for the given task. */
+    fun getRunningAppState(taskId: Int): RunningAppState {
+        return when (taskId) {
+            in minimizedTaskIds -> RunningAppState.MINIMIZED
+            in runningTaskIds -> RunningAppState.RUNNING
+            else -> RunningAppState.NOT_RUNNING
+        }
+    }
+
+    @VisibleForTesting
     val runningTaskIds: Set<Int>
         /**
          * Returns the task IDs of apps that should be indicated as "running" to the user.
@@ -88,6 +126,7 @@
             return tasks.map { task -> task.key.id }.toSet()
         }
 
+    @VisibleForTesting
     val minimizedTaskIds: Set<Int>
         /**
          * Returns the task IDs for the tasks that should be indicated as "minimized" to the user.
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java
index bf086b4..4a7e4f0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java
@@ -155,7 +155,7 @@
     }
 
     private void onClick() {
-        SystemUiProxy.INSTANCE.get(mActivity).onBackPressed();
+        SystemUiProxy.INSTANCE.get(mActivity).onBackEvent(null);
     }
 
     @Override
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index 87e19be..494c472 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -623,12 +623,10 @@
      * Updates which icons are marked as running or minimized given the Sets of currently running
      * and minimized tasks.
      */
-    public void updateIconViewsRunningStates(Set<Integer> runningTaskIds,
-            Set<Integer> minimizedTaskIds) {
+    public void updateIconViewsRunningStates() {
         for (View iconView : getIconViews()) {
             if (iconView instanceof BubbleTextView btv) {
-                btv.updateRunningState(
-                        getRunningAppState(btv, runningTaskIds, minimizedTaskIds));
+                btv.updateRunningState(getRunningAppState(btv));
             }
         }
     }
@@ -651,26 +649,15 @@
         return pinnedAppsWithTasks;
     }
 
-    private BubbleTextView.RunningAppState getRunningAppState(
-            BubbleTextView btv,
-            Set<Integer> runningTaskIds,
-            Set<Integer> minimizedTaskIds) {
+    private BubbleTextView.RunningAppState getRunningAppState(BubbleTextView btv) {
         Object tag = btv.getTag();
         if (tag instanceof TaskItemInfo itemInfo) {
-            if (minimizedTaskIds.contains(itemInfo.getTaskId())) {
-                return BubbleTextView.RunningAppState.MINIMIZED;
-            }
-            if (runningTaskIds.contains(itemInfo.getTaskId())) {
-                return BubbleTextView.RunningAppState.RUNNING;
-            }
+            return mControllers.taskbarRecentAppsController.getRunningAppState(
+                    itemInfo.getTaskId());
         }
         if (tag instanceof GroupTask groupTask && !groupTask.hasMultipleTasks()) {
-            if (minimizedTaskIds.contains(groupTask.task1.key.id)) {
-                return BubbleTextView.RunningAppState.MINIMIZED;
-            }
-            if (runningTaskIds.contains(groupTask.task1.key.id)) {
-                return BubbleTextView.RunningAppState.RUNNING;
-            }
+            return mControllers.taskbarRecentAppsController.getRunningAppState(
+                    groupTask.task1.key.id);
         }
         return BubbleTextView.RunningAppState.NOT_RUNNING;
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index 30e4e47..334ba6e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -297,9 +297,7 @@
                 Log.e(TAG, "Could not instantiate BubbleBarBubble for " + bubbleInfos.get(i));
                 continue;
             }
-            addBubbleInternally(bubble,  /* showAppBadge = */
-                    mBubbleBarViewController.isExpanded() || i == 0,
-                    /* isExpanding = */ false,  /* suppressAnimation = */ true);
+            addBubbleInternally(bubble, /* isExpanding= */ false, /* suppressAnimation= */ true);
         }
     }
 
@@ -390,8 +388,7 @@
             for (int i = update.currentBubbles.size() - 1; i >= 0; i--) {
                 BubbleBarBubble bubble = update.currentBubbles.get(i);
                 if (bubble != null) {
-                    addBubbleInternally(bubble, /* showAppBadge = */ !isCollapsed || i == 0,
-                            isExpanding, suppressAnimation);
+                    addBubbleInternally(bubble, isExpanding, suppressAnimation);
                     if (isCollapsed) {
                         // If we're collapsed, the most recently added bubble will be selected.
                         bubbleToSelect = bubble;
@@ -563,10 +560,8 @@
         }
     }
 
-    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);
+    private void addBubbleInternally(BubbleBarBubble bubble, boolean isExpanding,
+            boolean suppressAnimation) {
         mBubbles.put(bubble.getKey(), bubble);
         mBubbleBarViewController.addBubble(bubble, isExpanding, suppressAnimation);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
index d91d10a..c0a76a8 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarView.java
@@ -791,6 +791,7 @@
         updateLayoutParams();
         updateBubbleAccessibilityStates();
         updateContentDescription();
+        updateDotsAndBadgesIfCollapsed();
     }
 
     /** Removes the given bubble from the bubble bar. */
@@ -856,7 +857,7 @@
         updateBubbleAccessibilityStates();
         updateContentDescription();
         mDismissedByDragBubbleView = null;
-        updateNotificationDotsIfCollapsed();
+        updateDotsAndBadgesIfCollapsed();
     }
 
     /**
@@ -886,17 +887,23 @@
         return childViews;
     }
 
-    private void updateNotificationDotsIfCollapsed() {
+    private void updateDotsAndBadgesIfCollapsed() {
         if (isExpanded()) {
             return;
         }
         for (int i = 0; i < getChildCount(); i++) {
             BubbleView bubbleView = (BubbleView) getChildAt(i);
-            // when we're collapsed, the first bubble should show the dot if it has it. the rest of
-            // the bubbles should hide their dots.
-            if (i == 0 && bubbleView.hasUnseenContent()) {
-                bubbleView.showDotIfNeeded(/* animate= */ true);
+            // when we're collapsed, the first bubble should show the badge and the dot if it has
+            // it. the rest of the bubbles should hide their badges and dots.
+            if (i == 0) {
+                bubbleView.showBadge();
+                if (bubbleView.hasUnseenContent()) {
+                    bubbleView.showDotIfNeeded(/* animate= */ true);
+                } else {
+                    bubbleView.hideDot();
+                }
             } else {
+                bubbleView.hideBadge();
                 bubbleView.hideDot();
             }
         }
@@ -1100,7 +1107,7 @@
             }
             updateBubblesLayoutProperties(mBubbleBarLocation);
             updateContentDescription();
-            updateNotificationDotsIfCollapsed();
+            updateDotsAndBadgesIfCollapsed();
         }
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
index 114edf4..0ea4222 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleView.java
@@ -49,6 +49,8 @@
 public class BubbleView extends ConstraintLayout {
 
     public static final int DEFAULT_PATH_SIZE = 100;
+    /** Duration for animating the scale of the dot and badge. */
+    private static final int SCALE_ANIMATION_DURATION_MS = 200;
 
     private final ImageView mBubbleIcon;
     private final ImageView mAppIcon;
@@ -316,12 +318,37 @@
     }
 
     void setBadgeScale(float fraction) {
-        if (mAppIcon.getVisibility() == VISIBLE) {
+        if (hasBadge()) {
             mAppIcon.setScaleX(fraction);
             mAppIcon.setScaleY(fraction);
         }
     }
 
+    void showBadge() {
+        animateBadgeScale(1);
+    }
+
+    void hideBadge() {
+        animateBadgeScale(0);
+    }
+
+    private boolean hasBadge() {
+        return mAppIcon.getVisibility() == VISIBLE;
+    }
+
+    private void animateBadgeScale(float scale) {
+        if (!hasBadge()) {
+            return;
+        }
+        mAppIcon.clearAnimation();
+        mAppIcon.animate()
+                .setDuration(SCALE_ANIMATION_DURATION_MS)
+                .setInterpolator(Interpolators.FAST_OUT_SLOW_IN)
+                .scaleX(scale)
+                .scaleY(scale)
+                .start();
+    }
+
     /** Suppresses or un-suppresses drawing the dot due to an update for this bubble. */
     public void suppressDotForBubbleUpdate(boolean suppress) {
         mDotSuppressedForBubbleUpdate = suppress;
@@ -409,7 +436,7 @@
 
         clearAnimation();
         animate()
-                .setDuration(200)
+                .setDuration(SCALE_ANIMATION_DURATION_MS)
                 .setInterpolator(Interpolators.FAST_OUT_SLOW_IN)
                 .setUpdateListener((valueAnimator) -> {
                     float fraction = valueAnimator.getAnimatedFraction();
diff --git a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
index 7eb34a5..79cb748 100644
--- a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
@@ -35,6 +35,7 @@
 
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.DeviceProfile;
+import com.android.launcher3.Flags;
 import com.android.launcher3.taskbar.TaskbarActivityContext;
 import com.android.launcher3.taskbar.TaskbarControllers;
 import com.android.systemui.shared.system.TaskStackChangeListener;
@@ -216,6 +217,13 @@
         @Override
         protected void handleClose(boolean animate) {
             if (!mIsOpen) return;
+            if (Flags.taskbarOverflow()) {
+                // Mark the view closed before attempting to remove it, so the drag layer does not
+                // schedule another call to close. Needed for taskbar overflow in case the KQS
+                // view shown for taskbar overflow needs to be reshown - delayed close call would
+                // would result in reshown KQS view getting hidden.
+                mIsOpen = false;
+            }
             mTaskbarContext.getDragLayer().removeView(this);
             Optional.ofNullable(mOverlayContext).ifPresent(c -> {
                 if (canCloseWindow()) {
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 73e22bb..c618422 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -43,6 +43,7 @@
 import android.os.UserHandle;
 import android.util.Log;
 import android.view.IRemoteAnimationRunner;
+import android.view.KeyEvent;
 import android.view.MotionEvent;
 import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
@@ -210,10 +211,10 @@
     }
 
     @Override
-    public void onBackPressed() {
+    public void onBackEvent(KeyEvent backEvent) {
         if (mSystemUiProxy != null) {
             try {
-                mSystemUiProxy.onBackPressed();
+                mSystemUiProxy.onBackEvent(backEvent);
             } catch (RemoteException e) {
                 Log.w(TAG, "Failed call onBackPressed", e);
             }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
index 253d921..4b04dba 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
@@ -115,7 +115,7 @@
     @Test
     public void testPressBack() {
         mNavButtonController.onButtonClick(BUTTON_BACK, mockView);
-        verify(mockSystemUiProxy, times(1)).onBackPressed();
+        verify(mockSystemUiProxy, times(1)).onBackEvent(null);
     }
 
     @Test
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
index 3912051..436dfd3 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarScrimViewControllerTest.kt
@@ -17,6 +17,7 @@
 package com.android.launcher3.taskbar
 
 import android.animation.AnimatorTestRule
+import android.view.KeyEvent
 import android.view.View.GONE
 import android.view.View.VISIBLE
 import androidx.test.platform.app.InstrumentationRegistry.getInstrumentation
@@ -47,8 +48,8 @@
         TaskbarWindowSandboxContext.create { builder ->
             builder.bindSystemUiProxy(
                 object : SystemUiProxy(this) {
-                    override fun onBackPressed() {
-                        super.onBackPressed()
+                    override fun onBackEvent(backEvent: KeyEvent?) {
+                        super.onBackEvent(backEvent)
                         backPressed = true
                     }
                 }
diff --git a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
index 7b57c81..c53c177 100644
--- a/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
+++ b/quickstep/tests/src/com/android/launcher3/model/WidgetsPredicationUpdateTaskTest.java
@@ -33,6 +33,7 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.spy;
 
 import android.app.prediction.AppTarget;
 import android.app.prediction.AppTargetId;
@@ -42,6 +43,8 @@
 import android.content.pm.ApplicationInfo;
 import android.content.pm.LauncherApps;
 import android.os.UserHandle;
+import android.platform.test.annotations.DisableFlags;
+import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
 import android.text.TextUtils;
 
@@ -62,9 +65,13 @@
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
 
 import java.util.Arrays;
 import java.util.List;
+import java.util.Set;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 
 @SmallTest
@@ -72,6 +79,9 @@
 public final class WidgetsPredicationUpdateTaskTest {
 
     @Rule
+    public final MockitoRule mocks = MockitoJUnit.rule();
+
+    @Rule
     public final SetFlagsRule mSetFlagsRule = new SetFlagsRule();
 
     private AppWidgetProviderInfo mApp1Provider1;
@@ -145,6 +155,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER) // Flag off
     public void widgetsRecommendationRan_shouldOnlyReturnNotAddedWidgetsInAppPredictionOrder() {
         // Run on model executor so that no other task runs in the middle.
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
@@ -184,6 +195,7 @@
     }
 
     @Test
+    @DisableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER) // Flag off
     public void widgetsRecommendationRan_shouldReturnEmptyWidgetsWhenEmpty() {
         runOnExecutorSync(MODEL_EXECUTOR, () -> {
 
@@ -213,6 +225,50 @@
         });
     }
 
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_TIERED_WIDGETS_BY_DEFAULT_IN_PICKER)
+    public void widgetsRecommendationRan_keepsWidgetsNotOnWorkspace_addsWidgetsFromEligibleApps() {
+        runOnExecutorSync(MODEL_EXECUTOR, () -> {
+            WidgetsFilterDataProvider spiedFilterProvider = spy(
+                    mModelHelper.getModel().getWidgetsFilterDataProvider());
+            doAnswer(i -> new Predicate<WidgetItem>() {
+                @Override
+                public boolean test(WidgetItem widgetItem) {
+                    // app5's widget is already on workspace, but, app2 is not.
+                    // And app4's second widget is also not on workspace.
+                    return Set.of("app5", "app2", "app4").contains(
+                            widgetItem.componentName.getPackageName());
+                }
+            }).when(spiedFilterProvider).getPredictedWidgetsFilter();
+            mModelHelper.getBgDataModel().widgetsModel.updateWidgetFilters(spiedFilterProvider);
+            // App5's widget that's already on workspace.
+            AppTarget widget1 = new AppTarget(new AppTargetId("app5"), "app5", "provider1",
+                    mUserHandle);
+            // App4's widget eligible and not on workspace.
+            AppTarget widget2 = new AppTarget(new AppTargetId("app4"), "app4", "provider2",
+                    mUserHandle);
+
+            mCallback.mRecommendedWidgets = null;
+            mModelHelper.getModel().enqueueModelUpdateTask(
+                    newWidgetsPredicationTask(List.of(widget1, widget2)));
+            runOnExecutorSync(MAIN_EXECUTOR, () -> {
+            });
+
+            List<PendingAddWidgetInfo> recommendedWidgets = mCallback.mRecommendedWidgets.items
+                    .stream()
+                    .map(itemInfo -> (PendingAddWidgetInfo) itemInfo)
+                    .collect(Collectors.toList());
+            assertThat(recommendedWidgets).hasSize(2);
+            List<ComponentName> componentNames = recommendedWidgets.stream().map(
+                    w -> w.componentName).toList();
+            assertThat(componentNames).containsExactly(
+                    // Locally added, not on workspace, eligible app per filter
+                    mApp2Provider1.provider,
+                    // From prediction service, not on workspace, eligible app per filter
+                    mApp4Provider2.provider);
+        });
+    }
+
     private void assertWidgetInfo(
             LauncherAppWidgetProviderInfo actual, AppWidgetProviderInfo expected) {
         assertThat(actual.provider).isEqualTo(expected.provider);
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
index b67bc5a..066ddc0 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarRecentAppsControllerTest.kt
@@ -26,11 +26,13 @@
 import android.platform.test.rule.TestWatcher
 import android.testing.AndroidTestingRunner
 import com.android.internal.R
+import com.android.launcher3.BubbleTextView.RunningAppState
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
 import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION
 import com.android.launcher3.model.data.AppInfo
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.TaskItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.quickstep.RecentsModel
 import com.android.quickstep.RecentsModel.RecentTasksChangedListener
 import com.android.quickstep.TaskIconCache
@@ -77,7 +79,9 @@
     private var taskListChangeId: Int = 1
 
     private lateinit var recentAppsController: TaskbarRecentAppsController
-    private lateinit var userHandle: UserHandle
+    private lateinit var myUserHandle: UserHandle
+    private val USER_HANDLE_1 = UserHandle.of(1)
+    private val USER_HANDLE_2 = UserHandle.of(2)
 
     private var canShowRunningAndRecentAppsAtInit = true
     private var recentTasksChangedListener: RecentTasksChangedListener? = null
@@ -85,7 +89,7 @@
     @Before
     fun setUp() {
         super.setup()
-        userHandle = Process.myUserHandle()
+        myUserHandle = Process.myUserHandle()
 
         // Set desktop mode supported
         whenever(mockContext.getResources()).thenReturn(mockResources)
@@ -148,6 +152,115 @@
     }
 
     @Test
+    fun getDesktopItemState_nullItemInfo_returnsNotRunning() {
+        setInDesktopMode(true)
+        assertThat(recentAppsController.getDesktopItemState(/* itemInfo= */ null))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_noItemPackage_returnsNotRunning() {
+        setInDesktopMode(true)
+        assertThat(recentAppsController.getDesktopItemState(ItemInfo()))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_noMatchingTasks_returnsNotRunning() {
+        setInDesktopMode(true)
+        val itemInfo = createItemInfo("package")
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_matchingVisibleTask_returnsVisible() {
+        setInDesktopMode(true)
+        val visibleTask = createTask(id = 1, "visiblePackage", isVisible = true)
+        updateRecentTasks(runningTasks = listOf(visibleTask), recentTaskPackages = emptyList())
+        val itemInfo = createItemInfo("visiblePackage")
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_matchingMinimizedTask_returnsMinimized() {
+        setInDesktopMode(true)
+        val minimizedTask = createTask(id = 1, "minimizedPackage", isVisible = false)
+        updateRecentTasks(runningTasks = listOf(minimizedTask), recentTaskPackages = emptyList())
+        val itemInfo = createItemInfo("minimizedPackage")
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.MINIMIZED)
+    }
+
+    @Test
+    fun getDesktopItemState_matchingMinimizedAndRunningTask_returnsVisible() {
+        setInDesktopMode(true)
+        updateRecentTasks(
+            runningTasks =
+                listOf(
+                    createTask(id = 1, "package", isVisible = false),
+                    createTask(id = 2, "package", isVisible = true),
+                ),
+            recentTaskPackages = emptyList(),
+        )
+        val itemInfo = createItemInfo("package")
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.RUNNING)
+    }
+
+    @Test
+    fun getDesktopItemState_noMatchingUserId_returnsNotRunning() {
+        setInDesktopMode(true)
+        updateRecentTasks(
+            runningTasks =
+                listOf(
+                    createTask(id = 1, "package", isVisible = false, USER_HANDLE_1),
+                    createTask(id = 2, "package", isVisible = true, USER_HANDLE_1),
+                ),
+            recentTaskPackages = emptyList(),
+        )
+        val itemInfo = createItemInfo("package", USER_HANDLE_2)
+
+        assertThat(recentAppsController.getDesktopItemState(itemInfo))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getRunningAppState_taskNotRunningOrMinimized_returnsNotRunning() {
+        setInDesktopMode(true)
+        updateRecentTasks(runningTasks = emptyList(), recentTaskPackages = emptyList())
+
+        assertThat(recentAppsController.getRunningAppState(taskId = 1))
+            .isEqualTo(RunningAppState.NOT_RUNNING)
+    }
+
+    @Test
+    fun getRunningAppState_taskNotVisible_returnsMinimized() {
+        setInDesktopMode(true)
+        val task1 = createTask(id = 1, packageName = RUNNING_APP_PACKAGE_1, isVisible = false)
+        val task2 = createTask(id = 2, packageName = RUNNING_APP_PACKAGE_1, isVisible = true)
+        updateRecentTasks(runningTasks = listOf(task1, task2), recentTaskPackages = emptyList())
+
+        assertThat(recentAppsController.getRunningAppState(taskId = 1))
+            .isEqualTo(RunningAppState.MINIMIZED)
+    }
+
+    @Test
+    fun getRunningAppState_taskVisible_returnsRunning() {
+        setInDesktopMode(true)
+        val task1 = createTask(id = 1, packageName = RUNNING_APP_PACKAGE_1, isVisible = false)
+        val task2 = createTask(id = 2, packageName = RUNNING_APP_PACKAGE_1, isVisible = true)
+        updateRecentTasks(runningTasks = listOf(task1, task2), recentTaskPackages = emptyList())
+
+        assertThat(recentAppsController.getRunningAppState(taskId = 2))
+            .isEqualTo(RunningAppState.RUNNING)
+    }
+
+    @Test
     fun updateHotseatItemInfos_cantShowRunning_inDesktopMode_returnsAllHotseatItems() {
         recentAppsController.canShowRunningApps = false
         setInDesktopMode(true)
@@ -782,7 +895,13 @@
     private fun createTestAppInfo(
         packageName: String = "testPackageName",
         className: String = "testClassName",
-    ) = AppInfo(ComponentName(packageName, className), className /* title */, userHandle, Intent())
+    ) =
+        AppInfo(
+            ComponentName(packageName, className),
+            className /* title */,
+            myUserHandle,
+            Intent(),
+        )
 
     private fun createRecentTasksFromPackageNames(packageNames: List<String>): List<GroupTask> {
         return packageNames.map { packageName ->
@@ -801,14 +920,19 @@
         }
     }
 
-    private fun createTask(id: Int, packageName: String, isVisible: Boolean = true): Task {
+    private fun createTask(
+        id: Int,
+        packageName: String,
+        isVisible: Boolean = true,
+        localUserHandle: UserHandle? = null,
+    ): Task {
         return Task(
                 Task.TaskKey(
                     id,
                     WINDOWING_MODE_FREEFORM,
                     Intent().apply { `package` = packageName },
                     ComponentName(packageName, "TestActivity"),
-                    userHandle.identifier,
+                    localUserHandle?.identifier ?: myUserHandle.identifier,
                     0,
                 )
             )
@@ -820,6 +944,16 @@
             .thenReturn(inDesktopMode)
     }
 
+    private fun createItemInfo(
+        packageName: String,
+        userHandle: UserHandle = myUserHandle,
+    ): ItemInfo {
+        val appInfo = AppInfo()
+        appInfo.intent = Intent().setComponent(ComponentName(packageName, "className"))
+        appInfo.user = userHandle
+        return WorkspaceItemInfo(appInfo)
+    }
+
     private val GroupTask.packageNames: List<String>
         get() = tasks.map { task -> task.key.packageName }
 
diff --git a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
index 44c23ba..6a7b6f8 100644
--- a/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
+++ b/quickstep/tests/src/com/android/quickstep/AbstractQuickStepTest.java
@@ -26,6 +26,7 @@
 import com.android.launcher3.tapl.LaunchedAppState;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.uioverrides.QuickstepLauncher;
+import com.android.launcher3.util.TestUtil;
 import com.android.quickstep.views.RecentsView;
 
 import org.junit.rules.RuleChain;
@@ -56,7 +57,7 @@
     protected void assertTestActivityIsRunning(int activityNumber, String message) {
         assertTrue(message, mDevice.wait(
                 Until.hasObject(By.pkg(getAppPackageName()).text("TestActivity" + activityNumber)),
-                DEFAULT_UI_TIMEOUT));
+                TestUtil.DEFAULT_UI_TIMEOUT));
     }
 
     protected LaunchedAppState getAndAssertLaunchedApp() {
diff --git a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
index 1f11c14..aa105f9 100644
--- a/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/FallbackRecentsTest.java
@@ -22,9 +22,7 @@
 import static com.android.launcher3.tapl.LauncherInstrumentation.WAIT_TIME_MS;
 import static com.android.launcher3.tapl.TestHelpers.getHomeIntentInPackage;
 import static com.android.launcher3.tapl.TestHelpers.getLauncherInMyProcess;
-import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_ACTIVITY_TIMEOUT;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_BROADCAST_TIMEOUT_SECS;
-import static com.android.launcher3.ui.AbstractLauncherUiTest.DEFAULT_UI_TIMEOUT;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.resolveSystemApp;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.startAppFast;
 import static com.android.launcher3.ui.AbstractLauncherUiTest.startTestActivity;
@@ -56,6 +54,7 @@
 import com.android.launcher3.tapl.TestHelpers;
 import com.android.launcher3.testcomponent.TestCommandReceiver;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ExtendedLongPressTimeoutRule;
 import com.android.launcher3.util.rule.FailureWatcher;
@@ -214,7 +213,7 @@
             }
             result[0] = f.apply(activity);
             return true;
-        }).get(), DEFAULT_UI_TIMEOUT, mLauncher);
+        }).get(), mLauncher);
         return (T) result[0];
     }
 
@@ -244,7 +243,7 @@
 
         Wait.atMost("Recents activity didn't stop",
                 () -> getFromRecents(recents -> !recents.isStarted()),
-                DEFAULT_UI_TIMEOUT, mLauncher);
+                mLauncher);
     }
 
     @Test
@@ -254,7 +253,8 @@
         startTestActivity(2);
         waitForRecentsActivityStop();
         Wait.atMost("Expected three apps in the task list",
-                () -> mLauncher.getRecentTasks().size() >= 3, DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+                () -> mLauncher.getRecentTasks().size() >= 3,
+                mLauncher);
 
         checkTestLauncher();
         BaseOverview overview = mLauncher.getLaunchedAppState().switchToOverview();
@@ -282,7 +282,7 @@
         assertNotNull("OverviewTask.open returned null", task.open());
         assertTrue("Test activity didn't open from Overview", TestHelpers.wait(Until.hasObject(
                 By.pkg(getAppPackageName()).text("TestActivity2")),
-                DEFAULT_UI_TIMEOUT));
+                TestUtil.DEFAULT_UI_TIMEOUT));
 
 
         // Test dismissing a task.
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 4459ed6..77f4c05 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -57,8 +57,6 @@
 
     static final String TAG = "QuickStepOnOffRule";
 
-    public static final int WAIT_TIME_MS = 10000;
-
     public enum Mode {
         THREE_BUTTON, ZERO_BUTTON, ALL
     }
@@ -179,12 +177,13 @@
         }
 
         Wait.atMost("Couldn't switch to " + overlayPackage,
-                () -> launcher.getNavigationModel() == expectedMode, WAIT_TIME_MS, launcher);
+                () -> launcher.getNavigationModel() == expectedMode,
+                launcher);
 
         Wait.atMost(() -> "Switching nav mode: "
                         + launcher.getNavigationModeMismatchError(false),
                 () -> launcher.getNavigationModeMismatchError(false) == null,
-                WAIT_TIME_MS, launcher);
+                launcher);
         AbstractLauncherUiTest.checkDetectedLeaks(launcher, false);
         return true;
     }
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt b/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt
index 120a89b..f58c84e 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsOverviewDesktop.kt
@@ -26,6 +26,7 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape
 import com.android.launcher3.uioverrides.QuickstepLauncher
+import com.android.launcher3.util.TestUtil
 import com.google.common.truth.Truth.assertWithMessage
 import org.junit.Before
 import org.junit.Test
@@ -172,7 +173,7 @@
             .that(
                 mDevice.wait(
                     Until.hasObject(By.pkg(getAppPackageName()).text("TestActivity$index")),
-                    DEFAULT_UI_TIMEOUT,
+                    TestUtil.DEFAULT_UI_TIMEOUT,
                 )
             )
             .isTrue()
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
index 5ff2af7..f1fe2d2 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsQuickstep.java
@@ -47,6 +47,7 @@
 import com.android.launcher3.tapl.SelectModeButtons;
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.Wait;
 import com.android.launcher3.util.rule.ScreenRecordRule.ScreenRecord;
 import com.android.launcher3.util.rule.TestStabilityRule;
@@ -145,7 +146,7 @@
         assertNotNull("OverviewTask.open returned null", task.open());
         assertTrue("Test activity didn't open from Overview", mDevice.wait(Until.hasObject(
                         By.pkg(getAppPackageName()).text("TestActivity2")),
-                DEFAULT_UI_TIMEOUT));
+                TestUtil.DEFAULT_UI_TIMEOUT));
         executeOnLauncher(launcher -> assertTrue(
                 "Launcher activity is the top activity; expecting another activity to be the top "
                         + "one",
@@ -448,7 +449,7 @@
                 mDevice.wait(Until.hasObject(By.pkg(getAppPackageName()).text(
                                 mLauncher.isGridOnlyOverviewEnabled() ? "TestActivity12"
                                         : "TestActivity13")),
-                        DEFAULT_UI_TIMEOUT));
+                        TestUtil.DEFAULT_UI_TIMEOUT));
 
         // Scroll the task offscreen as it is now first
         overview = mLauncher.goHome().switchToOverview();
@@ -563,7 +564,7 @@
             mLauncher.getDevice().setOrientationLeft();
             startTestActivity(7);
             Wait.atMost("Device should not be in natural orientation",
-                    () -> !mDevice.isNaturalOrientation(), DEFAULT_UI_TIMEOUT, mLauncher);
+                    () -> !mDevice.isNaturalOrientation(), mLauncher);
             mLauncher.goHome();
         } finally {
             mLauncher.setExpectedRotationCheckEnabled(true);
diff --git a/res/drawable/ic_more_horiz_24.xml b/res/drawable/ic_more_horiz_24.xml
new file mode 100644
index 0000000..d46827c
--- /dev/null
+++ b/res/drawable/ic_more_horiz_24.xml
@@ -0,0 +1,26 @@
+<!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<vector xmlns:android="http://schemas.android.com/apk/res/android"
+    android:width="24dp"
+    android:height="24dp"
+    android:tint="#000000"
+    android:viewportHeight="24"
+    android:viewportWidth="24">
+    <path
+        android:fillColor="@android:color/white"
+        android:pathData="M6,10c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2zM18,10c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2zM12,10c-1.1,0 -2,0.9 -2,2s0.9,2 2,2 2,-0.9 2,-2 -0.9,-2 -2,-2z" />
+</vector>
diff --git a/res/drawable/widgets_list_expand_button_background.xml b/res/drawable/widgets_list_expand_button_background.xml
new file mode 100644
index 0000000..068b26d
--- /dev/null
+++ b/res/drawable/widgets_list_expand_button_background.xml
@@ -0,0 +1,25 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+<inset xmlns:android="http://schemas.android.com/apk/res/android">
+    <ripple android:color="?android:attr/colorControlHighlight">
+        <item>
+            <shape android:shape="rectangle">
+                <corners android:radius="50dp" />
+                <solid android:color="?attr/widgetPickerExpandButtonBackgroundColor" />
+            </shape>
+        </item>
+    </ripple>
+</inset>
\ No newline at end of file
diff --git a/res/layout/widgets_list_expand_button.xml b/res/layout/widgets_list_expand_button.xml
new file mode 100644
index 0000000..17c19ac
--- /dev/null
+++ b/res/layout/widgets_list_expand_button.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+  ~ Copyright (C) 2024 The Android Open Source Project
+  ~
+  ~ Licensed under the Apache License, Version 2.0 (the "License");
+  ~ you may not use this file except in compliance with the License.
+  ~ You may obtain a copy of the License at
+  ~
+  ~      http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing, software
+  ~ distributed under the License is distributed on an "AS IS" BASIS,
+  ~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  ~ See the License for the specific language governing permissions and
+  ~ limitations under the License.
+  -->
+
+<Button xmlns:android="http://schemas.android.com/apk/res/android"
+    style="@style/Button.Rounded.Colored"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:layout_marginTop="@dimen/widgets_list_expand_button_top_margin"
+    android:background="@drawable/widgets_list_expand_button_background"
+    android:drawablePadding="@dimen/widgets_list_expand_button_drawable_padding"
+    android:drawableStart="@drawable/ic_more_horiz_24"
+    android:drawableTint="?attr/widgetPickerExpandButtonTextColor"
+    android:maxLines="1"
+    android:minHeight="48dp"
+    android:paddingEnd="@dimen/widgets_list_expand_button_end_padding"
+    android:paddingStart="@dimen/widgets_list_expand_button_start_padding"
+    android:paddingVertical="@dimen/widgets_list_expand_button_vertical_padding"
+    android:text="@string/widgets_list_expand_button_label"
+    android:contentDescription="@string/widgets_list_expand_button_content_description"
+    android:textColor="?attr/widgetPickerExpandButtonTextColor" />
\ No newline at end of file
diff --git a/res/layout/widgets_two_pane_sheet.xml b/res/layout/widgets_two_pane_sheet.xml
index 8235875..5dc1b47 100644
--- a/res/layout/widgets_two_pane_sheet.xml
+++ b/res/layout/widgets_two_pane_sheet.xml
@@ -70,7 +70,7 @@
                 android:layout_height="match_parent"
                 android:clipChildren="false"
                 android:clipToPadding="false"
-                android:paddingBottom="24dp"
+                android:paddingBottom="8dp"
                 android:layout_gravity="start"
                 android:layout_weight="0.33">
                 <TextView
@@ -84,14 +84,6 @@
                     android:layout_width="@dimen/fastscroll_width"
                     android:layout_height="match_parent"
                     android:layout_marginEnd="@dimen/fastscroll_end_margin" />
-
-                <com.android.launcher3.widget.picker.WidgetsRecyclerView
-                    android:id="@+id/search_widgets_list_view"
-                    android:layout_width="match_parent"
-                    android:layout_height="wrap_content"
-                    android:clipToPadding="false"
-                    android:paddingHorizontal="@dimen/widget_list_horizontal_margin_two_pane"
-                    android:visibility="gone" />
             </FrameLayout>
 
             <FrameLayout
diff --git a/res/layout/widgets_two_pane_sheet_paged_view.xml b/res/layout/widgets_two_pane_sheet_paged_view.xml
index 33a50b0..0528d3b 100644
--- a/res/layout/widgets_two_pane_sheet_paged_view.xml
+++ b/res/layout/widgets_two_pane_sheet_paged_view.xml
@@ -143,5 +143,13 @@
                 android:clipToPadding="false" />
 
         </com.android.launcher3.widget.picker.WidgetPagedView>
+
+        <com.android.launcher3.widget.picker.WidgetsRecyclerView
+            android:id="@+id/search_widgets_list_view"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:clipToPadding="false"
+            android:paddingHorizontal="@dimen/widget_list_horizontal_margin_two_pane"
+            android:visibility="gone" />
     </LinearLayout>
 </merge>
diff --git a/res/layout/widgets_two_pane_sheet_recyclerview.xml b/res/layout/widgets_two_pane_sheet_recyclerview.xml
index 94f141b..45a9ac0 100644
--- a/res/layout/widgets_two_pane_sheet_recyclerview.xml
+++ b/res/layout/widgets_two_pane_sheet_recyclerview.xml
@@ -85,5 +85,13 @@
             android:layout_height="match_parent"
             android:layout_marginHorizontal="@dimen/widget_list_horizontal_margin_two_pane"
             android:clipToPadding="false" />
+
+        <com.android.launcher3.widget.picker.WidgetsRecyclerView
+            android:id="@+id/search_widgets_list_view"
+            android:layout_width="match_parent"
+            android:layout_height="match_parent"
+            android:clipToPadding="false"
+            android:paddingHorizontal="@dimen/widget_list_horizontal_margin_two_pane"
+            android:visibility="gone" />
     </LinearLayout>
 </merge>
\ No newline at end of file
diff --git a/res/values/attrs.xml b/res/values/attrs.xml
index 4dddb9a..8a805c3 100644
--- a/res/values/attrs.xml
+++ b/res/values/attrs.xml
@@ -126,6 +126,8 @@
     <attr name="widgetPickerCollapseHandleColor" format="color"/>
     <attr name="widgetPickerAddButtonBackgroundColor" format="color"/>
     <attr name="widgetPickerAddButtonTextColor" format="color"/>
+    <attr name="widgetPickerExpandButtonBackgroundColor" format="color"/>
+    <attr name="widgetPickerExpandButtonTextColor" format="color"/>
     <attr name="widgetCellTitleColor" format="color" />
     <attr name="widgetCellSubtitleColor" format="color" />
 
diff --git a/res/values/colors.xml b/res/values/colors.xml
index fa1626e..4549b86 100644
--- a/res/values/colors.xml
+++ b/res/values/colors.xml
@@ -117,6 +117,12 @@
     <color name="widget_picker_collapse_handle_color_light">#C4C7C5</color>
     <color name="widget_picker_add_button_background_color_light">#0B57D0</color>
     <color name="widget_picker_add_button_text_color_light">#0B57D0</color>
+    <color name="widget_picker_expand_button_background_color_light">
+        @color/widget_picker_secondary_surface_color_light
+    </color>
+    <color name="widget_picker_expand_button_text_color_light">
+        @color/widget_picker_header_app_title_color_light
+    </color>
     <color name="widget_cell_title_color_light">@color/system_on_surface_light</color>
     <color name="widget_cell_subtitle_color_light">@color/system_on_surface_variant_light</color>
 
@@ -138,6 +144,12 @@
     <color name="widget_picker_collapse_handle_color_dark">#444746</color>
     <color name="widget_picker_add_button_background_color_dark">#062E6F</color>
     <color name="widget_picker_add_button_text_color_dark">#FFFFFF</color>
+    <color name="widget_picker_expand_button_background_color_dark">
+        @color/widget_picker_secondary_surface_color_dark
+    </color>
+    <color name="widget_picker_expand_button_text_color_dark">
+        @color/widget_picker_header_app_title_color_dark
+    </color>
     <color name="widget_cell_title_color_dark">@color/system_on_surface_dark</color>
     <color name="widget_cell_subtitle_color_dark">@color/system_on_surface_variant_dark</color>
 
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index d4773c3..d1dde3f 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -230,6 +230,12 @@
 
     <dimen name="widget_list_top_bottom_corner_radius">28dp</dimen>
     <dimen name="widget_list_content_corner_radius">4dp</dimen>
+    <!-- Button that expands the widget apps list in the widget picker. -->
+    <dimen name="widgets_list_expand_button_drawable_padding">8dp</dimen>
+    <dimen name="widgets_list_expand_button_start_padding">16dp</dimen>
+    <dimen name="widgets_list_expand_button_end_padding">24dp</dimen>
+    <dimen name="widgets_list_expand_button_vertical_padding">16dp</dimen>
+    <dimen name="widgets_list_expand_button_top_margin">14dp</dimen>
 
     <dimen name="widget_list_header_view_vertical_padding">20dp</dimen>
     <dimen name="widget_list_entry_spacing">2dp</dimen>
diff --git a/res/values/id.xml b/res/values/id.xml
index 28496b5..67692d8 100644
--- a/res/values/id.xml
+++ b/res/values/id.xml
@@ -19,6 +19,7 @@
     <item type="id" name="view_type_widgets_space" />
     <item type="id" name="view_type_widgets_list" />
     <item type="id" name="view_type_widgets_header" />
+    <item type="id" name="view_type_widgets_list_expand" />
 
     <!-- Accessibility actions -->
     <item type="id" name="action_remove" />
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 746bd7e..f5af339 100644
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -157,6 +157,12 @@
     <!-- Accessibility content description for the button that adds a widget to the home screen. The
          placeholder text is the widget name. [CHAR_LIMIT=none] -->
     <string name="widget_add_button_content_description">Add <xliff:g id="widget_name" example="Calendar month view">%1$s</xliff:g> widget</string>
+    <!-- Text on the button that enables users to expand the widgets list to see all widget apps besides the default ones displayed. [CHAR_LIMIT=15] -->
+    <string name="widgets_list_expand_button_label">Show all</string>
+    <!-- Accessibility content description for the button that enables users to expand the widgets list to see all widget apps besides the default ones displayed. [CHAR_LIMIT=none]  -->
+    <string name="widgets_list_expand_button_content_description">Show all widgets</string>
+    <!-- Accessibility announcement to indicate to the users that widgets list is now expanded -->
+    <string name="widgets_list_expanded">Showing all widgets</string>
 
     <!-- Text on an educational tip on widget informing users that they can change widget settings.
          [CHAR_LIMIT=NONE] -->
diff --git a/res/values/styles.xml b/res/values/styles.xml
index 728c523..6d3579b 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -275,6 +275,12 @@
             @color/widget_picker_add_button_background_color_light</item>
         <item name="widgetPickerAddButtonTextColor">
             @color/widget_picker_add_button_text_color_light</item>
+        <item name="widgetPickerExpandButtonBackgroundColor">
+            @color/widget_picker_expand_button_background_color_light
+        </item>
+        <item name="widgetPickerExpandButtonTextColor">
+            @color/widget_picker_expand_button_text_color_light
+        </item>
         <item name="widgetCellTitleColor">
             @color/widget_cell_title_color_light</item>
         <item name="widgetCellSubtitleColor">
@@ -316,6 +322,12 @@
             @color/widget_picker_add_button_background_color_dark</item>
         <item name="widgetPickerAddButtonTextColor">
             @color/widget_picker_add_button_text_color_dark</item>
+        <item name="widgetPickerExpandButtonBackgroundColor">
+            @color/widget_picker_expand_button_background_color_dark
+        </item>
+        <item name="widgetPickerExpandButtonTextColor">
+            @color/widget_picker_expand_button_text_color_dark
+        </item>
         <item name="widgetCellTitleColor">
             @color/widget_cell_title_color_dark</item>
         <item name="widgetCellSubtitleColor">
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index 7112a1b..d8a898d 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -132,6 +132,7 @@
     public int iconBitmapSize;
     public int fillResIconDpi;
     public @DeviceType int deviceType;
+    public Info displayInfo;
 
     public PointF[] minCellSize;
 
@@ -267,7 +268,7 @@
         @DeviceType int defaultDeviceType = defaultInfo.getDeviceType();
         DisplayOption defaultDisplayOption = invDistWeightedInterpolate(
                 defaultInfo,
-                getPredefinedDeviceProfiles(context, gridName, defaultDeviceType,
+                getPredefinedDeviceProfiles(context, gridName, defaultInfo,
                         /*allowDisabledGrid=*/false),
                 defaultDeviceType);
 
@@ -276,7 +277,7 @@
         @DeviceType int deviceType = myInfo.getDeviceType();
         DisplayOption myDisplayOption = invDistWeightedInterpolate(
                 myInfo,
-                getPredefinedDeviceProfiles(context, gridName, deviceType,
+                getPredefinedDeviceProfiles(context, gridName, myInfo,
                         /*allowDisabledGrid=*/false),
                 deviceType);
 
@@ -294,7 +295,7 @@
         System.arraycopy(defaultDisplayOption.borderSpaces, 0, result.borderSpaces, 0,
                 COUNT_SIZES);
 
-        initGrid(context, myInfo, result, deviceType);
+        initGrid(context, myInfo, result);
     }
 
     @Override
@@ -308,11 +309,9 @@
 
     private String initGrid(Context context, String gridName) {
         Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
-        @DeviceType int deviceType = displayInfo.getDeviceType();
 
-        List<DisplayOption> allOptions =
-                getPredefinedDeviceProfiles(context, gridName, deviceType,
-                        RestoreDbTask.isPending(context));
+        List<DisplayOption> allOptions = getPredefinedDeviceProfiles(context, gridName,
+                displayInfo, RestoreDbTask.isPending(context));
 
         // Filter out options that don't have the same number of columns as the grid
         DeviceGridState deviceGridState = new DeviceGridState(context);
@@ -322,8 +321,9 @@
         DisplayOption displayOption =
                 invDistWeightedInterpolate(displayInfo, allOptionsFilteredByColCount.isEmpty()
                         ? new ArrayList<>(allOptions)
-                        : new ArrayList<>(allOptionsFilteredByColCount), deviceType);
-        initGrid(context, displayInfo, displayOption, deviceType);
+                        : new ArrayList<>(allOptionsFilteredByColCount),
+                        displayInfo.getDeviceType());
+        initGrid(context, displayInfo, displayOption);
         return displayOption.grid.name;
     }
 
@@ -347,8 +347,7 @@
         return new InvariantDeviceProfile().initGrid(context, null);
     }
 
-    private void initGrid(Context context, Info displayInfo, DisplayOption displayOption,
-            @DeviceType int deviceType) {
+    private void initGrid(Context context, Info displayInfo, DisplayOption displayOption) {
         DisplayMetrics metrics = context.getResources().getDisplayMetrics();
         GridOption closestProfile = displayOption.grid;
         numRows = closestProfile.numRows;
@@ -381,7 +380,8 @@
         allAppsCellSpecsTwoPanelId = closestProfile.mAllAppsCellSpecsTwoPanelId;
         numAllAppsRowsForCellHeightCalculation =
                 closestProfile.mNumAllAppsRowsForCellHeightCalculation;
-        this.deviceType = deviceType;
+        this.deviceType = displayInfo.getDeviceType();
+        this.displayInfo = displayInfo;
 
         inlineNavButtonsEndSpacing = closestProfile.inlineNavButtonsEndSpacing;
 
@@ -510,8 +510,8 @@
         }
     }
 
-    private static List<DisplayOption> getPredefinedDeviceProfiles(Context context,
-            String gridName, @DeviceType int deviceType, boolean allowDisabledGrid) {
+    private static List<DisplayOption> getPredefinedDeviceProfiles(Context context, String gridName,
+            Info displayInfo, boolean allowDisabledGrid) {
         ArrayList<DisplayOption> profiles = new ArrayList<>();
 
         try (XmlResourceParser parser = context.getResources().getXml(R.xml.device_profiles)) {
@@ -522,9 +522,10 @@
                 if ((type == XmlPullParser.START_TAG)
                         && GridOption.TAG_NAME.equals(parser.getName())) {
 
-                    GridOption gridOption = new GridOption(context, Xml.asAttributeSet(parser));
-                    if ((gridOption.isEnabled(deviceType) || allowDisabledGrid)
-                            && gridOption.filterByFlag(deviceType)) {
+                    GridOption gridOption = new GridOption(
+                            context, Xml.asAttributeSet(parser), displayInfo);
+                    if ((gridOption.isEnabled(displayInfo.getDeviceType()) || allowDisabledGrid)
+                            && gridOption.filterByFlag(displayInfo.getDeviceType())) {
                         final int displayDepth = parser.getDepth();
                         while (((type = parser.next()) != XmlPullParser.END_TAG
                                 || parser.getDepth() > displayDepth)
@@ -545,8 +546,8 @@
         ArrayList<DisplayOption> filteredProfiles = new ArrayList<>();
         if (!TextUtils.isEmpty(gridName)) {
             for (DisplayOption option : profiles) {
-                if (gridName.equals(option.grid.name)
-                        && (option.grid.isEnabled(deviceType) || allowDisabledGrid)) {
+                if (gridName.equals(option.grid.name) && (option.grid.isEnabled(
+                        displayInfo.getDeviceType()) || allowDisabledGrid)) {
                     filteredProfiles.add(option);
                 }
             }
@@ -572,10 +573,10 @@
      * Parses through the xml to find NumRows specs. Then calls findBestRowCount to get the correct
      * row count for this GridOption.
      *
-     * @return the result of {@link #findBestRowCount(List, Context, int)}.
+     * @return the result of {@link #findBestRowCount(List, Info)}.
      */
     public static NumRows getRowCount(ResourceHelper resourceHelper, Context context,
-            int deviceType) {
+            Info displayInfo) {
         ArrayList<NumRows> rowCounts = new ArrayList<>();
 
         try (XmlResourceParser parser = resourceHelper.getXml()) {
@@ -592,21 +593,19 @@
             throw new RuntimeException(e);
         }
 
-        return findBestRowCount(rowCounts, context, deviceType);
+        return findBestRowCount(rowCounts, displayInfo);
     }
 
     /**
      * @return the biggest row count that fits the display dimensions spec using NumRows to
      * determine that. If no best row count is found, return -1.
      */
-    public static NumRows findBestRowCount(List<NumRows> list, Context context,
-            @DeviceType int deviceType) {
-        Info displayInfo = DisplayController.INSTANCE.get(context).getInfo();
+    public static NumRows findBestRowCount(List<NumRows> list, Info displayInfo) {
         int minWidthPx = Integer.MAX_VALUE;
         int minHeightPx = Integer.MAX_VALUE;
         for (WindowBounds bounds : displayInfo.supportedBounds) {
             boolean isTablet = displayInfo.isTablet(bounds);
-            if (isTablet && deviceType == TYPE_MULTI_DISPLAY) {
+            if (isTablet && displayInfo.getDeviceType() == TYPE_MULTI_DISPLAY) {
                 // For split displays, take half width per page
                 minWidthPx = Math.min(minWidthPx, bounds.availableSize.x / 2);
                 minHeightPx = Math.min(minHeightPx, bounds.availableSize.y);
@@ -677,7 +676,7 @@
      * @return all the grid options that can be shown on the device
      */
     public List<GridOption> parseAllGridOptions(Context context) {
-        return parseAllDefinedGridOptions(context)
+        return parseAllDefinedGridOptions(context, displayInfo)
                 .stream()
                 .filter(go -> go.isEnabled(deviceType))
                 .filter(go -> go.filterByFlag(deviceType))
@@ -687,9 +686,8 @@
     /**
      * @return all the grid options that can be shown on the device
      */
-    public static List<GridOption> parseAllDefinedGridOptions(Context context) {
+    public static List<GridOption> parseAllDefinedGridOptions(Context context, Info displayInfo) {
         List<GridOption> result = new ArrayList<>();
-
         try (XmlResourceParser parser = context.getResources().getXml(R.xml.device_profiles)) {
             final int depth = parser.getDepth();
             int type;
@@ -697,7 +695,7 @@
                     || parser.getDepth() > depth) && type != XmlPullParser.END_DOCUMENT) {
                 if ((type == XmlPullParser.START_TAG)
                         && GridOption.TAG_NAME.equals(parser.getName())) {
-                    result.add(new GridOption(context, Xml.asAttributeSet(parser)));
+                    result.add(new GridOption(context, Xml.asAttributeSet(parser), displayInfo));
                 }
             }
         } catch (IOException | XmlPullParserException e) {
@@ -953,7 +951,7 @@
         private final int mAllAppsCellSpecsTwoPanelId;
         private final int mRowCountSpecsId;
 
-        public GridOption(Context context, AttributeSet attrs) {
+        public GridOption(Context context, AttributeSet attrs, Info displayInfo) {
             TypedArray a = context.obtainStyledAttributes(
                     attrs, R.styleable.GridDisplayOption);
             name = a.getString(R.styleable.GridDisplayOption_name);
@@ -965,7 +963,7 @@
             mIsDualGrid = a.getBoolean(R.styleable.GridDisplayOption_isDualGrid, false);
             if (mRowCountSpecsId != INVALID_RESOURCE_HANDLE) {
                 ResourceHelper resourceHelper = new ResourceHelper(context, mRowCountSpecsId);
-                NumRows numR = getRowCount(resourceHelper, context, deviceCategory);
+                NumRows numR = getRowCount(resourceHelper, context, displayInfo);
                 numRows = numR.mNumRowsNew;
                 dbFile = numR.mDbFile;
             } else {
diff --git a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
index fc8465d..1b58987 100644
--- a/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
+++ b/src/com/android/launcher3/allapps/ActivityAllAppsContainerView.java
@@ -95,6 +95,7 @@
 import com.android.launcher3.views.ScrimView;
 import com.android.launcher3.views.SpringRelativeLayout;
 import com.android.launcher3.workprofile.PersonalWorkSlidingTabStrip;
+import com.android.systemui.plugins.AllAppsRow;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -152,6 +153,7 @@
     private final RectF mTmpRectF = new RectF();
     protected AllAppsPagedView mViewPager;
     protected FloatingHeaderView mHeader;
+    protected final List<AllAppsRow> mAdditionalHeaderRows = new ArrayList<>();
     protected View mBottomSheetBackground;
     protected RecyclerViewFastScroller mFastScroller;
     private ConstraintLayout mFastScrollLetterLayout;
@@ -262,6 +264,8 @@
 
         getLayoutInflater().inflate(R.layout.all_apps_content, this);
         mHeader = findViewById(R.id.all_apps_header);
+        mAdditionalHeaderRows.clear();
+        mAdditionalHeaderRows.addAll(getAdditionalHeaderRows());
         mBottomSheetBackground = findViewById(R.id.bottom_sheet_background);
         mBottomSheetHandleArea = findViewById(R.id.bottom_sheet_handle_area);
         mSearchRecyclerView = findViewById(R.id.search_results_list_view);
@@ -288,6 +292,10 @@
         mSearchUiManager = (SearchUiManager) mSearchContainer;
     }
 
+    public List<AllAppsRow> getAdditionalHeaderRows() {
+        return List.of();
+    }
+
     @Override
     protected void onFinishInflate() {
         super.onFinishInflate();
@@ -714,6 +722,8 @@
     }
 
     void setupHeader() {
+        mAdditionalHeaderRows.forEach(row -> mHeader.onPluginDisconnected(row));
+
         mHeader.setVisibility(View.VISIBLE);
         boolean tabsHidden = !mUsingTabs;
         mHeader.setup(
@@ -731,6 +741,7 @@
                 adapterHolder.mRecyclerView.scrollToTop();
             }
         });
+        mAdditionalHeaderRows.forEach(row -> mHeader.onPluginConnected(row, mActivityContext));
 
         removeCustomRules(mHeader);
         if (isSearchBarFloating()) {
diff --git a/src/com/android/launcher3/allapps/FloatingHeaderView.java b/src/com/android/launcher3/allapps/FloatingHeaderView.java
index ac06ab4..8193511 100644
--- a/src/com/android/launcher3/allapps/FloatingHeaderView.java
+++ b/src/com/android/launcher3/allapps/FloatingHeaderView.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.allapps;
 
+import static com.android.launcher3.allapps.FloatingHeaderRow.NO_ROWS;
+
 import android.animation.ValueAnimator;
 import android.content.Context;
 import android.graphics.Point;
@@ -109,11 +111,11 @@
 
     // This is initialized once during inflation and stays constant after that. Fixed views
     // cannot be added or removed dynamically.
-    private FloatingHeaderRow[] mFixedRows = FloatingHeaderRow.NO_ROWS;
+    private FloatingHeaderRow[] mFixedRows = NO_ROWS;
 
     // Array of all fixed rows and plugin rows. This is initialized every time a plugin is
     // enabled or disabled, and represent the current set of all rows.
-    private FloatingHeaderRow[] mAllRows = FloatingHeaderRow.NO_ROWS;
+    private FloatingHeaderRow[] mAllRows = NO_ROWS;
 
     public FloatingHeaderView(@NonNull Context context) {
         this(context, null);
@@ -180,6 +182,10 @@
 
     @Override
     public void onPluginConnected(AllAppsRow allAppsRowPlugin, Context context) {
+        if (mPluginRows.containsKey(allAppsRowPlugin)) {
+            // Plugin has already been connected
+            return;
+        }
         PluginHeaderRow headerRow = new PluginHeaderRow(allAppsRowPlugin, this);
         addView(headerRow.mView, indexOfChild(mTabLayout));
         mPluginRows.put(allAppsRowPlugin, headerRow);
@@ -211,6 +217,9 @@
     @Override
     public void onPluginDisconnected(AllAppsRow plugin) {
         PluginHeaderRow row = mPluginRows.get(plugin);
+        if (row == null) {
+            return;
+        }
         removeView(row.mView);
         mPluginRows.remove(plugin);
         recreateAllRowsArray();
diff --git a/src/com/android/launcher3/icons/CacheableShortcutInfo.kt b/src/com/android/launcher3/icons/CacheableShortcutInfo.kt
index f8a5552..a78da23 100644
--- a/src/com/android/launcher3/icons/CacheableShortcutInfo.kt
+++ b/src/com/android/launcher3/icons/CacheableShortcutInfo.kt
@@ -121,7 +121,10 @@
         item: CacheableShortcutInfo,
         provider: IconProvider,
     ): String? =
-        item.shortcutInfo.lastChangedTimestamp.toString() +
+        // Manifest shortcuts get updated on every reboot. Don't include their change timestamp as
+        // it gets covered by the app's version
+        (if (item.shortcutInfo.isDeclaredInManifest) ""
+        else item.shortcutInfo.lastChangedTimestamp.toString()) +
             "-" +
             provider.getStateForApp(getApplicationInfo(item))
 }
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index e5cd76a..2550ebb 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -175,6 +175,10 @@
         @UiEvent(doc = "User searched for a widget in the widget picker.")
         LAUNCHER_WIDGETSTRAY_SEARCHED(819),
 
+        @UiEvent(doc = "User clicked on view all button to expand the displayed list in the "
+                + "widget picker.")
+        LAUNCHER_WIDGETSTRAY_EXPAND_PRESS(1978),
+
         @UiEvent(doc = "A dragged item is dropped on 'Remove' button in the target bar")
         LAUNCHER_ITEM_DROPPED_ON_REMOVE(465),
 
diff --git a/src/com/android/launcher3/widget/model/WidgetsListExpandActionEntry.java b/src/com/android/launcher3/widget/model/WidgetsListExpandActionEntry.java
new file mode 100644
index 0000000..8c84030
--- /dev/null
+++ b/src/com/android/launcher3/widget/model/WidgetsListExpandActionEntry.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget.model;
+
+import android.os.Process;
+
+import com.android.launcher3.model.data.PackageItemInfo;
+
+import java.util.Collections;
+
+/**
+ * Binds the section to be displayed at the bottom of the widgets list that enables user to expand
+ * and view all the widget apps including non-default. Bound when
+ * {@link WidgetsListExpandActionEntry} exists in the list on adapter.
+ */
+public class WidgetsListExpandActionEntry extends WidgetsListBaseEntry {
+
+    public WidgetsListExpandActionEntry() {
+        super(/*pkgItem=*/ new PackageItemInfo(/* packageName= */ "", Process.myUserHandle()),
+                /*titleSectionName=*/ "",
+                /*items=*/ Collections.EMPTY_LIST);
+        mPkgItem.title = "";
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 1860977..2f64ab1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -16,12 +16,16 @@
 package com.android.launcher3.widget.picker;
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.Flags.enableUnfoldedTwoPanePicker;
 import static com.android.launcher3.allapps.ActivityAllAppsContainerView.AdapterHolder.SEARCH;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_EXPAND_PRESS;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WIDGETSTRAY_SEARCHED;
 import static com.android.launcher3.testing.shared.TestProtocol.NORMAL_STATE_ORDINAL;
 import static com.android.launcher3.views.RecyclerViewFastScroller.FastScrollerLocation.WIDGET_SCROLLER;
 
+import static java.util.Collections.emptyList;
+
 import android.animation.Animator;
 import android.content.Context;
 import android.content.res.Resources;
@@ -68,6 +72,7 @@
 import com.android.launcher3.widget.BaseWidgetSheet;
 import com.android.launcher3.widget.WidgetCell;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
+import com.android.launcher3.widget.picker.model.data.WidgetPickerData;
 import com.android.launcher3.widget.picker.search.SearchModeListener;
 import com.android.launcher3.widget.picker.search.WidgetsSearchBar;
 import com.android.launcher3.widget.picker.search.WidgetsSearchBar.WidgetsSearchDataProvider;
@@ -87,7 +92,8 @@
  */
 public class WidgetsFullSheet extends BaseWidgetSheet
         implements OnActivePageChangedListener,
-        WidgetsRecyclerView.HeaderViewDimensionsProvider, SearchModeListener {
+        WidgetsRecyclerView.HeaderViewDimensionsProvider, SearchModeListener,
+        WidgetsListAdapter.ExpandButtonClickListener {
 
     private static final long FADE_IN_DURATION = 150;
 
@@ -257,7 +263,13 @@
         mSearchBar.initialize(new WidgetsSearchDataProvider() {
             @Override
             public List<WidgetsListBaseEntry> getWidgets() {
-                return getWidgetsToDisplay();
+                if (enableTieredWidgetsByDefaultInPicker()) {
+                    // search all
+                    return mActivityContext.getWidgetPickerDataProvider().get().getAllWidgets();
+                } else {
+                    // Can be removed when inlining enableTieredWidgetsByDefaultInPicker flag
+                    return getWidgetsToDisplay();
+                }
             }
         }, /* searchModeListener= */ this);
     }
@@ -482,6 +494,9 @@
     /**
      * Returns all displayable widgets.
      */
+    // Used by the two pane sheet to show 3-dot menu to toggle between default lists and all lists
+    // when enableTieredWidgetsByDefaultInPicker is OFF. This code path and the 3-dot menu can be
+    // safely deleted when it's alternative "enableTieredWidgetsByDefaultInPicker" flag is inlined.
     protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
         return mActivityContext.getWidgetPickerDataProvider().get().getAllWidgets();
     }
@@ -491,16 +506,27 @@
         if (mIsInSearchMode) {
             return;
         }
-        List<WidgetsListBaseEntry> widgets = getWidgetsToDisplay();
+        List<WidgetsListBaseEntry> widgets;
+        List<WidgetsListBaseEntry> defaultWidgets = emptyList();
+
+        if (enableTieredWidgetsByDefaultInPicker()) {
+            WidgetPickerData dataProvider =
+                    mActivityContext.getWidgetPickerDataProvider().get();
+            widgets = dataProvider.getAllWidgets();
+            defaultWidgets = dataProvider.getDefaultWidgets();
+        } else {
+            // This code path can be deleted once enableTieredWidgetsByDefaultInPicker is inlined.
+            widgets = getWidgetsToDisplay();
+        }
 
         AdapterHolder primaryUserAdapterHolder = mAdapters.get(AdapterHolder.PRIMARY);
-        primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets);
+        primaryUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets, defaultWidgets);
 
         if (mHasWorkProfile) {
             mViewPager.setVisibility(VISIBLE);
             mTabBar.setVisibility(VISIBLE);
             AdapterHolder workUserAdapterHolder = mAdapters.get(AdapterHolder.WORK);
-            workUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets);
+            workUserAdapterHolder.mWidgetsListAdapter.setWidgets(widgets, defaultWidgets);
             onActivePageChanged(mViewPager.getCurrentPage());
         } else {
             onActivePageChanged(0);
@@ -518,6 +544,16 @@
     }
 
     @Override
+    public void onWidgetsListExpandButtonClick(View v) {
+        AdapterHolder currentAdapterHolder = mAdapters.get(getCurrentAdapterHolderType());
+        currentAdapterHolder.mWidgetsListAdapter.useExpandedList();
+        onWidgetsBound();
+        currentAdapterHolder.mWidgetsRecyclerView.announceForAccessibility(
+                mActivityContext.getString(R.string.widgets_list_expanded));
+        mActivityContext.getStatsLogManager().logger().log(LAUNCHER_WIDGETSTRAY_EXPAND_PRESS);
+    }
+
+    @Override
     public void enterSearchMode(boolean shouldLog) {
         if (mIsInSearchMode) return;
         setViewVisibilityBasedOnSearch(/*isInSearchMode= */ true);
@@ -832,7 +868,7 @@
                 + marginLayoutParams.topMargin;
     }
 
-    private int getCurrentAdapterHolderType() {
+    protected int getCurrentAdapterHolderType() {
         if (mIsInSearchMode) {
             return SEARCH;
         } else if (mViewPager != null) {
@@ -861,6 +897,7 @@
             WidgetsFullSheet sheet = show(BaseActivity.fromContext(getContext()), false);
             sheet.restoreRecommendations(mRecommendedWidgets, mRecommendedWidgetsMap);
             sheet.restoreHierarchyState(widgetsState);
+            sheet.restoreAdapterStates(mAdapters);
             sheet.restorePreviousAdapterHolderType(getCurrentAdapterHolderType());
         } else if (!isTwoPane()) {
             reset();
@@ -876,6 +913,17 @@
         mRecommendedWidgetsMap = recommendedWidgetsMap;
     }
 
+    private void restoreAdapterStates(SparseArray<AdapterHolder> adapters) {
+        if (adapters.contains(AdapterHolder.WORK)) {
+            mAdapters.get(AdapterHolder.WORK).mWidgetsListAdapter.restoreState(
+                    adapters.get(AdapterHolder.WORK).mWidgetsListAdapter);
+        }
+        mAdapters.get(AdapterHolder.PRIMARY).mWidgetsListAdapter.restoreState(
+                adapters.get(AdapterHolder.PRIMARY).mWidgetsListAdapter);
+        mAdapters.get(AdapterHolder.SEARCH).mWidgetsListAdapter.restoreState(
+                adapters.get(AdapterHolder.SEARCH).mWidgetsListAdapter);
+    }
+
     /**
      * Indicates if layout should be re-created on device profile change - so that a different
      * layout can be displayed.
@@ -1045,6 +1093,7 @@
                     this::getEmptySpaceHeight,
                     /* iconClickListener= */ WidgetsFullSheet.this,
                     /* iconLongClickListener= */ WidgetsFullSheet.this,
+                    /* expandButtonClickListener= */ WidgetsFullSheet.this,
                     isTwoPane());
             mWidgetsListAdapter.setHasStableIds(true);
             switch (mAdapterType) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
index 3d3a669..3c67538 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListAdapter.java
@@ -49,6 +49,7 @@
 import com.android.launcher3.widget.model.WidgetListSpaceEntry;
 import com.android.launcher3.widget.model.WidgetsListBaseEntry;
 import com.android.launcher3.widget.model.WidgetsListContentEntry;
+import com.android.launcher3.widget.model.WidgetsListExpandActionEntry;
 import com.android.launcher3.widget.model.WidgetsListHeaderEntry;
 import com.android.launcher3.widget.util.WidgetSizes;
 
@@ -82,6 +83,7 @@
     public static final int VIEW_TYPE_WIDGETS_SPACE = R.id.view_type_widgets_space;
     public static final int VIEW_TYPE_WIDGETS_LIST = R.id.view_type_widgets_list;
     public static final int VIEW_TYPE_WIDGETS_HEADER = R.id.view_type_widgets_header;
+    public static final int VIEW_TYPE_WIDGETS_EXPAND = R.id.view_type_widgets_list_expand;
 
     private final Context mContext;
     private final SparseArray<ViewHolderBinder> mViewHolderBinders = new SparseArray<>();
@@ -90,7 +92,9 @@
     @Nullable private WidgetsTwoPaneSheet.HeaderChangeListener mHeaderChangeListener;
 
     private final List<WidgetsListBaseEntry> mAllEntries = new ArrayList<>();
+    private final List<WidgetsListBaseEntry> mAllDefaultEntries = new ArrayList<>();
     private ArrayList<WidgetsListBaseEntry> mVisibleEntries = new ArrayList<>();
+
     @Nullable private PackageUserKey mWidgetsContentVisiblePackageUserKey = null;
 
     private Predicate<WidgetsListBaseEntry> mHeaderAndSelectedContentFilter = entry ->
@@ -102,9 +106,12 @@
     @Nullable private PackageUserKey mPendingClickHeader;
     @Px private int mMaxHorizontalSpan;
 
+    private boolean mShowOnlyDefaultList = true;
+
     public WidgetsListAdapter(Context context, LayoutInflater layoutInflater,
             IntSupplier emptySpaceHeightProvider, OnClickListener iconClickListener,
             OnLongClickListener iconLongClickListener,
+            ExpandButtonClickListener expandButtonClickListener,
             boolean isTwoPane) {
         mContext = context;
         mMaxHorizontalSpan = WidgetSizes.getWidgetSizePx(
@@ -123,6 +130,16 @@
         mViewHolderBinders.put(
                 VIEW_TYPE_WIDGETS_SPACE,
                 new WidgetsSpaceViewHolderBinder(emptySpaceHeightProvider));
+        mViewHolderBinders.put(VIEW_TYPE_WIDGETS_EXPAND,
+                new WidgetsListExpandActionViewHolderBinder(layoutInflater,
+                        expandButtonClickListener::onWidgetsListExpandButtonClick));
+    }
+
+    /**
+     * Copies state info from another adapter.
+     */
+    public void restoreState(WidgetsListAdapter adapter) {
+        mShowOnlyDefaultList = adapter.mShowOnlyDefaultList;
     }
 
     public void setHeaderChangeListener(WidgetsTwoPaneSheet.HeaderChangeListener
@@ -168,10 +185,21 @@
     }
 
     /** Updates the widget list based on {@code tempEntries}. */
-    public void setWidgets(List<WidgetsListBaseEntry> tempEntries) {
+    public void setWidgets(List<WidgetsListBaseEntry> tempEntries,
+            List<WidgetsListBaseEntry> tempDefaultEntries) {
         mAllEntries.clear();
         mAllEntries.add(new WidgetListSpaceEntry());
         tempEntries.stream().sorted(mRowComparator).forEach(mAllEntries::add);
+
+        mAllDefaultEntries.clear();
+
+        if (mShowOnlyDefaultList && !tempDefaultEntries.isEmpty()) {
+            mAllDefaultEntries.add(new WidgetListSpaceEntry());
+            tempDefaultEntries.stream().sorted(mRowComparator).forEach(mAllDefaultEntries::add);
+            // Include view all action when default entries exist.
+            mAllDefaultEntries.add(new WidgetsListExpandActionEntry());
+        }
+
         updateVisibleEntries();
     }
 
@@ -179,7 +207,8 @@
     public void setWidgetsOnSearch(List<WidgetsListBaseEntry> searchResults) {
         // Forget the expanded package every time widget list is refreshed in search mode.
         mWidgetsContentVisiblePackageUserKey = null;
-        setWidgets(searchResults);
+        mShowOnlyDefaultList = false;
+        setWidgets(searchResults, /*tempDefaultEntries=*/ List.of());
     }
 
     private void updateVisibleEntries() {
@@ -190,10 +219,11 @@
         OptionalInt topForPackageUserKey =
                 getOffsetForPosition(previousPositionForPackageUserKey);
 
-        List<WidgetsListBaseEntry> newVisibleEntries = mAllEntries.stream()
+        List<WidgetsListBaseEntry> newVisibleEntries = getAllEntries().stream()
                 .filter(entry -> (((mFilter == null || mFilter.test(entry))
                         && mHeaderAndSelectedContentFilter.test(entry))
-                        || entry instanceof WidgetListSpaceEntry)
+                        || entry instanceof WidgetListSpaceEntry
+                        || entry instanceof WidgetsListExpandActionEntry)
                         && (mHeaderChangeListener == null
                         || !(entry instanceof WidgetsListContentEntry)))
                 .map(entry -> {
@@ -227,6 +257,11 @@
         }
     }
 
+    private List<WidgetsListBaseEntry> getAllEntries() {
+        return (mShowOnlyDefaultList && !mAllDefaultEntries.isEmpty()) ? mAllDefaultEntries
+                : mAllEntries;
+    }
+
     /** Returns whether {@code entry} matches {@code key}. */
     private static boolean isHeaderForPackageUserKey(
             @NonNull WidgetsListBaseEntry entry, @Nullable PackageUserKey key) {
@@ -262,7 +297,13 @@
 
         // The first entry has an empty space, count from second entries.
         int listPos = (pos > 1) ? POSITION_DEFAULT : POSITION_FIRST;
-        if (pos == (getItemCount() - 1)) {
+        int lastIndex = getItemCount() - 1;
+        // Last index may be the view all entry
+        int actualLastItemIndex = (mVisibleEntries.get(
+                lastIndex) instanceof WidgetsListExpandActionEntry) ? getItemCount() - 2
+                : getItemCount() - 1;
+
+        if (pos == (actualLastItemIndex)) {
             listPos |= POSITION_LAST;
         }
         viewHolderBinder.bindViewHolder(holder, mVisibleEntries.get(pos), listPos, payloads);
@@ -319,6 +360,8 @@
             return VIEW_TYPE_WIDGETS_HEADER;
         } else if (entry instanceof WidgetListSpaceEntry) {
             return VIEW_TYPE_WIDGETS_SPACE;
+        } else if (entry instanceof WidgetsListExpandActionEntry) {
+            return VIEW_TYPE_WIDGETS_EXPAND;
         }
         throw new UnsupportedOperationException("ViewHolderBinder not found for " + entry);
     }
@@ -412,6 +455,23 @@
         updateVisibleEntries();
     }
 
+    /**
+     * Returns the widget content {@link WidgetsListContentEntry} for a selected header.
+     */
+    public WidgetsListContentEntry getContentEntry(PackageUserKey selectedHeader) {
+        return getAllEntries().stream().filter(entry -> entry instanceof WidgetsListContentEntry)
+                .map(entry -> (WidgetsListContentEntry) entry)
+                .filter(entry -> PackageUserKey.fromPackageItemInfo(entry.mPkgItem).equals(
+                        selectedHeader)).findFirst().orElse(null);
+    }
+
+    /**
+     * Sets adapter to use expanded list when updating widgets.
+     */
+    public void useExpandedList() {
+        mShowOnlyDefaultList = false;
+    }
+
     /** Comparator for sorting WidgetListRowEntry based on package title. */
     public static class WidgetListBaseRowEntryComparator implements
             Comparator<WidgetsListBaseEntry> {
@@ -430,4 +490,10 @@
             return 1;
         }
     }
+
+    /** Callback interface for the interaction with the expand button */
+    public interface ExpandButtonClickListener {
+        /** Called when user clicks the button at end of widget apps list to expand it. */
+        void onWidgetsListExpandButtonClick(View view);
+    }
 }
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListExpandActionViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListExpandActionViewHolderBinder.java
new file mode 100644
index 0000000..288c456
--- /dev/null
+++ b/src/com/android/launcher3/widget/picker/WidgetsListExpandActionViewHolderBinder.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.widget.picker;
+
+import android.view.LayoutInflater;
+import android.view.View;
+import android.view.ViewGroup;
+
+import androidx.annotation.NonNull;
+import androidx.recyclerview.widget.RecyclerView;
+
+import com.android.launcher3.R;
+import com.android.launcher3.recyclerview.ViewHolderBinder;
+import com.android.launcher3.widget.model.WidgetsListExpandActionEntry;
+
+import java.util.List;
+
+/**
+ * Creates and populates views for the {@link WidgetsListExpandActionEntry}.
+ */
+public class WidgetsListExpandActionViewHolderBinder implements
+        ViewHolderBinder<WidgetsListExpandActionEntry, RecyclerView.ViewHolder> {
+    @NonNull
+    View.OnClickListener mExpandListClickListener;
+    private final LayoutInflater mLayoutInflater;
+
+    public WidgetsListExpandActionViewHolderBinder(
+            @NonNull LayoutInflater layoutInflater,
+            @NonNull View.OnClickListener expandListClickListener) {
+        mLayoutInflater = layoutInflater;
+        mExpandListClickListener = expandListClickListener;
+    }
+
+    @Override
+    public RecyclerView.ViewHolder newViewHolder(ViewGroup parent) {
+        return new RecyclerView.ViewHolder(mLayoutInflater.inflate(
+                R.layout.widgets_list_expand_button, parent, false)) {
+        };
+    }
+
+    @Override
+    public void bindViewHolder(RecyclerView.ViewHolder viewHolder,
+            WidgetsListExpandActionEntry data, int position, List<Object> payloads) {
+        viewHolder.itemView.setOnClickListener(mExpandListClickListener);
+    }
+}
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index f4b99a0..f9bd5f1 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.widget.picker;
 
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.Flags.enableUnfoldedTwoPanePicker;
 import static com.android.launcher3.UtilitiesKt.CLIP_CHILDREN_FALSE_MODIFIER;
 import static com.android.launcher3.UtilitiesKt.CLIP_TO_PADDING_FALSE_MODIFIER;
@@ -147,7 +148,9 @@
         mHeaderDescription = mContent.findViewById(R.id.widget_picker_description);
 
         mWidgetOptionsMenu = mContent.findViewById(R.id.widget_picker_widget_options_menu);
-        setupWidgetOptionsMenu();
+        if (!enableTieredWidgetsByDefaultInPicker()) {
+            setupWidgetOptionsMenu();
+        }
 
         mRightPane = mContent.findViewById(R.id.right_pane);
         mRightPaneScrollView = mContent.findViewById(R.id.right_pane_scroll_view);
@@ -286,6 +289,9 @@
         }
     }
 
+    // Used by the two pane sheet to show 3-dot menu to toggle between default lists and all lists
+    // when enableTieredWidgetsByDefaultInPicker is OFF. This code path and the 3-dot menu can be
+    // safely deleted when it's alternative "enableTieredWidgetsByDefaultInPicker" flag is inlined.
     @Override
     protected List<WidgetsListBaseEntry> getWidgetsToDisplay() {
         List<WidgetsListBaseEntry> allWidgets =
@@ -319,6 +325,15 @@
     }
 
     @Override
+    public void onWidgetsListExpandButtonClick(View v) {
+        super.onWidgetsListExpandButtonClick(v);
+        // Refresh right pane with updated data for the selected header.
+        if (mSelectedHeader != null && mSelectedHeader != mSuggestedWidgetsPackageUserKey) {
+            getHeaderChangeListener().onHeaderChanged(mSelectedHeader);
+        }
+    }
+
+    @Override
     public void onRecommendedWidgetsBound() {
         super.onRecommendedWidgetsBound();
 
@@ -511,11 +526,20 @@
                         && !mOpenCloseAnimation.getAnimationPlayer().isRunning()
                         && !getAccessibilityInitialFocusView().isAccessibilityFocused();
                 mSelectedHeader = selectedHeader;
-                final boolean showDefaultWidgets = mWidgetOptionsMenuState != null
-                        && !mWidgetOptionsMenuState.showAllWidgets;
-                WidgetsListContentEntry contentEntry = findContentEntryForPackageUser(
-                        mActivityContext.getWidgetPickerDataProvider().get(),
-                        selectedHeader, showDefaultWidgets);
+
+                WidgetsListContentEntry contentEntry;
+                if (enableTieredWidgetsByDefaultInPicker()) {
+                    contentEntry = mAdapters.get(
+                            getCurrentAdapterHolderType()).mWidgetsListAdapter.getContentEntry(
+                            selectedHeader);
+                } else { // Can be deleted when inlining the "enableTieredWidgetsByDefaultInPicker"
+                    // flag
+                    final boolean showDefaultWidgets = mWidgetOptionsMenuState != null
+                            && !mWidgetOptionsMenuState.showAllWidgets;
+                    contentEntry = findContentEntryForPackageUser(
+                            mActivityContext.getWidgetPickerDataProvider().get(),
+                            selectedHeader, showDefaultWidgets);
+                }
 
                 if (contentEntry == null || mRightPane == null) {
                     return;
diff --git a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
index ed9a080..bae74c8 100644
--- a/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/icons/IconCacheUpdateHandlerTest.kt
@@ -19,6 +19,7 @@
 import android.content.ComponentName
 import android.content.pm.ApplicationInfo
 import android.database.MatrixCursor
+import android.os.Handler
 import android.os.Process.myUserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
@@ -34,9 +35,18 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Captor
 import org.mockito.Mock
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.anyOrNull
+import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.eq
+import org.mockito.kotlin.never
+import org.mockito.kotlin.times
+import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 
 @SmallTest
@@ -45,14 +55,25 @@
 
     @Mock private lateinit var iconProvider: IconProvider
     @Mock private lateinit var baseIconCache: BaseIconCache
+    @Mock private lateinit var cacheDb: IconDB
+    @Mock private lateinit var workerHandler: Handler
 
-    private var cursor: MatrixCursor? = null
-    private var cachingLogic = CachedObjectCachingLogic
+    @Captor private lateinit var deleteCaptor: ArgumentCaptor<String>
+
+    private var cursor =
+        MatrixCursor(
+            arrayOf(IconDB.COLUMN_ROWID, IconDB.COLUMN_COMPONENT, IconDB.COLUMN_FRESHNESS_ID)
+        )
+
+    private lateinit var updateHandlerUnderTest: IconCacheUpdateHandler
 
     @Before
     fun setup() {
         MockitoAnnotations.initMocks(this)
         doReturn(iconProvider).whenever(baseIconCache).iconProvider
+        doReturn(cursor).whenever(cacheDb).query(any(), any(), any())
+
+        updateHandlerUnderTest = IconCacheUpdateHandler(baseIconCache, cacheDb, workerHandler)
     }
 
     @After
@@ -61,63 +82,136 @@
     }
 
     @Test
-    fun `IconCacheUpdateHandler returns null if the component name is malformed`() {
-        val updateHandlerUnderTest = IconCacheUpdateHandler(baseIconCache)
-        val cn = ComponentName.unflattenFromString("com.android.fake/.FakeActivity")!!
+    fun `keeps correct icons irrespective of call order`() {
+        val obj1 = TestCachedObject(1).apply { addToCursor(cursor) }
+        val obj2 = TestCachedObject(2).apply { addToCursor(cursor) }
 
-        val result =
-            updateHandlerUnderTest.updateOrDeleteIcon(
-                createCursor(1, cn.flattenToString() + "#", "freshId-old"),
-                hashMapOf(cn to TestCachedObject(cn, "freshId")),
-                setOf(),
-                myUserHandle(),
-                cachingLogic,
-            )
-        assertThat(result).isNull()
+        updateHandlerUnderTest.updateIcons(obj1)
+        updateHandlerUnderTest.updateIcons(obj2)
+        updateHandlerUnderTest.finish()
+
+        verify(cacheDb, never()).delete(any(), anyOrNull())
     }
 
     @Test
-    fun `IconCacheUpdateHandler returns null if the freshId match`() {
-        val updateHandlerUnderTest = IconCacheUpdateHandler(baseIconCache)
-        val cn = ComponentName.unflattenFromString("com.android.fake/.FakeActivity")!!
+    fun `removes missing entries in single call`() {
+        TestCachedObject(1).addToCursor(cursor)
+        TestCachedObject(2).addToCursor(cursor)
+        TestCachedObject(3).addToCursor(cursor)
+        TestCachedObject(4).addToCursor(cursor)
+        TestCachedObject(5).addToCursor(cursor)
 
-        val result =
-            updateHandlerUnderTest.updateOrDeleteIcon(
-                createCursor(1, cn.flattenToString(), "freshId"),
-                hashMapOf(cn to TestCachedObject(cn, "freshId")),
-                setOf(),
-                myUserHandle(),
-                cachingLogic,
-            )
-        assertThat(result).isNull()
+        updateHandlerUnderTest.updateIcons(TestCachedObject(1), TestCachedObject(4))
+        updateHandlerUnderTest.finish()
+
+        verifyItemsDeleted(2, 3, 5)
     }
 
     @Test
-    fun `IconCacheUpdateHandler returns non-null if the freshId do not match`() {
-        val updateHandlerUnderTest = IconCacheUpdateHandler(baseIconCache)
-        val cn = ComponentName.unflattenFromString("com.android.fake/.FakeActivity")!!
-        val testObj = TestCachedObject(cn, "freshId")
+    fun `removes missing entries in multiple calls`() {
+        TestCachedObject(1).addToCursor(cursor)
+        TestCachedObject(2).addToCursor(cursor)
+        TestCachedObject(3).addToCursor(cursor)
+        TestCachedObject(4).addToCursor(cursor)
+        TestCachedObject(5).addToCursor(cursor)
+        TestCachedObject(6).addToCursor(cursor)
 
-        val result =
-            updateHandlerUnderTest.updateOrDeleteIcon(
-                createCursor(1, cn.flattenToString(), "freshId-old"),
-                hashMapOf(cn to testObj),
-                setOf(),
-                myUserHandle(),
-                cachingLogic,
-            )
-        assertThat(result).isEqualTo(testObj)
+        updateHandlerUnderTest.updateIcons(TestCachedObject(1), TestCachedObject(2))
+        updateHandlerUnderTest.updateIcons(TestCachedObject(4), TestCachedObject(5))
+        updateHandlerUnderTest.finish()
+
+        verifyItemsDeleted(3, 6)
     }
 
-    private fun createCursor(row: Long, component: String, appState: String) =
-        MatrixCursor(
-                arrayOf(IconDB.COLUMN_ROWID, IconDB.COLUMN_COMPONENT, IconDB.COLUMN_FRESHNESS_ID)
-            )
-            .apply { addRow(arrayOf(row, component, appState)) }
-            .apply {
-                cursor = this
-                moveToNext()
+    @Test
+    fun `keeps valid app infos`() {
+        val appInfo = ApplicationInfo()
+        doReturn("app-fresh").whenever(iconProvider).getStateForApp(eq(appInfo))
+
+        TestCachedObject(1).addToCursor(cursor)
+        TestCachedObject(2).addToCursor(cursor)
+        cursor.addRow(arrayOf(33, TestCachedObject(1).getPackageKey(), "app-fresh"))
+
+        updateHandlerUnderTest.updateIcons(
+            TestCachedObject(1, appInfo = appInfo),
+            TestCachedObject(2),
+        )
+        updateHandlerUnderTest.finish()
+
+        verify(cacheDb, never()).delete(any(), anyOrNull())
+    }
+
+    @Test
+    fun `deletes stale app infos`() {
+        val appInfo1 = ApplicationInfo()
+        doReturn("app1-fresh").whenever(iconProvider).getStateForApp(eq(appInfo1))
+
+        val appInfo2 = ApplicationInfo()
+        doReturn("app2-fresh").whenever(iconProvider).getStateForApp(eq(appInfo2))
+
+        TestCachedObject(1).addToCursor(cursor)
+        TestCachedObject(2).addToCursor(cursor)
+        cursor.addRow(arrayOf(33, TestCachedObject(1).getPackageKey(), "app1-not-fresh"))
+        cursor.addRow(arrayOf(34, TestCachedObject(2).getPackageKey(), "app2-fresh"))
+
+        updateHandlerUnderTest.updateIcons(
+            TestCachedObject(1, appInfo = appInfo1),
+            TestCachedObject(2, appInfo = appInfo2),
+        )
+        updateHandlerUnderTest.finish()
+
+        verifyItemsDeleted(33)
+    }
+
+    @Test
+    fun `updates stale entries`() {
+        doAnswer { i ->
+                (i.arguments[0] as Runnable).run()
+                true
             }
+            .whenever(workerHandler)
+            .postAtTime(any(), anyOrNull(), any())
+
+        TestCachedObject(1).addToCursor(cursor)
+        TestCachedObject(2).addToCursor(cursor)
+        TestCachedObject(3).addToCursor(cursor)
+
+        var updatedPackages = mutableSetOf<String>()
+        updateHandlerUnderTest.updateIcons(
+            listOf(
+                TestCachedObject(1, freshnessId = "not-fresh"),
+                TestCachedObject(2, freshnessId = "not-fresh"),
+                TestCachedObject(3),
+            ),
+            CachedObjectCachingLogic,
+        ) { apps, _ ->
+            updatedPackages.addAll(apps)
+        }
+        updateHandlerUnderTest.finish()
+
+        assertThat(updatedPackages)
+            .isEqualTo(
+                mutableSetOf(TestCachedObject(1).cn.packageName, TestCachedObject(2).cn.packageName)
+            )
+    }
+
+    private fun IconCacheUpdateHandler.updateIcons(vararg items: TestCachedObject) {
+        updateIcons(items.toList(), CachedObjectCachingLogic) { _, _ -> }
+    }
+
+    private fun verifyItemsDeleted(vararg rowIds: Long) {
+        verify(cacheDb, times(1)).delete(deleteCaptor.capture(), anyOrNull())
+        val actual =
+            deleteCaptor.value
+                .split('(')
+                ?.get(1)
+                ?.split(')')
+                ?.get(0)
+                ?.split(",")
+                ?.map { it.trim().toLong() }!!
+                .sorted()
+        assertThat(actual).isEqualTo(rowIds.toList().sorted())
+    }
 }
 
 /** Utility method to wait for the icon update handler to finish */
@@ -135,7 +229,13 @@
     }
 }
 
-class TestCachedObject(val cn: ComponentName, val freshnessId: String) : CachedObject {
+class TestCachedObject(
+    val rowId: Long,
+    val cn: ComponentName =
+        ComponentName.unflattenFromString("com.android.fake$rowId/.FakeActivity")!!,
+    val freshnessId: String = "fresh-$rowId",
+    val appInfo: ApplicationInfo? = null,
+) : CachedObject {
 
     override fun getComponent() = cn
 
@@ -143,7 +243,13 @@
 
     override fun getLabel(): CharSequence? = null
 
-    override fun getApplicationInfo(): ApplicationInfo? = null
+    override fun getApplicationInfo(): ApplicationInfo? = appInfo
 
     override fun getFreshnessIdentifier(iconProvider: IconProvider): String? = freshnessId
+
+    fun addToCursor(cursor: MatrixCursor) =
+        cursor.addRow(arrayOf(rowId, cn.flattenToString(), freshnessId))
+
+    fun getPackageKey() =
+        BaseIconCache.getPackageKey(cn.packageName, user).componentName.flattenToString()
 }
diff --git a/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt b/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt
index dee98e7..0f212eb 100644
--- a/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/util/VibratorWrapperTest.kt
@@ -66,6 +66,7 @@
 
     @Test
     fun init_register_onChangeListener() {
+        TestUtil.runOnExecutorSync(Executors.MAIN_EXECUTOR) {}
         verify(settingsCache).register(HAPTIC_FEEDBACK_URI, underTest.mHapticChangeListener)
     }
 
diff --git a/tests/src/com/android/launcher3/LauncherIntentTest.java b/tests/src/com/android/launcher3/LauncherIntentTest.java
index aeeb42a..3e16713 100644
--- a/tests/src/com/android/launcher3/LauncherIntentTest.java
+++ b/tests/src/com/android/launcher3/LauncherIntentTest.java
@@ -23,7 +23,7 @@
 import android.platform.test.annotations.LargeTest;
 import android.view.KeyEvent;
 
-import androidx.test.runner.AndroidJUnit4;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
 
 import com.android.launcher3.allapps.ActivityAllAppsContainerView;
 import com.android.launcher3.allapps.SearchRecyclerView;
diff --git a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
index 44b8ff8..1816030 100644
--- a/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
+++ b/tests/src/com/android/launcher3/dragging/TaplUninstallRemoveTest.java
@@ -69,8 +69,7 @@
     private void verifyAppUninstalledFromAllApps(Workspace workspace, String appName) {
         final HomeAllApps allApps = workspace.switchToAllApps();
         Wait.atMost(appName + " app was found on all apps after being uninstalled",
-                () -> allApps.tryGetAppIcon(appName) == null,
-                DEFAULT_UI_TIMEOUT, mLauncher);
+                () -> allApps.tryGetAppIcon(appName) == null, mLauncher);
     }
 
     private void installDummyAppAndWaitForUIUpdate() throws IOException {
diff --git a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
index 206647a..1fbdceb 100644
--- a/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
+++ b/tests/src/com/android/launcher3/ui/AbstractLauncherUiTest.java
@@ -32,8 +32,6 @@
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.ActivityInfo;
-import android.content.pm.LauncherActivityInfo;
-import android.content.pm.LauncherApps;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.graphics.Point;
@@ -97,10 +95,8 @@
  */
 public abstract class AbstractLauncherUiTest<LAUNCHER_TYPE extends Launcher> {
 
-    public static final long DEFAULT_ACTIVITY_TIMEOUT = TimeUnit.SECONDS.toMillis(10);
     public static final long DEFAULT_BROADCAST_TIMEOUT_SECS = 10;
 
-    public static final long DEFAULT_UI_TIMEOUT = TestUtil.DEFAULT_UI_TIMEOUT;
     private static final String TAG = "AbstractLauncherUiTest";
 
     private static final long BYTES_PER_MEGABYTE = 1 << 20;
@@ -151,7 +147,7 @@
                     launcher.forceGc();
                     return MAIN_EXECUTOR.submit(
                             () -> launcher.noLeakedActivities(requireOneActiveActivity)).get();
-                }, DEFAULT_UI_TIMEOUT, launcher);
+                }, launcher);
     }
 
     public static String getAppPackageName() {
@@ -443,7 +439,7 @@
      */
     protected <T> T getOnUiThread(final Callable<T> callback) {
         try {
-            return mMainThreadExecutor.submit(callback).get(DEFAULT_UI_TIMEOUT,
+            return mMainThreadExecutor.submit(callback).get(TestUtil.DEFAULT_UI_TIMEOUT,
                     TimeUnit.MILLISECONDS);
         } catch (TimeoutException e) {
             Log.e(TAG, "Timeout in getOnUiThread, sending SIGABRT", e);
@@ -498,13 +494,7 @@
     // flakiness.
     protected void waitForLauncherCondition(String
             message, Function<LAUNCHER_TYPE, Boolean> condition) {
-        waitForLauncherCondition(message, condition, DEFAULT_ACTIVITY_TIMEOUT);
-    }
-
-    // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
-    // flakiness.
-    protected <O> O getOnceNotNull(String message, Function<LAUNCHER_TYPE, O> f) {
-        return getOnceNotNull(message, f, DEFAULT_ACTIVITY_TIMEOUT);
+        waitForLauncherCondition(message, condition, TestUtil.DEFAULT_UI_TIMEOUT);
     }
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
@@ -513,12 +503,12 @@
             String message, Function<LAUNCHER_TYPE, Boolean> condition, long timeout) {
         verifyKeyguardInvisible();
         if (!TestHelpers.isInLauncherProcess()) return;
-        Wait.atMost(message, () -> getFromLauncher(condition), timeout, mLauncher);
+        Wait.atMost(message, () -> getFromLauncher(condition), mLauncher, timeout);
     }
 
     // Cannot be used in TaplTests after injecting any gesture using Tapl because this can hide
     // flakiness.
-    protected <T> T getOnceNotNull(String message, Function<LAUNCHER_TYPE, T> f, long timeout) {
+    protected <T> T getOnceNotNull(String message, Function<LAUNCHER_TYPE, T> f) {
         if (!TestHelpers.isInLauncherProcess()) return null;
 
         final Object[] output = new Object[1];
@@ -526,7 +516,7 @@
             final Object fromLauncher = getFromLauncher(f);
             output[0] = fromLauncher;
             return fromLauncher != null;
-        }, timeout, mLauncher);
+        }, mLauncher);
         return (T) output[0];
     }
 
@@ -540,12 +530,7 @@
         Wait.atMost(message, () -> {
             testThreadAction.run();
             return getFromLauncher(condition);
-        }, timeout, mLauncher);
-    }
-
-    protected LauncherActivityInfo getSettingsApp() {
-        return mTargetContext.getSystemService(LauncherApps.class)
-                .getActivityList("com.android.settings", Process.myUserHandle()).get(0);
+        }, mLauncher, timeout);
     }
 
     /**
@@ -633,13 +618,13 @@
         }
         getInstrumentation().getTargetContext().startActivity(intent);
         assertTrue("App didn't start: " + selector,
-                TestHelpers.wait(Until.hasObject(selector), DEFAULT_UI_TIMEOUT));
+                TestHelpers.wait(Until.hasObject(selector), TestUtil.DEFAULT_UI_TIMEOUT));
 
         // Wait for the Launcher to stop.
         final LauncherInstrumentation launcherInstrumentation = new LauncherInstrumentation();
         Wait.atMost("Launcher activity didn't stop",
                 () -> !launcherInstrumentation.isLauncherActivityStarted(),
-                DEFAULT_ACTIVITY_TIMEOUT, launcherInstrumentation);
+                launcherInstrumentation);
     }
 
     public static ActivityInfo resolveSystemAppInfo(String category) {
@@ -662,8 +647,7 @@
                 launcher.finish();
             }
         });
-        waitForLauncherCondition(
-                "Launcher still active", launcher -> launcher == null, DEFAULT_UI_TIMEOUT);
+        waitForLauncherCondition("Launcher still active", launcher -> launcher == null);
     }
 
     protected boolean isInLaunchedApp(LAUNCHER_TYPE launcher) {
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
index 7ff4f22..7845222 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddConfigWidgetTest.java
@@ -103,12 +103,12 @@
 
         setResultAndWaitForAnimation(acceptConfig);
         if (acceptConfig) {
-            Wait.atMost("", new WidgetSearchCondition(), DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+            Wait.atMost("", new WidgetSearchCondition(), mLauncher);
             assertNotNull(mAppWidgetManager.getAppWidgetInfo(mWidgetId));
         } else {
             // Verify that the widget id is deleted.
             Wait.atMost("", () -> mAppWidgetManager.getAppWidgetInfo(mWidgetId) == null,
-                    DEFAULT_ACTIVITY_TIMEOUT, mLauncher);
+                    mLauncher);
         }
     }
 
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
index 9a2147a..460ffc4 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplAddWidgetTest.java
@@ -29,6 +29,7 @@
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.PortraitLandscapeRunner.PortraitLandscape;
 import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 
@@ -68,7 +69,7 @@
         resizeFrame.dismiss();
 
         final Widget widget = mLauncher.getWorkspace().tryGetWidget(widgetInfo.label,
-                DEFAULT_UI_TIMEOUT);
+                TestUtil.DEFAULT_UI_TIMEOUT);
         assertNotNull("Widget not found on the workspace", widget);
         widget.launch(getAppPackageName());
         mLauncher.disableDebugTracing(); // b/289161193
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
index d40d3bc..4cdbd96 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplBindWidgetTest.java
@@ -53,6 +53,7 @@
 import com.android.launcher3.tapl.Workspace;
 import com.android.launcher3.ui.AbstractLauncherUiTest;
 import com.android.launcher3.ui.TestViewHelpers;
+import com.android.launcher3.util.TestUtil;
 import com.android.launcher3.util.rule.ShellCommandRule;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
@@ -233,13 +234,15 @@
     }
 
     private void verifyWidgetPresent(LauncherAppWidgetProviderInfo info) {
-        final Widget widget = mLauncher.getWorkspace().tryGetWidget(info.label, DEFAULT_UI_TIMEOUT);
+        final Widget widget = mLauncher.getWorkspace().tryGetWidget(info.label,
+                TestUtil.DEFAULT_UI_TIMEOUT);
         assertTrue("Widget is not present",
                 widget != null);
     }
 
     private void verifyPendingWidgetPresent() {
-        final Widget widget = mLauncher.getWorkspace().tryGetPendingWidget(DEFAULT_UI_TIMEOUT);
+        final Widget widget = mLauncher.getWorkspace().tryGetPendingWidget(
+                TestUtil.DEFAULT_UI_TIMEOUT);
         assertTrue("Pending widget is not present",
                 widget != null);
     }
diff --git a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java b/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
index 35c7cab..fe3b2ee 100644
--- a/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/TaplRequestPinItemTest.java
@@ -169,8 +169,7 @@
 
         // Go back to home
         mLauncher.goHome();
-        Wait.atMost("", new ItemSearchCondition(itemMatcher), DEFAULT_ACTIVITY_TIMEOUT,
-                mLauncher);
+        Wait.atMost("", new ItemSearchCondition(itemMatcher), mLauncher);
     }
 
     /**
diff --git a/tests/src/com/android/launcher3/util/Wait.java b/tests/src/com/android/launcher3/util/Wait.java
deleted file mode 100644
index 50bc32e..0000000
--- a/tests/src/com/android/launcher3/util/Wait.java
+++ /dev/null
@@ -1,66 +0,0 @@
-package com.android.launcher3.util;
-
-import android.os.SystemClock;
-import android.util.Log;
-
-import com.android.launcher3.tapl.LauncherInstrumentation;
-
-import org.junit.Assert;
-
-import java.util.function.Supplier;
-
-/**
- * A utility class for waiting for a condition to be true.
- */
-public class Wait {
-
-    private static final long DEFAULT_SLEEP_MS = 200;
-
-    public static void atMost(String message, Condition condition, long timeout,
-            LauncherInstrumentation launcher) {
-        atMost(() -> message, condition, timeout, DEFAULT_SLEEP_MS, launcher);
-    }
-
-    public static void atMost(Supplier<String> message, Condition condition, long timeout,
-            LauncherInstrumentation launcher) {
-        atMost(message, condition, timeout, DEFAULT_SLEEP_MS, launcher);
-    }
-
-    public static void atMost(Supplier<String> message, Condition condition, long timeout,
-            long sleepMillis,
-            LauncherInstrumentation launcher) {
-        final long startTime = SystemClock.uptimeMillis();
-        long endTime = startTime + timeout;
-        Log.d("Wait", "atMost: " + startTime + " - " + endTime);
-        while (SystemClock.uptimeMillis() < endTime) {
-            try {
-                if (condition.isTrue()) {
-                    return;
-                }
-            } catch (Throwable t) {
-                throw new RuntimeException(t);
-            }
-            SystemClock.sleep(sleepMillis);
-        }
-
-        // Check once more before returning false.
-        try {
-            if (condition.isTrue()) {
-                return;
-            }
-        } catch (Throwable t) {
-            throw new RuntimeException(t);
-        }
-        Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis());
-        launcher.checkForAnomaly(false, false);
-        Assert.fail(message.get());
-    }
-
-    /**
-     * Interface representing a generic condition
-     */
-    public interface Condition {
-
-        boolean isTrue() throws Throwable;
-    }
-}
diff --git a/tests/src/com/android/launcher3/util/Wait.kt b/tests/src/com/android/launcher3/util/Wait.kt
new file mode 100644
index 0000000..1e5af54
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/Wait.kt
@@ -0,0 +1,79 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.util
+
+import android.os.SystemClock
+import android.util.Log
+import com.android.launcher3.tapl.LauncherInstrumentation
+import java.util.function.Supplier
+import org.junit.Assert
+
+/** A utility class for waiting for a condition to be true. */
+object Wait {
+    private const val DEFAULT_SLEEP_MS: Long = 200
+
+    @JvmStatic
+    @JvmOverloads
+    fun atMost(
+        message: String,
+        condition: Condition,
+        launcherInstrumentation: LauncherInstrumentation? = null,
+        timeout: Long = TestUtil.DEFAULT_UI_TIMEOUT,
+    ) {
+        atMost({ message }, condition, launcherInstrumentation, timeout)
+    }
+
+    @JvmStatic
+    @JvmOverloads
+    fun atMost(
+        message: Supplier<String>,
+        condition: Condition,
+        launcherInstrumentation: LauncherInstrumentation? = null,
+        timeout: Long = TestUtil.DEFAULT_UI_TIMEOUT,
+    ) {
+        val startTime = SystemClock.uptimeMillis()
+        val endTime = startTime + timeout
+        Log.d("Wait", "atMost: $startTime - $endTime")
+        while (SystemClock.uptimeMillis() < endTime) {
+            try {
+                if (condition.isTrue()) {
+                    return
+                }
+            } catch (t: Throwable) {
+                throw RuntimeException(t)
+            }
+            SystemClock.sleep(DEFAULT_SLEEP_MS)
+        }
+
+        // Check once more before returning false.
+        try {
+            if (condition.isTrue()) {
+                return
+            }
+        } catch (t: Throwable) {
+            throw RuntimeException(t)
+        }
+        Log.d("Wait", "atMost: timed out: " + SystemClock.uptimeMillis())
+        launcherInstrumentation?.checkForAnomaly(false, false)
+        Assert.fail(message.get())
+    }
+
+    /** Interface representing a generic condition */
+    fun interface Condition {
+
+        @Throws(Throwable::class) fun isTrue(): Boolean
+    }
+}