diff --git a/aconfig/launcher_overview.aconfig b/aconfig/launcher_overview.aconfig
index 30b0d40..f379e22 100644
--- a/aconfig/launcher_overview.aconfig
+++ b/aconfig/launcher_overview.aconfig
@@ -89,3 +89,10 @@
       purpose: PURPOSE_BUGFIX
     }
 }
+
+flag {
+    name: "enable_overview_on_connected_displays"
+    namespace: "launcher_overview"
+    description: "Enable overview on connected displays."
+    bug: "363251602"
+}
diff --git a/quickstep/res/values/strings.xml b/quickstep/res/values/strings.xml
index 0474000..d2a7029 100644
--- a/quickstep/res/values/strings.xml
+++ b/quickstep/res/values/strings.xml
@@ -304,10 +304,6 @@
     <string name="taskbar_a11y_shown_with_bubbles_left_title">Taskbar &#38; bubbles left shown</string>
     <!-- Accessibility title for the Taskbar window appearing together with bubble bar on right. [CHAR_LIMIT=30] -->
     <string name="taskbar_a11y_shown_with_bubbles_right_title">Taskbar &#38; bubbles right shown</string>
-    <!-- Accessibility title for the Taskbar window being closed. [CHAR_LIMIT=30] -->
-    <string name="taskbar_a11y_hidden_title">Taskbar hidden</string>
-    <!-- Accessibility title for the Taskbar window being closed together with bubble bar. [CHAR_LIMIT=30] -->
-    <string name="taskbar_a11y_hidden_with_bubbles_title">Taskbar &#38; bubbles hidden</string>
     <!-- Accessibility title for the Taskbar window on phones. [CHAR_LIMIT=NONE] -->
     <string name="taskbar_phone_a11y_title">Navigation bar</string>
     <!-- Text in popup dialog for user to switch between always showing Taskbar or not. [CHAR LIMIT=30] -->
diff --git a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
index 56945ba..70868c5 100644
--- a/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
+++ b/quickstep/src/com/android/launcher3/hybridhotseat/HotseatPredictionModel.java
@@ -17,7 +17,6 @@
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.model.PredictionHelper.getAppTargetFromItemInfo;
-import static com.android.launcher3.model.PredictionHelper.isTrackedForHotseatPrediction;
 import static com.android.launcher3.model.PredictionHelper.wrapAppTargetWithItemLocation;
 
 import android.app.prediction.AppTarget;
@@ -27,9 +26,12 @@
 
 import com.android.launcher3.model.BgDataModel;
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
+import com.android.launcher3.model.PredictionHelper;
 import com.android.launcher3.model.data.ItemInfo;
 
 import java.util.ArrayList;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
 /**
  * Model helper for app predictions in workspace
@@ -43,13 +45,18 @@
      */
     public static Bundle convertDataModelToAppTargetBundle(Context context, BgDataModel dataModel) {
         Bundle bundle = new Bundle();
-        ArrayList<AppTargetEvent> events = new ArrayList<>();
-        ArrayList<ItemInfo> workspaceItems = dataModel.getAllWorkspaceItems();
-        for (ItemInfo item : workspaceItems) {
-            AppTarget target = getAppTargetFromItemInfo(context, item);
-            if (target != null && !isTrackedForHotseatPrediction(item)) continue;
-            events.add(wrapAppTargetWithItemLocation(target, AppTargetEvent.ACTION_PIN, item));
-        }
+        ArrayList<AppTargetEvent> events = dataModel.itemsIdMap
+                .stream()
+                .filter(PredictionHelper::isTrackedForHotseatPrediction)
+                .map(item -> {
+                    AppTarget target = getAppTargetFromItemInfo(context, item);
+                    return target != null
+                            ? wrapAppTargetWithItemLocation(target, AppTargetEvent.ACTION_PIN, item)
+                            : null;
+                })
+                .filter(Objects::nonNull)
+                .collect(Collectors.toCollection(ArrayList::new));
+
         ArrayList<AppTarget> currentTargets = new ArrayList<>();
         FixedContainerItems hotseatItems = dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
         if (hotseatItems != null) {
diff --git a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
index 25e1813..40e8fc2 100644
--- a/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
+++ b/quickstep/src/com/android/launcher3/model/QuickstepModelDelegate.java
@@ -462,7 +462,7 @@
     private Bundle getBundleForWidgetsOnWorkspace(Context context, BgDataModel dataModel) {
         Bundle bundle = new Bundle();
         ArrayList<AppTargetEvent> widgetEvents =
-                dataModel.getAllWorkspaceItems().stream()
+                dataModel.itemsIdMap.stream()
                         .filter(PredictionHelper::isTrackedForWidgetPrediction)
                         .map(item -> {
                             AppTarget target = getAppTargetFromItemInfo(context, item);
diff --git a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
index 9d9054e..40e1c10 100644
--- a/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
+++ b/quickstep/src/com/android/launcher3/model/WidgetsPredictionUpdateTask.java
@@ -18,6 +18,7 @@
 import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.Flags.enableTieredWidgetsByDefaultInPicker;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
 
 import static java.util.stream.Collectors.groupingBy;
 import static java.util.stream.Collectors.toMap;
@@ -69,9 +70,11 @@
             @NonNull AllAppsList apps) {
         Predicate<WidgetItem> predictedWidgetsFilter = enableTieredWidgetsByDefaultInPicker()
                 ? dataModel.widgetsModel.getPredictedWidgetsFilter() : null;
-        Set<ComponentKey> widgetsInWorkspace = dataModel.appWidgets.stream().map(
-                widget -> new ComponentKey(widget.providerName, widget.user)).collect(
-                Collectors.toSet());
+        Set<ComponentKey> widgetsInWorkspace = dataModel.itemsIdMap
+                .stream()
+                .filter(WIDGET_FILTER)
+                .map(item -> new ComponentKey(item.getTargetComponent(), item.user))
+                .collect(Collectors.toSet());
 
         // Widgets (excluding shortcuts & already added widgets) that belong to apps eligible for
         // being in predictions.
diff --git a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
index 289e720..2e48910 100644
--- a/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/NavbarButtonsViewController.java
@@ -40,9 +40,10 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_A11Y_BUTTON_CLICKABLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_A11Y_BUTTON_LONG_CLICKABLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BACK_DISABLED;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BACK_DISMISS_IME;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NAV_BAR_HIDDEN;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_OVERVIEW_DISABLED;
@@ -133,7 +134,14 @@
     private static final int FLAG_IME_SWITCHER_BUTTON_VISIBLE = 1 << 0;
     /** Whether the IME is visible. */
     private static final int FLAG_IME_VISIBLE = 1 << 1;
-    private static final int FLAG_ROTATION_BUTTON_VISIBLE = 1 << 2;
+    /**
+     * The back button is visually adjusted to indicate that it will dismiss the IME when pressed.
+     * This only takes effect while the IME is visible. By default, it is set while the IME is
+     * visible, but may be overridden by the
+     * {@link android.inputmethodservice.InputMethodService.BackDispositionMode backDispositionMode}
+     * set by the IME.
+     */
+    private static final int FLAG_BACK_DISMISS_IME = 1 << 2;
     private static final int FLAG_A11Y_VISIBLE = 1 << 3;
     private static final int FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE = 1 << 4;
     private static final int FLAG_KEYGUARD_VISIBLE = 1 << 5;
@@ -275,13 +283,14 @@
     }
 
     protected void setupController() {
-        boolean isThreeButtonNav = mContext.isThreeButtonNav();
+        final boolean isThreeButtonNav = mContext.isThreeButtonNav();
+        final boolean isPhoneMode = mContext.isPhoneMode();
         DeviceProfile deviceProfile = mContext.getDeviceProfile();
         Resources resources = mContext.getResources();
 
         int setupSize = mControllers.taskbarActivityContext.getSetupWindowSize();
-        Point p = DimensionUtils.getTaskbarPhoneDimensions(deviceProfile, resources,
-                mContext.isPhoneMode(), mContext.isGestureNav());
+        Point p = DimensionUtils.getTaskbarPhoneDimensions(deviceProfile, resources, isPhoneMode,
+                mContext.isGestureNav());
         ViewGroup.LayoutParams navButtonsViewLayoutParams = mNavButtonsView.getLayoutParams();
         navButtonsViewLayoutParams.width = p.x;
         if (!mContext.isUserSetupComplete()) {
@@ -302,9 +311,10 @@
             mImeSwitcherButton = addButton(switcherResId, BUTTON_IME_SWITCH,
                     isThreeButtonNav ? mStartContextualContainer : mEndContextualContainer,
                     mControllers.navButtonController, R.id.ime_switcher);
+            // A11y and IME Switcher buttons overlap on phone mode, show only a11y if both visible.
             mPropertyHolders.add(new StatePropertyHolder(mImeSwitcherButton,
-                    flags -> ((flags & FLAG_IME_SWITCHER_BUTTON_VISIBLE) != 0)
-                            && ((flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0)));
+                    flags -> (flags & FLAG_IME_SWITCHER_BUTTON_VISIBLE) != 0
+                            && !(isPhoneMode && (flags & FLAG_A11Y_VISIBLE) != 0)));
         }
 
         mPropertyHolders.add(new StatePropertyHolder(
@@ -318,7 +328,7 @@
                         .get(ALPHA_INDEX_SMALL_SCREEN),
                 flags -> (flags & FLAG_SMALL_SCREEN) == 0));
 
-        if (!mContext.isPhoneMode()) {
+        if (!isPhoneMode) {
             mPropertyHolders.add(new StatePropertyHolder(mControllers.taskbarDragLayerController
                     .getKeyguardBgTaskbar(), flags -> (flags & FLAG_KEYGUARD_VISIBLE) == 0));
         }
@@ -333,10 +343,10 @@
 
         // Make sure to remove nav bar buttons translation when any of the following occur:
         // - Notification shade is expanded
-        // - IME is showing (add separate translation for IME)
+        // - IME is visible (add separate translation for IME)
         // - VoiceInteractionWindow (assistant) is showing
         // - Keyboard shortcuts helper is showing
-        if (!mContext.isPhoneMode()) {
+        if (!isPhoneMode) {
             int flagsToRemoveTranslation = FLAG_NOTIFICATION_SHADE_EXPANDED | FLAG_IME_VISIBLE
                     | FLAG_VOICE_INTERACTION_WINDOW_SHOWING | FLAG_KEYBOARD_SHORTCUT_HELPER_SHOWING;
             mPropertyHolders.add(new StatePropertyHolder(mNavButtonInAppDisplayProgressForSysui,
@@ -364,9 +374,9 @@
             initButtons(mNavButtonContainer, mEndContextualContainer,
                     mControllers.navButtonController);
             updateButtonLayoutSpacing();
-            updateStateForFlag(FLAG_SMALL_SCREEN, mContext.isPhoneMode());
+            updateStateForFlag(FLAG_SMALL_SCREEN, isPhoneMode);
 
-            if (!mContext.isPhoneMode()) {
+            if (!isPhoneMode) {
                 mPropertyHolders.add(new StatePropertyHolder(
                         mControllers.taskbarDragLayerController.getNavbarBackgroundAlpha(),
                         flags -> (flags & FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE) != 0));
@@ -393,7 +403,7 @@
                 R.bool.floating_rotation_button_position_left);
         mControllers.rotationButtonController.setRotationButton(mFloatingRotationButton,
                 mRotationButtonListener);
-        if (mContext.isPhoneMode()) {
+        if (isPhoneMode) {
             mTaskbarTransitions.init();
         }
 
@@ -449,7 +459,7 @@
                     flags -> (flags & FLAG_IME_VISIBLE) == 0));
         }
         mPropertyHolders.add(new StatePropertyHolder(mBackButton,
-                flags -> (flags & FLAG_IME_VISIBLE) != 0,
+                flags -> (flags & FLAG_BACK_DISMISS_IME) != 0,
                 ROTATION_DRAWABLE_PERCENT, 1f, 0f));
         // Translate back button to be at end/start of other buttons for keyguard (only after SUW
         // since it is laid to align with SUW actions while in that state)
@@ -498,8 +508,7 @@
                 endContainer, navButtonController, R.id.accessibility_button,
                 R.layout.taskbar_contextual_button);
         mPropertyHolders.add(new StatePropertyHolder(mA11yButton,
-                flags -> (flags & FLAG_A11Y_VISIBLE) != 0
-                        && (flags & FLAG_ROTATION_BUTTON_VISIBLE) == 0));
+                flags -> (flags & FLAG_A11Y_VISIBLE) != 0));
 
         mSpace = new Space(mNavButtonsView.getContext());
         mSpace.setOnClickListener(view -> navButtonController.onButtonClick(BUTTON_SPACE, view));
@@ -510,8 +519,9 @@
     private void parseSystemUiFlags(@SystemUiStateFlags long sysUiStateFlags) {
         mSysuiStateFlags = sysUiStateFlags;
         boolean isImeSwitcherButtonVisible =
-                (sysUiStateFlags & SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING) != 0;
-        boolean isImeVisible = (sysUiStateFlags & SYSUI_STATE_IME_SHOWING) != 0;
+                (sysUiStateFlags & SYSUI_STATE_IME_SWITCHER_BUTTON_VISIBLE) != 0;
+        boolean isImeVisible = (sysUiStateFlags & SYSUI_STATE_IME_VISIBLE) != 0;
+        boolean isBackDismissIme = (sysUiStateFlags & SYSUI_STATE_BACK_DISMISS_IME) != 0;
         boolean a11yVisible = (sysUiStateFlags & SYSUI_STATE_A11Y_BUTTON_CLICKABLE) != 0;
         boolean isHomeDisabled = (sysUiStateFlags & SYSUI_STATE_HOME_DISABLED) != 0;
         boolean isRecentsDisabled = (sysUiStateFlags & SYSUI_STATE_OVERVIEW_DISABLED) != 0;
@@ -527,6 +537,7 @@
 
         updateStateForFlag(FLAG_IME_SWITCHER_BUTTON_VISIBLE, isImeSwitcherButtonVisible);
         updateStateForFlag(FLAG_IME_VISIBLE, isImeVisible);
+        updateStateForFlag(FLAG_BACK_DISMISS_IME, isBackDismissIme);
         updateStateForFlag(FLAG_A11Y_VISIBLE, a11yVisible);
         updateStateForFlag(FLAG_DISABLE_HOME, isHomeDisabled);
         updateStateForFlag(FLAG_DISABLE_RECENTS, isRecentsDisabled);
@@ -1231,7 +1242,7 @@
         appendFlag(str, flags, FLAG_IME_SWITCHER_BUTTON_VISIBLE,
                 "FLAG_IME_SWITCHER_BUTTON_VISIBLE");
         appendFlag(str, flags, FLAG_IME_VISIBLE, "FLAG_IME_VISIBLE");
-        appendFlag(str, flags, FLAG_ROTATION_BUTTON_VISIBLE, "FLAG_ROTATION_BUTTON_VISIBLE");
+        appendFlag(str, flags, FLAG_BACK_DISMISS_IME, "FLAG_BACK_DISMISS_IME");
         appendFlag(str, flags, FLAG_A11Y_VISIBLE, "FLAG_A11Y_VISIBLE");
         appendFlag(str, flags, FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE,
                 "FLAG_ONLY_BACK_FOR_BOUNCER_VISIBLE");
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index e8a0c45..a109a40 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -43,6 +43,7 @@
 import static com.android.quickstep.util.AnimUtils.completeRunnableListCallback;
 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;
+import static com.android.window.flags.Flags.enableStartLaunchTransitionFromTaskbarBugfix;
 import static com.android.wm.shell.Flags.enableTinyTaskbar;
 
 import static java.lang.invoke.MethodHandles.Lookup.PROTECTED;
@@ -1689,7 +1690,11 @@
         }
         // There is no task associated with this launch - launch a new task through an intent
         ActivityOptionsWrapper opts = getActivityLaunchDesktopOptions();
-        mSysUiProxy.startLaunchIntentTransition(intent, opts.options.toBundle(), displayId);
+        if (enableStartLaunchTransitionFromTaskbarBugfix()) {
+            mSysUiProxy.startLaunchIntentTransition(intent, opts.options.toBundle(), displayId);
+        } else {
+            startActivity(intent, opts.options.toBundle());
+        }
     }
 
     /** Expands a folder icon when it is clicked */
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index b46b0dc..ebd4ee5 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -723,6 +723,11 @@
     public void onDisplayRemoved(int displayId) {
     }
 
+    /**
+     * Signal from SysUI indicating that system decorations should be removed from the display.
+     */
+    public void onDisplayRemoveSystemDecorations(int displayId) {}
+
     private void removeActivityCallbacksAndListeners() {
         if (mActivity != null) {
             mActivity.removeOnDeviceProfileChangeListener(mDebugActivityDeviceProfileChanged);
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
index 68114f1..1ca3dfb 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarStashController.java
@@ -34,8 +34,7 @@
 import static com.android.quickstep.util.SystemUiFlagUtils.isTaskbarHidden;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DIALOG_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_SCREEN_PINNING;
 
@@ -260,9 +259,7 @@
     private @Nullable AnimatorSet mAnimator;
     private boolean mIsSystemGestureInProgress;
     /** Whether the IME is visible. */
-    private boolean mIsImeShowing;
-    /** Whether the IME Switcher button is visible. */
-    private boolean mIsImeSwitcherButtonShowing;
+    private boolean mIsImeVisible;
 
     private final Alarm mTimeoutAlarm = new Alarm();
     private boolean mEnableBlockingTimeoutDuringTests = false;
@@ -1122,7 +1119,7 @@
      */
     @VisibleForTesting
     long getTaskbarStashStartDelayForIme() {
-        if (mIsImeShowing) {
+        if (mIsImeVisible) {
             // Only delay when IME is exiting, not entering.
             return 0;
         }
@@ -1148,9 +1145,7 @@
         updateStateForFlag(FLAG_STASHED_DEVICE_LOCKED,
                 SystemUiFlagUtils.isLocked(systemUiStateFlags));
 
-        mIsImeShowing = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_SHOWING);
-        mIsImeSwitcherButtonShowing =
-                hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING);
+        mIsImeVisible = hasAnyFlag(systemUiStateFlags, SYSUI_STATE_IME_VISIBLE);
         if (updateStateForFlag(FLAG_STASHED_IME, shouldStashForIme())) {
             animDuration = TASKBAR_STASH_DURATION_FOR_IME;
             startDelay = getTaskbarStashStartDelayForIme();
@@ -1166,7 +1161,7 @@
     }
 
     /**
-     * We stash when IME or IME switcher is showing.
+     * We stash when the IME is visible.
      *
      * <p>Do not stash if in small screen, with 3 button nav, and in landscape (or seascape).
      * <p>Do not stash if taskbar is transient.
@@ -1202,7 +1197,7 @@
             return false;
         }
 
-        return mIsImeShowing || mIsImeSwitcherButtonShowing;
+        return mIsImeVisible;
     }
 
     /**
@@ -1377,8 +1372,7 @@
         pw.println(prefix + "\tappliedState=" + getStateString(mStatePropertyHolder.mPrevFlags));
         pw.println(prefix + "\tmState=" + getStateString(mState));
         pw.println(prefix + "\tmIsSystemGestureInProgress=" + mIsSystemGestureInProgress);
-        pw.println(prefix + "\tmIsImeShowing=" + mIsImeShowing);
-        pw.println(prefix + "\tmIsImeSwitcherButtonShowing=" + mIsImeSwitcherButtonShowing);
+        pw.println(prefix + "\tmIsImeVisible=" + mIsImeVisible);
     }
 
     private static String getStateString(long flags) {
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index e4e97e5..457ba3d 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -33,7 +33,6 @@
 import android.graphics.Canvas;
 import android.graphics.Rect;
 import android.graphics.drawable.Drawable;
-import android.os.Bundle;
 import android.util.ArraySet;
 import android.util.AttributeSet;
 import android.view.DisplayCutout;
@@ -41,7 +40,6 @@
 import android.view.LayoutInflater;
 import android.view.MotionEvent;
 import android.view.View;
-import android.view.accessibility.AccessibilityNodeInfo;
 import android.widget.FrameLayout;
 
 import androidx.annotation.LayoutRes;
@@ -311,16 +309,6 @@
         mShouldTryStartAlign = mActivityContext.shouldStartAlignTaskbar();
     }
 
-    @Override
-    public boolean performAccessibilityActionInternal(int action, Bundle arguments) {
-        if (action == AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS) {
-            announceTaskbarShown();
-        } else if (action == AccessibilityNodeInfo.ACTION_CLEAR_ACCESSIBILITY_FOCUS) {
-            announceTaskbarHidden();
-        }
-        return super.performAccessibilityActionInternal(action, arguments);
-    }
-
     private void announceTaskbarShown() {
         BubbleBarLocation bubbleBarLocation = mControllerCallbacks.getBubbleBarLocationIfVisible();
         if (bubbleBarLocation == null) {
@@ -334,21 +322,12 @@
         }
     }
 
-    private void announceTaskbarHidden() {
-        BubbleBarLocation bubbleBarLocation = mControllerCallbacks.getBubbleBarLocationIfVisible();
-        if (bubbleBarLocation == null) {
-            announceForAccessibility(mContext.getString(R.string.taskbar_a11y_hidden_title));
-        } else {
-            announceForAccessibility(
-                    mContext.getString(R.string.taskbar_a11y_hidden_with_bubbles_title));
-        }
-    }
-
     protected void announceAccessibilityChanges() {
-        this.performAccessibilityAction(
-                isVisibleToUser() ? AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS
-                        : AccessibilityNodeInfo.ACTION_CLEAR_ACCESSIBILITY_FOCUS, null);
-
+        // Only announce taskbar window shown. Window disappearing is generally not announce.
+        // This also aligns with talkback guidelines and unnecessary announcement to users.
+        if (isVisibleToUser()) {
+            announceTaskbarShown();
+        }
         ActivityContext.lookupContext(getContext()).getDragLayer()
                 .sendAccessibilityEvent(TYPE_WINDOW_CONTENT_CHANGED);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
index aff879e..4e029e3 100644
--- a/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/bubbles/BubbleBarController.java
@@ -19,8 +19,7 @@
 
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BOUNCER_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_QUICK_SETTINGS_EXPANDED;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING;
@@ -92,10 +91,9 @@
     private static final long MASK_HIDE_BUBBLE_BAR = SYSUI_STATE_BOUNCER_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING_OCCLUDED
-            | SYSUI_STATE_IME_SHOWING
+            | SYSUI_STATE_IME_VISIBLE
             | SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED
-            | SYSUI_STATE_QUICK_SETTINGS_EXPANDED
-            | SYSUI_STATE_IME_SWITCHER_BUTTON_SHOWING;
+            | SYSUI_STATE_QUICK_SETTINGS_EXPANDED;
 
     private static final long MASK_HIDE_HANDLE_VIEW = SYSUI_STATE_BOUNCER_SHOWING
             | SYSUI_STATE_STATUS_BAR_KEYGUARD_SHOWING
@@ -239,7 +237,7 @@
 
         boolean sysuiLocked = (flags & MASK_SYSUI_LOCKED) != 0;
         mBubbleStashController.setSysuiLocked(sysuiLocked);
-        mIsImeVisible = (flags & SYSUI_STATE_IME_SHOWING) != 0;
+        mIsImeVisible = (flags & SYSUI_STATE_IME_VISIBLE) != 0;
         if (mIsImeVisible) {
             mBubbleBarViewController.onImeVisible();
         }
diff --git a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
index d60dab6..914855b 100644
--- a/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
+++ b/quickstep/src/com/android/quickstep/DesktopSystemShortcut.kt
@@ -26,6 +26,7 @@
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.systemui.shared.system.InteractionJankMonitorWrapper
+import com.android.wm.shell.shared.desktopmode.DesktopModeCompatPolicy
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 
@@ -70,16 +71,31 @@
                     container: RecentsViewContainer,
                     taskContainer: TaskContainer,
                 ): List<DesktopSystemShortcut>? {
-                    return if (!DesktopModeStatus.canEnterDesktopMode(container.asContext())) null
-                    else if (!taskContainer.task.isDockable) null
-                    else
-                        listOf(
-                            DesktopSystemShortcut(
-                                container,
-                                taskContainer,
-                                abstractFloatingViewHelper,
+                    val context = container.asContext()
+                    val taskKey = taskContainer.task.key
+                    val desktopModeCompatPolicy = DesktopModeCompatPolicy(context)
+                    return when {
+                        !DesktopModeStatus.canEnterDesktopMode(context) -> null
+
+                        desktopModeCompatPolicy.isTopActivityExemptFromDesktopWindowing(
+                            taskKey.baseActivity?.packageName,
+                            taskKey.numActivities,
+                            taskKey.isTopActivityNoDisplay,
+                            taskKey.isActivityStackTransparent,
+                        ) -> null
+
+                        !taskContainer.task.isDockable -> null
+
+                        else -> {
+                            listOf(
+                                DesktopSystemShortcut(
+                                    container,
+                                    taskContainer,
+                                    abstractFloatingViewHelper,
+                                )
                             )
-                        )
+                        }
+                    }
                 }
 
                 override fun showForGroupedTask() = true
diff --git a/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt b/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
index 46c4f36..f97cf9c 100644
--- a/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
+++ b/quickstep/src/com/android/quickstep/ExternalDisplaySystemShortcut.kt
@@ -25,6 +25,7 @@
 import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.window.flags.Flags
+import com.android.wm.shell.shared.desktopmode.DesktopModeCompatPolicy
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 
 /** A menu item that allows the user to move the current app into external display. */
@@ -66,18 +67,31 @@
                     container: RecentsViewContainer,
                     taskContainer: TaskContainer,
                 ): List<ExternalDisplaySystemShortcut>? {
-                    return if (
-                        DesktopModeStatus.canEnterDesktopMode(container.asContext()) &&
-                            Flags.moveToExternalDisplayShortcut()
-                    )
-                        listOf(
-                            ExternalDisplaySystemShortcut(
-                                container,
-                                abstractFloatingViewHelper,
-                                taskContainer,
+                    val context = container.asContext()
+                    val taskKey = taskContainer.task.key
+                    val desktopModeCompatPolicy = DesktopModeCompatPolicy(context)
+                    return when {
+                        !DesktopModeStatus.canEnterDesktopMode(context) -> null
+
+                        !Flags.moveToExternalDisplayShortcut() -> null
+
+                        desktopModeCompatPolicy.isTopActivityExemptFromDesktopWindowing(
+                            taskKey.baseActivity?.packageName,
+                            taskKey.numActivities,
+                            taskKey.isTopActivityNoDisplay,
+                            taskKey.isActivityStackTransparent,
+                        ) -> null
+
+                        else -> {
+                            listOf(
+                                ExternalDisplaySystemShortcut(
+                                    container,
+                                    abstractFloatingViewHelper,
+                                    taskContainer,
+                                )
                             )
-                        )
-                    else null
+                        }
+                    }
                 }
 
                 override fun showForGroupedTask() = true
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
index 66f307c..6ad9a2c 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
@@ -28,7 +28,9 @@
 import androidx.annotation.UiThread
 import androidx.annotation.VisibleForTesting
 import com.android.internal.jank.Cuj
+import com.android.launcher3.Flags.enableFallbackOverviewInWindow
 import com.android.launcher3.Flags.enableLargeDesktopWindowingTile
+import com.android.launcher3.Flags.enableLauncherOverviewInWindow
 import com.android.launcher3.Flags.enableOverviewCommandHelperTimeout
 import com.android.launcher3.PagedView
 import com.android.launcher3.logger.LauncherAtom
@@ -344,9 +346,12 @@
             return false
         }
 
-        val activity = containerInterface.getCreatedContainer()
-        if (activity != null) {
-            InteractionJankMonitorWrapper.begin(activity.rootView, Cuj.CUJ_LAUNCHER_QUICK_SWITCH)
+        val recentsInWindowFlagSet =
+            enableFallbackOverviewInWindow() || enableLauncherOverviewInWindow()
+        if (!recentsInWindowFlagSet) {
+            containerInterface.getCreatedContainer()?.rootView?.let { view ->
+                InteractionJankMonitorWrapper.begin(view, Cuj.CUJ_LAUNCHER_QUICK_SWITCH)
+            }
         }
 
         val gestureState =
@@ -373,6 +378,12 @@
                     transitionInfo: TransitionInfo,
                 ) {
                     Log.d(TAG, "recents animation started: $command")
+                    if (recentsInWindowFlagSet) {
+                        containerInterface.getCreatedContainer()?.rootView?.let { view ->
+                            InteractionJankMonitorWrapper.begin(view, Cuj.CUJ_LAUNCHER_QUICK_SWITCH)
+                        }
+                    }
+
                     updateRecentsViewFocus(command)
                     logShowOverviewFrom(command.type)
                     containerInterface.runOnInitBackgroundStateUI {
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
index 0c89a80..e1e962a 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationDeviceState.java
@@ -37,7 +37,7 @@
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DISABLE_GESTURE_PIP_ANIMATING;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_DISABLE_GESTURE_SPLIT_INVOCATION;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_HOME_DISABLED;
-import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING;
+import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_MAGNIFICATION_OVERLAP;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NAV_BAR_HIDDEN;
 import static com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_NOTIFICATION_PANEL_EXPANDED;
@@ -586,7 +586,7 @@
     /** Returns whether IME is rendering nav buttons, and IME is currently showing. */
     public boolean isImeRenderingNavButtons() {
         return mCanImeRenderGesturalNavButtons && mMode == NO_BUTTON
-                && ((mSystemUiStateFlags & SYSUI_STATE_IME_SHOWING) != 0);
+                && ((mSystemUiStateFlags & SYSUI_STATE_IME_VISIBLE) != 0);
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.kt b/quickstep/src/com/android/quickstep/SystemUiProxy.kt
index 87953c7..d5382ad 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.kt
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.kt
@@ -126,7 +126,7 @@
     private val systemUiProxyDeathRecipient =
         IBinder.DeathRecipient { Executors.MAIN_EXECUTOR.execute { clearProxy() } }
 
-    // Save the listeners passed into the proxy since OverviewProxyService may not have been bound
+    // Save the listeners passed into the proxy since LauncherProxyService may not have been bound
     // yet, and we'll need to set/register these listeners with SysUI when they do.  Note that it is
     // up to the caller to clear the listeners to prevent leaks as these can be held indefinitely
     // in case SysUI needs to rebind.
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 15f320d..516b24c 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -100,7 +100,7 @@
 import com.android.quickstep.util.ContextualSearchInvoker;
 import com.android.quickstep.util.ContextualSearchStateManager;
 import com.android.quickstep.views.RecentsViewContainer;
-import com.android.systemui.shared.recents.IOverviewProxy;
+import com.android.systemui.shared.recents.ILauncherProxy;
 import com.android.systemui.shared.recents.ISystemUiProxy;
 import com.android.systemui.shared.statusbar.phone.BarTransitions;
 import com.android.systemui.shared.system.InputChannelCompat.InputEventReceiver;
@@ -141,9 +141,9 @@
     private final TISBinder mTISBinder = new TISBinder(this);
 
     /**
-     * Local IOverviewProxy implementation with some methods for local components
+     * Local ILauncherProxy implementation with some methods for local components
      */
-    public static class TISBinder extends IOverviewProxy.Stub {
+    public static class TISBinder extends ILauncherProxy.Stub {
 
         private final WeakReference<TouchInteractionService> mTis;
 
@@ -317,6 +317,17 @@
 
         @BinderThread
         @Override
+        public void onDisplayRemoveSystemDecorations(int displayId) {
+            // TODO(b/391786915): Replace all
+            // `executeForTouchInteractionService(executeForTaskbarManager())` with just
+            // `executeForTaskbarManager` directly (since `tis` is unused).
+            MAIN_EXECUTOR.execute(() -> executeForTouchInteractionService(
+                    tis -> executeForTaskbarManager(taskbarManager -> taskbarManager
+                            .onDisplayRemoveSystemDecorations(displayId))));
+        }
+
+        @BinderThread
+        @Override
         public void updateWallpaperVisibility(int displayId, boolean visible) {
             MAIN_EXECUTOR.execute(() -> executeForTouchInteractionService(
                     tis -> executeForTaskbarManager(
@@ -412,7 +423,7 @@
                 try {
                     reply.sendResult(null);
                 } catch (RemoteException e) {
-                    Log.w(TAG, "onUnbind: Failed to reply to OverviewProxyService", e);
+                    Log.w(TAG, "onUnbind: Failed to reply to LauncherProxyService", e);
                 }
             });
         }
diff --git a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
index 9bfe71f..b2a30ca 100644
--- a/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
+++ b/quickstep/src/com/android/quickstep/logging/StatsLogCompatManager.java
@@ -31,15 +31,15 @@
 import static com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers.ContainerCase.DEVICE_SEARCH_RESULT_CONTAINER;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_WORKSPACE_SNAPSHOT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_0;
-import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_90;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_180;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_270;
+import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DISPLAY_ROTATION__ROTATION_90;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__ALLAPPS;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__BACKGROUND;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__HOME;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__DST_STATE__OVERVIEW;
-import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__PORTRAIT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__LANDSCAPE;
+import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__PORTRAIT;
 import static com.android.systemui.shared.system.SysUiStatsLog.LAUNCHER_UICHANGED__RECENTS_ORIENTATION_HANDLER__SEASCAPE;
 
 import android.content.Context;
@@ -69,6 +69,7 @@
 import com.android.launcher3.logger.LauncherAtomExtensions.ExtendedContainers;
 import com.android.launcher3.logging.InstanceId;
 import com.android.launcher3.logging.StatsLogManager;
+import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.util.DisplayController;
 import com.android.launcher3.util.Executors;
@@ -386,7 +387,8 @@
                 // and then write to StatsLog.
                 app.getModel().enqueueModelUpdateTask((taskController, dataModel, apps) ->
                         write(event, applyOverwrites(mItemInfo.buildProto(
-                                dataModel.collections.get(mItemInfo.container), mContext))));
+                                (CollectionInfo) dataModel.itemsIdMap.get(mItemInfo.container),
+                                mContext))));
             })) {
                 // Write log on the model thread so that logs do not go out of order
                 // (for eg: drop comes after drag)
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 2f95413..002a4e8 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -17,7 +17,6 @@
 package com.android.quickstep.recents.data
 
 import android.graphics.drawable.Drawable
-import android.graphics.drawable.ShapeDrawable
 import android.util.Log
 import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.data.TaskVisualsChangedDelegate.TaskIconChangedCallback
@@ -52,37 +51,29 @@
     override fun getAllTaskData(forceRefresh: Boolean): Flow<List<Task>> {
         if (forceRefresh) {
             recentsModel.getTasks { newTaskList ->
-                val oldTaskMap = tasks.value
                 val recentTasks =
                     newTaskList
                         .flatMap { groupTask -> groupTask.tasks }
                         .associateBy { it.key.id }
                         .also { newTaskMap ->
                             // Clean tasks that are not in the latest group tasks list.
-                            val tasksNoLongerVisible = oldTaskMap.keys.subtract(newTaskMap.keys)
+                            val tasksNoLongerVisible = tasks.value.keys.subtract(newTaskMap.keys)
                             removeTasks(tasksNoLongerVisible)
-
-                            // Use pre-loaded thumbnail data and icon from the previous list.
-                            // This reduces the Thumbnail loading time in the Overview and prevent
-                            // empty thumbnail and icon.
-                            val cache =
-                                taskRequests.keys
-                                    .mapNotNull { key ->
-                                        val task = oldTaskMap[key] ?: return@mapNotNull null
-                                        key to Pair(task.thumbnail, task.icon)
-                                    }
-                                    .toMap()
-
-                            newTaskMap.values.forEach { task ->
-                                task.thumbnail = task.thumbnail ?: cache[task.key.id]?.first
-                                task.icon = task.icon ?: cache[task.key.id]?.second
-                            }
                         }
-                tasks.value = MapForStateFlow(recentTasks)
                 Log.d(
                     TAG,
-                    "getAllTaskData: oldTasks ${oldTaskMap.keys}, newTasks: ${recentTasks.keys}",
+                    "getAllTaskData: oldTasks ${tasks.value.keys}, newTasks: ${recentTasks.keys}",
                 )
+                tasks.value = MapForStateFlow(recentTasks)
+
+                // Request data for completed tasks to prevent stale data.
+                // This will prevent thumbnail and icon from being replaced and
+                // null due to race condition.
+                taskRequests.values.forEach { (taskKey, job) ->
+                    if (job.isCompleted) {
+                        requestTaskData(taskKey.id)
+                    }
+                }
             }
         }
         return tasks.map { it.values.toList() }
@@ -202,13 +193,11 @@
     private suspend fun getIconFromDataSource(task: Task) =
         withContext(dispatcherProvider.background) {
             val iconCacheEntry = taskIconDataSource.getIcon(task)
-            val icon = iconCacheEntry.icon.constantState?.newDrawable()?.mutate() ?: EMPTY_DRAWABLE
-            IconData(icon, iconCacheEntry.contentDescription, iconCacheEntry.title)
+            IconData(iconCacheEntry.icon, iconCacheEntry.contentDescription, iconCacheEntry.title)
         }
 
     companion object {
         private const val TAG = "TasksRepository"
-        private val EMPTY_DRAWABLE = ShapeDrawable()
     }
 
     /** Helper class to support StateFlow emissions when using a Map with a MutableStateFlow. */
diff --git a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
index 358537c..0a47338 100644
--- a/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
+++ b/quickstep/src/com/android/quickstep/recents/di/RecentsDependencies.kt
@@ -26,6 +26,7 @@
 import com.android.quickstep.recents.data.TaskVisualsChangedDelegate
 import com.android.quickstep.recents.data.TaskVisualsChangedDelegateImpl
 import com.android.quickstep.recents.data.TasksRepository
+import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
 import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.GetThumbnailUseCase
 import com.android.quickstep.recents.usecase.SysUiStatusNavFlagsUseCase
@@ -181,8 +182,6 @@
                         taskContainerData = inject(scopeId),
                         dispatcherProvider = inject(),
                         getThumbnailPositionUseCase = inject(),
-                        tasksRepository = inject(),
-                        deviceProfileRepository = inject(),
                         splashAlphaUseCase = inject(scopeId),
                     )
                 TaskOverlayViewModel::class.java -> {
@@ -195,6 +194,7 @@
                         dispatcherProvider = inject(),
                     )
                 }
+                GetTaskUseCase::class.java -> GetTaskUseCase(repository = inject())
                 GetThumbnailUseCase::class.java -> GetThumbnailUseCase(taskRepository = inject())
                 SysUiStatusNavFlagsUseCase::class.java ->
                     SysUiStatusNavFlagsUseCase(taskRepository = inject())
diff --git a/quickstep/src/com/android/quickstep/recents/domain/model/TaskModel.kt b/quickstep/src/com/android/quickstep/recents/domain/model/TaskModel.kt
index 3823100..bf29b1d 100644
--- a/quickstep/src/com/android/quickstep/recents/domain/model/TaskModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/domain/model/TaskModel.kt
@@ -38,7 +38,7 @@
  */
 data class TaskModel(
     val id: TaskId,
-    val title: String,
+    val title: String?,
     val titleDescription: String?,
     val icon: Drawable?,
     val thumbnail: ThumbnailData?,
diff --git a/quickstep/src/com/android/quickstep/recents/ui/mapper/TaskUiStateMapper.kt b/quickstep/src/com/android/quickstep/recents/ui/mapper/TaskUiStateMapper.kt
new file mode 100644
index 0000000..fb62268
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/recents/ui/mapper/TaskUiStateMapper.kt
@@ -0,0 +1,86 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.ui.mapper
+
+import com.android.quickstep.recents.ui.viewmodel.TaskData
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.ThumbnailHeader
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
+
+object TaskUiStateMapper {
+
+    /**
+     * Converts a [TaskData] object into a [TaskThumbnailUiState] for display in the UI.
+     *
+     * This function handles different types of [TaskData] and determines the appropriate UI state
+     * based on the data and provided flags.
+     *
+     * @param taskData The [TaskData] to convert. Can be null or a specific subclass.
+     * @param isLiveTile A flag indicating whether the task data represents live tile.
+     * @param hasHeader A flag indicating whether the UI should display a header.
+     * @return A [TaskThumbnailUiState] representing the UI state for the given task data.
+     */
+    fun toTaskThumbnailUiState(
+        taskData: TaskData?,
+        isLiveTile: Boolean,
+        hasHeader: Boolean,
+    ): TaskThumbnailUiState =
+        when {
+            taskData !is TaskData.Data -> Uninitialized
+            isLiveTile -> createLiveTileState(taskData, hasHeader)
+            isBackgroundOnly(taskData) -> BackgroundOnly(taskData.backgroundColor)
+            isSnapshotSplash(taskData) ->
+                SnapshotSplash(createSnapshotState(taskData, hasHeader), taskData.icon)
+            else -> Uninitialized
+        }
+
+    private fun createSnapshotState(taskData: TaskData.Data, hasHeader: Boolean): Snapshot =
+        if (canHeaderBeCreated(taskData, hasHeader)) {
+            Snapshot.WithHeader(
+                taskData.thumbnailData?.thumbnail!!,
+                taskData.thumbnailData.rotation,
+                taskData.backgroundColor,
+                ThumbnailHeader(taskData.icon!!, taskData.titleDescription!!),
+            )
+        } else {
+            Snapshot.WithoutHeader(
+                taskData.thumbnailData?.thumbnail!!,
+                taskData.thumbnailData.rotation,
+                taskData.backgroundColor,
+            )
+        }
+
+    private fun isBackgroundOnly(taskData: TaskData.Data) =
+        taskData.isLocked || taskData.thumbnailData == null
+
+    private fun isSnapshotSplash(taskData: TaskData.Data) =
+        taskData.thumbnailData?.thumbnail != null && !taskData.isLocked
+
+    private fun canHeaderBeCreated(taskData: TaskData.Data, hasHeader: Boolean) =
+        hasHeader && taskData.icon != null && taskData.titleDescription != null
+
+    private fun createLiveTileState(taskData: TaskData.Data, hasHeader: Boolean) =
+        if (canHeaderBeCreated(taskData, hasHeader)) {
+            // TODO(http://b/353965691): figure out what to do when `icon` or `titleDescription` is
+            //  null.
+            LiveTile.WithHeader(ThumbnailHeader(taskData.icon!!, taskData.titleDescription!!))
+        } else LiveTile.WithoutHeader
+}
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
index 5f98479..54b2389 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskTileUiState.kt
@@ -30,28 +30,36 @@
  * @property isLiveTile Indicates whether this data is intended for a live tile. If `true`, the
  *   running app will be displayed instead of the thumbnail.
  */
-data class TaskTileUiState(val tasks: List<TaskData>, val isLiveTile: Boolean)
+data class TaskTileUiState(
+    val tasks: List<TaskData>,
+    val isLiveTile: Boolean,
+    val hasHeader: Boolean,
+)
 
-sealed interface TaskData {
+sealed class TaskData {
+    abstract val taskId: Int
+
     /** When no data was found for the TaskId provided */
-    data class NoData(val taskId: Int) : TaskData
+    data class NoData(override val taskId: Int) : TaskData()
 
     /**
      * This class provides UI information related to a Task (App) to be displayed within a TaskView.
      *
      * @property taskId Identifier of the task
      * @property title App title
+     * @property titleDescription App content description
      * @property icon App icon
      * @property thumbnailData Information related to the last snapshot retrieved from the app
      * @property backgroundColor The background color of the task.
      * @property isLocked Indicates whether the task is locked or not.
      */
     data class Data(
-        val taskId: Int,
-        val title: String,
+        override val taskId: Int,
+        val title: String?,
+        val titleDescription: String?,
         val icon: Drawable?,
         val thumbnailData: ThumbnailData?,
         val backgroundColor: Int,
         val isLocked: Boolean,
-    ) : TaskData
+    ) : TaskData()
 }
diff --git a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
index 2e51a8a..b2806f0 100644
--- a/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
+++ b/quickstep/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModel.kt
@@ -16,12 +16,15 @@
 
 package com.android.quickstep.recents.ui.viewmodel
 
+import android.annotation.ColorInt
 import android.util.Log
+import androidx.core.graphics.ColorUtils
 import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.quickstep.recents.domain.model.TaskId
 import com.android.quickstep.recents.domain.model.TaskModel
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.views.TaskViewType
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -37,6 +40,7 @@
  */
 @OptIn(ExperimentalCoroutinesApi::class)
 class TaskViewModel(
+    private val taskViewType: TaskViewType,
     recentsViewData: RecentsViewData,
     private val getTaskUseCase: GetTaskUseCase,
     dispatcherProvider: DispatcherProvider,
@@ -62,7 +66,14 @@
                     ::mapToUiState,
                 )
             }
-            .combine(isLiveTile) { tasks, isLiveTile -> TaskTileUiState(tasks, isLiveTile) }
+            .combine(isLiveTile) { tasks, isLiveTile ->
+                TaskTileUiState(
+                    tasks = tasks,
+                    isLiveTile = isLiveTile,
+                    hasHeader = taskViewType == TaskViewType.DESKTOP,
+                )
+            }
+            .distinctUntilChanged()
             .flowOn(dispatcherProvider.background)
 
     fun bind(vararg taskId: TaskId) {
@@ -78,13 +89,16 @@
             TaskData.Data(
                 taskId = taskId,
                 title = result.title,
+                titleDescription = result.titleDescription,
                 icon = result.icon,
                 thumbnailData = result.thumbnail,
-                backgroundColor = result.backgroundColor,
+                backgroundColor = result.backgroundColor.removeAlpha(),
                 isLocked = result.isLocked,
             )
         } ?: TaskData.NoData(taskId)
 
+    @ColorInt private fun Int.removeAlpha(): Int = ColorUtils.setAlphaComponent(this, 0xff)
+
     private companion object {
         const val TAG = "TaskViewModel"
     }
diff --git a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
index 02baa39..4a990b3 100644
--- a/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
+++ b/quickstep/src/com/android/quickstep/task/thumbnail/TaskThumbnailView.kt
@@ -20,6 +20,7 @@
 import android.graphics.Color
 import android.graphics.Outline
 import android.graphics.Rect
+import android.graphics.drawable.ShapeDrawable
 import android.util.AttributeSet
 import android.util.Log
 import android.view.LayoutInflater
@@ -108,21 +109,6 @@
         viewData = RecentsDependencies.get(this)
         updateViewDataValues()
         viewModel = RecentsDependencies.get(this)
-        viewModel.uiState
-            .dropWhile { it == Uninitialized }
-            .flowOn(dispatcherProvider.background)
-            .onEach { viewModelUiState ->
-                Log.d(TAG, "viewModelUiState changed from: $uiState to: $viewModelUiState")
-                uiState = viewModelUiState
-                resetViews()
-                when (viewModelUiState) {
-                    is Uninitialized -> {}
-                    is LiveTile -> drawLiveWindow(viewModelUiState)
-                    is SnapshotSplash -> drawSnapshotSplash(viewModelUiState)
-                    is BackgroundOnly -> drawBackground(viewModelUiState.backgroundColor)
-                }
-            }
-            .launchIn(viewAttachedScope)
         viewModel.dimProgress
             .dropWhile { it == 0f }
             .flowOn(dispatcherProvider.background)
@@ -166,6 +152,19 @@
         }
     }
 
+    fun setState(state: TaskThumbnailUiState) {
+        Log.d(TAG, "viewModelUiState changed from: $uiState to: $state")
+        if (uiState == state) return
+        uiState = state
+        resetViews()
+        when (state) {
+            is Uninitialized -> {}
+            is LiveTile -> drawLiveWindow(state)
+            is SnapshotSplash -> drawSnapshotSplash(state)
+            is BackgroundOnly -> drawBackground(state.backgroundColor)
+        }
+    }
+
     private fun updateViewDataValues() {
         viewData.width.value = width
         viewData.height.value = height
@@ -219,7 +218,8 @@
         drawSnapshot(snapshotSplash.snapshot)
 
         splashBackground.setBackgroundColor(snapshotSplash.snapshot.backgroundColor)
-        splashIcon.setImageDrawable(snapshotSplash.splash)
+        val icon = snapshotSplash.splash?.constantState?.newDrawable()?.mutate() ?: ShapeDrawable()
+        splashIcon.setImageDrawable(icon)
     }
 
     private fun drawSnapshot(snapshot: Snapshot) {
@@ -238,10 +238,6 @@
         thumbnailView.imageMatrix = viewModel.getThumbnailPositionState(width, height, isLayoutRtl)
     }
 
-    private companion object {
-        const val TAG = "TaskThumbnailView"
-    }
-
     private fun maybeCreateHeader() {
         if (enableDesktopExplodedView() && taskThumbnailViewHeader == null) {
             taskThumbnailViewHeader =
@@ -251,4 +247,8 @@
             addView(taskThumbnailViewHeader)
         }
     }
+
+    private companion object {
+        const val TAG = "TaskThumbnailView"
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
index a048a1d..a9fdaa5 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModel.kt
@@ -17,7 +17,6 @@
 package com.android.quickstep.task.viewmodel
 
 import android.graphics.Matrix
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
 import kotlinx.coroutines.flow.Flow
 
 /** ViewModel for representing TaskThumbnails */
@@ -28,9 +27,6 @@
     /** Provides the alpha of the splash icon */
     val splashAlpha: Flow<Float>
 
-    /** Provides the UiState by which the task thumbnail can be represented */
-    val uiState: Flow<TaskThumbnailUiState>
-
     /** Attaches this ViewModel to a specific task id for it to provide data from. */
     fun bind(taskId: Int)
 
diff --git a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
index a154c3c..4e4e225 100644
--- a/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
+++ b/quickstep/src/com/android/quickstep/task/viewmodel/TaskThumbnailViewModelImpl.kt
@@ -16,50 +16,31 @@
 
 package com.android.quickstep.task.viewmodel
 
-import android.annotation.ColorInt
 import android.app.ActivityTaskManager.INVALID_TASK_ID
-import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
 import android.graphics.Matrix
 import android.util.Log
-import androidx.core.graphics.ColorUtils
-import com.android.launcher3.Flags.enableDesktopExplodedView
 import com.android.launcher3.util.coroutines.DispatcherProvider
-import com.android.quickstep.recents.data.RecentTasksRepository
-import com.android.quickstep.recents.data.RecentsDeviceProfileRepository
 import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.ThumbnailPositionState
 import com.android.quickstep.recents.viewmodel.RecentsViewData
 import com.android.quickstep.task.thumbnail.SplashAlphaUseCase
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.ThumbnailHeader
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
-import com.android.systemui.shared.recents.model.Task
 import kotlin.math.max
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.flowOn
-import kotlinx.coroutines.flow.map
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class TaskThumbnailViewModelImpl(
     recentsViewData: RecentsViewData,
     taskContainerData: TaskContainerData,
     dispatcherProvider: DispatcherProvider,
-    private val tasksRepository: RecentTasksRepository,
-    private val deviceProfileRepository: RecentsDeviceProfileRepository,
     private val getThumbnailPositionUseCase: GetThumbnailPositionUseCase,
     private val splashAlphaUseCase: SplashAlphaUseCase,
 ) : TaskThumbnailViewModel {
-    private val task = MutableStateFlow<Flow<Task?>>(flowOf(null))
     private val splashProgress = MutableStateFlow(flowOf(0f))
     private var taskId: Int = INVALID_TASK_ID
 
@@ -74,42 +55,9 @@
     override val splashAlpha =
         splashProgress.flatMapLatest { it }.flowOn(dispatcherProvider.background)
 
-    private val isLiveTile =
-        combine(
-                task.flatMapLatest { it }.map { it?.key?.id }.distinctUntilChanged(),
-                recentsViewData.runningTaskIds,
-                recentsViewData.runningTaskShowScreenshot,
-            ) { taskId, runningTaskIds, runningTaskShowScreenshot ->
-                runningTaskIds.contains(taskId) && !runningTaskShowScreenshot
-            }
-            .distinctUntilChanged()
-
-    override val uiState: Flow<TaskThumbnailUiState> =
-        combine(task.flatMapLatest { it }, isLiveTile) { taskVal, isRunning ->
-                // TODO(b/369339561) This log is firing a lot. Reduce emissions from TasksRepository
-                //  then re-enable this log.
-                //                Log.d(
-                //                    TAG,
-                //                    "Received task and / or live tile update. taskVal: $taskVal"
-                //                    + " isRunning: $isRunning.",
-                //                )
-                when {
-                    taskVal == null -> Uninitialized
-                    isRunning -> createLiveTileState(taskVal)
-                    isBackgroundOnly(taskVal) ->
-                        BackgroundOnly(taskVal.colorBackground.removeAlpha())
-                    isSnapshotSplashState(taskVal) ->
-                        SnapshotSplash(createSnapshotState(taskVal), taskVal.icon)
-                    else -> Uninitialized
-                }
-            }
-            .distinctUntilChanged()
-            .flowOn(dispatcherProvider.background)
-
     override fun bind(taskId: Int) {
         Log.d(TAG, "bind taskId: $taskId")
         this.taskId = taskId
-        task.value = tasksRepository.getTaskDataById(taskId)
         splashProgress.value = splashAlphaUseCase.execute(taskId)
     }
 
@@ -122,62 +70,6 @@
             is ThumbnailPositionState.MissingThumbnail -> Matrix.IDENTITY_MATRIX
         }
 
-    private fun isBackgroundOnly(task: Task): Boolean = task.isLocked || task.thumbnail == null
-
-    private fun isSnapshotSplashState(task: Task): Boolean {
-        val thumbnailPresent = task.thumbnail?.thumbnail != null
-        val taskLocked = task.isLocked
-
-        return thumbnailPresent && !taskLocked
-    }
-
-    private fun createSnapshotState(task: Task): Snapshot {
-        val thumbnailData = task.thumbnail
-        val bitmap = thumbnailData?.thumbnail!!
-        var thumbnailHeader = maybeCreateHeader(task)
-        return if (thumbnailHeader != null)
-            Snapshot.WithHeader(
-                bitmap,
-                thumbnailData.rotation,
-                task.colorBackground.removeAlpha(),
-                thumbnailHeader,
-            )
-        else
-            Snapshot.WithoutHeader(
-                bitmap,
-                thumbnailData.rotation,
-                task.colorBackground.removeAlpha(),
-            )
-    }
-
-    private fun shouldHaveThumbnailHeader(task: Task): Boolean {
-        return deviceProfileRepository.getRecentsDeviceProfile().canEnterDesktopMode &&
-            enableDesktopExplodedView() &&
-            task.key.windowingMode == WINDOWING_MODE_FREEFORM
-    }
-
-    private fun maybeCreateHeader(task: Task): ThumbnailHeader? {
-        // Header is only needed when this task is a desktop task and Overivew exploded view is
-        // enabled.
-        if (!shouldHaveThumbnailHeader(task)) {
-            return null
-        }
-
-        // TODO(http://b/353965691): figure out what to do when `icon` or `titleDescription` is
-        // null.
-        val icon = task.icon ?: return null
-        val titleDescription = task.titleDescription ?: return null
-        return ThumbnailHeader(icon, titleDescription)
-    }
-
-    private fun createLiveTileState(task: Task): LiveTile {
-        val thumbnailHeader = maybeCreateHeader(task)
-        return if (thumbnailHeader != null) LiveTile.WithHeader(thumbnailHeader)
-        else LiveTile.WithoutHeader
-    }
-
-    @ColorInt private fun Int.removeAlpha(): Int = ColorUtils.setAlphaComponent(this, 0xff)
-
     private companion object {
         const val MAX_SCRIM_ALPHA = 0.4f
         const val TAG = "TaskThumbnailViewModel"
diff --git a/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt b/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt
index 1dab18a..e353160 100644
--- a/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt
+++ b/quickstep/src/com/android/quickstep/views/AddDesktopButton.kt
@@ -21,7 +21,9 @@
 import android.graphics.drawable.shapes.RoundRectShape
 import android.util.AttributeSet
 import android.widget.ImageButton
+import com.android.launcher3.LauncherAnimUtils.VIEW_TRANSLATE_X
 import com.android.launcher3.R
+import com.android.launcher3.util.MultiPropertyFactory
 
 /**
  * Button for supporting multiple desktop sessions. The button will be next to the first TaskView
@@ -30,6 +32,29 @@
 class AddDesktopButton @JvmOverloads constructor(context: Context, attrs: AttributeSet? = null) :
     ImageButton(context, attrs) {
 
+    private enum class TranslationX {
+        GRID,
+        OFFSET,
+    }
+
+    private val multiTranslationX =
+        MultiPropertyFactory(this, VIEW_TRANSLATE_X, TranslationX.entries.size) { a: Float, b: Float
+            ->
+            a + b
+        }
+
+    var gridTranslationX
+        get() = multiTranslationX[TranslationX.GRID.ordinal].value
+        set(value) {
+            multiTranslationX[TranslationX.GRID.ordinal].value = value
+        }
+
+    var offsetTranslationX
+        get() = multiTranslationX[TranslationX.OFFSET.ordinal].value
+        set(value) {
+            multiTranslationX[TranslationX.OFFSET.ordinal].value = value
+        }
+
     override fun onFinishInflate() {
         super.onFinishInflate()
 
diff --git a/quickstep/src/com/android/quickstep/views/IconView.kt b/quickstep/src/com/android/quickstep/views/IconView.kt
index 2e6c4bf..6da52d6 100644
--- a/quickstep/src/com/android/quickstep/views/IconView.kt
+++ b/quickstep/src/com/android/quickstep/views/IconView.kt
@@ -45,11 +45,11 @@
     private var msdlPlayerWrapper: MSDLPlayerWrapper? = null
 
     constructor(context: Context) : super(context) {
-        msdlPlayerWrapper = MSDLPlayerWrapper.INSTANCE.get(context)
+        setUpHaptics()
     }
 
     constructor(context: Context, attrs: AttributeSet?) : super(context, attrs) {
-        msdlPlayerWrapper = MSDLPlayerWrapper.INSTANCE.get(context)
+        setUpHaptics()
     }
 
     constructor(
@@ -57,11 +57,15 @@
         attrs: AttributeSet?,
         defStyleAttr: Int,
     ) : super(context, attrs, defStyleAttr) {
-        msdlPlayerWrapper = MSDLPlayerWrapper.INSTANCE.get(context)
+        setUpHaptics()
     }
 
     init {
         multiValueAlpha.setUpdateVisibility(true)
+    }
+
+    private fun setUpHaptics() {
+        msdlPlayerWrapper = MSDLPlayerWrapper.INSTANCE.get(context)
         // Haptics are handled by the MSDLPlayerWrapper
         isHapticFeedbackEnabled = !Flags.msdlFeedback() || msdlPlayerWrapper == null
     }
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index c0b026b..b93a2f0 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -3467,7 +3467,7 @@
                 // `mAddDesktopButton`, shift `mAddDesktopButton` to accommodate.
                 translationX += largeTaskWidthAndSpacing;
             }
-            mAddDesktopButton.setTranslationX(translationX);
+            mAddDesktopButton.setGridTranslationX(translationX);
         }
 
         final TaskView runningTask = getRunningTaskView();
@@ -4973,8 +4973,8 @@
             } else if (child instanceof ClearAllButton) {
                 getPagedOrientationHandler().getPrimaryViewTranslate().set(child,
                         totalTranslationX);
-            } else {
-                // TODO(b/389209581): Handle the page offsets update of the 'mAddDesktopButton'.
+            } else if (child instanceof AddDesktopButton addDesktopButton) {
+                addDesktopButton.setOffsetTranslationX(totalTranslationX);
             }
             if (mEnableDrawingLiveTile && i == getRunningTaskIndex()) {
                 runActionOnRemoteHandles(
@@ -6157,7 +6157,7 @@
         if (addDesktopButtonIndex != -1 && addDesktopButtonIndex < outPageScrolls.length) {
             outPageScrolls[addDesktopButtonIndex] =
                     newPageScrolls[addDesktopButtonIndex] + Math.round(
-                            mAddDesktopButton.getTranslationX());
+                            mAddDesktopButton.getGridTranslationX());
         }
 
         int lastTaskScroll = getLastTaskScroll(clearAllScroll, clearAllWidth);
diff --git a/quickstep/src/com/android/quickstep/views/TaskContainer.kt b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
index 7ac0946..6b5d8dd 100644
--- a/quickstep/src/com/android/quickstep/views/TaskContainer.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskContainer.kt
@@ -28,6 +28,8 @@
 import com.android.quickstep.recents.di.get
 import com.android.quickstep.recents.di.getScope
 import com.android.quickstep.recents.di.inject
+import com.android.quickstep.recents.ui.mapper.TaskUiStateMapper
+import com.android.quickstep.recents.ui.viewmodel.TaskData
 import com.android.quickstep.recents.viewmodel.TaskContainerViewModel
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.task.viewmodel.TaskContainerData
@@ -163,4 +165,8 @@
         digitalWellBeingToast?.let { addAccessibleChildToList(it, outChildren) }
         overlay.addChildForAccessibility(outChildren)
     }
+
+    fun setState(state: TaskData?, liveTile: Boolean, hasHeader: Boolean) {
+        thumbnailView.setState(TaskUiStateMapper.toTaskThumbnailUiState(state, liveTile, hasHeader))
+    }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.kt b/quickstep/src/com/android/quickstep/views/TaskView.kt
index e7a395f..741297d 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskView.kt
@@ -70,6 +70,7 @@
 import com.android.launcher3.util.TraceHelper
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.launcher3.util.ViewPool
+import com.android.launcher3.util.coroutines.DispatcherProvider
 import com.android.launcher3.util.rects.set
 import com.android.quickstep.FullscreenDrawParams
 import com.android.quickstep.RecentsModel
@@ -77,6 +78,11 @@
 import com.android.quickstep.TaskOverlayFactory
 import com.android.quickstep.TaskViewUtils
 import com.android.quickstep.orientation.RecentsPagedOrientationHandler
+import com.android.quickstep.recents.di.RecentsDependencies
+import com.android.quickstep.recents.di.get
+import com.android.quickstep.recents.di.inject
+import com.android.quickstep.recents.ui.viewmodel.TaskTileUiState
+import com.android.quickstep.recents.ui.viewmodel.TaskViewModel
 import com.android.quickstep.util.ActiveGestureErrorDetector
 import com.android.quickstep.util.ActiveGestureLog
 import com.android.quickstep.util.BorderAnimator
@@ -88,6 +94,12 @@
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.android.systemui.shared.system.ActivityManagerWrapper
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.flow.collectLatest
+import kotlinx.coroutines.launch
 
 /** A task in the Recents view. */
 open class TaskView
@@ -448,6 +460,11 @@
     private val settledProgressDismiss =
         settledProgressPropertyFactory.get(SETTLED_PROGRESS_INDEX_DISMISS)
 
+    private var viewModel: TaskViewModel? = null
+    private val dispatcherProvider: DispatcherProvider by RecentsDependencies.inject()
+    private val coroutineScope by lazy { CoroutineScope(SupervisorJob() + dispatcherProvider.main) }
+    private val coroutineJobs = mutableListOf<Job>()
+
     /**
      * Returns an animator of [settledProgressDismiss] that transition in with a built-in
      * interpolator.
@@ -601,6 +618,8 @@
 
     override fun onRecycle() {
         resetPersistentViewTransforms()
+
+        viewModel = null
         attachAlpha = 1f
         splitAlpha = 1f
         // Clear any references to the thumbnail (it will be re-read either from the cache or the
@@ -715,6 +734,44 @@
             ?.inflate()
     }
 
+    override fun onAttachedToWindow() {
+        super.onAttachedToWindow()
+        if (enableRefactorTaskThumbnail()) {
+            // The TaskView lifecycle is starts the ViewModel during onBind, and cleans it in
+            // onRecycle. So it should be initialized at this point. TaskView Lifecycle:
+            // `bind` -> `onBind` ->  onAttachedToWindow() -> onDetachFromWindow -> onRecycle
+            coroutineJobs +=
+                coroutineScope.launch {
+                    viewModel!!.state.collectLatest(::updateTaskContainerState)
+                }
+        }
+    }
+
+    private fun updateTaskContainerState(state: TaskTileUiState) {
+        val mapOfTasks = state.tasks.associateBy { it.taskId }
+        taskContainers.forEach { container ->
+            container.setState(
+                state = mapOfTasks[container.task.key.id],
+                liveTile = state.isLiveTile,
+                hasHeader = type == TaskViewType.DESKTOP,
+            )
+        }
+    }
+
+    override fun onDetachedFromWindow() {
+        super.onDetachedFromWindow()
+        if (enableRefactorTaskThumbnail()) {
+            // The jobs are being cancelled in the background thread. So we make a copy of the list
+            // to prevent cleaning a new job that might be added to this list during onAttach
+            // or another moment in the lifecycle.
+            val coroutineJobsToCancel = coroutineJobs.toList()
+            coroutineJobs.clear()
+            coroutineScope.launch(dispatcherProvider.background) {
+                coroutineJobsToCancel.forEach { it.cancel("TaskView detaching from window") }
+            }
+        }
+    }
+
     /** Updates this task view to the given {@param task}. */
     open fun bind(
         task: Task,
@@ -738,6 +795,17 @@
     }
 
     open fun onBind(orientedState: RecentsOrientedState) {
+        if (enableRefactorTaskThumbnail()) {
+            viewModel =
+                TaskViewModel(
+                        taskViewType = type,
+                        recentsViewData = RecentsDependencies.get(),
+                        getTaskUseCase = RecentsDependencies.get(),
+                        dispatcherProvider = RecentsDependencies.get(),
+                    )
+                    .apply { bind(*taskIds) }
+        }
+
         taskContainers.forEach {
             it.bind()
             if (enableRefactorTaskThumbnail()) {
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
index 47d2bfc..e033e7b 100644
--- a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/FakeTaskThumbnailViewModel.kt
@@ -17,14 +17,12 @@
 package com.android.quickstep.task.thumbnail
 
 import android.graphics.Matrix
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
 import com.android.quickstep.task.viewmodel.TaskThumbnailViewModel
 import kotlinx.coroutines.flow.MutableStateFlow
 
 class FakeTaskThumbnailViewModel : TaskThumbnailViewModel {
     override val dimProgress = MutableStateFlow(0f)
     override val splashAlpha = MutableStateFlow(0f)
-    override val uiState = MutableStateFlow<TaskThumbnailUiState>(Uninitialized)
 
     override fun bind(taskId: Int) {
         // no-op
diff --git a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
index a76f83c..3b28afd 100644
--- a/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
+++ b/quickstep/tests/multivalentScreenshotTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewScreenshotTest.kt
@@ -60,36 +60,27 @@
         screenshotRule.screenshotTest("taskThumbnailView_uninitialized") { activity ->
             activity.actionBar?.hide()
             val taskThumbnailView = createTaskThumbnailView(activity)
-            taskThumbnailViewModel.uiState.value = BackgroundOnly(Color.YELLOW)
-            taskThumbnailViewModel.uiState.value = Uninitialized
+            taskThumbnailView.setState(Uninitialized)
             taskThumbnailView
         }
     }
 
     @Test
     fun taskThumbnailView_recyclesToUninitialized() {
-        screenshotRule.screenshotTest(
-            "taskThumbnailView_uninitialized",
-            viewProvider = { activity ->
-                activity.actionBar?.hide()
-                val taskThumbnailView = createTaskThumbnailView(activity)
-                taskThumbnailViewModel.uiState.value = BackgroundOnly(Color.YELLOW)
-                taskThumbnailView
-            },
-            checkView = { _, taskThumbnailView ->
-                // Call onRecycle() after View is attached (end of block above)
-                (taskThumbnailView as TaskThumbnailView).onRecycle()
-                false
-            },
-        )
+        screenshotRule.screenshotTest("taskThumbnailView_uninitialized") { activity ->
+            activity.actionBar?.hide()
+            val taskThumbnailView = createTaskThumbnailView(activity)
+            taskThumbnailView.setState(BackgroundOnly(Color.YELLOW))
+            taskThumbnailView.onRecycle()
+            taskThumbnailView
+        }
     }
 
     @Test
     fun taskThumbnailView_backgroundOnly() {
         screenshotRule.screenshotTest("taskThumbnailView_backgroundOnly") { activity ->
             activity.actionBar?.hide()
-            taskThumbnailViewModel.uiState.value = BackgroundOnly(Color.YELLOW)
-            createTaskThumbnailView(activity)
+            createTaskThumbnailView(activity).apply { setState(BackgroundOnly(Color.YELLOW)) }
         }
     }
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
index 588c22c..021e1e4 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/TaskbarStashControllerTest.kt
@@ -53,7 +53,7 @@
 import com.android.launcher3.util.LauncherMultivalentJUnit
 import com.android.launcher3.util.LauncherMultivalentJUnit.EmulatedDevices
 import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_BUBBLES_EXPANDED
-import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_SHOWING
+import com.android.systemui.shared.system.QuickStepContract.SYSUI_STATE_IME_VISIBLE
 import com.android.wm.shell.Flags.FLAG_ENABLE_BUBBLE_BAR
 import com.google.common.truth.Truth.assertThat
 import com.google.common.truth.TruthJUnit.assume
@@ -542,7 +542,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, false)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, false)
             animatorTestRule.advanceTimeBy(TASKBAR_STASH_DURATION_FOR_IME)
         }
         assertThat(viewController.areIconsVisible()).isFalse()
@@ -555,7 +555,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
             animatorTestRule.advanceTimeBy(0)
         }
 
@@ -574,7 +574,7 @@
 
         // Start with IME shown.
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
             animatorTestRule.advanceTimeBy(0)
         }
 
@@ -600,7 +600,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, false)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, false)
             animatorTestRule.advanceTimeBy(TASKBAR_STASH_DURATION_FOR_IME)
         }
         assertThat(viewController.areIconsVisible()).isFalse()
@@ -613,7 +613,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, false)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, false)
             animatorTestRule.advanceTimeBy(TASKBAR_STASH_DURATION_FOR_IME)
         }
 
@@ -633,7 +633,7 @@
         assume().that(activityContext.isHardwareKeyboard).isFalse()
 
         getInstrumentation().runOnMainSync {
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
             animatorTestRule.advanceTimeBy(0)
         }
 
@@ -653,7 +653,7 @@
 
         getInstrumentation().runOnMainSync {
             stashController.updateStateForFlag(FLAG_IN_APP, true)
-            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_SHOWING, true)
+            stashController.updateStateForSysuiFlags(SYSUI_STATE_IME_VISIBLE, true)
         }
 
         assertThat(stashController.isStashed).isFalse()
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/mapper/TaskUiStateMapperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/mapper/TaskUiStateMapperTest.kt
new file mode 100644
index 0000000..124045f
--- /dev/null
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/mapper/TaskUiStateMapperTest.kt
@@ -0,0 +1,231 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.recents.ui.mapper
+
+import android.graphics.Bitmap
+import android.graphics.Color
+import android.graphics.drawable.ShapeDrawable
+import android.platform.test.annotations.EnableFlags
+import android.view.Surface
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import com.android.launcher3.Flags
+import com.android.quickstep.recents.ui.viewmodel.TaskData
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
+import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.ThumbnailHeader
+import com.android.systemui.shared.recents.model.ThumbnailData
+import com.google.common.truth.Truth.assertThat
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@RunWith(AndroidJUnit4::class)
+class TaskUiStateMapperTest {
+
+    @Test
+    fun taskData_isNull_returns_Uninitialized() {
+        val result =
+            TaskUiStateMapper.toTaskThumbnailUiState(
+                taskData = null,
+                isLiveTile = false,
+                hasHeader = false,
+            )
+        assertThat(result).isEqualTo(TaskThumbnailUiState.Uninitialized)
+    }
+
+    @Test
+    fun taskData_isLiveTile_returns_LiveTile() {
+        val inputs =
+            listOf(TASK_DATA, TASK_DATA.copy(thumbnailData = null), TASK_DATA.copy(isLocked = true))
+        inputs.forEach { input ->
+            val result =
+                TaskUiStateMapper.toTaskThumbnailUiState(
+                    taskData = input,
+                    isLiveTile = true,
+                    hasHeader = false,
+                )
+            assertThat(result).isEqualTo(LiveTile.WithoutHeader)
+        }
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_EXPLODED_VIEW)
+    @Test
+    fun taskData_isLiveTileWithHeader_returns_LiveTileWithHeader() {
+        val inputs =
+            listOf(
+                TASK_DATA,
+                TASK_DATA.copy(thumbnailData = null),
+                TASK_DATA.copy(isLocked = true),
+                TASK_DATA.copy(title = null),
+            )
+        val expected =
+            LiveTile.WithHeader(header = ThumbnailHeader(TASK_ICON, TASK_TITLE_DESCRIPTION))
+        inputs.forEach { taskData ->
+            val result =
+                TaskUiStateMapper.toTaskThumbnailUiState(
+                    taskData = taskData,
+                    isLiveTile = true,
+                    hasHeader = true,
+                )
+            assertThat(result).isEqualTo(expected)
+        }
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_EXPLODED_VIEW)
+    @Test
+    fun taskData_isLiveTileWithHeader_missingHeaderData_returns_LiveTileWithoutHeader() {
+        val inputs =
+            listOf(
+                TASK_DATA.copy(icon = null),
+                TASK_DATA.copy(titleDescription = null),
+                TASK_DATA.copy(icon = null, titleDescription = null),
+            )
+
+        inputs.forEach { taskData ->
+            val result =
+                TaskUiStateMapper.toTaskThumbnailUiState(
+                    taskData = taskData,
+                    isLiveTile = true,
+                    hasHeader = true,
+                )
+            assertThat(result).isEqualTo(LiveTile.WithoutHeader)
+        }
+    }
+
+    @Test
+    fun taskData_isStaticTile_returns_SnapshotSplash() {
+        val result =
+            TaskUiStateMapper.toTaskThumbnailUiState(
+                taskData = TASK_DATA,
+                isLiveTile = false,
+                hasHeader = false,
+            )
+
+        val expected =
+            TaskThumbnailUiState.SnapshotSplash(
+                snapshot =
+                    Snapshot.WithoutHeader(
+                        backgroundColor = TASK_BACKGROUND_COLOR,
+                        bitmap = TASK_THUMBNAIL,
+                        thumbnailRotation = Surface.ROTATION_0,
+                    ),
+                splash = TASK_ICON,
+            )
+
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_EXPLODED_VIEW)
+    @Test
+    fun taskData_isStaticTile_withHeader_returns_SnapshotSplashWithHeader() {
+        val inputs = listOf(TASK_DATA, TASK_DATA.copy(title = null))
+        val expected =
+            TaskThumbnailUiState.SnapshotSplash(
+                snapshot =
+                    Snapshot.WithHeader(
+                        backgroundColor = TASK_BACKGROUND_COLOR,
+                        bitmap = TASK_THUMBNAIL,
+                        thumbnailRotation = Surface.ROTATION_0,
+                        header = ThumbnailHeader(TASK_ICON, TASK_TITLE_DESCRIPTION),
+                    ),
+                splash = TASK_ICON,
+            )
+        inputs.forEach { taskData ->
+            val result =
+                TaskUiStateMapper.toTaskThumbnailUiState(
+                    taskData = taskData,
+                    isLiveTile = false,
+                    hasHeader = true,
+                )
+            assertThat(result).isEqualTo(expected)
+        }
+    }
+
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_EXPLODED_VIEW)
+    @Test
+    fun taskData_isStaticTile_missingHeaderData_returns_SnapshotSplashWithoutHeader() {
+        val inputs =
+            listOf(
+                TASK_DATA.copy(titleDescription = null, icon = null),
+                TASK_DATA.copy(titleDescription = null),
+                TASK_DATA.copy(icon = null),
+            )
+        val expected =
+            Snapshot.WithoutHeader(
+                backgroundColor = TASK_BACKGROUND_COLOR,
+                thumbnailRotation = Surface.ROTATION_0,
+                bitmap = TASK_THUMBNAIL,
+            )
+        inputs.forEach { taskData ->
+            val result =
+                TaskUiStateMapper.toTaskThumbnailUiState(
+                    taskData = taskData,
+                    isLiveTile = false,
+                    hasHeader = true,
+                )
+
+            assertThat(result).isInstanceOf(TaskThumbnailUiState.SnapshotSplash::class.java)
+            result as TaskThumbnailUiState.SnapshotSplash
+            assertThat(result.snapshot).isEqualTo(expected)
+        }
+    }
+
+    @Test
+    fun taskData_thumbnailIsNull_returns_BackgroundOnly() {
+        val result =
+            TaskUiStateMapper.toTaskThumbnailUiState(
+                taskData = TASK_DATA.copy(thumbnailData = null),
+                isLiveTile = false,
+                hasHeader = false,
+            )
+
+        val expected = TaskThumbnailUiState.BackgroundOnly(TASK_BACKGROUND_COLOR)
+        assertThat(result).isEqualTo(expected)
+    }
+
+    @Test
+    fun taskData_isLocked_returns_BackgroundOnly() {
+        val result =
+            TaskUiStateMapper.toTaskThumbnailUiState(
+                taskData = TASK_DATA.copy(isLocked = true),
+                isLiveTile = false,
+                hasHeader = false,
+            )
+
+        val expected = TaskThumbnailUiState.BackgroundOnly(TASK_BACKGROUND_COLOR)
+        assertThat(result).isEqualTo(expected)
+    }
+
+    private companion object {
+        const val TASK_TITLE_DESCRIPTION = "Title Description 1"
+        val TASK_ICON = ShapeDrawable()
+        val TASK_THUMBNAIL = Bitmap.createBitmap(1, 1, Bitmap.Config.ARGB_8888)
+        val TASK_THUMBNAIL_DATA =
+            ThumbnailData(thumbnail = TASK_THUMBNAIL, rotation = Surface.ROTATION_0)
+        val TASK_BACKGROUND_COLOR = Color.rgb(1, 2, 3)
+        val TASK_DATA =
+            TaskData.Data(
+                1,
+                title = "Task 1",
+                titleDescription = TASK_TITLE_DESCRIPTION,
+                icon = TASK_ICON,
+                thumbnailData = TASK_THUMBNAIL_DATA,
+                backgroundColor = TASK_BACKGROUND_COLOR,
+                isLocked = false,
+            )
+    }
+}
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
index 54a27e9..7a4b5f2 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/recents/ui/viewmodel/TaskViewModelTest.kt
@@ -23,6 +23,7 @@
 import com.android.quickstep.recents.domain.model.TaskModel
 import com.android.quickstep.recents.domain.usecase.GetTaskUseCase
 import com.android.quickstep.recents.viewmodel.RecentsViewData
+import com.android.quickstep.views.TaskViewType
 import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -45,15 +46,17 @@
 
     private val recentsViewData = RecentsViewData()
     private val getTaskUseCase = mock<GetTaskUseCase>()
-    private val sut =
-        TaskViewModel(
-            recentsViewData = recentsViewData,
-            getTaskUseCase = getTaskUseCase,
-            dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
-        )
+    private lateinit var sut: TaskViewModel
 
     @Before
     fun setUp() {
+        sut =
+            TaskViewModel(
+                taskViewType = TaskViewType.SINGLE,
+                recentsViewData = recentsViewData,
+                getTaskUseCase = getTaskUseCase,
+                dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
+            )
         whenever(getTaskUseCase.invoke(TASK_MODEL_1.id)).thenReturn(flow { emit(TASK_MODEL_1) })
         whenever(getTaskUseCase.invoke(TASK_MODEL_2.id)).thenReturn(flow { emit(TASK_MODEL_2) })
         whenever(getTaskUseCase.invoke(TASK_MODEL_3.id)).thenReturn(flow { emit(TASK_MODEL_3) })
@@ -65,11 +68,39 @@
     fun singleTaskRetrieved_when_validTaskId() =
         testScope.runTest {
             sut.bind(TASK_MODEL_1.id)
-            val expectedResult = TaskTileUiState(listOf(TASK_MODEL_1.toUiState()), false)
+            val expectedResult =
+                TaskTileUiState(
+                    tasks = listOf(TASK_MODEL_1.toUiState()),
+                    isLiveTile = false,
+                    hasHeader = false,
+                )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
 
     @Test
+    fun hasHeader_when_taskViewTypeIsDesktop() =
+        testScope.runTest {
+            val expectedResults =
+                mapOf(
+                    TaskViewType.SINGLE to false,
+                    TaskViewType.GROUPED to false,
+                    TaskViewType.DESKTOP to true,
+                )
+
+            expectedResults.forEach { (type, expectedResult) ->
+                sut =
+                    TaskViewModel(
+                        taskViewType = type,
+                        recentsViewData = recentsViewData,
+                        getTaskUseCase = getTaskUseCase,
+                        dispatcherProvider = TestDispatcherProvider(unconfinedTestDispatcher),
+                    )
+                sut.bind(TASK_MODEL_1.id)
+                assertThat(sut.state.first().hasHeader).isEqualTo(expectedResult)
+            }
+        }
+
+    @Test
     fun multipleTasksRetrieved_when_validTaskIds() =
         testScope.runTest {
             sut.bind(TASK_MODEL_1.id, TASK_MODEL_2.id, TASK_MODEL_3.id, INVALID_TASK_ID)
@@ -83,6 +114,7 @@
                             TaskData.NoData(INVALID_TASK_ID),
                         ),
                     isLiveTile = false,
+                    hasHeader = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -103,6 +135,7 @@
                             TASK_MODEL_3.toUiState(),
                         ),
                     isLiveTile = true,
+                    hasHeader = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -123,6 +156,7 @@
                             TASK_MODEL_3.toUiState(),
                         ),
                     isLiveTile = false,
+                    hasHeader = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -142,6 +176,7 @@
                             TASK_MODEL_3.toUiState(),
                         ),
                     isLiveTile = false,
+                    hasHeader = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -157,6 +192,7 @@
                 TaskTileUiState(
                     tasks = listOf(TASK_MODEL_1.toUiState(), TASK_MODEL_2.toUiState()),
                     isLiveTile = false,
+                    hasHeader = false,
                 )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
@@ -166,7 +202,11 @@
         testScope.runTest {
             sut.bind(INVALID_TASK_ID)
             val expectedResult =
-                TaskTileUiState(listOf(TaskData.NoData(INVALID_TASK_ID)), isLiveTile = false)
+                TaskTileUiState(
+                    listOf(TaskData.NoData(INVALID_TASK_ID)),
+                    isLiveTile = false,
+                    hasHeader = false,
+                )
             assertThat(sut.state.first()).isEqualTo(expectedResult)
         }
 
@@ -174,6 +214,7 @@
         TaskData.Data(
             taskId = id,
             title = title,
+            titleDescription = titleDescription,
             icon = icon!!,
             thumbnailData = thumbnail,
             backgroundColor = backgroundColor,
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
index a956c9c..22636b0 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/task/thumbnail/TaskThumbnailViewModelImplTest.kt
@@ -16,37 +16,16 @@
 
 package com.android.quickstep.task.thumbnail
 
-import android.app.WindowConfiguration.WINDOWING_MODE_FREEFORM
-import android.app.WindowConfiguration.WINDOWING_MODE_FULLSCREEN
-import android.content.ComponentName
-import android.content.Intent
-import android.graphics.Bitmap
-import android.graphics.Color
 import android.graphics.Matrix
-import android.graphics.drawable.Drawable
-import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
-import android.view.Surface
 import androidx.test.ext.junit.runners.AndroidJUnit4
-import com.android.launcher3.Flags
 import com.android.launcher3.util.TestDispatcherProvider
-import com.android.quickstep.recents.data.FakeRecentsDeviceProfileRepository
-import com.android.quickstep.recents.data.FakeTasksRepository
-import com.android.quickstep.recents.data.RecentsDeviceProfile
 import com.android.quickstep.recents.usecase.GetThumbnailPositionUseCase
 import com.android.quickstep.recents.usecase.ThumbnailPositionState.MatrixScaling
 import com.android.quickstep.recents.usecase.ThumbnailPositionState.MissingThumbnail
 import com.android.quickstep.recents.viewmodel.RecentsViewData
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.BackgroundOnly
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.LiveTile
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Snapshot
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.SnapshotSplash
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.ThumbnailHeader
-import com.android.quickstep.task.thumbnail.TaskThumbnailUiState.Uninitialized
 import com.android.quickstep.task.viewmodel.TaskContainerData
 import com.android.quickstep.task.viewmodel.TaskThumbnailViewModelImpl
-import com.android.systemui.shared.recents.model.Task
-import com.android.systemui.shared.recents.model.ThumbnailData
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.test.StandardTestDispatcher
@@ -69,8 +48,6 @@
     private val recentsViewData = RecentsViewData()
     private val taskContainerData = TaskContainerData()
     private val dispatcherProvider = TestDispatcherProvider(dispatcher)
-    private val tasksRepository = FakeTasksRepository()
-    private val deviceProfileRepository = FakeRecentsDeviceProfileRepository()
     private val mGetThumbnailPositionUseCase = mock<GetThumbnailPositionUseCase>()
     private val splashAlphaUseCase: SplashAlphaUseCase = mock()
 
@@ -79,208 +56,11 @@
             recentsViewData,
             taskContainerData,
             dispatcherProvider,
-            tasksRepository,
-            deviceProfileRepository,
             mGetThumbnailPositionUseCase,
             splashAlphaUseCase,
         )
     }
 
-    private val fullscreenTaskIdRange: IntRange = 0..5
-    private val freeformTaskIdRange: IntRange = 6..10
-
-    private val fullscreenTasks = fullscreenTaskIdRange.map(::createTaskWithId)
-    private val freeformTasks = freeformTaskIdRange.map(::createFreeformTaskWithId)
-    private val tasks = fullscreenTasks + freeformTasks
-
-    @Test
-    fun initialStateIsUninitialized() =
-        testScope.runTest { assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized) }
-
-    @Test
-    fun bindRunningTask_thenStateIs_LiveTile() =
-        testScope.runTest {
-            val taskId = 1
-            tasksRepository.seedTasks(tasks)
-            tasksRepository.setVisibleTasks(setOf(taskId))
-            recentsViewData.runningTaskIds.value = setOf(taskId)
-            systemUnderTest.bind(taskId)
-
-            assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile.WithoutHeader)
-        }
-
-    @Test
-    fun bindRunningTaskShouldShowScreenshot_thenStateIs_SnapshotSplash() =
-        testScope.runTest {
-            val taskId = 1
-            val expectedThumbnailData = createThumbnailData()
-            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
-            val expectedIconData = mock<Drawable>()
-            tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
-            tasksRepository.seedTasks(tasks)
-            tasksRepository.setVisibleTasks(setOf(taskId))
-            recentsViewData.runningTaskIds.value = setOf(taskId)
-            recentsViewData.runningTaskShowScreenshot.value = true
-            systemUnderTest.bind(taskId)
-
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(
-                    SnapshotSplash(
-                        Snapshot.WithoutHeader(
-                            backgroundColor = Color.rgb(1, 1, 1),
-                            bitmap = expectedThumbnailData.thumbnail!!,
-                            thumbnailRotation = Surface.ROTATION_0,
-                        ),
-                        expectedIconData,
-                    )
-                )
-        }
-
-    @Test
-    fun bindRunningTaskThenStoppedTaskWithoutThumbnail_thenStateChangesToBackgroundOnly() =
-        testScope.runTest {
-            val runningTaskId = 1
-            val stoppedTaskId = 2
-            tasksRepository.seedTasks(tasks)
-            tasksRepository.setVisibleTasks(setOf(runningTaskId, stoppedTaskId))
-            recentsViewData.runningTaskIds.value = setOf(runningTaskId)
-            systemUnderTest.bind(runningTaskId)
-            assertThat(systemUnderTest.uiState.first()).isEqualTo(LiveTile.WithoutHeader)
-
-            systemUnderTest.bind(stoppedTaskId)
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
-        }
-
-    @Test
-    fun bindStoppedTaskWithoutThumbnail_thenStateIs_BackgroundOnly_withAlphaRemoved() =
-        testScope.runTest {
-            val stoppedTaskId = 2
-            tasksRepository.seedTasks(tasks)
-            tasksRepository.setVisibleTasks(setOf(stoppedTaskId))
-
-            systemUnderTest.bind(stoppedTaskId)
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
-        }
-
-    @Test
-    fun bindLockedTaskWithThumbnail_thenStateIs_BackgroundOnly() =
-        testScope.runTest {
-            val taskId = 2
-            tasksRepository.seedThumbnailData(mapOf(taskId to createThumbnailData()))
-            tasks[taskId].isLocked = true
-            tasksRepository.seedTasks(tasks)
-            tasksRepository.setVisibleTasks(setOf(taskId))
-
-            systemUnderTest.bind(taskId)
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(BackgroundOnly(backgroundColor = Color.rgb(2, 2, 2)))
-        }
-
-    @Test
-    fun bindStoppedTaskWithThumbnail_thenStateIs_SnapshotSplash_withAlphaRemoved() =
-        testScope.runTest {
-            val taskId = 2
-            val expectedThumbnailData = createThumbnailData(rotation = Surface.ROTATION_270)
-            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
-            val expectedIconData = mock<Drawable>()
-            tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
-            tasksRepository.seedTasks(tasks)
-            tasksRepository.setVisibleTasks(setOf(taskId))
-
-            systemUnderTest.bind(taskId)
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(
-                    SnapshotSplash(
-                        Snapshot.WithoutHeader(
-                            backgroundColor = Color.rgb(2, 2, 2),
-                            bitmap = expectedThumbnailData.thumbnail!!,
-                            thumbnailRotation = Surface.ROTATION_270,
-                        ),
-                        expectedIconData,
-                    )
-                )
-        }
-
-    @Test
-    fun bindNonVisibleStoppedTask_whenMadeVisible_thenStateIsSnapshotSplash() =
-        testScope.runTest {
-            val taskId = 2
-            val expectedThumbnailData = createThumbnailData()
-            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
-            val expectedIconData = mock<Drawable>()
-            tasksRepository.seedIconData(taskId, "Task $taskId", "", expectedIconData)
-            tasksRepository.seedTasks(tasks)
-
-            systemUnderTest.bind(taskId)
-            assertThat(systemUnderTest.uiState.first()).isEqualTo(Uninitialized)
-
-            tasksRepository.setVisibleTasks(setOf(taskId))
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(
-                    SnapshotSplash(
-                        Snapshot.WithoutHeader(
-                            backgroundColor = Color.rgb(2, 2, 2),
-                            bitmap = expectedThumbnailData.thumbnail!!,
-                            thumbnailRotation = Surface.ROTATION_0,
-                        ),
-                        expectedIconData,
-                    )
-                )
-        }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_EXPLODED_VIEW)
-    fun bindRunningTask_inDesktop_thenStateIs_LiveTile_withHeader() =
-        testScope.runTest {
-            deviceProfileRepository.setRecentsDeviceProfile(
-                RecentsDeviceProfile(isLargeScreen = true, canEnterDesktopMode = true)
-            )
-
-            val taskId = freeformTaskIdRange.first
-            val expectedIconData = mock<Drawable>()
-            tasksRepository.seedIconData(taskId, "Task $taskId", "Task $taskId", expectedIconData)
-            tasksRepository.seedTasks(freeformTasks)
-            tasksRepository.setVisibleTasks(setOf(taskId))
-            recentsViewData.runningTaskIds.value = setOf(taskId)
-            systemUnderTest.bind(taskId)
-
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(LiveTile.WithHeader(ThumbnailHeader(expectedIconData, "Task $taskId")))
-        }
-
-    @Test
-    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_EXPLODED_VIEW)
-    fun bindStoppedTaskWithThumbnail_inDesktop_thenStateIs_SnapshotSplash_withHeader() =
-        testScope.runTest {
-            deviceProfileRepository.setRecentsDeviceProfile(
-                RecentsDeviceProfile(isLargeScreen = true, canEnterDesktopMode = true)
-            )
-
-            val taskId = freeformTaskIdRange.first
-            val expectedThumbnailData = createThumbnailData(rotation = Surface.ROTATION_0)
-            tasksRepository.seedThumbnailData(mapOf(taskId to expectedThumbnailData))
-            val expectedIconData = mock<Drawable>()
-            tasksRepository.seedIconData(taskId, "Task $taskId", "Task $taskId", expectedIconData)
-            tasksRepository.seedTasks(freeformTasks)
-            tasksRepository.setVisibleTasks(setOf(taskId))
-
-            systemUnderTest.bind(taskId)
-            assertThat(systemUnderTest.uiState.first())
-                .isEqualTo(
-                    SnapshotSplash(
-                        Snapshot.WithHeader(
-                            backgroundColor = Color.rgb(taskId, taskId, taskId),
-                            bitmap = expectedThumbnailData.thumbnail!!,
-                            thumbnailRotation = Surface.ROTATION_0,
-                            header = ThumbnailHeader(expectedIconData, "Task $taskId"),
-                        ),
-                        expectedIconData,
-                    )
-                )
-        }
-
     @Test
     fun getSnapshotMatrix_MissingThumbnail() =
         testScope.runTest {
@@ -337,51 +117,7 @@
             assertThat(systemUnderTest.dimProgress.first()).isEqualTo(0f)
         }
 
-    private fun createTaskWithId(taskId: Int) =
-        Task(
-                Task.TaskKey(
-                    taskId,
-                    WINDOWING_MODE_FULLSCREEN,
-                    Intent(),
-                    ComponentName("", ""),
-                    0,
-                    2000,
-                )
-            )
-            .apply {
-                colorBackground = Color.argb(taskId, taskId, taskId, taskId)
-                titleDescription = "Task $taskId"
-                icon = mock<Drawable>()
-            }
-
-    private fun createFreeformTaskWithId(taskId: Int) =
-        Task(
-                Task.TaskKey(
-                    taskId,
-                    WINDOWING_MODE_FREEFORM,
-                    Intent(),
-                    ComponentName("", ""),
-                    0,
-                    2000,
-                )
-            )
-            .apply {
-                colorBackground = Color.argb(taskId, taskId, taskId, taskId)
-                titleDescription = "Task $taskId"
-                icon = mock<Drawable>()
-            }
-
-    private fun createThumbnailData(rotation: Int = Surface.ROTATION_0): ThumbnailData {
-        val bitmap = mock<Bitmap>()
-        whenever(bitmap.width).thenReturn(THUMBNAIL_WIDTH)
-        whenever(bitmap.height).thenReturn(THUMBNAIL_HEIGHT)
-
-        return ThumbnailData(thumbnail = bitmap, rotation = rotation)
-    }
-
-    companion object {
-        const val THUMBNAIL_WIDTH = 100
-        const val THUMBNAIL_HEIGHT = 200
+    private companion object {
         const val CANVAS_WIDTH = 300
         const val CANVAS_HEIGHT = 600
         val MATRIX =
diff --git a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
index c152ee1..52bd2ea 100644
--- a/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/DesktopSystemShortcutTest.kt
@@ -17,28 +17,33 @@
 package com.android.quickstep
 
 import android.content.ComponentName
+import android.content.Context
 import android.content.Intent
+import android.platform.test.annotations.EnableFlags
+import android.view.Display.DEFAULT_DISPLAY
 import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.internal.R
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingViewHelper
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
 import com.android.launcher3.model.data.TaskViewItemInfo
-import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory.TaskOverlay
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.views.LauncherRecentsView
+import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
 import com.android.quickstep.views.TaskViewIcon
 import com.android.systemui.shared.recents.model.Task
 import com.android.systemui.shared.recents.model.Task.TaskKey
+import com.android.window.flags.Flags
 import com.android.wm.shell.shared.desktopmode.DesktopModeStatus
 import com.android.wm.shell.shared.desktopmode.DesktopModeTransitionSource
 import com.google.common.truth.Truth.assertThat
@@ -58,7 +63,7 @@
 /** Test for [DesktopSystemShortcut] */
 class DesktopSystemShortcutTest {
 
-    private val launcher: QuickstepLauncher = mock()
+    private val launcher: RecentsViewContainer = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
@@ -67,6 +72,7 @@
     private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         DesktopSystemShortcut.createFactory(abstractFloatingViewHelper)
+    private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -79,6 +85,7 @@
                 .startMocking()
         whenever(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(true)
         whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
+        whenever(launcher.asContext()).thenReturn(context)
     }
 
     @After
@@ -97,6 +104,53 @@
     }
 
     @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
+    fun createDesktopTaskShortcutFactory_transparentTask() {
+        val baseComponent = ComponentName("", /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ true,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
+    @EnableFlags(Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY)
+    fun createDesktopTaskShortcutFactory_systemUiTask() {
+        val sysUiPackageName: String = context.resources.getString(R.string.config_systemUi)
+        val baseComponent = ComponentName(sysUiPackageName, /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
     fun createDesktopTaskShortcutFactory_undockable() {
         val unDockableTask = createTask().apply { isDockable = false }
         val taskContainer = createTaskContainer(unDockableTask)
@@ -114,8 +168,7 @@
         whenever(launcher.statsLogManager).thenReturn(statsLogManager)
         whenever(statsLogManager.logger()).thenReturn(statsLogger)
         whenever(statsLogger.withItemInfo(any())).thenReturn(statsLogger)
-        whenever(taskView.context)
-            .thenReturn(InstrumentationRegistry.getInstrumentation().targetContext)
+        whenever(taskView.context).thenReturn(context)
         whenever(recentsView.moveTaskToDesktop(any(), any(), any())).thenAnswer {
             val successCallback = it.getArgument<Runnable>(2)
             successCallback.run()
@@ -145,7 +198,22 @@
     }
 
     private fun createTask() =
-        Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000)).apply { isDockable = true }
+        Task(
+                TaskKey(
+                    /* id */ 1,
+                    /* windowingMode */ 0,
+                    Intent(),
+                    ComponentName("", ""),
+                    /* userId */ 0,
+                    /* lastActiveTime */ 2000,
+                    DEFAULT_DISPLAY,
+                    ComponentName("", ""),
+                    /* numActivities */ 1,
+                    /* isTopActivityNoDisplay */ false,
+                    /* isActivityStackTransparent */ false,
+                )
+            )
+            .apply { isDockable = true }
 
     private fun createTaskContainer(task: Task) =
         TaskContainer(
diff --git a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
index 9c2c13c..4111dec 100644
--- a/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
+++ b/quickstep/tests/src/com/android/quickstep/ExternalDisplaySystemShortcutTest.kt
@@ -17,23 +17,27 @@
 package com.android.quickstep
 
 import android.content.ComponentName
+import android.content.Context
 import android.content.Intent
 import android.platform.test.annotations.EnableFlags
 import android.platform.test.flag.junit.SetFlagsRule
+import android.view.Display.DEFAULT_DISPLAY
+import androidx.test.platform.app.InstrumentationRegistry
 import com.android.dx.mockito.inline.extended.ExtendedMockito.mockitoSession
 import com.android.dx.mockito.inline.extended.StaticMockitoSession
+import com.android.internal.R
 import com.android.launcher3.AbstractFloatingView
 import com.android.launcher3.AbstractFloatingViewHelper
 import com.android.launcher3.Flags.enableRefactorTaskThumbnail
 import com.android.launcher3.logging.StatsLogManager
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent
 import com.android.launcher3.model.data.TaskViewItemInfo
-import com.android.launcher3.uioverrides.QuickstepLauncher
 import com.android.launcher3.util.SplitConfigurationOptions
 import com.android.launcher3.util.TransformingTouchDelegate
 import com.android.quickstep.TaskOverlayFactory.TaskOverlay
 import com.android.quickstep.task.thumbnail.TaskThumbnailView
 import com.android.quickstep.views.LauncherRecentsView
+import com.android.quickstep.views.RecentsViewContainer
 import com.android.quickstep.views.TaskContainer
 import com.android.quickstep.views.TaskThumbnailViewDeprecated
 import com.android.quickstep.views.TaskView
@@ -62,7 +66,7 @@
 
     @get:Rule val setFlagsRule = SetFlagsRule(SetFlagsRule.DefaultInitValueType.DEVICE_DEFAULT)
 
-    private val launcher: QuickstepLauncher = mock()
+    private val launcher: RecentsViewContainer = mock()
     private val statsLogManager: StatsLogManager = mock()
     private val statsLogger: StatsLogManager.StatsLogger = mock()
     private val recentsView: LauncherRecentsView = mock()
@@ -71,6 +75,7 @@
     private val overlayFactory: TaskOverlayFactory = mock()
     private val factory: TaskShortcutFactory =
         ExternalDisplaySystemShortcut.createFactory(abstractFloatingViewHelper)
+    private val context: Context = InstrumentationRegistry.getInstrumentation().targetContext
 
     private lateinit var mockitoSession: StaticMockitoSession
 
@@ -83,6 +88,7 @@
                 .startMocking()
         whenever(DesktopModeStatus.canEnterDesktopMode(any())).thenReturn(true)
         whenever(overlayFactory.createOverlay(any())).thenReturn(mock<TaskOverlay<*>>())
+        whenever(launcher.asContext()).thenReturn(context)
     }
 
     @After
@@ -102,6 +108,59 @@
     }
 
     @Test
+    @EnableFlags(
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+    )
+    fun createExternalDisplayTaskShortcut_transparentTask() {
+        val baseComponent = ComponentName("", /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ true,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
+    @EnableFlags(
+        Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT,
+        Flags.FLAG_ENABLE_DESKTOP_WINDOWING_MODALS_POLICY,
+    )
+    fun createExternalDisplayTaskShortcut_systemUiTask() {
+        val sysUiPackageName: String = context.resources.getString(R.string.config_systemUi)
+        val baseComponent = ComponentName(sysUiPackageName, /* class */ "")
+        val taskKey =
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                baseComponent,
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                baseComponent,
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        val taskContainer = createTaskContainer(Task(taskKey))
+        val shortcuts = factory.getShortcuts(launcher, taskContainer)
+        assertThat(shortcuts).isNull()
+    }
+
+    @Test
     @EnableFlags(Flags.FLAG_MOVE_TO_EXTERNAL_DISPLAY_SHORTCUT)
     fun externalDisplaySystemShortcutClicked() {
         val task = createTask()
@@ -134,7 +193,22 @@
         verify(statsLogger).log(LauncherEvent.LAUNCHER_SYSTEM_SHORTCUT_EXTERNAL_DISPLAY_TAP)
     }
 
-    private fun createTask() = Task(TaskKey(1, 0, Intent(), ComponentName("", ""), 0, 2000))
+    private fun createTask() =
+        Task(
+            TaskKey(
+                /* id */ 1,
+                /* windowingMode */ 0,
+                Intent(),
+                ComponentName("", ""),
+                /* userId */ 0,
+                /* lastActiveTime */ 2000,
+                DEFAULT_DISPLAY,
+                ComponentName("", ""),
+                /* numActivities */ 1,
+                /* isTopActivityNoDisplay */ false,
+                /* isActivityStackTransparent */ false,
+            )
+        )
 
     private fun createTaskContainer(task: Task) =
         TaskContainer(
diff --git a/src/com/android/launcher3/CellLayout.java b/src/com/android/launcher3/CellLayout.java
index c38444c..988d164 100644
--- a/src/com/android/launcher3/CellLayout.java
+++ b/src/com/android/launcher3/CellLayout.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3;
 
+import static com.android.launcher3.LauncherState.NORMAL;
 import static com.android.launcher3.dragndrop.DraggableView.DRAGGABLE_ICON;
 import static com.android.launcher3.icons.IconNormalizer.ICON_VISIBLE_AREA_FACTOR;
 import static com.android.launcher3.util.MultiTranslateDelegate.INDEX_REORDER_PREVIEW_OFFSET;
@@ -47,6 +48,7 @@
 import android.view.ViewDebug;
 import android.view.ViewGroup;
 import android.view.accessibility.AccessibilityEvent;
+import android.view.accessibility.AccessibilityNodeInfo;
 
 import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
@@ -69,6 +71,7 @@
 import com.android.launcher3.folder.PreviewBackground;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
+import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.util.CellAndSpan;
 import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.MSDLPlayerWrapper;
@@ -785,6 +788,22 @@
             }
             mShortcutsAndWidgets.addView(child, index, lp);
 
+            // Whenever an app is added, if Accessibility service is enabled, focus on that app.
+            if (mActivity instanceof Launcher) {
+                Launcher.cast(mActivity).getStateManager().addStateListener(
+                        new StateManager.StateListener<LauncherState>() {
+                            @Override
+                            public void onStateTransitionComplete(LauncherState finalState) {
+                                if (finalState == NORMAL) {
+                                    child.performAccessibilityAction(
+                                            AccessibilityNodeInfo.ACTION_ACCESSIBILITY_FOCUS, null);
+                                    Launcher.cast(mActivity).getStateManager()
+                                            .removeStateListener(this);
+                                }
+                            }
+                        });
+            }
+
             if (markCells) markCellsAsOccupiedForView(child);
 
             return true;
diff --git a/src/com/android/launcher3/Utilities.java b/src/com/android/launcher3/Utilities.java
index 7563493..d93c07f 100644
--- a/src/com/android/launcher3/Utilities.java
+++ b/src/com/android/launcher3/Utilities.java
@@ -659,7 +659,8 @@
                 // Only fetch badge if the icon is on workspace
                 if (info.id != ItemInfo.NO_ID && badge == null) {
                     badge = appState.getIconCache().getShortcutInfoBadge(si)
-                            .newIcon(context, FLAG_THEMED);
+                            .newIcon(context, ThemeManager.INSTANCE.get(context)
+                                    .isMonoThemeEnabled() ? FLAG_THEMED : 0);
                 }
             }
         } else if (info.itemType == LauncherSettings.Favorites.ITEM_TYPE_FOLDER) {
diff --git a/src/com/android/launcher3/folder/FolderNameProvider.java b/src/com/android/launcher3/folder/FolderNameProvider.java
index be5f8f7..8a1f96d 100644
--- a/src/com/android/launcher3/folder/FolderNameProvider.java
+++ b/src/com/android/launcher3/folder/FolderNameProvider.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.folder;
 
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER;
+
 import android.annotation.SuppressLint;
 import android.app.admin.DevicePolicyManager;
 import android.content.ComponentName;
@@ -37,6 +39,7 @@
 import com.android.launcher3.model.StringCache;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.CollectionInfo;
+import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.Preconditions;
@@ -197,9 +200,18 @@
         @Override
         public void execute(@NonNull ModelTaskController taskController,
                 @NonNull BgDataModel dataModel, @NonNull AllAppsList apps) {
-            mCollectionInfos = dataModel.collections.clone();
+            mCollectionInfos = getCollectionForSuggestions(dataModel);
             mAppInfos = Arrays.asList(apps.copyData());
         }
     }
 
+    public static IntSparseArrayMap<CollectionInfo> getCollectionForSuggestions(
+            BgDataModel dataModel) {
+        IntSparseArrayMap<CollectionInfo> result = new IntSparseArrayMap<>();
+        dataModel.itemsIdMap.stream()
+                .filter(item -> item.itemType == ITEM_TYPE_FOLDER)
+                .forEach(item -> result.put(item.id, (FolderInfo) item));
+        return result;
+    }
+
 }
diff --git a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
index f0e4fc4..a0b73ae 100644
--- a/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
+++ b/src/com/android/launcher3/graphics/LauncherPreviewRenderer.java
@@ -24,10 +24,10 @@
 import static com.android.launcher3.BubbleTextView.DISPLAY_WORKSPACE;
 import static com.android.launcher3.DeviceProfile.DEFAULT_SCALE;
 import static com.android.launcher3.Hotseat.ALPHA_CHANNEL_PREVIEW_RENDERER;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT_PREDICTION;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
-import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
-import static com.android.launcher3.model.ModelUtils.getMissingHotseatRanks;
+import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
 
 import android.app.Fragment;
 import android.app.WallpaperColors;
@@ -105,7 +105,9 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Map.Entry;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 /**
  * Utility class for generating the preview of Launcher for a given InvariantDeviceProfile.
@@ -456,54 +458,48 @@
 
     private void populate(BgDataModel dataModel,
             Map<ComponentKey, AppWidgetProviderInfo> widgetProviderInfoMap) {
-        // Separate the items that are on the current screen, and the other remaining items.
-        ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
-        ArrayList<ItemInfo> otherWorkspaceItems = new ArrayList<>();
-        ArrayList<LauncherAppWidgetInfo> currentAppWidgets = new ArrayList<>();
-        ArrayList<LauncherAppWidgetInfo> otherAppWidgets = new ArrayList<>();
+        IntSet missingHotseatRank = new IntSet();
+        IntStream.range(0, mDp.numShownHotseatIcons).forEach(missingHotseatRank::add);
 
-        IntSet currentScreenIds = IntSet.wrap(mWorkspaceScreens.keySet());
-        filterCurrentWorkspaceItems(currentScreenIds, dataModel.workspaceItems,
-                currentWorkspaceItems, otherWorkspaceItems);
-        filterCurrentWorkspaceItems(currentScreenIds, dataModel.appWidgets, currentAppWidgets,
-                otherAppWidgets);
-        for (ItemInfo itemInfo : currentWorkspaceItems) {
-            switch (itemInfo.itemType) {
-                case Favorites.ITEM_TYPE_APPLICATION:
-                case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
-                    inflateAndAddIcon((WorkspaceItemInfo) itemInfo);
-                    break;
-                case Favorites.ITEM_TYPE_FOLDER:
-                case Favorites.ITEM_TYPE_APP_PAIR:
-                    inflateAndAddCollectionIcon((CollectionInfo) itemInfo);
-                    break;
-                default:
-                    break;
-            }
-        }
-        Map<ComponentKey, AppWidgetProviderInfo> widgetsMap = widgetProviderInfoMap;
-        for (ItemInfo itemInfo : currentAppWidgets) {
-            switch (itemInfo.itemType) {
-                case Favorites.ITEM_TYPE_APPWIDGET:
-                case Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                    if (widgetsMap == null) {
-                        widgetsMap = dataModel.widgetsModel.getWidgetsByComponentKey()
-                                .entrySet()
-                                .stream()
-                                .filter(entry -> entry.getValue().widgetInfo != null)
-                                .collect(Collectors.toMap(
-                                        Map.Entry::getKey,
-                                        entry -> entry.getValue().widgetInfo
-                                ));
+        Map<ComponentKey, AppWidgetProviderInfo>[] widgetsMap = new Map[] { widgetProviderInfoMap};
+
+        // Separate the items that are on the current screen, and the other remaining items.
+        dataModel.itemsIdMap.stream()
+                .filter(currentScreenContentFilter(IntSet.wrap(mWorkspaceScreens.keySet())))
+                .forEach(itemInfo -> {
+                    switch (itemInfo.itemType) {
+                        case Favorites.ITEM_TYPE_APPLICATION:
+                        case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
+                            inflateAndAddIcon((WorkspaceItemInfo) itemInfo);
+                            break;
+                        case Favorites.ITEM_TYPE_FOLDER:
+                        case Favorites.ITEM_TYPE_APP_PAIR:
+                            inflateAndAddCollectionIcon((CollectionInfo) itemInfo);
+                            break;
+                        case Favorites.ITEM_TYPE_APPWIDGET:
+                        case Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
+                            if (widgetsMap[0] == null) {
+                                widgetsMap[0] = dataModel.widgetsModel.getWidgetsByComponentKey()
+                                        .entrySet()
+                                        .stream()
+                                        .filter(entry -> entry.getValue().widgetInfo != null)
+                                        .collect(Collectors.toMap(
+                                                Entry::getKey,
+                                                entry -> entry.getValue().widgetInfo
+                                        ));
+                            }
+                            inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo, widgetsMap[0]);
+                            break;
+                        default:
+                            break;
                     }
-                    inflateAndAddWidgets((LauncherAppWidgetInfo) itemInfo, widgetsMap);
-                    break;
-                default:
-                    break;
-            }
-        }
-        IntArray ranks = getMissingHotseatRanks(currentWorkspaceItems,
-                mDp.numShownHotseatIcons);
+
+                    if (itemInfo.container == CONTAINER_HOTSEAT) {
+                        missingHotseatRank.remove(itemInfo.screenId);
+                    }
+                });
+
+        IntArray ranks = missingHotseatRank.getArray();
         FixedContainerItems hotseatPredictions =
                 dataModel.extraItems.get(CONTAINER_HOTSEAT_PREDICTION);
         List<ItemInfo> predictions = hotseatPredictions == null
diff --git a/src/com/android/launcher3/model/BaseLauncherBinder.java b/src/com/android/launcher3/model/BaseLauncherBinder.java
index c251114..de74ae8 100644
--- a/src/com/android/launcher3/model/BaseLauncherBinder.java
+++ b/src/com/android/launcher3/model/BaseLauncherBinder.java
@@ -19,8 +19,10 @@
 import static com.android.launcher3.BuildConfig.WIDGETS_ENABLED;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
 import static com.android.launcher3.Flags.enableWorkspaceInflation;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
 import static com.android.launcher3.model.ItemInstallQueue.FLAG_LOADER_RUNNING;
-import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
 import static com.android.launcher3.util.Executors.MAIN_EXECUTOR;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 
@@ -44,11 +46,11 @@
 import com.android.launcher3.model.BgDataModel.FixedContainerItems;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.PackageItemInfo;
 import com.android.launcher3.util.ComponentKey;
 import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
+import com.android.launcher3.util.IntSparseArrayMap;
 import com.android.launcher3.util.ItemInflater;
 import com.android.launcher3.util.LooperExecutor;
 import com.android.launcher3.util.LooperIdleLock;
@@ -102,14 +104,12 @@
         Trace.beginSection("BaseLauncherBinder#bindWorkspace");
         try {
             // Save a copy of all the bg-thread collections
-            ArrayList<ItemInfo> workspaceItems = new ArrayList<>();
-            ArrayList<LauncherAppWidgetInfo> appWidgets = new ArrayList<>();
+            IntSparseArrayMap<ItemInfo> itemsIdMap;
             final IntArray orderedScreenIds = new IntArray();
             ArrayList<FixedContainerItems> extraItems = new ArrayList<>();
             final int workspaceItemCount;
             synchronized (mBgDataModel) {
-                workspaceItems.addAll(mBgDataModel.workspaceItems);
-                appWidgets.addAll(mBgDataModel.appWidgets);
+                itemsIdMap = mBgDataModel.itemsIdMap.clone();
                 orderedScreenIds.addAll(mBgDataModel.collectWorkspaceScreens());
                 mBgDataModel.extraItems.forEach(extraItems::add);
                 if (incrementBindId) {
@@ -122,7 +122,7 @@
 
             for (Callbacks cb : mCallbacksList) {
                 new UnifiedWorkspaceBinder(cb, mUiExecutor, mApp, mBgDataModel, mMyBindingId,
-                        workspaceItems, appWidgets, extraItems, orderedScreenIds)
+                        itemsIdMap, extraItems, orderedScreenIds)
                         .bind(isBindSync, workspaceItemCount);
             }
         } finally {
@@ -258,8 +258,7 @@
         private final BgDataModel mBgDataModel;
 
         private final int mMyBindingId;
-        private final ArrayList<ItemInfo> mWorkspaceItems;
-        private final ArrayList<LauncherAppWidgetInfo> mAppWidgets;
+        private final IntSparseArrayMap<ItemInfo> mItemIdMap;
         private final IntArray mOrderedScreenIds;
         private final ArrayList<FixedContainerItems> mExtraItems;
 
@@ -268,8 +267,7 @@
                 LauncherAppState app,
                 BgDataModel bgDataModel,
                 int myBindingId,
-                ArrayList<ItemInfo> workspaceItems,
-                ArrayList<LauncherAppWidgetInfo> appWidgets,
+                IntSparseArrayMap<ItemInfo> itemIdMap,
                 ArrayList<FixedContainerItems> extraItems,
                 IntArray orderedScreenIds) {
             mCallbacks = callbacks;
@@ -277,8 +275,7 @@
             mApp = app;
             mBgDataModel = bgDataModel;
             mMyBindingId = myBindingId;
-            mWorkspaceItems = workspaceItems;
-            mAppWidgets = appWidgets;
+            mItemIdMap = itemIdMap;
             mExtraItems = extraItems;
             mOrderedScreenIds = orderedScreenIds;
         }
@@ -294,10 +291,15 @@
             ArrayList<ItemInfo> currentAppWidgets = new ArrayList<>();
             ArrayList<ItemInfo> otherAppWidgets = new ArrayList<>();
 
-            filterCurrentWorkspaceItems(currentScreenIds, mWorkspaceItems, currentWorkspaceItems,
-                    otherWorkspaceItems);
-            filterCurrentWorkspaceItems(currentScreenIds, mAppWidgets, currentAppWidgets,
-                    otherAppWidgets);
+            Predicate<ItemInfo> currentScreenCheck = currentScreenContentFilter(currentScreenIds);
+            mItemIdMap.forEach(item -> {
+                if (currentScreenCheck.test(item)) {
+                    (WIDGET_FILTER.test(item) ? currentAppWidgets : currentWorkspaceItems)
+                            .add(item);
+                } else if (item.container == CONTAINER_DESKTOP) {
+                    (WIDGET_FILTER.test(item) ? otherAppWidgets : otherWorkspaceItems).add(item);
+                }
+            });
             final InvariantDeviceProfile idp = mApp.getInvariantDeviceProfile();
             sortWorkspaceItemsSpatially(idp, currentWorkspaceItems);
             sortWorkspaceItemsSpatially(idp, otherWorkspaceItems);
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index b9b1e98..a04cbfb 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -20,6 +20,11 @@
 import static com.android.launcher3.BuildConfig.QSB_ON_FIRST_SCREEN;
 import static com.android.launcher3.BuildConfig.WIDGETS_ENABLED;
 import static com.android.launcher3.Flags.enableSmartspaceRemovalToggle;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.shortcuts.ShortcutRequest.PINNED;
 
@@ -31,7 +36,6 @@
 import android.content.pm.ShortcutInfo;
 import android.os.UserHandle;
 import android.text.TextUtils;
-import android.util.ArraySet;
 import android.util.Log;
 import android.util.Pair;
 import android.view.View;
@@ -39,14 +43,11 @@
 import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.LauncherSettings.Favorites;
+import com.android.launcher3.BuildConfig;
 import com.android.launcher3.Workspace;
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.model.data.AppInfo;
-import com.android.launcher3.model.data.AppPairInfo;
 import com.android.launcher3.model.data.CollectionInfo;
-import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.LauncherAppWidgetInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -93,22 +94,6 @@
     public final IntSparseArrayMap<ItemInfo> itemsIdMap = new IntSparseArrayMap<>();
 
     /**
-     * List of all the folders and shortcuts directly on the home screen (no widgets
-     * or shortcuts within folders).
-     */
-    public final ArrayList<ItemInfo> workspaceItems = new ArrayList<>();
-
-    /**
-     * All LauncherAppWidgetInfo created by LauncherModel.
-     */
-    public final ArrayList<LauncherAppWidgetInfo> appWidgets = new ArrayList<>();
-
-    /**
-     * Map of id to CollectionInfos of all the folders or app pairs created by LauncherModel
-     */
-    public final IntSparseArrayMap<CollectionInfo> collections = new IntSparseArrayMap<>();
-
-    /**
      * Extra container based items
      */
     public final IntSparseArrayMap<FixedContainerItems> extraItems = new IntSparseArrayMap<>();
@@ -144,9 +129,6 @@
      * Clears all the data
      */
     public synchronized void clear() {
-        workspaceItems.clear();
-        appWidgets.clear();
-        collections.clear();
         itemsIdMap.clear();
         deepShortcutMap.clear();
         extraItems.clear();
@@ -158,7 +140,7 @@
     public synchronized IntArray collectWorkspaceScreens() {
         IntSet screenSet = new IntSet();
         for (ItemInfo item: itemsIdMap) {
-            if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
+            if (item.container == CONTAINER_DESKTOP) {
                 screenSet.add(item.screenId);
             }
         }
@@ -173,26 +155,14 @@
     public synchronized void dump(String prefix, FileDescriptor fd, PrintWriter writer,
             String[] args) {
         writer.println(prefix + "Data Model:");
-        writer.println(prefix + " ---- workspace items ");
-        for (int i = 0; i < workspaceItems.size(); i++) {
-            writer.println(prefix + '\t' + workspaceItems.get(i).toString());
-        }
-        writer.println(prefix + " ---- appwidget items ");
-        for (int i = 0; i < appWidgets.size(); i++) {
-            writer.println(prefix + '\t' + appWidgets.get(i).toString());
-        }
-        writer.println(prefix + " ---- collection items ");
-        for (int i = 0; i < collections.size(); i++) {
-            writer.println(prefix + '\t' + collections.valueAt(i).toString());
+        writer.println(prefix + " ---- items id map ");
+        for (int i = 0; i < itemsIdMap.size(); i++) {
+            writer.println(prefix + '\t' + itemsIdMap.valueAt(i).toString());
         }
         writer.println(prefix + " ---- extra items ");
         for (int i = 0; i < extraItems.size(); i++) {
             writer.println(prefix + '\t' + extraItems.valueAt(i).toString());
         }
-        writer.println(prefix + " ---- items id map ");
-        for (int i = 0; i < itemsIdMap.size(); i++) {
-            writer.println(prefix + '\t' + itemsIdMap.valueAt(i).toString());
-        }
 
         if (args.length > 0 && TextUtils.equals(args[0], "--all")) {
             writer.println(prefix + "shortcut counts ");
@@ -207,94 +177,38 @@
         removeItem(context, Arrays.asList(items));
     }
 
-    public synchronized void removeItem(Context context, Iterable<? extends ItemInfo> items) {
-        ArraySet<UserHandle> updatedDeepShortcuts = new ArraySet<>();
-        for (ItemInfo item : items) {
-            switch (item.itemType) {
-                case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
-                case LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR:
-                    collections.remove(item.id);
-                    if (FeatureFlags.IS_STUDIO_BUILD) {
-                        for (ItemInfo info : itemsIdMap) {
-                            if (info.container == item.id) {
-                                // We are deleting a collection which still contains items that
-                                // think they are contained by that collection.
-                                String msg = "deleting a collection (" + item + ") which still "
-                                        + "contains items (" + info + ")";
-                                Log.e(TAG, msg);
-                            }
-                        }
-                    }
-                    workspaceItems.remove(item);
-                    break;
-                case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT: {
-                    updatedDeepShortcuts.add(item.user);
-                    // Fall through.
-                }
-                case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
-                    workspaceItems.remove(item);
-                    break;
-                case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
-                case LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                    appWidgets.remove(item);
-                    break;
-            }
-            itemsIdMap.remove(item.id);
+    public synchronized void removeItem(Context context, List<? extends ItemInfo> items) {
+        if (BuildConfig.IS_STUDIO_BUILD) {
+            items.stream()
+                    .filter(item -> item.itemType == ITEM_TYPE_FOLDER
+                            || item.itemType == ITEM_TYPE_APP_PAIR)
+                    .forEach(item -> itemsIdMap.stream()
+                            .filter(info -> info.container == item.id)
+                            // We are deleting a collection which still contains items that
+                            // think they are contained by that collection.
+                            .forEach(info -> Log.e(TAG,
+                                    "deleting a collection (" + item + ") which still contains"
+                                            + " items (" + info + ")")));
         }
-        updatedDeepShortcuts.forEach(user -> updateShortcutPinnedState(context, user));
+
+        items.forEach(item -> itemsIdMap.remove(item.id));
+        items.stream().map(info -> info.user).distinct().forEach(
+                user -> updateShortcutPinnedState(context, user));
     }
 
     public synchronized void addItem(Context context, ItemInfo item, boolean newItem) {
-        addItem(context, item, newItem, null);
-    }
-
-    public synchronized void addItem(
-            Context context, ItemInfo item, boolean newItem, @Nullable LoaderMemoryLogger logger) {
-        if (logger != null) {
-            logger.addLog(
-                    Log.DEBUG,
-                    TAG,
-                    String.format("Adding item to ID map: %s", item.toString()),
-                    /* stackTrace= */ null);
-        }
         itemsIdMap.put(item.id, item);
-        switch (item.itemType) {
-            case LauncherSettings.Favorites.ITEM_TYPE_FOLDER:
-                collections.put(item.id, (FolderInfo) item);
-                workspaceItems.add(item);
-                break;
-            case LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR:
-                collections.put(item.id, (AppPairInfo) item);
-                // Fall through here. App pairs are both containers (like folders) and containable
-                // items (can be placed in folders). So we need to add app pairs to the folders
-                // array (above) but also verify the existence of their container, like regular
-                // apps (below).
-            case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT:
-            case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
-                if (item.container == LauncherSettings.Favorites.CONTAINER_DESKTOP ||
-                        item.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
-                    workspaceItems.add(item);
-                } else {
-                    if (newItem) {
-                        if (!collections.containsKey(item.container)) {
-                            // Adding an item to a nonexistent collection.
-                            String msg = "attempted to add item: " + item + " to a nonexistent app"
-                                    + " collection";
-                            Log.e(TAG, msg);
-                        }
-                    } else {
-                        findOrMakeFolder(item.container).add(item);
-                    }
-                }
-                break;
-            case LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET:
-            case LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET:
-                appWidgets.add((LauncherAppWidgetInfo) item);
-                break;
-        }
-        if (newItem && item.itemType == LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT) {
+        if (newItem && item.itemType == ITEM_TYPE_DEEP_SHORTCUT) {
             updateShortcutPinnedState(context, item.user);
         }
+        if (BuildConfig.IS_DEBUG_DEVICE
+                && newItem
+                && item.container != CONTAINER_DESKTOP
+                && item.container != CONTAINER_HOTSEAT
+                && !(itemsIdMap.get(item.container) instanceof CollectionInfo)) {
+            // Adding an item to a nonexistent collection.
+            Log.e(TAG, "attempted to add item: " + item + " to a nonexistent app collection");
+        }
     }
 
     /**
@@ -334,7 +248,7 @@
         Map<String, Set<String>> modelMap = Stream.concat(
                     // Model shortcuts
                     itemStream.build()
-                        .filter(wi -> wi.itemType == Favorites.ITEM_TYPE_DEEP_SHORTCUT)
+                        .filter(wi -> wi.itemType == ITEM_TYPE_DEEP_SHORTCUT)
                         .map(ShortcutKey::fromItemInfo),
                     // Pending shortcuts
                     ItemInstallQueue.INSTANCE.get(context).getPendingShortcuts(user))
@@ -375,24 +289,6 @@
     }
 
     /**
-     * Return an existing FolderInfo object if we have encountered this ID previously,
-     * or make a new one.
-     */
-    public synchronized CollectionInfo findOrMakeFolder(int id) {
-        // See if a placeholder was created for us already
-        CollectionInfo collectionInfo = collections.get(id);
-        if (collectionInfo == null) {
-            // No placeholder -- create a new blank folder instance. At this point, we don't know
-            // if the desired container is supposed to be a folder or an app pair. In the case that
-            // it is an app pair, the blank folder will be replaced by a blank app pair when the app
-            // pair is getting processed, in WorkspaceItemProcessor.processFolderOrAppPair().
-            collectionInfo = new FolderInfo();
-            collections.put(id, collectionInfo);
-        }
-        return collectionInfo;
-    }
-
-    /**
      * Clear all the deep shortcut counts for the given package, and re-add the new shortcut counts.
      */
     public synchronized void updateDeepShortcutCounts(
@@ -424,16 +320,6 @@
     }
 
     /**
-     * Returns a list containing all workspace items including widgets.
-     */
-    public synchronized ArrayList<ItemInfo> getAllWorkspaceItems() {
-        ArrayList<ItemInfo> items = new ArrayList<>(workspaceItems.size() + appWidgets.size());
-        items.addAll(workspaceItems);
-        items.addAll(appWidgets);
-        return items;
-    }
-
-    /**
      * Calls the provided {@code op} for all workspaceItems in the in-memory model (both persisted
      * items and dynamic/predicted items for the provided {@code userHandle}.
      * Note the call is not synchronized over the model, that should be handled by the called.
diff --git a/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt b/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt
index aa62c32..6ad52ea 100644
--- a/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt
+++ b/src/com/android/launcher3/model/FirstScreenBroadcastHelper.kt
@@ -30,7 +30,6 @@
 import com.android.launcher3.model.data.ItemInfo
 import com.android.launcher3.model.data.LauncherAppWidgetInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
-import com.android.launcher3.pm.InstallSessionHelper
 import com.android.launcher3.util.Executors
 import com.android.launcher3.util.PackageManagerHelper
 import com.android.launcher3.util.PackageUserKey
@@ -80,21 +79,22 @@
         packageManagerHelper: PackageManagerHelper,
         firstScreenItems: List<ItemInfo>,
         userKeyToSessionMap: Map<PackageUserKey, SessionInfo>,
-        allWidgets: List<LauncherAppWidgetInfo>
+        allWidgets: List<ItemInfo>,
     ): List<FirstScreenBroadcastModel> {
 
         // installers for installing items
-        val pendingItemInstallerMap: Map<String, MutableSet<String>> =
+        val pendingItemInstallerMap: Map<String, Set<String>> =
             createPendingItemsMap(userKeyToSessionMap)
+
         val installingPackages = pendingItemInstallerMap.values.flatten().toSet()
 
         // installers for installed items on first screen
-        val installedItemInstallerMap: Map<String, MutableSet<ItemInfo>> =
+        val installedItemInstallerMap: Map<String, List<ItemInfo>> =
             createInstalledItemsMap(firstScreenItems, installingPackages, packageManagerHelper)
 
         // installers for widgets on all screens
-        val allInstalledWidgetsMap: Map<String, MutableSet<LauncherAppWidgetInfo>> =
-            createAllInstalledWidgetsMap(allWidgets, installingPackages, packageManagerHelper)
+        val allInstalledWidgetsMap: Map<String, List<ItemInfo>> =
+            createInstalledItemsMap(allWidgets, installingPackages, packageManagerHelper)
 
         val allInstallers: Set<String> =
             pendingItemInstallerMap.keys +
@@ -131,39 +131,39 @@
                             context,
                             0 /* requestCode */,
                             Intent(),
-                            PendingIntent.FLAG_ONE_SHOT or PendingIntent.FLAG_IMMUTABLE
-                        )
+                            PendingIntent.FLAG_ONE_SHOT or PendingIntent.FLAG_IMMUTABLE,
+                        ),
                     )
                     .putStringArrayListExtra(
                         PENDING_COLLECTION_ITEM_EXTRA,
-                        ArrayList(model.pendingCollectionItems)
+                        ArrayList(model.pendingCollectionItems),
                     )
                     .putStringArrayListExtra(
                         PENDING_WORKSPACE_ITEM_EXTRA,
-                        ArrayList(model.pendingWorkspaceItems)
+                        ArrayList(model.pendingWorkspaceItems),
                     )
                     .putStringArrayListExtra(
                         PENDING_HOTSEAT_ITEM_EXTRA,
-                        ArrayList(model.pendingHotseatItems)
+                        ArrayList(model.pendingHotseatItems),
                     )
                     .putStringArrayListExtra(
                         PENDING_WIDGET_ITEM_EXTRA,
-                        ArrayList(model.pendingWidgetItems)
+                        ArrayList(model.pendingWidgetItems),
                     )
                     .putStringArrayListExtra(
                         INSTALLED_WORKSPACE_ITEMS_EXTRA,
-                        ArrayList(model.installedWorkspaceItems)
+                        ArrayList(model.installedWorkspaceItems),
                     )
                     .putStringArrayListExtra(
                         INSTALLED_HOTSEAT_ITEMS_EXTRA,
-                        ArrayList(model.installedHotseatItems)
+                        ArrayList(model.installedHotseatItems),
                     )
                     .putStringArrayListExtra(
                         ALL_INSTALLED_WIDGETS_ITEM_EXTRA,
                         ArrayList(
                             model.firstScreenInstalledWidgets +
                                 model.secondaryScreenInstalledWidgets
-                        )
+                        ),
                     )
             context.sendBroadcast(intent)
         }
@@ -172,66 +172,46 @@
     /** Maps Installer packages to Set of app packages from install sessions */
     private fun createPendingItemsMap(
         userKeyToSessionMap: Map<PackageUserKey, SessionInfo>
-    ): Map<String, MutableSet<String>> {
+    ): Map<String, Set<String>> {
         val myUser = Process.myUserHandle()
-        val result = mutableMapOf<String, MutableSet<String>>()
-        userKeyToSessionMap.forEach { entry ->
-            if (!myUser.equals(InstallSessionHelper.getUserHandle(entry.value))) return@forEach
-            val installer = entry.value.installerPackageName
-            val appPackage = entry.value.appPackageName
-            if (installer.isNullOrEmpty() || appPackage.isNullOrEmpty()) return@forEach
-            result.getOrPut(installer) { mutableSetOf() }.add(appPackage)
-        }
-        return result
-    }
-
-    /**
-     * Maps Installer packages to Set of ItemInfo from first screen. Filter out installing packages.
-     */
-    private fun createInstalledItemsMap(
-        firstScreenItems: List<ItemInfo>,
-        installingPackages: Set<String>,
-        packageManagerHelper: PackageManagerHelper
-    ): Map<String, MutableSet<ItemInfo>> {
-        val result = mutableMapOf<String, MutableSet<ItemInfo>>()
-        firstScreenItems.forEach { item ->
-            val appPackage = getPackageName(item) ?: return@forEach
-            if (installingPackages.contains(appPackage)) return@forEach
-            val installer = packageManagerHelper.getAppInstallerPackage(appPackage)
-            if (installer.isNullOrEmpty()) return@forEach
-            result.getOrPut(installer) { mutableSetOf() }.add(item)
-        }
-        return result
-    }
-
-    /**
-     * Maps Installer packages to Set of AppWidget packages installed on all screens. Filter out
-     * installing packages.
-     */
-    private fun createAllInstalledWidgetsMap(
-        allWidgets: List<LauncherAppWidgetInfo>,
-        installingPackages: Set<String>,
-        packageManagerHelper: PackageManagerHelper
-    ): Map<String, MutableSet<LauncherAppWidgetInfo>> {
-        val result = mutableMapOf<String, MutableSet<LauncherAppWidgetInfo>>()
-        allWidgets
-            .sortedBy { widget -> widget.screenId }
-            .forEach { widget ->
-                val appPackage = getPackageName(widget) ?: return@forEach
-                if (installingPackages.contains(appPackage)) return@forEach
-                val installer = packageManagerHelper.getAppInstallerPackage(appPackage)
-                if (installer.isNullOrEmpty()) return@forEach
-                result.getOrPut(installer) { mutableSetOf() }.add(widget)
+        return userKeyToSessionMap.values
+            .filter {
+                it.user == myUser &&
+                    !it.installerPackageName.isNullOrEmpty() &&
+                    !it.appPackageName.isNullOrEmpty()
             }
-        return result
+            .groupBy(
+                keySelector = { it.installerPackageName },
+                valueTransform = { it.appPackageName },
+            )
+            .mapValues { it.value.filterNotNull().toSet() } as Map<String, Set<String>>
     }
 
+    /** Maps Installer packages to Set of ItemInfos. Filter out installing packages. */
+    private fun createInstalledItemsMap(
+        allItems: Iterable<ItemInfo>,
+        installingPackages: Set<String>,
+        packageManagerHelper: PackageManagerHelper,
+    ): Map<String, List<ItemInfo>> =
+        allItems
+            .sortedBy { it.screenId }
+            .groupByTo(mutableMapOf()) {
+                getPackageName(it)?.let { pkg ->
+                    if (installingPackages.contains(pkg)) {
+                        null
+                    } else {
+                        packageManagerHelper.getAppInstallerPackage(pkg)
+                    }
+                }
+            }
+            .apply { remove(null) } as Map<String, List<ItemInfo>>
+
     /**
      * Add first screen Pending Items from Map to [FirstScreenBroadcastModel] for given installer
      */
     private fun FirstScreenBroadcastModel.addPendingItems(
         installingItems: Set<String>?,
-        firstScreenItems: List<ItemInfo>
+        firstScreenItems: List<ItemInfo>,
     ) {
         if (installingItems == null) return
         for (info in firstScreenItems) {
@@ -251,7 +231,7 @@
      */
     private fun FirstScreenBroadcastModel.addInstalledItems(
         installer: String,
-        installedItemInstallerMap: Map<String, Set<ItemInfo>>,
+        installedItemInstallerMap: Map<String, List<ItemInfo>>,
     ) {
         installedItemInstallerMap[installer]?.forEach { info ->
             val packageName: String = getPackageName(info) ?: return@forEach
@@ -265,7 +245,7 @@
     /** Add Widgets on every screen from Map to [FirstScreenBroadcastModel] for given installer */
     private fun FirstScreenBroadcastModel.addAllScreenWidgets(
         installer: String,
-        allInstalledWidgetsMap: Map<String, Set<LauncherAppWidgetInfo>>
+        allInstalledWidgetsMap: Map<String, List<ItemInfo>>,
     ) {
         allInstalledWidgetsMap[installer]?.forEach { widget ->
             val packageName: String = getPackageName(widget) ?: return@forEach
@@ -279,7 +259,7 @@
 
     private fun FirstScreenBroadcastModel.addCollectionItems(
         info: ItemInfo,
-        installingPackages: Set<String>
+        installingPackages: Set<String>,
     ) {
         if (info !is CollectionInfo) return
         pendingCollectionItems.addAll(
@@ -336,7 +316,7 @@
             Log.d(
                 TAG,
                 "Sending First Screen Broadcast for installer=$installerPackage" +
-                    ", total packages=${getTotalItemCount()}"
+                    ", total packages=${getTotalItemCount()}",
             )
             pendingCollectionItems.forEach {
                 Log.d(TAG, "$installerPackage:Pending Collection item:$it")
@@ -361,15 +341,7 @@
         }
     }
 
-    private fun getPackageName(info: ItemInfo): String? {
-        var packageName: String? = null
-        if (info is LauncherAppWidgetInfo) {
-            info.providerName?.let { packageName = info.providerName.packageName }
-        } else if (info.targetComponent != null) {
-            packageName = info.targetComponent?.packageName
-        }
-        return packageName
-    }
+    private fun getPackageName(info: ItemInfo): String? = info.targetComponent?.packageName
 
     /**
      * Clone the provided list on UI thread. This is used for [FolderInfo.getContents] which is
diff --git a/src/com/android/launcher3/model/LoaderCursor.java b/src/com/android/launcher3/model/LoaderCursor.java
index 536d4c9..1623881 100644
--- a/src/com/android/launcher3/model/LoaderCursor.java
+++ b/src/com/android/launcher3/model/LoaderCursor.java
@@ -16,6 +16,11 @@
 
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPLICATION;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 import static com.android.launcher3.LauncherSettings.Favorites.TABLE_NAME;
 import static com.android.launcher3.Utilities.SHOULD_SHOW_FIRST_PAGE_WIDGET;
 import static com.android.launcher3.icons.cache.CacheLookupFlag.DEFAULT_LOOKUP_FLAG;
@@ -48,6 +53,8 @@
 import com.android.launcher3.icons.IconCache;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
+import com.android.launcher3.model.data.CollectionInfo;
+import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
@@ -84,6 +91,11 @@
     private final IntArray mRestoredRows = new IntArray();
     private final IntSparseArrayMap<GridOccupancy> mOccupied = new IntSparseArrayMap<>();
 
+    // CollectionInfo objects, which have not yet been loaded from the DB, but are expected to
+    // found eventually as the loading progresses
+    private final IntSparseArrayMap<CollectionInfo> mPendingCollectionInfo =
+            new IntSparseArrayMap<>();
+
     private final int mIconIndex;
     public final int mTitleIndex;
 
@@ -479,8 +491,26 @@
         info.cellY = getInt(mCellYIndex);
     }
 
-    public void checkAndAddItem(ItemInfo info, BgDataModel dataModel) {
-        checkAndAddItem(info, dataModel, null);
+    /**
+     * Return an existing FolderInfo object if we have encountered this ID previously,
+     * or make a new one.
+     */
+    public CollectionInfo findOrMakeFolder(int id, BgDataModel dataModel) {
+        // See if a placeholder was created for us already
+        ItemInfo info = dataModel.itemsIdMap.get(id);
+        if (info instanceof CollectionInfo c) return c;
+
+        CollectionInfo pending = mPendingCollectionInfo.get(id);
+        if (pending != null) return pending;
+
+        // No placeholder -- create a new blank folder instance. At this point, we don't know
+        // if the desired container is supposed to be a folder or an app pair. In the case that
+        // it is an app pair, the blank folder will be replaced by a blank app pair when the app
+        // pair is getting processed, in WorkspaceItemProcessor.processFolderOrAppPair().
+        pending = new FolderInfo();
+        pending.id = id;
+        mPendingCollectionInfo.put(id, pending);
+        return pending;
     }
 
     /**
@@ -495,7 +525,21 @@
             ShortcutKey.fromItemInfo(info);
         }
         if (checkItemPlacement(info, dataModel.isFirstPagePinnedItemEnabled)) {
-            dataModel.addItem(mContext, info, false, logger);
+            if (logger != null) {
+                logger.addLog(
+                        Log.DEBUG,
+                        TAG,
+                        String.format("Adding item to ID map: %s", info),
+                        /* stackTrace= */ null);
+            }
+            dataModel.addItem(mContext, info, false);
+            if ((info.itemType == ITEM_TYPE_APP_PAIR
+                    || info.itemType == ITEM_TYPE_DEEP_SHORTCUT
+                    || info.itemType == ITEM_TYPE_APPLICATION)
+                    && info.container != CONTAINER_DESKTOP
+                    && info.container != CONTAINER_HOTSEAT) {
+                findOrMakeFolder(info.container, dataModel).add(info);
+            }
             if (mRestoreEventLogger != null) {
                 mRestoreEventLogger.logSingleFavoritesItemRestored(itemType);
             }
diff --git a/src/com/android/launcher3/model/LoaderTask.java b/src/com/android/launcher3/model/LoaderTask.java
index 44b7e8b..fee9696 100644
--- a/src/com/android/launcher3/model/LoaderTask.java
+++ b/src/com/android/launcher3/model/LoaderTask.java
@@ -31,7 +31,8 @@
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_CHANGE_PERMISSION;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_WORK_PROFILE_QUIET_MODE_ENABLED;
-import static com.android.launcher3.model.ModelUtils.filterCurrentWorkspaceItems;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+import static com.android.launcher3.model.ModelUtils.currentScreenContentFilter;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_INSTALL_SESSION_ACTIVE;
 import static com.android.launcher3.util.Executors.MODEL_EXECUTOR;
 import static com.android.launcher3.util.PackageManagerHelper.hasShortcutsPermission;
@@ -82,7 +83,6 @@
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.AppPairInfo;
-import com.android.launcher3.model.data.CollectionInfo;
 import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.IconRequestInfo;
 import com.android.launcher3.model.data.ItemInfo;
@@ -210,10 +210,10 @@
         final int firstScreen = allScreens.get(0);
         IntSet firstScreens = IntSet.wrap(firstScreen);
 
-        ArrayList<ItemInfo> allItems = mBgDataModel.getAllWorkspaceItems();
-        ArrayList<ItemInfo> firstScreenItems = new ArrayList<>();
-        filterCurrentWorkspaceItems(firstScreens, allItems, firstScreenItems,
-                new ArrayList<>() /* otherScreenItems are ignored */);
+        List<ItemInfo> firstScreenItems =
+                mBgDataModel.itemsIdMap.stream()
+                        .filter(currentScreenContentFilter(firstScreens))
+                        .toList();
         final int disableArchivingLauncherBroadcast = Settings.Secure.getInt(
                 mApp.getContext().getContentResolver(),
                 "disable_launcher_broadcast_installed_apps",
@@ -227,7 +227,7 @@
                             mPmHelper,
                             firstScreenItems,
                             mInstallingPkgsCached,
-                            mBgDataModel.appWidgets
+                            mBgDataModel.itemsIdMap.stream().filter(WIDGET_FILTER).toList()
                     );
             logASplit("Sending first screen broadcast with additional archiving Extras");
             FirstScreenBroadcastHelper.sendBroadcastsForModels(mApp.getContext(), broadcastModels);
@@ -523,14 +523,13 @@
      * requests high-res icons for the items that are part of an app pair.
      */
     private void processAppPairItems() {
-        for (CollectionInfo collection : mBgDataModel.collections) {
-            if (!(collection instanceof AppPairInfo appPair)) {
-                continue;
-            }
-
-            appPair.getContents().sort(Folder.ITEM_POS_COMPARATOR);
-            appPair.fetchHiResIconsIfNeeded(mIconCache);
-        }
+        mBgDataModel.itemsIdMap.stream()
+                .filter(item -> item instanceof AppPairInfo)
+                .forEach(item -> {
+                    AppPairInfo appPair = (AppPairInfo) item;
+                    appPair.getContents().sort(Folder.ITEM_POS_COMPARATOR);
+                    appPair.fetchHiResIconsIfNeeded(mIconCache);
+                });
     }
 
     /**
@@ -586,8 +585,8 @@
         // Sort the folder items, update ranks, and make sure all preview items are high res.
         List<FolderGridOrganizer> verifiers = mApp.getInvariantDeviceProfile().supportedProfiles
                 .stream().map(FolderGridOrganizer::createFolderGridOrganizer).toList();
-        for (CollectionInfo collection : mBgDataModel.collections) {
-            if (!(collection instanceof FolderInfo folder)) {
+        for (ItemInfo itemInfo : mBgDataModel.itemsIdMap) {
+            if (!(itemInfo instanceof FolderInfo folder)) {
                 continue;
             }
 
@@ -657,8 +656,6 @@
             IntArray deletedFolderIds = mApp.getModel().getModelDbController().deleteEmptyFolders();
             synchronized (mBgDataModel) {
                 for (int folderId : deletedFolderIds) {
-                    mBgDataModel.workspaceItems.remove(mBgDataModel.collections.get(folderId));
-                    mBgDataModel.collections.remove(folderId);
                     mBgDataModel.itemsIdMap.remove(folderId);
                 }
             }
@@ -676,8 +673,6 @@
 
         synchronized (mBgDataModel) {
             for (int id : deleted) {
-                mBgDataModel.workspaceItems.remove(mBgDataModel.collections.get(id));
-                mBgDataModel.collections.remove(id);
                 mBgDataModel.itemsIdMap.remove(id);
             }
         }
@@ -819,18 +814,19 @@
 
     private void loadFolderNames() {
         FolderNameProvider provider = FolderNameProvider.newInstance(mApp.getContext(),
-                mBgAllAppsList.data, mBgDataModel.collections);
+                mBgAllAppsList.data, FolderNameProvider.getCollectionForSuggestions(mBgDataModel));
 
         synchronized (mBgDataModel) {
-            for (int i = 0; i < mBgDataModel.collections.size(); i++) {
-                FolderNameInfos suggestionInfos = new FolderNameInfos();
-                CollectionInfo info = mBgDataModel.collections.valueAt(i);
-                if (info instanceof FolderInfo fi && fi.suggestedFolderNames == null) {
-                    provider.getSuggestedFolderName(mApp.getContext(), fi.getAppContents(),
-                            suggestionInfos);
-                    fi.suggestedFolderNames = suggestionInfos;
-                }
-            }
+            mBgDataModel.itemsIdMap.stream()
+                    .filter(item ->
+                            item instanceof FolderInfo fi && fi.suggestedFolderNames == null)
+                    .forEach(info -> {
+                        FolderInfo fi = (FolderInfo) info;
+                        FolderNameInfos suggestionInfos = new FolderNameInfos();
+                        provider.getSuggestedFolderName(mApp.getContext(), fi.getAppContents(),
+                                suggestionInfos);
+                        fi.suggestedFolderNames = suggestionInfos;
+                    });
         }
     }
 
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
index 9e72e28..da79982 100644
--- a/src/com/android/launcher3/model/ModelUtils.java
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -15,15 +15,15 @@
  */
 package com.android.launcher3.model;
 
-import com.android.launcher3.LauncherSettings;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_CUSTOM_APPWIDGET;
+
 import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.util.IntArray;
 import com.android.launcher3.util.IntSet;
 
-import java.util.Collections;
-import java.util.List;
-import java.util.Objects;
-import java.util.stream.IntStream;
+import java.util.function.Predicate;
 
 /**
  * Utils class for {@link com.android.launcher3.LauncherModel}.
@@ -31,54 +31,17 @@
 public class ModelUtils {
 
     /**
-     * Filters the set of items who are directly or indirectly (via another container) on the
-     * specified screen.
+     * Returns a filter for items on hotseat or current screens
      */
-    public static <T extends ItemInfo> void filterCurrentWorkspaceItems(
-            final IntSet currentScreenIds,
-            List<? extends T> allWorkspaceItems,
-            List<T> currentScreenItems,
-            List<T> otherScreenItems) {
-        // Purge any null ItemInfos
-        allWorkspaceItems.removeIf(Objects::isNull);
-        // Order the set of items by their containers first, this allows use to walk through the
-        // list sequentially, build up a list of containers that are in the specified screen,
-        // as well as all items in those containers.
-        IntSet itemsOnScreen = new IntSet();
-        Collections.sort(allWorkspaceItems,
-                (lhs, rhs) -> Integer.compare(lhs.container, rhs.container));
-        for (T info : allWorkspaceItems) {
-            if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-                if (currentScreenIds.contains(info.screenId)) {
-                    currentScreenItems.add(info);
-                    itemsOnScreen.add(info.id);
-                } else {
-                    otherScreenItems.add(info);
-                }
-            } else if (info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT) {
-                currentScreenItems.add(info);
-                itemsOnScreen.add(info.id);
-            } else {
-                if (itemsOnScreen.contains(info.container)) {
-                    currentScreenItems.add(info);
-                    itemsOnScreen.add(info.id);
-                } else {
-                    otherScreenItems.add(info);
-                }
-            }
-        }
+    public static Predicate<ItemInfo> currentScreenContentFilter(IntSet currentScreenIds) {
+        return item -> item.container == CONTAINER_HOTSEAT
+                || (item.container == CONTAINER_DESKTOP
+                        && currentScreenIds.contains(item.screenId));
     }
 
     /**
-     * Iterates though current workspace items and returns available hotseat ranks for prediction.
+     * Returns a filter for widget items
      */
-    public static IntArray getMissingHotseatRanks(List<ItemInfo> items, int len) {
-        IntSet seen = new IntSet();
-        items.stream().filter(
-                info -> info.container == LauncherSettings.Favorites.CONTAINER_HOTSEAT)
-                .forEach(i -> seen.add(i.screenId));
-        IntArray result = new IntArray(len);
-        IntStream.range(0, len).filter(i -> !seen.contains(i)).forEach(result::add);
-        return result;
-    }
+    public static final Predicate<ItemInfo> WIDGET_FILTER = item ->
+            item.itemType == ITEM_TYPE_APPWIDGET || item.itemType == ITEM_TYPE_CUSTOM_APPWIDGET;
 }
diff --git a/src/com/android/launcher3/model/ModelWriter.java b/src/com/android/launcher3/model/ModelWriter.java
index b477cb1..0332775 100644
--- a/src/com/android/launcher3/model/ModelWriter.java
+++ b/src/com/android/launcher3/model/ModelWriter.java
@@ -459,37 +459,14 @@
                 if (item.container != Favorites.CONTAINER_DESKTOP &&
                         item.container != Favorites.CONTAINER_HOTSEAT) {
                     // Item is in a collection, make sure this collection exists
-                    if (!mBgDataModel.collections.containsKey(item.container)) {
+                    if (!(mBgDataModel.itemsIdMap.get(item.container) instanceof CollectionInfo)) {
                         // An items container is being set to a that of an item which is not in
-                        // the list of Folders.
+                        // the list of collections.
                         String msg = "item: " + item + " container being set to: " +
                                 item.container + ", not in the list of collections";
                         Log.e(TAG, msg);
                     }
                 }
-
-                // Items are added/removed from the corresponding FolderInfo elsewhere, such
-                // as in Workspace.onDrop. Here, we just add/remove them from the list of items
-                // that are on the desktop, as appropriate
-                ItemInfo modelItem = mBgDataModel.itemsIdMap.get(itemId);
-                if (modelItem != null &&
-                        (modelItem.container == Favorites.CONTAINER_DESKTOP ||
-                                modelItem.container == Favorites.CONTAINER_HOTSEAT)) {
-                    switch (modelItem.itemType) {
-                        case Favorites.ITEM_TYPE_APPLICATION:
-                        case Favorites.ITEM_TYPE_DEEP_SHORTCUT:
-                        case Favorites.ITEM_TYPE_FOLDER:
-                        case Favorites.ITEM_TYPE_APP_PAIR:
-                            if (!mBgDataModel.workspaceItems.contains(modelItem)) {
-                                mBgDataModel.workspaceItems.add(modelItem);
-                            }
-                            break;
-                        default:
-                            break;
-                    }
-                } else {
-                    mBgDataModel.workspaceItems.remove(modelItem);
-                }
                 mVerifier.verifyModel();
             }
         }
diff --git a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
index d238213..4103937 100644
--- a/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
+++ b/src/com/android/launcher3/model/PackageInstallStateChangedTask.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.model;
 
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
+
 import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageManager;
@@ -85,12 +87,16 @@
                 }
             });
 
-            for (LauncherAppWidgetInfo widget : dataModel.appWidgets) {
-                if (widget.providerName.getPackageName().equals(mInstallInfo.packageName)) {
-                    widget.installProgress = mInstallInfo.progress;
-                    updates.add(widget);
-                }
-            }
+            dataModel.itemsIdMap.stream()
+                    .filter(WIDGET_FILTER)
+                    .filter(item -> mInstallInfo.user.equals(item.user))
+                    .map(item -> (LauncherAppWidgetInfo) item)
+                    .filter(widget -> widget.providerName.getPackageName()
+                            .equals(mInstallInfo.packageName))
+                    .forEach(widget -> {
+                        widget.installProgress = mInstallInfo.progress;
+                        updates.add(widget);
+                    });
 
             if (!updates.isEmpty()) {
                 taskController.scheduleCallbackTask(
diff --git a/src/com/android/launcher3/model/PackageUpdatedTask.java b/src/com/android/launcher3/model/PackageUpdatedTask.java
index d619965..1153f48 100644
--- a/src/com/android/launcher3/model/PackageUpdatedTask.java
+++ b/src/com/android/launcher3/model/PackageUpdatedTask.java
@@ -18,7 +18,9 @@
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_PRIVATE_PROFILE_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_QUIET_MODE_ENABLED;
 import static com.android.launcher3.model.BgDataModel.Callbacks.FLAG_WORK_PROFILE_QUIET_MODE_ENABLED;
+import static com.android.launcher3.model.ModelUtils.WIDGET_FILTER;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED;
+import static com.android.launcher3.model.data.LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY;
 import static com.android.launcher3.model.data.WorkspaceItemInfo.FLAG_AUTOINSTALL_ICON;
 import static com.android.launcher3.model.data.WorkspaceItemInfo.FLAG_RESTORED_ICON;
 
@@ -347,24 +349,25 @@
                     }
                 });
 
-                for (LauncherAppWidgetInfo widgetInfo : dataModel.appWidgets) {
-                    if (mUser.equals(widgetInfo.user)
-                            && widgetInfo.hasRestoreFlag(
-                            LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY)
-                            && packageSet.contains(widgetInfo.providerName.getPackageName())) {
-                        widgetInfo.restoreStatus &=
-                                ~LauncherAppWidgetInfo.FLAG_PROVIDER_NOT_READY
-                                        & ~LauncherAppWidgetInfo.FLAG_RESTORE_STARTED;
+                dataModel.itemsIdMap.stream()
+                        .filter(WIDGET_FILTER)
+                        .filter(item -> mUser.equals(item.user))
+                        .map(item -> (LauncherAppWidgetInfo) item)
+                        .filter(widget -> widget.hasRestoreFlag(FLAG_PROVIDER_NOT_READY)
+                                && packageSet.contains(widget.providerName.getPackageName()))
+                        .forEach(widgetInfo -> {
+                            widgetInfo.restoreStatus &=
+                                    ~FLAG_PROVIDER_NOT_READY
+                                            & ~LauncherAppWidgetInfo.FLAG_RESTORE_STARTED;
 
-                        // adding this flag ensures that launcher shows 'click to setup'
-                        // if the widget has a config activity. In case there is no config
-                        // activity, it will be marked as 'restored' during bind.
-                        widgetInfo.restoreStatus |= LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
+                            // adding this flag ensures that launcher shows 'click to setup'
+                            // if the widget has a config activity. In case there is no config
+                            // activity, it will be marked as 'restored' during bind.
+                            widgetInfo.restoreStatus |= LauncherAppWidgetInfo.FLAG_UI_NOT_READY;
 
-                        widgets.add(widgetInfo);
-                        taskController.getModelWriter().updateItemInDatabase(widgetInfo);
-                    }
-                }
+                            widgets.add(widgetInfo);
+                            taskController.getModelWriter().updateItemInDatabase(widgetInfo);
+                        });
             }
 
             taskController.bindUpdatedWorkspaceItems(updatedWorkspaceItems);
diff --git a/src/com/android/launcher3/model/WorkspaceItemProcessor.kt b/src/com/android/launcher3/model/WorkspaceItemProcessor.kt
index dad78dd..de1df2e 100644
--- a/src/com/android/launcher3/model/WorkspaceItemProcessor.kt
+++ b/src/com/android/launcher3/model/WorkspaceItemProcessor.kt
@@ -404,18 +404,14 @@
      * stored in the BgDataModel.
      */
     private fun processFolderOrAppPair() {
-        var collection = bgDataModel.findOrMakeFolder(c.id)
+        var collection = c.findOrMakeFolder(c.id, bgDataModel)
         // If we generated a placeholder Folder before this point, it may need to be replaced with
         // an app pair.
         if (c.itemType == Favorites.ITEM_TYPE_APP_PAIR && collection is FolderInfo) {
-            val folderInfo: FolderInfo = collection
             val newAppPair = AppPairInfo()
             // Move the placeholder's contents over to the new app pair.
-            folderInfo.getContents().forEach(newAppPair::add)
+            collection.getContents().forEach(newAppPair::add)
             collection = newAppPair
-            // Remove the placeholder and add the app pair into the data model.
-            bgDataModel.collections.remove(c.id)
-            bgDataModel.collections.put(c.id, collection)
         }
 
         c.applyCommonProperties(collection)
@@ -569,7 +565,7 @@
                 logWidgetInfo(app.invariantDeviceProfile, lapi)
             }
         }
-        c.checkAndAddItem(appWidgetInfo, bgDataModel)
+        c.checkAndAddItem(appWidgetInfo, bgDataModel, memoryLogger)
     }
 
     companion object {
diff --git a/src/com/android/launcher3/model/data/ItemInfoWithIcon.java b/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
index 772ea7f..7fb0152 100644
--- a/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
+++ b/src/com/android/launcher3/model/data/ItemInfoWithIcon.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.model.data;
 
+import static com.android.launcher3.icons.BitmapInfo.FLAG_THEMED;
+
 import android.content.Context;
 import android.content.Intent;
 import android.os.Process;
@@ -23,6 +25,7 @@
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.Flags;
+import com.android.launcher3.graphics.ThemeManager;
 import com.android.launcher3.icons.BitmapInfo;
 import com.android.launcher3.icons.BitmapInfo.DrawableCreationFlags;
 import com.android.launcher3.icons.FastBitmapDrawable;
@@ -320,6 +323,9 @@
      * Returns a FastBitmapDrawable with the icon and context theme applied
      */
     public FastBitmapDrawable newIcon(Context context, @DrawableCreationFlags int creationFlags) {
+        if (!ThemeManager.INSTANCE.get(context).isMonoThemeEnabled()) {
+            creationFlags &= ~FLAG_THEMED;
+        }
         FastBitmapDrawable drawable = bitmap.newIcon(context, creationFlags);
         drawable.setIsDisabled(isDisabled());
         return drawable;
diff --git a/src/com/android/launcher3/util/IntSparseArrayMap.java b/src/com/android/launcher3/util/IntSparseArrayMap.java
index 9d5391b..70f74e3 100644
--- a/src/com/android/launcher3/util/IntSparseArrayMap.java
+++ b/src/com/android/launcher3/util/IntSparseArrayMap.java
@@ -19,6 +19,8 @@
 import android.util.SparseArray;
 
 import java.util.Iterator;
+import java.util.stream.Stream;
+import java.util.stream.StreamSupport;
 
 /**
  * Extension of {@link SparseArray} with some utility methods.
@@ -43,6 +45,10 @@
         return new ValueIterator();
     }
 
+    public Stream<E> stream() {
+        return StreamSupport.stream(spliterator(), false);
+    }
+
     @Thunk class ValueIterator implements Iterator<E> {
 
         private int mNextIndex = 0;
diff --git a/src/com/android/launcher3/util/LayoutImportExportHelper.kt b/src/com/android/launcher3/util/LayoutImportExportHelper.kt
index 0df9dae..8559f3b 100644
--- a/src/com/android/launcher3/util/LayoutImportExportHelper.kt
+++ b/src/com/android/launcher3/util/LayoutImportExportHelper.kt
@@ -56,7 +56,7 @@
 
         model.enqueueModelUpdateTask { _, dataModel, _ ->
             val builder = LauncherLayoutBuilder()
-            dataModel.workspaceItems.forEach { info ->
+            dataModel.itemsIdMap.forEach { info ->
                 val loc =
                     when (info.container) {
                         CONTAINER_DESKTOP ->
@@ -67,9 +67,6 @@
                     }
                 loc.addItem(context, info)
             }
-            dataModel.appWidgets.forEach { info ->
-                builder.atWorkspace(info.cellX, info.cellY, info.screenId).addItem(context, info)
-            }
 
             val layoutXml = builder.build()
             callback(layoutXml)
diff --git a/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt b/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
index 553d08c..15accbd 100644
--- a/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/folder/PreviewItemManagerTest.kt
@@ -17,11 +17,14 @@
 package com.android.launcher3.folder
 
 import android.R
-import android.graphics.Bitmap
 import android.os.Process
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherPrefs
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
+import com.android.launcher3.dagger.LauncherAppComponent
+import com.android.launcher3.dagger.LauncherAppSingleton
 import com.android.launcher3.graphics.PreloadIconDrawable
 import com.android.launcher3.graphics.ThemeManager
 import com.android.launcher3.icons.BitmapInfo
@@ -30,13 +33,14 @@
 import com.android.launcher3.icons.IconCache.ItemInfoUpdateReceiver
 import com.android.launcher3.icons.PlaceHolderIconDrawable
 import com.android.launcher3.icons.UserBadgeDrawable
-import com.android.launcher3.icons.mono.MonoThemedBitmap
 import com.android.launcher3.model.data.FolderInfo
 import com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_ARCHIVED
 import com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_INSTALL_SESSION_ACTIVE
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.util.ActivityContextWrapper
+import com.android.launcher3.util.AllModulesForTest
 import com.android.launcher3.util.Executors
+import com.android.launcher3.util.FakePrefsModule
 import com.android.launcher3.util.FlagOp
 import com.android.launcher3.util.LauncherLayoutBuilder
 import com.android.launcher3.util.LauncherModelHelper
@@ -44,10 +48,19 @@
 import com.android.launcher3.util.TestUtil
 import com.android.launcher3.util.UserIconInfo
 import com.google.common.truth.Truth.assertThat
+import dagger.Component
+import kotlin.annotation.AnnotationRetention.RUNTIME
+import kotlin.annotation.AnnotationTarget.FUNCTION
+import kotlin.annotation.AnnotationTarget.PROPERTY_GETTER
+import kotlin.annotation.AnnotationTarget.PROPERTY_SETTER
 import org.junit.After
 import org.junit.Before
+import org.junit.Rule
 import org.junit.Test
+import org.junit.rules.TestRule
+import org.junit.runner.Description
 import org.junit.runner.RunWith
+import org.junit.runners.model.Statement
 import org.mockito.kotlin.any
 import org.mockito.kotlin.argumentCaptor
 import org.mockito.kotlin.doReturn
@@ -61,6 +74,8 @@
 @RunWith(AndroidJUnit4::class)
 class PreviewItemManagerTest {
 
+    @get:Rule val theseStateRule = ThemeStateRule()
+
     private lateinit var previewItemManager: PreviewItemManager
     private lateinit var context: SandboxModelContext
     private lateinit var folderItems: ArrayList<WorkspaceItemInfo>
@@ -68,15 +83,14 @@
     private lateinit var folderIcon: FolderIcon
     private lateinit var iconCache: IconCache
 
-    private var defaultThemedIcons = false
-
-    private val themeManager: ThemeManager
-        get() = ThemeManager.INSTANCE.get(context)
-
     @Before
     fun setup() {
         modelHelper = LauncherModelHelper()
         context = modelHelper.sandboxContext
+        context.initDaggerComponent(DaggerPreviewItemManagerTestComponent.builder())
+        theseStateRule.themeState?.let {
+            LauncherPrefs.get(context).putSync(ThemeManager.THEMED_ICONS.to(it))
+        }
         folderIcon = FolderIcon(ActivityContextWrapper(context))
 
         val app = spy(LauncherAppState.getInstance(context))
@@ -99,27 +113,16 @@
             )
             .loadModelSync()
 
+        folderIcon.mInfo =
+            modelHelper.bgDataModel.itemsIdMap.find { it.itemType == ITEM_TYPE_FOLDER }
+                as FolderInfo
         // Use getAppContents() to "cast" contents to WorkspaceItemInfo so we can set bitmaps
-        folderItems = modelHelper.bgDataModel.collections.valueAt(0).getAppContents()
-        folderIcon.mInfo = modelHelper.bgDataModel.collections.valueAt(0) as FolderInfo
-        folderIcon.mInfo.getContents().addAll(folderItems)
-
-        // Set first icon to be themed.
-        folderItems[0].bitmap.themedBitmap =
-            MonoThemedBitmap(
-                folderItems[0].bitmap.icon,
-                Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888),
-            )
+        folderItems = folderIcon.mInfo.getAppContents()
 
         // Set second icon to be non-themed.
         folderItems[1].bitmap.themedBitmap = null
 
         // Set third icon to be themed with badge.
-        folderItems[2].bitmap.themedBitmap =
-            MonoThemedBitmap(
-                folderItems[2].bitmap.icon,
-                Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888),
-            )
         folderItems[2].bitmap =
             folderItems[2].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
 
@@ -127,20 +130,17 @@
         folderItems[3].bitmap =
             folderItems[3].bitmap.withFlags(profileFlagOp(UserIconInfo.TYPE_WORK))
         folderItems[3].bitmap.themedBitmap = null
-
-        defaultThemedIcons = themeManager.isMonoThemeEnabled
     }
 
     @After
     @Throws(Exception::class)
     fun tearDown() {
-        themeManager.isMonoThemeEnabled = defaultThemedIcons
         modelHelper.destroy()
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkThemedIconWithThemingOn_iconShouldBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[0])
@@ -149,8 +149,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(false)
     fun checkThemedIconWithThemingOff_iconShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = false
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[0])
@@ -159,8 +159,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkUnthemedIconWithThemingOn_iconShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[1])
@@ -169,8 +169,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(false)
     fun checkUnthemedIconWithThemingOff_iconShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = false
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[1])
@@ -179,8 +179,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkThemedIconWithBadgeWithThemingOn_iconAndBadgeShouldBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[2])
@@ -192,8 +192,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(true)
     fun checkUnthemedIconWithBadgeWithThemingOn_badgeShouldBeThemed() {
-        themeManager.isMonoThemeEnabled = true
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[3])
@@ -205,8 +205,8 @@
     }
 
     @Test
+    @MonoThemeEnabled(false)
     fun checkUnthemedIconWithBadgeWithThemingOff_iconAndBadgeShouldNotBeThemed() {
-        themeManager.isMonoThemeEnabled = false
         val drawingParams = PreviewItemDrawingParams(0f, 0f, 0f)
 
         previewItemManager.setDrawable(drawingParams, folderItems[3])
@@ -278,3 +278,28 @@
     private fun profileFlagOp(type: Int) =
         UserIconInfo(Process.myUserHandle(), type).applyBitmapInfoFlags(FlagOp.NO_OP)
 }
+
+class ThemeStateRule : TestRule {
+
+    var themeState: Boolean? = null
+
+    override fun apply(base: Statement, description: Description): Statement {
+        themeState = description.getAnnotation(MonoThemeEnabled::class.java)?.value
+        return base
+    }
+}
+
+// Annotation for tests that need to be run with quickstep enabled and disabled.
+@Retention(RUNTIME)
+@Target(FUNCTION, PROPERTY_GETTER, PROPERTY_SETTER)
+annotation class MonoThemeEnabled(val value: Boolean = false)
+
+@LauncherAppSingleton
+@Component(modules = [AllModulesForTest::class, FakePrefsModule::class])
+interface PreviewItemManagerTestComponent : LauncherAppComponent {
+
+    @Component.Builder
+    interface Builder : LauncherAppComponent.Builder {
+        override fun build(): PreviewItemManagerTestComponent
+    }
+}
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java b/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
index 1e2431f..0ae4d00 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
+++ b/tests/multivalentTests/src/com/android/launcher3/model/DefaultLayoutProviderTest.java
@@ -16,6 +16,8 @@
 
 package com.android.launcher3.model;
 
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_ACTIVITY;
 import static com.android.launcher3.util.LauncherModelHelper.TEST_PACKAGE;
 
@@ -42,6 +44,8 @@
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
+import java.util.List;
+
 /**
  * Tests for layout parser for remote layout
  */
@@ -63,14 +67,23 @@
         mModelHelper.destroy();
     }
 
+    private List<ItemInfo> getWorkspaceItems() {
+        return mModelHelper
+                .getBgDataModel()
+                .itemsIdMap
+                .stream()
+                .filter(i -> i.container == CONTAINER_DESKTOP || i.container == CONTAINER_HOTSEAT)
+                .toList();
+    }
+
     @Test
     public void testCustomProfileLoaded_with_icon_on_hotseat() throws Exception {
         writeLayoutAndLoad(new LauncherLayoutBuilder().atHotseat(0)
                 .putApp(TEST_PACKAGE, TEST_ACTIVITY));
 
         // Verify one item in hotseat
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.CONTAINER_HOTSEAT, info.container);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_APPLICATION, info.itemType);
     }
@@ -84,8 +97,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_FOLDER, info.itemType);
         assertEquals(3, ((FolderInfo) info).getContents().size());
     }
@@ -99,8 +112,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_FOLDER, info.itemType);
         assertEquals(3, ((FolderInfo) info).getContents().size());
         assertEquals("CustomFolder", info.title.toString());
@@ -124,8 +137,8 @@
                 .putWidget(pendingAppPkg, "PlaceholderWidget", 2, 2));
 
         // Verify widget
-        assertEquals(1, mModelHelper.getBgDataModel().appWidgets.size());
-        ItemInfo info = mModelHelper.getBgDataModel().appWidgets.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_APPWIDGET, info.itemType);
         assertEquals(2, info.spanX);
         assertEquals(2, info.spanY);
@@ -138,8 +151,8 @@
                 .putShortcut(TEST_PACKAGE, "shortcut2"));
 
         // Verify one item in hotseat
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        ItemInfo info = mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        ItemInfo info = getWorkspaceItems().get(0);
         assertEquals(LauncherSettings.Favorites.CONTAINER_HOTSEAT, info.container);
         assertEquals(LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT, info.itemType);
     }
@@ -154,8 +167,8 @@
                 .build());
 
         // Verify folder
-        assertEquals(1, mModelHelper.getBgDataModel().workspaceItems.size());
-        FolderInfo info = (FolderInfo) mModelHelper.getBgDataModel().workspaceItems.get(0);
+        assertEquals(1, getWorkspaceItems().size());
+        FolderInfo info = (FolderInfo) getWorkspaceItems().get(0);
         assertEquals(3, info.getContents().size());
 
         // Verify last icon
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
index e8f778f..f357487 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/FolderIconLoadTest.kt
@@ -18,8 +18,10 @@
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
 import com.android.launcher3.icons.BitmapInfo
 import com.android.launcher3.icons.waitForUpdateHandlerToFinish
+import com.android.launcher3.model.data.FolderInfo
 import com.android.launcher3.model.data.WorkspaceItemInfo
 import com.android.launcher3.util.Executors
 import com.android.launcher3.util.LauncherLayoutBuilder
@@ -149,11 +151,13 @@
         // Reload again with correct icon state
         app.model.forceReload()
         modelHelper.loadModelSync()
-        val collections = modelHelper.getBgDataModel().collections
-
-        assertThat(collections.size()).isEqualTo(1)
-        assertThat(collections.valueAt(0).getAppContents().size).isEqualTo(itemCount)
-        return collections.valueAt(0).getAppContents()
+        val collections =
+            modelHelper.bgDataModel.itemsIdMap
+                .filter { it.itemType == ITEM_TYPE_FOLDER }
+                .map { it as FolderInfo }
+        assertThat(collections.size).isEqualTo(1)
+        assertThat(collections[0].getAppContents().size).isEqualTo(itemCount)
+        return collections[0].getAppContents()
     }
 
     private fun verifyHighRes(items: ArrayList<WorkspaceItemInfo>, vararg indices: Int) {
diff --git a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
index d699eee..da87dfc 100644
--- a/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
+++ b/tests/multivalentTests/src/com/android/launcher3/model/WorkspaceItemProcessorTest.kt
@@ -494,8 +494,7 @@
     @Test
     fun `When processing Folder then create FolderInfo and mark restored`() {
         val actualFolderInfo = FolderInfo()
-        mockBgDataModel =
-            mock<BgDataModel>().apply { whenever(findOrMakeFolder(1)).thenReturn(actualFolderInfo) }
+        mockBgDataModel = mock<BgDataModel>()
         mockCursor =
             mock<LoaderCursor>().apply {
                 user = UserHandle(0)
@@ -509,6 +508,7 @@
                 whenever(getColumnIndex(Favorites.TITLE)).thenReturn(4)
                 whenever(getString(4)).thenReturn("title")
                 whenever(options).thenReturn(5)
+                whenever(findOrMakeFolder(eq(1), any())).thenReturn(actualFolderInfo)
             }
         val expectedFolderInfo =
             FolderInfo().apply {
@@ -600,7 +600,8 @@
 
         // Then
         val widgetInfoCaptor = ArgumentCaptor.forClass(LauncherAppWidgetInfo::class.java)
-        verify(mockCursor).checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel))
+        verify(mockCursor)
+            .checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel), anyOrNull())
         val actualWidgetInfo = widgetInfoCaptor.value
         with(actualWidgetInfo) {
             assertThat(providerName).isEqualTo(expectedWidgetInfo.providerName)
@@ -655,7 +656,7 @@
         itemProcessorUnderTest.processItem()
 
         // Then
-        verify(mockCursor).checkAndAddItem(any(), any())
+        verify(mockCursor).checkAndAddItem(any(), any(), anyOrNull())
     }
 
     @Test
diff --git a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
index d2229c4..f04688d 100644
--- a/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
+++ b/tests/src/com/android/launcher3/model/LoaderTaskTest.kt
@@ -19,6 +19,10 @@
 import com.android.launcher3.LauncherPrefs
 import com.android.launcher3.LauncherPrefs.Companion.IS_FIRST_LOAD_AFTER_RESTORE
 import com.android.launcher3.LauncherPrefs.Companion.RESTORE_DEVICE
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_DESKTOP
+import com.android.launcher3.LauncherSettings.Favorites.CONTAINER_HOTSEAT
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_APP_PAIR
+import com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_FOLDER
 import com.android.launcher3.icons.IconCache
 import com.android.launcher3.icons.cache.CachingLogic
 import com.android.launcher3.icons.cache.IconCacheUpdateHandler
@@ -155,9 +159,24 @@
                     widgetsFilterDataProvider,
                 )
                 .runSyncOnBackgroundThread()
-            Truth.assertThat(workspaceItems.size).isAtLeast(25)
-            Truth.assertThat(appWidgets.size).isAtLeast(7)
-            Truth.assertThat(collections.size()).isAtLeast(8)
+            Truth.assertThat(
+                    itemsIdMap
+                        .filter {
+                            it.container == CONTAINER_DESKTOP || it.container == CONTAINER_HOTSEAT
+                        }
+                        .size
+                )
+                .isAtLeast(32)
+            Truth.assertThat(itemsIdMap.filter { ModelUtils.WIDGET_FILTER.test(it) }.size)
+                .isAtLeast(7)
+            Truth.assertThat(
+                    itemsIdMap
+                        .filter {
+                            it.itemType == ITEM_TYPE_FOLDER || it.itemType == ITEM_TYPE_APP_PAIR
+                        }
+                        .size
+                )
+                .isAtLeast(8)
             Truth.assertThat(itemsIdMap.size()).isAtLeast(40)
             Truth.assertThat(widgetsModel.defaultWidgetsFilter).isNotNull()
         }
diff --git a/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt b/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt
index d553f47..8db049c 100644
--- a/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt
+++ b/tests/src/com/android/launcher3/model/WorkspaceItemProcessorExtraTest.kt
@@ -58,6 +58,7 @@
 import org.mockito.Mockito.RETURNS_DEEP_STUBS
 import org.mockito.Mockito.verify
 import org.mockito.kotlin.any
+import org.mockito.kotlin.anyOrNull
 import org.mockito.kotlin.doAnswer
 import org.mockito.kotlin.eq
 import org.mockito.kotlin.mock
@@ -220,7 +221,8 @@
             )
             .commit()
         val widgetInfoCaptor = ArgumentCaptor.forClass(LauncherAppWidgetInfo::class.java)
-        verify(mockCursor).checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel))
+        verify(mockCursor)
+            .checkAndAddItem(widgetInfoCaptor.capture(), eq(mockBgDataModel), anyOrNull())
         val actualWidgetInfo = widgetInfoCaptor.value
         with(actualWidgetInfo) {
             assertThat(providerName).isEqualTo(expectedWidgetInfo.providerName)
@@ -271,7 +273,7 @@
         itemProcessorUnderTest.processItem()
 
         // Then
-        verify(mockCursor).checkAndAddItem(any(), any())
+        verify(mockCursor).checkAndAddItem(any(), any(), anyOrNull())
     }
 
     private fun createWorkspaceItemProcessorUnderTest(