Merge "Revert "Reset the frozen recents list state when switching to bu..."" into main
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index b024418..f32f204 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -306,7 +306,7 @@
     <dimen name="taskbar_home_button_left_margin_kids">48dp</dimen>
     <dimen name="taskbar_icon_size_kids">32dp</dimen>
     <dimen name="taskbar_all_apps_button_translation_x_offset">6dp</dimen>
-
+    <dimen name="taskbar_all_apps_search_button_translation_x_offset">4.5dp</dimen>
 
     <!-- Transient taskbar -->
     <dimen name="transient_taskbar_padding">12dp</dimen>
diff --git a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
index d7a4f76..42e6809 100644
--- a/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
+++ b/quickstep/src/com/android/launcher3/statehandlers/DesktopVisibilityController.java
@@ -160,7 +160,9 @@
             if (mInOverviewState) {
                 setLauncherViewsVisibility(View.VISIBLE);
                 markLauncherResumed();
-            } else if (mFreeformTasksVisible) {
+            } else if (mFreeformTasksVisible && !mGestureInProgress) {
+                // Switching out of overview state and gesture finished.
+                // If freeform tasks are still visible, hide launcher again.
                 setLauncherViewsVisibility(View.INVISIBLE);
                 markLauncherPaused();
             }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 0b83a88..80ac9be 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -27,11 +27,13 @@
 import static com.android.launcher3.AbstractFloatingView.TYPE_REBIND_SAFE;
 import static com.android.launcher3.AbstractFloatingView.TYPE_TASKBAR_OVERLAY_PROXY;
 import static com.android.launcher3.Utilities.isRunningInTestHarness;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_NO_RECREATION;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_OPEN;
 import static com.android.launcher3.taskbar.TaskbarAutohideSuspendController.FLAG_AUTOHIDE_SUSPEND_DRAGGING;
 import static com.android.launcher3.taskbar.TaskbarAutohideSuspendController.FLAG_AUTOHIDE_SUSPEND_FULLSCREEN;
 import static com.android.launcher3.taskbar.TaskbarManager.FLAG_HIDE_NAVBAR_WINDOW;
 import static com.android.launcher3.testing.shared.ResourceUtils.getBoolByName;
+import static com.android.launcher3.util.VibratorWrapper.EFFECT_CLICK;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_VOICE_INTERACTION_WINDOW_SHOWING;
 
@@ -57,6 +59,7 @@
 import android.view.RoundedCorner;
 import android.view.Surface;
 import android.view.View;
+import android.view.WindowInsets;
 import android.view.WindowManager;
 import android.widget.FrameLayout;
 import android.widget.Toast;
@@ -108,6 +111,7 @@
 import com.android.launcher3.util.SettingsCache;
 import com.android.launcher3.util.SplitConfigurationOptions.SplitSelectSource;
 import com.android.launcher3.util.TraceHelper;
+import com.android.launcher3.util.VibratorWrapper;
 import com.android.launcher3.util.ViewCache;
 import com.android.launcher3.views.ActivityContext;
 import com.android.quickstep.views.RecentsView;
@@ -323,11 +327,11 @@
             mIsDestroyed = false;
         }
 
-        if (!mAddedWindow) {
+        if (!ENABLE_TASKBAR_NO_RECREATION.get() && !mAddedWindow) {
             mWindowManager.addView(mDragLayer, mWindowLayoutParams);
             mAddedWindow = true;
         } else {
-            mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+            notifyUpdateLayoutParams();
         }
     }
 
@@ -674,7 +678,7 @@
         mIsDestroyed = true;
         setUIController(TaskbarUIController.DEFAULT);
         mControllers.onDestroy();
-        if (!FLAG_HIDE_NAVBAR_WINDOW) {
+        if (!ENABLE_TASKBAR_NO_RECREATION.get() && !FLAG_HIDE_NAVBAR_WINDOW) {
             mWindowManager.removeViewImmediate(mDragLayer);
             mAddedWindow = false;
         }
@@ -806,7 +810,7 @@
         }
         mWindowLayoutParams.height = height;
         mControllers.taskbarInsetsController.onTaskbarOrBubblebarWindowHeightOrInsetsChanged();
-        mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+        notifyUpdateLayoutParams();
     }
 
     /**
@@ -849,7 +853,22 @@
         } else {
             mWindowLayoutParams.flags |= FLAG_NOT_FOCUSABLE;
         }
-        mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+        notifyUpdateLayoutParams();
+    }
+
+    /**
+     * Applies forcibly show flag to taskbar window iff transient taskbar is unstashed.
+     */
+    public void applyForciblyShownFlagWhileTransientTaskbarUnstashed(boolean shouldForceShow) {
+        if (!DisplayController.isTransientTaskbar(this)) {
+            return;
+        }
+        if (shouldForceShow) {
+            mWindowLayoutParams.forciblyShownTypes |= WindowInsets.Type.navigationBars();
+        } else {
+            mWindowLayoutParams.forciblyShownTypes &= ~WindowInsets.Type.navigationBars();
+        }
+        notifyUpdateLayoutParams();
     }
 
     /**
@@ -958,8 +977,8 @@
                         }
 
                     } catch (NullPointerException
-                            | ActivityNotFoundException
-                            | SecurityException e) {
+                             | ActivityNotFoundException
+                             | SecurityException e) {
                         Toast.makeText(this, R.string.activity_not_found, Toast.LENGTH_SHORT)
                                 .show();
                         Log.e(TAG, "Unable to launch. tag=" + info + " intent=" + intent, e);
@@ -1053,6 +1072,7 @@
 
     /**
      * Called when we detect a long press in the nav region before passing the gesture slop.
+     *
      * @return Whether taskbar handled the long press, and thus should cancel the gesture.
      */
     public boolean onLongPressToUnstashTaskbar() {
@@ -1063,6 +1083,7 @@
      * Called when we want to unstash taskbar when user performs swipes up gesture.
      */
     public void onSwipeToUnstashTaskbar() {
+        VibratorWrapper.INSTANCE.get(this).vibrate(EFFECT_CLICK);
         mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(/* stash= */ false);
         mControllers.taskbarEduTooltipController.hide();
     }
@@ -1121,7 +1142,7 @@
      * Called when we detect a motion down or up/cancel in the nav region while stashed.
      *
      * @param animateForward Whether to animate towards the unstashed hint state or back to stashed.
-     * @param forceUnstash Whether we force the unstash hint.
+     * @param forceUnstash   Whether we force the unstash hint.
      */
     public void startTaskbarUnstashHint(boolean animateForward, boolean forceUnstash) {
         // TODO(b/270395798): Clean up forceUnstash after removing long-press unstashing code.
@@ -1229,12 +1250,16 @@
             mWindowLayoutParams.privateFlags &=
                     ~WindowManager.LayoutParams.PRIVATE_FLAG_EXCLUDE_FROM_SCREEN_MAGNIFICATION;
         }
-        mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+        notifyUpdateLayoutParams();
     }
 
     void notifyUpdateLayoutParams() {
         if (mDragLayer.isAttachedToWindow()) {
-            mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+            if (ENABLE_TASKBAR_NO_RECREATION.get()) {
+                mWindowManager.updateViewLayout(mDragLayer.getRootView(), mWindowLayoutParams);
+            } else {
+                mWindowManager.updateViewLayout(mDragLayer, mWindowLayoutParams);
+            }
         }
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
index c51a7ec..12fa17c 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarInsetsController.kt
@@ -15,9 +15,9 @@
  */
 package com.android.launcher3.taskbar
 
-import android.inputmethodservice.InputMethodService.ENABLE_HIDE_IME_CAPTION_BAR
 import android.graphics.Insets
 import android.graphics.Region
+import android.inputmethodservice.InputMethodService.ENABLE_HIDE_IME_CAPTION_BAR
 import android.os.Binder
 import android.os.IBinder
 import android.view.Gravity
@@ -41,6 +41,7 @@
 import com.android.launcher3.DeviceProfile
 import com.android.launcher3.R
 import com.android.launcher3.anim.AlphaUpdateListener
+import com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_NO_RECREATION
 import com.android.launcher3.taskbar.TaskbarControllers.LoggableTaskbarController
 import com.android.launcher3.util.DisplayController
 import java.io.PrintWriter
@@ -97,11 +98,18 @@
                 0
             }
 
-        windowLayoutParams.providedInsets = getProvidedInsets(insetsRoundedCornerFlag)
+        windowLayoutParams.providedInsets =
+            if (ENABLE_TASKBAR_NO_RECREATION.get()) {
+                getProvidedInsets(controllers.sharedState!!.insetsFrameProviders!!,
+                        insetsRoundedCornerFlag)
+            } else {
+                getProvidedInsets(insetsRoundedCornerFlag)
+            }
+
         if (!context.isGestureNav) {
             if (windowLayoutParams.paramsForRotation != null) {
                 for (layoutParams in windowLayoutParams.paramsForRotation) {
-                    layoutParams.providedInsets = getProvidedInsets(insetsRoundedCornerFlag)
+                    layoutParams.providedInsets = windowLayoutParams.providedInsets
                 }
             }
         }
@@ -154,6 +162,26 @@
     }
 
     /**
+     * This is for when ENABLE_TASKBAR_NO_RECREATION is enabled. We generate one instance of
+     * providedInsets and use it across the entire lifecycle of TaskbarManager. The only thing
+     * we need to reset is nav bar flags based on insetsRoundedCornerFlag.
+     */
+    private fun getProvidedInsets(providedInsets: Array<InsetsFrameProvider>,
+                                  insetsRoundedCornerFlag: Int): Array<InsetsFrameProvider> {
+        val navBarsFlag =
+                (if (context.isGestureNav) FLAG_SUPPRESS_SCRIM else 0) or insetsRoundedCornerFlag
+        for (provider in providedInsets) {
+            if (provider.type == navigationBars()) {
+                provider.setFlags(
+                        navBarsFlag,
+                        FLAG_SUPPRESS_SCRIM or FLAG_INSETS_ROUNDED_CORNER
+                )
+            }
+        }
+        return providedInsets
+    }
+
+    /**
      * The inset types and number of insets provided have to match for both gesture nav and button
      * nav. The values and the order of the elements in array are allowed to differ.
      * Reason being WM does not allow types and number of insets changing for a given window once it
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
index 90f7bea..88ae349 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarLauncherStateController.java
@@ -18,6 +18,7 @@
 import static com.android.app.animation.Interpolators.EMPHASIZED;
 import static com.android.launcher3.taskbar.TaskbarKeyguardController.MASK_ANY_SYSUI_LOCKED;
 import static com.android.launcher3.taskbar.TaskbarStashController.FLAG_IN_APP;
+import static com.android.launcher3.taskbar.TaskbarStashController.FLAG_IN_OVERVIEW;
 import static com.android.launcher3.taskbar.TaskbarStashController.FLAG_IN_STASHED_LAUNCHER_STATE;
 import static com.android.launcher3.taskbar.TaskbarViewController.ALPHA_INDEX_HOME;
 import static com.android.launcher3.util.FlagDebugUtils.appendFlag;
@@ -416,6 +417,9 @@
             controllers.bubbleStashController.setBubblesShowingOnOverview(onOverview);
         });
 
+        mControllers.taskbarStashController.updateStateForFlag(FLAG_IN_OVERVIEW,
+                mLauncherState == LauncherState.OVERVIEW);
+
         AnimatorSet animatorSet = new AnimatorSet();
 
         if (hasAnyFlag(changedFlags, FLAG_LAUNCHER_IN_STATE_TRANSITION)) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index c423fb3..55d56f0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -23,6 +23,7 @@
 import static com.android.launcher3.LauncherPrefs.TASKBAR_PINNING;
 import static com.android.launcher3.LauncherPrefs.TASKBAR_PINNING_KEY;
 import static com.android.launcher3.LauncherState.OVERVIEW;
+import static com.android.launcher3.config.FeatureFlags.ENABLE_TASKBAR_NO_RECREATION;
 import static com.android.launcher3.util.DisplayController.TASKBAR_NOT_DESTROYED_TAG;
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.FlagDebugUtils.formatFlagChange;
@@ -45,6 +46,8 @@
 import android.provider.Settings;
 import android.util.Log;
 import android.view.Display;
+import android.view.WindowManager;
+import android.widget.FrameLayout;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -64,6 +67,7 @@
 import com.android.quickstep.RecentsActivity;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TouchInteractionService;
+import com.android.quickstep.util.AssistUtils;
 import com.android.systemui.shared.system.QuickStepContract;
 import com.android.systemui.unfold.UnfoldTransitionProgressProvider;
 import com.android.systemui.unfold.util.ScopedUnfoldTransitionProgressProvider;
@@ -104,6 +108,9 @@
             Settings.Secure.NAV_BAR_KIDS_MODE);
 
     private final Context mContext;
+    private WindowManager mWindowManager;
+    private FrameLayout mTaskbarRootLayout;
+    private boolean mAddedWindow;
     private final TaskbarNavButtonController mNavButtonController;
     private final ComponentCallbacks mComponentCallbacks;
 
@@ -175,8 +182,13 @@
         Display display =
                 service.getSystemService(DisplayManager.class).getDisplay(DEFAULT_DISPLAY);
         mContext = service.createWindowContext(display, TYPE_NAVIGATION_BAR_PANEL, null);
+        if (ENABLE_TASKBAR_NO_RECREATION.get()) {
+            mWindowManager = mContext.getSystemService(WindowManager.class);
+            mTaskbarRootLayout = new FrameLayout(mContext);
+        }
         mNavButtonController = new TaskbarNavButtonController(service,
-                SystemUiProxy.INSTANCE.get(mContext), new Handler());
+                SystemUiProxy.INSTANCE.get(mContext), new Handler(),
+                AssistUtils.newInstance(mContext));
         mComponentCallbacks = new ComponentCallbacks() {
             private Configuration mOldConfig = mContext.getResources().getConfiguration();
 
@@ -256,10 +268,15 @@
             LauncherPrefs.get(mContext).removeListener(mTaskbarPinningPreferenceChangeListener,
                     TASKBAR_PINNING);
             mTaskbarActivityContext.onDestroy();
-            if (!FLAG_HIDE_NAVBAR_WINDOW) {
+            if (!FLAG_HIDE_NAVBAR_WINDOW || ENABLE_TASKBAR_NO_RECREATION.get()) {
                 mTaskbarActivityContext = null;
             }
         }
+        DeviceProfile dp = mUserUnlocked ?
+                LauncherAppState.getIDP(mContext).getDeviceProfile(mContext) : null;
+        if (dp == null || !isTaskbarPresent(dp)) {
+            removeTaskbarRootViewFromWindow();
+        }
     }
 
     /**
@@ -308,6 +325,7 @@
         mUserUnlocked = true;
         LauncherAppState.getIDP(mContext).addOnChangeListener(mIdpChangeListener);
         recreateTaskbar();
+        addTaskbarRootViewToWindow();
     }
 
     /**
@@ -388,10 +406,9 @@
                 return;
             }
 
-            if (mTaskbarActivityContext == null) {
+            if (ENABLE_TASKBAR_NO_RECREATION.get() || mTaskbarActivityContext == null) {
                 mTaskbarActivityContext = new TaskbarActivityContext(mContext, dp,
-                    mNavButtonController,
-                    mUnfoldProgressProvider);
+                        mNavButtonController, mUnfoldProgressProvider);
             } else {
                 mTaskbarActivityContext.updateDeviceProfile(dp);
             }
@@ -402,6 +419,13 @@
                     createTaskbarUIControllerForActivity(mActivity));
             }
 
+            if (ENABLE_TASKBAR_NO_RECREATION.get()) {
+                addTaskbarRootViewToWindow();
+                mTaskbarRootLayout.removeAllViews();
+                mTaskbarRootLayout.addView(mTaskbarActivityContext.getDragLayer());
+                mTaskbarActivityContext.notifyUpdateLayoutParams();
+            }
+
             // We to wait until user unlocks the device to attach listener.
             LauncherPrefs.get(mContext).addListener(mTaskbarPinningPreferenceChangeListener,
                 TASKBAR_PINNING);
@@ -523,6 +547,22 @@
         }
     }
 
+    private void addTaskbarRootViewToWindow() {
+        if (ENABLE_TASKBAR_NO_RECREATION.get() && !mAddedWindow
+                && mTaskbarActivityContext != null) {
+            mWindowManager.addView(mTaskbarRootLayout,
+                    mTaskbarActivityContext.getWindowLayoutParams());
+            mAddedWindow = true;
+        }
+    }
+
+    private void removeTaskbarRootViewFromWindow() {
+        if (ENABLE_TASKBAR_NO_RECREATION.get() && mAddedWindow) {
+            mWindowManager.removeViewImmediate(mTaskbarRootLayout);
+            mAddedWindow = false;
+        }
+    }
+
     /** Temp logs for b/254119092. */
     public void debugWhyTaskbarNotDestroyed(String debugReason) {
         StringJoiner log = new StringJoiner("\n");
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
index 6d86b1e..533785f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarNavButtonController.java
@@ -109,15 +109,17 @@
     private final TouchInteractionService mService;
     private final SystemUiProxy mSystemUiProxy;
     private final Handler mHandler;
+    private final AssistUtils mAssistUtils;
     @Nullable private StatsLogManager mStatsLogManager;
 
     private final Runnable mResetLongPress = this::resetScreenUnpin;
 
     public TaskbarNavButtonController(TouchInteractionService service,
-            SystemUiProxy systemUiProxy, Handler handler) {
+            SystemUiProxy systemUiProxy, Handler handler, AssistUtils assistUtils) {
         mService = service;
         mSystemUiProxy = systemUiProxy;
         mHandler = handler;
+        mAssistUtils = assistUtils;
     }
 
     public void onButtonClick(@TaskbarButton int buttonType, View view) {
@@ -313,8 +315,7 @@
             return;
         }
         // Attempt to start Assist with AssistUtils, otherwise fall back to SysUi's implementation.
-        if (!AssistUtils.newInstance(mService.getApplicationContext()).tryStartAssistOverride(
-                INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS)) {
+        if (!mAssistUtils.tryStartAssistOverride(INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS)) {
             Bundle args = new Bundle();
             args.putInt(INVOCATION_TYPE_KEY, INVOCATION_TYPE_HOME_BUTTON_LONG_PRESS);
             mSystemUiProxy.startAssistant(args);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
index 66ca7d9..abbd18b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarSharedState.java
@@ -15,15 +15,28 @@
  */
 package com.android.launcher3.taskbar;
 
+import static android.view.InsetsFrameProvider.SOURCE_DISPLAY;
+import static android.view.WindowInsets.Type.mandatorySystemGestures;
+import static android.view.WindowInsets.Type.navigationBars;
+import static android.view.WindowInsets.Type.systemGestures;
+import static android.view.WindowInsets.Type.tappableElement;
+
 import static com.android.launcher3.taskbar.LauncherTaskbarUIController.DISPLAY_PROGRESS_COUNT;
 
 import android.app.PendingIntent;
+import android.os.Binder;
+import android.os.IBinder;
+import android.view.InsetsFrameProvider;
 
 /**
  * State shared across different taskbar instance
  */
 public class TaskbarSharedState {
 
+    private final IBinder mInsetsOwner = new Binder();
+    private static int INDEX_LEFT = 0;
+    private static int INDEX_RIGHT = 1;
+
     // TaskbarManager#onSystemUiFlagsChanged
     public int sysuiStateFlags;
 
@@ -48,4 +61,14 @@
 
     // Taskbar System Action
     public PendingIntent taskbarSystemActionPendingIntent;
+
+    public final InsetsFrameProvider[] insetsFrameProviders = new InsetsFrameProvider[] {
+            new InsetsFrameProvider(mInsetsOwner, 0, navigationBars()),
+            new InsetsFrameProvider(mInsetsOwner, 0, tappableElement()),
+            new InsetsFrameProvider(mInsetsOwner, 0, mandatorySystemGestures()),
+            new InsetsFrameProvider(mInsetsOwner, INDEX_LEFT, systemGestures())
+                    .setSource(SOURCE_DISPLAY),
+            new InsetsFrameProvider(mInsetsOwner, INDEX_RIGHT, systemGestures())
+                    .setSource(SOURCE_DISPLAY)
+    };
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index b5b453b..7b12733 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -33,6 +33,7 @@
 import static com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR;
 import static com.android.launcher3.util.FlagDebugUtils.appendFlag;
 import static com.android.launcher3.util.FlagDebugUtils.formatFlagChange;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_SHOWING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE;
@@ -96,6 +97,7 @@
     public static final int FLAG_STASHED_IN_APP_AUTO = 1 << 9; // Autohide (transient taskbar).
     public static final int FLAG_STASHED_SYSUI = 1 << 10; //  app pinning,...
     public static final int FLAG_STASHED_DEVICE_LOCKED = 1 << 11; // device is locked: keyguard, ...
+    public static final int FLAG_IN_OVERVIEW = 1 << 12; // launcher is in overview
 
     // If any of these flags are enabled, isInApp should return true.
     private static final int FLAGS_IN_APP = FLAG_IN_APP | FLAG_IN_SETUP;
@@ -1015,9 +1017,11 @@
 
         updateStateForFlag(FLAG_STASHED_IN_APP_SYSUI, hasAnyFlag(systemUiStateFlags,
                 SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE));
-        updateStateForFlag(FLAG_STASHED_SYSUI,
-                hasAnyFlag(systemUiStateFlags, SYSUI_STATE_SCREEN_PINNING));
 
+        boolean bubblesOnOverview = hasAnyFlag(FLAG_IN_OVERVIEW)
+                && hasAnyFlag(systemUiStateFlags, SYSUI_STATE_BUBBLES_EXPANDED);
+        updateStateForFlag(FLAG_STASHED_SYSUI,
+                hasAnyFlag(systemUiStateFlags, SYSUI_STATE_SCREEN_PINNING) || bubblesOnOverview);
         boolean isLocked = hasAnyFlag(systemUiStateFlags, MASK_ANY_SYSUI_LOCKED)
                 && !hasAnyFlag(systemUiStateFlags, SYSUI_STATE_STATUS_BAR_KEYGUARD_GOING_AWAY);
         updateStateForFlag(FLAG_STASHED_DEVICE_LOCKED, isLocked);
@@ -1097,6 +1101,7 @@
                     TaskbarAutohideSuspendController.FLAG_AUTOHIDE_SUSPEND_TRANSIENT_TASKBAR,
                     !hasAnyFlag(FLAG_STASHED_IN_APP_AUTO));
         }
+        mActivity.applyForciblyShownFlagWhileTransientTaskbarUnstashed(!isStashedInApp());
     }
 
     private void notifyStashChange(boolean visible, boolean stashed) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index fa5a1ae..0e5ab71 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -18,6 +18,7 @@
 import static android.content.pm.PackageManager.FEATURE_PC;
 import static android.view.accessibility.AccessibilityEvent.TYPE_WINDOW_CONTENT_CHANGED;
 
+import static com.android.launcher3.config.FeatureFlags.ENABLE_ALL_APPS_SEARCH_IN_TASKBAR;
 import static com.android.launcher3.config.FeatureFlags.ENABLE_CURSOR_HOVER_STATES;
 import static com.android.launcher3.icons.IconNormalizer.ICON_VISIBLE_AREA_FACTOR;
 
@@ -33,6 +34,8 @@
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.FrameLayout;
 
+import androidx.annotation.DimenRes;
+import androidx.annotation.DrawableRes;
 import androidx.annotation.LayoutRes;
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
@@ -90,11 +93,11 @@
     // Only non-null when device supports having an All Apps button.
     private @Nullable IconButtonView mTaskbarDivider;
 
-    private View mQsb;
+    private final View mQsb;
 
-    private float mTransientTaskbarMinWidth;
+    private final float mTransientTaskbarMinWidth;
 
-    private float mTransientTaskbarAllAppsButtonTranslationXOffset;
+    private final float mTaskbarAllAppsButtonTranslationXOffset;
 
     private boolean mShouldTryStartAlign;
 
@@ -120,12 +123,9 @@
         boolean isTransientTaskbar = DisplayController.isTransientTaskbar(mActivityContext)
                 && !TaskbarManager.isPhoneMode(mActivityContext.getDeviceProfile());
         mIsRtl = Utilities.isRtl(resources);
-        mTransientTaskbarMinWidth = mContext.getResources().getDimension(
-                R.dimen.transient_taskbar_min_width);
-        mTransientTaskbarAllAppsButtonTranslationXOffset =
-                resources.getDimension(isTransientTaskbar
-                        ? R.dimen.transient_taskbar_all_apps_button_translation_x_offset
-                        : R.dimen.taskbar_all_apps_button_translation_x_offset);
+        mTransientTaskbarMinWidth = resources.getDimension(R.dimen.transient_taskbar_min_width);
+        mTaskbarAllAppsButtonTranslationXOffset =
+                resources.getDimension(getAllAppsButtonTranslationXOffset(isTransientTaskbar));
 
         onDeviceProfileChanged(mActivityContext.getDeviceProfile());
 
@@ -149,9 +149,8 @@
         if (!mActivityContext.getPackageManager().hasSystemFeature(FEATURE_PC)) {
             mAllAppsButton = (IconButtonView) LayoutInflater.from(context)
                     .inflate(R.layout.taskbar_all_apps_button, this, false);
-            mAllAppsButton.setIconDrawable(resources.getDrawable(isTransientTaskbar
-                    ? R.drawable.ic_transient_taskbar_all_apps_button
-                    : R.drawable.ic_taskbar_all_apps_button));
+            mAllAppsButton.setIconDrawable(resources.getDrawable(
+                    getAllAppsButton(isTransientTaskbar)));
             mAllAppsButton.setScaleX(mIsRtl ? -1 : 1);
             mAllAppsButton.setPadding(mItemPadding, mItemPadding, mItemPadding, mItemPadding);
             mAllAppsButton.setForegroundTint(
@@ -171,6 +170,30 @@
         mQsb = LayoutInflater.from(context).inflate(R.layout.search_container_hotseat, this, false);
     }
 
+    @DrawableRes
+    private int getAllAppsButton(boolean isTransientTaskbar) {
+        if (ENABLE_ALL_APPS_SEARCH_IN_TASKBAR.get()) {
+            return isTransientTaskbar
+                    ? R.drawable.ic_transient_taskbar_all_apps_search_button
+                    : R.drawable.ic_taskbar_all_apps_search_button;
+        } else {
+            return isTransientTaskbar
+                    ? R.drawable.ic_transient_taskbar_all_apps_button
+                    : R.drawable.ic_taskbar_all_apps_button;
+        }
+    }
+
+    @DimenRes
+    private int getAllAppsButtonTranslationXOffset(boolean isTransientTaskbar) {
+        if (isTransientTaskbar) {
+            return R.dimen.transient_taskbar_all_apps_button_translation_x_offset;
+        } else {
+            return ENABLE_ALL_APPS_SEARCH_IN_TASKBAR.get()
+                    ? R.dimen.taskbar_all_apps_search_button_translation_x_offset
+                    : R.dimen.taskbar_all_apps_button_translation_x_offset;
+        }
+    }
+
     @Override
     protected void onAttachedToWindow() {
         super.onAttachedToWindow();
@@ -336,7 +359,7 @@
 
         if (mAllAppsButton != null) {
             mAllAppsButton.setTranslationXForTaskbarAllAppsIcon(getChildCount() > 0
-                    ? mTransientTaskbarAllAppsButtonTranslationXOffset : 0f);
+                    ? mTaskbarAllAppsButtonTranslationXOffset : 0f);
             addView(mAllAppsButton, mIsRtl ? getChildCount() : 0);
 
             // if only all apps button present, don't include divider view.
diff --git a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsSlideInView.java b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsSlideInView.java
index 537d2c6..001c3bc 100644
--- a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsSlideInView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsSlideInView.java
@@ -21,12 +21,16 @@
 import android.content.Context;
 import android.graphics.Canvas;
 import android.graphics.Rect;
+import android.os.Handler;
+import android.os.Looper;
 import android.util.AttributeSet;
 import android.view.MotionEvent;
 import android.view.View;
 import android.view.animation.Interpolator;
 import android.window.OnBackInvokedDispatcher;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Insettable;
 import com.android.launcher3.R;
@@ -40,8 +44,11 @@
 /** Wrapper for taskbar all apps with slide-in behavior. */
 public class TaskbarAllAppsSlideInView extends AbstractSlideInView<TaskbarOverlayContext>
         implements Insettable, DeviceProfile.OnDeviceProfileChangeListener {
+    private final Handler mHandler;
+
     private TaskbarAllAppsContainerView mAppsView;
     private float mShiftRange;
+    private @Nullable Runnable mShowOnFullyAttachedToWindowRunnable;
 
     // Initialized in init.
     private TaskbarAllAppsCallbacks mAllAppsCallbacks;
@@ -53,6 +60,7 @@
     public TaskbarAllAppsSlideInView(Context context, AttributeSet attrs,
             int defStyleAttr) {
         super(context, attrs, defStyleAttr);
+        mHandler = new Handler(Looper.myLooper());
     }
 
     void init(TaskbarAllAppsCallbacks callbacks) {
@@ -65,14 +73,14 @@
             return;
         }
         mIsOpen = true;
-        attachToContainer();
 
         addOnAttachStateChangeListener(new OnAttachStateChangeListener() {
             @Override
             public void onViewAttachedToWindow(View v) {
                 removeOnAttachStateChangeListener(this);
                 // Wait for view and its descendants to be fully attached before starting open.
-                post(() -> showOnFullyAttachedToWindow(animate));
+                mShowOnFullyAttachedToWindowRunnable = () -> showOnFullyAttachedToWindow(animate);
+                mHandler.post(mShowOnFullyAttachedToWindowRunnable);
             }
 
             @Override
@@ -80,6 +88,7 @@
                 removeOnAttachStateChangeListener(this);
             }
         });
+        attachToContainer();
     }
 
     private void showOnFullyAttachedToWindow(boolean animate) {
@@ -114,6 +123,10 @@
 
     @Override
     protected void handleClose(boolean animate) {
+        if (mShowOnFullyAttachedToWindowRunnable != null) {
+            mHandler.removeCallbacks(mShowOnFullyAttachedToWindowRunnable);
+            mShowOnFullyAttachedToWindowRunnable = null;
+        }
         if (mIsOpen) {
             mAllAppsCallbacks.onAllAppsTransitionStart(false);
         }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index 24db380..90f4748 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -137,6 +137,7 @@
     private static class BubbleBarViewUpdate {
         boolean expandedChanged;
         boolean expanded;
+        boolean shouldShowEducation;
         String selectedBubbleKey;
         String suppressedBubbleKey;
         String unsuppressedBubbleKey;
@@ -151,6 +152,7 @@
         BubbleBarViewUpdate(BubbleBarUpdate update) {
             expandedChanged = update.expandedChanged;
             expanded = update.expanded;
+            shouldShowEducation = update.shouldShowEducation;
             selectedBubbleKey = update.selectedBubbleKey;
             suppressedBubbleKey = update.suppressedBubbleKey;
             unsuppressedBubbleKey = update.unsupressedBubbleKey;
@@ -366,7 +368,9 @@
                 mBubbleStashController.animateToInitialState(update.expanded);
             }
         }
-
+        if (update.shouldShowEducation) {
+            mBubbleBarViewController.prepareToShowEducation();
+        }
         if (update.expandedChanged) {
             if (update.expanded != mBubbleBarViewController.isExpanded()) {
                 mBubbleBarViewController.setExpandedFromSysui(update.expanded);
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
index 20b8e3b..5c607c9 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarViewController.java
@@ -18,6 +18,7 @@
 import static android.view.View.INVISIBLE;
 import static android.view.View.VISIBLE;
 
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.util.Log;
 import android.view.MotionEvent;
@@ -75,6 +76,7 @@
     private boolean mHiddenForSysui;
     // Whether the bar is hidden because there are no bubbles.
     private boolean mHiddenForNoBubbles;
+    private boolean mShouldShowEducation;
 
     public BubbleBarViewController(TaskbarActivityContext activity, BubbleBarView barView) {
         mActivity = activity;
@@ -98,7 +100,7 @@
         mBarView.getLayoutParams().height = mActivity.getDeviceProfile().taskbarHeight;
         mBubbleBarScale.updateValue(1f);
         mBubbleClickListener = v -> onBubbleClicked(v);
-        mBubbleBarClickListener = v -> setExpanded(true);
+        mBubbleBarClickListener = v -> onBubbleBarClicked();
         mBubbleDragController.setupBubbleBarView(mBarView);
         mBarView.setOnClickListener(mBubbleBarClickListener);
         mBarView.addOnLayoutChangeListener((view, i, i1, i2, i3, i4, i5, i6, i7) ->
@@ -121,6 +123,21 @@
         }
     }
 
+    private void onBubbleBarClicked() {
+        if (mShouldShowEducation) {
+            mShouldShowEducation = false;
+            // Get the bubble bar bounds on screen
+            Rect bounds = new Rect();
+            mBarView.getBoundsOnScreen(bounds);
+            // Calculate user education reference position in Screen coordinates
+            Point position = new Point(bounds.centerX(), bounds.top);
+            // Show user education relative to the reference point
+            mSystemUiProxy.showUserEducation(position);
+        } else {
+            setExpanded(true);
+        }
+    }
+
     //
     // The below animators are exposed to BubbleStashController so it can manage the stashing
     // animation.
@@ -326,6 +343,12 @@
         }
     }
 
+    /** Marks as should show education and shows the bubble bar in a collapsed state */
+    public void prepareToShowEducation() {
+        mShouldShowEducation = true;
+        mBubbleStashController.showBubbleBar(false /* expand the bubbles */);
+    }
+
     /**
      * Updates the dragged bubble view in the bubble bar view, and notifies SystemUI
      * that a bubble is being dragged to dismiss.
diff --git a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
index d4e2be9..9126c4b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
@@ -62,12 +62,30 @@
         @Override
         public void onTaskCreated(int taskId, ComponentName componentName) {
             // Created task will be below existing overlay, so move out of the way.
-            hideWindow();
+            hideWindowOnTaskStackChange();
         }
 
         @Override
         public void onTaskMovedToFront(int taskId) {
             // New front task will be below existing overlay, so move out of the way.
+            hideWindowOnTaskStackChange();
+        }
+
+        @Override
+        public void onTaskStackChanged() {
+            // The other callbacks are insufficient for All Apps, because there are many cases where
+            // it can relaunch the same task already behind it. However, this callback needs to be a
+            // no-op when only EDU is shown, because going between the EDU steps invokes this
+            // callback.
+            if (mControllers.getSharedState() != null
+                    && mControllers.getSharedState().allAppsVisible) {
+                hideWindowOnTaskStackChange();
+            }
+        }
+
+        private void hideWindowOnTaskStackChange() {
+            // A task was launched while overlay window was open, so stash Taskbar.
+            mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(true);
             hideWindow();
         }
     };
@@ -199,8 +217,10 @@
 
         @Override
         protected void handleClose(boolean animate) {
-            mTaskbarContext.getDragLayer().removeView(this);
-            Optional.ofNullable(mOverlayContext).ifPresent(c -> closeAllOpenViews(c, animate));
+            if (mIsOpen) {
+                mTaskbarContext.getDragLayer().removeView(this);
+                Optional.ofNullable(mOverlayContext).ifPresent(c -> closeAllOpenViews(c, animate));
+            }
         }
 
         @Override
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
index f7bef03..f66bc60 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepWidgetHolder.java
@@ -37,6 +37,7 @@
 import com.android.launcher3.widget.LauncherAppWidgetHostView;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.LauncherWidgetHolder;
+import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -237,6 +238,14 @@
     @Override
     public LauncherAppWidgetHostView createView(@NonNull Context context, int appWidgetId,
             @NonNull LauncherAppWidgetProviderInfo appWidget) {
+
+        if (appWidget.isCustomWidget()) {
+            LauncherAppWidgetHostView lahv = new LauncherAppWidgetHostView(context);
+            lahv.setAppWidget(appWidgetId, appWidget);
+            CustomWidgetManager.INSTANCE.get(context).onViewCreated(lahv);
+            return lahv;
+        }
+
         LauncherAppWidgetHostView widgetView = getPendingView(appWidgetId);
         if (widgetView != null) {
             removePendingView(appWidgetId);
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index 99dd634..5da41ae 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -31,6 +31,7 @@
 import android.content.pm.ActivityInfo;
 import android.content.pm.PackageManager;
 import android.content.pm.ShortcutInfo;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.Bundle;
 import android.os.Handler;
@@ -731,6 +732,18 @@
         }
     }
 
+    /**
+     * Tells SysUI to show user education relative to the reference point provided.
+     * @param position the bubble bar top center position in Screen coordinates.
+     */
+    public void showUserEducation(Point position) {
+        try {
+            mBubbles.showUserEducation(position.x, position.y);
+        } catch (RemoteException e) {
+            Log.w(TAG, "Failed call showUserEducation");
+        }
+    }
+
     //
     // Splitscreen
     //
diff --git a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
index eb7598d..437009f 100644
--- a/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/LauncherRecentsView.java
@@ -16,7 +16,6 @@
 package com.android.quickstep.views;
 
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
-
 import static com.android.launcher3.LauncherState.ALL_APPS;
 import static com.android.launcher3.LauncherState.CLEAR_ALL_BUTTON;
 import static com.android.launcher3.LauncherState.EDIT_MODE;
@@ -136,6 +135,10 @@
     @Override
     public void onStateTransitionStart(LauncherState toState) {
         setOverviewStateEnabled(toState.overviewUi);
+        if (toState.overviewUi) {
+            // If overview is enabled, we want to update at the start
+            updateOverviewStateForDesktop(true);
+        }
         setOverviewGridEnabled(toState.displayOverviewTasksAsGrid(mActivity.getDeviceProfile()));
         setOverviewFullscreenEnabled(toState.getOverviewFullscreenProgress() == 1);
         if (toState == OVERVIEW_MODAL_TASK) {
@@ -162,6 +165,11 @@
             runActionOnRemoteHandles(remoteTargetHandle ->
                     remoteTargetHandle.getTaskViewSimulator().setDrawsBelowRecents(true));
         }
+
+        if (!finalState.overviewUi) {
+            // If overview is disabled, we want to update at the end
+            updateOverviewStateForDesktop(false);
+        }
     }
 
     @Override
@@ -273,4 +281,11 @@
             SystemUiProxy.INSTANCE.get(mActivity).showDesktopApps(mActivity.getDisplayId());
         }
     }
+
+    private void updateOverviewStateForDesktop(boolean enabled) {
+        DesktopVisibilityController controller = mActivity.getDesktopVisibilityController();
+        if (controller != null) {
+            controller.setOverviewStateEnabled(enabled);
+        }
+    }
 }
diff --git a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
index b3d04c6..58be345 100644
--- a/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
+++ b/quickstep/tests/src/com/android/launcher3/taskbar/TaskbarNavButtonControllerTest.java
@@ -17,6 +17,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
@@ -33,6 +34,7 @@
 import com.android.quickstep.OverviewCommandHelper;
 import com.android.quickstep.SystemUiProxy;
 import com.android.quickstep.TouchInteractionService;
+import com.android.quickstep.util.AssistUtils;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -54,6 +56,8 @@
     @Mock
     Handler mockHandler;
     @Mock
+    AssistUtils mockAssistUtils;
+    @Mock
     StatsLogManager mockStatsLogManager;
     @Mock
     StatsLogManager.StatsLogger mockStatsLogger;
@@ -79,7 +83,7 @@
                 .thenReturn(mockTaskbarActivityContext);
         doReturn(mockStatsLogManager).when(mockTaskbarActivityContext).getStatsLogManager();
         mNavButtonController = new TaskbarNavButtonController(mockService,
-                mockSystemUiProxy, mockHandler);
+                mockSystemUiProxy, mockHandler, mockAssistUtils);
     }
 
     @Test
@@ -108,16 +112,42 @@
     }
 
     @Test
-    public void testLongPressHome_enabled() {
+    public void testLongPressHome_enabled_withoutOverride() {
         mNavButtonController.setAssistantLongPressEnabled(true /*assistantLongPressEnabled*/);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(false);
+
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
+        verify(mockAssistUtils, times(1)).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, times(1)).startAssistant(any());
     }
 
     @Test
-    public void testLongPressHome_disabled() {
-        mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
+    public void testLongPressHome_enabled_withOverride() {
+        mNavButtonController.setAssistantLongPressEnabled(true /*assistantLongPressEnabled*/);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(true);
+
         mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
+        verify(mockAssistUtils, times(1)).tryStartAssistOverride(anyInt());
+        verify(mockSystemUiProxy, never()).startAssistant(any());
+    }
+
+    @Test
+    public void testLongPressHome_disabled_withoutOverride() {
+        mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(false);
+
+        mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
+        verify(mockAssistUtils, never()).tryStartAssistOverride(anyInt());
+        verify(mockSystemUiProxy, never()).startAssistant(any());
+    }
+
+    @Test
+    public void testLongPressHome_disabled_withOverride() {
+        mNavButtonController.setAssistantLongPressEnabled(false /*assistantLongPressEnabled*/);
+        when(mockAssistUtils.tryStartAssistOverride(anyInt())).thenReturn(true);
+
+        mNavButtonController.onButtonLongClick(BUTTON_HOME, mockView);
+        verify(mockAssistUtils, never()).tryStartAssistOverride(anyInt());
         verify(mockSystemUiProxy, never()).startAssistant(any());
     }
 
diff --git a/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java b/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java
index a90c326..cc56faf 100644
--- a/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java
+++ b/quickstep/tests/src/com/android/quickstep/TaplTestsSplitscreen.java
@@ -37,6 +37,7 @@
 
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -107,6 +108,7 @@
     }
 
     @Test
+    @Ignore("Enable once App Pairs flagged on. These cause memory leaks b/297135374")
     public void testSaveAppPairMenuItemExistsOnSplitPair() throws Exception {
         assumeTrue("App pairs feature is currently not enabled, no test needed",
                 FeatureFlags.ENABLE_APP_PAIRS.get());
@@ -122,6 +124,7 @@
     }
 
     @Test
+    @Ignore("Enable once App Pairs flagged on. These cause memory leaks b/297135374")
     public void testSaveAppPairMenuItemDoesNotExistOnSingleTask() throws Exception {
         assumeTrue("App pairs feature is currently not enabled, no test needed",
                 FeatureFlags.ENABLE_APP_PAIRS.get());
diff --git a/res/drawable-sw720dp/ic_transient_taskbar_all_apps_search_button.xml b/res/drawable-sw720dp/ic_transient_taskbar_all_apps_search_button.xml
new file mode 100644
index 0000000..22f5d6b
--- /dev/null
+++ b/res/drawable-sw720dp/ic_transient_taskbar_all_apps_search_button.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2023 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="52dp"
+    android:height="52dp"
+    android:viewportWidth="52"
+    android:viewportHeight="52">
+  <path
+      android:pathData="M42.647,43.855L37.215,38.423C36.784,38.768 36.288,39.041 35.727,39.242C35.167,39.443 34.57,39.544 33.938,39.544C32.372,39.544 31.042,39.005 29.95,37.927C28.872,36.835 28.333,35.505 28.333,33.939C28.333,32.372 28.872,31.05 29.95,29.972C31.042,28.88 32.372,28.334 33.938,28.334C35.505,28.334 36.827,28.88 37.904,29.972C38.997,31.05 39.543,32.372 39.543,33.939C39.543,34.571 39.442,35.167 39.241,35.728C39.04,36.289 38.767,36.784 38.422,37.215L43.854,42.648L42.647,43.855ZM33.938,37.819C35.016,37.819 35.929,37.445 36.676,36.698C37.437,35.936 37.818,35.017 37.818,33.939C37.818,32.861 37.437,31.948 36.676,31.201C35.929,30.439 35.016,30.059 33.938,30.059C32.86,30.059 31.94,30.439 31.179,31.201C30.431,31.948 30.058,32.861 30.058,33.939C30.058,35.017 30.431,35.936 31.179,36.698C31.94,37.445 32.86,37.819 33.938,37.819Z"
+      android:fillColor="#48473A"/>
+  <path
+      android:pathData="M39.42,17.543C39.42,20.605 36.938,23.086 33.876,23.086C30.815,23.086 28.333,20.605 28.333,17.543C28.333,14.482 30.815,12 33.876,12C36.938,12 39.42,14.482 39.42,17.543Z"
+      android:fillColor="#48473A"/>
+  <path
+      android:pathData="M23.086,17.543C23.086,20.605 20.605,23.086 17.543,23.086C14.482,23.086 12,20.605 12,17.543C12,14.482 14.482,12 17.543,12C20.605,12 23.086,14.482 23.086,17.543Z"
+      android:fillColor="#48473A"/>
+  <path
+      android:pathData="M17.543,33.877m-5.543,0a5.543,5.543 0,1 1,11.086 0a5.543,5.543 0,1 1,-11.086 0"
+      android:fillColor="#48473A"/>
+</vector>
diff --git a/res/drawable/ic_taskbar_all_apps_search_button.xml b/res/drawable/ic_taskbar_all_apps_search_button.xml
new file mode 100644
index 0000000..8a42dc7
--- /dev/null
+++ b/res/drawable/ic_taskbar_all_apps_search_button.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2023 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="44dp"
+    android:height="44dp"
+    android:viewportWidth="44"
+    android:viewportHeight="44">
+  <path
+      android:pathData="M36.938,38L32.164,33.225C31.785,33.528 31.349,33.769 30.856,33.945C30.363,34.122 29.839,34.211 29.283,34.211C27.906,34.211 26.738,33.737 25.778,32.79C24.83,31.83 24.357,30.661 24.357,29.284C24.357,27.907 24.83,26.745 25.778,25.798C26.738,24.837 27.906,24.357 29.283,24.357C30.66,24.357 31.822,24.837 32.77,25.798C33.73,26.745 34.21,27.907 34.21,29.284C34.21,29.84 34.121,30.364 33.945,30.857C33.768,31.349 33.528,31.785 33.225,32.164L38,36.939L36.938,38ZM29.283,32.695C30.231,32.695 31.033,32.366 31.69,31.709C32.359,31.04 32.694,30.232 32.694,29.284C32.694,28.337 32.359,27.535 31.69,26.878C31.033,26.208 30.231,25.873 29.283,25.873C28.336,25.873 27.527,26.208 26.858,26.878C26.201,27.535 25.873,28.337 25.873,29.284C25.873,30.232 26.201,31.04 26.858,31.709C27.527,32.366 28.336,32.695 29.283,32.695Z"
+      android:fillColor="#52443C"/>
+  <path
+      android:pathData="M34.102,14.873C34.102,17.563 31.92,19.745 29.229,19.745C26.538,19.745 24.357,17.563 24.357,14.873C24.357,12.182 26.538,10 29.229,10C31.92,10 34.102,12.182 34.102,14.873Z"
+      android:fillColor="#52443C"/>
+  <path
+      android:pathData="M19.745,14.873C19.745,17.563 17.563,19.745 14.873,19.745C12.182,19.745 10,17.563 10,14.873C10,12.182 12.182,10 14.873,10C17.563,10 19.745,12.182 19.745,14.873Z"
+      android:fillColor="#52443C"/>
+  <path
+      android:pathData="M14.873,29.23m-4.872,0a4.872,4.872 0,1 1,9.745 0a4.872,4.872 0,1 1,-9.745 0"
+      android:fillColor="#52443C"/>
+</vector>
diff --git a/res/drawable/ic_transient_taskbar_all_apps_search_button.xml b/res/drawable/ic_transient_taskbar_all_apps_search_button.xml
new file mode 100644
index 0000000..52d6818
--- /dev/null
+++ b/res/drawable/ic_transient_taskbar_all_apps_search_button.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2023 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="48dp"
+    android:height="48dp"
+    android:viewportWidth="48"
+    android:viewportHeight="48">
+  <path
+      android:pathData="M40.647,41.855L35.215,36.423C34.784,36.768 34.288,37.041 33.727,37.242C33.167,37.443 32.57,37.544 31.938,37.544C30.372,37.544 29.042,37.005 27.95,35.927C26.872,34.835 26.333,33.505 26.333,31.939C26.333,30.372 26.872,29.05 27.95,27.972C29.042,26.88 30.372,26.334 31.938,26.334C33.505,26.334 34.827,26.88 35.904,27.972C36.997,29.05 37.543,30.372 37.543,31.939C37.543,32.571 37.442,33.167 37.241,33.728C37.04,34.289 36.767,34.784 36.422,35.215L41.854,40.648L40.647,41.855ZM31.938,35.819C33.016,35.819 33.929,35.445 34.676,34.698C35.437,33.936 35.818,33.017 35.818,31.939C35.818,30.861 35.437,29.948 34.676,29.201C33.929,28.439 33.016,28.059 31.938,28.059C30.86,28.059 29.94,28.439 29.179,29.201C28.431,29.948 28.058,30.861 28.058,31.939C28.058,33.017 28.431,33.936 29.179,34.698C29.94,35.445 30.86,35.819 31.938,35.819Z"
+      android:fillColor="#52443C"/>
+  <path
+      android:pathData="M37.42,15.543C37.42,18.605 34.938,21.086 31.876,21.086C28.815,21.086 26.333,18.605 26.333,15.543C26.333,12.482 28.815,10 31.876,10C34.938,10 37.42,12.482 37.42,15.543Z"
+      android:fillColor="#52443C"/>
+  <path
+      android:pathData="M21.086,15.543C21.086,18.605 18.605,21.086 15.543,21.086C12.482,21.086 10,18.605 10,15.543C10,12.482 12.482,10 15.543,10C18.605,10 21.086,12.482 21.086,15.543Z"
+      android:fillColor="#52443C"/>
+  <path
+      android:pathData="M15.543,31.877m-5.543,0a5.543,5.543 0,1 1,11.086 0a5.543,5.543 0,1 1,-11.086 0"
+      android:fillColor="#52443C"/>
+</vector>
diff --git a/res/values/config.xml b/res/values/config.xml
index 8f9731c..cd9c9de 100644
--- a/res/values/config.xml
+++ b/res/values/config.xml
@@ -241,4 +241,7 @@
         <item>@dimen/iconSize66dp</item>
         <item>@dimen/iconSize72dp</item>
     </integer-array>
+
+    <!--  Used for custom widgets  -->
+    <array name="custom_widget_providers"/>
 </resources>
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 606b2c4..1fef2f8 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -68,6 +68,7 @@
 import static com.android.launcher3.logging.StatsLogManager.StatsLatencyLogger.LatencyType.WARM;
 import static com.android.launcher3.model.ItemInstallQueue.FLAG_ACTIVITY_PAUSED;
 import static com.android.launcher3.model.ItemInstallQueue.FLAG_DRAG_AND_DROP;
+import static com.android.launcher3.model.data.LauncherAppWidgetInfo.CUSTOM_WIDGET_ID;
 import static com.android.launcher3.popup.SystemShortcut.APP_INFO;
 import static com.android.launcher3.popup.SystemShortcut.INSTALL;
 import static com.android.launcher3.popup.SystemShortcut.WIDGETS;
@@ -144,6 +145,7 @@
 import com.android.launcher3.allapps.AllAppsTransitionController;
 import com.android.launcher3.allapps.BaseSearchConfig;
 import com.android.launcher3.allapps.DiscoveryBounce;
+import com.android.launcher3.anim.AnimationSuccessListener;
 import com.android.launcher3.anim.PropertyListBuilder;
 import com.android.launcher3.apppairs.AppPairIcon;
 import com.android.launcher3.celllayout.CellPosMapper;
@@ -885,7 +887,7 @@
                 if (widgetInfo != null) {
                     // Since the view was just bound, also launch the configure activity if needed
                     LauncherAppWidgetProviderInfo provider = mAppWidgetManager
-                            .getLauncherAppWidgetInfo(widgetId);
+                            .getLauncherAppWidgetInfo(widgetId, info.getTargetComponent());
                     if (provider != null) {
                         new WidgetAddFlowHandler(provider)
                                 .startConfigActivity(this, widgetInfo,
@@ -1499,7 +1501,8 @@
             AppWidgetHostView hostView, LauncherAppWidgetProviderInfo appWidgetInfo) {
 
         if (appWidgetInfo == null) {
-            appWidgetInfo = mAppWidgetManager.getLauncherAppWidgetInfo(appWidgetId);
+            appWidgetInfo = mAppWidgetManager.getLauncherAppWidgetInfo(appWidgetId,
+                    itemInfo.getTargetComponent());
         }
 
         if (hostView == null) {
@@ -1727,7 +1730,16 @@
         if (getStateManager().isInStableState(ALL_APPS)) {
             getStateManager().goToState(NORMAL, alreadyOnHome);
         } else {
-            showAllAppsFromIntent(alreadyOnHome);
+            AbstractFloatingView.closeAllOpenViews(this);
+            getStateManager().goToState(ALL_APPS, true /* animated */,
+                    new AnimationSuccessListener() {
+                        @Override
+                        public void onAnimationSuccess(Animator animator) {
+                            if (mAppsView.getSearchUiManager().getEditText() != null) {
+                                mAppsView.getSearchUiManager().getEditText().requestFocus();
+                            }
+                        }
+                    });
         }
     }
 
@@ -1985,8 +1997,8 @@
             // In this case, we either need to start an activity to get permission to bind
             // the widget, or we need to start an activity to configure the widget, or both.
             if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET) {
-                appWidgetId = CustomWidgetManager.INSTANCE.get(this).getWidgetIdForCustomProvider(
-                        info.componentName);
+                appWidgetId = CustomWidgetManager.INSTANCE.get(this)
+                        .allocateCustomAppWidgetId(info.componentName);
             } else {
                 appWidgetId = getAppWidgetHolder().allocateAppWidgetId();
             }
@@ -2632,9 +2644,10 @@
                     }
                 }
             } else {
-                appWidgetInfo = mAppWidgetManager.getLauncherAppWidgetInfo(item.appWidgetId);
+                appWidgetInfo = mAppWidgetManager.getLauncherAppWidgetInfo(item.appWidgetId,
+                        item.getTargetComponent());
                 if (appWidgetInfo == null) {
-                    if (item.appWidgetId <= LauncherAppWidgetInfo.CUSTOM_WIDGET_ID) {
+                    if (item.appWidgetId <= CUSTOM_WIDGET_ID) {
                         removalReason =
                                 "CustomWidgetManager cannot find provider from that widget id.";
                     } else {
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 8be8fed..a1a3974 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -3413,7 +3413,8 @@
             if (item.hasRestoreFlag(LauncherAppWidgetInfo.FLAG_ID_NOT_VALID)) {
                 widgetInfo = widgetHelper.findProvider(item.providerName, item.user);
             } else {
-                widgetInfo = widgetHelper.getLauncherAppWidgetInfo(item.appWidgetId);
+                widgetInfo = widgetHelper.getLauncherAppWidgetInfo(item.appWidgetId,
+                        item.getTargetComponent());
             }
 
             if (widgetInfo != null) {
diff --git a/src/com/android/launcher3/config/FeatureFlags.java b/src/com/android/launcher3/config/FeatureFlags.java
index 4285755..c0075e0 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -141,7 +141,7 @@
 
     // TODO(Block 6): Clean up flags
     public static final BooleanFlag ENABLE_ALL_APPS_SEARCH_IN_TASKBAR = getDebugFlag(270393900,
-            "ENABLE_ALL_APPS_SEARCH_IN_TASKBAR", DISABLED,
+            "ENABLE_ALL_APPS_SEARCH_IN_TASKBAR", TEAMFOOD,
             "Enables Search box in Taskbar All Apps.");
 
     public static final BooleanFlag SECONDARY_DRAG_N_DROP_TO_PIN = getDebugFlag(270395140,
@@ -175,6 +175,9 @@
             "MULTI_SELECT_EDIT_MODE", DISABLED, "Enable new multi-select edit mode "
                     + "for home screen");
 
+    public static final BooleanFlag SMARTSPACE_AS_A_WIDGET = getDebugFlag(299181941,
+            "SMARTSPACE_AS_A_WIDGET", DISABLED, "Enable SmartSpace as a widget");
+
     // TODO(Block 10): Clean up flags
     public static final BooleanFlag ENABLE_BACK_SWIPE_LAUNCHER_ANIMATION = getDebugFlag(270614790,
             "ENABLE_BACK_SWIPE_LAUNCHER_ANIMATION", DISABLED,
@@ -213,6 +216,10 @@
     public static final BooleanFlag ENABLE_TRANSIENT_TASKBAR = getDebugFlag(270395798,
             "ENABLE_TRANSIENT_TASKBAR", ENABLED, "Enables transient taskbar.");
 
+    public static final BooleanFlag ENABLE_TASKBAR_NO_RECREATION = getDebugFlag(299193589,
+            "ENABLE_TASKBAR_NO_RECREATION", DISABLED,
+            "Enables taskbar with no recreation from lifecycle changes of TaskbarActivityContext.");
+
     // TODO(Block 16): Clean up flags
     // When enabled the promise icon is visible in all apps while installation an app.
     public static final BooleanFlag PROMISE_APPS_IN_ALL_APPS = getDebugFlag(270390012,
@@ -261,7 +268,7 @@
 
     // TODO(Block 17): Clean up flags
     public static final BooleanFlag ENABLE_TASKBAR_PINNING = getDebugFlag(270396583,
-            "ENABLE_TASKBAR_PINNING", DISABLED,
+            "ENABLE_TASKBAR_PINNING", TEAMFOOD,
             "Enables taskbar pinning to allow user to switch between transient and persistent "
                     + "taskbar flavors");
 
diff --git a/src/com/android/launcher3/model/GridSizeMigrationUtil.java b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
index c233872..78875a3 100644
--- a/src/com/android/launcher3/model/GridSizeMigrationUtil.java
+++ b/src/com/android/launcher3/model/GridSizeMigrationUtil.java
@@ -539,8 +539,8 @@
                             verifyPackage(cn.getPackageName());
 
                             int widgetId = c.getInt(indexAppWidgetId);
-                            LauncherAppWidgetProviderInfo pInfo =
-                                    widgetManagerHelper.getLauncherAppWidgetInfo(widgetId);
+                            LauncherAppWidgetProviderInfo pInfo = widgetManagerHelper
+                                    .getLauncherAppWidgetInfo(widgetId, cn);
                             Point spans = null;
                             if (pInfo != null) {
                                 spans = pInfo.getMinSpans();
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 0e68db2..08a6cf0 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -95,6 +95,7 @@
 import com.android.launcher3.util.TraceHelper;
 import com.android.launcher3.widget.LauncherAppWidgetProviderInfo;
 import com.android.launcher3.widget.WidgetManagerHelper;
+import com.android.launcher3.widget.custom.CustomWidgetManager;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -740,8 +741,13 @@
 
                     ComponentKey providerKey = new ComponentKey(component, c.user);
                     if (!mWidgetProvidersMap.containsKey(providerKey)) {
-                        mWidgetProvidersMap.put(providerKey,
-                                widgetHelper.findProvider(component, c.user));
+                        if (customWidget) {
+                            mWidgetProvidersMap.put(providerKey, CustomWidgetManager.INSTANCE
+                                    .get(mApp.getContext()).getWidgetProvider(component));
+                        } else {
+                            mWidgetProvidersMap.put(providerKey,
+                                    widgetHelper.findProvider(component, c.user));
+                        }
                     }
                     final AppWidgetProviderInfo provider = mWidgetProvidersMap.get(providerKey);
 
@@ -814,7 +820,8 @@
                             return;
                         }
                         LauncherAppWidgetProviderInfo widgetProviderInfo =
-                                widgetHelper.getLauncherAppWidgetInfo(appWidgetId);
+                                widgetHelper.getLauncherAppWidgetInfo(appWidgetId,
+                                        appWidgetInfo.getTargetComponent());
                         if (widgetProviderInfo != null
                                 && (appWidgetInfo.spanX < widgetProviderInfo.minSpanX
                                 || appWidgetInfo.spanY < widgetProviderInfo.minSpanY)) {
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
index 76a044d..f64fd05 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetHostView.java
@@ -23,9 +23,11 @@
 import android.graphics.Rect;
 import android.os.Build;
 import android.os.Handler;
+import android.os.Parcelable;
 import android.os.SystemClock;
 import android.os.Trace;
 import android.util.Log;
+import android.util.SparseArray;
 import android.util.SparseBooleanArray;
 import android.util.SparseIntArray;
 import android.view.MotionEvent;
@@ -498,4 +500,13 @@
          */
         void notifyBoundChangeOnPreLayout(View v, int left, int top, int right, int bottom);
     }
+
+    @Override
+    protected void dispatchRestoreInstanceState(SparseArray<Parcelable> container) {
+        try {
+            super.dispatchRestoreInstanceState(container);
+        } catch (Exception e) {
+            Log.i(TAG, "Exception: " + e);
+        }
+    }
 }
diff --git a/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
index 10aef9a..ef51d15 100644
--- a/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/LauncherAppWidgetProviderInfo.java
@@ -67,7 +67,7 @@
      */
     public int maxSpanY;
 
-    private boolean mIsMinSizeFulfilled;
+    protected boolean mIsMinSizeFulfilled;
 
     public static LauncherAppWidgetProviderInfo fromProviderInfo(Context context,
             AppWidgetProviderInfo info) {
diff --git a/src/com/android/launcher3/widget/WidgetManagerHelper.java b/src/com/android/launcher3/widget/WidgetManagerHelper.java
index 737cdbd..0860e72 100644
--- a/src/com/android/launcher3/widget/WidgetManagerHelper.java
+++ b/src/com/android/launcher3/widget/WidgetManagerHelper.java
@@ -57,10 +57,15 @@
     /**
      * @see AppWidgetManager#getAppWidgetInfo(int)
      */
-    public LauncherAppWidgetProviderInfo getLauncherAppWidgetInfo(int appWidgetId) {
-        if (appWidgetId <= LauncherAppWidgetInfo.CUSTOM_WIDGET_ID) {
-            return CustomWidgetManager.INSTANCE.get(mContext).getWidgetProvider(appWidgetId);
+    public LauncherAppWidgetProviderInfo getLauncherAppWidgetInfo(
+            int appWidgetId, ComponentName componentName) {
+
+        // For custom widgets.
+        if (appWidgetId <= LauncherAppWidgetInfo.CUSTOM_WIDGET_ID && !CustomWidgetManager
+                .INSTANCE.get(mContext).getWidgetIdForCustomProvider(componentName).equals("")) {
+            return CustomWidgetManager.INSTANCE.get(mContext).getWidgetProvider(componentName);
         }
+
         AppWidgetProviderInfo info = mAppWidgetManager.getAppWidgetInfo(appWidgetId);
         return info == null ? null : LauncherAppWidgetProviderInfo.fromProviderInfo(mContext, info);
     }
diff --git a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
index 8b3bbce..44571a6 100644
--- a/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
+++ b/src/com/android/launcher3/widget/custom/CustomAppWidgetProviderInfo.java
@@ -33,14 +33,15 @@
 public class CustomAppWidgetProviderInfo extends LauncherAppWidgetProviderInfo
         implements Parcelable {
 
-    public final int providerId;
+    public final String providerId;
 
-    protected CustomAppWidgetProviderInfo(Parcel parcel, boolean readSelf, int providerId) {
+    protected CustomAppWidgetProviderInfo(Parcel parcel, boolean readSelf, String providerId) {
         super(parcel);
         if (readSelf) {
-            this.providerId = parcel.readInt();
+            this.providerId = parcel.readString();
 
-            provider = new ComponentName(parcel.readString(), CLS_CUSTOM_WIDGET_PREFIX + providerId);
+            provider = new ComponentName(parcel.readString(),
+                    CLS_CUSTOM_WIDGET_PREFIX + parcel.readString());
 
             label = parcel.readString();
             initialLayout = parcel.readInt();
@@ -58,7 +59,10 @@
     }
 
     @Override
-    public void initSpans(Context context, InvariantDeviceProfile idp) { }
+    public void initSpans(Context context, InvariantDeviceProfile idp) {
+        mIsMinSizeFulfilled = Math.min(spanX, minSpanX) <= idp.numColumns
+                && Math.min(spanY, minSpanY) <= idp.numRows;
+    }
 
     @Override
     public String getLabel(PackageManager packageManager) {
@@ -73,8 +77,9 @@
     @Override
     public void writeToParcel(Parcel out, int flags) {
         super.writeToParcel(out, flags);
-        out.writeInt(providerId);
+        out.writeString(providerId);
         out.writeString(provider.getPackageName());
+        out.writeString(provider.getClassName());
 
         out.writeString(label);
         out.writeInt(initialLayout);
@@ -93,7 +98,7 @@
 
         @Override
         public CustomAppWidgetProviderInfo createFromParcel(Parcel parcel) {
-            return new CustomAppWidgetProviderInfo(parcel, true, 0);
+            return new CustomAppWidgetProviderInfo(parcel, true, "");
         }
 
         @Override
diff --git a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
index 2e2a968..7cf0221 100644
--- a/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
+++ b/src/com/android/launcher3/widget/custom/CustomWidgetManager.java
@@ -16,7 +16,8 @@
 
 package com.android.launcher3.widget.custom;
 
-import static com.android.launcher3.widget.LauncherAppWidgetProviderInfo.CLS_CUSTOM_WIDGET_PREFIX;
+import static com.android.launcher3.config.FeatureFlags.SMARTSPACE_AS_A_WIDGET;
+import static com.android.launcher3.model.data.LauncherAppWidgetInfo.CUSTOM_WIDGET_ID;
 
 import android.appwidget.AppWidgetManager;
 import android.appwidget.AppWidgetProviderInfo;
@@ -24,12 +25,12 @@
 import android.content.Context;
 import android.os.Parcel;
 import android.os.Process;
-import android.util.SparseArray;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.R;
 import com.android.launcher3.uioverrides.plugins.PluginManagerWrapper;
 import com.android.launcher3.util.MainThreadInitializedObject;
 import com.android.launcher3.util.PackageUserKey;
@@ -39,8 +40,11 @@
 import com.android.systemui.plugins.CustomWidgetPlugin;
 import com.android.systemui.plugins.PluginListener;
 
+import java.lang.reflect.InvocationTargetException;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
 
@@ -52,24 +56,37 @@
     public static final MainThreadInitializedObject<CustomWidgetManager> INSTANCE =
             new MainThreadInitializedObject<>(CustomWidgetManager::new);
 
+    private static final String TAG = "CustomWidgetManager";
     private final Context mContext;
-    /**
-     * auto provider Id is an ever-increasing number that serves as the providerId whenever a new
-     * custom widget has been connected.
-     */
-    private int mAutoProviderId = 0;
-    private final SparseArray<CustomWidgetPlugin> mPlugins;
+    private final HashMap<String, CustomWidgetPlugin> mPlugins;
     private final List<CustomAppWidgetProviderInfo> mCustomWidgets;
-    private final SparseArray<ComponentName> mWidgetsIdMap;
+    private final HashMap<ComponentName, String> mWidgetsIdMap;
     private Consumer<PackageUserKey> mWidgetRefreshCallback;
 
     private CustomWidgetManager(Context context) {
         mContext = context;
-        mPlugins = new SparseArray<>();
+        mPlugins = new HashMap<>();
         mCustomWidgets = new ArrayList<>();
-        mWidgetsIdMap = new SparseArray<>();
+        mWidgetsIdMap = new HashMap<>();
         PluginManagerWrapper.INSTANCE.get(context)
                 .addPluginListener(this, CustomWidgetPlugin.class, true);
+
+        if (SMARTSPACE_AS_A_WIDGET.get()) {
+            for (String s: context.getResources()
+                    .getStringArray(R.array.custom_widget_providers)) {
+                try {
+                    Class<?> cls = Class.forName(s);
+                    CustomWidgetPlugin plugin = (CustomWidgetPlugin)
+                            cls.getDeclaredConstructor(Context.class).newInstance(context);
+                    mPlugins.put(plugin.getId(), plugin);
+                    onPluginConnected(mPlugins.get(plugin.getId()), context);
+                } catch (ClassNotFoundException | InstantiationException | IllegalAccessException
+                         | ClassCastException | NoSuchMethodException
+                         | InvocationTargetException e) {
+                    Log.e(TAG, "Exception found when trying to add custom widgets: " + e);
+                }
+            }
+        }
     }
 
     @Override
@@ -79,28 +96,41 @@
 
     @Override
     public void onPluginConnected(CustomWidgetPlugin plugin, Context context) {
-        mPlugins.put(mAutoProviderId, plugin);
         List<AppWidgetProviderInfo> providers = AppWidgetManager.getInstance(context)
                 .getInstalledProvidersForProfile(Process.myUserHandle());
         if (providers.isEmpty()) return;
         Parcel parcel = Parcel.obtain();
         providers.get(0).writeToParcel(parcel, 0);
         parcel.setDataPosition(0);
-        CustomAppWidgetProviderInfo info = newInfo(mAutoProviderId, plugin, parcel, context);
+        CustomAppWidgetProviderInfo info = newInfo(plugin.getId(), plugin, parcel, context);
         parcel.recycle();
         mCustomWidgets.add(info);
-        mWidgetsIdMap.put(mAutoProviderId, info.provider);
-        mWidgetRefreshCallback.accept(null);
-        mAutoProviderId++;
+        mWidgetsIdMap.put(info.provider, plugin.getId());
     }
 
     @Override
     public void onPluginDisconnected(CustomWidgetPlugin plugin) {
-        int providerId = findProviderId(plugin);
-        if (providerId == -1) return;
-        mPlugins.remove(providerId);
-        mCustomWidgets.remove(getWidgetProvider(providerId));
-        mWidgetsIdMap.remove(providerId);
+        String providerId = plugin.getId();
+        if (mPlugins.containsKey(providerId)) {
+            mPlugins.remove(providerId);
+        }
+
+        ComponentName cn = null;
+        for (Map.Entry entry: mWidgetsIdMap.entrySet()) {
+            if (entry.getValue().equals(providerId)) {
+                cn = (ComponentName) entry.getKey();
+            }
+        }
+
+        if (cn != null) {
+            mWidgetsIdMap.remove(cn);
+            for (int i = 0; i < mCustomWidgets.size(); i++) {
+                if (mCustomWidgets.get(i).getComponent().equals(cn)) {
+                    mCustomWidgets.remove(i);
+                    return;
+                }
+            }
+        }
     }
 
     /**
@@ -131,12 +161,11 @@
     /**
      * Returns the widget id for a specific provider.
      */
-    public int getWidgetIdForCustomProvider(@NonNull ComponentName provider) {
-        int index = mWidgetsIdMap.indexOfValue(provider);
-        if (index >= 0) {
-            return LauncherAppWidgetInfo.CUSTOM_WIDGET_ID - mWidgetsIdMap.keyAt(index);
+    public String getWidgetIdForCustomProvider(@NonNull ComponentName provider) {
+        if (mWidgetsIdMap.containsKey(provider)) {
+            return mWidgetsIdMap.get(provider);
         } else {
-            return AppWidgetManager.INVALID_APPWIDGET_ID;
+            return "";
         }
     }
 
@@ -144,38 +173,26 @@
      * Returns the widget provider in respect to given widget id.
      */
     @Nullable
-    public LauncherAppWidgetProviderInfo getWidgetProvider(int widgetId) {
-        ComponentName cn = mWidgetsIdMap.get(LauncherAppWidgetInfo.CUSTOM_WIDGET_ID - widgetId);
+    public LauncherAppWidgetProviderInfo getWidgetProvider(ComponentName componentName) {
         for (LauncherAppWidgetProviderInfo info : mCustomWidgets) {
-            if (info.provider.equals(cn)) return info;
+            if (info.provider.equals(componentName)) return info;
         }
         return null;
     }
 
-    private static CustomAppWidgetProviderInfo newInfo(int providerId, CustomWidgetPlugin plugin,
+    private static CustomAppWidgetProviderInfo newInfo(String providerId, CustomWidgetPlugin plugin,
             Parcel parcel, Context context) {
         CustomAppWidgetProviderInfo info = new CustomAppWidgetProviderInfo(
                 parcel, false, providerId);
-        info.provider = new ComponentName(
-                context.getPackageName(), CLS_CUSTOM_WIDGET_PREFIX + providerId);
-
-        info.label = plugin.getLabel();
-        info.resizeMode = plugin.getResizeMode();
-
-        info.spanX = plugin.getSpanX();
-        info.spanY = plugin.getSpanY();
-        info.minSpanX = plugin.getMinSpanX();
-        info.minSpanY = plugin.getMinSpanY();
+        plugin.updateWidgetInfo(info, context);
         return info;
     }
 
-    private int findProviderId(CustomWidgetPlugin plugin) {
-        for (int i = 0; i < mPlugins.size(); i++) {
-            int providerId = mPlugins.keyAt(i);
-            if (mPlugins.get(providerId) == plugin) {
-                return providerId;
-            }
-        }
-        return -1;
+    /**
+     * Returns an id to set as the appWidgetId for a custom widget.
+     */
+    public int allocateCustomAppWidgetId(ComponentName componentName) {
+        return CUSTOM_WIDGET_ID - mCustomWidgets.indexOf(getWidgetProvider(componentName));
     }
+
 }
diff --git a/src_plugins/com/android/systemui/plugins/CustomWidgetPlugin.java b/src_plugins/com/android/systemui/plugins/CustomWidgetPlugin.java
index 56ebcc5..af4f22c 100644
--- a/src_plugins/com/android/systemui/plugins/CustomWidgetPlugin.java
+++ b/src_plugins/com/android/systemui/plugins/CustomWidgetPlugin.java
@@ -17,6 +17,8 @@
 package com.android.systemui.plugins;
 
 import android.appwidget.AppWidgetHostView;
+import android.appwidget.AppWidgetProviderInfo;
+import android.content.Context;
 
 import com.android.systemui.plugins.annotations.ProvidesInterface;
 
@@ -30,42 +32,20 @@
     int VERSION = 1;
 
     /**
-     * The label to display to the user in the AppWidget picker.
-     */
-    String getLabel();
-
-    /**
-     * The default width of the widget when added to a host, in dp. The widget will get
-     * at least this width, and will often be given more, depending on the host.
-     */
-    int getSpanX();
-
-    /**
-     * The default height of the widget when added to a host, in dp. The widget will get
-     * at least this height, and will often be given more, depending on the host.
-     */
-    int getSpanY();
-
-    /**
-     * Minimum width (in dp) which the widget can be resized to. This field has no effect if it
-     * is greater than minWidth or if horizontal resizing isn't enabled.
-     */
-    int getMinSpanX();
-
-    /**
-     * Minimum height (in dp) which the widget can be resized to. This field has no effect if it
-     * is greater than minHeight or if vertical resizing isn't enabled.
-     */
-    int getMinSpanY();
-
-    /**
-     * The rules by which a widget can be resized.
-     */
-    int getResizeMode();
-
-    /**
      * Notify the plugin that container of the widget has been rendered, where the custom widget
      * can be attached to.
      */
     void onViewCreated(AppWidgetHostView parent);
+
+    /**
+     * Get the UUID for the custom widget.
+     */
+    String getId();
+
+    /**
+     * Used to modify a widgets' info.
+     */
+    default void updateWidgetInfo(AppWidgetProviderInfo info, Context context) {
+
+    }
 }
diff --git a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
index dd79ca8..ed8e324 100644
--- a/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
+++ b/tests/src/com/android/launcher3/AbstractDeviceProfileTest.kt
@@ -26,6 +26,7 @@
 import com.android.launcher3.util.DisplayController
 import com.android.launcher3.util.NavigationMode
 import com.android.launcher3.util.WindowBounds
+import com.android.launcher3.util.rule.TestStabilityRule
 import com.android.launcher3.util.window.CachedDisplayInfo
 import com.android.launcher3.util.window.WindowManagerProxy
 import java.io.BufferedReader
@@ -36,6 +37,7 @@
 import kotlin.math.min
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.mockito.ArgumentMatchers
 import org.mockito.Mockito.mock
 import org.mockito.Mockito.`when` as whenever
@@ -54,6 +56,8 @@
     private lateinit var originalDisplayController: DisplayController
     private lateinit var originalWindowManagerProxy: WindowManagerProxy
 
+    @Rule @JvmField val testStabilityRule = TestStabilityRule()
+
     @Before
     open fun setUp() {
         val appContext: Context = ApplicationProvider.getApplicationContext()