diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index aca1b3b..226d717 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -13,3 +13,10 @@
     description: "Enables two line label inside all apps."
     bug: "270390937"
 }
+
+flag {
+    name: "enable_grid_only_overview"
+    namespace: "launcher"
+    description: "Enable a grid-only overview without a focused task."
+    bug: "270397206"
+}
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 80ac9be..3a1ed0b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -347,6 +347,11 @@
         mControllers.taskbarAllAppsController.toggle();
     }
 
+    /** Toggles Taskbar All Apps overlay with keyboard ready for search. */
+    public void toggleAllAppsSearch() {
+        mControllers.taskbarAllAppsController.toggleSearch();
+    }
+
     @Override
     public DeviceProfile getDeviceProfile() {
         return mDeviceProfile;
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 55d56f0..ba2116d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -304,7 +304,7 @@
             return;
         }
 
-        mTaskbarActivityContext.toggleAllApps();
+        mTaskbarActivityContext.toggleAllAppsSearch();
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java
index 1c250bf..3a733bf 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarScrimViewController.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.taskbar.bubbles.BubbleBarController.BUBBLE_BAR_ENABLED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_MANAGE_MENU_EXPANDED;
+import static com.android.wm.shell.common.bubbles.BubbleConstants.BUBBLE_EXPANDED_SCRIM_ALPHA;
 
 import android.animation.ObjectAnimator;
 import android.view.animation.Interpolator;
@@ -35,8 +36,6 @@
 public class TaskbarScrimViewController implements TaskbarControllers.LoggableTaskbarController,
         TaskbarControllers.BackgroundRendererController {
 
-    private static final float SCRIM_ALPHA = 0.6f;
-
     private static final Interpolator SCRIM_ALPHA_IN = new PathInterpolator(0.4f, 0f, 1f, 1f);
     private static final Interpolator SCRIM_ALPHA_OUT = new PathInterpolator(0f, 0f, 0.8f, 1f);
 
@@ -78,8 +77,9 @@
         final float scrimAlpha = manageMenuExpanded
                 // When manage menu shows there's the first scrim and second scrim so figure out
                 // what the total transparency would be.
-                ? (SCRIM_ALPHA + (SCRIM_ALPHA * (1 - SCRIM_ALPHA)))
-                : showScrim ? SCRIM_ALPHA : 0;
+                ? (BUBBLE_EXPANDED_SCRIM_ALPHA + (BUBBLE_EXPANDED_SCRIM_ALPHA
+                * (1 - BUBBLE_EXPANDED_SCRIM_ALPHA)))
+                : showScrim ? BUBBLE_EXPANDED_SCRIM_ALPHA : 0;
         showScrim(showScrim, scrimAlpha, skipAnim);
     }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsController.java b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsController.java
index d93954c..c54bb7e 100644
--- a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsController.java
@@ -128,10 +128,19 @@
 
     /** Toggles visibility of {@link TaskbarAllAppsContainerView} in the overlay window. */
     public void toggle() {
+        toggle(false);
+    }
+
+    /** Toggles visibility of {@link TaskbarAllAppsContainerView} with the keyboard for search. */
+    public void toggleSearch() {
+        toggle(true);
+    }
+
+    private void toggle(boolean showKeyboard) {
         if (isOpen()) {
             mSlideInView.close(true);
         } else {
-            show(true);
+            show(true, showKeyboard);
         }
     }
 
@@ -141,6 +150,10 @@
     }
 
     private void show(boolean animate) {
+        show(animate, false);
+    }
+
+    private void show(boolean animate, boolean showKeyboard) {
         if (mAppsView != null) {
             return;
         }
@@ -169,7 +182,11 @@
             cleanUpOverlay();
         });
         TaskbarAllAppsViewController viewController = new TaskbarAllAppsViewController(
-                mOverlayContext, mSlideInView, mControllers, mSearchSessionController);
+                mOverlayContext,
+                mSlideInView,
+                mControllers,
+                mSearchSessionController,
+                showKeyboard);
 
         viewController.show(animate);
         mAppsView = mOverlayContext.getAppsView();
diff --git a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewController.java b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewController.java
index 85633e9..3a6d613 100644
--- a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarAllAppsViewController.java
@@ -41,12 +41,14 @@
     private final TaskbarStashController mTaskbarStashController;
     private final NavbarButtonsViewController mNavbarButtonsViewController;
     private final TaskbarOverlayController mOverlayController;
+    private final boolean mShowKeyboard;
 
     TaskbarAllAppsViewController(
             TaskbarOverlayContext context,
             TaskbarAllAppsSlideInView slideInView,
             TaskbarControllers taskbarControllers,
-            TaskbarSearchSessionController searchSessionController) {
+            TaskbarSearchSessionController searchSessionController,
+            boolean showKeyboard) {
 
         mContext = context;
         mSlideInView = slideInView;
@@ -54,6 +56,7 @@
         mTaskbarStashController = taskbarControllers.taskbarStashController;
         mNavbarButtonsViewController = taskbarControllers.navbarButtonsViewController;
         mOverlayController = taskbarControllers.taskbarOverlayController;
+        mShowKeyboard = showKeyboard;
 
         mSlideInView.init(new TaskbarAllAppsCallbacks(searchSessionController));
         setUpAppDivider();
@@ -120,6 +123,11 @@
         @Override
         public void onAllAppsTransitionEnd(boolean toAllApps) {
             mSearchSessionController.onAllAppsTransitionEnd(toAllApps);
+            if (toAllApps
+                    && mShowKeyboard
+                    && mAppsView.getSearchUiManager().getEditText() != null) {
+                mAppsView.getSearchUiManager().getEditText().requestFocus();
+            }
         }
 
         /** Invoked on back press, returning {@code true} if the search session handled it. */
@@ -128,7 +136,7 @@
         }
 
         void onAllAppsAnimationPending(PendingAnimation animation, boolean toAllApps) {
-            mSearchSessionController.onAllAppsAnimationPending(animation, toAllApps);
+            mSearchSessionController.onAllAppsAnimationPending(animation, toAllApps, mShowKeyboard);
         }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarSearchSessionController.kt b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarSearchSessionController.kt
index 8a2041f..3d15fbd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarSearchSessionController.kt
+++ b/quickstep/src/com/android/launcher3/taskbar/allapps/TaskbarSearchSessionController.kt
@@ -51,7 +51,11 @@
 
     open fun handleBackInvoked(): Boolean = false
 
-    open fun onAllAppsAnimationPending(animation: PendingAnimation, toAllApps: Boolean) = Unit
+    open fun onAllAppsAnimationPending(
+        animation: PendingAnimation,
+        toAllApps: Boolean,
+        showKeyboard: Boolean,
+    ) = Unit
 
     companion object {
         @JvmStatic
diff --git a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
index 2dba263..9126c4b 100644
--- a/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/overlay/TaskbarOverlayController.java
@@ -62,13 +62,13 @@
         @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.
-            hideWindow();
+            hideWindowOnTaskStackChange();
         }
 
         @Override
@@ -79,9 +79,15 @@
             // callback.
             if (mControllers.getSharedState() != null
                     && mControllers.getSharedState().allAppsVisible) {
-                hideWindow();
+                hideWindowOnTaskStackChange();
             }
         }
+
+        private void hideWindowOnTaskStackChange() {
+            // A task was launched while overlay window was open, so stash Taskbar.
+            mControllers.taskbarStashController.updateAndAnimateTransientTaskbar(true);
+            hideWindow();
+        }
     };
 
     private DeviceProfile mLauncherDeviceProfile;
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/launcher3/uioverrides/states/OverviewModalTaskState.java b/quickstep/src/com/android/launcher3/uioverrides/states/OverviewModalTaskState.java
index b9221ee..3e7d45e 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/states/OverviewModalTaskState.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/states/OverviewModalTaskState.java
@@ -72,7 +72,7 @@
 
     @Override
     public boolean isTaskbarStashed(Launcher launcher) {
-        if (FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get()) {
+        if (FeatureFlags.enableGridOnlyOverview()) {
             return true;
         }
         return super.isTaskbarStashed(launcher);
diff --git a/quickstep/src/com/android/quickstep/BaseActivityInterface.java b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
index 5a9d80d..e316175 100644
--- a/quickstep/src/com/android/quickstep/BaseActivityInterface.java
+++ b/quickstep/src/com/android/quickstep/BaseActivityInterface.java
@@ -239,7 +239,7 @@
     public final void calculateTaskSize(Context context, DeviceProfile dp, Rect outRect,
             PagedOrientationHandler orientedState) {
         if (dp.isTablet) {
-            if (FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get()) {
+            if (FeatureFlags.enableGridOnlyOverview()) {
                 calculateGridTaskSize(context, dp, outRect, orientedState);
             } else {
                 calculateFocusTaskSize(context, dp, outRect);
@@ -336,7 +336,7 @@
             PagedOrientationHandler orientedState) {
         Resources res = context.getResources();
         Rect potentialTaskRect = new Rect();
-        if (FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get()) {
+        if (FeatureFlags.enableGridOnlyOverview()) {
             calculateGridSize(dp, potentialTaskRect);
         } else {
             calculateFocusTaskSize(context, dp, potentialTaskRect);
@@ -368,7 +368,7 @@
     public final void calculateModalTaskSize(Context context, DeviceProfile dp, Rect outRect,
             PagedOrientationHandler orientedState) {
         calculateTaskSize(context, dp, outRect, orientedState);
-        boolean isGridOnlyOverview = dp.isTablet && FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get();
+        boolean isGridOnlyOverview = dp.isTablet && FeatureFlags.enableGridOnlyOverview();
         int claimedSpaceBelow = isGridOnlyOverview
                 ? dp.overviewActionsTopMarginPx + dp.overviewActionsHeight + dp.stashedTaskbarHeight
                 : (dp.heightPx - outRect.bottom - dp.getInsets().bottom);
diff --git a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
index 076f4b1..bd3ccb7 100644
--- a/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskOverlayFactory.java
@@ -80,7 +80,7 @@
         boolean isInLandscape = orientedState.getTouchRotation() != ROTATION_0;
         boolean isTablet = activity.getDeviceProfile().isTablet;
 
-        boolean isGridOnlyOverview = isTablet && FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get();
+        boolean isGridOnlyOverview = isTablet && FeatureFlags.enableGridOnlyOverview();
         // Add overview actions to the menu when in in-place rotate landscape mode, or in
         // grid-only overview.
         if ((!canLauncherRotate && isInLandscape) || isGridOnlyOverview) {
diff --git a/quickstep/src/com/android/quickstep/views/ClearAllButton.java b/quickstep/src/com/android/quickstep/views/ClearAllButton.java
index 19ac1f8..cfb4d0d 100644
--- a/quickstep/src/com/android/quickstep/views/ClearAllButton.java
+++ b/quickstep/src/com/android/quickstep/views/ClearAllButton.java
@@ -16,7 +16,7 @@
 
 package com.android.quickstep.views;
 
-import static com.android.launcher3.config.FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW;
+import static com.android.launcher3.config.FeatureFlags.enableGridOnlyOverview;
 
 import android.content.Context;
 import android.util.AttributeSet;
@@ -251,7 +251,7 @@
     private float getOriginalTranslationY() {
         DeviceProfile deviceProfile = mActivity.getDeviceProfile();
         if (deviceProfile.isTablet) {
-            if (ENABLE_GRID_ONLY_OVERVIEW.get()) {
+            if (enableGridOnlyOverview()) {
                 return (getRecentsView().getLastComputedTaskSize().height()
                         + deviceProfile.overviewTaskThumbnailTopMarginPx) / 2.0f
                         + deviceProfile.overviewRowSpacing;
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/src/com/android/quickstep/views/OverviewActionsView.java b/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
index b31791a..9141c99 100644
--- a/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
+++ b/quickstep/src/com/android/quickstep/views/OverviewActionsView.java
@@ -289,7 +289,7 @@
             return 0;
         }
 
-        if (mDp.isTablet && FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get()) {
+        if (mDp.isTablet && FeatureFlags.enableGridOnlyOverview()) {
             return mDp.stashedTaskbarHeight;
         }
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 7cf47ce..fd82f50 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -42,7 +42,7 @@
 import static com.android.launcher3.Utilities.mapToRange;
 import static com.android.launcher3.Utilities.squaredHypot;
 import static com.android.launcher3.Utilities.squaredTouchSlop;
-import static com.android.launcher3.config.FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW;
+import static com.android.launcher3.config.FeatureFlags.enableGridOnlyOverview;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OVERVIEW_ACTIONS_SPLIT;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_CLEAR_ALL;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_TASK_DISMISS_SWIPE_UP;
@@ -1782,7 +1782,7 @@
                 newFocusedTaskView = getTaskViewAt(1);
             }
         }
-        mFocusedTaskViewId = newFocusedTaskView != null && !ENABLE_GRID_ONLY_OVERVIEW.get()
+        mFocusedTaskViewId = newFocusedTaskView != null && !enableGridOnlyOverview()
                 ? newFocusedTaskView.getTaskViewId() : INVALID_TASK_ID;
         updateTaskSize();
         if (newFocusedTaskView != null) {
@@ -1900,7 +1900,7 @@
      * Returns the number of tasks in the bottom row of the overview grid.
      */
     public int getBottomRowTaskCountForTablet() {
-        return getTaskViewCount() - mTopRowIdSet.size() - (ENABLE_GRID_ONLY_OVERVIEW.get() ? 0 : 1);
+        return getTaskViewCount() - mTopRowIdSet.size() - (enableGridOnlyOverview() ? 0 : 1);
     }
 
     protected void onTaskStackUpdated() {
@@ -1978,7 +1978,7 @@
         DeviceProfile dp = mActivity.getDeviceProfile();
         setOverviewGridEnabled(
                 mActivity.getStateManager().getState().displayOverviewTasksAsGrid(dp));
-        if (ENABLE_GRID_ONLY_OVERVIEW.get()) {
+        if (enableGridOnlyOverview()) {
             mActionsView.updateHiddenFlags(HIDDEN_ACTIONS_IN_MENU, dp.isTablet);
         }
         setPageSpacing(dp.overviewPageSpacing);
@@ -2140,7 +2140,7 @@
         Rect outRect = new Rect(mLastComputedTaskSize);
         outRect.offset(
                 -(primaryScroll - (selectedPageScroll + getOffsetFromScrollPosition(selectedPage))),
-                (int) (showAsGrid() && ENABLE_GRID_ONLY_OVERVIEW.get() && !isTopRow
+                (int) (showAsGrid() && enableGridOnlyOverview() && !isTopRow
                         ? mTopBottomRowHeightDiff : 0));
         return outRect;
     }
@@ -2726,7 +2726,7 @@
 
         boolean runningTaskTileHidden = mRunningTaskTileHidden;
         setCurrentTask(runningTaskViewId);
-        mFocusedTaskViewId = ENABLE_GRID_ONLY_OVERVIEW.get() ? INVALID_TASK_ID : runningTaskViewId;
+        mFocusedTaskViewId = enableGridOnlyOverview() ? INVALID_TASK_ID : runningTaskViewId;
         runOnPageScrollsInitialized(() -> setCurrentPage(getRunningTaskIndex()));
         setRunningTaskViewShowScreenshot(false);
         setRunningTaskHidden(runningTaskTileHidden);
@@ -2919,7 +2919,7 @@
                     // Desktop task view is hidden, skip it from grid calculations
                     continue;
                 }
-                if (!ENABLE_GRID_ONLY_OVERVIEW.get()) {
+                if (!enableGridOnlyOverview()) {
                     // Only apply x-translation when using legacy overview grid
                     gridTranslations[i] += mIsRtl ? taskWidthAndSpacing : -taskWidthAndSpacing;
                 }
@@ -3432,7 +3432,7 @@
 
             int topGridRowSize = mTopRowIdSet.size();
             int bottomGridRowSize = taskCount - mTopRowIdSet.size()
-                    - (ENABLE_GRID_ONLY_OVERVIEW.get() ? 0 : 1);
+                    - (enableGridOnlyOverview() ? 0 : 1);
             boolean topRowLonger = topGridRowSize > bottomGridRowSize;
             boolean bottomRowLonger = bottomGridRowSize > topGridRowSize;
             boolean dismissedTaskFromTop = mTopRowIdSet.contains(dismissedTaskViewId);
@@ -3445,7 +3445,7 @@
             }
             int longRowWidth = Math.max(topGridRowSize, bottomGridRowSize)
                     * (mLastComputedGridTaskSize.width() + mPageSpacing);
-            if (!ENABLE_GRID_ONLY_OVERVIEW.get() && !isStagingFocusedTask) {
+            if (!enableGridOnlyOverview() && !isStagingFocusedTask) {
                 longRowWidth += mLastComputedTaskSize.width() + mPageSpacing;
             }
 
@@ -3853,7 +3853,7 @@
                     } else {
                         // Update focus task and its size.
                         if (finalIsFocusedTaskDismissed && finalNextFocusedTaskView != null) {
-                            mFocusedTaskViewId = ENABLE_GRID_ONLY_OVERVIEW.get()
+                            mFocusedTaskViewId = enableGridOnlyOverview()
                                     ? INVALID_TASK_ID
                                     : finalNextFocusedTaskView.getTaskViewId();
                             mTopRowIdSet.remove(mFocusedTaskViewId);
@@ -4358,7 +4358,7 @@
         int midpoint = runningTask == null ? -1 : indexOfChild(runningTask);
         int modalMidpoint = getCurrentPage();
         boolean isModalGridWithoutFocusedTask =
-                showAsGrid && ENABLE_GRID_ONLY_OVERVIEW.get() && mTaskModalness > 0;
+                showAsGrid && enableGridOnlyOverview() && mTaskModalness > 0;
         if (isModalGridWithoutFocusedTask) {
             modalMidpoint = indexOfChild(mSelectedTask);
         }
@@ -4420,7 +4420,7 @@
                 redrawLiveTile();
             }
 
-            if (showAsGrid && ENABLE_GRID_ONLY_OVERVIEW.get() && child instanceof TaskView) {
+            if (showAsGrid && enableGridOnlyOverview() && child instanceof TaskView) {
                 float totalTranslationY = getVerticalOffsetSize(i, modalOffset);
                 FloatProperty translationPropertyY =
                         ((TaskView) child).getSecondaryTaskOffsetTranslationProperty();
@@ -4529,7 +4529,7 @@
      * @param offsetProgress From 0 to 1 where 0 means no offset and 1 means offset offscreen.
      */
     private float getVerticalOffsetSize(int childIndex, float offsetProgress) {
-        if (offsetProgress == 0 || !(showAsGrid() && ENABLE_GRID_ONLY_OVERVIEW.get())
+        if (offsetProgress == 0 || !(showAsGrid() && enableGridOnlyOverview())
                 || mSelectedTask == null) {
             // Don't bother calculating everything below if we won't offset vertically.
             return 0;
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index a2976a8..fcaa0d7 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -1137,7 +1137,7 @@
         if (dp.isTablet) {
             int alignedOptionIndex = 0;
             if (getRecentsView().isOnGridBottomRow(menuContainer.getTaskView()) && dp.isLandscape) {
-                if (FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get()) {
+                if (FeatureFlags.enableGridOnlyOverview()) {
                     // With no focused task, there is less available space below the tasks, so align
                     // the arrow to the third option in the menu.
                     alignedOptionIndex = 2;
diff --git a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
index 2752003..eded1c9 100644
--- a/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
+++ b/quickstep/tests/src/com/android/quickstep/NavigationModeSwitchRule.java
@@ -26,10 +26,7 @@
 
 import android.content.Context;
 import android.content.pm.PackageManager;
-import android.graphics.Point;
-import android.os.SystemClock;
 import android.util.Log;
-import android.view.MotionEvent;
 
 import androidx.test.uiautomator.UiDevice;
 
@@ -118,18 +115,6 @@
                 private void evaluateWithThreeButtons() throws Throwable {
                     if (setActiveOverlay(mLauncher, NAV_BAR_MODE_3BUTTON_OVERLAY,
                             LauncherInstrumentation.NavigationModel.THREE_BUTTON, description)) {
-                        // After switching to three button, ensure that we interact with the screen
-                        // within the app area to reset the frozen recents state (if a quickswitch
-                        // was just performed), otherwise the list may be in the wrong order
-                        // spatially when executing the next test
-                        final Point center = new Point(mLauncher.getDevice().getDisplayWidth() / 2,
-                                mLauncher.getDevice().getDisplayHeight() / 2);
-                        final long clickTime = SystemClock.uptimeMillis();
-                        mLauncher.sendPointer(clickTime, clickTime, MotionEvent.ACTION_DOWN, center,
-                                LauncherInstrumentation.GestureScope.DONT_EXPECT_PILFER);
-                        mLauncher.sendPointer(clickTime, clickTime, MotionEvent.ACTION_UP, center,
-                                LauncherInstrumentation.GestureScope.DONT_EXPECT_PILFER);
-
                         base.evaluate();
                     }
                 }
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/AppWidgetsRestoredReceiver.java b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
index 55b8fcc..ec874b9 100644
--- a/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
+++ b/src/com/android/launcher3/AppWidgetsRestoredReceiver.java
@@ -1,12 +1,28 @@
 package com.android.launcher3;
 
+import static android.os.Process.myUserHandle;
+
+import android.appwidget.AppWidgetHost;
 import android.appwidget.AppWidgetManager;
+import android.appwidget.AppWidgetProviderInfo;
 import android.content.BroadcastReceiver;
 import android.content.Context;
 import android.content.Intent;
+import android.database.Cursor;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
+import androidx.annotation.WorkerThread;
+
+import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.model.LoaderTask;
+import com.android.launcher3.model.ModelDbController;
+import com.android.launcher3.model.WidgetsModel;
+import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.RestoreDbTask;
+import com.android.launcher3.util.ContentWriter;
+import com.android.launcher3.util.IntArray;
 import com.android.launcher3.widget.LauncherWidgetHolder;
 
 public class AppWidgetsRestoredReceiver extends BroadcastReceiver {
@@ -31,4 +47,131 @@
             }
         }
     }
-}
\ No newline at end of file
+
+    /**
+     * Updates the app widgets whose id has changed during the restore process.
+     */
+    @WorkerThread
+    public static void restoreAppWidgetIds(Context context, ModelDbController controller,
+            int[] oldWidgetIds, int[] newWidgetIds, @NonNull AppWidgetHost host) {
+        if (WidgetsModel.GO_DISABLE_WIDGETS) {
+            Log.e(TAG, "Skipping widget ID remap as widgets not supported");
+            host.deleteHost();
+            return;
+        }
+        if (!RestoreDbTask.isPending(context)) {
+            // Someone has already gone through our DB once, probably LoaderTask. Skip any further
+            // modifications of the DB.
+            Log.e(TAG, "Skipping widget ID remap as DB already in use");
+            for (int widgetId : newWidgetIds) {
+                Log.d(TAG, "Deleting widgetId: " + widgetId);
+                host.deleteAppWidgetId(widgetId);
+            }
+            return;
+        }
+
+        final AppWidgetManager widgets = AppWidgetManager.getInstance(context);
+
+        Log.d(TAG, "restoreAppWidgetIds: "
+                + "oldWidgetIds=" + IntArray.wrap(oldWidgetIds).toConcatString()
+                + ", newWidgetIds=" + IntArray.wrap(newWidgetIds).toConcatString());
+
+        // TODO(b/234700507): Remove the logs after the bug is fixed
+        logDatabaseWidgetInfo(controller);
+
+        for (int i = 0; i < oldWidgetIds.length; i++) {
+            Log.i(TAG, "Widget state restore id " + oldWidgetIds[i] + " => " + newWidgetIds[i]);
+
+            final AppWidgetProviderInfo provider = widgets.getAppWidgetInfo(newWidgetIds[i]);
+            final int state;
+            if (LoaderTask.isValidProvider(provider)) {
+                // This will ensure that we show 'Click to setup' UI if required.
+                state = LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
+            } else {
+                state = LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
+            }
+
+            // b/135926478: Work profile widget restore is broken in platform. This forces us to
+            // recreate the widget during loading with the correct host provider.
+            long mainProfileId = UserCache.INSTANCE.get(context)
+                    .getSerialNumberForUser(myUserHandle());
+            long controllerProfileId = controller.getSerialNumberForUser(myUserHandle());
+            String oldWidgetId = Integer.toString(oldWidgetIds[i]);
+            final String where = "appWidgetId=? and (restored & 1) = 1 and profileId=?";
+            String profileId = Long.toString(mainProfileId);
+            final String[] args = new String[] { oldWidgetId, profileId };
+            Log.d(TAG, "restoreAppWidgetIds: querying profile id=" + profileId
+                    + " with controller profile ID=" + controllerProfileId);
+            int result = new ContentWriter(context,
+                            new ContentWriter.CommitParams(controller, where, args))
+                    .put(LauncherSettings.Favorites.APPWIDGET_ID, newWidgetIds[i])
+                    .put(LauncherSettings.Favorites.RESTORED, state)
+                    .commit();
+            if (result == 0) {
+                // TODO(b/234700507): Remove the logs after the bug is fixed
+                Log.e(TAG, "restoreAppWidgetIds: remapping failed since the widget is not in"
+                        + " the database anymore");
+                try (Cursor cursor = controller.getDb().query(
+                        Favorites.TABLE_NAME,
+                        new String[]{Favorites.APPWIDGET_ID},
+                        "appWidgetId=?", new String[]{oldWidgetId}, null, null, null)) {
+                    if (!cursor.moveToFirst()) {
+                        // The widget no long exists.
+                        Log.d(TAG, "Deleting widgetId: " + newWidgetIds[i] + " with old id: "
+                                + oldWidgetId);
+                        host.deleteAppWidgetId(newWidgetIds[i]);
+                    }
+                }
+            }
+        }
+
+        LauncherAppState app = LauncherAppState.getInstanceNoCreate();
+        if (app != null) {
+            app.getModel().forceReload();
+        }
+    }
+
+    private static void logDatabaseWidgetInfo(ModelDbController controller) {
+        try (Cursor cursor = controller.getDb().query(Favorites.TABLE_NAME,
+                new String[]{Favorites.APPWIDGET_ID, Favorites.RESTORED, Favorites.PROFILE_ID},
+                Favorites.APPWIDGET_ID + "!=" + LauncherAppWidgetInfo.NO_ID, null,
+                null, null, null)) {
+            IntArray widgetIdList = new IntArray();
+            IntArray widgetRestoreList = new IntArray();
+            IntArray widgetProfileIdList = new IntArray();
+
+            if (cursor.moveToFirst()) {
+                final int widgetIdColumnIndex = cursor.getColumnIndex(Favorites.APPWIDGET_ID);
+                final int widgetRestoredColumnIndex = cursor.getColumnIndex(Favorites.RESTORED);
+                final int widgetProfileIdIndex = cursor.getColumnIndex(Favorites.PROFILE_ID);
+                while (!cursor.isAfterLast()) {
+                    int widgetId = cursor.getInt(widgetIdColumnIndex);
+                    int widgetRestoredFlag = cursor.getInt(widgetRestoredColumnIndex);
+                    int widgetProfileId = cursor.getInt(widgetProfileIdIndex);
+
+                    widgetIdList.add(widgetId);
+                    widgetRestoreList.add(widgetRestoredFlag);
+                    widgetProfileIdList.add(widgetProfileId);
+                    cursor.moveToNext();
+                }
+            }
+
+            StringBuilder builder = new StringBuilder();
+            builder.append("[");
+            for (int i = 0; i < widgetIdList.size(); i++) {
+                builder.append("[")
+                        .append(widgetIdList.get(i))
+                        .append(", ")
+                        .append(widgetRestoreList.get(i))
+                        .append(", ")
+                        .append(widgetProfileIdList.get(i))
+                        .append("]");
+            }
+            builder.append("]");
+            Log.d(TAG, "restoreAppWidgetIds: all widget ids in database: "
+                    + builder.toString());
+        } catch (Exception ex) {
+            Log.e(TAG, "Getting widget ids from the database failed", ex);
+        }
+    }
+}
diff --git a/src/com/android/launcher3/DeviceProfile.java b/src/com/android/launcher3/DeviceProfile.java
index e7d1db4..986ec73 100644
--- a/src/com/android/launcher3/DeviceProfile.java
+++ b/src/com/android/launcher3/DeviceProfile.java
@@ -40,6 +40,7 @@
 import android.graphics.PointF;
 import android.graphics.Rect;
 import android.util.DisplayMetrics;
+import android.util.Log;
 import android.util.SparseArray;
 import android.view.Surface;
 
@@ -363,6 +364,12 @@
                 mTypeIndex = INDEX_DEFAULT;
             }
         }
+        if (InvariantDeviceProfile.sDebug) {
+            Log.d("b/298077774",
+                    "DeviceProfile - mTypeIndex: " + mTypeIndex + ": " + dpPointFToString(
+                            "minCellSize",
+                            inv.minCellSize[mTypeIndex]));
+        }
 
         isTransientTaskbar = DisplayController.isTransientTaskbar(context);
         if (isTransientTaskbar) {
@@ -1715,7 +1722,7 @@
 
     /** Gets the space that the overview actions will take, including bottom margin. */
     public int getOverviewActionsClaimedSpace() {
-        int overviewActionsSpace = isTablet && FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW.get()
+        int overviewActionsSpace = isTablet && FeatureFlags.enableGridOnlyOverview()
                 ? 0
                 : (overviewActionsTopMarginPx + overviewActionsHeight);
         return overviewActionsSpace + getOverviewActionsClaimedSpaceBelow();
diff --git a/src/com/android/launcher3/InvariantDeviceProfile.java b/src/com/android/launcher3/InvariantDeviceProfile.java
index e3de79e..3332f0e 100644
--- a/src/com/android/launcher3/InvariantDeviceProfile.java
+++ b/src/com/android/launcher3/InvariantDeviceProfile.java
@@ -72,6 +72,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Locale;
 import java.util.stream.Collectors;
 
 public class InvariantDeviceProfile {
@@ -81,6 +82,9 @@
     public static final MainThreadInitializedObject<InvariantDeviceProfile> INSTANCE =
             new MainThreadInitializedObject<>(InvariantDeviceProfile::new);
 
+    @VisibleForTesting
+    public static boolean sDebug = false;
+
     @Retention(RetentionPolicy.SOURCE)
     @IntDef({TYPE_PHONE, TYPE_MULTI_DISPLAY, TYPE_TABLET})
     public @interface DeviceType {}
@@ -683,10 +687,29 @@
         for (int i = 0; i < points.size() && i < KNEARESTNEIGHBOR; ++i) {
             DisplayOption p = points.get(i);
             float w = weight(width, height, p.minWidthDps, p.minHeightDps, WEIGHT_POWER);
+            if (sDebug) {
+                for (int j = INDEX_DEFAULT; j < COUNT_SIZES; j++) {
+                    Log.d("b/298077774",
+                            "invDistWeightedInterpolate - points[" + i + "]: " + String.format(
+                                    Locale.ENGLISH, "\t%s: PointF(%.1f, %.1f)dp",
+                                    "minCellSize[" + j + "]",
+                                    p.minCellSize[j].x, p.minCellSize[j].y));
+                }
+                Log.d("b/298077774", "invDistWeightedInterpolate - w[" + i + "]: " + w);
+            }
             weights += w;
             out.add(new DisplayOption().add(p).multiply(w));
         }
         out.multiply(1.0f / weights);
+        if (sDebug) {
+            Log.d("b/298077774", "invDistWeightedInterpolate - weights: " + weights);
+            for (int j = INDEX_DEFAULT; j < COUNT_SIZES; j++) {
+                Log.d("b/298077774",
+                        "invDistWeightedInterpolate - out: " + String.format(Locale.ENGLISH,
+                                "\t%s: PointF(%.1f, %.1f)dp", "minCellSize[" + j + "]",
+                                out.minCellSize[j].x, out.minCellSize[j].y));
+            }
+        }
 
         // Since the bitmaps are persisted, ensure that all bitmap sizes are not larger than
         // predefined size to avoid cache invalidation
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 5653c3c..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;
@@ -886,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,
@@ -1500,7 +1501,8 @@
             AppWidgetHostView hostView, LauncherAppWidgetProviderInfo appWidgetInfo) {
 
         if (appWidgetInfo == null) {
-            appWidgetInfo = mAppWidgetManager.getLauncherAppWidgetInfo(appWidgetId);
+            appWidgetInfo = mAppWidgetManager.getLauncherAppWidgetInfo(appWidgetId,
+                    itemInfo.getTargetComponent());
         }
 
         if (hostView == null) {
@@ -1995,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();
             }
@@ -2642,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 03a7d87..1ade51a 100644
--- a/src/com/android/launcher3/config/FeatureFlags.java
+++ b/src/com/android/launcher3/config/FeatureFlags.java
@@ -27,6 +27,7 @@
 import androidx.annotation.VisibleForTesting;
 
 import com.android.launcher3.BuildConfig;
+import com.android.launcher3.Flags;
 import com.android.launcher3.Utilities;
 
 import java.util.function.Predicate;
@@ -141,7 +142,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 +176,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,
@@ -312,9 +316,14 @@
                     + "start receiving the events");
 
     // TODO(Block 23): Clean up flags
+    // Aconfig migration complete for ENABLE_GRID_ONLY_OVERVIEW.
+    @VisibleForTesting
     public static final BooleanFlag ENABLE_GRID_ONLY_OVERVIEW = getDebugFlag(270397206,
             "ENABLE_GRID_ONLY_OVERVIEW", TEAMFOOD,
             "Enable a grid-only overview without a focused task.");
+    public static boolean enableGridOnlyOverview() {
+        return ENABLE_GRID_ONLY_OVERVIEW.get() || Flags.enableGridOnlyOverview();
+    }
 
     public static final BooleanFlag ENABLE_CURSOR_HOVER_STATES = getDebugFlag(243191650,
             "ENABLE_CURSOR_HOVER_STATES", TEAMFOOD,
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/provider/RestoreDbTask.java b/src/com/android/launcher3/provider/RestoreDbTask.java
index 10005e5..4725dd1 100644
--- a/src/com/android/launcher3/provider/RestoreDbTask.java
+++ b/src/com/android/launcher3/provider/RestoreDbTask.java
@@ -28,8 +28,6 @@
 
 import android.app.backup.BackupManager;
 import android.appwidget.AppWidgetHost;
-import android.appwidget.AppWidgetManager;
-import android.appwidget.AppWidgetProviderInfo;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
@@ -44,26 +42,20 @@
 
 import androidx.annotation.NonNull;
 import androidx.annotation.VisibleForTesting;
-import androidx.annotation.WorkerThread;
 
+import com.android.launcher3.AppWidgetsRestoredReceiver;
 import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherPrefs;
-import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.DeviceGridState;
-import com.android.launcher3.model.LoaderTask;
 import com.android.launcher3.model.ModelDbController;
-import com.android.launcher3.model.WidgetsModel;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.pm.UserCache;
 import com.android.launcher3.provider.LauncherDbUtils.SQLiteTransaction;
 import com.android.launcher3.uioverrides.ApiWrapper;
-import com.android.launcher3.util.ContentWriter;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.LogConfig;
 
@@ -385,13 +377,11 @@
                 .putSync(RESTORE_DEVICE.to(new DeviceGridState(context).getDeviceType()));
     }
 
-    @WorkerThread
-    @VisibleForTesting
-    void restoreAppWidgetIdsIfExists(Context context, ModelDbController controller) {
+    private void restoreAppWidgetIdsIfExists(Context context, ModelDbController controller) {
         LauncherPrefs lp = LauncherPrefs.get(context);
         if (lp.has(APP_WIDGET_IDS, OLD_APP_WIDGET_IDS)) {
             AppWidgetHost host = new AppWidgetHost(context, APPWIDGET_HOST_ID);
-            restoreAppWidgetIds(context, controller,
+            AppWidgetsRestoredReceiver.restoreAppWidgetIds(context, controller,
                     IntArray.fromConcatString(lp.get(OLD_APP_WIDGET_IDS)).toArray(),
                     IntArray.fromConcatString(lp.get(APP_WIDGET_IDS)).toArray(),
                     host);
@@ -402,133 +392,6 @@
         lp.remove(APP_WIDGET_IDS, OLD_APP_WIDGET_IDS);
     }
 
-    /**
-     * Updates the app widgets whose id has changed during the restore process.
-     */
-    @WorkerThread
-    private void restoreAppWidgetIds(Context context, ModelDbController controller,
-            int[] oldWidgetIds, int[] newWidgetIds, @NonNull AppWidgetHost host) {
-        if (WidgetsModel.GO_DISABLE_WIDGETS) {
-            Log.e(TAG, "Skipping widget ID remap as widgets not supported");
-            host.deleteHost();
-            return;
-        }
-        if (!RestoreDbTask.isPending(context)) {
-            // Someone has already gone through our DB once, probably LoaderTask. Skip any further
-            // modifications of the DB.
-            Log.e(TAG, "Skipping widget ID remap as DB already in use");
-            for (int widgetId : newWidgetIds) {
-                Log.d(TAG, "Deleting widgetId: " + widgetId);
-                host.deleteAppWidgetId(widgetId);
-            }
-            return;
-        }
-
-        final AppWidgetManager widgets = AppWidgetManager.getInstance(context);
-
-        Log.d(TAG, "restoreAppWidgetIds: "
-                + "oldWidgetIds=" + IntArray.wrap(oldWidgetIds).toConcatString()
-                + ", newWidgetIds=" + IntArray.wrap(newWidgetIds).toConcatString());
-
-        // TODO(b/234700507): Remove the logs after the bug is fixed
-        logDatabaseWidgetInfo(controller);
-
-        for (int i = 0; i < oldWidgetIds.length; i++) {
-            Log.i(TAG, "Widget state restore id " + oldWidgetIds[i] + " => " + newWidgetIds[i]);
-
-            final AppWidgetProviderInfo provider = widgets.getAppWidgetInfo(newWidgetIds[i]);
-            final int state;
-            if (LoaderTask.isValidProvider(provider)) {
-                // This will ensure that we show 'Click to setup' UI if required.
-                state = LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
-            } else {
-                state = LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
-            }
-
-            // b/135926478: Work profile widget restore is broken in platform. This forces us to
-            // recreate the widget during loading with the correct host provider.
-            long mainProfileId = UserCache.INSTANCE.get(context)
-                    .getSerialNumberForUser(myUserHandle());
-            long controllerProfileId = controller.getSerialNumberForUser(myUserHandle());
-            String oldWidgetId = Integer.toString(oldWidgetIds[i]);
-            final String where = "appWidgetId=? and (restored & 1) = 1 and profileId=?";
-            String profileId = Long.toString(mainProfileId);
-            final String[] args = new String[] { oldWidgetId, profileId };
-            Log.d(TAG, "restoreAppWidgetIds: querying profile id=" + profileId
-                    + " with controller profile ID=" + controllerProfileId);
-            int result = new ContentWriter(context,
-                    new ContentWriter.CommitParams(controller, where, args))
-                    .put(LauncherSettings.Favorites.APPWIDGET_ID, newWidgetIds[i])
-                    .put(LauncherSettings.Favorites.RESTORED, state)
-                    .commit();
-            if (result == 0) {
-                // TODO(b/234700507): Remove the logs after the bug is fixed
-                Log.e(TAG, "restoreAppWidgetIds: remapping failed since the widget is not in"
-                        + " the database anymore");
-                try (Cursor cursor = controller.getDb().query(
-                        Favorites.TABLE_NAME,
-                        new String[]{Favorites.APPWIDGET_ID},
-                        "appWidgetId=?", new String[]{oldWidgetId}, null, null, null)) {
-                    if (!cursor.moveToFirst()) {
-                        // The widget no long exists.
-                        Log.d(TAG, "Deleting widgetId: " + newWidgetIds[i] + " with old id: "
-                                + oldWidgetId);
-                        host.deleteAppWidgetId(newWidgetIds[i]);
-                    }
-                }
-            }
-        }
-
-        LauncherAppState app = LauncherAppState.getInstanceNoCreate();
-        if (app != null) {
-            app.getModel().forceReload();
-        }
-    }
-
-    private static void logDatabaseWidgetInfo(ModelDbController controller) {
-        try (Cursor cursor = controller.getDb().query(Favorites.TABLE_NAME,
-                new String[]{Favorites.APPWIDGET_ID, Favorites.RESTORED, Favorites.PROFILE_ID},
-                Favorites.APPWIDGET_ID + "!=" + LauncherAppWidgetInfo.NO_ID, null,
-                null, null, null)) {
-            IntArray widgetIdList = new IntArray();
-            IntArray widgetRestoreList = new IntArray();
-            IntArray widgetProfileIdList = new IntArray();
-
-            if (cursor.moveToFirst()) {
-                final int widgetIdColumnIndex = cursor.getColumnIndex(Favorites.APPWIDGET_ID);
-                final int widgetRestoredColumnIndex = cursor.getColumnIndex(Favorites.RESTORED);
-                final int widgetProfileIdIndex = cursor.getColumnIndex(Favorites.PROFILE_ID);
-                while (!cursor.isAfterLast()) {
-                    int widgetId = cursor.getInt(widgetIdColumnIndex);
-                    int widgetRestoredFlag = cursor.getInt(widgetRestoredColumnIndex);
-                    int widgetProfileId = cursor.getInt(widgetProfileIdIndex);
-
-                    widgetIdList.add(widgetId);
-                    widgetRestoreList.add(widgetRestoredFlag);
-                    widgetProfileIdList.add(widgetProfileId);
-                    cursor.moveToNext();
-                }
-            }
-
-            StringBuilder builder = new StringBuilder();
-            builder.append("[");
-            for (int i = 0; i < widgetIdList.size(); i++) {
-                builder.append("[")
-                        .append(widgetIdList.get(i))
-                        .append(", ")
-                        .append(widgetRestoreList.get(i))
-                        .append(", ")
-                        .append(widgetProfileIdList.get(i))
-                        .append("]");
-            }
-            builder.append("]");
-            Log.d(TAG, "restoreAppWidgetIds: all widget ids in database: "
-                    + builder.toString());
-        } catch (Exception ex) {
-            Log.e(TAG, "Getting widget ids from the database failed", ex);
-        }
-    }
-
     public static void setRestoredAppWidgetIds(Context context, @NonNull int[] oldIds,
             @NonNull int[] newIds) {
         LauncherPrefs.get(context).putSync(
diff --git a/src/com/android/launcher3/testing/TestInformationHandler.java b/src/com/android/launcher3/testing/TestInformationHandler.java
index 07a1b82..a2991c7 100644
--- a/src/com/android/launcher3/testing/TestInformationHandler.java
+++ b/src/com/android/launcher3/testing/TestInformationHandler.java
@@ -16,7 +16,7 @@
 package com.android.launcher3.testing;
 
 import static com.android.launcher3.allapps.AllAppsStore.DEFER_UPDATES_TEST;
-import static com.android.launcher3.config.FeatureFlags.ENABLE_GRID_ONLY_OVERVIEW;
+import static com.android.launcher3.config.FeatureFlags.enableGridOnlyOverview;
 import static com.android.launcher3.config.FeatureFlags.FOLDABLE_SINGLE_PAGE;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 
@@ -249,7 +249,7 @@
 
             case TestProtocol.REQUEST_FLAG_ENABLE_GRID_ONLY_OVERVIEW: {
                 response.putBoolean(TestProtocol.TEST_INFO_RESPONSE_FIELD,
-                        ENABLE_GRID_ONLY_OVERVIEW.get());
+                        enableGridOnlyOverview());
                 return response;
             }
 
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/provider/RestoreDbTaskTest.java b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
index 15f6177..73bb586 100644
--- a/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
+++ b/tests/src/com/android/launcher3/provider/RestoreDbTaskTest.java
@@ -19,30 +19,17 @@
 
 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
 
-import static com.android.launcher3.LauncherPrefs.APP_WIDGET_IDS;
-import static com.android.launcher3.LauncherPrefs.OLD_APP_WIDGET_IDS;
-import static com.android.launcher3.LauncherPrefs.RESTORE_DEVICE;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
 import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
-import static com.android.launcher3.widget.LauncherWidgetHolder.APPWIDGET_HOST_ID;
-
-import static com.google.common.truth.Truth.assertThat;
 
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doReturn;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
-import static org.mockito.Mockito.verifyZeroInteractions;
-import static org.mockito.Mockito.when;
 
 import android.app.backup.BackupManager;
-import android.appwidget.AppWidgetHost;
 import android.content.ContentValues;
 import android.content.Context;
 import android.content.Intent;
@@ -56,7 +43,6 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherPrefs;
 import com.android.launcher3.LauncherSettings;
 import com.android.launcher3.LauncherSettings.Favorites;
 import com.android.launcher3.model.ModelDbController;
@@ -66,10 +52,6 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.mockito.Mockito;
-
-import java.util.Arrays;
-import java.util.stream.IntStream;
 
 /**
  * Tests for {@link RestoreDbTask}
@@ -84,21 +66,11 @@
 
     private LauncherModelHelper mModelHelper;
     private Context mContext;
-    private RestoreDbTask mTask;
-    private ModelDbController mMockController;
-    private SQLiteDatabase mMockDb;
-    private Cursor mMockCursor;
-    private LauncherPrefs mPrefs;
 
     @Before
     public void setup() {
         mModelHelper = new LauncherModelHelper();
         mContext = mModelHelper.sandboxContext;
-        mTask = new RestoreDbTask();
-        mMockController = Mockito.mock(ModelDbController.class);
-        mMockDb = mock(SQLiteDatabase.class);
-        mMockCursor = mock(Cursor.class);
-        mPrefs = new LauncherPrefs(mContext);
     }
 
     @After
@@ -176,7 +148,8 @@
         assertEquals(10, getItemCountForProfile(db, myProfileId_old));
         assertEquals(6, getItemCountForProfile(db, workProfileId_old));
 
-        mTask.sanitizeDB(mContext, controller, controller.getDb(), bm);
+        RestoreDbTask task = new RestoreDbTask();
+        task.sanitizeDB(mContext, controller, controller.getDb(), bm);
 
         // All the data has been migrated to the new user ids
         assertEquals(0, getItemCountForProfile(db, myProfileId_old));
@@ -205,7 +178,8 @@
         assertEquals(10, getItemCountForProfile(db, myProfileId_old));
         assertEquals(6, getItemCountForProfile(db, workProfileId_old));
 
-        mTask.sanitizeDB(mContext, controller, controller.getDb(), bm);
+        RestoreDbTask task = new RestoreDbTask();
+        task.sanitizeDB(mContext, controller, controller.getDb(), bm);
 
         // All the data has been migrated to the new user ids
         assertEquals(0, getItemCountForProfile(db, myProfileId_old));
@@ -214,83 +188,6 @@
         assertEquals(10, getCount(db, "select * from favorites"));
     }
 
-    @Test
-    public void givenLauncherPrefsHasNoIds_whenRestoreAppWidgetIdsIfExists_thenIdsAreRemoved() {
-        // When
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
-        // Then
-        assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
-    }
-
-    @Test
-    public void givenNoPendingRestore_WhenRestoreAppWidgetIds_ThenRemoveNewWidgetIds() {
-        // Given
-        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
-        int[] expectedOldIds = generateOldWidgetIds(expectedHost);
-        int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
-        when(mMockController.getDb()).thenReturn(mMockDb);
-        mPrefs.remove(RESTORE_DEVICE);
-
-        // When
-        RestoreDbTask.setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
-
-        // Then
-        assertThat(expectedHost.getAppWidgetIds()).isEqualTo(expectedOldIds);
-        assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
-        verifyZeroInteractions(mMockController);
-    }
-
-    @Test
-    public void givenRestoreWithNonExistingWidgets_WhenRestoreAppWidgetIds_ThenRemoveNewIds() {
-        // Given
-        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
-        int[] expectedOldIds = generateOldWidgetIds(expectedHost);
-        int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
-        when(mMockController.getDb()).thenReturn(mMockDb);
-        when(mMockDb.query(any(), any(), any(), any(), any(), any(), any())).thenReturn(
-                mMockCursor);
-        when(mMockCursor.moveToFirst()).thenReturn(false);
-        RestoreDbTask.setPending(mContext);
-
-        // When
-        RestoreDbTask.setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
-
-        // Then
-        assertThat(expectedHost.getAppWidgetIds()).isEqualTo(expectedOldIds);
-        assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
-        verify(mMockController, times(expectedOldIds.length)).update(any(), any(), any(), any());
-    }
-
-    @Test
-    public void givenRestore_WhenRestoreAppWidgetIds_ThenAddNewIds() {
-        // Given
-        AppWidgetHost expectedHost = new AppWidgetHost(mContext, APPWIDGET_HOST_ID);
-        int[] expectedOldIds = generateOldWidgetIds(expectedHost);
-        int[] expectedNewIds = generateNewWidgetIds(expectedHost, expectedOldIds);
-        int[] allExpectedIds = IntStream.concat(
-                Arrays.stream(expectedOldIds),
-                Arrays.stream(expectedNewIds)
-        ).toArray();
-
-        when(mMockController.getDb()).thenReturn(mMockDb);
-        when(mMockDb.query(any(), any(), any(), any(), any(), any(), any()))
-                .thenReturn(mMockCursor);
-        when(mMockCursor.moveToFirst()).thenReturn(true);
-        when(mMockCursor.isAfterLast()).thenReturn(true);
-        RestoreDbTask.setPending(mContext);
-
-        // When
-        RestoreDbTask.setRestoredAppWidgetIds(mContext, expectedOldIds, expectedNewIds);
-        mTask.restoreAppWidgetIdsIfExists(mContext, mMockController);
-
-        // Then
-        assertThat(expectedHost.getAppWidgetIds()).isEqualTo(allExpectedIds);
-        assertThat(mPrefs.has(OLD_APP_WIDGET_IDS, APP_WIDGET_IDS)).isFalse();
-        verify(mMockController, times(expectedOldIds.length)).update(any(), any(), any(), any());
-    }
-
     private void addIconsBulk(MyModelDbController controller,
             int count, int screen, long profileId) {
         int columns = LauncherAppState.getIDP(mContext).numColumns;
@@ -373,19 +270,6 @@
         }
     }
 
-    private int[] generateOldWidgetIds(AppWidgetHost host) {
-        // generate some widget ids in case there are none
-        host.allocateAppWidgetId();
-        host.allocateAppWidgetId();
-        return host.getAppWidgetIds();
-    }
-
-    private int[] generateNewWidgetIds(AppWidgetHost host, int[] oldWidgetIds) {
-        // map as many new ids as old ids
-        return Arrays.stream(oldWidgetIds)
-                .map(id -> host.allocateAppWidgetId()).toArray();
-    }
-
     private class MyModelDbController extends ModelDbController {
 
         public final LongSparseArray<UserHandle> users = new LongSparseArray<>();
diff --git a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
index 9dca24b..64ab206 100644
--- a/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
+++ b/tests/src/com/android/launcher3/ui/widget/AddWidgetTest.java
@@ -16,7 +16,6 @@
 package com.android.launcher3.ui.widget;
 
 import static com.android.launcher3.ui.TaplTestsLauncher3.getAppPackageName;
-
 import static org.junit.Assert.assertNotNull;
 
 import android.platform.test.annotations.PlatinumTest;
@@ -96,4 +95,26 @@
         mLauncher.getWorkspace().getWorkspaceAppIcon("Shortcut")
                 .launch(getAppPackageName());
     }
+
+    /**
+     * Test dragging a widget to the workspace and resize it.
+     */
+    @Test
+    public void testResizeWidget() throws Throwable {
+        new FavoriteItemsTransaction(mTargetContext).commitAndLoadHome(mLauncher);
+
+        waitForLauncherCondition("Workspace didn't finish loading", l -> !l.isWorkspaceLoading());
+
+        final LauncherAppWidgetProviderInfo widgetInfo =
+                TestViewHelpers.findWidgetProvider(false /* hasConfigureScreen */);
+
+        WidgetResizeFrame resizeFrame = mLauncher
+                .getWorkspace()
+                .openAllWidgets()
+                .getWidget(widgetInfo.getLabel(mTargetContext.getPackageManager()))
+                .dragWidgetToWorkspace();
+
+        assertNotNull("Widget resize frame not shown after widget add", resizeFrame);
+        resizeFrame.resize();
+    }
 }
diff --git a/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java b/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java
index 8f51d04..ec1cbd8 100644
--- a/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java
+++ b/tests/tapl/com/android/launcher3/tapl/WidgetResizeFrame.java
@@ -15,6 +15,16 @@
  */
 package com.android.launcher3.tapl;
 
+import static com.android.launcher3.tapl.Launchable.DEFAULT_DRAG_STEPS;
+import static org.junit.Assert.assertTrue;
+
+import android.graphics.Point;
+import android.graphics.Rect;
+import android.os.SystemClock;
+import android.view.MotionEvent;
+
+import androidx.test.uiautomator.UiObject2;
+
 /** The resize frame that is shown for a widget on the workspace. */
 public class WidgetResizeFrame {
 
@@ -34,4 +44,36 @@
             mLauncher.getDevice().pressHome();
         }
     }
+
+    /** Resizes the widget to double its height, and returns the resize frame. */
+    public WidgetResizeFrame resize() {
+        try (LauncherInstrumentation.Closable e = mLauncher.eventsCheck();
+             LauncherInstrumentation.Closable c = mLauncher.addContextLayer(
+                     "want to resize the widget frame.")) {
+            UiObject2 widget = mLauncher.waitForLauncherObject("widget_resize_frame");
+            UiObject2 bottomResizeHandle =
+                    mLauncher.waitForLauncherObject("widget_resize_bottom_handle");
+            Rect originalWidgetSize = widget.getVisibleBounds();
+            Point targetStart = bottomResizeHandle.getVisibleCenter();
+            Point targetDest = bottomResizeHandle.getVisibleCenter();
+            targetDest.offset(0, originalWidgetSize.height());
+
+            final long downTime = SystemClock.uptimeMillis();
+            mLauncher.sendPointer(downTime, downTime, MotionEvent.ACTION_DOWN, targetStart,
+                    LauncherInstrumentation.GestureScope.DONT_EXPECT_PILFER);
+            mLauncher.movePointer(targetStart, targetDest, DEFAULT_DRAG_STEPS,
+                    true, downTime, downTime, true,
+                    LauncherInstrumentation.GestureScope.DONT_EXPECT_PILFER);
+            mLauncher.sendPointer(downTime, downTime, MotionEvent.ACTION_UP, targetDest,
+                    LauncherInstrumentation.GestureScope.DONT_EXPECT_PILFER);
+
+            try (LauncherInstrumentation.Closable c2 = mLauncher.addContextLayer(
+                         "want to return resized widget resize frame")) {
+                float newHeight = mLauncher.waitForLauncherObject(
+                        "widget_resize_frame").getVisibleBounds().height();
+                assertTrue("Widget not resized.", newHeight >= originalWidgetSize.height() * 2);
+                return this;
+            }
+        }
+    }
 }